codex-grok-mcp 0.2.0-beta.1

package/relay/package.json

@@ -0,0 +1,17 @@
+{
+  "name": "@codex-grok-mcp/relay",
+  "version": "0.1.0",
+  "private": true,
+  "type": "module",
+  "scripts": {
+    "check": "tsc --noEmit && tsc --noEmit -p test/tsconfig.json",
+    "test": "vitest run",
+    "types": "wrangler types"
+  },
+  "devDependencies": {
+    "@cloudflare/vitest-plugin": "1.1.3",
+    "typescript": "7.0.2",
+    "vitest": "4.1.11",
+    "wrangler": "4.128.0"
+  }
+}

package/relay/src/index.ts

@@ -0,0 +1,182 @@
+import { DurableObject } from "cloudflare:workers";
+
+const CHANNEL_PATTERN = /^[A-Za-z0-9_-]{21}[AQgw]$/;
+const MAX_FRAME_BYTES = 128 * 1024;
+const INVALID_FRAME_CLOSE_CODE = 4400;
+const PEER_UNAVAILABLE_CLOSE_CODE = 4404;
+const TEXT_ENCODER = new TextEncoder();
+
+type Role = "codex" | "bridge";
+type RelayEnv = Env & { RELAY_ACCESS_TOKEN: string };
+
+function isRole(value: string | null): value is Role {
+  return value === "codex" || value === "bridge";
+}
+
+function close(socket: WebSocket, code: number, reason: string): void {
+  try {
+    socket.close(code, reason);
+  } catch {
+    // The peer may already be disconnected.
+  }
+}
+
+function frameSize(message: string | ArrayBuffer): number {
+  return typeof message === "string"
+    ? TEXT_ENCODER.encode(message).byteLength
+    : message.byteLength;
+}
+
+function encodeBase64Url(bytes: Uint8Array): string {
+  return btoa(String.fromCharCode(...bytes))
+    .replaceAll("+", "-")
+    .replaceAll("/", "_")
+    .replace(/=+$/, "");
+}
+
+function decodeMasterToken(value: unknown): Uint8Array | undefined {
+  if (typeof value !== "string" || !/^[A-Za-z0-9_-]+$/.test(value)) return undefined;
+  try {
+    const padded = `${value.replaceAll("-", "+").replaceAll("_", "/")}${"=".repeat(
+      (4 - (value.length % 4)) % 4,
+    )}`;
+    const bytes = Uint8Array.from(atob(padded), (character) => character.charCodeAt(0));
+    return bytes.byteLength === 32 && encodeBase64Url(bytes) === value ? bytes : undefined;
+  } catch {
+    return undefined;
+  }
+}
+
+async function channelAccessToken(master: Uint8Array, channel: string): Promise<string> {
+  const key = await crypto.subtle.importKey(
+    "raw",
+    master,
+    { name: "HMAC", hash: "SHA-256" },
+    false,
+    ["sign"],
+  );
+  const signature = await crypto.subtle.sign(
+    "HMAC",
+    key,
+    TEXT_ENCODER.encode(`codex-grok-mcp-relay:v1\0${channel}`),
+  );
+  return encodeBase64Url(new Uint8Array(signature));
+}
+
+async function isAuthorized(
+  request: Request,
+  master: Uint8Array,
+  channel: string,
+): Promise<boolean> {
+  const accessToken = await channelAccessToken(master, channel);
+  const [provided, expected] = await Promise.all([
+    crypto.subtle.digest("SHA-256", TEXT_ENCODER.encode(request.headers.get("Authorization") ?? "")),
+    crypto.subtle.digest("SHA-256", TEXT_ENCODER.encode(`Bearer ${accessToken}`)),
+  ]);
+  return crypto.subtle.timingSafeEqual(provided, expected);
+}
+
+export class Relay extends DurableObject<RelayEnv> {
+  fetch(request: Request): Response {
+    const role = new URL(request.url).searchParams.get("role");
+    if (!isRole(role) || request.headers.get("Upgrade")?.toLowerCase() !== "websocket") {
+      return new Response(null, { status: 400 });
+    }
+
+    for (const socket of this.ctx.getWebSockets(role)) {
+      close(socket, 4000, "replaced");
+    }
+
+    const [client, server] = Object.values(new WebSocketPair());
+    this.ctx.acceptWebSocket(server, [role]);
+    server.serializeAttachment(role);
+    return new Response(null, { status: 101, webSocket: client });
+  }
+
+  webSocketMessage(socket: WebSocket, message: string | ArrayBuffer): void {
+    const role: unknown = socket.deserializeAttachment();
+    if (!isRole(typeof role === "string" ? role : null)) {
+      close(socket, 1011, "invalid connection state");
+      return;
+    }
+
+    const peerRole: Role = role === "codex" ? "bridge" : "codex";
+    const peer = this.ctx
+      .getWebSockets(peerRole)
+      .find((candidate) => candidate.readyState === WebSocket.OPEN);
+
+    if (!peer) {
+      close(socket, PEER_UNAVAILABLE_CLOSE_CODE, "peer unavailable");
+      return;
+    }
+
+    if (frameSize(message) > MAX_FRAME_BYTES) {
+      close(socket, 1009, "frame too large");
+      return;
+    }
+
+    peer.send(message);
+  }
+
+  webSocketClose(socket: WebSocket, code: number, _reason: string, _wasClean: boolean): void {
+    const role: unknown = socket.deserializeAttachment();
+    if (role !== "bridge") return;
+
+    const peer = this.ctx
+      .getWebSockets("codex")
+      .find((candidate) => candidate.readyState === WebSocket.OPEN);
+    if (!peer) return;
+
+    const rejectedFrame = code === INVALID_FRAME_CLOSE_CODE;
+    close(
+      peer,
+      rejectedFrame ? INVALID_FRAME_CLOSE_CODE : PEER_UNAVAILABLE_CLOSE_CODE,
+      rejectedFrame ? "peer rejected frame" : "peer unavailable",
+    );
+  }
+
+  webSocketError(socket: WebSocket, _error: unknown): void {
+    close(socket, 1011, "websocket error");
+  }
+}
+
+export default {
+  async fetch(request: Request, env: RelayEnv): Promise<Response> {
+    if (request.headers.has("Origin")) {
+      return new Response(null, { status: 403 });
+    }
+
+    const url = new URL(request.url);
+    const match = /^\/v1\/connect\/([^/]+)$/.exec(url.pathname);
+
+    if (!match) {
+      return new Response(null, { status: 404 });
+    }
+    if (request.method !== "GET") {
+      return new Response(null, { status: 405 });
+    }
+    if (!CHANNEL_PATTERN.test(match[1])) {
+      return new Response(null, { status: 400 });
+    }
+
+    const roles = url.searchParams.getAll("role");
+    if (roles.length !== 1 || !isRole(roles[0])) {
+      return new Response(null, { status: 400 });
+    }
+    const master = decodeMasterToken(env.RELAY_ACCESS_TOKEN);
+    if (master === undefined) {
+      return new Response(null, { status: 503 });
+    }
+    if (!(await isAuthorized(request, master, match[1]))) {
+      return new Response(null, {
+        status: 401,
+        headers: { "WWW-Authenticate": 'Bearer realm="codex-grok-relay"' },
+      });
+    }
+    if (request.headers.get("Upgrade")?.toLowerCase() !== "websocket") {
+      return new Response(null, { status: 426 });
+    }
+
+    return env.RELAY.getByName(match[1]).fetch(request);
+  },
+} satisfies ExportedHandler<RelayEnv>;

package/relay/test/relay.test.ts

@@ -0,0 +1,127 @@
+import { env, exports } from "cloudflare:workers";
+import { listDurableObjectIds } from "cloudflare:test";
+import { expect, it } from "vitest";
+
+const CHANNEL = "AAAAAAAAAAAAAAAAAAAAAA";
+const UNAVAILABLE_CHANNEL = "AAAAAAAAAAAAAAAAAAAAAQ";
+const CAP_CHANNEL = "AAAAAAAAAAAAAAAAAAAABA";
+const AUTHORIZATIONS: Record<string, string> = {
+  [CHANNEL]: "Bearer oCZGsv3_Wa-CwwkEYxnF9Zvl9nzxLC3mTUDx4R23H2E",
+  [UNAVAILABLE_CHANNEL]: "Bearer chjihmYNxBTJTqJG2OdkC2-BY78eR8KObXxeN1rmwQc",
+  [CAP_CHANNEL]: "Bearer NhP4BGrMM7x-D5Rf7GSLEj3uIJenewrPcW_Vv3JC02w",
+};
+const AUTHORIZATION = AUTHORIZATIONS[CHANNEL] ?? "";
+
+function event(socket: WebSocket, type: "close"): Promise<CloseEvent>;
+function event(socket: WebSocket, type: "message"): Promise<MessageEvent>;
+function event(socket: WebSocket, type: "close" | "message"): Promise<CloseEvent | MessageEvent> {
+  return new Promise((resolve) => socket.addEventListener(type, resolve, { once: true }));
+}
+
+async function connect(channel: string, role: "codex" | "bridge"): Promise<WebSocket> {
+  const response = await exports.default.fetch(`https://relay.test/v1/connect/${channel}?role=${role}`, {
+    headers: { Authorization: AUTHORIZATIONS[channel] ?? "", Upgrade: "websocket" },
+  });
+  expect(response.status).toBe(101);
+  expect(response.webSocket).not.toBeNull();
+  const socket = response.webSocket;
+  if (!socket) throw new Error("missing WebSocket");
+  socket.accept();
+  socket.binaryType = "arraybuffer";
+  return socket;
+}
+
+it("rejects browser and unauthenticated requests before creating a Durable Object", async () => {
+  const url = `https://relay.test/v1/connect/${CHANNEL}?role=codex`;
+
+  const missing = await exports.default.fetch(url, { headers: { Upgrade: "websocket" } });
+  expect(missing.status).toBe(401);
+  expect(missing.headers.get("WWW-Authenticate")).toBe('Bearer realm="codex-grok-relay"');
+
+  expect(
+    (await exports.default.fetch(url, {
+      headers: { Authorization: "Bearer wrong-token", Upgrade: "websocket" },
+    })).status,
+  ).toBe(401);
+  expect(
+    (await exports.default.fetch(url, {
+      headers: { Authorization: AUTHORIZATION, Origin: "https://example.test", Upgrade: "websocket" },
+    })).status,
+  ).toBe(403);
+  expect(
+    (await exports.default.fetch(
+      `https://relay.test/v1/connect/${UNAVAILABLE_CHANNEL}?role=codex`,
+      { headers: { Authorization: AUTHORIZATION, Upgrade: "websocket" } },
+    )).status,
+  ).toBe(401);
+  expect(await listDurableObjectIds(env.RELAY)).toHaveLength(0);
+});
+
+it("validates, routes opaque frames, caps frames, reports availability, and replaces roles", async () => {
+  expect((await exports.default.fetch("https://relay.test/nope")).status).toBe(404);
+  expect(
+    (await exports.default.fetch("https://relay.test/v1/connect/not-a-channel?role=codex", {
+      headers: { Authorization: AUTHORIZATION, Upgrade: "websocket" },
+    })).status,
+  ).toBe(400);
+  expect(
+    (await exports.default.fetch("https://relay.test/v1/connect/AAAAAAAAAAAAAAAAAAAAAB?role=codex", {
+      headers: { Authorization: AUTHORIZATION, Upgrade: "websocket" },
+    })).status,
+  ).toBe(400);
+  expect(
+    (await exports.default.fetch(`https://relay.test/v1/connect/${CHANNEL}?role=nope`, {
+      headers: { Authorization: AUTHORIZATION, Upgrade: "websocket" },
+    })).status,
+  ).toBe(400);
+  expect(
+    (await exports.default.fetch(`https://relay.test/v1/connect/${CHANNEL}?role=codex`, {
+      headers: { Authorization: AUTHORIZATION },
+    })).status,
+  ).toBe(426);
+
+  const codex = await connect(CHANNEL, "codex");
+  const bridge = await connect(CHANNEL, "bridge");
+
+  const stringMessage = event(bridge, "message");
+  codex.send("opaque");
+  expect((await stringMessage).data).toBe("opaque");
+
+  const binaryMessage = event(codex, "message");
+  bridge.send(new Uint8Array([1, 2, 3]).buffer);
+  expect(Array.from(new Uint8Array((await binaryMessage).data as ArrayBuffer))).toEqual([1, 2, 3]);
+
+  const replaced = event(codex, "close");
+  const replacement = await connect(CHANNEL, "codex");
+  expect((await replaced).code).toBe(4000);
+
+  const replacementMessage = event(replacement, "message");
+  bridge.send("new socket");
+  expect((await replacementMessage).data).toBe("new socket");
+
+  const unavailable = await connect(UNAVAILABLE_CHANNEL, "codex");
+  const unavailableClose = event(unavailable, "close");
+  unavailable.send(new Uint8Array(128 * 1024 + 1));
+  const unavailableEvent = await unavailableClose;
+  expect(unavailableEvent.code).toBe(4404);
+  expect(unavailableEvent.reason).toBe("peer unavailable");
+
+  const cappedSender = await connect(CAP_CHANNEL, "codex");
+  const cappedPeer = await connect(CAP_CHANNEL, "bridge");
+  const cappedClose = event(cappedSender, "close");
+  let oversizedForwarded = false;
+  cappedPeer.addEventListener("message", () => {
+    oversizedForwarded = true;
+  });
+  cappedSender.send(new Uint8Array(128 * 1024 + 1));
+  expect((await cappedClose).code).toBe(1009);
+  expect(oversizedForwarded).toBe(false);
+
+  const rejectedByCompanion = event(replacement, "close");
+  bridge.close(4400, "untrusted detail");
+  const rejection = await rejectedByCompanion;
+  expect(rejection.code).toBe(4400);
+  expect(rejection.reason).toBe("peer rejected frame");
+
+  cappedPeer.close(1000, "done");
+});

package/relay/test/tsconfig.json

@@ -0,0 +1,13 @@
+{
+  "extends": "../tsconfig.json",
+  "compilerOptions": {
+    "moduleResolution": "Bundler",
+    "skipLibCheck": true,
+    "types": ["@cloudflare/vitest-plugin/types"]
+  },
+  "include": [
+    "./**/*.ts",
+    "../vitest.config.ts",
+    "../worker-configuration.d.ts"
+  ]
+}

package/relay/tsconfig.json

@@ -0,0 +1,15 @@
+{
+  "compilerOptions": {
+    "target": "ES2022",
+    "lib": ["ES2022"],
+    "module": "ESNext",
+    "moduleResolution": "Bundler",
+    "strict": true,
+    "noEmit": true,
+    "types": ["./worker-configuration.d.ts"]
+  },
+  "include": [
+    "src/**/*.ts",
+    "worker-configuration.d.ts"
+  ]
+}

package/relay/vitest.config.ts

@@ -0,0 +1,15 @@
+import { cloudflareTest } from "@cloudflare/vitest-plugin";
+import { defineConfig } from "vitest/config";
+
+export default defineConfig({
+  plugins: [
+    cloudflareTest({
+      miniflare: {
+        bindings: {
+          RELAY_ACCESS_TOKEN: "AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA",
+        },
+      },
+      wrangler: { configPath: "./wrangler.jsonc" },
+    }),
+  ],
+});