tech-log-frontend/tests/unit/realtime/websocket-protocol.test.ts

import { describe, expect, it, vi } from "vitest";

import {
  REALTIME_WEBSOCKET_PROTOCOL,
  decodeWebSocketServerFrame,
  encodeWebSocketClientFrame,
  nextUnsignedSequence,
  type WebSocketAdvertisedLimits,
  type WebSocketSubscribeFrame,
} from "../../../src/adapters/realtime/websocket/websocket-protocol.ts";

const LIMITS: WebSocketAdvertisedLimits = Object.freeze({
  maxFrameBytes: 65_536,
  maxSubscriptions: 32,
  maxInboundQueueCount: 256,
  maxInboundQueueBytes: 4 * 1_024 * 1_024,
  maxOutboundQueueCount: 128,
  maxOutboundQueueBytes: 256 * 1_024,
  maxBufferedAmountBytes: 256 * 1_024,
  maxEventsPerSecond: 128,
});

function encode(value: unknown): string {
  return JSON.stringify(value);
}

describe("realtime WebSocket protocol", () => {

  it("rejects oversized text before allocating a full UTF-8 copy", () => {
    const encoderSpy = vi.spyOn(TextEncoder.prototype, "encode");
    try {
      const oversize = "a".repeat(64);
      expect(decodeWebSocketServerFrame(oversize, 8)).toMatchObject({
        ok: false,
        error: { code: "FRAME_TOO_LARGE" },
      });
      expect(encoderSpy).not.toHaveBeenCalled();
    } finally {
      encoderSpy.mockRestore();
    }
  });

  it("counts multibyte and lone-surrogate bytes like TextEncoder", () => {
    const samples = [
      "abc",
      "\u00e9\u00e9",
      "\u20ac\u20ac",
      "\u{1f600}",
      "a\ud800b",
      "\udc00",
    ];
    for (const sample of samples) {
      const expected = new TextEncoder().encode(sample).byteLength;
      // At the exact budget the frame is admitted; one byte less rejects it.
      expect(
        decodeWebSocketServerFrame(sample, expected).ok ||
          decodeWebSocketServerFrame(sample, expected),
      ).toBeTruthy();
      expect(decodeWebSocketServerFrame(sample, expected - 1)).toMatchObject({
        ok: false,
        error: { code: "FRAME_TOO_LARGE" },
      });
    }
  });
  it("decodes and freezes an exact WELCOME frame", () => {
    const result = decodeWebSocketServerFrame(
      encode({
        type: "WELCOME",
        protocol: REALTIME_WEBSOCKET_PROTOCOL,
        connectionId: "connection.0001",
        heartbeatMs: 15_000,
        heartbeatAckTimeoutMs: 5_000,
        limits: LIMITS,
      }),
      65_536,
    );

    expect(result).toMatchObject({
      ok: true,
      value: {
        type: "WELCOME",
        limits: { maxSubscriptions: 32 },
      },
    });
    if (!result.ok) throw new Error("Expected WELCOME to decode.");
    if (result.value.type !== "WELCOME") {
      throw new Error("Expected the WELCOME discriminant.");
    }
    expect(Object.isFrozen(result.value)).toBe(true);
    expect(Object.isFrozen(result.value.limits)).toBe(true);
  });

  it("rejects extra keys, unknown frames, wrong versions and binary data", () => {
    const welcome = {
      type: "WELCOME",
      protocol: REALTIME_WEBSOCKET_PROTOCOL,
      connectionId: "connection.0001",
      heartbeatMs: 15_000,
      heartbeatAckTimeoutMs: 5_000,
      limits: LIMITS,
    };

    expect(
      decodeWebSocketServerFrame(
        encode({ ...welcome, credential: "must-not-cross" }),
        65_536,
      ),
    ).toEqual({
      ok: false,
      error: { code: "MALFORMED_FRAME" },
    });
    expect(
      decodeWebSocketServerFrame(
        encode({ ...welcome, protocol: "realtime.v2" }),
        65_536,
      ),
    ).toEqual({
      ok: false,
      error: { code: "PROTOCOL_MISMATCH" },
    });
    expect(
      decodeWebSocketServerFrame(
        encode({
          type: "COMMAND",
          protocol: REALTIME_WEBSOCKET_PROTOCOL,
        }),
        65_536,
      ),
    ).toEqual({
      ok: false,
      error: { code: "UNKNOWN_FRAME" },
    });
    expect(
      decodeWebSocketServerFrame(new Uint8Array([1, 2, 3]), 65_536),
    ).toEqual({
      ok: false,
      error: { code: "BINARY_FRAME" },
    });

    const duplicateTopLevel = encode(welcome).replace(
      '{"type":"WELCOME",',
      '{"\\u0074ype":"WELCOME","type":"WELCOME",',
    );
    expect(
      decodeWebSocketServerFrame(duplicateTopLevel, 65_536),
    ).toEqual({
      ok: false,
      error: { code: "MALFORMED_FRAME" },
    });
    const duplicateNested = encode({
      type: "EVENT",
      protocol: REALTIME_WEBSOCKET_PROTOCOL,
      subscriptionId: "subscription.0001",
      envelope: { streamId: "orders.v1" },
    }).replace(
      '"streamId":"orders.v1"',
      '"streamId":"orders.v1","\\u0073treamId":"shadowed"',
    );
    expect(
      decodeWebSocketServerFrame(duplicateNested, 65_536),
    ).toEqual({
      ok: false,
      error: { code: "MALFORMED_FRAME" },
    });
  });

  it("enforces frame bytes and canonical uint64 sequences", () => {
    expect(
      decodeWebSocketServerFrame(
        encode({
          type: "SUBSCRIBED",
          protocol: REALTIME_WEBSOCKET_PROTOCOL,
          subscriptionId: "subscription.0001",
          streamEpoch: "stream-epoch.0001",
          acceptedCursor: "cursor.0001",
          nextExpectedSequence: "01",
        }),
        65_536,
      ),
    ).toEqual({
      ok: false,
      error: { code: "MALFORMED_FRAME" },
    });
    expect(
      decodeWebSocketServerFrame(
        encode({
          type: "HEARTBEAT_ACK",
          protocol: REALTIME_WEBSOCKET_PROTOCOL,
          nonce: "nonce.0001",
        }),
        8,
      ),
    ).toEqual({
      ok: false,
      error: { code: "FRAME_TOO_LARGE" },
    });
    expect(nextUnsignedSequence("0")).toBe("1");
    expect(nextUnsignedSequence("18446744073709551614")).toBe(
      "18446744073709551615",
    );
    expect(nextUnsignedSequence("18446744073709551615")).toBeNull();
    expect(nextUnsignedSequence("01")).toBeNull();
  });

  it("decodes only an exact UNSUBSCRIBED acknowledgement", () => {
    const acknowledgement = {
      type: "UNSUBSCRIBED",
      protocol: REALTIME_WEBSOCKET_PROTOCOL,
      subscriptionId: "subscription.0001",
    };
    const decoded = decodeWebSocketServerFrame(
      encode(acknowledgement),
      65_536,
    );

    expect(decoded).toMatchObject({
      ok: true,
      value: acknowledgement,
    });
    if (!decoded.ok) {
      throw new Error("Expected UNSUBSCRIBED to decode.");
    }
    expect(Object.isFrozen(decoded.value)).toBe(true);
    expect(
      decodeWebSocketServerFrame(
        encode({ ...acknowledgement, released: true }),
        65_536,
      ),
    ).toEqual({
      ok: false,
      error: { code: "MALFORMED_FRAME" },
    });
    expect(
      decodeWebSocketServerFrame(
        encode({ ...acknowledgement, subscriptionId: "" }),
        65_536,
      ),
    ).toEqual({
      ok: false,
      error: { code: "MALFORMED_FRAME" },
    });
  });

  it("rejects deeply nested or structurally excessive event envelopes without recursion", () => {
    let nested: unknown = "leaf";
    for (let depth = 0; depth < 40; depth += 1) {
      nested = [nested];
    }
    const event = (envelope: unknown) =>
      encode({
        type: "EVENT",
        protocol: REALTIME_WEBSOCKET_PROTOCOL,
        subscriptionId: "subscription.0001",
        envelope,
      });

    expect(
      decodeWebSocketServerFrame(
        event({ payload: nested }),
        65_536,
      ),
    ).toEqual({
      ok: false,
      error: { code: "MALFORMED_FRAME" },
    });
    expect(
      decodeWebSocketServerFrame(
        event({
          payload: Array.from({ length: 4_097 }, () => ({})),
        }),
        65_536,
      ),
    ).toEqual({
      ok: false,
      error: { code: "MALFORMED_FRAME" },
    });
  });

  it("encodes only closed client frames without leaking arbitrary commands", () => {
    const frame: WebSocketSubscribeFrame = {
      type: "SUBSCRIBE",
      protocol: REALTIME_WEBSOCKET_PROTOCOL,
      subscriptionId: "subscription.0001",
      streamId: "orders.v1",
      cursor: "cursor.0001",
      scopeBinding: "scope-binding.0001",
    };
    const result = encodeWebSocketClientFrame(frame, 65_536);

    expect(result.ok).toBe(true);
    if (!result.ok) throw new Error("Expected SUBSCRIBE to encode.");
    expect(JSON.parse(result.value)).toEqual(frame);
    expect(
      encodeWebSocketClientFrame(
        { ...frame, payload: "arbitrary" } as WebSocketSubscribeFrame,
        65_536,
      ),
    ).toEqual({
      ok: false,
      error: { code: "MALFORMED_FRAME" },
    });
  });
});