@atproto/xrpc-server 0.7.18 → 0.8.0

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, {
@@ -21,20 +22,18 @@ import {
   Lexicons,
   lexToJson,
 } from '@atproto/lexicon'
-import log from './logger'
-import { consumeMany, resetMany } from './rate-limiter'
+import log, { LOGGER_NAME } from './logger'
+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 })
   }
 }
 
@@ -542,24 +487,69 @@ function createErrorMiddleware({
   return (err, req, res, next) => {
     const locals: RequestLocals | undefined = req[kRequestLocals]
     const methodSuffix = locals ? ` method ${locals.nsid}` : ''
+
     const xrpcError = errorParser(err)
-    if (xrpcError instanceof InternalServerError) {
-      // log trace for unhandled exceptions
-      log.error({ err }, `unhandled exception in xrpc${methodSuffix}`)
-    } else {
-      // do not log trace for known xrpc errors
-      log.error(
-        {
-          status: xrpcError.type,
-          message: xrpcError.message,
-          name: xrpcError.customErrorName,
-        },
-        `error in xrpc${methodSuffix}`,
-      )
-    }
+
+    // Use the request's logger (if available) to benefit from request context
+    // (id, timing) and logging configuration (serialization, etc.).
+    const logger = isPinoHttpRequest(req) ? req.log : log
+
+    const isInternalError = xrpcError instanceof InternalServerError
+
+    logger.error(
+      {
+        // @NOTE Computation of error stack is an expensive operation, so
+        // we strip it for expected errors.
+        err:
+          isInternalError || process.env.NODE_ENV === 'development'
+            ? err
+            : toSimplifiedErrorLike(err),
+
+        // XRPC specific properties, for easier browsing of logs
+        nsid: locals?.nsid,
+        type: xrpcError.type,
+        status: xrpcError.statusCode,
+        payload: xrpcError.payload,
+
+        // Ensure that the logged item's name is set to LOGGER_NAME, instead of
+        // the name of the pino-http logger, to ensure consistency across logs.
+        name: LOGGER_NAME,
+      },
+      isInternalError
+        ? `unhandled exception in xrpc${methodSuffix}`
+        : `error in xrpc${methodSuffix}`,
+    )
+
     if (res.headersSent) {
       return next(err)
     }
+
     return res.status(xrpcError.statusCode).json(xrpcError.payload)
   }
 }
+
+function isPinoHttpRequest(req: Request): req is Request & {
+  log: { error: (obj: unknown, msg: string) => void }
+} {
+  return typeof (req as { log?: any }).log?.error === 'function'
+}
+
+function toSimplifiedErrorLike(err: unknown): unknown {
+  if (err instanceof Error) {
+    // Transform into an "ErrorLike" for pino's std "err" serializer
+    return {
+      ...err,
+      // Carry over non-enumerable properties
+      message: err.message,
+      name:
+        !Object.hasOwn(err, 'name') &&
+        Object.prototype.toString.call(err.constructor) === '[object Function]'
+          ? err.constructor.name // extract the class name for sub-classes of Error
+          : err.name,
+      // @NOTE Error.stack, Error.cause and AggregateError.error are non
+      // enumerable properties so they won't be spread to the ErrorLike
+    }
+  }
+
+  return err
+}

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',

package/tsconfig.tests.tsbuildinfo

@@ -1 +1 @@
-{"root":["./tests/_util.ts","./tests/auth.test.ts","./tests/bodies.test.ts","./tests/errors.test.ts","./tests/frames.test.ts","./tests/ipld.test.ts","./tests/parameters.test.ts","./tests/parsing.test.ts","./tests/procedures.test.ts","./tests/queries.test.ts","./tests/rate-limiter.test.ts","./tests/responses.test.ts","./tests/stream.test.ts","./tests/subscriptions.test.ts"],"version":"5.8.2"}
+{"root":["./tests/_util.ts","./tests/auth.test.ts","./tests/bodies.test.ts","./tests/errors.test.ts","./tests/frames.test.ts","./tests/ipld.test.ts","./tests/parameters.test.ts","./tests/parsing.test.ts","./tests/procedures.test.ts","./tests/queries.test.ts","./tests/rate-limiter.test.ts","./tests/responses.test.ts","./tests/stream.test.ts","./tests/subscriptions.test.ts"],"version":"5.8.3"}