server.ts

import { createServer } from "node:http";
import type { IncomingMessage } from "node:http";
import { timingSafeEqual } from "node:crypto";
import { createClerkClient, verifyToken } from "@clerk/backend";
import { WebSocketServer, WebSocket, type RawData } from "ws";
import { z } from "zod";

const production = process.argv.includes("--production");
const realtimeOnly = process.argv.includes("--realtime");
const port = Number(process.env.PORT ?? (realtimeOnly ? 3001 : 3000));
const wss = new WebSocketServer({ noServer: true, maxPayload: 16_384 });
const clerk = createClerkClient({
  secretKey: process.env.CLERK_SECRET_KEY,
  publishableKey: process.env.NEXT_PUBLIC_CLERK_PUBLISHABLE_KEY,
});

type Session = {
  userId: string;
  subscriptions: Set<string>;
  watchedPeers: Map<string, string>;
  watchedTeams: Map<string, Set<string>>;
};
const sessions = new Map<WebSocket, Session>();
const MAX_CONNECTIONS = 1_000;
const MAX_CONNECTIONS_PER_USER = 5;
const MAX_SUBSCRIPTIONS = 20;
type ChatStore = Pick<typeof import("./src/lib/teams/writes"), "getMessagesForUser" | "postMessageForUser" | "canReadConversation">;
let chatStore: ChatStore;

const send = (socket: WebSocket, payload: unknown) => {
  if (socket.readyState === WebSocket.OPEN) socket.send(JSON.stringify(payload));
};

const fanOut = async (convoId: string, payload: unknown, except?: WebSocket) => {
  await Promise.all([...sessions].map(async ([socket, session]) => {
    if (socket === except || !session.subscriptions.has(convoId)) return;
    try {
      if (await chatStore.canReadConversation(convoId, session.userId)) send(socket, payload);
      else session.subscriptions.delete(convoId);
    } catch {
      session.subscriptions.delete(convoId);
    }
  }));
};

const broadcastPresence = (convoId: string, userId: string, online: boolean, except?: WebSocket) => {
  for (const [socket, session] of sessions) {
    if (socket !== except && session.subscriptions.has(convoId)) {
      send(socket, { t: "presence", convoId, userId, online });
    }
  }
};

const hasUserConnection = (userId: string) =>
  [...sessions.values()].some((session) => session.userId === userId);

/** Notify every open DM and group chat that watches this user, regardless of which page the
 * watched user is viewing. Presence is dashboard-level activity, not
 * conversation-level activity. */
const broadcastUserPresence = (userId: string, online: boolean) => {
  for (const [socket, session] of sessions) {
    for (const [convoId, peerId] of session.watchedPeers) {
      if (peerId === userId) send(socket, { t: "presence", convoId, userId, online });
    }
    for (const [convoId, teamPeers] of session.watchedTeams) {
      if (teamPeers.has(userId)) send(socket, { t: "presence", convoId, userId, online });
    }
  }
};

const notifyUser = (userId: string, frame: unknown, excludeSubscribedTo?: string) => {
  for (const [socket, session] of sessions) {
    if (session.userId === userId) {
      if (excludeSubscribedTo && session.subscriptions.has(excludeSubscribedTo)) {
        continue;
      }
      send(socket, frame);
    }
  }
};

const EVENT_TYPES = [
  "request.received",
  "request.accepted",
  "request.rejected",
  "invite.received",
  "invite.responded",
  "request.withdrawn",
  "team.updated",
  "team.disbanded",
  "member.left",
  "member.removed",
  "leadership.transferred",
  "message.received",
] as const;

const NotifyPayloadSchema = z.object({
  id: z.string().min(1).max(100),
  type: z.enum(EVENT_TYPES),
  userIds: z.array(z.string().min(1).max(200)).min(1).max(50),
  data: z.record(z.string(), z.unknown()),
  at: z.string(),
  excludeSubscribedTo: z.string().optional(),
});

const notifyWindows = new Map<string, { count: number; windowStart: number }>();
const NOTIFY_RATE_LIMIT = 300;
const NOTIFY_WINDOW_MS = 60_000;

function verifyNotifySecret(authHeader: string | undefined): boolean {
  const secret = process.env.REALTIME_EVENT_SECRET;
  if (!secret || !authHeader || !authHeader.startsWith("Bearer ")) return false;
  const token = authHeader.slice(7).trim();
  if (token.length !== secret.length) return false;
  try {
    return timingSafeEqual(Buffer.from(token), Buffer.from(secret));
  } catch {
    return false;
  }
}

const parseFrame = (raw: RawData): Record<string, unknown> | null => {
  try {
    const value = JSON.parse(raw.toString());
    return value && typeof value === "object" ? value as Record<string, unknown> : null;
  } catch {
    return null;
  }
};

const validConvoId = (value: unknown): value is string =>
  typeof value === "string" && value.length > 0 && value.length <= 200;

async function authenticate(requestUrl: string | undefined): Promise<string | null> {
  try {
    const token = new URL(requestUrl ?? "/", "http://localhost").searchParams.get("token");
    if (!token) return null;
    const claims = await verifyToken(token, {
      secretKey: process.env.CLERK_SECRET_KEY,
      clockSkewInMs: 120_000,
    });
    return typeof claims?.sub === "string" ? claims.sub : null;
  } catch {
    return null;
  }
}

async function authenticateRequestHeaders(request: IncomingMessage): Promise<string | null> {
  try {
    const headers = new Headers();
    for (const [name, value] of Object.entries(request.headers)) {
      if (value) headers.set(name, Array.isArray(value) ? value.join(", ") : value);
    }
    const protocol = request.headers["x-forwarded-proto"] === "https" ? "https" : "http";
    const state = await clerk.authenticateRequest(new Request(`${protocol}://${request.headers.host ?? "localhost"}/ws`, { headers }), {
      acceptsToken: "session_token",
      clockSkewInMs: 120_000,
    });
    return state.isAuthenticated ? state.toAuth().userId : null;
  } catch {
    return null;
  }
}

wss.on("connection", (socket: WebSocket, _request: IncomingMessage, session: Session) => {
  const firstConnection = !hasUserConnection(session.userId);
  sessions.set(socket, session);
  if (firstConnection) broadcastUserPresence(session.userId, true);
  const subscriptionVersions = new Map<string, number>();
  let lastActivity = Date.now();
  let messageWindow = 0;
  let windowStarted = Date.now();
  const idle = setInterval(() => {
    if (Date.now() - lastActivity > 10 * 60_000) socket.close(1000, "idle");
  }, 60_000);
  send(socket, { t: "ready" });

  socket.on("message", async (raw: RawData) => {
    if (raw.toString().length > 16_384) {
      socket.close(1009, "frame too large");
      return;
    }
    lastActivity = Date.now();
    if (Date.now() - windowStarted >= 10_000) { windowStarted = Date.now(); messageWindow = 0; }
    if (++messageWindow > 100) { socket.close(1008, "rate limit"); return; }
    const frame = parseFrame(raw);
    if (!frame || typeof frame.t !== "string") return;
    const convoId = frame.convoId;

    try {
      if (frame.t === "ping") return;
      if (frame.t === "unsubscribe") {
        if (!validConvoId(convoId)) return;
        session.subscriptions.delete(convoId);
        session.watchedPeers.delete(convoId);
        session.watchedTeams.delete(convoId);
        subscriptionVersions.set(convoId, (subscriptionVersions.get(convoId) ?? 0) + 1);
        return;
      }
      if (frame.t === "subscribe") {
        if (!validConvoId(convoId)) return;
        if (!session.subscriptions.has(convoId) && session.subscriptions.size >= MAX_SUBSCRIPTIONS) {
          socket.close(1008, "subscription limit");
          return;
        }
        const version = (subscriptionVersions.get(convoId) ?? 0) + 1;
        subscriptionVersions.set(convoId, version);
        const messages = await chatStore.getMessagesForUser(convoId, session.userId);
        if (subscriptionVersions.get(convoId) !== version || socket.readyState !== WebSocket.OPEN) return;
        session.subscriptions.add(convoId);
        send(socket, { t: "history", convoId, messages, nextCursor: messages.length === 50 ? messages[0]?.id ?? null : null });
        const peerId = typeof frame.peerId === "string" ? frame.peerId : null;
        if (peerId && peerId !== session.userId) session.watchedPeers.set(convoId, peerId);
        if (peerId && peerId !== session.userId && hasUserConnection(peerId)) {
          send(socket, { t: "presence", convoId, userId: peerId, online: true });
        }
        const peerIds = Array.isArray(frame.peerIds)
          ? (frame.peerIds as unknown[]).filter((id): id is string => typeof id === "string")
          : [];
        if (peerIds.length > 0) {
          session.watchedTeams.set(convoId, new Set(peerIds));
          for (const pid of peerIds) {
            if (pid !== session.userId && hasUserConnection(pid)) {
              send(socket, { t: "presence", convoId, userId: pid, online: true });
            }
          }
        }
        for (const [peerSocket, peerSession] of sessions) {
          if (peerSocket !== socket && peerSession.userId !== session.userId && peerSession.subscriptions.has(convoId)) {
            send(socket, { t: "presence", convoId, userId: peerSession.userId, online: true });
          }
        }
        broadcastPresence(convoId, session.userId, true, socket);
        return;
      }
      if (frame.t === "typing") {
        if (!validConvoId(convoId)) return;
        if (!session.subscriptions.has(convoId)) return;
        if (!await chatStore.canReadConversation(convoId, session.userId)) {
          session.subscriptions.delete(convoId);
          return;
        }
        void fanOut(convoId, { t: "typing", convoId, userId: session.userId, isTyping: Boolean(frame.isTyping) }, socket);
        return;
      }
      if (frame.t === "message") {
        if (!validConvoId(convoId)) return;
        if (!session.subscriptions.has(convoId)) return;
        const body = typeof frame.body === "string" ? frame.body : "";
        const clientId = typeof frame.clientId === "string" ? frame.clientId : "";
        if (!body.trim() || body.length > 2000 || clientId.length < 1 || clientId.length > 200) return;
        const { message } = await chatStore.postMessageForUser(convoId, session.userId, body, clientId);
        send(socket, { t: "ack", convoId, clientId, id: message.id });
        void fanOut(convoId, { t: "message", convoId, message });
        return;
      }
    } catch {
      // Do not expose authorization/database details over a public socket.
      send(socket, {
        t: "error",
        convoId,
        clientId: typeof frame.clientId === "string" ? frame.clientId : undefined,
        code: "REJECTED",
      });
    }
  });
  socket.on("close", () => {
    sessions.delete(socket);
    if (!hasUserConnection(session.userId)) {
      broadcastUserPresence(session.userId, false);
    }
    clearInterval(idle);
  });
});

async function main() {
  // Next aliases `server-only` to an empty module in its server compiler. A
  // custom Node host does not go through that compiler, so install the same
  // alias before loading the shared server-only DAL.
  let handle: ((request: IncomingMessage, response: import("node:http").ServerResponse) => unknown) | null = null;
  if (!realtimeOnly) {
    const { addHookAliases } = await import("next/dist/server/require-hook");
    addHookAliases([["server-only", "next/dist/compiled/server-only/empty"]]);
    const next = (await import("next")).default;
    const app = next({ dev: !production });
    await app.prepare();
    handle = app.getRequestHandler();
  }
  chatStore = await import("./src/lib/teams/writes");
  const server = createServer((request, response) => {
    if (request.method === "POST" && request.url === "/notify") {
      const ip = request.socket.remoteAddress ?? "unknown";
      const now = Date.now();
      const rate = notifyWindows.get(ip) ?? { count: 0, windowStart: now };
      if (now - rate.windowStart > NOTIFY_WINDOW_MS) {
        rate.count = 0;
        rate.windowStart = now;
      }
      if (++rate.count > NOTIFY_RATE_LIMIT) {
        notifyWindows.set(ip, rate);
        response.writeHead(429, { "Content-Type": "application/json" });
        response.end(JSON.stringify({ error: "rate limit exceeded" }));
        return;
      }
      notifyWindows.set(ip, rate);

      if (!verifyNotifySecret(request.headers.authorization)) {
        response.writeHead(401, { "Content-Type": "application/json" });
        response.end(JSON.stringify({ error: "unauthorized" }));
        return;
      }

      let body = "";
      request.on("data", (chunk) => {
        body += chunk;
        if (body.length > 8192) {
          response.writeHead(413, { "Content-Type": "application/json" });
          response.end(JSON.stringify({ error: "payload too large" }));
          request.destroy();
        }
      });
      request.on("end", () => {
        try {
          const parsed = JSON.parse(body);
          const validated = NotifyPayloadSchema.safeParse(parsed);
          if (!validated.success) {
            response.writeHead(400, { "Content-Type": "application/json" });
            response.end(JSON.stringify({ error: "invalid payload", details: validated.error.issues }));
            return;
          }
          const { id, type, userIds, data, at, excludeSubscribedTo } = validated.data;
          const frame = { t: "event", id, type, data, at };
          for (const uid of userIds) {
            notifyUser(uid, frame, excludeSubscribedTo);
          }
          response.writeHead(204);
          response.end();
        } catch {
          response.writeHead(400, { "Content-Type": "application/json" });
          response.end(JSON.stringify({ error: "malformed JSON" }));
        }
      });
      return;
    }

    if (realtimeOnly) {
      if (request.url === "/" || request.url?.startsWith("/health") || request.url === "/config") {
        response.writeHead(200, { "Content-Type": "text/html" });
        response.end(`<!DOCTYPE html>
<html>
<head><title>SIHJantaParty Realtime Chat Service</title></head>
<body style="font-family: monospace; background: #0b0c10; color: #66fcf1; padding: 40px;">
  <h2>● SIHJantaParty Realtime Chat Server</h2>
  <p>Status: ONLINE</p>
  <p>WebSocket Endpoint: /ws</p>
</body>
</html>`);
        return;
      }
      response.writeHead(404);
      response.end();
      return;
    }
    return handle!(request, response);
  });
  server.on("upgrade", async (request, socket, head) => {
    // Only /ws is ours. Anything else (notably Next's dev HMR socket) gets a
    // clean close rather than a silent hang. NOTE: this host cannot serve
    // Next's HMR upgrade at all on Next 16 Turbopack, which is exactly why
    // `npm run dev` uses plain `next dev` and this file is reserved for
    // production and `npm run dev:realtime`.
    if (!request.url?.startsWith("/ws")) {
      socket.write("HTTP/1.1 426 Upgrade Required\r\n\r\n");
      socket.destroy();
      return;
    }
    const userId = await authenticateRequestHeaders(request) ?? await authenticate(request.url);
    if (!userId) {
      socket.write("HTTP/1.1 401 Unauthorized\r\n\r\n");
      socket.destroy();
      return;
    }
    const userConnections = [...sessions.values()].filter((session) => session.userId === userId).length;
    if (sessions.size >= MAX_CONNECTIONS || userConnections >= MAX_CONNECTIONS_PER_USER) {
      socket.write("HTTP/1.1 429 Too Many Requests\r\n\r\n");
      socket.destroy();
      return;
    }
    wss.handleUpgrade(request, socket, head, (websocket) => {
      console.log(`[WS] Authenticated connection: ${userId}`);
      wss.emit("connection", websocket, request, {
        userId,
        subscriptions: new Set<string>(),
        watchedPeers: new Map<string, string>(),
        watchedTeams: new Map<string, Set<string>>(),
      } satisfies Session);
    });
  });

  server.listen(port, "0.0.0.0", () => {
     console.log(`SIHJantaParty ${realtimeOnly ? "realtime" : production ? "production" : "development"} server on http://0.0.0.0:${port}`);
  });
}

void main().catch((error) => {
  console.error(error);
  process.exitCode = 1;
});