@atproto/xrpc-server 0.3.0 → 0.3.2

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,
+    })
+  }
+}

package/src/types.ts

@@ -15,6 +15,11 @@ export type Options = {
     blobLimit?: number
     textLimit?: number
   }
+  rateLimits?: {
+    creator: RateLimiterCreator
+    global?: ServerRateLimitDescription[]
+    shared?: ServerRateLimitDescription[]
+  }
 }
 
 export type UndecodedParams = typeof express.request['query']
@@ -50,13 +55,17 @@ export type HandlerError = zod.infer<typeof handlerError>
 
 export type HandlerOutput = HandlerSuccess | HandlerError
 
-export type XRPCHandler = (ctx: {
+export type XRPCReqContext = {
   auth: HandlerAuth | undefined
   params: Params
   input: HandlerInput | undefined
   req: express.Request
   res: express.Response
-}) => Promise<HandlerOutput> | HandlerOutput | undefined
+}
+
+export type XRPCHandler = (
+  ctx: XRPCReqContext,
+) => Promise<HandlerOutput> | HandlerOutput | undefined
 
 export type XRPCStreamHandler = (ctx: {
   auth: HandlerAuth | undefined
@@ -76,7 +85,66 @@ export type StreamAuthVerifier = (ctx: {
   req: IncomingMessage
 }) => Promise<AuthOutput> | AuthOutput
 
+export type CalcKeyFn = (ctx: XRPCReqContext) => string
+export type CalcPointsFn = (ctx: XRPCReqContext) => number
+
+export interface RateLimiterI {
+  consume: RateLimiterConsume
+}
+
+export type RateLimiterConsume = (
+  ctx: XRPCReqContext,
+  opts?: { calcKey?: CalcKeyFn; calcPoints?: CalcPointsFn },
+) => Promise<RateLimiterStatus | RateLimitExceededError | null>
+
+export type RateLimiterCreator = (opts: {
+  keyPrefix: string
+  durationMs: number
+  points: number
+  calcKey?: (ctx: XRPCReqContext) => string
+  calcPoints?: (ctx: XRPCReqContext) => number
+}) => RateLimiterI
+
+export type ServerRateLimitDescription = {
+  name: string
+  durationMs: number
+  points: number
+  calcKey?: (ctx: XRPCReqContext) => string
+  calcPoints?: (ctx: XRPCReqContext) => number
+}
+
+export type SharedRateLimitOpts = {
+  name: string
+  calcKey?: (ctx: XRPCReqContext) => string
+  calcPoints?: (ctx: XRPCReqContext) => number
+}
+
+export type RouteRateLimitOpts = {
+  durationMs: number
+  points: number
+  calcKey?: (ctx: XRPCReqContext) => string
+  calcPoints?: (ctx: XRPCReqContext) => number
+}
+
+export type HandlerRateLimitOpts = SharedRateLimitOpts | RouteRateLimitOpts
+
+export const isShared = (
+  opts: HandlerRateLimitOpts,
+): opts is SharedRateLimitOpts => {
+  return typeof opts['name'] === 'string'
+}
+
+export type RateLimiterStatus = {
+  limit: number
+  duration: number
+  remainingPoints: number
+  msBeforeNext: number
+  consumedPoints: number
+  isFirstInDuration: boolean
+}
+
 export type XRPCHandlerConfig = {
+  rateLimit?: HandlerRateLimitOpts | HandlerRateLimitOpts[]
   auth?: AuthVerifier
   handler: XRPCHandler
 }
@@ -133,7 +201,13 @@ export class XRPCError extends Error {
 }
 
 export function isHandlerError(v: unknown): v is HandlerError {
-  return handlerError.safeParse(v).success
+  return (
+    !!v &&
+    typeof v === 'object' &&
+    typeof v['status'] === 'number' &&
+    (v['error'] === undefined || typeof v['error'] === 'string') &&
+    (v['message'] === undefined || typeof v['message'] === 'string')
+  )
 }
 
 export class InvalidRequestError extends XRPCError {
@@ -154,6 +228,16 @@ export class ForbiddenError extends XRPCError {
   }
 }
 
+export class RateLimitExceededError extends XRPCError {
+  constructor(
+    public status: RateLimiterStatus,
+    errorMessage?: string,
+    customErrorName?: string,
+  ) {
+    super(ResponseType.RateLimitExceeded, errorMessage, customErrorName)
+  }
+}
+
 export class InternalServerError extends XRPCError {
   constructor(errorMessage?: string, customErrorName?: string) {
     super(ResponseType.InternalServerError, errorMessage, customErrorName)
@@ -166,9 +250,9 @@ export class UpstreamFailureError extends XRPCError {
   }
 }
 
-export class NotEnoughResoucesError extends XRPCError {
+export class NotEnoughResourcesError extends XRPCError {
   constructor(errorMessage?: string, customErrorName?: string) {
-    super(ResponseType.NotEnoughResouces, errorMessage, customErrorName)
+    super(ResponseType.NotEnoughResources, errorMessage, customErrorName)
   }
 }