Add basic application rate limits to appview (#1902)

* cache did docs in redis * drop table * expire from redis * fix tests * add cache class * update api * refactor * filter negative labels * fix up dev-env * refactor did cache to use new redis cache class * tidy * ensure caching negatives * redis cache tests * remove timeout on did cache * fix ns in test * rename driver * add timeout & fail open * add test for timeout & fail open * add basic rate limits * tidy * small pr feedback * refactor caches * bugfixg * test for caching negative values * little more to cache * wire up cache cfg * switch from redis scratch to redis * fix build issues * use different redis clients for tests * fix test * fix flaky test * fix build issue * use separate db for redis cache
bluesky-social · Dec 5, 2023 · af3e33a · af3e33a
1 parent 3110a2a
commit af3e33a
Show file tree

Hide file tree

Showing 5 changed files with 72 additions and 11 deletions.
diff --git a/packages/bsky/src/config.ts b/packages/bsky/src/config.ts
@@ -36,6 +36,9 @@ export interface ServerConfigValues {
   moderatorPassword?: string
   triagePassword?: string
   moderationPushUrl?: string
+  rateLimitsEnabled: boolean
+  rateLimitsBypassKey?: string
+  rateLimitsBypassIps?: string[]
 }
 
 export class ServerConfig {
@@ -115,6 +118,12 @@ export class ServerConfig {
       overrides?.moderationPushUrl ||
       process.env.MODERATION_PUSH_URL ||
       undefined
+    const rateLimitsEnabled = process.env.RATE_LIMITS_ENABLED === 'true'
+    const rateLimitsBypassKey = process.env.RATE_LIMITS_BYPASS_KEY
+    const rateLimitsBypassIps = process.env.RATE_LIMITS_BYPASS_IPS
+      ? process.env.RATE_LIMITS_BYPASS_IPS.split(',')
+      : undefined
+
     return new ServerConfig({
       version,
       debugMode,
@@ -144,6 +153,9 @@ export class ServerConfig {
       moderatorPassword,
       triagePassword,
       moderationPushUrl,
+      rateLimitsEnabled,
+      rateLimitsBypassKey,
+      rateLimitsBypassIps,
       ...stripUndefineds(overrides ?? {}),
     })
   }
@@ -272,6 +284,18 @@ export class ServerConfig {
   get moderationPushUrl() {
     return this.cfg.moderationPushUrl
   }
+
+  get rateLimitsEnabled() {
+    return this.cfg.rateLimitsEnabled
+  }
+
+  get rateLimitsBypassKey() {
+    return this.cfg.rateLimitsBypassKey
+  }
+
+  get rateLimitsBypassIps() {
+    return this.cfg.rateLimitsBypassIps
+  }
 }
 
 function getTagIdxs(str?: string): number[] {

diff --git a/packages/bsky/src/index.ts b/packages/bsky/src/index.ts
@@ -6,6 +6,12 @@ import { createHttpTerminator, HttpTerminator } from 'http-terminator'
 import cors from 'cors'
 import compression from 'compression'
 import { IdResolver } from '@atproto/identity'
+import {
+  RateLimiter,
+  RateLimiterOpts,
+  Options as XrpcServerOptions,
+} from '@atproto/xrpc-server'
+import { MINUTE } from '@atproto/common'
 import API, { health, wellKnown, blobResolver } from './api'
 import { DatabaseCoordinator } from './db'
 import * as error from './error'
@@ -135,14 +141,41 @@ export class BskyAppView {
       notifServer,
     })
 
-    let server = createServer({
+    const xrpcOpts: XrpcServerOptions = {
       validateResponse: config.debugMode,
       payload: {
         jsonLimit: 100 * 1024, // 100kb
         textLimit: 100 * 1024, // 100kb
         blobLimit: 5 * 1024 * 1024, // 5mb
       },
-    })
+    }
+    if (config.rateLimitsEnabled) {
+      const rlCreator = (opts: RateLimiterOpts) =>
+        RateLimiter.redis(redis.driver, {
+          bypassSecret: config.rateLimitsBypassKey,
+          bypassIps: config.rateLimitsBypassIps,
+          ...opts,
+        })
+      xrpcOpts['rateLimits'] = {
+        creator: rlCreator,
+        global: [
+          {
+            name: 'global-unauthed-ip',
+            durationMs: 5 * MINUTE,
+            points: 3000,
+            calcKey: (ctx) => (ctx.auth ? null : ctx.req.ip),
+          },
+          {
+            name: 'global-authed-did',
+            durationMs: 5 * MINUTE,
+            points: 3000,
+            calcKey: (ctx) => ctx.auth?.credentials?.did ?? null,
+          },
+        ],
+      }
+    }
+
+    let server = createServer(xrpcOpts)
 
     server = API(server, ctx)
 

diff --git a/packages/dev-env/src/bsky.ts b/packages/dev-env/src/bsky.ts
@@ -51,6 +51,7 @@ export class TestBsky {
       triagePassword: TRIAGE_PASSWORD,
       labelerDid: 'did:example:labeler',
       feedGenDid: 'did:example:feedGen',
+      rateLimitsEnabled: false,
     })
 
     // shared across server, ingester, and indexer in order to share pool, avoid too many pg connections.

diff --git a/packages/xrpc-server/src/rate-limiter.ts b/packages/xrpc-server/src/rate-limiter.ts
@@ -75,6 +75,9 @@ export class RateLimiter implements RateLimiterI {
       return null
     }
     const key = opts?.calcKey ? opts.calcKey(ctx) : this.calcKey(ctx)
+    if (key === null) {
+      return null
+    }
     const points = opts?.calcPoints
       ? opts.calcPoints(ctx)
       : this.calcPoints(ctx)

diff --git a/packages/xrpc-server/src/types.ts b/packages/xrpc-server/src/types.ts
@@ -85,7 +85,7 @@ export type StreamAuthVerifier = (ctx: {
   req: IncomingMessage
 }) => Promise<AuthOutput> | AuthOutput
 
-export type CalcKeyFn = (ctx: XRPCReqContext) => string
+export type CalcKeyFn = (ctx: XRPCReqContext) => string | null
 export type CalcPointsFn = (ctx: XRPCReqContext) => number
 
 export interface RateLimiterI {
@@ -101,29 +101,29 @@ export type RateLimiterCreator = (opts: {
   keyPrefix: string
   durationMs: number
   points: number
-  calcKey?: (ctx: XRPCReqContext) => string
-  calcPoints?: (ctx: XRPCReqContext) => number
+  calcKey?: CalcKeyFn
+  calcPoints?: CalcPointsFn
 }) => RateLimiterI
 
 export type ServerRateLimitDescription = {
   name: string
   durationMs: number
   points: number
-  calcKey?: (ctx: XRPCReqContext) => string
-  calcPoints?: (ctx: XRPCReqContext) => number
+  calcKey?: CalcKeyFn
+  calcPoints?: CalcPointsFn
 }
 
 export type SharedRateLimitOpts = {
   name: string
-  calcKey?: (ctx: XRPCReqContext) => string
-  calcPoints?: (ctx: XRPCReqContext) => number
+  calcKey?: CalcKeyFn
+  calcPoints?: CalcPointsFn
 }
 
 export type RouteRateLimitOpts = {
   durationMs: number
   points: number
-  calcKey?: (ctx: XRPCReqContext) => string
-  calcPoints?: (ctx: XRPCReqContext) => number
+  calcKey?: CalcKeyFn
+  calcPoints?: CalcPointsFn
 }
 
 export type HandlerRateLimitOpts = SharedRateLimitOpts | RouteRateLimitOpts