@atproto/xrpc-server 0.7.19 → 0.8.0

package/src/rate-limiter.ts

@@ -10,18 +10,19 @@ import {
   CalcPointsFn,
   RateLimitExceededError,
   RateLimiterConsume,
+  RateLimiterConsumeOptions,
   RateLimiterI,
   RateLimiterReset,
+  RateLimiterResetOptions,
   RateLimiterStatus,
   XRPCReqContext,
 } from './types'
+import { setHeaders } from './util'
 
 export type RateLimiterOpts = {
   keyPrefix: string
   durationMs: number
   points: number
-  bypassSecret?: string
-  bypassIps?: string[]
   calcKey?: CalcKeyFn
   calcPoints?: CalcPointsFn
   failClosed?: boolean
@@ -29,16 +30,13 @@ export type RateLimiterOpts = {
 
 export class RateLimiter implements RateLimiterI {
   public limiter: RateLimiterAbstract
-  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.bypassIps = opts.bypassIps
     this.calcKey = opts.calcKey ?? defaultKey
     this.calcPoints = opts.calcPoints ?? defaultPoints
   }
@@ -64,17 +62,8 @@ export class RateLimiter implements RateLimiterI {
 
   async consume(
     ctx: XRPCReqContext,
-    opts?: { calcKey?: CalcKeyFn; calcPoints?: CalcPointsFn },
+    opts?: RateLimiterConsumeOptions,
   ): Promise<RateLimiterStatus | RateLimitExceededError | null> {
-    if (
-      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)
     if (key === null) {
       return null
@@ -113,7 +102,7 @@ export class RateLimiter implements RateLimiterI {
 
   async reset(
     ctx: XRPCReqContext,
-    opts?: { calcKey?: CalcKeyFn },
+    opts?: RateLimiterResetOptions,
   ): Promise<void> {
     const key = opts?.calcKey ? opts.calcKey(ctx) : this.calcKey(ctx)
     if (key === null) {
@@ -142,46 +131,72 @@ export const formatLimiterStatus = (
   }
 }
 
-export const consumeMany = async (
-  ctx: XRPCReqContext,
-  fns: RateLimiterConsume[],
-): Promise<RateLimiterStatus | RateLimitExceededError | null> => {
-  if (fns.length === 0) return null
-  const results = await Promise.all(fns.map((fn) => fn(ctx)))
-  const tightestLimit = getTightestLimit(results)
-  if (tightestLimit === null) {
-    return null
-  } else if (tightestLimit instanceof RateLimitExceededError) {
-    setResHeaders(ctx, tightestLimit.status)
-    return tightestLimit
-  } else {
-    setResHeaders(ctx, tightestLimit)
-    return tightestLimit
-  }
+export type WrappedRateLimiterOptions = {
+  calcKey?: CalcKeyFn
+  calcPoints?: CalcPointsFn
 }
 
-export const resetMany = async (
-  ctx: XRPCReqContext,
-  fns: RateLimiterReset[],
-): Promise<void> => {
-  if (fns.length === 0) return
-  await Promise.all(fns.map((fn) => fn(ctx)))
+/**
+ * Wraps a {@link RateLimiterI} instance with custom key and points calculation
+ * functions.
+ */
+export class WrappedRateLimiter implements RateLimiterI {
+  private constructor(
+    private readonly rateLimiter: RateLimiterI,
+    private readonly options: Readonly<WrappedRateLimiterOptions>,
+  ) {}
+
+  async consume(ctx: XRPCReqContext, opts?: RateLimiterConsumeOptions) {
+    return this.rateLimiter.consume(ctx, {
+      calcKey: opts?.calcKey ?? this.options.calcKey,
+      calcPoints: opts?.calcPoints ?? this.options.calcPoints,
+    })
+  }
+
+  async reset(ctx: XRPCReqContext, opts?: RateLimiterResetOptions) {
+    return this.rateLimiter.reset(ctx, {
+      calcKey: opts?.calcKey ?? this.options.calcKey,
+    })
+  }
+
+  static from(
+    rateLimiter: RateLimiterI,
+    { calcKey, calcPoints }: WrappedRateLimiterOptions = {},
+  ): RateLimiterI {
+    if (!calcKey && !calcPoints) return rateLimiter
+    return new WrappedRateLimiter(rateLimiter, { calcKey, calcPoints })
+  }
 }
 
-export const setResHeaders = (
-  ctx: XRPCReqContext,
-  status: RateLimiterStatus,
-) => {
-  ctx.res.setHeader('RateLimit-Limit', status.limit)
-  ctx.res.setHeader('RateLimit-Remaining', status.remainingPoints)
-  ctx.res.setHeader(
-    'RateLimit-Reset',
-    Math.floor((Date.now() + status.msBeforeNext) / 1000),
-  )
-  ctx.res.setHeader('RateLimit-Policy', `${status.limit};w=${status.duration}`)
+/**
+ * Combines multiple rate limiters into one.
+ *
+ * The combined rate limiter will return the tightest (most restrictive) of all
+ * the provided rate limiters.
+ */
+export class CombinedRateLimiter implements RateLimiterI {
+  private constructor(private readonly rateLimiters: readonly RateLimiterI[]) {}
+
+  async consume(ctx: XRPCReqContext, opts?: RateLimiterConsumeOptions) {
+    const promises: ReturnType<RateLimiterConsume>[] = []
+    for (const rl of this.rateLimiters) promises.push(rl.consume(ctx, opts))
+    return Promise.all(promises).then(getTightestLimit)
+  }
+
+  async reset(ctx: XRPCReqContext, opts?: RateLimiterResetOptions) {
+    const promises: ReturnType<RateLimiterReset>[] = []
+    for (const rl of this.rateLimiters) promises.push(rl.reset(ctx, opts))
+    await Promise.all(promises)
+  }
+
+  static from(rateLimiters: readonly RateLimiterI[]): RateLimiterI | undefined {
+    if (rateLimiters.length === 0) return undefined
+    if (rateLimiters.length === 1) return rateLimiters[0]
+    return new CombinedRateLimiter(rateLimiters)
+  }
 }
 
-export const getTightestLimit = (
+const getTightestLimit = (
   resps: (RateLimiterStatus | RateLimitExceededError | null)[],
 ): RateLimiterStatus | RateLimitExceededError | null => {
   let lowest: RateLimiterStatus | null = null
@@ -195,6 +210,65 @@ export const getTightestLimit = (
   return lowest
 }
 
+export type RouteRateLimiterOptions = {
+  bypass?: (ctx: XRPCReqContext) => boolean
+}
+
+/**
+ * Wraps a {@link RateLimiterI} interface into a class that will apply the
+ * appropriate headers to the response if a limit is exceeded.
+ */
+export class RouteRateLimiter implements RateLimiterI {
+  constructor(
+    private readonly rateLimiter: RateLimiterI,
+    private readonly options: Readonly<RouteRateLimiterOptions> = {},
+  ) {}
+
+  async handle(ctx: XRPCReqContext): Promise<RateLimiterStatus | null> {
+    const { bypass } = this.options
+    if (bypass && bypass(ctx)) {
+      return null
+    }
+
+    const result = await this.consume(ctx)
+    if (result instanceof RateLimitExceededError) {
+      setStatusHeaders(ctx, result.status)
+      throw result
+    } else if (result != null) {
+      setStatusHeaders(ctx, result)
+    }
+
+    return result
+  }
+
+  async consume(...args: Parameters<RateLimiterConsume>) {
+    return this.rateLimiter.consume(...args)
+  }
+
+  async reset(...args: Parameters<RateLimiterReset>) {
+    return this.rateLimiter.reset(...args)
+  }
+
+  static from(
+    rateLimiters: readonly RateLimiterI[],
+    { bypass }: RouteRateLimiterOptions = {},
+  ): RouteRateLimiter | undefined {
+    const rateLimiter = CombinedRateLimiter.from(rateLimiters)
+    if (!rateLimiter) return undefined
+
+    return new RouteRateLimiter(rateLimiter, { bypass })
+  }
+}
+
+function setStatusHeaders(ctx: XRPCReqContext, status: RateLimiterStatus) {
+  setHeaders(ctx.res, {
+    'RateLimit-Limit': status.limit,
+    'RateLimit-Reset': Math.floor((Date.now() + status.msBeforeNext) / 1000),
+    'RateLimit-Remaining': status.remainingPoints,
+    'RateLimit-Policy': `${status.limit};w=${status.duration}`,
+  })
+}
+
 // when using a proxy, ensure headers are getting forwarded correctly: `app.set('trust proxy', true)`
 // https://expressjs.com/en/guide/behind-proxies.html
 const defaultKey: CalcKeyFn = (ctx: XRPCReqContext) => ctx.req.ip

package/src/server.ts

@@ -1,3 +1,4 @@
+import assert from 'node:assert'
 import { Readable } from 'node:stream'
 import { pipeline } from 'node:stream/promises'
 import express, {
@@ -22,19 +23,17 @@ import {
   lexToJson,
 } from '@atproto/lexicon'
 import log, { LOGGER_NAME } from './logger'
-import { consumeMany, resetMany } from './rate-limiter'
+import { RouteRateLimiter, WrappedRateLimiter } from './rate-limiter'
 import { ErrorFrame, Frame, MessageFrame, XrpcStreamServer } from './stream'
 import {
   AuthVerifier,
   HandlerAuth,
-  HandlerPipeThrough,
   HandlerSuccess,
   InternalServerError,
   InvalidRequestError,
   MethodNotImplementedError,
   Options,
   Params,
-  RateLimitExceededError,
   RateLimiterI,
   XRPCError,
   XRPCHandler,
@@ -48,8 +47,10 @@ import {
   isShared,
 } from './types'
 import {
+  asArray,
   decodeQueryParams,
   getQueryParams,
+  setHeaders,
   validateInput,
   validateOutput,
 } from './util'
@@ -65,9 +66,8 @@ export class Server {
   lex = new Lexicons()
   options: Options
   middleware: Record<'json' | 'text', RequestHandler>
-  globalRateLimiters: RateLimiterI[]
-  sharedRateLimiters: Record<string, RateLimiterI>
-  routeRateLimiters: Record<string, RateLimiterI[]>
+  globalRateLimiter?: RouteRateLimiter
+  sharedRateLimiters?: Map<string, RateLimiterI>
 
   constructor(lexicons?: LexiconDoc[], opts: Options = {}) {
     if (lexicons) {
@@ -84,25 +84,26 @@ export class Server {
       json: jsonParser({ limit: opts?.payload?.jsonLimit }),
       text: textParser({ limit: opts?.payload?.textLimit }),
     }
-    this.globalRateLimiters = []
-    this.sharedRateLimiters = {}
-    this.routeRateLimiters = {}
-    if (opts?.rateLimits?.global) {
-      for (const limit of opts.rateLimits.global) {
-        const rateLimiter = opts.rateLimits.creator({
-          ...limit,
-          keyPrefix: `rl-${limit.name}`,
-        })
-        this.globalRateLimiters.push(rateLimiter)
+
+    if (opts.rateLimits) {
+      const { global, shared, creator, bypass } = opts.rateLimits
+
+      if (global) {
+        this.globalRateLimiter = RouteRateLimiter.from(
+          global.map(({ name, ...limit }) =>
+            creator({ ...limit, keyPrefix: `rl-${name}` }),
+          ),
+          { bypass },
+        )
       }
-    }
-    if (opts?.rateLimits?.shared) {
-      for (const limit of opts.rateLimits.shared) {
-        const rateLimiter = opts.rateLimits.creator({
-          ...limit,
-          keyPrefix: `rl-${limit.name}`,
-        })
-        this.sharedRateLimiters[limit.name] = rateLimiter
+
+      if (shared) {
+        this.sharedRateLimiters = new Map(
+          shared.map(({ name, ...limit }) => [
+            name,
+            creator({ ...limit, keyPrefix: `rl-${name}` }),
+          ]),
+        )
       }
     }
   }
@@ -177,7 +178,6 @@ export class Server {
       middleware.push(this.middleware.json)
       middleware.push(this.middleware.text)
     }
-    this.setupRouteRateLimits(nsid, config)
     this.routes[verb](
       `/xrpc/${nsid}`,
       ...middleware,
@@ -186,52 +186,43 @@ export class Server {
   }
 
   async catchall(req: Request, res: Response, next: NextFunction) {
-    if (this.globalRateLimiters) {
+    if (this.globalRateLimiter) {
       try {
-        const rlRes = await consumeMany(
-          {
-            req,
-            res,
-            auth: undefined,
-            params: {},
-            input: undefined,
-            async resetRouteRateLimits() {},
-          },
-          this.globalRateLimiters.map(
-            (rl) => (ctx: XRPCReqContext) => rl.consume(ctx),
-          ),
-        )
-        if (rlRes instanceof RateLimitExceededError) {
-          return next(rlRes)
-        }
+        await this.globalRateLimiter.handle({
+          req,
+          res,
+          auth: undefined,
+          params: {},
+          input: undefined,
+          async resetRouteRateLimits() {},
+        })
       } catch (err) {
         return next(err)
       }
     }
 
-    if (this.options.catchall) {
-      return this.options.catchall(req, res, next)
-    }
-
+    // Ensure that known XRPC methods are only called with the correct HTTP
+    // method.
     const def = this.lex.getDef(req.params.methodId)
-    if (!def) {
-      return next(new MethodNotImplementedError())
+    if (def) {
+      const expectedMethod =
+        def.type === 'procedure' ? 'POST' : def.type === 'query' ? 'GET' : null
+      if (expectedMethod != null && expectedMethod !== req.method) {
+        return next(
+          new InvalidRequestError(
+            `Incorrect HTTP method (${req.method}) expected ${expectedMethod}`,
+          ),
+        )
+      }
     }
-    // validate method
-    if (def.type === 'query' && req.method !== 'GET') {
-      return next(
-        new InvalidRequestError(
-          `Incorrect HTTP method (${req.method}) expected GET`,
-        ),
-      )
-    } else if (def.type === 'procedure' && req.method !== 'POST') {
-      return next(
-        new InvalidRequestError(
-          `Incorrect HTTP method (${req.method}) expected POST`,
-        ),
-      )
+
+    if (this.options.catchall) {
+      this.options.catchall.call(null, req, res, next)
+    } else if (!def) {
+      next(new MethodNotImplementedError())
+    } else {
+      next()
     }
-    return next()
   }
 
   createHandler(
@@ -251,18 +242,8 @@ export class Server {
             validateOutput(nsid, def, output, this.lex)
     const assertValidXrpcParams = (params: unknown) =>
       this.lex.assertValidXrpcParams(nsid, params)
-    const rls = this.routeRateLimiters[nsid] ?? []
-    const consumeRateLimit = (reqCtx: XRPCReqContext) =>
-      consumeMany(
-        reqCtx,
-        rls.map((rl) => (ctx: XRPCReqContext) => rl.consume(ctx)),
-      )
-
-    const resetRateLimit = (reqCtx: XRPCReqContext) =>
-      resetMany(
-        reqCtx,
-        rls.map((rl) => (ctx: XRPCReqContext) => rl.reset(ctx)),
-      )
+
+    const routeLimiter = this.createRouteRateLimiter(nsid, routeCfg)
 
     return async function (req, res, next) {
       try {
@@ -283,14 +264,11 @@ export class Server {
           auth: locals.auth,
           req,
           res,
-          resetRouteRateLimits: async () => resetRateLimit(reqCtx),
+          resetRouteRateLimits: async () => routeLimiter?.reset(reqCtx),
         }
 
         // handle rate limits
-        const result = await consumeRateLimit(reqCtx)
-        if (result instanceof RateLimitExceededError) {
-          return next(result)
-        }
+        if (routeLimiter) await routeLimiter.handle(reqCtx)
 
         // run the handler
         const output = await routeCfg.handler(reqCtx)
@@ -300,12 +278,12 @@ export class Server {
           res.status(200)
           res.end()
         } else if (isHandlerPipeThroughStream(output)) {
-          setHeaders(res, output)
+          setHeaders(res, output.headers)
           res.status(200)
           res.header('Content-Type', output.encoding)
           await pipeline(output.stream, res)
         } else if (isHandlerPipeThroughBuffer(output)) {
-          setHeaders(res, output)
+          setHeaders(res, output.headers)
           res.status(200)
           res.header('Content-Type', output.encoding)
           res.end(output.buffer)
@@ -315,7 +293,7 @@ export class Server {
           validateResOutput?.(output)
 
           res.status(200)
-          setHeaders(res, output)
+          setHeaders(res, output.headers)
 
           if (
             output.encoding === 'application/json' ||
@@ -432,76 +410,43 @@ export class Server {
     }
   }
 
-  private setupRouteRateLimits(nsid: string, config: XRPCHandlerConfig) {
-    this.routeRateLimiters[nsid] = []
-    for (const limit of this.globalRateLimiters) {
-      this.routeRateLimiters[nsid].push({
-        consume: (ctx: XRPCReqContext) => limit.consume(ctx),
-        reset: (ctx: XRPCReqContext) => limit.reset(ctx),
-      })
-    }
+  private createRouteRateLimiter(
+    nsid: string,
+    config: XRPCHandlerConfig,
+  ): RouteRateLimiter | undefined {
+    // No route specific rate limiting configured, use the global rate limiter.
+    if (!config.rateLimit) return this.globalRateLimiter
 
-    if (config.rateLimit) {
-      const limits = Array.isArray(config.rateLimit)
-        ? config.rateLimit
-        : [config.rateLimit]
-      this.routeRateLimiters[nsid] = []
-      for (let i = 0; i < limits.length; i++) {
-        const limit = limits[i]
-        const { calcKey, calcPoints } = limit
-        if (isShared(limit)) {
-          const rateLimiter = this.sharedRateLimiters[limit.name]
-          if (rateLimiter) {
-            this.routeRateLimiters[nsid].push({
-              consume: (ctx: XRPCReqContext) =>
-                rateLimiter.consume(ctx, {
-                  calcKey,
-                  calcPoints,
-                }),
-              reset: (ctx: XRPCReqContext) =>
-                rateLimiter.reset(ctx, {
-                  calcKey,
-                }),
-            })
-          }
-        } else {
-          const { durationMs, points } = limit
-          const rateLimiter = this.options.rateLimits?.creator({
-            keyPrefix: `nsid-${i}`,
-            durationMs,
-            points,
-            calcKey,
-            calcPoints,
-          })
-          if (rateLimiter) {
-            this.sharedRateLimiters[nsid] = rateLimiter
-            this.routeRateLimiters[nsid].push({
-              consume: (ctx: XRPCReqContext) =>
-                rateLimiter.consume(ctx, {
-                  calcKey,
-                  calcPoints,
-                }),
-              reset: (ctx: XRPCReqContext) =>
-                rateLimiter.reset(ctx, {
-                  calcKey,
-                }),
-            })
-          }
-        }
+    const { rateLimits } = this.options
+
+    // @NOTE Silently ignore creation of route specific rate limiter if the
+    // `rateLimits` options was not provided to the constructor.
+    if (!rateLimits) return this.globalRateLimiter
+
+    const { creator, bypass } = rateLimits
+
+    const rateLimiters = asArray(config.rateLimit).map((options, i) => {
+      if (isShared(options)) {
+        const rateLimiter = this.sharedRateLimiters?.get(options.name)
+
+        // The route config references a shared rate limiter that does not
+        // exist. This is a configuration error.
+        assert(rateLimiter, `Shared rate limiter "${options.name}" not defined`)
+
+        return WrappedRateLimiter.from(rateLimiter, options)
+      } else {
+        return creator({ ...options, keyPrefix: `${nsid}-${i}` })
       }
-    }
-  }
-}
+    })
 
-function setHeaders(
-  res: Response,
-  result: HandlerSuccess | HandlerPipeThrough,
-) {
-  const { headers } = result
-  if (headers) {
-    for (const [name, val] of Object.entries(headers)) {
-      if (val != null) res.header(name, val)
-    }
+    // If the route config contains an empty array, use global rate limiter.
+    if (!rateLimiters.length) return this.globalRateLimiter
+
+    // The global rate limiter (if present) should be applied in addition to
+    // the route specific rate limiters.
+    if (this.globalRateLimiter) rateLimiters.push(this.globalRateLimiter)
+
+    return RouteRateLimiter.from(rateLimiters, { bypass })
   }
 }
 

package/src/types.ts

@@ -29,6 +29,7 @@ export type Options = {
     creator: RateLimiterCreator
     global?: ServerRateLimitDescription[]
     shared?: ServerRateLimitDescription[]
+    bypass?: (ctx: XRPCReqContext) => boolean
   }
   /**
    * By default, errors are converted to {@link XRPCError} using
@@ -148,14 +149,23 @@ export interface RateLimiterI {
   reset: RateLimiterReset
 }
 
+export type RateLimiterConsumeOptions = {
+  calcKey?: CalcKeyFn
+  calcPoints?: CalcPointsFn
+}
+
 export type RateLimiterConsume = (
   ctx: XRPCReqContext,
-  opts?: { calcKey?: CalcKeyFn; calcPoints?: CalcPointsFn },
+  opts?: RateLimiterConsumeOptions,
 ) => Promise<RateLimiterStatus | RateLimitExceededError | null>
 
+export type RateLimiterResetOptions = {
+  calcKey?: CalcKeyFn
+}
+
 export type RateLimiterReset = (
   ctx: XRPCReqContext,
-  opts?: { calcKey?: CalcKeyFn },
+  opts?: RateLimiterResetOptions,
 ) => Promise<void>
 
 export type RateLimiterCreator = (opts: {

package/src/util.ts

@@ -1,5 +1,5 @@
 import assert from 'node:assert'
-import { IncomingMessage } from 'node:http'
+import { IncomingMessage, OutgoingMessage } from 'node:http'
 import { Duplex, Readable, pipeline } from 'node:stream'
 import express from 'express'
 import mime from 'mime-types'
@@ -24,6 +24,20 @@ import {
   handlerSuccess,
 } from './types'
 
+export const asArray = <T>(arr: T | T[]): T[] =>
+  Array.isArray(arr) ? arr : [arr]
+
+export function setHeaders(
+  res: OutgoingMessage,
+  headers?: Record<string, string | number>,
+) {
+  if (headers) {
+    for (const [name, val] of Object.entries(headers)) {
+      if (val != null) res.setHeader(name, val)
+    }
+  }
+}
+
 export function decodeQueryParams(
   def: LexXrpcProcedure | LexXrpcQuery | LexXrpcSubscription,
   params: UndecodedParams,

package/tests/rate-limiter.test.ts

@@ -132,11 +132,8 @@ describe('Parameters', () => {
   let s: http.Server
   const server = xrpcServer.createServer(LEXICONS, {
     rateLimits: {
-      creator: (opts: xrpcServer.RateLimiterOpts) =>
-        RateLimiter.memory({
-          bypassSecret: 'bypass',
-          ...opts,
-        }),
+      creator: (opts) => RateLimiter.memory(opts),
+      bypass: ({ req }) => req.headers['x-ratelimit-bypass'] === 'bypass',
       shared: [
         {
           name: 'shared-limit',