platform/src/servers/api/vault/websocket.ts

import type { ServerWebSocket } from 'bun';
import { resolveAuthToken } from '../../auth-token';
import { isTokenBlacklisted } from 'officerdb';
import { isSuperAdmin } from '../../super-admin';
import { getVaultServerWsUrl } from './sidecar-server';
import { getValidAccessToken } from './token-store';

// Platform side of the Bitwarden notifications WebSocket. The device connects with its platform JWT (via
// ?access_token=, how SignalR carries the token); we validate that session in `open`, then pipe the socket
// to the officer-vault sidecar with the stored Vaultwarden token injected — the device never holds it.
// Dumb pipe: text + binary frames both ways, no inspection. SignalR's HTTP long-poll fallback rides the
// HTTP proxy (vaultRouter) instead.

export type VaultWSData = {
  provider: 'vault';
  platformToken: string; // the platform JWT the device presented on the upgrade
  vaultWsPath: string; // path + query after /api/vault, e.g. /notifications/hub?access_token=<jwt>
  vaultWsProtocol?: string; // requested Sec-WebSocket-Protocol, forwarded to the sidecar
};

type UpstreamState = {
  ws: WebSocket | null;
  queue: (string | Uint8Array<ArrayBuffer>)[];
  ready: boolean;
  closed: boolean;
};
const upstreams = new Map<ServerWebSocket<VaultWSData>, UpstreamState>();

// Bun hands frames over as `string | Buffer`; a Buffer is a Uint8Array at runtime, so forward as-is.
const asPayload = (raw: string | Buffer): string | Uint8Array<ArrayBuffer> =>
  typeof raw === 'string' ? raw : (raw as Uint8Array<ArrayBuffer>);

function closeClient(ws: ServerWebSocket<VaultWSData>, state: UpstreamState, code: number, reason: string) {
  state.closed = true;
  upstreams.delete(ws);
  try {
    ws.close(code, reason);
  } catch {
    /* already closed */
  }
}

// Swap the platform JWT in the path for the Vaultwarden access token before dialing upstream.
function injectToken(wsBase: string, vaultWsPath: string, vaultToken: string): string {
  const qIdx = vaultWsPath.indexOf('?');
  const path = qIdx >= 0 ? vaultWsPath.slice(0, qIdx) : vaultWsPath;
  const params = new URLSearchParams(qIdx >= 0 ? vaultWsPath.slice(qIdx + 1) : '');
  params.delete('token');
  params.set('access_token', vaultToken);
  return `${wsBase}${path}?${params.toString()}`;
}

export const vaultWebsocket = {
  async open(ws: ServerWebSocket<VaultWSData>) {
    // Register state synchronously so frames sent during the async setup below are buffered, not dropped.
    const state: UpstreamState = { ws: null, queue: [], ready: false, closed: false };
    upstreams.set(ws, state);

    // Deferred session validation (Bun requires the upgrade itself to be synchronous).
    let userId: number;
    try {
      const payload = await resolveAuthToken(ws.data.platformToken);
      if (!payload?.id) return closeClient(ws, state, 4001, 'Unauthorized');
      if (payload.jti && (await isTokenBlacklisted(payload.jti))) return closeClient(ws, state, 4001, 'Unauthorized');
      if (!(await isSuperAdmin(payload))) return closeClient(ws, state, 4001, 'Forbidden');
      userId = payload.id;
    } catch {
      return closeClient(ws, state, 4001, 'Unauthorized');
    }
    if (state.closed) return;

    const vaultToken = await getValidAccessToken(userId);
    if (!vaultToken) return closeClient(ws, state, 4001, 'No vault session');
    const wsBase = getVaultServerWsUrl();
    if (!wsBase) return closeClient(ws, state, 1011, 'Vault sidecar not available');
    if (state.closed) return;

    const protocols = ws.data.vaultWsProtocol
      ? ws.data.vaultWsProtocol
          .split(',')
          .map((s) => s.trim())
          .filter(Boolean)
      : undefined;
    const url = injectToken(wsBase, ws.data.vaultWsPath, vaultToken);
    const upstream = protocols?.length ? new WebSocket(url, protocols) : new WebSocket(url);
    upstream.binaryType = 'arraybuffer';
    state.ws = upstream;

    upstream.addEventListener('open', () => {
      state.ready = true;
      for (const m of state.queue) upstream.send(m);
      state.queue.length = 0;
    });
    upstream.addEventListener('message', (ev) => {
      try {
        ws.send(ev.data as string | ArrayBuffer);
      } catch {
        /* client gone */
      }
    });
    upstream.addEventListener('close', (ev) => {
      upstreams.delete(ws);
      try {
        ws.close(ev.code || 1000, ev.reason || '');
      } catch {
        /* already closed */
      }
    });
    upstream.addEventListener('error', () => {
      upstreams.delete(ws);
      try {
        ws.close(1011, 'upstream error');
      } catch {
        /* already closed */
      }
    });
  },
  message(ws: ServerWebSocket<VaultWSData>, raw: string | Buffer) {
    const state = upstreams.get(ws);
    if (!state) return;
    const payload = asPayload(raw);
    if (state.ready && state.ws) state.ws.send(payload);
    else state.queue.push(payload); // buffer until the upstream socket opens
  },
  close(ws: ServerWebSocket<VaultWSData>) {
    const state = upstreams.get(ws);
    if (state) {
      state.closed = true;
      try {
        state.ws?.close();
      } catch {
        /* already closed */
      }
      upstreams.delete(ws);
    }
  },
  drain() {},
};

const PREFIX = '/api/vault';

// Serve-level upgrade for /api/vault/notifications/* WebSockets. The platform JWT rides the query
// (?access_token= for SignalR, or ?token=) and the session is validated in `open` (deferred). The device
// never sends a Vaultwarden token — we inject the stored one upstream.
//
// There was an isOriginAllowed gate here until 2026-08-13, removed with the rest of origin validation.
// It had defaulted to allow-everything, so it refused nothing on a real install.
export function upgradeVaultWs(req: Request, server: any): Response | undefined {
  const url = new URL(req.url);
  const platformToken = url.searchParams.get('access_token') || url.searchParams.get('token') || '';
  if (!platformToken) return new Response('Unauthorized', { status: 401 });

  const data: VaultWSData = {
    provider: 'vault',
    platformToken,
    vaultWsPath: url.pathname.slice(PREFIX.length) + url.search,
    vaultWsProtocol: req.headers.get('sec-websocket-protocol') ?? undefined,
  };
  const ok = server.upgrade(req, { data });
  if (!ok) return new Response('Upgrade failed', { status: 500 });
  return undefined;
}