browser/sidecar/relay.ts

// Based on OpenClaw Browser Extension Relay (MIT License)
// https://github.com/nicepkg/openclaw
//
// Rewritten for Bun.serve() with multi-user support.

import type { ServerWebSocket } from 'bun';
import type {
  AttachedToTargetEvent,
  CdpCommand,
  CdpEvent,
  CdpResponse,
  ConnectedTarget,
  DetachedFromTargetEvent,
  ExtensionForwardCommandMessage,
  ExtensionForwardEventMessage,
  ExtensionMessage,
  ExtensionPongMessage,
} from './relay-types';
import { allPairings } from '../db/queries';
import { getRelayAuthHeader, resolveUserFromToken, restoreUserTokens } from './relay-auth';

const RELAY_AUTH_HEADER = getRelayAuthHeader();
const DEFAULT_EXTENSION_RECONNECT_GRACE_MS = 5_000;
const DEFAULT_EXTENSION_COMMAND_RECONNECT_WAIT_MS = 3_000;

type WSData = {
  kind: 'extension' | 'cdp';
  userId: number;
  token: string;
};

type UserRelay = {
  extensionWs: ServerWebSocket<WSData> | null;
  cdpClients: Set<ServerWebSocket<WSData>>;
  connectedTargets: Map<string, ConnectedTarget>;
  pingInterval: Timer | null;
  gracePeriodTimer: Timer | null;
  pendingExtension: Map<number, { resolve: (v: unknown) => void; reject: (e: Error) => void; timer: Timer }>;
  nextExtensionId: number;
  extensionReconnectWaiters: Set<(connected: boolean) => void>;
};

const userRelays = new Map<number, UserRelay>();

function getUserRelay(userId: number): UserRelay {
  let relay = userRelays.get(userId);
  if (!relay) {
    relay = {
      extensionWs: null,
      cdpClients: new Set(),
      connectedTargets: new Map(),
      pingInterval: null,
      gracePeriodTimer: null,
      pendingExtension: new Map(),
      nextExtensionId: 1,
      extensionReconnectWaiters: new Set(),
    };
    userRelays.set(userId, relay);
  }
  return relay;
}

function extensionConnected(relay: UserRelay): boolean {
  return relay.extensionWs !== null && relay.extensionWs.readyState === WebSocket.OPEN;
}

function flushExtensionReconnectWaiters(relay: UserRelay, connected: boolean) {
  if (relay.extensionReconnectWaiters.size === 0) return;
  const waiters = Array.from(relay.extensionReconnectWaiters);
  relay.extensionReconnectWaiters.clear();
  for (const waiter of waiters) waiter(connected);
}

function clearGracePeriodTimer(relay: UserRelay) {
  if (!relay.gracePeriodTimer) return;
  clearTimeout(relay.gracePeriodTimer);
  relay.gracePeriodTimer = null;
}

function closeCdpClientsAfterExtensionDisconnect(relay: UserRelay) {
  relay.connectedTargets.clear();
  for (const client of relay.cdpClients) {
    try {
      client.close(1011, 'extension disconnected');
    } catch {
      // ignore
    }
  }
  relay.cdpClients.clear();
  flushExtensionReconnectWaiters(relay, false);
}

function scheduleExtensionDisconnectCleanup(relay: UserRelay) {
  clearGracePeriodTimer(relay);
  relay.gracePeriodTimer = setTimeout(() => {
    relay.gracePeriodTimer = null;
    if (extensionConnected(relay)) return;
    closeCdpClientsAfterExtensionDisconnect(relay);
  }, DEFAULT_EXTENSION_RECONNECT_GRACE_MS);
}

async function waitForExtensionReconnect(relay: UserRelay, timeoutMs: number): Promise<boolean> {
  if (extensionConnected(relay)) return true;
  return await new Promise<boolean>((resolve) => {
    let settled = false;
    const waiter = (connected: boolean) => {
      if (settled) return;
      settled = true;
      clearTimeout(timer);
      relay.extensionReconnectWaiters.delete(waiter);
      resolve(connected);
    };
    const timer = setTimeout(() => waiter(false), timeoutMs);
    relay.extensionReconnectWaiters.add(waiter);
  });
}

function sendToExtension(relay: UserRelay, payload: ExtensionForwardCommandMessage): Promise<unknown> {
  const ws = relay.extensionWs;
  if (!ws || ws.readyState !== WebSocket.OPEN) {
    throw new Error('Chrome extension not connected');
  }
  ws.send(JSON.stringify(payload));
  return new Promise<unknown>((resolve, reject) => {
    const timer = setTimeout(() => {
      relay.pendingExtension.delete(payload.id);
      reject(new Error(`extension request timeout: ${payload.params.method}`));
    }, 30_000);
    relay.pendingExtension.set(payload.id, { resolve, reject, timer });
  });
}

function broadcastToCdpClients(relay: UserRelay, evt: CdpEvent) {
  const msg = JSON.stringify(evt);
  for (const ws of relay.cdpClients) {
    if (ws.readyState !== WebSocket.OPEN) continue;
    ws.send(msg);
  }
}

function sendResponseToCdp(ws: ServerWebSocket<WSData>, res: CdpResponse) {
  if (ws.readyState !== WebSocket.OPEN) return;
  ws.send(JSON.stringify(res));
}

function ensureTargetEventsForClient(relay: UserRelay, ws: ServerWebSocket<WSData>, mode: 'autoAttach' | 'discover') {
  for (const target of relay.connectedTargets.values()) {
    if (mode === 'autoAttach') {
      ws.send(
        JSON.stringify({
          method: 'Target.attachedToTarget',
          params: {
            sessionId: target.sessionId,
            targetInfo: { ...target.targetInfo, attached: true },
            waitingForDebugger: false,
          },
        } satisfies CdpEvent),
      );
    } else {
      ws.send(
        JSON.stringify({
          method: 'Target.targetCreated',
          params: { targetInfo: { ...target.targetInfo, attached: true } },
        } satisfies CdpEvent),
      );
    }
  }
}

async function routeCdpCommand(relay: UserRelay, cmd: CdpCommand): Promise<unknown> {
  switch (cmd.method) {
    case 'Browser.getVersion':
      return {
        protocolVersion: '1.3',
        product: 'Chrome/Officer-Browser-Relay',
        revision: '0',
        userAgent: 'Officer-Browser-Relay',
        jsVersion: 'V8',
      };
    case 'Browser.setDownloadBehavior':
      return {};
    case 'Target.setAutoAttach':
    case 'Target.setDiscoverTargets':
      return {};
    case 'Target.getTargets':
      return {
        targetInfos: Array.from(relay.connectedTargets.values()).map((t) => ({
          ...t.targetInfo,
          attached: true,
        })),
      };
    case 'Target.getTargetInfo': {
      const params = (cmd.params ?? {}) as { targetId?: string };
      const targetId = typeof params.targetId === 'string' ? params.targetId : undefined;
      if (targetId) {
        for (const t of relay.connectedTargets.values()) {
          if (t.targetId === targetId) return { targetInfo: t.targetInfo };
        }
      }
      if (cmd.sessionId && relay.connectedTargets.has(cmd.sessionId)) {
        const t = relay.connectedTargets.get(cmd.sessionId);
        if (t) return { targetInfo: t.targetInfo };
      }
      const first = Array.from(relay.connectedTargets.values())[0];
      return { targetInfo: first?.targetInfo };
    }
    case 'Target.attachToTarget': {
      const params = (cmd.params ?? {}) as { targetId?: string };
      const targetId = typeof params.targetId === 'string' ? params.targetId : undefined;
      if (!targetId) throw new Error('targetId required');
      for (const t of relay.connectedTargets.values()) {
        if (t.targetId === targetId) return { sessionId: t.sessionId };
      }
      throw new Error('target not found');
    }
    default: {
      const id = relay.nextExtensionId++;
      return await sendToExtension(relay, {
        id,
        method: 'forwardCDPCommand',
        params: {
          method: cmd.method,
          sessionId: cmd.sessionId,
          params: cmd.params,
        },
      });
    }
  }
}

function handleExtensionMessage(relay: UserRelay, data: string) {
  let parsed: ExtensionMessage | null = null;
  try {
    parsed = JSON.parse(data) as ExtensionMessage;
  } catch {
    return;
  }

  if (parsed && typeof parsed === 'object' && 'id' in parsed && typeof parsed.id === 'number') {
    const pending = relay.pendingExtension.get(parsed.id);
    if (!pending) return;
    relay.pendingExtension.delete(parsed.id);
    clearTimeout(pending.timer);
    if ('error' in parsed && typeof parsed.error === 'string' && parsed.error.trim()) {
      pending.reject(new Error(parsed.error));
    } else {
      pending.resolve(parsed.result);
    }
    return;
  }

  if (parsed && typeof parsed === 'object' && 'method' in parsed) {
    if ((parsed as ExtensionPongMessage).method === 'pong') return;
    if ((parsed as ExtensionForwardEventMessage).method !== 'forwardCDPEvent') return;

    const evt = parsed as ExtensionForwardEventMessage;
    const method = evt.params?.method;
    const params = evt.params?.params;
    const sessionId = evt.params?.sessionId;
    if (!method || typeof method !== 'string') return;

    if (method === 'Target.attachedToTarget') {
      const attached = (params ?? {}) as AttachedToTargetEvent;
      const targetType = attached?.targetInfo?.type ?? 'page';
      if (targetType !== 'page') return;
      if (attached?.sessionId && attached?.targetInfo?.targetId) {
        const prev = relay.connectedTargets.get(attached.sessionId);
        const nextTargetId = attached.targetInfo.targetId;
        const prevTargetId = prev?.targetId;
        const changedTarget = Boolean(prev && prevTargetId && prevTargetId !== nextTargetId);
        relay.connectedTargets.set(attached.sessionId, {
          sessionId: attached.sessionId,
          targetId: nextTargetId,
          targetInfo: attached.targetInfo,
        });
        if (changedTarget && prevTargetId) {
          broadcastToCdpClients(relay, {
            method: 'Target.detachedFromTarget',
            params: { sessionId: attached.sessionId, targetId: prevTargetId },
            sessionId: attached.sessionId,
          });
        }
        if (!prev || changedTarget) {
          broadcastToCdpClients(relay, { method, params, sessionId });
        }
        return;
      }
    }

    if (method === 'Target.detachedFromTarget') {
      const detached = (params ?? {}) as DetachedFromTargetEvent;
      if (detached?.sessionId) {
        relay.connectedTargets.delete(detached.sessionId);
      }
      broadcastToCdpClients(relay, { method, params, sessionId });
      return;
    }

    if (method === 'Target.targetInfoChanged') {
      const changed = (params ?? {}) as { targetInfo?: { targetId?: string; type?: string } };
      const targetInfo = changed?.targetInfo;
      const targetId = targetInfo?.targetId;
      if (targetId && (targetInfo?.type ?? 'page') === 'page') {
        for (const [sid, target] of relay.connectedTargets) {
          if (target.targetId !== targetId) continue;
          relay.connectedTargets.set(sid, {
            ...target,
            targetInfo: { ...target.targetInfo, ...(targetInfo as object) },
          });
        }
      }
    }

    broadcastToCdpClients(relay, { method, params, sessionId });
  }
}

function handleExtensionClose(relay: UserRelay, ws: ServerWebSocket<WSData>) {
  if (relay.extensionWs !== ws) return;
  if (relay.pingInterval) {
    clearInterval(relay.pingInterval);
    relay.pingInterval = null;
  }
  relay.extensionWs = null;
  for (const [, pending] of relay.pendingExtension) {
    clearTimeout(pending.timer);
    pending.reject(new Error('extension disconnected'));
  }
  relay.pendingExtension.clear();
  scheduleExtensionDisconnectCleanup(relay);
}

function handleExtensionOpen(relay: UserRelay, ws: ServerWebSocket<WSData>) {
  relay.extensionWs = ws;
  clearGracePeriodTimer(relay);
  flushExtensionReconnectWaiters(relay, true);

  relay.pingInterval = setInterval(() => {
    if (ws.readyState !== WebSocket.OPEN) return;
    ws.send(JSON.stringify({ method: 'ping' }));
  }, 5000);
}

async function handleCdpMessage(relay: UserRelay, ws: ServerWebSocket<WSData>, data: string) {
  let cmd: CdpCommand | null = null;
  try {
    cmd = JSON.parse(data) as CdpCommand;
  } catch {
    return;
  }
  if (!cmd || typeof cmd !== 'object') return;
  if (typeof cmd.id !== 'number' || typeof cmd.method !== 'string') return;

  if (!extensionConnected(relay)) {
    const reconnected = await waitForExtensionReconnect(relay, DEFAULT_EXTENSION_COMMAND_RECONNECT_WAIT_MS);
    if (!reconnected || !extensionConnected(relay)) {
      sendResponseToCdp(ws, {
        id: cmd.id,
        sessionId: cmd.sessionId,
        error: { message: 'Extension not connected' },
      });
      return;
    }
  }

  try {
    const result = await routeCdpCommand(relay, cmd);

    if (cmd.method === 'Target.setAutoAttach' && !cmd.sessionId) {
      ensureTargetEventsForClient(relay, ws, 'autoAttach');
    }
    if (cmd.method === 'Target.setDiscoverTargets') {
      const discover = (cmd.params ?? {}) as { discover?: boolean };
      if (discover.discover === true) {
        ensureTargetEventsForClient(relay, ws, 'discover');
      }
    }
    if (cmd.method === 'Target.attachToTarget') {
      const params = (cmd.params ?? {}) as { targetId?: string };
      const targetId = typeof params.targetId === 'string' ? params.targetId : undefined;
      if (targetId) {
        const target = Array.from(relay.connectedTargets.values()).find((t) => t.targetId === targetId);
        if (target) {
          ws.send(
            JSON.stringify({
              method: 'Target.attachedToTarget',
              params: {
                sessionId: target.sessionId,
                targetInfo: { ...target.targetInfo, attached: true },
                waitingForDebugger: false,
              },
            } satisfies CdpEvent),
          );
        }
      }
    }

    sendResponseToCdp(ws, { id: cmd.id, sessionId: cmd.sessionId, result });
  } catch (err) {
    sendResponseToCdp(ws, {
      id: cmd.id,
      sessionId: cmd.sessionId,
      error: { message: err instanceof Error ? err.message : String(err) },
    });
  }
}

function getTokenFromUrl(url: string): string | null {
  try {
    const parsed = new URL(url, 'http://localhost');
    return parsed.searchParams.get('token')?.trim() ?? null;
  } catch {
    return null;
  }
}

function getTokenFromHeaders(headers: Headers): string | null {
  return headers.get(RELAY_AUTH_HEADER)?.trim() ?? null;
}

function getTokenFromRequest(req: Request): string | null {
  return getTokenFromHeaders(req.headers) ?? getTokenFromUrl(req.url);
}

function authenticateRequest(req: Request): number | null {
  const token = getTokenFromRequest(req);
  if (!token) return null;
  return resolveUserFromToken(token);
}

let relayServer: ReturnType<typeof Bun.serve> | null = null;
let relayPort = 0;

export function getRelayPort(): number {
  return relayPort;
}

export function getUserRelayStatus(userId: number): { extensionConnected: boolean; targetCount: number } {
  const relay = userRelays.get(userId);
  if (!relay) return { extensionConnected: false, targetCount: 0 };
  return {
    extensionConnected: extensionConnected(relay),
    targetCount: relay.connectedTargets.size,
  };
}

export function getUserTargets(userId: number): ConnectedTarget[] {
  const relay = userRelays.get(userId);
  if (!relay) return [];
  return Array.from(relay.connectedTargets.values());
}

export async function startBrowserRelay(port: number) {
  relayPort = port;

  // Rebuild the token map from the plugin's own table. Every paired browser's token is derived rather
  // than stored, so this is what makes a restart transparent to an extension that is already configured.
  const entries = await allPairings();
  if (entries.length > 0) {
    restoreUserTokens(entries, port);
  }

  const server = Bun.serve<WSData>({
    port,
    hostname: '0.0.0.0',
    idleTimeout: 60,

    routes: {
      '/': (req: Request) => {
        if (req.method === 'HEAD') return new Response(null, { status: 200 });
        return new Response('OK', { status: 200, headers: { 'Content-Type': 'text/plain' } });
      },
    },

    fetch(req, server) {
      const url = new URL(req.url);
      const path = url.pathname;

      // CORS for chrome-extension origins
      const origin = req.headers.get('origin');
      const corsHeaders: Record<string, string> = {};
      if (origin) {
        corsHeaders['Access-Control-Allow-Origin'] = origin;
        corsHeaders['Vary'] = 'Origin';
      }

      if (req.method === 'OPTIONS') {
        return new Response(null, {
          status: 204,
          headers: {
            ...corsHeaders,
            'Access-Control-Allow-Methods': 'GET, PUT, POST, OPTIONS',
            'Access-Control-Allow-Headers': `Content-Type, ${RELAY_AUTH_HEADER}`,
            'Access-Control-Max-Age': '86400',
          },
        });
      }

      if (path === '/extension/status') {
        const connectedCount = Array.from(userRelays.values()).filter((r) => extensionConnected(r)).length;
        return Response.json({ connected: connectedCount > 0, users: connectedCount }, { headers: corsHeaders });
      }

      // Auth-protected endpoints
      const token = getTokenFromRequest(req);
      const userId = token ? resolveUserFromToken(token) : null;

      if (path === '/json/version' || path === '/json/version/') {
        if (userId === null) return new Response('Unauthorized', { status: 401, headers: corsHeaders });
        const relay = userRelays.get(userId);
        const payload: Record<string, unknown> = {
          Browser: 'Officer/browser-relay',
          'Protocol-Version': '1.3',
        };
        if (relay && extensionConnected(relay)) {
          const hostHeader = req.headers.get('host')?.trim() || `0.0.0.0:${port}`;
          payload.webSocketDebuggerUrl = `ws://${hostHeader}/cdp?token=${token}`;
        }
        return Response.json(payload, { headers: corsHeaders });
      }

      const listPaths = new Set(['/json', '/json/', '/json/list', '/json/list/']);
      if (listPaths.has(path)) {
        if (userId === null) return new Response('Unauthorized', { status: 401, headers: corsHeaders });
        const relay = userRelays.get(userId);
        const hostHeader = req.headers.get('host')?.trim() || `0.0.0.0:${port}`;
        const cdpWsUrl = `ws://${hostHeader}/cdp?token=${token}`;
        const list = relay
          ? Array.from(relay.connectedTargets.values()).map((t) => ({
              id: t.targetId,
              type: t.targetInfo.type ?? 'page',
              title: t.targetInfo.title ?? '',
              description: t.targetInfo.title ?? '',
              url: t.targetInfo.url ?? '',
              webSocketDebuggerUrl: cdpWsUrl,
            }))
          : [];
        return Response.json(list, { headers: corsHeaders });
      }

      const activateMatch = path.match(/^\/json\/activate\/(.+)$/);
      if (activateMatch) {
        if (userId === null) return new Response('Unauthorized', { status: 401, headers: corsHeaders });
        const targetId = decodeURIComponent(activateMatch[1]!).trim();
        if (!targetId) return new Response('targetId required', { status: 400, headers: corsHeaders });
        const relay = userRelays.get(userId);
        if (relay && extensionConnected(relay)) {
          void sendToExtension(relay, {
            id: relay.nextExtensionId++,
            method: 'forwardCDPCommand',
            params: { method: 'Target.activateTarget', params: { targetId } },
          }).catch(() => {});
        }
        return new Response('OK', { status: 200, headers: corsHeaders });
      }

      const closeMatch = path.match(/^\/json\/close\/(.+)$/);
      if (closeMatch) {
        if (userId === null) return new Response('Unauthorized', { status: 401, headers: corsHeaders });
        const targetId = decodeURIComponent(closeMatch[1]!).trim();
        if (!targetId) return new Response('targetId required', { status: 400, headers: corsHeaders });
        const relay = userRelays.get(userId);
        if (relay && extensionConnected(relay)) {
          void sendToExtension(relay, {
            id: relay.nextExtensionId++,
            method: 'forwardCDPCommand',
            params: { method: 'Target.closeTarget', params: { targetId } },
          }).catch(() => {});
        }
        return new Response('OK', { status: 200, headers: corsHeaders });
      }

      // WebSocket upgrades
      if (path === '/extension') {
        if (userId === null) return new Response('Unauthorized', { status: 401, headers: corsHeaders });
        const relay = getUserRelay(userId);
        // Close existing extension WS (stale or active) — last connection wins
        if (relay.extensionWs) {
          try {
            relay.extensionWs.close(4000, 'replaced by new connection');
          } catch {
            // ignore
          }
          relay.extensionWs = null;
        }
        const ok = server.upgrade(req, { data: { kind: 'extension', userId, token: token! } });
        if (!ok) return new Response('Upgrade failed', { status: 500, headers: corsHeaders });
        return undefined as unknown as Response;
      }

      if (path === '/cdp') {
        if (userId === null) return new Response('Unauthorized', { status: 401, headers: corsHeaders });
        const relay = userRelays.get(userId);
        if (!relay || !extensionConnected(relay)) {
          return new Response('Extension not connected', { status: 503, headers: corsHeaders });
        }
        const ok = server.upgrade(req, { data: { kind: 'cdp', userId, token: token! } });
        if (!ok) return new Response('Upgrade failed', { status: 500, headers: corsHeaders });
        return undefined as unknown as Response;
      }

      return new Response('Not found', { status: 404, headers: corsHeaders });
    },

    websocket: {
      open(ws) {
        const { kind, userId } = ws.data;
        const relay = getUserRelay(userId);

        if (kind === 'extension') {
          handleExtensionOpen(relay, ws);
        } else {
          relay.cdpClients.add(ws);
        }
      },

      message(ws, raw) {
        const { kind, userId } = ws.data;
        const relay = userRelays.get(userId);
        if (!relay) return;
        const data = typeof raw === 'string' ? raw : new TextDecoder().decode(raw as unknown as ArrayBuffer);

        if (kind === 'extension') {
          handleExtensionMessage(relay, data);
        } else {
          void handleCdpMessage(relay, ws, data);
        }
      },

      close(ws) {
        const { kind, userId } = ws.data;
        const relay = userRelays.get(userId);
        if (!relay) return;

        if (kind === 'extension') {
          handleExtensionClose(relay, ws);
        } else {
          relay.cdpClients.delete(ws);
        }
      },
    },
  });

  relayServer = server;
  return server;
}

export function stopBrowserRelay() {
  if (!relayServer) return;
  for (const [, relay] of userRelays) {
    clearGracePeriodTimer(relay);
    if (relay.pingInterval) clearInterval(relay.pingInterval);
    flushExtensionReconnectWaiters(relay, false);
    for (const [, pending] of relay.pendingExtension) {
      clearTimeout(pending.timer);
      pending.reject(new Error('server stopping'));
    }
    relay.pendingExtension.clear();
    try {
      relay.extensionWs?.close(1001, 'server stopping');
    } catch {
      // ignore
    }
    for (const ws of relay.cdpClients) {
      try {
        ws.close(1001, 'server stopping');
      } catch {
        // ignore
      }
    }
  }
  userRelays.clear();
  relayServer.stop();
  relayServer = null;
}