Submatter/apps/collab-server/src/chat.ts

import { createHmac, timingSafeEqual } from "node:crypto";
import { createServer } from "node:http";
import { and, eq } from "drizzle-orm";
import IoRedis from "ioredis";
import { WebSocketServer, type WebSocket } from "ws";
import { z } from "zod";
import { db } from "@tasks/database/client";
import {
  chatChannelMembers,
  chatChannels,
  workspaceMembers,
} from "@tasks/database/schema";

const uuid = z.string().uuid();
const ticketClaims = z
  .object({
    iss: z.literal("echodo-web"),
    aud: z.literal("echodo-chat-v1"),
    scope: z.literal("chat:subscribe"),
    sub: uuid,
    workspaceId: uuid,
    exp: z.number().int().positive(),
  })
  .strict();
export type ChatTicket = z.infer<typeof ticketClaims>;
const topicPattern =
  /^chat:v1:workspace:([0-9a-f-]{36}):channel:([0-9a-f-]{36})$/;
const subscribe = z
  .object({ type: z.literal("subscribe"), channelId: uuid })
  .strict();
const unsubscribe = z
  .object({ type: z.literal("unsubscribe"), channelId: uuid })
  .strict();
const clientMessage = z.discriminatedUnion("type", [subscribe, unsubscribe]);
const MAX_CONNECTIONS = 2000;
const MAX_CHANNELS = 64;
const MAX_MESSAGE_BYTES = 512;
const MAX_BUFFERED_BYTES = 64 * 1024;
const HEARTBEAT_MS = 15_000;
const AUTH_TIMEOUT_MS = 5_000;

export function verifyTicket(value: string, secret: string): ChatTicket | null {
  if (secret.length < 32 || value.length > 2048) return null;
  const parts = value.split(".");
  if (
    parts.length !== 2 ||
    !parts[0] ||
    !/^[A-Za-z0-9_-]+$/.test(parts[0]) ||
    !/^[A-Za-z0-9_-]+$/.test(parts[1])
  )
    return null;
  const expected = createHmac("sha256", secret).update(parts[0]).digest();
  const supplied = Buffer.from(parts[1], "base64url");
  if (
    supplied.length !== expected.length ||
    !timingSafeEqual(supplied, expected)
  )
    return null;
  try {
    const parsed = ticketClaims.safeParse(
      JSON.parse(Buffer.from(parts[0], "base64url").toString("utf8")),
    );
    if (
      !parsed.success ||
      parsed.data.exp <= Math.floor(Date.now() / 1000) ||
      parsed.data.exp > Math.floor(Date.now() / 1000) + 60
    )
      return null;
    return parsed.data;
  } catch {
    return null;
  }
}

export function parseTopic(
  topic: string,
): { workspaceId: string; channelId: string } | null {
  const match = topicPattern.exec(topic);
  if (
    !match ||
    !uuid.safeParse(match[1]).success ||
    !uuid.safeParse(match[2]).success
  )
    return null;
  return { workspaceId: match[1], channelId: match[2] };
}

type Lookup = (
  workspaceId: string,
  channelId: string,
  userId: string,
) => Promise<{
  kind: string;
  workspaceMember: boolean;
  channelMember: boolean;
} | null>;
export async function authorizeChannel(
  lookup: Lookup,
  workspaceId: string,
  channelId: string,
  userId: string,
): Promise<boolean> {
  try {
    let timeout: ReturnType<typeof setTimeout> | undefined;
    const row = await Promise.race([
      lookup(workspaceId, channelId, userId),
      new Promise<null>((resolve) => {
        timeout = setTimeout(() => resolve(null), AUTH_TIMEOUT_MS);
      }),
    ]).finally(() => clearTimeout(timeout));
    return (
      !!row &&
      row.workspaceMember &&
      (row.kind === "public" ||
        ((row.kind === "private" || row.kind === "direct") &&
          row.channelMember))
    );
  } catch {
    return false;
  }
}

// Always query current membership; never cache authorization across notifications.
const lookupChannel: Lookup = async (workspaceId, channelId, userId) => {
  const [row] = await db
    .select({
      kind: chatChannels.kind,
      workspaceMember: workspaceMembers.userId,
      channelMember: chatChannelMembers.userId,
    })
    .from(chatChannels)
    .leftJoin(
      workspaceMembers,
      and(
        eq(workspaceMembers.workspaceId, chatChannels.workspaceId),
        eq(workspaceMembers.userId, userId),
      ),
    )
    .leftJoin(
      chatChannelMembers,
      and(
        eq(chatChannelMembers.workspaceId, chatChannels.workspaceId),
        eq(chatChannelMembers.channelId, chatChannels.id),
        eq(chatChannelMembers.userId, userId),
      ),
    )
    .where(
      and(
        eq(chatChannels.workspaceId, workspaceId),
        eq(chatChannels.id, channelId),
      ),
    )
    .limit(1);
  return row
    ? {
        kind: row.kind,
        workspaceMember: !!row.workspaceMember,
        channelMember: !!row.channelMember,
      }
    : null;
};

type Peer = {
  socket: WebSocket;
  ticket: ChatTicket;
  channels: Set<string>;
  busy: boolean;
};

/** Separate port: Hocuspocus owns its existing HTTP/WebSocket upgrade handler. */
export async function startChatServer(): Promise<void> {
  const secret = process.env.CHAT_TICKET_SECRET?.trim();
  const redisUrl = process.env.REDIS_URL?.trim();
  if (!secret || secret.length < 32 || !redisUrl) {
    console.warn(
      "[chat] disabled: CHAT_TICKET_SECRET (>=32 characters) and REDIS_URL required",
    );
    return;
  }
  const port = Number(process.env.CHAT_PORT ?? 1235);
  if (!Number.isInteger(port) || port < 1 || port > 65535)
    throw new Error("Invalid CHAT_PORT");
  const subscriber = new IoRedis(redisUrl, { maxRetriesPerRequest: 1 });
  await subscriber.ping(); // Fail closed: never accept subscribers without Redis.
  const http = createServer((_req, res) => {
    res.writeHead(404).end();
  });
  const wss = new WebSocketServer({
    noServer: true,
    maxPayload: MAX_MESSAGE_BYTES,
    perMessageDeflate: false,
  });
  const peers = new Set<Peer>();
  const topicPeers = new Map<string, Set<Peer>>();
  const send = (peer: Peer, value: object) => {
    if (peer.socket.readyState !== 1) return;
    if (peer.socket.bufferedAmount > MAX_BUFFERED_BYTES) {
      peer.socket.close(1013);
      return;
    }
    peer.socket.send(JSON.stringify(value));
  };
  const remove = (peer: Peer, topic: string) => {
    peer.channels.delete(topic);
    const set = topicPeers.get(topic);
    set?.delete(peer);
    if (set?.size === 0) {
      topicPeers.delete(topic);
      void subscriber.unsubscribe(topic).catch(() => {});
    }
  };
  const revoke = (peer: Peer, topic: string, channelId: string) => {
    remove(peer, topic);
    send(peer, { type: "revoked", channelId });
  };
  const deliver = async (topic: string) => {
    const ids = parseTopic(topic);
    if (!ids) return;
    for (const peer of [...(topicPeers.get(topic) ?? [])]) {
      if (!peer.channels.has(topic)) continue;
      if (
        peer.ticket.exp <= Date.now() / 1000 ||
        !(await authorizeChannel(
          lookupChannel,
          ids.workspaceId,
          ids.channelId,
          peer.ticket.sub,
        ))
      ) {
        revoke(peer, topic, ids.channelId);
        peer.socket.close(1008);
        continue;
      }
      // Deliberately ignore Redis envelope including message bodies and keys.
      if (peer.channels.has(topic))
        send(peer, { type: "invalidate", channelId: ids.channelId });
    }
  };
  subscriber.on("message", (topic) => {
    void deliver(topic).catch(() => {});
  });
  subscriber.on("error", () => {
    for (const peer of peers) peer.socket.close(1011);
  });
  subscriber.on("end", () => {
    for (const peer of peers) peer.socket.close(1011);
  });
  http.on("upgrade", (req, socket, head) => {
    if (
      peers.size >= MAX_CONNECTIONS ||
      req.url?.length === undefined ||
      req.url.length > 2300
    ) {
      socket.destroy();
      return;
    }
    let token: string | null;
    try {
      const url = new URL(req.url, "http://localhost");
      token = url.pathname === "/chat" ? url.searchParams.get("ticket") : null;
    } catch {
      token = null;
    }
    const ticket = token ? verifyTicket(token, secret) : null;
    if (!ticket) {
      socket.destroy();
      return;
    }
    wss.handleUpgrade(req, socket, head, (ws) =>
      wss.emit("connection", ws, ticket),
    );
  });
  wss.on("connection", (socket: WebSocket, ticket: ChatTicket) => {
    const peer: Peer = { socket, ticket, channels: new Set(), busy: false };
    peers.add(peer);
    const expiry = setTimeout(
      () => socket.close(1008),
      Math.max(0, ticket.exp * 1000 - Date.now()),
    );
    expiry.unref();
    socket.on("close", () => {
      clearTimeout(expiry);
      peers.delete(peer);
      for (const topic of [...peer.channels]) remove(peer, topic);
    });
    socket.on("message", (bytes) => {
      if (peer.busy) {
        socket.close(1008);
        return;
      }
      peer.busy = true;
      void (async () => {
        if (ticket.exp <= Date.now() / 1000) {
          socket.close(1008);
          return;
        }
        let data: unknown;
        try {
          data = JSON.parse(bytes.toString());
        } catch {
          socket.close(1008);
          return;
        }
        const parsed = clientMessage.safeParse(data);
        if (!parsed.success) {
          socket.close(1008);
          return;
        }
        const { channelId } = parsed.data;
        const topic = `chat:v1:workspace:${ticket.workspaceId}:channel:${channelId}`;
        if (parsed.data.type === "unsubscribe") {
          remove(peer, topic);
          return;
        }
        if (peer.channels.size >= MAX_CHANNELS && !peer.channels.has(topic)) {
          socket.close(1008);
          return;
        }
        if (
          !(await authorizeChannel(
            lookupChannel,
            ticket.workspaceId,
            channelId,
            ticket.sub,
          ))
        ) {
          socket.close(1008);
          return;
        }
        if (peer.channels.has(topic)) return;
        if (!topicPeers.has(topic)) {
          try {
            await subscriber.subscribe(topic);
          } catch {
            socket.close(1011);
            return;
          }
          if (!topicPeers.has(topic)) topicPeers.set(topic, new Set());
        }
        // Recheck after the Redis subscribe await: access may have been revoked meanwhile.
        if (
          !(await authorizeChannel(
            lookupChannel,
            ticket.workspaceId,
            channelId,
            ticket.sub,
          ))
        ) {
          socket.close(1008);
          return;
        }
        if (socket.readyState !== 1) return;
        topicPeers.get(topic)?.add(peer);
        peer.channels.add(topic);
        send(peer, { type: "subscribed", channelId });
      })()
        .catch(() => socket.close(1011))
        .finally(() => {
          peer.busy = false;
        });
    });
  });
  setInterval(() => {
    for (const peer of peers) {
      if (peer.busy) continue;
      peer.busy = true;
      void (async () => {
        if (peer.ticket.exp <= Date.now() / 1000) {
          peer.socket.close(1008);
          return;
        }
        for (const topic of [...peer.channels]) {
          const ids = parseTopic(topic);
          if (
            !ids ||
            !(await authorizeChannel(
              lookupChannel,
              ids.workspaceId,
              ids.channelId,
              peer.ticket.sub,
            ))
          ) {
            revoke(peer, topic, ids?.channelId ?? "");
            peer.socket.close(1008);
            return;
          }
        }
      })()
        .catch(() => peer.socket.close(1011))
        .finally(() => {
          peer.busy = false;
        });
    }
  }, HEARTBEAT_MS).unref();
  await new Promise<void>((resolve, reject) => {
    http.once("error", reject);
    http.listen(port, "0.0.0.0", resolve);
  });
  console.log(`[chat] listening on port ${port}`);
}