feat: rate limiting

.github/workflows/sub-ci.yml

@@ -52,6 +52,10 @@ jobs:
           --health-interval 10s
           --health-timeout 5s
           --health-retries 5
+      redis:
+        image: redis:7-alpine
+        ports:
+          - 6379:6379
     steps:
       - name: Checkout
         uses: actions/checkout@v3

Cargo.toml

@@ -88,6 +88,7 @@ once_cell = "1.18.0"
 lazy_static = "1.4.0"
 rmp-serde = "1.1.1"
 deadpool-redis = "0.12.0"
+redis = { version = "0.23.3", default-features = false, features = ["script"] }
 rand_chacha = "0.3.1"
 sqlx = { version = "0.7.1", features = ["runtime-tokio-native-tls", "postgres", "chrono", "uuid"] }
 wiremock = "0.5.19"

docker-compose.storage.yml

@@ -7,6 +7,11 @@ services:
 #     ports:
 #       - "3001:16686"
 
+  redis:
+    image: redis:7-alpine
+    ports:
+    - "6379:6379"
+
   postgres:
     image: postgres:16
     environment:

src/auth.rs

@@ -425,7 +425,7 @@ pub struct Authorization {
     pub domain: String,
 }
 
-#[derive(Serialize, Deserialize, Debug)]
+#[derive(Serialize, Deserialize, Debug, PartialEq)]
 pub enum AuthorizedApp {
     Limited(String),
     Unlimited,
@@ -490,21 +490,7 @@ pub async fn verify_identity(
     let app = {
         let statement = cacao.p.statement.ok_or(AuthError::CacaoStatementMissing)?;
         info!("CACAO statement: {statement}");
-        if statement.contains("DAPP")
-            || statement == STATEMENT_THIS_DOMAIN_IDENTITY
-            || statement == STATEMENT_THIS_DOMAIN
-        {
-            AuthorizedApp::Limited(cacao.p.domain.clone())
-        } else if statement.contains("WALLET")
-            || statement == STATEMENT
-            || statement == STATEMENT_ALL_DOMAINS_IDENTITY
-            || statement == STATEMENT_ALL_DOMAINS_OLD
-            || statement == STATEMENT_ALL_DOMAINS
-        {
-            AuthorizedApp::Unlimited
-        } else {
-            return Err(AuthError::CacaoStatementInvalid)?;
-        }
+        parse_cacao_statement(&statement, &cacao.p.domain)?
     };
 
     if cacao.p.iss != sub {
@@ -541,6 +527,24 @@ pub async fn verify_identity(
     })
 }
 
+fn parse_cacao_statement(statement: &str, domain: &str) -> Result<AuthorizedApp> {
+    if statement.contains("DAPP")
+        || statement == STATEMENT_THIS_DOMAIN_IDENTITY
+        || statement == STATEMENT_THIS_DOMAIN
+    {
+        Ok(AuthorizedApp::Limited(domain.to_owned()))
+    } else if statement.contains("WALLET")
+        || statement == STATEMENT
+        || statement == STATEMENT_ALL_DOMAINS_IDENTITY
+        || statement == STATEMENT_ALL_DOMAINS_OLD
+        || statement == STATEMENT_ALL_DOMAINS
+    {
+        Ok(AuthorizedApp::Unlimited)
+    } else {
+        return Err(AuthError::CacaoStatementInvalid)?;
+    }
+}
+
 #[derive(Debug, Clone, Serialize, Deserialize)]
 struct KeyServerResponse {
     status: String,
@@ -568,3 +572,69 @@ pub fn encode_subscribe_private_key(subscribe_key: &StaticSecret) -> String {
 pub fn encode_subscribe_public_key(subscribe_key: &StaticSecret) -> String {
     hex::encode(PublicKey::from(subscribe_key))
 }
+
+#[cfg(test)]
+mod test {
+    use super::*;
+
+    #[test]
+    fn notify_all_domains() {
+        assert_eq!(
+            parse_cacao_statement(STATEMENT_ALL_DOMAINS, "app.example.com").unwrap(),
+            AuthorizedApp::Unlimited
+        );
+    }
+
+    #[test]
+    fn notify_all_domains_old() {
+        assert_eq!(
+            parse_cacao_statement(STATEMENT_ALL_DOMAINS_OLD, "app.example.com").unwrap(),
+            AuthorizedApp::Unlimited
+        );
+    }
+
+    #[test]
+    fn notify_this_domain() {
+        assert_eq!(
+            parse_cacao_statement(STATEMENT_THIS_DOMAIN, "app.example.com").unwrap(),
+            AuthorizedApp::Limited("app.example.com".to_owned())
+        );
+    }
+
+    #[test]
+    fn notify_all_domains_identity() {
+        assert_eq!(
+            parse_cacao_statement(STATEMENT_ALL_DOMAINS_IDENTITY, "app.example.com").unwrap(),
+            AuthorizedApp::Unlimited
+        );
+    }
+
+    #[test]
+    fn notify_this_domain_identity() {
+        assert_eq!(
+            parse_cacao_statement(STATEMENT_THIS_DOMAIN_IDENTITY, "app.example.com").unwrap(),
+            AuthorizedApp::Limited("app.example.com".to_owned())
+        );
+    }
+
+    #[test]
+    fn old_siwe_compatible() {
+        assert_eq!(
+            parse_cacao_statement(STATEMENT, "app.example.com").unwrap(),
+            AuthorizedApp::Unlimited
+        );
+    }
+
+    #[test]
+    fn old_old_siwe_compatible() {
+        assert_eq!(
+            parse_cacao_statement(
+                "I further authorize this DAPP to send and receive messages on my behalf for \
+    this domain using my WalletConnect identity.",
+                "app.example.com"
+            )
+            .unwrap(),
+            AuthorizedApp::Limited("app.example.com".to_owned())
+        );
+    }
+}

src/config/local.rs

@@ -25,6 +25,8 @@ pub struct LocalConfiguration {
     pub postgres_url: String,
     #[serde(default = "default_postgres_max_connections")]
     pub postgres_max_connections: u32,
+    #[serde(default = "default_redis_url")]
+    pub redis_url: String,
     #[serde(default = "default_keypair_seed")]
     pub keypair_seed: String,
     #[serde(default = "default_relay_url")]
@@ -49,6 +51,10 @@ pub fn default_postgres_url() -> String {
     "postgres://postgres:postgres@localhost:5432/postgres".to_owned()
 }
 
+pub fn default_redis_url() -> String {
+    "redis://localhost:6379/0".to_owned()
+}
+
 pub fn default_postgres_max_connections() -> u32 {
     10
 }

src/error.rs

@@ -183,6 +183,15 @@ pub enum Error {
 
     #[error("App domain in-use by another project")]
     AppDomainInUseByAnotherProject,
+
+    #[error("Redis pool error: {0}")]
+    RedisPoolError(#[from] deadpool_redis::PoolError),
+
+    #[error("Redis error: {0}")]
+    RedisError(#[from] redis::RedisError),
+
+    #[error("Rate limit exceeded. Try again at {0}")]
+    TooManyRequests(u64),
 }
 
 impl IntoResponse for Error {
@@ -201,6 +210,9 @@ impl IntoResponse for Error {
             Self::AppDomainInUseByAnotherProject => {
                 (StatusCode::CONFLICT, self.to_string()).into_response()
             }
+            Self::TooManyRequests(_) => {
+                (StatusCode::TOO_MANY_REQUESTS, self.to_string()).into_response()
+            }
             error => {
                 warn!("Error does not have response clause: {:?}", error);
                 (StatusCode::INTERNAL_SERVER_ERROR, "Internal server error.").into_response()

src/lib.rs

@@ -3,6 +3,7 @@ use {
         config::Configuration,
         error::Result,
         metrics::Metrics,
+        registry::storage::redis::Redis,
         relay_client_helpers::create_http_client,
         services::{
             private_http_server, public_http_server, publisher_service, watcher_expiration_job,
@@ -34,6 +35,7 @@ pub mod model;
 mod notify_keys;
 pub mod notify_message;
 pub mod publish_relay_message;
+pub mod rate_limit;
 pub mod registry;
 pub mod relay_client_helpers;
 pub mod services;
@@ -81,10 +83,19 @@ pub async fn bootstrap(mut shutdown: broadcast::Receiver<()>, config: Configurat
 
     let metrics = Some(Metrics::default());
 
+    let redis = if let Some(redis_addr) = &config.auth_redis_addr() {
+        Some(Arc::new(Redis::new(
+            redis_addr,
+            config.redis_pool_size as usize,
+        )?))
+    } else {
+        None
+    };
+
     let registry = Arc::new(registry::Registry::new(
         config.registry_url.clone(),
         &config.registry_auth_token,
-        &config,
+        redis.clone(),
         metrics.clone(),
     )?);
 
@@ -97,6 +108,7 @@ pub async fn bootstrap(mut shutdown: broadcast::Receiver<()>, config: Configurat
         relay_ws_client.clone(),
         relay_http_client.clone(),
         metrics.clone(),
+        redis,
         registry,
     )?);
 

src/rate_limit/mod.rs

@@ -0,0 +1,2 @@
+mod token_bucket;
+pub use token_bucket::*;