server.ts

import { randomUUID } from "node:crypto";
import * as http from "node:http";
import { URL } from "node:url";
import { WebSocketServer, WebSocket } from "ws";
import { attachActionsProxy } from "./cdp-proxy-actions.js";
import { validateApiKey, type UserContext } from "./auth-cache.js";
import { createDriverSession, stopDriverSession, type DriverSessionConfig } from "./driver-api.js";
import { createSessionToken, verifySessionToken } from "./session-token.js";
import {
  type SessionRecord,
  type ProxyMode,
  connectRedis,
  disconnectRedis,
  saveSession,
  getSession,
  deleteSession,
  updateSession,
  checkConcurrencyLimit,
  checkRateLimit,
  countAllSessions,
  findStaleSessions,
  isOwnedByThisReplica,
  getSessionReplicaId,
  REPLICA_ID,
} from "./session-store.js";

/* ------------------------------------------------------------------ */
/*  Config                                                             */
/* ------------------------------------------------------------------ */

const PORT = parseInt(process.env.PORT ?? "3000", 10);
const HOST = process.env.HOST?.trim() || "0.0.0.0";
const DEFAULT_COUNTRY = process.env.DEFAULT_DRIVER_COUNTRY?.trim() || "CA";
const DEFAULT_NODE_TYPE = process.env.DEFAULT_DRIVER_NODE_TYPE?.trim() || "hosted";
const DEFAULT_CAPTCHA_SOLVER =
  (process.env.DEFAULT_DRIVER_CAPTCHA_SOLVER ?? "true").trim().toLowerCase() !== "false";
const DEFAULT_IDLE_TIMEOUT_MS = parseInt(
  process.env.DEFAULT_SESSION_IDLE_TIMEOUT_MS ?? "120000",
  10,
);
const DEFAULT_MAX_DURATION_MS = parseInt(
  process.env.DEFAULT_SESSION_MAX_DURATION_MS ?? "1800000",
  10,
);
const REQUIRE_AUTH = (process.env.REQUIRE_AUTH ?? "true").trim() === "true";

/* ------------------------------------------------------------------ */
/*  Local session state (this replica only)                            */
/*  Keeps idle timers and live client counts that can't live in Redis. */
/* ------------------------------------------------------------------ */

interface LocalSession {
  id: string;
  orgId: string;
  apiKey: string;
  upstreamSessionId: string;
  upstreamCdpUrl: string;
  expiresAtMs: number;
  idleTimeoutMs: number;
  maxDurationMs: number;
  connectedClients: number;
  lastActivityMs: number;
  idleTimer?: NodeJS.Timeout;
  maxDurationTimer?: NodeJS.Timeout;
}

const localSessions = new Map<string, LocalSession>();

/* ------------------------------------------------------------------ */
/*  HTTP helpers                                                       */
/* ------------------------------------------------------------------ */

const wsServer = new WebSocketServer({ noServer: true });

function json(
  res: http.ServerResponse,
  statusCode: number,
  payload: Record<string, unknown>,
): void {
  const body = JSON.stringify(payload);
  res.writeHead(statusCode, {
    "Content-Type": "application/json",
    "Content-Length": Buffer.byteLength(body),
  });
  res.end(body);
}

function getBearerToken(req: http.IncomingMessage): string | null {
  const authHeader = req.headers.authorization;
  if (!authHeader?.startsWith("Bearer ")) return null;
  return authHeader.slice(7);
}

async function readJsonBody<T>(req: http.IncomingMessage): Promise<T> {
  const chunks: Buffer[] = [];
  for await (const chunk of req) {
    chunks.push(Buffer.isBuffer(chunk) ? chunk : Buffer.from(chunk));
  }
  const body = Buffer.concat(chunks).toString("utf8");
  return (body ? JSON.parse(body) : {}) as T;
}

async function authenticate(
  req: http.IncomingMessage,
): Promise<{ apiKey: string | null; user: UserContext | null }> {
  const apiKey = getBearerToken(req);
  if (!apiKey) {
    if (REQUIRE_AUTH) return { apiKey: null, user: null };
    return { apiKey: null, user: null };
  }
  const user = await validateApiKey(apiKey);
  return { apiKey, user };
}

function buildPublicWsUrl(req: http.IncomingMessage, sessionId: string, token: string): string {
  const host = req.headers.host ?? `127.0.0.1:${PORT}`;
  const forwardedProto = (req.headers["x-forwarded-proto"] as string | undefined)?.split(",")[0]?.trim();
  const proto = forwardedProto === "https" ? "wss" : "ws";
  return `${proto}://${host}/v1/proxy/connect/${sessionId}?token=${encodeURIComponent(token)}`;
}

function buildSessionResponse(req: http.IncomingMessage, session: SessionRecord) {
  const token = createSessionToken(session.id, session.owner.orgId, session.expiresAtMs);
  return {
    success: true as const,
    data: {
      sessionId: session.id,
      upstreamSessionId: session.upstreamSessionId,
      mode: session.mode,
      cdpUrl: buildPublicWsUrl(req, session.id, token),
      expiresAt: new Date(session.expiresAtMs).toISOString(),
      idleTimeoutMs: session.idleTimeoutMs,
      maxDurationMs: session.maxDurationMs,
    },
  };
}

/* ------------------------------------------------------------------ */
/*  Session lifecycle                                                  */
/* ------------------------------------------------------------------ */

function scheduleIdleReap(local: LocalSession): void {
  if (local.idleTimer) clearTimeout(local.idleTimer);
  local.idleTimer = setTimeout(() => {
    void destroySession(local.id, "idle_timeout");
  }, local.idleTimeoutMs);
  local.idleTimer.unref();
}

async function destroySession(sessionId: string, reason: string): Promise<void> {
  const local = localSessions.get(sessionId);
  if (!local) {
    // Not owned by this replica — just remove from Redis if it exists
    const session = await getSession(sessionId);
    if (session) {
      await deleteSession(sessionId, session.owner.orgId);
      try {
        await stopDriverSession(session.apiKey, session.upstreamSessionId);
      } catch (error) {
        console.error(
          `[hosted-proxy] stop upstream failed for ${sessionId}: ${
            error instanceof Error ? error.message : String(error)
          }`,
        );
      }
    }
    return;
  }

  localSessions.delete(sessionId);
  if (local.idleTimer) clearTimeout(local.idleTimer);
  if (local.maxDurationTimer) clearTimeout(local.maxDurationTimer);
  await deleteSession(sessionId, local.orgId);

  console.log(
    `[hosted-proxy] destroy session=${sessionId} upstream=${local.upstreamSessionId} reason=${reason} replica=${REPLICA_ID}`,
  );
  try {
    await stopDriverSession(local.apiKey, local.upstreamSessionId);
  } catch (error) {
    console.error(
      `[hosted-proxy] stop upstream failed for ${sessionId}: ${
        error instanceof Error ? error.message : String(error)
      }`,
    );
  }
}

/* ------------------------------------------------------------------ */
/*  Handlers                                                           */
/* ------------------------------------------------------------------ */

function validateCreatePayload(payload: any): {
  mode: ProxyMode;
  config: DriverSessionConfig;
  idleTimeoutMs: number;
  maxDurationMs: number;
} {
  const mode: ProxyMode = "actions";
  const idleTimeoutMs =
    typeof payload?.idleTimeoutMs === "number" ? payload.idleTimeoutMs : DEFAULT_IDLE_TIMEOUT_MS;
  const maxDurationMs =
    typeof payload?.maxDurationMs === "number" ? payload.maxDurationMs : DEFAULT_MAX_DURATION_MS;
  return {
    mode,
    config: {
      country: payload?.country ?? DEFAULT_COUNTRY,
      type: payload?.type ?? DEFAULT_NODE_TYPE,
      captchaSolver:
        typeof payload?.captchaSolver === "boolean"
          ? payload.captchaSolver
          : DEFAULT_CAPTCHA_SOLVER,
    },
    idleTimeoutMs: Math.max(10_000, idleTimeoutMs),
    maxDurationMs: Math.max(60_000, maxDurationMs),
  };
}

async function createHostedSession(
  req: http.IncomingMessage,
  res: http.ServerResponse,
  apiKey: string,
  user: UserContext,
): Promise<void> {
  // Check per-org limits before creating upstream session
  const [concurrency, rateLimit] = await Promise.all([
    checkConcurrencyLimit(user.orgId),
    checkRateLimit(user.orgId),
  ]);

  if (!concurrency.allowed) {
    json(res, 429, {
      success: false,
      error: `Concurrent session limit reached (${concurrency.current}/${concurrency.limit})`,
    });
    return;
  }
  if (!rateLimit.allowed) {
    json(res, 429, {
      success: false,
      error: `Hourly session creation limit reached (${rateLimit.current}/${rateLimit.limit})`,
    });
    return;
  }

  const payload = await readJsonBody<Record<string, unknown>>(req);
  const { mode, config, idleTimeoutMs, maxDurationMs } = validateCreatePayload(payload);
  const remote = await createDriverSession(apiKey, config);
  const sessionId = randomUUID();
  const now = Date.now();

  // Save to Redis (source of truth)
  const sessionRecord: SessionRecord = {
    id: sessionId,
    owner: user,
    apiKey,
    upstreamSessionId: remote.sessionId,
    upstreamCdpUrl: remote.cdpUrl,
    mode,
    createdAtMs: now,
    expiresAtMs: now + maxDurationMs,
    idleTimeoutMs,
    maxDurationMs,
    connectedClients: 0,
    lastActivityMs: now,
    replicaId: REPLICA_ID,
  };
  await saveSession(sessionRecord);

  // Track locally for timers and WebSocket management
  const local: LocalSession = {
    id: sessionId,
    orgId: user.orgId,
    apiKey,
    upstreamSessionId: remote.sessionId,
    upstreamCdpUrl: remote.cdpUrl,
    expiresAtMs: now + maxDurationMs,
    idleTimeoutMs,
    maxDurationMs,
    connectedClients: 0,
    lastActivityMs: now,
  };
  local.maxDurationTimer = setTimeout(() => {
    void destroySession(sessionId, "max_duration");
  }, maxDurationMs);
  local.maxDurationTimer.unref();
  localSessions.set(sessionId, local);
  scheduleIdleReap(local);

  json(res, 201, buildSessionResponse(req, sessionRecord));
}

async function handleDeleteSession(
  req: http.IncomingMessage,
  res: http.ServerResponse,
  apiKey: string | null,
  user: UserContext | null,
  sessionId: string,
): Promise<void> {
  const session = await getSession(sessionId);
  if (!session) {
    json(res, 404, { success: false, error: "Session not found" });
    return;
  }
  if (!apiKey || !user || session.owner.orgId !== user.orgId) {
    json(res, 403, { success: false, error: "Forbidden" });
    return;
  }
  await destroySession(sessionId, "client_delete");
  json(res, 200, { success: true });
}

async function handleGetSession(
  res: http.ServerResponse,
  user: UserContext | null,
  sessionId: string,
): Promise<void> {
  const session = await getSession(sessionId);
  if (!session) {
    json(res, 404, { success: false, error: "Session not found" });
    return;
  }
  if (!user || session.owner.orgId !== user.orgId) {
    json(res, 403, { success: false, error: "Forbidden" });
    return;
  }
  json(res, 200, {
    success: true,
    data: {
      sessionId: session.id,
      upstreamSessionId: session.upstreamSessionId,
      mode: session.mode,
      connectedClients: session.connectedClients,
      createdAt: new Date(session.createdAtMs).toISOString(),
      expiresAt: new Date(session.expiresAtMs).toISOString(),
      lastActivityAt: new Date(session.lastActivityMs).toISOString(),
      replicaId: session.replicaId,
    },
  });
}

/* ------------------------------------------------------------------ */
/*  HTTP server                                                        */
/* ------------------------------------------------------------------ */

const server = http.createServer(async (req, res) => {
  try {
    const url = new URL(req.url ?? "/", `http://${req.headers.host ?? `127.0.0.1:${PORT}`}`);

    if (req.method === "GET" && url.pathname === "/health") {
      const activeSessions = await countAllSessions();
      json(res, 200, {
        status: "ok",
        activeSessions,
        localSessions: localSessions.size,
        replicaId: REPLICA_ID,
      });
      return;
    }

    if (req.method === "POST" && url.pathname === "/v1/proxy/session") {
      const { apiKey, user } = await authenticate(req);
      if (!apiKey || !user) {
        json(res, 401, { success: false, error: "Unauthorized" });
        return;
      }
      await createHostedSession(req, res, apiKey, user);
      return;
    }

    const match = url.pathname.match(/^\/v1\/proxy\/session\/([^/]+)$/);
    if (match && req.method === "DELETE") {
      const { apiKey, user } = await authenticate(req);
      await handleDeleteSession(req, res, apiKey, user, match[1]);
      return;
    }
    if (match && req.method === "GET") {
      const { user } = await authenticate(req);
      await handleGetSession(res, user, match[1]);
      return;
    }

    json(res, 404, { success: false, error: "Not found" });
  } catch (error) {
    json(res, 500, {
      success: false,
      error: error instanceof Error ? error.message : "Unknown error",
    });
  }
});

/* ------------------------------------------------------------------ */
/*  WebSocket upgrade                                                  */
/* ------------------------------------------------------------------ */

server.on("upgrade", async (req, socket, head) => {
  try {
    const upgradeSocket = socket as any;
    if (typeof upgradeSocket.setNoDelay === "function") {
      upgradeSocket.setNoDelay(true);
    }
    const url = new URL(req.url ?? "/", `http://${req.headers.host ?? `127.0.0.1:${PORT}`}`);
    const match = url.pathname.match(/^\/v1\/proxy\/connect\/([^/]+)$/);
    if (!match) {
      socket.destroy();
      return;
    }
    const sessionId = match[1];
    const token = url.searchParams.get("token");
    if (!token) {
      socket.destroy();
      return;
    }
    const verified = verifySessionToken(token);
    if (!verified || verified.sessionId !== sessionId) {
      socket.destroy();
      return;
    }

    // Check local sessions first (fast path for sessions on this replica)
    const local = localSessions.get(sessionId);
    if (!local) {
      // Session not on this replica — could be on another replica or not exist
      const session = await getSession(sessionId);
      if (!session || verified.orgId !== session.owner.orgId) {
        socket.destroy();
        return;
      }
      // Session exists but is on a different replica
      if (session.replicaId !== REPLICA_ID) {
        // Return a 421 Misdirected Request so the router knows to redirect
        socket.write(
          "HTTP/1.1 421 Misdirected Request\r\n" +
          `X-Session-Replica: ${session.replicaId}\r\n` +
          "\r\n",
        );
        socket.destroy();
        return;
      }
      // Session is supposedly ours but not in local map — stale
      socket.destroy();
      return;
    }

    if (verified.orgId !== local.orgId) {
      socket.destroy();
      return;
    }
    if (Date.now() > local.expiresAtMs) {
      void destroySession(sessionId, "expired_before_connect");
      socket.destroy();
      return;
    }

    wsServer.handleUpgrade(req, socket, head, (clientWs) => {
      const upstreamWs = attachActionsProxy(clientWs, local.upstreamCdpUrl, {
        label: "cdp-proxy-actions",
      });
      local.connectedClients += 1;
      local.lastActivityMs = Date.now();
      if (local.idleTimer) clearTimeout(local.idleTimer);
      void updateSession(sessionId, {
        connectedClients: local.connectedClients,
        lastActivityMs: local.lastActivityMs,
      });

      clientWs.on("message", () => {
        local.lastActivityMs = Date.now();
      });
      upstreamWs.on("message", () => {
        local.lastActivityMs = Date.now();
      });

      const closeBoth = (reason: string) => {
        local.connectedClients = Math.max(0, local.connectedClients - 1);
        local.lastActivityMs = Date.now();
        if (local.connectedClients === 0) scheduleIdleReap(local);
        void updateSession(sessionId, {
          connectedClients: local.connectedClients,
          lastActivityMs: local.lastActivityMs,
        });
        console.log(`[hosted-proxy] ws close session=${sessionId} reason=${reason} replica=${REPLICA_ID}`);
      };

      clientWs.on("close", () => closeBoth("client_closed"));
      upstreamWs.on("close", () => closeBoth("upstream_closed"));
      clientWs.on("error", () => closeBoth("client_error"));
      upstreamWs.on("error", () => closeBoth("upstream_error"));
    });
  } catch {
    socket.destroy();
  }
});

/* ------------------------------------------------------------------ */
/*  Graceful shutdown                                                  */
/* ------------------------------------------------------------------ */

let shuttingDown = false;

async function gracefulShutdown(signal: string): Promise<void> {
  if (shuttingDown) return;
  shuttingDown = true;
  console.log(`[hosted-proxy] received ${signal}, starting graceful shutdown...`);

  // Stop accepting new connections
  server.close();

  // Destroy all local sessions (this stops upstream browser sessions)
  const destroyPromises: Promise<void>[] = [];
  for (const [sessionId] of localSessions) {
    destroyPromises.push(destroySession(sessionId, "shutdown"));
  }
  await Promise.allSettled(destroyPromises);

  // Disconnect Redis
  await disconnectRedis();

  console.log("[hosted-proxy] shutdown complete");
  process.exit(0);
}

process.on("SIGTERM", () => void gracefulShutdown("SIGTERM"));
process.on("SIGINT", () => void gracefulShutdown("SIGINT"));

/* ------------------------------------------------------------------ */
/*  Startup                                                            */
/* ------------------------------------------------------------------ */

async function cleanupStaleSessions(): Promise<void> {
  const stale = await findStaleSessions();
  if (stale.length === 0) return;
  console.log(`[hosted-proxy] cleaning up ${stale.length} stale session(s) from previous instance`);
  for (const session of stale) {
    await deleteSession(session.id, session.owner.orgId);
    try {
      await stopDriverSession(session.apiKey, session.upstreamSessionId);
    } catch {
      // Best effort
    }
  }
}

async function start(): Promise<void> {
  await connectRedis();
  await cleanupStaleSessions();
  server.listen(PORT, HOST, () => {
    console.log(`[hosted-proxy] listening on http://${HOST}:${PORT} replica=${REPLICA_ID}`);
  });
}

start().catch((err) => {
  console.error(`[hosted-proxy] failed to start: ${err.message}`);
  process.exit(1);
});