@atproto/xrpc-server 0.2.0 → 0.3.1

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']
@@ -37,6 +42,7 @@ export type HandlerAuth = zod.infer<typeof handlerAuth>
 export const handlerSuccess = zod.object({
   encoding: zod.string(),
   body: zod.any(),
+  headers: zod.record(zod.string()).optional(),
 })
 export type HandlerSuccess = zod.infer<typeof handlerSuccess>
 
@@ -49,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
@@ -75,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
 }
@@ -153,6 +222,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)
@@ -165,9 +244,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)
   }
 }
 

package/src/util.ts

@@ -1,3 +1,4 @@
+import assert from 'assert'
 import { Readable, Transform } from 'stream'
 import { createDeflate, createGunzip } from 'zlib'
 import express from 'express'
@@ -266,3 +267,35 @@ function decodeBodyStream(
 
   return stream
 }
+
+export function serverTimingHeader(timings: ServerTiming[]) {
+  return timings
+    .map((timing) => {
+      let header = timing.name
+      if (timing.duration) header += `;dur=${timing.duration}`
+      if (timing.description) header += `;desc="${timing.description}"`
+      return header
+    })
+    .join(', ')
+}
+
+export class ServerTimer implements ServerTiming {
+  public duration?: number
+  private startMs?: number
+  constructor(public name: string, public description?: string) {}
+  start() {
+    this.startMs = Date.now()
+    return this
+  }
+  stop() {
+    assert(this.startMs, "timer hasn't been started")
+    this.duration = Date.now() - this.startMs
+    return this
+  }
+}
+
+export interface ServerTiming {
+  name: string
+  duration?: number
+  description?: string
+}

package/tests/bodies.test.ts

@@ -43,7 +43,7 @@ const LEXICONS = [
   },
   {
     lexicon: 1,
-    id: 'io.example.validationTest2',
+    id: 'io.example.validationTestTwo',
     defs: {
       main: {
         type: 'query',
@@ -101,7 +101,7 @@ describe('Bodies', () => {
       body: ctx.input?.body,
     }),
   )
-  server.method('io.example.validationTest2', () => ({
+  server.method('io.example.validationTestTwo', () => ({
     encoding: 'json',
     body: { wrong: 'data' },
   }))
@@ -175,7 +175,7 @@ describe('Bodies', () => {
       return logger.error(obj, ...args)
     }
 
-    await expect(client.call('io.example.validationTest2')).rejects.toThrow(
+    await expect(client.call('io.example.validationTestTwo')).rejects.toThrow(
       'Internal Server Error',
     )
     expect(error).toEqual(`Output must have the property "foo"`)

package/tests/procedures.test.ts

@@ -8,7 +8,7 @@ import * as xrpcServer from '../src'
 const LEXICONS = [
   {
     lexicon: 1,
-    id: 'io.example.ping1',
+    id: 'io.example.pingOne',
     defs: {
       main: {
         type: 'procedure',
@@ -26,7 +26,7 @@ const LEXICONS = [
   },
   {
     lexicon: 1,
-    id: 'io.example.ping2',
+    id: 'io.example.pingTwo',
     defs: {
       main: {
         type: 'procedure',
@@ -41,7 +41,7 @@ const LEXICONS = [
   },
   {
     lexicon: 1,
-    id: 'io.example.ping3',
+    id: 'io.example.pingThree',
     defs: {
       main: {
         type: 'procedure',
@@ -56,7 +56,7 @@ const LEXICONS = [
   },
   {
     lexicon: 1,
-    id: 'io.example.ping4',
+    id: 'io.example.pingFour',
     defs: {
       main: {
         type: 'procedure',
@@ -84,17 +84,17 @@ const LEXICONS = [
 describe('Procedures', () => {
   let s: http.Server
   const server = xrpcServer.createServer(LEXICONS)
-  server.method('io.example.ping1', (ctx: { params: xrpcServer.Params }) => {
+  server.method('io.example.pingOne', (ctx: { params: xrpcServer.Params }) => {
     return { encoding: 'text/plain', body: ctx.params.message }
   })
   server.method(
-    'io.example.ping2',
+    'io.example.pingTwo',
     (ctx: { params: xrpcServer.Params; input?: xrpcServer.HandlerInput }) => {
       return { encoding: 'text/plain', body: ctx.input?.body }
     },
   )
   server.method(
-    'io.example.ping3',
+    'io.example.pingThree',
     async (ctx: {
       params: xrpcServer.Params
       input?: xrpcServer.HandlerInput
@@ -112,7 +112,7 @@ describe('Procedures', () => {
     },
   )
   server.method(
-    'io.example.ping4',
+    'io.example.pingFour',
     (ctx: { params: xrpcServer.Params; input?: xrpcServer.HandlerInput }) => {
       return {
         encoding: 'application/json',
@@ -133,14 +133,14 @@ describe('Procedures', () => {
   })
 
   it('serves requests', async () => {
-    const res1 = await client.call('io.example.ping1', {
+    const res1 = await client.call('io.example.pingOne', {
       message: 'hello world',
     })
     expect(res1.success).toBeTruthy()
     expect(res1.headers['content-type']).toBe('text/plain; charset=utf-8')
     expect(res1.data).toBe('hello world')
 
-    const res2 = await client.call('io.example.ping2', {}, 'hello world', {
+    const res2 = await client.call('io.example.pingTwo', {}, 'hello world', {
       encoding: 'text/plain',
     })
     expect(res2.success).toBeTruthy()
@@ -148,7 +148,7 @@ describe('Procedures', () => {
     expect(res2.data).toBe('hello world')
 
     const res3 = await client.call(
-      'io.example.ping3',
+      'io.example.pingThree',
       {},
       new TextEncoder().encode('hello world'),
       { encoding: 'application/octet-stream' },
@@ -158,7 +158,7 @@ describe('Procedures', () => {
     expect(new TextDecoder().decode(res3.data)).toBe('hello world')
 
     const res4 = await client.call(
-      'io.example.ping4',
+      'io.example.pingFour',
       {},
       { message: 'hello world' },
     )

package/tests/queries.test.ts

@@ -7,7 +7,7 @@ import * as xrpcServer from '../src'
 const LEXICONS = [
   {
     lexicon: 1,
-    id: 'io.example.ping1',
+    id: 'io.example.pingOne',
     defs: {
       main: {
         type: 'query',
@@ -25,7 +25,7 @@ const LEXICONS = [
   },
   {
     lexicon: 1,
-    id: 'io.example.ping2',
+    id: 'io.example.pingTwo',
     defs: {
       main: {
         type: 'query',
@@ -43,7 +43,7 @@ const LEXICONS = [
   },
   {
     lexicon: 1,
-    id: 'io.example.ping3',
+    id: 'io.example.pingThree',
     defs: {
       main: {
         type: 'query',
@@ -69,21 +69,25 @@ const LEXICONS = [
 describe('Queries', () => {
   let s: http.Server
   const server = xrpcServer.createServer(LEXICONS)
-  server.method('io.example.ping1', (ctx: { params: xrpcServer.Params }) => {
+  server.method('io.example.pingOne', (ctx: { params: xrpcServer.Params }) => {
     return { encoding: 'text/plain', body: ctx.params.message }
   })
-  server.method('io.example.ping2', (ctx: { params: xrpcServer.Params }) => {
+  server.method('io.example.pingTwo', (ctx: { params: xrpcServer.Params }) => {
     return {
       encoding: 'application/octet-stream',
       body: new TextEncoder().encode(String(ctx.params.message)),
     }
   })
-  server.method('io.example.ping3', (ctx: { params: xrpcServer.Params }) => {
-    return {
-      encoding: 'application/json',
-      body: { message: ctx.params.message },
-    }
-  })
+  server.method(
+    'io.example.pingThree',
+    (ctx: { params: xrpcServer.Params }) => {
+      return {
+        encoding: 'application/json',
+        body: { message: ctx.params.message },
+        headers: { 'x-test-header-name': 'test-value' },
+      }
+    },
+  )
   xrpc.addLexicons(LEXICONS)
 
   let client: ServiceClient
@@ -97,25 +101,26 @@ describe('Queries', () => {
   })
 
   it('serves requests', async () => {
-    const res1 = await client.call('io.example.ping1', {
+    const res1 = await client.call('io.example.pingOne', {
       message: 'hello world',
     })
     expect(res1.success).toBeTruthy()
     expect(res1.headers['content-type']).toBe('text/plain; charset=utf-8')
     expect(res1.data).toBe('hello world')
 
-    const res2 = await client.call('io.example.ping2', {
+    const res2 = await client.call('io.example.pingTwo', {
       message: 'hello world',
     })
     expect(res2.success).toBeTruthy()
     expect(res2.headers['content-type']).toBe('application/octet-stream')
     expect(new TextDecoder().decode(res2.data)).toBe('hello world')
 
-    const res3 = await client.call('io.example.ping3', {
+    const res3 = await client.call('io.example.pingThree', {
       message: 'hello world',
     })
     expect(res3.success).toBeTruthy()
     expect(res3.headers['content-type']).toBe('application/json; charset=utf-8')
     expect(res3.data?.message).toBe('hello world')
+    expect(res3.headers['x-test-header-name']).toEqual('test-value')
   })
 })

package/tests/rate-limiter.test.ts

@@ -0,0 +1,249 @@
+import * as http from 'http'
+import getPort from 'get-port'
+import xrpc, { ServiceClient } from '@atproto/xrpc'
+import { createServer, closeServer } from './_util'
+import * as xrpcServer from '../src'
+import { RateLimiter } from '../src'
+import { MINUTE } from '@atproto/common'
+
+const LEXICONS = [
+  {
+    lexicon: 1,
+    id: 'io.example.routeLimit',
+    defs: {
+      main: {
+        type: 'query',
+        parameters: {
+          type: 'params',
+          required: ['str'],
+          properties: {
+            str: { type: 'string' },
+          },
+        },
+        output: {
+          encoding: 'application/json',
+        },
+      },
+    },
+  },
+  {
+    lexicon: 1,
+    id: 'io.example.sharedLimitOne',
+    defs: {
+      main: {
+        type: 'query',
+        parameters: {
+          type: 'params',
+          required: ['points'],
+          properties: {
+            points: { type: 'integer' },
+          },
+        },
+        output: {
+          encoding: 'application/json',
+        },
+      },
+    },
+  },
+  {
+    lexicon: 1,
+    id: 'io.example.sharedLimitTwo',
+    defs: {
+      main: {
+        type: 'query',
+        parameters: {
+          type: 'params',
+          required: ['points'],
+          properties: {
+            points: { type: 'integer' },
+          },
+        },
+        output: {
+          encoding: 'application/json',
+        },
+      },
+    },
+  },
+  {
+    lexicon: 1,
+    id: 'io.example.toggleLimit',
+    defs: {
+      main: {
+        type: 'query',
+        parameters: {
+          type: 'params',
+          properties: {
+            shouldCount: { type: 'boolean' },
+          },
+        },
+        output: {
+          encoding: 'application/json',
+        },
+      },
+    },
+  },
+  {
+    lexicon: 1,
+    id: 'io.example.noLimit',
+    defs: {
+      main: {
+        type: 'query',
+        output: {
+          encoding: 'application/json',
+        },
+      },
+    },
+  },
+]
+
+describe('Parameters', () => {
+  let s: http.Server
+  const server = xrpcServer.createServer(LEXICONS, {
+    rateLimits: {
+      creator: (opts: xrpcServer.RateLimiterOpts) =>
+        RateLimiter.memory({
+          bypassSecret: 'bypass',
+          ...opts,
+        }),
+      shared: [
+        {
+          name: 'shared-limit',
+          durationMs: 5 * MINUTE,
+          points: 6,
+        },
+      ],
+      global: [
+        {
+          name: 'global-ip',
+          durationMs: 5 * MINUTE,
+          points: 100,
+        },
+      ],
+    },
+  })
+  server.method('io.example.routeLimit', {
+    rateLimit: {
+      durationMs: 5 * MINUTE,
+      points: 5,
+      calcKey: ({ params }) => params.str as string,
+    },
+    handler: (ctx: { params: xrpcServer.Params }) => ({
+      encoding: 'json',
+      body: ctx.params,
+    }),
+  })
+
+  server.method('io.example.sharedLimitOne', {
+    rateLimit: {
+      name: 'shared-limit',
+      calcPoints: ({ params }) => params.points as number,
+    },
+    handler: (ctx: { params: xrpcServer.Params }) => ({
+      encoding: 'json',
+      body: ctx.params,
+    }),
+  })
+  server.method('io.example.sharedLimitTwo', {
+    rateLimit: {
+      name: 'shared-limit',
+      calcPoints: ({ params }) => params.points as number,
+    },
+    handler: (ctx: { params: xrpcServer.Params }) => ({
+      encoding: 'json',
+      body: ctx.params,
+    }),
+  })
+  server.method('io.example.toggleLimit', {
+    rateLimit: [
+      {
+        durationMs: 5 * MINUTE,
+        points: 5,
+        calcPoints: ({ params }) => (params.shouldCount ? 1 : 0),
+      },
+      {
+        durationMs: 5 * MINUTE,
+        points: 10,
+      },
+    ],
+    handler: (ctx: { params: xrpcServer.Params }) => ({
+      encoding: 'json',
+      body: ctx.params,
+    }),
+  })
+  server.method('io.example.noLimit', {
+    handler: () => ({
+      encoding: 'json',
+      body: {},
+    }),
+  })
+
+  xrpc.addLexicons(LEXICONS)
+
+  let client: ServiceClient
+  beforeAll(async () => {
+    const port = await getPort()
+    s = await createServer(port, server)
+    client = xrpc.service(`http://localhost:${port}`)
+  })
+  afterAll(async () => {
+    await closeServer(s)
+  })
+
+  it('rate limits a given route', async () => {
+    const makeCall = () => client.call('io.example.routeLimit', { str: 'test' })
+    for (let i = 0; i < 5; i++) {
+      await makeCall()
+    }
+    await expect(makeCall).rejects.toThrow('Rate Limit Exceeded')
+  })
+
+  it('rate limits on a shared route', async () => {
+    await client.call('io.example.sharedLimitOne', { points: 1 })
+    await client.call('io.example.sharedLimitTwo', { points: 1 })
+    await client.call('io.example.sharedLimitOne', { points: 2 })
+    await client.call('io.example.sharedLimitTwo', { points: 2 })
+    await expect(
+      client.call('io.example.sharedLimitOne', { points: 1 }),
+    ).rejects.toThrow('Rate Limit Exceeded')
+    await expect(
+      client.call('io.example.sharedLimitTwo', { points: 1 }),
+    ).rejects.toThrow('Rate Limit Exceeded')
+  })
+
+  it('applies multiple rate-limits', async () => {
+    const makeCall = (shouldCount: boolean) =>
+      client.call('io.example.toggleLimit', { shouldCount })
+    for (let i = 0; i < 5; i++) {
+      await makeCall(true)
+    }
+    await expect(() => makeCall(true)).rejects.toThrow('Rate Limit Exceeded')
+    for (let i = 0; i < 4; i++) {
+      await makeCall(false)
+    }
+    await expect(() => makeCall(false)).rejects.toThrow('Rate Limit Exceeded')
+  })
+
+  it('applies global limits', async () => {
+    const makeCall = () => client.call('io.example.noLimit')
+    const calls: Promise<unknown>[] = []
+    for (let i = 0; i < 110; i++) {
+      calls.push(makeCall())
+    }
+    await expect(Promise.all(calls)).rejects.toThrow('Rate Limit Exceeded')
+  })
+
+  it('can bypass rate limits', async () => {
+    const makeCall = () =>
+      client.call(
+        'io.example.noLimit',
+        {},
+        {},
+        { headers: { 'X-RateLimit-Bypass': 'bypass' } },
+      )
+    const calls: Promise<unknown>[] = []
+    for (let i = 0; i < 110; i++) {
+      calls.push(makeCall())
+    }
+    await Promise.all(calls)
+  })
+})