@atproto/xrpc-server 0.2.0 → 0.3.1

package/src/server.ts

@@ -30,6 +30,11 @@ import {
   XRPCStreamHandler,
   Params,
   InternalServerError,
+  XRPCReqContext,
+  RateLimiterI,
+  RateLimiterConsume,
+  isShared,
+  RateLimitExceededError,
 } from './types'
 import {
   decodeQueryParams,
@@ -38,6 +43,7 @@ import {
   validateOutput,
 } from './util'
 import log from './logger'
+import { consumeMany } from './rate-limiter'
 
 export function createServer(lexicons?: unknown[], options?: Options) {
   return new Server(lexicons, options)
@@ -50,6 +56,9 @@ export class Server {
   lex = new Lexicons()
   options: Options
   middleware: Record<'json' | 'text', RequestHandler>
+  globalRateLimiters: RateLimiterI[]
+  sharedRateLimiters: Record<string, RateLimiterI>
+  routeRateLimiterFns: Record<string, RateLimiterConsume[]>
 
   constructor(lexicons?: unknown[], opts?: Options) {
     if (lexicons) {
@@ -66,6 +75,27 @@ export class Server {
       json: express.json({ limit: opts?.payload?.jsonLimit }),
       text: express.text({ limit: opts?.payload?.textLimit }),
     }
+    this.globalRateLimiters = []
+    this.sharedRateLimiters = {}
+    this.routeRateLimiterFns = {}
+    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?.shared) {
+      for (const limit of opts.rateLimits.shared) {
+        const rateLimiter = opts.rateLimits.creator({
+          ...limit,
+          keyPrefix: `rl-${limit.name}`,
+        })
+        this.sharedRateLimiters[limit.name] = rateLimiter
+      }
+    }
   }
 
   // handlers
@@ -138,6 +168,7 @@ export class Server {
       middleware.push(this.middleware.json)
       middleware.push(this.middleware.text)
     }
+    this.setupRouteRateLimits(nsid, config)
     this.routes[verb](
       `/xrpc/${nsid}`,
       ...middleware,
@@ -185,6 +216,10 @@ export class Server {
             validateOutput(nsid, def, output, this.lex)
     const assertValidXrpcParams = (params: unknown) =>
       this.lex.assertValidXrpcParams(nsid, params)
+    const rlFns = this.routeRateLimiterFns[nsid] ?? []
+    const consumeRateLimit = (reqCtx: XRPCReqContext) =>
+      consumeMany(reqCtx, rlFns)
+
     return async function (req, res, next) {
       try {
         // validate request
@@ -203,14 +238,24 @@ export class Server {
 
         const locals: RequestLocals = req[kRequestLocals]
 
-        // run the handler
-        const outputUnvalidated = await handler({
+        const reqCtx: XRPCReqContext = {
           params,
           input,
           auth: locals.auth,
           req,
           res,
-        })
+        }
+
+        // handle rate limits
+        if (consumeRateLimit) {
+          const result = await consumeRateLimit(reqCtx)
+          if (result instanceof RateLimitExceededError) {
+            return next(result)
+          }
+        }
+
+        // run the handler
+        const outputUnvalidated = await handler(reqCtx)
 
         if (isHandlerError(outputUnvalidated)) {
           throw XRPCError.fromError(outputUnvalidated)
@@ -219,6 +264,12 @@ export class Server {
         if (!outputUnvalidated || isHandlerSuccess(outputUnvalidated)) {
           // validate response
           const output = validateResOutput(outputUnvalidated)
+          // set headers
+          if (output?.headers) {
+            Object.entries(output.headers).forEach(([name, val]) => {
+              res.header(name, val)
+            })
+          }
           // send response
           if (
             output?.encoding === 'application/json' ||
@@ -229,6 +280,7 @@ export class Server {
           } else if (output?.body instanceof Readable) {
             res.header('Content-Type', output.encoding)
             res.status(200)
+            res.once('error', (err) => res.destroy(err))
             forwardStreamErrors(output.body, res)
             output.body.pipe(res)
           } else if (output) {
@@ -338,6 +390,55 @@ export class Server {
       return httpServer
     }
   }
+
+  private setupRouteRateLimits(nsid: string, config: XRPCHandlerConfig) {
+    this.routeRateLimiterFns[nsid] = []
+    for (const limit of this.globalRateLimiters) {
+      const consumeFn = async (ctx: XRPCReqContext) => {
+        return limit.consume(ctx)
+      }
+      this.routeRateLimiterFns[nsid].push(consumeFn)
+    }
+
+    if (config.rateLimit) {
+      const limits = Array.isArray(config.rateLimit)
+        ? config.rateLimit
+        : [config.rateLimit]
+      this.routeRateLimiterFns[nsid] = []
+      for (const limit of limits) {
+        const { calcKey, calcPoints } = limit
+        if (isShared(limit)) {
+          const rateLimiter = this.sharedRateLimiters[limit.name]
+          if (rateLimiter) {
+            const consumeFn = (ctx: XRPCReqContext) =>
+              rateLimiter.consume(ctx, {
+                calcKey,
+                calcPoints,
+              })
+            this.routeRateLimiterFns[nsid].push(consumeFn)
+          }
+        } else {
+          const { durationMs, points } = limit
+          const rateLimiter = this.options.rateLimits?.creator({
+            keyPrefix: nsid,
+            durationMs,
+            points,
+            calcKey,
+            calcPoints,
+          })
+          if (rateLimiter) {
+            this.sharedRateLimiters[nsid] = rateLimiter
+            const consumeFn = (ctx: XRPCReqContext) =>
+              rateLimiter.consume(ctx, {
+                calcKey,
+                calcPoints,
+              })
+            this.routeRateLimiterFns[nsid].push(consumeFn)
+          }
+        }
+      }
+    }
+  }
 }
 
 function isHandlerSuccess(v: HandlerOutput): v is HandlerSuccess {
@@ -378,14 +479,23 @@ function createAuthMiddleware(verifier: AuthVerifier): RequestHandler {
 const errorMiddleware: ErrorRequestHandler = function (err, req, res, next) {
   const locals: RequestLocals | undefined = req[kRequestLocals]
   const methodSuffix = locals ? ` method ${locals.nsid}` : ''
-  if (err instanceof XRPCError) {
-    log.error(err, `error in xrpc${methodSuffix}`)
-  } else {
+  const xrpcError = XRPCError.fromError(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}`,
+    )
   }
   if (res.headersSent) {
     return next(err)
   }
-  const xrpcError = XRPCError.fromError(err)
   return res.status(xrpcError.type).json(xrpcError.payload)
 }

package/src/stream/logger.ts

@@ -1,5 +1,6 @@
 import { subsystemLogger } from '@atproto/common'
 
-export const logger = subsystemLogger('xrpc-stream')
+export const logger: ReturnType<typeof subsystemLogger> =
+  subsystemLogger('xrpc-stream')
 
 export default logger

package/src/stream/stream.ts

@@ -1,26 +1,39 @@
 import { XRPCError, ResponseType } from '@atproto/xrpc'
 import { DuplexOptions } from 'stream'
 import { createWebSocketStream, WebSocket } from 'ws'
-import { Frame } from './frames'
+import { Frame, MessageFrame } from './frames'
 
-export async function* byFrame(ws: WebSocket, options?: DuplexOptions) {
-  const wsStream = createWebSocketStream(ws, {
+export function streamByteChunks(ws: WebSocket, options?: DuplexOptions) {
+  return createWebSocketStream(ws, {
     ...options,
     readableObjectMode: true, // Ensures frame bytes don't get buffered/combined together
   })
+}
+
+export async function* byFrame(ws: WebSocket, options?: DuplexOptions) {
+  const wsStream = streamByteChunks(ws, options)
   for await (const chunk of wsStream) {
     yield Frame.fromBytes(chunk)
   }
 }
 
 export async function* byMessage(ws: WebSocket, options?: DuplexOptions) {
-  for await (const frame of byFrame(ws, options)) {
-    if (frame.isMessage()) {
-      yield frame
-    } else if (frame.isError()) {
-      throw new XRPCError(-1, frame.code, frame.message)
-    } else {
-      throw new XRPCError(ResponseType.Unknown, undefined, 'Unknown frame type')
-    }
+  const wsStream = streamByteChunks(ws, options)
+  for await (const chunk of wsStream) {
+    const msg = ensureChunkIsMessage(chunk)
+    yield msg
+  }
+}
+
+export function ensureChunkIsMessage(chunk: Uint8Array): MessageFrame<unknown> {
+  const frame = Frame.fromBytes(chunk)
+  if (frame.isMessage()) {
+    return frame
+  } else if (frame.isError()) {
+    // @TODO work -1 error code into XRPCError
+    // @ts-ignore
+    throw new XRPCError(-1, frame.code, frame.message)
+  } else {
+    throw new XRPCError(ResponseType.Unknown, undefined, 'Unknown frame type')
   }
 }

package/src/stream/subscription.ts

@@ -1,7 +1,6 @@
-import { wait } from '@atproto/common'
-import { WebSocket, ClientOptions } from 'ws'
-import { byMessage } from './stream'
-import { CloseCode, DisconnectError } from './types'
+import { ClientOptions } from 'ws'
+import { WebSocketKeepAlive } from './websocket-keepalive'
+import { ensureChunkIsMessage } from './stream'
 
 export class Subscription<T = unknown> {
   constructor(
@@ -9,6 +8,7 @@ export class Subscription<T = unknown> {
       service: string
       method: string
       maxReconnectSeconds?: number
+      heartbeatIntervalMs?: number
       signal?: AbortSignal
       validate: (obj: unknown) => T | undefined
       onReconnectError?: (
@@ -24,107 +24,31 @@ export class Subscription<T = unknown> {
   ) {}
 
   async *[Symbol.asyncIterator](): AsyncGenerator<T> {
-    let initialSetup = true
-    let reconnects: number | null = null
-    const maxReconnectMs = 1000 * (this.opts.maxReconnectSeconds ?? 64)
-    while (true) {
-      if (reconnects !== null) {
-        const duration = initialSetup
-          ? Math.min(1000, maxReconnectMs)
-          : backoffMs(reconnects++, maxReconnectMs)
-        await wait(duration)
-      }
-      const ws = await this.getSocket()
-      const ac = new AbortController()
-      if (this.opts.signal) {
-        forwardSignal(this.opts.signal, ac)
+    const ws = new WebSocketKeepAlive({
+      ...this.opts,
+      getUrl: async () => {
+        const params = (await this.opts.getParams?.()) ?? {}
+        const query = encodeQueryParams(params)
+        return `${this.opts.service}/xrpc/${this.opts.method}?${query}`
+      },
+    })
+    for await (const chunk of ws) {
+      const message = await ensureChunkIsMessage(chunk)
+      const t = message.header.t
+      const clone = message.body !== undefined ? { ...message.body } : undefined
+      if (clone !== undefined && t !== undefined) {
+        clone['$type'] = t.startsWith('#') ? this.opts.method + t : t
       }
-      ws.once('open', () => {
-        initialSetup = false
-        reconnects = 0
-      })
-      ws.once('close', (code, reason) => {
-        if (code === CloseCode.Abnormal) {
-          // Forward into an error to distinguish from a clean close
-          ac.abort(
-            new AbnormalCloseError(`Abnormal ws close: ${reason.toString()}`),
-          )
-        }
-      })
-      try {
-        const cancelable = { signal: ac.signal }
-        for await (const message of byMessage(ws, cancelable)) {
-          const t = message.header.t
-          const clone =
-            message.body !== undefined ? { ...message.body } : undefined
-          if (clone !== undefined && t !== undefined) {
-            clone['$type'] = t.startsWith('#') ? this.opts.method + t : t
-          }
-          const result = this.opts.validate(clone)
-          if (result !== undefined) {
-            yield result
-          }
-        }
-      } catch (_err) {
-        const err = _err?.['code'] === 'ABORT_ERR' ? _err['cause'] : _err
-        if (err instanceof DisconnectError) {
-          // We cleanly end the connection
-          ws.close(err.wsCode)
-          break
-        }
-        ws.close() // No-ops if already closed or closing
-        if (isReconnectable(err)) {
-          reconnects ??= 0 // Never reconnect with a null
-          this.opts.onReconnectError?.(err, reconnects, initialSetup)
-          continue
-        } else {
-          throw err
-        }
+      const result = this.opts.validate(clone)
+      if (result !== undefined) {
+        yield result
       }
-      break // Other side cleanly ended stream and disconnected
     }
   }
-
-  private async getSocket() {
-    const params = (await this.opts.getParams?.()) ?? {}
-    const query = encodeQueryParams(params)
-    const url = `${this.opts.service}/xrpc/${this.opts.method}?${query}`
-    return new WebSocket(url, this.opts)
-  }
 }
 
 export default Subscription
 
-class AbnormalCloseError extends Error {
-  code = 'EWSABNORMALCLOSE'
-}
-
-function isReconnectable(err: unknown): boolean {
-  // Network errors are reconnectable.
-  // AuthenticationRequired and InvalidRequest XRPCErrors are not reconnectable.
-  // @TODO method-specific XRPCErrors may be reconnectable, need to consider. Receiving
-  // an invalid message is not current reconnectable, but the user can decide to skip them.
-  if (!err || typeof err['code'] !== 'string') return false
-  return networkErrorCodes.includes(err['code'])
-}
-
-const networkErrorCodes = [
-  'EWSABNORMALCLOSE',
-  'ECONNRESET',
-  'ECONNREFUSED',
-  'ECONNABORTED',
-  'EPIPE',
-  'ETIMEDOUT',
-  'ECANCELED',
-]
-
-function backoffMs(n: number, maxMs: number) {
-  const baseSec = Math.pow(2, n) // 1, 2, 4, ...
-  const randSec = Math.random() - 0.5 // Random jitter between -.5 and .5 seconds
-  const ms = 1000 * (baseSec + randSec)
-  return Math.min(ms, maxMs)
-}
-
 function encodeQueryParams(obj: Record<string, unknown>): string {
   const params = new URLSearchParams()
   Object.entries(obj).forEach(([key, value]) => {
@@ -163,13 +87,3 @@ function encodeQueryParam(value: unknown): string | string[] {
   }
   throw new Error(`Cannot encode ${typeof value}s into query params`)
 }
-
-function forwardSignal(signal: AbortSignal, ac: AbortController) {
-  if (signal.aborted) {
-    return ac.abort(signal.reason)
-  } else {
-    signal.addEventListener('abort', () => ac.abort(signal.reason), {
-      signal: ac.signal,
-    })
-  }
-}

package/src/stream/websocket-keepalive.ts

@@ -0,0 +1,151 @@
+import { SECOND, wait } from '@atproto/common'
+import { WebSocket, ClientOptions } from 'ws'
+import { streamByteChunks } from './stream'
+import { CloseCode, DisconnectError } from './types'
+
+export class WebSocketKeepAlive {
+  public ws: WebSocket | null = null
+  public initialSetup = true
+  public reconnects: number | null = null
+
+  constructor(
+    public opts: ClientOptions & {
+      getUrl: () => Promise<string>
+      maxReconnectSeconds?: number
+      signal?: AbortSignal
+      heartbeatIntervalMs?: number
+      onReconnectError?: (
+        error: unknown,
+        n: number,
+        initialSetup: boolean,
+      ) => void
+    },
+  ) {}
+
+  async *[Symbol.asyncIterator](): AsyncGenerator<Uint8Array> {
+    const maxReconnectMs = 1000 * (this.opts.maxReconnectSeconds ?? 64)
+    while (true) {
+      if (this.reconnects !== null) {
+        const duration = this.initialSetup
+          ? Math.min(1000, maxReconnectMs)
+          : backoffMs(this.reconnects++, maxReconnectMs)
+        await wait(duration)
+      }
+      const url = await this.opts.getUrl()
+      this.ws = new WebSocket(url, this.opts)
+      const ac = new AbortController()
+      if (this.opts.signal) {
+        forwardSignal(this.opts.signal, ac)
+      }
+      this.ws.once('open', () => {
+        this.initialSetup = false
+        this.reconnects = 0
+        if (this.ws) {
+          this.startHeartbeat(this.ws)
+        }
+      })
+      this.ws.once('close', (code, reason) => {
+        if (code === CloseCode.Abnormal) {
+          // Forward into an error to distinguish from a clean close
+          ac.abort(
+            new AbnormalCloseError(`Abnormal ws close: ${reason.toString()}`),
+          )
+        }
+      })
+
+      try {
+        const wsStream = streamByteChunks(this.ws, { signal: ac.signal })
+        for await (const chunk of wsStream) {
+          yield chunk
+        }
+      } catch (_err) {
+        const err = _err?.['code'] === 'ABORT_ERR' ? _err['cause'] : _err
+        if (err instanceof DisconnectError) {
+          // We cleanly end the connection
+          this.ws?.close(err.wsCode)
+          break
+        }
+        this.ws?.close() // No-ops if already closed or closing
+        if (isReconnectable(err)) {
+          this.reconnects ??= 0 // Never reconnect with a null
+          this.opts.onReconnectError?.(err, this.reconnects, this.initialSetup)
+          continue
+        } else {
+          throw err
+        }
+      }
+      break // Other side cleanly ended stream and disconnected
+    }
+  }
+
+  startHeartbeat(ws: WebSocket) {
+    let isAlive = true
+    let heartbeatInterval: NodeJS.Timer | null = null
+
+    const checkAlive = () => {
+      if (!isAlive) {
+        return ws.terminate()
+      }
+      isAlive = false // expect websocket to no longer be alive unless we receive a "pong" within the interval
+      ws.ping()
+    }
+
+    checkAlive()
+    heartbeatInterval = setInterval(
+      checkAlive,
+      this.opts.heartbeatIntervalMs ?? 10 * SECOND,
+    )
+
+    ws.on('pong', () => {
+      isAlive = true
+    })
+    ws.once('close', () => {
+      if (heartbeatInterval) {
+        clearInterval(heartbeatInterval)
+        heartbeatInterval = null
+      }
+    })
+  }
+}
+
+export default WebSocketKeepAlive
+
+class AbnormalCloseError extends Error {
+  code = 'EWSABNORMALCLOSE'
+}
+
+function isReconnectable(err: unknown): boolean {
+  // Network errors are reconnectable.
+  // AuthenticationRequired and InvalidRequest XRPCErrors are not reconnectable.
+  // @TODO method-specific XRPCErrors may be reconnectable, need to consider. Receiving
+  // an invalid message is not current reconnectable, but the user can decide to skip them.
+  if (!err || typeof err['code'] !== 'string') return false
+  return networkErrorCodes.includes(err['code'])
+}
+
+const networkErrorCodes = [
+  'EWSABNORMALCLOSE',
+  'ECONNRESET',
+  'ECONNREFUSED',
+  'ECONNABORTED',
+  'EPIPE',
+  'ETIMEDOUT',
+  'ECANCELED',
+]
+
+function backoffMs(n: number, maxMs: number) {
+  const baseSec = Math.pow(2, n) // 1, 2, 4, ...
+  const randSec = Math.random() - 0.5 // Random jitter between -.5 and .5 seconds
+  const ms = 1000 * (baseSec + randSec)
+  return Math.min(ms, maxMs)
+}
+
+function forwardSignal(signal: AbortSignal, ac: AbortController) {
+  if (signal.aborted) {
+    return ac.abort(signal.reason)
+  } else {
+    signal.addEventListener('abort', () => ac.abort(signal.reason), {
+      signal: ac.signal,
+    })
+  }
+}