@znight/websocket 0.0.0 → 0.6.1

package/src/rate-limit.ts

@@ -0,0 +1,27 @@
+/** The token bucket of `@Gateway({ rateLimit })` (RFC 0003 §8.6), one per connection. */
+export class TokenBucket {
+  readonly #capacity: number
+  readonly #perMs: number
+  #tokens: number
+  #updatedAt: number
+
+  /** `messages` per `intervalMs`, refilled continuously; full at first. */
+  constructor(messages: number, intervalMs: number, now: number = performance.now()) {
+    this.#capacity = messages
+    this.#perMs = messages / intervalMs
+    this.#tokens = messages
+    this.#updatedAt = now
+  }
+
+  /** Takes one token: `false` when the bucket is empty. */
+  take(now: number = performance.now()): boolean {
+    const elapsed = now - this.#updatedAt
+    if (elapsed > 0) {
+      this.#tokens = Math.min(this.#capacity, this.#tokens + elapsed * this.#perMs)
+      this.#updatedAt = now
+    }
+    if (this.#tokens < 1) return false
+    this.#tokens -= 1
+    return true
+  }
+}

package/src/rooms.ts

@@ -0,0 +1,185 @@
+/**
+ * Rooms and broadcasts of one gateway (RFC 0003 §5.4): rooms are strings scoped per gateway,
+ * created on first join and deleted when empty; a broadcast encodes its frame once, then sends it
+ * to each selected member, which applies its own backpressure policy.
+ */
+import type { WsCodec } from './envelope.ts'
+
+import { jsonCodec } from './envelope.ts'
+
+/** What a room holds: a connection, seen through the operations rooms need. */
+export interface ChannelMember {
+  readonly id: string
+  /** The rooms of the member, kept in sync by the channel. */
+  readonly roomSet: Set<string>
+  /** Sends an already encoded frame (`string`: text, `Uint8Array`: binary). */
+  sendFrame(frame: string | Uint8Array): unknown
+  close(code?: number, reason?: string): void
+}
+
+/** A target of `except()`: a member or its id. */
+export type MemberRef<M extends ChannelMember = ChannelMember> = M | string
+
+/** A selection of members of a gateway (RFC 0003 §5.4). */
+export interface Broadcast<E extends object = Record<string, unknown>, M extends ChannelMember = ChannelMember> {
+  to(room: string | readonly string[]): Broadcast<E, M>
+  except(target: MemberRef<M> | readonly MemberRef<M>[]): Broadcast<E, M>
+  exceptRoom(room: string | readonly string[]): Broadcast<E, M>
+  emit<K extends keyof E & string>(event: K, data: E[K]): void
+  send(frame: string | Uint8Array): void
+  join(room: string): void
+  leave(room: string): void
+  close(code?: number, reason?: string): void
+}
+
+interface Selection {
+  /** `undefined`: every member of the gateway. */
+  readonly rooms: readonly string[] | undefined
+  readonly except: ReadonlySet<string>
+  readonly exceptRooms: readonly string[]
+}
+
+function list(value: string | readonly string[]): readonly string[] {
+  return typeof value === 'string' ? [value] : value
+}
+
+/** The members and rooms of one gateway. */
+export class RoomRegistry<M extends ChannelMember = ChannelMember> {
+  readonly #members = new Map<string, M>()
+  readonly #rooms = new Map<string, Set<M>>()
+  /** Encodes the events of broadcasts (once per broadcast). */
+  readonly codec: WsCodec
+
+  constructor(codec: WsCodec = jsonCodec) {
+    this.codec = codec
+  }
+
+  /** Open members. */
+  get size(): number {
+    return this.#members.size
+  }
+
+  add(member: M): void {
+    this.#members.set(member.id, member)
+  }
+
+  /** Removes the member and leaves its rooms. */
+  remove(member: M): void {
+    for (const room of [...member.roomSet]) this.leave(member, room)
+    this.#members.delete(member.id)
+  }
+
+  get(id: string): M | undefined {
+    return this.#members.get(id)
+  }
+
+  members(): IterableIterator<M> {
+    return this.#members.values()
+  }
+
+  join(member: M, room: string): void {
+    if (!this.#members.has(member.id)) return
+    let set = this.#rooms.get(room)
+    if (set === undefined) {
+      set = new Set()
+      this.#rooms.set(room, set)
+    }
+    set.add(member)
+    member.roomSet.add(room)
+  }
+
+  leave(member: M, room: string): void {
+    const set = this.#rooms.get(room)
+    member.roomSet.delete(room)
+    if (set === undefined) return
+    set.delete(member)
+    if (set.size === 0) this.#rooms.delete(room)
+  }
+
+  /** The members of `room` (a read-only view). */
+  room(room: string): ReadonlySet<M> {
+    return this.#rooms.get(room) ?? new Set()
+  }
+
+  /** Room → size. */
+  rooms(): ReadonlyMap<string, number> {
+    const sizes = new Map<string, number>()
+    for (const [name, set] of this.#rooms) sizes.set(name, set.size)
+    return sizes
+  }
+
+  /** A broadcast over every member, or over the members of `rooms`. */
+  select<E extends object>(rooms?: string | readonly string[], except: readonly string[] = []): Broadcast<E, M> {
+    return new BroadcastImpl<E, M>(this, {
+      rooms: rooms === undefined ? undefined : list(rooms),
+      except: new Set(except),
+      exceptRooms: [],
+    })
+  }
+
+  /** @internal The members a selection designates, each once. */
+  resolve(selection: Selection): M[] {
+    const excluded = new Set(selection.except)
+    for (const room of selection.exceptRooms) for (const member of this.room(room)) excluded.add(member.id)
+    const result: M[] = []
+    const seen = new Set<string>()
+    const sources = selection.rooms === undefined ? [this.#members.values()] : selection.rooms.map(r => this.room(r))
+    for (const source of sources) {
+      for (const member of source) {
+        if (seen.has(member.id) || excluded.has(member.id)) continue
+        seen.add(member.id)
+        result.push(member)
+      }
+    }
+    return result
+  }
+}
+
+class BroadcastImpl<E extends object, M extends ChannelMember> implements Broadcast<E, M> {
+  readonly #registry: RoomRegistry<M>
+  readonly #selection: Selection
+
+  constructor(registry: RoomRegistry<M>, selection: Selection) {
+    this.#registry = registry
+    this.#selection = selection
+  }
+
+  to(room: string | readonly string[]): Broadcast<E, M> {
+    const rooms = [...(this.#selection.rooms ?? []), ...list(room)]
+    return new BroadcastImpl(this.#registry, { ...this.#selection, rooms })
+  }
+
+  except(target: MemberRef<M> | readonly MemberRef<M>[]): Broadcast<E, M> {
+    const refs: readonly MemberRef<M>[] = Array.isArray(target) ? target : [target as MemberRef<M>]
+    const except = new Set(this.#selection.except)
+    for (const ref of refs) except.add(typeof ref === 'string' ? ref : ref.id)
+    return new BroadcastImpl(this.#registry, { ...this.#selection, except })
+  }
+
+  exceptRoom(room: string | readonly string[]): Broadcast<E, M> {
+    return new BroadcastImpl(this.#registry, {
+      ...this.#selection,
+      exceptRooms: [...this.#selection.exceptRooms, ...list(room)],
+    })
+  }
+
+  emit<K extends keyof E & string>(event: K, data: E[K]): void {
+    this.send(this.#registry.codec.encode({ kind: 'event', event, data }))
+  }
+
+  send(frame: string | Uint8Array): void {
+    for (const member of this.#registry.resolve(this.#selection)) member.sendFrame(frame)
+  }
+
+  join(room: string): void {
+    for (const member of this.#registry.resolve(this.#selection)) this.#registry.join(member, room)
+  }
+
+  leave(room: string): void {
+    for (const member of this.#registry.resolve(this.#selection)) this.#registry.leave(member, room)
+  }
+
+  close(code?: number, reason?: string): void {
+    for (const member of this.#registry.resolve(this.#selection)) member.close(code, reason)
+  }
+}

package/src/server.ts

@@ -0,0 +1,125 @@
+/**
+ * `WsServer` and `GatewayChannel` (RFC 0003 §5.4): the injectable view of the gateways, their
+ * connections and rooms. `of()` is lazy: it never throws, a handle of an unknown gateway throws a
+ * `TypeError` when used.
+ */
+import type { WsConnection, WsConnectionImpl } from './connection.ts'
+import type { Broadcast, MemberRef, RoomRegistry } from './rooms.ts'
+
+/** A gateway class or its declared path (`'/chat'`, without the HTTP prefix). */
+export type GatewayRef = (abstract new (...args: never[]) => object) | string
+
+/** The connections and rooms of one gateway (this process). */
+export interface GatewayChannel<E extends object = Record<string, unknown>> {
+  /** The gateway path, prefix included. */
+  readonly path: string
+  readonly size: number
+  connection(id: string): WsConnection | undefined
+  connections(): IterableIterator<WsConnection>
+  members(room: string): ReadonlySet<WsConnection>
+  rooms(): ReadonlyMap<string, number>
+  to(room: string | readonly string[]): Broadcast<E>
+  except(target: MemberRef<WsConnectionImpl> | readonly MemberRef<WsConnectionImpl>[]): Broadcast<E>
+  emit<K extends keyof E & string>(event: K, data: E[K]): void
+  send(frame: string | Uint8Array): void
+}
+
+/** One served gateway, as registered by the plugin. */
+export interface ServedGateway {
+  /** The gateway class. */
+  readonly type: abstract new (
+    ...args: never[]
+  ) => object
+  /** The declared path, without the prefix. */
+  readonly declaredPath: string
+  /** The path, prefix included. */
+  readonly path: string
+  readonly rooms: RoomRegistry<WsConnectionImpl>
+}
+
+/** The gateways of the application: injectable singleton provided by `WebSocketPlugin`. */
+export class WsServer {
+  readonly #byType = new Map<unknown, ServedGateway>()
+  readonly #byPath = new Map<string, ServedGateway>()
+
+  /** @internal Called by the plugin at bootstrap. */
+  register(gateway: ServedGateway): void {
+    this.#byType.set(gateway.type, gateway)
+    this.#byPath.set(gateway.declaredPath, gateway)
+  }
+
+  /** @internal Called by the plugin when it serves a new application. */
+  reset(): void {
+    this.#byType.clear()
+    this.#byPath.clear()
+  }
+
+  /** Open connections of every gateway of this process. */
+  get connections(): number {
+    let total = 0
+    for (const gateway of this.#byType.values()) total += gateway.rooms.size
+    return total
+  }
+
+  /** A lazy handle on a gateway: its class or its declared path. */
+  of<E extends object = Record<string, unknown>>(gateway: GatewayRef): GatewayChannel<E> {
+    return new Channel<E>(() => {
+      const served = typeof gateway === 'string' ? this.#byPath.get(gateway) : this.#byType.get(gateway)
+      if (served === undefined) {
+        const name = typeof gateway === 'string' ? `"${gateway}"` : gateway.name
+        throw new TypeError(
+          `WsServer.of(${name}): no gateway is served for it (or WebSocketPlugin is not bootstrapped)`,
+        )
+      }
+      return served
+    })
+  }
+}
+
+class Channel<E extends object> implements GatewayChannel<E> {
+  readonly #resolve: () => ServedGateway
+
+  constructor(resolve: () => ServedGateway) {
+    this.#resolve = resolve
+  }
+
+  get path(): string {
+    return this.#resolve().path
+  }
+
+  get size(): number {
+    return this.#resolve().rooms.size
+  }
+
+  connection(id: string): WsConnection | undefined {
+    return this.#resolve().rooms.get(id)
+  }
+
+  connections(): IterableIterator<WsConnection> {
+    return this.#resolve().rooms.members()
+  }
+
+  members(room: string): ReadonlySet<WsConnection> {
+    return this.#resolve().rooms.room(room)
+  }
+
+  rooms(): ReadonlyMap<string, number> {
+    return this.#resolve().rooms.rooms()
+  }
+
+  to(room: string | readonly string[]): Broadcast<E> {
+    return this.#resolve().rooms.select<E>(room)
+  }
+
+  except(target: MemberRef<WsConnectionImpl> | readonly MemberRef<WsConnectionImpl>[]): Broadcast<E> {
+    return this.#resolve().rooms.select<E>().except(target)
+  }
+
+  emit<K extends keyof E & string>(event: K, data: E[K]): void {
+    this.#resolve().rooms.select<E>().emit(event, data)
+  }
+
+  send(frame: string | Uint8Array): void {
+    this.#resolve().rooms.select<E>().send(frame)
+  }
+}

package/src/testing/index.ts

@@ -0,0 +1,279 @@
+/**
+ * `@znight/websocket/testing` (RFC 0003 §5.10): fakes for unit tests of gateways and services.
+ * `createTestConnection()` is a `WsConnection` recording what the code under test sends to it;
+ * `createTestServer()` is a `WsServer` recording broadcasts instead of delivering them.
+ *
+ * ```ts
+ * const server = createTestServer({ paths: [[ChatGateway, '/chat']] })
+ * const { instance } = await TestBed.create(ChatGateway, { metadata, overrides: [[WsServer, server]] })
+ * const conn = createTestConnection<ChatData, ChatEvents>({ data: { user: ALICE }, rooms: ['general'] })
+ * await instance.send(conn, { room: 'general', text: 'hi' })
+ * expect(server.broadcasts).toEqual([{ gateway: '/chat', rooms: ['general'], except: [], event: 'message.new', data: MESSAGE }])
+ * ```
+ */
+import type { HttpContext } from '@znight/http'
+import type { WsCloseInfo, WsConnection, WsSendResult } from '../connection.ts'
+import type { Broadcast } from '../rooms.ts'
+import type { GatewayChannel, GatewayRef } from '../server.ts'
+
+import { WsServer } from '../server.ts'
+
+/** One recorded broadcast (or a `close`, `join` or `leave` of a selection). */
+export interface RecordedBroadcast {
+  /** The gateway path, or the class name when no path was given for it. */
+  readonly gateway: string
+  /** Selected rooms, `[]` for every connection of the gateway. */
+  readonly rooms: readonly string[]
+  /** Excluded connection ids. */
+  readonly except: readonly string[]
+  /** Excluded rooms (omitted when none). */
+  readonly exceptRooms?: readonly string[]
+  /** The emitted event, `undefined` for `send()` and the selection operations. */
+  readonly event?: string
+  readonly data?: unknown
+  /** A raw frame (`send()`). */
+  readonly frame?: string | Uint8Array
+  /** `join`, `leave` or `close` of the selection. */
+  readonly action?: 'join' | 'leave' | 'close'
+  readonly room?: string
+}
+
+/** Options of {@link createTestConnection}. */
+export interface TestConnectionOptions<TData extends object> {
+  readonly id?: string
+  /** Gateway path. Default: `'/'`. */
+  readonly gateway?: string
+  readonly data?: TData
+  readonly rooms?: readonly string[]
+  readonly params?: Readonly<Record<string, string>>
+  readonly protocol?: string
+  readonly requestId?: string
+  /** The handshake context. Default: a minimal object with `params`, `query`, `headers` and `state`. */
+  readonly handshake?: Partial<HttpContext>
+}
+
+/** A fake connection: a `WsConnection` with the record of what it was sent. */
+export interface TestConnection<
+  TData extends object = Record<string, unknown>,
+  E extends object = Record<string, unknown>,
+> extends WsConnection<TData, E> {
+  /** Events emitted to this connection only (`conn.emit`), in order. */
+  readonly emitted: { readonly event: string; readonly data: unknown }[]
+  /** Raw frames sent to this connection (`conn.send`). */
+  readonly sent: (string | Uint8Array)[]
+  /** Broadcasts made through `conn.to()` (the connection itself is in `except`). */
+  readonly broadcasts: RecordedBroadcast[]
+}
+
+/** A fake server: a `WsServer` recording broadcasts. */
+export interface TestServer extends WsServer {
+  readonly broadcasts: RecordedBroadcast[]
+}
+
+type Record_ = (entry: RecordedBroadcast) => void
+
+/** A selection recording what is done with it. */
+class RecordingBroadcast<E extends object> implements Broadcast<E> {
+  readonly #gateway: string
+  readonly #rooms: readonly string[]
+  readonly #except: readonly string[]
+  readonly #exceptRooms: readonly string[]
+  readonly #record: Record_
+
+  constructor(
+    gateway: string,
+    rooms: readonly string[],
+    except: readonly string[],
+    exceptRooms: readonly string[],
+    record: Record_,
+  ) {
+    this.#gateway = gateway
+    this.#rooms = rooms
+    this.#except = except
+    this.#exceptRooms = exceptRooms
+    this.#record = record
+  }
+
+  to(room: string | readonly string[]): Broadcast<E> {
+    return new RecordingBroadcast(
+      this.#gateway,
+      [...this.#rooms, ...list(room)],
+      this.#except,
+      this.#exceptRooms,
+      this.#record,
+    )
+  }
+
+  except(target: string | { readonly id: string } | readonly (string | { readonly id: string })[]): Broadcast<E> {
+    const refs = Array.isArray(target) ? target : [target]
+    const ids = refs.map(ref => (typeof ref === 'string' ? ref : (ref as { readonly id: string }).id))
+    return new RecordingBroadcast(
+      this.#gateway,
+      this.#rooms,
+      [...this.#except, ...ids],
+      this.#exceptRooms,
+      this.#record,
+    )
+  }
+
+  exceptRoom(room: string | readonly string[]): Broadcast<E> {
+    return new RecordingBroadcast(
+      this.#gateway,
+      this.#rooms,
+      this.#except,
+      [...this.#exceptRooms, ...list(room)],
+      this.#record,
+    )
+  }
+
+  emit<K extends keyof E & string>(event: K, data: E[K]): void {
+    this.#record({ ...this.#base(), event, data })
+  }
+
+  send(frame: string | Uint8Array): void {
+    this.#record({ ...this.#base(), frame })
+  }
+
+  join(room: string): void {
+    this.#record({ ...this.#base(), action: 'join', room })
+  }
+
+  leave(room: string): void {
+    this.#record({ ...this.#base(), action: 'leave', room })
+  }
+
+  close(): void {
+    this.#record({ ...this.#base(), action: 'close' })
+  }
+
+  #base(): RecordedBroadcast {
+    return {
+      gateway: this.#gateway,
+      rooms: [...this.#rooms],
+      except: [...this.#except],
+      ...(this.#exceptRooms.length === 0 ? {} : { exceptRooms: [...this.#exceptRooms] }),
+    }
+  }
+}
+
+function list(value: string | readonly string[]): readonly string[] {
+  return typeof value === 'string' ? [value] : value
+}
+
+/** Creates a fake connection for unit tests. */
+export function createTestConnection<
+  TData extends object = Record<string, unknown>,
+  E extends object = Record<string, unknown>,
+>(options: TestConnectionOptions<TData> = {}): TestConnection<TData, E> {
+  const id = options.id ?? crypto.randomUUID()
+  const gateway = options.gateway ?? '/'
+  const rooms = new Set(options.rooms ?? [])
+  const abort = new AbortController()
+  const params = options.params ?? {}
+  const handshake = {
+    params,
+    query: {},
+    headers: {},
+    state: new Map(),
+    requestId: options.requestId ?? 'test-request',
+    ...options.handshake,
+  } as unknown as HttpContext
+  let state: 'open' | 'closing' | 'closed' = 'open'
+  let closeInfo: WsCloseInfo | undefined
+  const emitted: { event: string; data: unknown }[] = []
+  const sent: (string | Uint8Array)[] = []
+  const broadcasts: RecordedBroadcast[] = []
+  const result = (): WsSendResult => (state === 'open' ? 'sent' : 'dropped')
+  const connection: TestConnection<TData, E> = {
+    id,
+    gateway,
+    requestId: options.requestId ?? 'test-request',
+    handshake,
+    params,
+    protocol: options.protocol ?? '',
+    data: (options.data ?? {}) as TData,
+    rooms,
+    get readyState() {
+      return state
+    },
+    bufferedAmount: 0,
+    signal: abort.signal,
+    get closeInfo() {
+      return closeInfo
+    },
+    emitted,
+    sent,
+    broadcasts,
+    emit(event, data) {
+      if (state === 'open') emitted.push({ event, data })
+      return result()
+    },
+    send(frame) {
+      if (state === 'open') sent.push(frame)
+      return result()
+    },
+    join(room) {
+      for (const name of list(room)) rooms.add(name)
+    },
+    leave(room) {
+      rooms.delete(room)
+    },
+    to(room) {
+      return new RecordingBroadcast<E>(gateway, list(room), [id], [], entry => broadcasts.push(entry))
+    },
+    drained() {
+      return state === 'open' ? Promise.resolve() : Promise.reject(new Error('The WebSocket connection is closed'))
+    },
+    close(code = 1000, reason = '') {
+      if (state !== 'open') return
+      state = 'closed'
+      closeInfo = { code, reason, initiator: 'server' }
+      abort.abort()
+    },
+    terminate() {
+      if (state !== 'open') return
+      state = 'closed'
+      closeInfo = { code: 1006, reason: '', initiator: 'server' }
+      abort.abort()
+    },
+  }
+  return connection
+}
+
+/** Options of {@link createTestServer}. */
+export interface TestServerOptions {
+  /** Paths of gateway classes, for the `gateway` of recorded broadcasts (`[[ChatGateway, '/chat']]`). */
+  readonly paths?: readonly (readonly [abstract new (...args: never[]) => object, string])[]
+}
+
+/** Creates a fake `WsServer` recording broadcasts (no connection, every room empty). */
+export function createTestServer(options: TestServerOptions = {}): TestServer {
+  const paths = new Map<unknown, string>(options.paths ?? [])
+  const broadcasts: RecordedBroadcast[] = []
+  const record: Record_ = entry => broadcasts.push(entry)
+  class FakeServer extends WsServer {
+    readonly broadcasts: RecordedBroadcast[] = broadcasts
+
+    override get connections(): number {
+      return 0
+    }
+
+    override of<E extends object = Record<string, unknown>>(gateway: GatewayRef): GatewayChannel<E> {
+      const path = typeof gateway === 'string' ? gateway : (paths.get(gateway) ?? gateway.name)
+      const all = (): RecordingBroadcast<E> => new RecordingBroadcast<E>(path, [], [], [], record)
+      return {
+        path,
+        size: 0,
+        connection: () => undefined,
+        connections: () => [][Symbol.iterator](),
+        members: () => new Set(),
+        rooms: () => new Map(),
+        to: room => all().to(room),
+        except: target => all().except(target as never),
+        emit: (event, data) => all().emit(event, data),
+        send: frame => all().send(frame),
+      }
+    }
+  }
+  return new FakeServer()
+}

package/src/tokens.ts

@@ -0,0 +1,23 @@
+import type { TypedToken } from '@znight/core'
+import type { WsConnection } from './connection.ts'
+import type { WebSocketPlugin } from './plugin.ts'
+import type { WsServer } from './server.ts'
+
+/**
+ * Tokens of the values `@znight/websocket` supplies to the container. `WsConnection` is a
+ * request-scoped external: the request scope of a WebSocket handshake lasts as long as its
+ * connection. Inject it in request-scoped providers (`@Gateway({ scope: 'request' })`) or through
+ * `Lazy<WsConnection>`.
+ */
+export const WsTokens: {
+  /** The {@link WsServer} of the application (singleton). */
+  readonly Server: TypedToken<WsServer>
+  /** The {@link WebSocketPlugin} instance (singleton). */
+  readonly Plugin: TypedToken<WebSocketPlugin>
+  /** The {@link WsConnection} of the current connection (request scope). */
+  readonly Connection: TypedToken<WsConnection>
+} = Object.freeze({
+  Server: '@znight/websocket#WsServer',
+  Plugin: '@znight/websocket#WebSocketPlugin',
+  Connection: '@znight/websocket#WsConnection',
+})