Allow bypass on ratelimit ip (bluesky-social#1613)

allow bypass on ratelimit ip
mloar · Sep 25, 2023 · 370b123 · 370b123
1 parent 3ac7534
commit 370b123
Show file tree

Hide file tree

Showing 3 changed files with 27 additions and 4 deletions.
diff --git a/packages/pds/src/config.ts b/packages/pds/src/config.ts
@@ -45,6 +45,7 @@ export interface ServerConfigValues {
 
   rateLimitsEnabled: boolean
   rateLimitBypassKey?: string
+  rateLimitBypassIps?: string[]
   redisScratchAddress?: string
   redisScratchPassword?: string
 
@@ -163,6 +164,15 @@ export class ServerConfig {
 
     const rateLimitsEnabled = process.env.RATE_LIMITS_ENABLED === 'true'
     const rateLimitBypassKey = nonemptyString(process.env.RATE_LIMIT_BYPASS_KEY)
+    const rateLimitBypassIpsStr = nonemptyString(
+      process.env.RATE_LIMIT_BYPASS_IPS,
+    )
+    const rateLimitBypassIps = rateLimitBypassIpsStr
+      ? rateLimitBypassIpsStr.split(',').map((ipOrCidr) => {
+          const ip = ipOrCidr.split('/')[0]
+          return ip.trim()
+        })
+      : undefined
     const redisScratchAddress = nonemptyString(
       process.env.REDIS_SCRATCH_ADDRESS,
     )
@@ -266,6 +276,7 @@ export class ServerConfig {
       blobCacheLocation,
       rateLimitsEnabled,
       rateLimitBypassKey,
+      rateLimitBypassIps,
       redisScratchAddress,
       redisScratchPassword,
       appUrlPasswordReset,
@@ -454,6 +465,10 @@ export class ServerConfig {
     return this.cfg.rateLimitBypassKey
   }
 
+  get rateLimitBypassIps() {
+    return this.cfg.rateLimitBypassIps
+  }
+
   get redisScratchAddress() {
     return this.cfg.redisScratchAddress
   }

diff --git a/packages/pds/src/index.ts b/packages/pds/src/index.ts
@@ -268,12 +268,14 @@ export class PDS {
         rlCreator = (opts: RateLimiterOpts) =>
           RateLimiter.redis(redisScratch, {
             bypassSecret: config.rateLimitBypassKey,
+            bypassIps: config.rateLimitBypassIps,
             ...opts,
           })
       } else {
         rlCreator = (opts: RateLimiterOpts) =>
           RateLimiter.memory({
             bypassSecret: config.rateLimitBypassKey,
+            bypassIps: config.rateLimitBypassIps,
             ...opts,
           })
       }

diff --git a/packages/xrpc-server/src/rate-limiter.ts b/packages/xrpc-server/src/rate-limiter.ts
@@ -20,21 +20,24 @@ export type RateLimiterOpts = {
   durationMs: number
   points: number
   bypassSecret?: string
+  bypassIps?: string[]
   calcKey?: CalcKeyFn
   calcPoints?: CalcPointsFn
   failClosed?: boolean
 }
 
 export class RateLimiter implements RateLimiterI {
   public limiter: RateLimiterAbstract
-  private byPassSecret?: string
+  private bypassSecret?: string
+  private bypassIps?: string[]
   private failClosed?: boolean
   public calcKey: CalcKeyFn
   public calcPoints: CalcPointsFn
 
   constructor(limiter: RateLimiterAbstract, opts: RateLimiterOpts) {
     this.limiter = limiter
-    this.byPassSecret = opts.bypassSecret
+    this.bypassSecret = opts.bypassSecret
+    this.bypassIps = opts.bypassIps
     this.calcKey = opts.calcKey ?? defaultKey
     this.calcPoints = opts.calcPoints ?? defaultPoints
   }
@@ -63,11 +66,14 @@ export class RateLimiter implements RateLimiterI {
     opts?: { calcKey?: CalcKeyFn; calcPoints?: CalcPointsFn },
   ): Promise<RateLimiterStatus | RateLimitExceededError | null> {
     if (
-      this.byPassSecret &&
-      ctx.req.header('x-ratelimit-bypass') === this.byPassSecret
+      this.bypassSecret &&
+      ctx.req.header('x-ratelimit-bypass') === this.bypassSecret
     ) {
       return null
     }
+    if (this.bypassIps && this.bypassIps.includes(ctx.req.ip)) {
+      return null
+    }
     const key = opts?.calcKey ? opts.calcKey(ctx) : this.calcKey(ctx)
     const points = opts?.calcPoints
       ? opts.calcPoints(ctx)