@atproto/xrpc-server 0.0.1 → 0.2.0

package/src/auth.ts

@@ -0,0 +1,111 @@
+import * as common from '@atproto/common'
+import { MINUTE } from '@atproto/common'
+import * as crypto from '@atproto/crypto'
+import * as ui8 from 'uint8arrays'
+import { AuthRequiredError } from './types'
+
+type ServiceJwtParams = {
+  iss: string
+  aud: string
+  exp?: number
+  keypair: crypto.Keypair
+}
+
+export const createServiceJwt = async (
+  params: ServiceJwtParams,
+): Promise<string> => {
+  const { iss, aud, keypair } = params
+  const exp = params.exp ?? Math.floor((Date.now() + MINUTE) / 1000)
+  const header = {
+    typ: 'JWT',
+    alg: keypair.jwtAlg,
+  }
+  const payload = {
+    iss,
+    aud,
+    exp,
+  }
+  const toSignStr = `${jsonToB64Url(header)}.${jsonToB64Url(payload)}`
+  const toSign = ui8.fromString(toSignStr, 'utf8')
+  const sig = await keypair.sign(toSign)
+  return `${toSignStr}.${ui8.toString(sig, 'base64url')}`
+}
+
+export const createServiceAuthHeaders = async (params: ServiceJwtParams) => {
+  const jwt = await createServiceJwt(params)
+  return {
+    headers: { authorization: `Bearer ${jwt}` },
+  }
+}
+
+const jsonToB64Url = (json: Record<string, unknown>): string => {
+  return common.utf8ToB64Url(JSON.stringify(json))
+}
+
+export const verifyJwt = async (
+  jwtStr: string,
+  ownDid: string,
+  getSigningKey: (did: string) => Promise<string>,
+): Promise<string> => {
+  const parts = jwtStr.split('.')
+  if (parts.length !== 3) {
+    throw new AuthRequiredError('poorly formatted jwt', 'BadJwt')
+  }
+  const payload = parsePayload(parts[1])
+  const sig = parts[2]
+
+  if (Date.now() / 1000 > payload.exp) {
+    throw new AuthRequiredError('jwt expired', 'JwtExpired')
+  }
+  if (payload.aud !== ownDid) {
+    throw new AuthRequiredError(
+      'jwt audience does not match service did',
+      'BadJwtAudience',
+    )
+  }
+
+  const msgBytes = ui8.fromString(parts.slice(0, 2).join('.'), 'utf8')
+  const sigBytes = ui8.fromString(sig, 'base64url')
+
+  const signingKey = await getSigningKey(payload.iss)
+
+  let validSig: boolean
+  try {
+    validSig = await crypto.verifySignature(signingKey, msgBytes, sigBytes)
+  } catch (err) {
+    throw new AuthRequiredError(
+      'could not verify jwt signature',
+      'BadJwtSignature',
+    )
+  }
+  if (!validSig) {
+    throw new AuthRequiredError(
+      'jwt signature does not match jwt issuer',
+      'BadJwtSignature',
+    )
+  }
+
+  return payload.iss
+}
+
+const parseB64UrlToJson = (b64: string) => {
+  return JSON.parse(common.b64UrlToUtf8(b64))
+}
+
+const parsePayload = (b64: string): JwtPayload => {
+  const payload = parseB64UrlToJson(b64)
+  if (!payload || typeof payload !== 'object') {
+    throw new AuthRequiredError('poorly formatted jwt', 'BadJwt')
+  } else if (typeof payload.exp !== 'number') {
+    throw new AuthRequiredError('poorly formatted jwt', 'BadJwt')
+  } else if (typeof payload.iss !== 'string') {
+    throw new AuthRequiredError('poorly formatted jwt', 'BadJwt')
+  }
+  return payload
+}
+
+type JwtPayload = {
+  iss: string
+  aud: string
+  exp: number
+}

package/src/index.ts

@@ -1,2 +1,4 @@
 export * from './types'
+export * from './auth'
 export * from './server'
+export * from './stream'

package/src/server.ts

@@ -4,7 +4,15 @@ import express, {
   NextFunction,
   RequestHandler,
 } from 'express'
-import { Lexicons, LexXrpcProcedure, LexXrpcQuery } from '@atproto/lexicon'
+import {
+  Lexicons,
+  lexToJson,
+  LexXrpcProcedure,
+  LexXrpcQuery,
+  LexXrpcSubscription,
+} from '@atproto/lexicon'
+import { check, forwardStreamErrors, schema } from '@atproto/common'
+import { ErrorFrame, Frame, MessageFrame, XrpcStreamServer } from './stream'
 import {
   XRPCHandler,
   XRPCError,
@@ -18,8 +26,17 @@ import {
   AuthVerifier,
   isHandlerError,
   Options,
+  XRPCStreamHandlerConfig,
+  XRPCStreamHandler,
+  Params,
+  InternalServerError,
 } from './types'
-import { decodeQueryParams, validateInput, validateOutput } from './util'
+import {
+  decodeQueryParams,
+  getQueryParams,
+  validateInput,
+  validateOutput,
+} from './util'
 import log from './logger'
 
 export function createServer(lexicons?: unknown[], options?: Options) {
@@ -27,8 +44,9 @@ export function createServer(lexicons?: unknown[], options?: Options) {
 }
 
 export class Server {
-  router = express.Router()
+  router = express()
   routes = express.Router()
+  subscriptions = new Map<string, XrpcStreamServer>()
   lex = new Lexicons()
   options: Options
   middleware: Record<'json' | 'text', RequestHandler>
@@ -40,6 +58,9 @@ export class Server {
     this.router.use(this.routes)
     this.router.use('/xrpc/:methodId', this.catchall.bind(this))
     this.router.use(errorMiddleware)
+    this.router.once('mount', (app: express.Application) => {
+      this.enableStreamingOnListen(app)
+    })
     this.options = opts ?? {}
     this.middleware = {
       json: express.json({ limit: opts?.payload?.jsonLimit }),
@@ -58,10 +79,32 @@ export class Server {
     const config =
       typeof configOrFn === 'function' ? { handler: configOrFn } : configOrFn
     const def = this.lex.getDef(nsid)
-    if (!def || (def.type !== 'query' && def.type !== 'procedure')) {
+    if (def?.type === 'query' || def?.type === 'procedure') {
+      this.addRoute(nsid, def, config)
+    } else {
       throw new Error(`Lex def for ${nsid} is not a query or a procedure`)
     }
-    this.addRoute(nsid, def, config)
+  }
+
+  streamMethod(
+    nsid: string,
+    configOrFn: XRPCStreamHandlerConfig | XRPCStreamHandler,
+  ) {
+    this.addStreamMethod(nsid, configOrFn)
+  }
+
+  addStreamMethod(
+    nsid: string,
+    configOrFn: XRPCStreamHandlerConfig | XRPCStreamHandler,
+  ) {
+    const config =
+      typeof configOrFn === 'function' ? { handler: configOrFn } : configOrFn
+    const def = this.lex.getDef(nsid)
+    if (def?.type === 'subscription') {
+      this.addSubscription(nsid, def, config)
+    } else {
+      throw new Error(`Lex def for ${nsid} is not a subscription`)
+    }
   }
 
   // schemas
@@ -145,9 +188,9 @@ export class Server {
     return async function (req, res, next) {
       try {
         // validate request
-        const params = decodeQueryParams(def, req.query)
+        let params = decodeQueryParams(def, req.query)
         try {
-          assertValidXrpcParams(params)
+          params = assertValidXrpcParams(params) as Params
         } catch (e) {
           throw new InvalidRequestError(String(e))
         }
@@ -181,10 +224,16 @@ export class Server {
             output?.encoding === 'application/json' ||
             output?.encoding === 'json'
           ) {
-            res.status(200).json(output.body)
-          } else if (output) {
+            const json = lexToJson(output.body)
+            res.status(200).json(json)
+          } else if (output?.body instanceof Readable) {
             res.header('Content-Type', output.encoding)
+            res.status(200)
+            forwardStreamErrors(output.body, res)
+            output.body.pipe(res)
+          } else if (output) {
             res
+              .header('Content-Type', output.encoding)
               .status(200)
               .send(
                 output.body instanceof Uint8Array
@@ -196,10 +245,99 @@ export class Server {
           }
         }
       } catch (err: unknown) {
-        next(err)
+        // Express will not call the next middleware (errorMiddleware in this case)
+        // if the value passed to next is falsy (e.g. null, undefined, 0).
+        // Hence we replace it with an InternalServerError.
+        if (!err) {
+          next(new InternalServerError())
+        } else {
+          next(err)
+        }
       }
     }
   }
+
+  protected async addSubscription(
+    nsid: string,
+    def: LexXrpcSubscription,
+    config: XRPCStreamHandlerConfig,
+  ) {
+    const assertValidXrpcParams = (params: unknown) =>
+      this.lex.assertValidXrpcParams(nsid, params)
+    this.subscriptions.set(
+      nsid,
+      new XrpcStreamServer({
+        noServer: true,
+        handler: async function* (req, signal) {
+          try {
+            // authenticate request
+            const auth = await config.auth?.({ req })
+            if (isHandlerError(auth)) {
+              throw XRPCError.fromError(auth)
+            }
+            // validate request
+            let params = decodeQueryParams(def, getQueryParams(req.url))
+            try {
+              params = assertValidXrpcParams(params) as Params
+            } catch (e) {
+              throw new InvalidRequestError(String(e))
+            }
+            // stream
+            const items = config.handler({ req, params, auth, signal })
+            for await (const item of items) {
+              if (item instanceof Frame) {
+                yield item
+                continue
+              }
+              const type = item?.['$type']
+              if (!check.is(item, schema.map) || typeof type !== 'string') {
+                yield new MessageFrame(item)
+                continue
+              }
+              const split = type.split('#')
+              let t: string
+              if (
+                split.length === 2 &&
+                (split[0] === '' || split[0] === nsid)
+              ) {
+                t = `#${split[1]}`
+              } else {
+                t = type
+              }
+              const clone = { ...item }
+              delete clone['$type']
+              yield new MessageFrame(clone, { type: t })
+            }
+          } catch (err) {
+            const xrpcErrPayload = XRPCError.fromError(err).payload
+            yield new ErrorFrame({
+              error: xrpcErrPayload.error ?? 'Unknown',
+              message: xrpcErrPayload.message,
+            })
+          }
+        },
+      }),
+    )
+  }
+
+  private enableStreamingOnListen(app: express.Application) {
+    const _listen = app.listen
+    app.listen = (...args) => {
+      // @ts-ignore the args spread
+      const httpServer = _listen.call(app, ...args)
+      httpServer.on('upgrade', (req, socket, head) => {
+        const url = new URL(req.url || '', 'http://x')
+        const sub = url.pathname.startsWith('/xrpc/')
+          ? this.subscriptions.get(url.pathname.replace('/xrpc/', ''))
+          : undefined
+        if (!sub) return socket.destroy()
+        sub.wss.handleUpgrade(req, socket, head, (ws) =>
+          sub.wss.emit('connection', ws, req),
+        )
+      })
+      return httpServer
+    }
+  }
 }
 
 function isHandlerSuccess(v: HandlerOutput): v is HandlerSuccess {

package/src/stream/frames.ts

@@ -0,0 +1,95 @@
+import * as uint8arrays from 'uint8arrays'
+import { cborEncode, cborDecodeMulti } from '@atproto/common'
+import {
+  frameHeader,
+  FrameHeader,
+  FrameType,
+  MessageFrameHeader,
+  ErrorFrameHeader,
+  ErrorFrameBody,
+  errorFrameBody,
+} from './types'
+
+export abstract class Frame {
+  header: FrameHeader
+  body: unknown
+  get op(): FrameType {
+    return this.header.op
+  }
+  toBytes(): Uint8Array {
+    return uint8arrays.concat([cborEncode(this.header), cborEncode(this.body)])
+  }
+  isMessage(): this is MessageFrame<unknown> {
+    return this.op === FrameType.Message
+  }
+  isError(): this is ErrorFrame {
+    return this.op === FrameType.Error
+  }
+  static fromBytes(bytes: Uint8Array) {
+    const decoded = cborDecodeMulti(bytes)
+    if (decoded.length > 2) {
+      throw new Error('Too many CBOR data items in frame')
+    }
+    const header = decoded[0]
+    let body: unknown = kUnset
+    if (decoded.length > 1) {
+      body = decoded[1]
+    }
+    const parsedHeader = frameHeader.safeParse(header)
+    if (!parsedHeader.success) {
+      throw new Error(`Invalid frame header: ${parsedHeader.error.message}`)
+    }
+    if (body === kUnset) {
+      throw new Error('Missing frame body')
+    }
+    const frameOp = parsedHeader.data.op
+    if (frameOp === FrameType.Message) {
+      return new MessageFrame(body, {
+        type: parsedHeader.data.t,
+      })
+    } else if (frameOp === FrameType.Error) {
+      const parsedBody = errorFrameBody.safeParse(body)
+      if (!parsedBody.success) {
+        throw new Error(`Invalid error frame body: ${parsedBody.error.message}`)
+      }
+      return new ErrorFrame(parsedBody.data)
+    } else {
+      const exhaustiveCheck: never = frameOp
+      throw new Error(`Unknown frame op: ${exhaustiveCheck}`)
+    }
+  }
+}
+
+export class MessageFrame<T = Record<string, unknown>> extends Frame {
+  header: MessageFrameHeader
+  body: T
+  constructor(body: T, opts?: { type?: string }) {
+    super()
+    this.header =
+      opts?.type !== undefined
+        ? { op: FrameType.Message, t: opts?.type }
+        : { op: FrameType.Message }
+    this.body = body
+  }
+  get type() {
+    return this.header.t
+  }
+}
+
+export class ErrorFrame<T extends string = string> extends Frame {
+  header: ErrorFrameHeader
+  body: ErrorFrameBody<T>
+  constructor(body: ErrorFrameBody<T>) {
+    super()
+    this.header = { op: FrameType.Error }
+    this.body = body
+  }
+  get code() {
+    return this.body.error
+  }
+  get message() {
+    return this.body.message
+  }
+}
+
+const kUnset = Symbol('unset')

package/src/stream/index.ts

@@ -0,0 +1,5 @@
+export * from './types'
+export * from './frames'
+export * from './stream'
+export * from './subscription'
+export * from './server'

package/src/stream/logger.ts

@@ -0,0 +1,5 @@
+import { subsystemLogger } from '@atproto/common'
+
+export const logger = subsystemLogger('xrpc-stream')
+
+export default logger

package/src/stream/server.ts

@@ -0,0 +1,65 @@
+import { IncomingMessage } from 'http'
+import { WebSocketServer, ServerOptions, WebSocket } from 'ws'
+import { ErrorFrame, Frame } from './frames'
+import logger from './logger'
+import { CloseCode, DisconnectError } from './types'
+
+export class XrpcStreamServer {
+  wss: WebSocketServer
+  constructor(opts: ServerOptions & { handler: Handler }) {
+    const { handler, ...serverOpts } = opts
+    this.wss = new WebSocketServer(serverOpts)
+    this.wss.on('connection', async (socket, req) => {
+      socket.on('error', (err) => logger.error(err, 'websocket error'))
+      try {
+        const ac = new AbortController()
+        const iterator = unwrapIterator(handler(req, ac.signal, socket, this))
+        socket.once('close', () => {
+          iterator.return?.()
+          ac.abort()
+        })
+        const safeFrames = wrapIterator(iterator)
+        for await (const frame of safeFrames) {
+          await new Promise((res, rej) => {
+            socket.send(frame.toBytes(), { binary: true }, (err) => {
+              // @TODO this callback may give more aggressive on backpressure than
+              // we ultimately want, but trying it out for the time being.
+              if (err) return rej(err)
+              res(undefined)
+            })
+          })
+          if (frame instanceof ErrorFrame) {
+            throw new DisconnectError(CloseCode.Policy, frame.body.error)
+          }
+        }
+      } catch (err) {
+        if (err instanceof DisconnectError) {
+          return socket.close(err.wsCode, err.xrpcCode)
+        } else {
+          logger.error(err, 'websocket server error')
+          return socket.terminate()
+        }
+      }
+      socket.close(CloseCode.Normal)
+    })
+  }
+}
+
+export type Handler = (
+  req: IncomingMessage,
+  signal: AbortSignal,
+  socket: WebSocket,
+  server: XrpcStreamServer,
+) => AsyncIterable<Frame>
+
+function unwrapIterator<T>(iterable: AsyncIterable<T>): AsyncIterator<T> {
+  return iterable[Symbol.asyncIterator]()
+}
+
+function wrapIterator<T>(iterator: AsyncIterator<T>): AsyncIterable<T> {
+  return {
+    [Symbol.asyncIterator]() {
+      return iterator
+    },
+  }
+}

package/src/stream/stream.ts

@@ -0,0 +1,26 @@
+import { XRPCError, ResponseType } from '@atproto/xrpc'
+import { DuplexOptions } from 'stream'
+import { createWebSocketStream, WebSocket } from 'ws'
+import { Frame } from './frames'
+
+export async function* byFrame(ws: WebSocket, options?: DuplexOptions) {
+  const wsStream = createWebSocketStream(ws, {
+    ...options,
+    readableObjectMode: true, // Ensures frame bytes don't get buffered/combined together
+  })
+  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')
+    }
+  }
+}