continuum-backend/tests/edge_upgrade.test.ts

import { describe, expect, it } from "bun:test";
import { Elysia } from "elysia";
import jwt from "@elysiajs/jwt";
import { UnauthorizedError } from "@continuum/ts-core";
import { farmFromEdgeUpgrade } from "../src/auth/edge_upgrade";

/**
 * The gateway protocol's authentication boundary.
 *
 * The edge JWT travels in `Authorization: Bearer` on the upgrade request and
 * is verified before any socket exists. It used to be a `token` field on the
 * `hello` message, which meant an unauthenticated peer got a live socket and
 * a parsed frame before proving anything. These tests are what stops that
 * coming back.
 */
describe("farmFromEdgeUpgrade", () => {
  const verifyOk = async (token: string) =>
    token === "good" ? { scope: "edge", farmId: "farm-1" } : false;

  it("returns the farm from a valid edge token", async () => {
    expect(await farmFromEdgeUpgrade("Bearer good", verifyOk)).toBe("farm-1");
  });

  it("refuses a missing or malformed Authorization header", async () => {
    for (const header of [undefined, null, "", "good", "Basic good", "Bearer"]) {
      await expect(farmFromEdgeUpgrade(header, verifyOk)).rejects.toThrow(UnauthorizedError);
    }
  });

  it("refuses a token the verifier rejects", async () => {
    await expect(farmFromEdgeUpgrade("Bearer forged", verifyOk)).rejects.toThrow(UnauthorizedError);
  });

  /**
   * A signature alone is not authorisation. The same secret could be made to
   * sign something else one day, and such a token must not open a farm's
   * gateway socket.
   */
  it("refuses a correctly signed token minted for a different scope", async () => {
    const verify = async () => ({ scope: "user", farmId: "farm-1" });
    await expect(farmFromEdgeUpgrade("Bearer good", verify)).rejects.toThrow(UnauthorizedError);
  });

  /**
   * Without this, the connection's `farmId` would be `undefined`, which to
   * anything doing a lookup reads as "every farm" rather than as an error.
   */
  it("refuses a token with no usable farm claim", async () => {
    for (const farmId of [undefined, null, "", 42, {}]) {
      const verify = async () => ({ scope: "edge", farmId }) as Record<string, unknown>;
      await expect(farmFromEdgeUpgrade("Bearer good", verify)).rejects.toThrow(UnauthorizedError);
    }
  });

  /**
   * The messages must not say which part failed - that tells the holder of a
   * forged token what to fix next. Asserted rather than assumed, because the
   * natural way to write this code is one specific message per branch.
   */
  it("says the same thing however the token fails", async () => {
    const messages = new Set<string>();
    for (const verify of [
      async () => false as const,
      async () => ({ scope: "user", farmId: "farm-1" }),
      async () => ({ scope: "edge" }) as Record<string, unknown>,
    ]) {
      await farmFromEdgeUpgrade("Bearer x", verify).catch((err: Error) =>
        messages.add(err.message),
      );
    }
    expect(messages.size).toBe(1);
  });
});

/**
 * That the check above actually gates the upgrade, which is the property the
 * whole design rests on and is a framework behaviour rather than ours: a
 * `resolve` that throws must fail the HTTP upgrade so `open` never runs.
 *
 * This builds its own Elysia app rather than importing `edgeWsRoute`, which
 * pulls in `../../db` and so a live Postgres. What it proves is that this
 * composition refuses a socket and that a resolved value reaches `ws.data`;
 * what it cannot prove is that `edge.ts` is wired this way. `scripts/
 * e2e-edge.sh` runs the real route against a real database and a real
 * gateway, and is the test for that.
 */
describe("a resolve that throws gates the websocket upgrade", () => {
  const SECRET = "edge-upgrade-test-secret";

  const app = new Elysia()
    .use(jwt({ name: "edgeJwt", secret: SECRET }))
    .get("/mint", ({ edgeJwt, query }) =>
      edgeJwt.sign({ scope: query.scope ?? "edge", farmId: "farm-1" }),
    )
    .resolve({ as: "scoped" }, async ({ headers, edgeJwt }): Promise<{ edgeFarmId: string }> => ({
      edgeFarmId: await farmFromEdgeUpgrade(headers.authorization, (t) => edgeJwt.verify(t)),
    }))
    .ws("/ws/edge/v1", {
      open(ws) {
        ws.send(JSON.stringify({ farm: (ws.data as { edgeFarmId: string }).edgeFarmId }));
      },
    })
    .listen(0);

  const base = `localhost:${app.server!.port}`;

  function attempt(headers: Record<string, string>): Promise<string> {
    return new Promise((resolve) => {
      const ws = new WebSocket(`ws://${base}/ws/edge/v1`, { headers } as never);
      const done = (r: string) => {
        try {
          ws.close();
        } catch {
          // Already closed by the server; nothing to do.
        }
        resolve(r);
      };
      ws.onmessage = (event) => done(`opened:${event.data}`);
      ws.onerror = () => done("refused");
      ws.onclose = () => done("refused");
      setTimeout(() => done("timed out"), 5_000);
    });
  }

  it("opens for a valid edge token and exposes its farm", async () => {
    const token = await (await fetch(`http://${base}/mint`)).text();
    expect(await attempt({ Authorization: `Bearer ${token}` })).toBe('opened:{"farm":"farm-1"}');
  });

  it("never opens a socket without a valid edge token", async () => {
    const wrongScope = await (await fetch(`http://${base}/mint?scope=user`)).text();
    expect(await attempt({})).toBe("refused");
    expect(await attempt({ Authorization: "Bearer not-a-jwt" })).toBe("refused");
    expect(await attempt({ Authorization: `Bearer ${wrongScope}` })).toBe("refused");
  });
});