Merge pull request #142 from dfinity/igor/rl-bypass-token

Add rate limiter bypass token

Author: blind-oracle

Cargo.lock

@@ -48,7 +48,7 @@ http = "1.3.1"
 http-body = "1.0.1"
 http-body-util = "0.1.2"
 humantime = "2.2.0"
-ic-bn-lib = { git = "https://github.com/dfinity/ic-bn-lib", rev = "3b99c44fae99d09cc6d78f0879d87dce0eb8768e", features = [
+ic-bn-lib = { git = "https://github.com/dfinity/ic-bn-lib", rev = "7f15e91888f76a6bfff9a59d478fdd11d76a5ff6", features = [
     "vector",
     "cert-providers",
     "clients-hyper",

Cargo.toml

@@ -83,6 +83,9 @@ pub struct Cli {
     #[command(flatten, next_help_heading = "CORS")]
     pub cors: Cors,
 
+    #[command(flatten, next_help_heading = "Rate limiting")]
+    pub rate_limit: RateLimit,
+
     #[command(flatten, next_help_heading = "Cache")]
     pub cache: CacheConfig,
 
@@ -557,6 +560,13 @@ pub struct Cors {
     pub cors_invalid_canisters_ttl: Duration,
 }
 
+#[derive(Args)]
+pub struct RateLimit {
+    /// Bypass token for rate-limiter that should be sent in `x-ratelimit-bypass-token` header
+    #[clap(env, long)]
+    pub rate_limit_bypass_token: Option<String>,
+}
+
 #[cfg(test)]
 mod test {
     use super::*;

src/cli.rs

@@ -489,7 +489,12 @@ pub async fn setup_router(
                     ])),
             )
             .route("/registrations", post(proxy::issuer_proxy).layer(cors_post))
-            .layer(rate_limiter::layer_by_ip(1, 2, RateLimitCause::Normal)?)
+            .layer(rate_limiter::layer_by_ip(
+                1,
+                2,
+                RateLimitCause::Normal,
+                cli.rate_limit.rate_limit_bypass_token.clone(),
+            )?)
             .with_state(state);
 
         Some(router)
@@ -590,7 +595,12 @@ pub async fn setup_router(
                     )
                     .context("unable to init SEV-SNP")?,
                 )
-                .layer(rate_limiter::layer_global(50, 100, RateLimitCause::Normal)?),
+                .layer(rate_limiter::layer_global(
+                    50,
+                    100,
+                    RateLimitCause::Normal,
+                    cli.rate_limit.rate_limit_bypass_token.clone(),
+                )?),
         );
 
         router = router.merge(router_sev_snp)