@trpc/server/src/adapters/ws.ts

The tRPC server library

import type { IncomingMessage } from 'http';
import type { WebSocketServer, WebSocket as WsWebSocket } from 'ws';
import type {
  AnyRouter,
  CreateContextCallback,
  inferRouterContext,
} from '../@trpc/server';
import {
  callTRPCProcedure,
  getErrorShape,
  getTRPCErrorFromUnknown,
  transformTRPCResponse,
  TRPCError,
} from '../@trpc/server';
import type { TRPCRequestInfo } from '../@trpc/server/http';
import { type BaseHandlerOptions } from '../@trpc/server/http';
import { parseTRPCMessage, TRPC_ERROR_CODES_BY_KEY } from '../@trpc/server/rpc';
// @trpc/server/rpc
import type {
  TRPCClientOutgoingMessage,
  TRPCConnectionParamsMessage,
  TRPCErrorShape,
  TRPCReconnectNotification,
  TRPCResponseMessage,
  TRPCResultMessage,
} from '../@trpc/server/rpc';
import { parseConnectionParamsFromUnknown } from '../http';
import { isObservable, observableToAsyncIterable } from '../observable';
// oxlint-disable-next-line no-restricted-imports
import {
  isAbortError,
  isAsyncIterable,
  isObject,
  isPromise,
  isTrackedEnvelope,
  makeResource,
  run,
  type MaybePromise,
} from '../unstable-core-do-not-import';
// oxlint-disable-next-line no-restricted-imports
import type { Result } from '../unstable-core-do-not-import';
// oxlint-disable-next-line no-restricted-imports
import { iteratorResource } from '../unstable-core-do-not-import/stream/utils/asyncIterable';
import { Unpromise } from '../vendor/unpromise';
import { createURL, type NodeHTTPCreateContextFnOptions } from './node-http';
import type { Encoder } from './wsEncoder';
import { jsonEncoder } from './wsEncoder';

/**
 * Importing ws causes a build error
 * @see https://github.com/trpc/trpc/pull/5279
 */
const WEBSOCKET_OPEN = 1; /* ws.WebSocket.OPEN */

function settle<TValue>(promise: Promise<TValue>): Promise<Result<TValue>> {
  return promise.then(
    (value): Result<TValue> => ({ ok: true, value }),
    (error: unknown): Result<TValue> => ({ ok: false, error }),
  );
}

/**
 * What `await using` rejects with when disposal fails after the scope threw:
 * `error` is the disposal failure and `suppressed` the original error.
 */
function isSuppressedError(
  cause: unknown,
): cause is { error: unknown; suppressed: unknown } {
  return (
    isObject(cause) &&
    cause['name'] === 'SuppressedError' &&
    'error' in cause &&
    'suppressed' in cause
  );
}

/**
 * What the client gets when its error can't be formatted or encoded. It says
 * nothing about the error, which is what `errorFormatter` would have hidden.
 */
const fallbackErrorShape: TRPCErrorShape = {
  code: TRPC_ERROR_CODES_BY_KEY.INTERNAL_SERVER_ERROR,
  message: 'Internal server error',
  data: { code: 'INTERNAL_SERVER_ERROR', httpStatus: 500 },
};

/**
 * A procedure that gives up because its subscription was cancelled rejects
 * with an `AbortError`, which is wrapped as an INTERNAL_SERVER_ERROR.
 */
function isCancelledSetup(error: TRPCError, signal: AbortSignal) {
  return (
    signal.aborted &&
    error.code === 'INTERNAL_SERVER_ERROR' &&
    isAbortError(error.cause)
  );
}

/**
 * What `createContext` gets in the WebSocket adapter: the upgrade request
 * as `req`, the socket as `res`, and `info`. It's called once per
 * connection.
 * @see https://trpc.io/docs/server/websockets#creating-a-websocket-server
 * @public
 */
export type CreateWSSContextFnOptions = NodeHTTPCreateContextFnOptions<
  IncomingMessage,
  WsWebSocket
>;

/**
 * The `createContext` function the WebSocket adapter takes.
 * @public
 */
export type CreateWSSContextFn<TRouter extends AnyRouter> = (
  opts: CreateWSSContextFnOptions,
) => MaybePromise<inferRouterContext<TRouter>>;

/** The options of `getWSConnectionHandler`. */
export type WSConnectionHandlerOptions<TRouter extends AnyRouter> =
  BaseHandlerOptions<TRouter, IncomingMessage> &
    CreateContextCallback<
      inferRouterContext<TRouter>,
      CreateWSSContextFn<TRouter>
    >;

/**
 * The options of `applyWSSHandler`.
 * @see https://trpc.io/docs/server/websockets#creating-a-websocket-server
 */
export type WSSHandlerOptions<TRouter extends AnyRouter> =
  WSConnectionHandlerOptions<TRouter> & {
    /** The `ws` server to serve the router on. */
    wss: WebSocketServer;
    /** Only connections whose URL starts with this are served. */
    prefix?: string;
    /**
     * Sends pings to each client and closes connections that stop
     * answering.
     */
    keepAlive?: {
      /**
       * Whether to send pings.
       * @default false
       */
      enabled: boolean;
      /**
       * How long a connection can be quiet before a ping, in milliseconds.
       * @default 30_000
       */
      pingMs?: number;
      /**
       * How long to wait for a pong before closing the connection, in
       * milliseconds.
       * @default 5_000
       */
      pongWaitMs?: number;
    };
    /**
     * Stops answering clients' pings. Meant for tests.
     * @default false
     */
    dangerouslyDisablePong?: boolean;
    /**
     * Encodes and decodes messages, for formats other than JSON.
     * Experimental: it may change in any release.
     * @default jsonEncoder
     */
    experimental_encoder?: Encoder;
  };

/**
 * `iterator.next()` for a subscription loop that races it against the
 * subscription's `signal`, with a rejection turned into a `TRPCError`.
 *
 * Once the subscription is stopped, the loop doesn't wait for this call, so a
 * rejection after the stop, such as a teardown that throws, goes to
 * `onErrorAfterStop` and resolves `'abort'`. An `AbortError` is how a stopped
 * subscription is expected to end, and isn't reported.
 */
function nextUnlessStopped<TValue, TReturn>(
  iterator: AsyncIterator<TValue, TReturn>,
  signal: AbortSignal,
  onErrorAfterStop: (error: TRPCError) => void,
): Promise<IteratorResult<TValue, TReturn> | TRPCError | 'abort'> {
  return iterator.next().catch((cause: unknown) => {
    const error = getTRPCErrorFromUnknown(cause);
    if (!signal.aborted) {
      return error;
    }
    if (!isAbortError(error.cause)) {
      onErrorAfterStop(error);
    }
    return 'abort';
  });
}

/**
 * Returns a `ws` connection listener that serves a router. `applyWSSHandler`
 * uses it; call it to attach tRPC to your own WebSocket server.
 */
export function getWSConnectionHandler<TRouter extends AnyRouter>(
  opts: WSSHandlerOptions<TRouter>,
) {
  const { createContext, router } = opts;
  const { transformer } = router._def._config;
  const encoder = opts.experimental_encoder ?? jsonEncoder;

  return (client: WsWebSocket, req: IncomingMessage) => {
    type Context = inferRouterContext<TRouter>;
    type ContextResult = Result<Context>;

    const clientSubscriptions = new Map<number | string, AbortController>();
    const abortController = new AbortController();

    function registerSubscription(
      id: number | string,
      subscription: AbortController,
    ) {
      clientSubscriptions.set(id, subscription);
      return makeResource({}, () => {
        clientSubscriptions.delete(id);
      });
    }

    if (opts.keepAlive?.enabled) {
      const { pingMs, pongWaitMs } = opts.keepAlive;
      handleKeepAlive(client, pingMs, pongWaitMs);
    }

    function respond(untransformedJSON: TRPCResponseMessage) {
      client.send(
        encoder.encode(
          transformTRPCResponse(router._def._config, untransformedJSON),
        ),
      );
    }

    type ErrorInfo = Parameters<NonNullable<typeof opts.onError>>[0];

    /**
     * `onError` runs in listeners and in promises nothing awaits, where a
     * throw or a rejection would crash the process, so either goes to
     * `onFailure` instead.
     */
    function guardOnError(
      info: ErrorInfo,
      onFailure: (cause: unknown) => void,
    ) {
      try {
        const result: unknown = opts.onError?.(info);
        if (isPromise(result)) {
          result.catch(onFailure);
        }
      } catch (cause) {
        onFailure(cause);
      }
    }

    /** Reports an error to `onError`, ignoring an `onError` that fails. */
    function callOnError(info: ErrorInfo) {
      guardOnError(info, () => {
        // ignored, see guardOnError
      });
    }

    /**
     * Reports a failure of the adapter's own error handling, such as an
     * `errorFormatter` that throws or a reply that can't be sent. It's logged
     * when `onError` is missing or fails, so it's never dropped silently.
     */
    function reportInternalError(info: ErrorInfo) {
      const log = () => {
        // oxlint-disable-next-line no-console
        console.error('[@trpc/server/adapters/ws]', info.error);
      };
      if (!opts.onError) {
        log();
        return;
      }
      guardOnError(info, log);
    }

    /**
     * The error shape the client gets, or a minimal one when `errorFormatter`
     * throws, so that the request is still answered.
     */
    function errorShape(call: Omit<ErrorInfo, 'ctx' | 'req'>): TRPCErrorShape {
      try {
        return getErrorShape({ config: router._def._config, ctx, ...call });
      } catch (cause) {
        reportInternalError({
          ...call,
          error: getTRPCErrorFromUnknown(cause),
          ctx,
          req,
        });
        return fallbackErrorShape;
      }
    }

    function reportConnectionError(cause: unknown) {
      reportInternalError({
        error: getTRPCErrorFromUnknown(cause),
        path: undefined,
        type: 'unknown',
        ctx,
        req,
        input: undefined,
      });
    }

    /**
     * Sends a response that isn't tied to a request, reporting it if it
     * can't be built or sent. One that can't be sent is replaced by the
     * fallback error, so that a client still learns it was rejected.
     */
    function respondOrReport(response: () => TRPCResponseMessage) {
      let message: TRPCResponseMessage | null = null;
      try {
        message = response();
        respond(message);
      } catch (cause) {
        reportConnectionError(cause);
        if (message === null) {
          return;
        }
        try {
          respond({
            id: message.id,
            jsonrpc: message.jsonrpc,
            error: fallbackErrorShape,
          });
        } catch {
          // the same failure, and it was reported above
        }
      }
    }

    async function createCtxPromise(
      getConnectionParams: () => TRPCRequestInfo['connectionParams'],
    ): Promise<ContextResult> {
      try {
        return await run(async (): Promise<ContextResult> => {
          ctx = await createContext?.({
            req,
            res: client,
            info: {
              connectionParams: getConnectionParams(),
              calls: [],
              isBatchCall: false,
              accept: null,
              type: 'unknown',
              signal: abortController.signal,
              url: null,
            },
          });

          return {
            ok: true,
            value: ctx,
          };
        });
      } catch (cause) {
        const error = getTRPCErrorFromUnknown(cause);
        callOnError({
          error,
          path: undefined,
          type: 'unknown',
          ctx,
          req,
          input: undefined,
        });
        respondOrReport(() => ({
          id: null,
          error: errorShape({
            error,
            type: 'unknown',
            path: undefined,
            input: undefined,
          }),
        }));

        // close in next tick
        (globalThis.setImmediate ?? globalThis.setTimeout)(() => {
          client.close();
        });
        return {
          ok: false,
          error,
        };
      }
    }

    let ctx: Context | undefined = undefined;

    /**
     * promise for initializing the context
     *
     * - the context promise will be created immediately on connection if no connectionParams are expected
     * - if connection params are expected, they will be created once received
     */
    let ctxPromise =
      createURL(req).searchParams.get('connectionParams') === '1'
        ? null
        : createCtxPromise(() => null);

    async function handleRequest(
      msg: TRPCClientOutgoingMessage,
      batchIndex: number,
      ctxResult: Promise<ContextResult>,
    ) {
      const { id, jsonrpc } = msg;

      if (id === null) {
        const error = getTRPCErrorFromUnknown(
          new TRPCError({
            code: 'PARSE_ERROR',
            message: '`id` is required',
          }),
        );
        callOnError({
          error,
          path: undefined,
          type: 'unknown',
          ctx,
          req,
          input: undefined,
        });
        respond({
          id,
          jsonrpc,
          error: errorShape({
            error,
            type: 'unknown',
            path: undefined,
            input: undefined,
          }),
        });
        return;
      }
      if (msg.method === 'subscription.stop') {
        clientSubscriptions.get(id)?.abort();
        return;
      }
      const { path, lastEventId } = msg.params;
      let { input } = msg.params;
      const type = msg.method;

      if (lastEventId !== undefined) {
        if (isObject(input)) {
          input = {
            ...input,
            lastEventId: lastEventId,
          };
        } else {
          input ??= {
            lastEventId: lastEventId,
          };
        }
      }
      const abortController = new AbortController();
      // `abort()` builds an AbortError with a stack trace even when the
      // signal is already aborted, which is the common case on a stop.
      const abort = () => {
        if (!abortController.signal.aborted) {
          abortController.abort();
        }
      };
      const reportError = (cause: unknown) => {
        const error = getTRPCErrorFromUnknown(cause);
        callOnError({ error, path, type, ctx, req, input });
        return error;
      };
      const respondWithError = (cause: unknown) => {
        const error = reportError(cause);
        respond({
          id,
          jsonrpc,
          error: errorShape({ error, type, path, input }),
        });
      };

      if (type === 'subscription' && clientSubscriptions.has(id)) {
        // duplicate request ids for client
        respondWithError(
          new TRPCError({
            message: `Duplicate id ${id}`,
            code: 'BAD_REQUEST',
          }),
        );
        return;
      }

      const execute = async () => {
        // Registered before any await so that a stop or a disconnect during
        // context creation or the procedure call cancels the subscription.
        // The id is freed once this has settled, after the iterator's
        // disposal and right before the client is told it may reuse it.
        using _subscription =
          type === 'subscription'
            ? registerSubscription(id, abortController)
            : null;

        const res = await ctxResult;
        if (!res.ok) {
          throw res.error;
        }

        if (abortController.signal.aborted) {
          return;
        }

        const result = await callTRPCProcedure({
          router,
          path,
          getRawInput: async () => input,
          ctx,
          type,
          signal: abortController.signal,
          batchIndex,
        });

        const isIterableResult =
          isAsyncIterable(result) || isObservable(result);

        if (type !== 'subscription') {
          if (isIterableResult) {
            throw new TRPCError({
              code: 'UNSUPPORTED_MEDIA_TYPE',
              message: `Cannot return an async iterable or observable from a ${type} procedure with WebSockets`,
            });
          }
          // send the value as data if the method is not a subscription
          respond({
            id,
            jsonrpc,
            result: {
              type: 'data',
              data: result,
            },
          });
          return;
        }

        if (!isIterableResult) {
          throw new TRPCError({
            message: `Subscription ${path} did not return an observable or a AsyncGenerator`,
            code: 'INTERNAL_SERVER_ERROR',
          });
        }

        const iterable = isObservable(result)
          ? observableToAsyncIterable(result, abortController.signal)
          : result;

        // A failed iterator cleanup is reported to the server only: the
        // client still gets `stopped`, so a subscription that completed or
        // was stopped doesn't turn into a failed one.
        if (abortController.signal.aborted) {
          try {
            // dispose the iterator without consuming it
            await using _iterator = iteratorResource(iterable);
          } catch (cause) {
            reportError(cause);
          }
          return;
        }

        respond({
          id,
          jsonrpc,
          result: {
            type: 'started',
          },
        });

        let ended = false;
        try {
          await using iterator = iteratorResource(iterable);
          let completed = false;
          // Disposed before the iterator, so cleanup that waits for
          // cancellation is woken when iteration ends other than by `done`.
          using _abortUnlessCompleted = makeResource({}, () => {
            if (!completed) {
              abort();
            }
          });

          const abortPromise = new Promise<'abort'>((resolve) => {
            abortController.signal.onabort = () => resolve('abort');
          });
          // We need those declarations outside the loop for garbage collection reasons. If they
          // were declared inside, they would not be freed until the next value is present.
          let next:
            | null
            | TRPCError
            | Awaited<
                typeof abortPromise | ReturnType<(typeof iterator)['next']>
              >;
          let result: null | TRPCResultMessage<unknown>['result'];

          while (true) {
            next = await Unpromise.race([
              nextUnlessStopped(iterator, abortController.signal, (error) =>
                callOnError({ error, path, type, ctx, req, input }),
              ),
              abortPromise,
            ]);

            if (next === 'abort') {
              break;
            }
            if (next instanceof Error) {
              // Report the error only after the iterator's disposal settles.
              throw next;
            }
            if (next.done) {
              completed = true;
              break;
            }

            result = {
              type: 'data',
              data: next.value,
            };

            if (isTrackedEnvelope(next.value)) {
              const [id, data] = next.value;
              result.id = id;
              result.data = {
                id,
                data,
              };
            }

            respond({
              id,
              jsonrpc,
              result,
            });

            // free up references for garbage collection
            next = null;
            result = null;
          }
          ended = true;
        } catch (cause) {
          abort();
          if (ended) {
            reportError(cause);
          } else if (isSuppressedError(cause)) {
            // the client gets the iteration error, the server both
            respondWithError(cause.suppressed);
            reportError(cause.error);
          } else {
            respondWithError(cause);
          }
        }
      };

      const outcome = await settle(execute());
      if (!outcome.ok) {
        const error = getTRPCErrorFromUnknown(outcome.error);
        if (!isCancelledSetup(error, abortController.signal)) {
          try {
            respondWithError(error);
          } catch (cause) {
            // the error reply couldn't be encoded or sent
            reportInternalError({
              error: getTRPCErrorFromUnknown(cause),
              path,
              type,
              ctx,
              req,
              input,
            });
            respondOrReport(() => ({ id, jsonrpc, error: fallbackErrorShape }));
          }
          return;
        }
      }
      if (type === 'subscription') {
        respond({ id, jsonrpc, result: { type: 'stopped' } });
      }
    }
    client.on('message', (rawData, isBinary) => {
      // Handle PING/PONG as text regardless of encoder
      if (!isBinary) {
        // oxlint-disable-next-line typescript/no-base-to-string
        const msgStr = rawData.toString();
        if (msgStr === 'PONG') {
          return;
        }
        if (msgStr === 'PING') {
          if (!opts.dangerouslyDisablePong) {
            client.send('PONG');
          }
          return;
        }
      }

      // Convert rawData to a format our encoder accepts
      // ws gives us Buffer for both text and binary frames
      if (!Buffer.isBuffer(rawData)) {
        const error = new TRPCError({
          code: 'UNPROCESSABLE_CONTENT',
          message: 'Unexpected WebSocket message format',
        });
        respondOrReport(() => ({
          id: null,
          error: errorShape({
            error,
            type: 'unknown',
            path: undefined,
            input: undefined,
          }),
        }));
        return;
      }
      const data: string | Uint8Array = isBinary
        ? rawData
        : rawData.toString('utf8');

      if (!ctxPromise) {
        // If the ctxPromise wasn't created immediately, we're expecting the first message to be a TRPCConnectionParamsMessage
        ctxPromise = createCtxPromise(() => {
          let msg;
          try {
            msg = encoder.decode(data) as TRPCConnectionParamsMessage;

            if (!isObject(msg)) {
              throw new Error('Message was not an object');
            }
          } catch (cause) {
            throw new TRPCError({
              code: 'PARSE_ERROR',
              message: `Malformed TRPCConnectionParamsMessage`,
              cause,
            });
          }

          const connectionParams = parseConnectionParamsFromUnknown(msg.data);

          return connectionParams;
        });
        return;
      }

      const parsedMsgs = run(() => {
        try {
          const msgJSON: unknown = encoder.decode(data);
          const msgs: unknown[] = Array.isArray(msgJSON) ? msgJSON : [msgJSON];

          if (
            Array.isArray(msgJSON) &&
            typeof opts.maxBatchSize === 'number' &&
            msgs.length > opts.maxBatchSize
          ) {
            const error = new TRPCError({
              code: 'BAD_REQUEST',
              message: `Batch call exceeds maximum size`,
            });
            callOnError({
              error,
              path: undefined,
              type: 'unknown',
              ctx,
              req,
              input: undefined,
            });
            const shape = errorShape({
              error,
              type: 'unknown',
              path: undefined,
              input: undefined,
            });

            // Checked before parsing so that no input is deserialized.
            // Every request in the batch gets the error, as a batched HTTP
            // call does, so that clients do not wait for a response forever.
            let responded = false;
            for (const raw of msgs) {
              if (
                isObject(raw) &&
                raw['method'] !== 'subscription.stop' &&
                (typeof raw['id'] === 'number' || typeof raw['id'] === 'string')
              ) {
                const id = raw['id'];
                const jsonrpc = raw['jsonrpc'] === '2.0' ? '2.0' : undefined;
                respondOrReport(() => ({ id, jsonrpc, error: shape }));
                responded = true;
              }
            }
            if (!responded) {
              respondOrReport(() => ({ id: null, error: shape }));
            }
            return [];
          }

          return msgs.map((raw) => parseTRPCMessage(raw, transformer));
        } catch (cause) {
          const error = new TRPCError({
            code: 'PARSE_ERROR',
            cause,
          });

          respondOrReport(() => ({
            id: null,
            error: errorShape({
              error,
              type: 'unknown',
              path: undefined,
              input: undefined,
            }),
          }));

          return [];
        }
      });

      for (const [index, msg] of parsedMsgs.entries()) {
        handleRequest(msg, index, ctxPromise).catch(reportConnectionError);
      }
    });

    // WebSocket errors should be handled, as otherwise unhandled exceptions will crash Node.js.
    // This line was introduced after the following error brought down production systems:
    // "RangeError: Invalid WebSocket frame: RSV2 and RSV3 must be clear"
    // Here is the relevant discussion: https://github.com/websockets/ws/issues/1354#issuecomment-774616962
    client.on('error', (cause) => {
      callOnError({
        ctx,
        error: getTRPCErrorFromUnknown(cause),
        input: undefined,
        path: undefined,
        type: 'unknown',
        req,
      });
    });

    client.once('close', () => {
      for (const sub of clientSubscriptions.values()) {
        sub.abort();
      }
      clientSubscriptions.clear();
      abortController.abort();
    });
  };
}

/**
 * Pings a client after `pingMs` of quiet, and closes the connection if no
 * message comes within `pongWaitMs`.
 */
export function handleKeepAlive(
  client: WsWebSocket,
  pingMs = 30_000,
  pongWaitMs = 5_000,
) {
  let timeout: NodeJS.Timeout | undefined = undefined;
  let ping: NodeJS.Timeout | undefined = undefined;

  const schedulePing = () => {
    const scheduleTimeout = () => {
      timeout = setTimeout(() => {
        client.terminate();
      }, pongWaitMs);
    };
    ping = setTimeout(() => {
      client.send('PING');

      scheduleTimeout();
    }, pingMs);
  };

  const onMessage = () => {
    clearTimeout(ping);
    clearTimeout(timeout);

    schedulePing();
  };

  client.on('message', onMessage);

  client.on('close', () => {
    clearTimeout(ping);
    clearTimeout(timeout);
  });

  schedulePing();
}

/**
 * Serves a router on a `ws` WebSocket server. Returns
 * `broadcastReconnectNotification()`, which tells every client to
 * reconnect, for example before a deploy.
 * @example
 * ```ts
 * const wss = new WebSocketServer({ port: 3001 });
 * const handler = applyWSSHandler({ wss, router: appRouter, createContext });
 * process.on('SIGTERM', () => {
 *   handler.broadcastReconnectNotification();
 *   wss.close();
 * });
 * ```
 * @see https://trpc.io/docs/server/websockets#creating-a-websocket-server
 */
export function applyWSSHandler<TRouter extends AnyRouter>(
  opts: WSSHandlerOptions<TRouter>,
) {
  const encoder = opts.experimental_encoder ?? jsonEncoder;
  const onConnection = getWSConnectionHandler(opts);
  opts.wss.on('connection', (client, req) => {
    if (opts.prefix && !req.url?.startsWith(opts.prefix)) {
      return;
    }

    onConnection(client, req);
  });

  return {
    broadcastReconnectNotification: () => {
      const response: TRPCReconnectNotification = {
        id: null,
        method: 'reconnect',
      };
      const data = encoder.encode(response);
      for (const client of opts.wss.clients) {
        if (client.readyState === WEBSOCKET_OPEN) {
          client.send(data);
        }
      }
    },
  };
}

export type { Encoder } from './wsEncoder';
export { jsonEncoder } from './wsEncoder';