anydrop/server/src/index.ts

import { createServer } from "node:http";
import { createHash } from "node:crypto";
import { networkInterfaces } from "node:os";
import { WebSocketServer, WebSocket } from "ws";
import { nanoid } from "nanoid";
import {
  type ClientMessage,
  type ServerMessage,
  type PeerInfo,
  type DeviceType,
} from "@anydrop/shared";
import { RoomManager, type Client } from "./rooms.js";

const PORT = parseInt(process.env.PORT || "3001", 10);
const BASE_URL = process.env.BASE_URL || "http://localhost:5173";

const RATE_LIMIT_CONNECTIONS_PER_IP = 10;
const RATE_LIMIT_WINDOW_MS = 60_000;
const RATE_LIMIT_MESSAGES = 100;
const MAX_AVATAR_SIZE = 100_000; // 100KB max for avatar data URLs

// Track connection rate per IP
const connectionCounts = new Map<string, { count: number; windowStart: number }>();

function hashIP(ip: string): string {
  return createHash("sha256").update(ip).digest("hex").slice(0, 16);
}

/** Strip ::ffff: prefix from IPv4-mapped IPv6 addresses */
function normalizeIP(ip: string): string {
  return ip.startsWith("::ffff:") ? ip.slice(7) : ip;
}

/** Check if an IPv4 address is private (RFC 1918) or loopback */
function isPrivateIP(ip: string): boolean {
  return (
    ip.startsWith("10.") ||
    ip.startsWith("192.168.") ||
    ip.startsWith("127.") ||
    /^172\.(1[6-9]|2\d|3[01])\./.test(ip)
  );
}

function isLoopback(ip: string): boolean {
  return ip === "::1" || ip === "127.0.0.1" || ip.startsWith("127.");
}

/** Detect the server's own LAN subnet at startup (e.g. "192.168.1") */
function detectLanSubnet(): string | null {
  const nets = networkInterfaces();
  for (const ifaces of Object.values(nets)) {
    if (!ifaces) continue;
    for (const iface of ifaces) {
      if (!iface.internal && iface.family === "IPv4" && isPrivateIP(iface.address)) {
        const parts = iface.address.split(".");
        return parts.slice(0, 3).join(".");
      }
    }
  }
  return null;
}

const serverLanSubnet = detectLanSubnet();
console.log(`[init] detected LAN subnet: ${serverLanSubnet ?? "none"}`);

function getLanGroupKey(rawIP: string): string {
  const ip = normalizeIP(rawIP);
  if (isLoopback(ip) || ip === "unknown") {
    return serverLanSubnet ?? "localhost";
  }
  if (isPrivateIP(ip)) {
    const parts = ip.split(".");
    return parts.slice(0, 3).join(".");
  }
  return ip;
}

function getClientIP(req: { headers: Record<string, string | string[] | undefined>; socket: { remoteAddress?: string } }): string {
  const forwarded = req.headers["x-forwarded-for"];
  if (forwarded) {
    const first = Array.isArray(forwarded) ? forwarded[0] : forwarded.split(",")[0];
    return first.trim();
  }
  return req.socket.remoteAddress || "unknown";
}

function send(ws: WebSocket, msg: ServerMessage): void {
  if (ws.readyState === WebSocket.OPEN) {
    ws.send(JSON.stringify(msg));
  }
}

function checkConnectionRate(ip: string): boolean {
  const now = Date.now();
  const entry = connectionCounts.get(ip);
  if (!entry || now - entry.windowStart > RATE_LIMIT_WINDOW_MS) {
    connectionCounts.set(ip, { count: 1, windowStart: now });
    return true;
  }
  entry.count++;
  return entry.count <= RATE_LIMIT_CONNECTIONS_PER_IP;
}

function checkMessageRate(client: Client): boolean {
  const now = Date.now();
  if (now - client.messageWindowStart > RATE_LIMIT_WINDOW_MS) {
    client.messageCount = 1;
    client.messageWindowStart = now;
    return true;
  }
  client.messageCount++;
  return client.messageCount <= RATE_LIMIT_MESSAGES;
}

/** Sanitize client-provided display name */
function sanitizeName(name: string): string {
  return name.trim().slice(0, 30) || "Appareil";
}

/** Validate device type */
function validateDeviceType(dt: string): DeviceType {
  const valid: DeviceType[] = ["phone", "tablet", "laptop", "desktop"];
  return valid.includes(dt as DeviceType) ? (dt as DeviceType) : "laptop";
}

/** Validate avatar (must be a small data URL or undefined) */
function validateAvatar(avatar: unknown): string | undefined {
  if (typeof avatar !== "string") return undefined;
  if (!avatar.startsWith("data:image/")) return undefined;
  if (avatar.length > MAX_AVATAR_SIZE) return undefined;
  return avatar;
}

// ── HTTP server (health check) ──

const httpServer = createServer((req, res) => {
  if (req.url === "/health") {
    res.writeHead(200, { "Content-Type": "application/json" });
    res.end(JSON.stringify({ status: "ok" }));
    return;
  }
  res.writeHead(404);
  res.end();
});

// ── WebSocket server ──

const wss = new WebSocketServer({ server: httpServer });
const roomManager = new RoomManager();
const clients = new Map<WebSocket, Client>();

wss.on("connection", (ws, req) => {
  const ip = getClientIP(req as any);
  const groupKey = getLanGroupKey(ip);
  console.log(`[connect] raw IP: ${ip} → groupKey: ${groupKey} → room: lan:${hashIP(groupKey)}`);

  if (!checkConnectionRate(ip)) {
    send(ws, { type: "error", code: "rate-limit", message: "Too many connections" });
    ws.close();
    return;
  }

  const peerId = nanoid(8);
  const lanRoomId = roomManager.getLanRoomId(hashIP(groupKey));

  // Client object created with placeholder name — filled in on "hello"
  const client: Client = {
    ws,
    peerId,
    displayName: "Appareil",
    deviceType: "laptop",
    ip,
    lanRoomId,
    publicRoomId: null,
    messageCount: 0,
    messageWindowStart: Date.now(),
  };

  clients.set(ws, client);

  ws.on("message", (raw) => {
    if (!checkMessageRate(client)) {
      send(ws, { type: "error", code: "rate-limit", message: "Too many messages" });
      return;
    }

    let msg: ClientMessage;
    try {
      msg = JSON.parse(raw.toString());
    } catch {
      return;
    }

    switch (msg.type) {
      case "hello":
        handleHello(client, msg);
        break;
      case "create-public-room":
        handleCreatePublicRoom(client);
        break;
      case "signal":
        handleSignal(client, msg.to, msg.data);
        break;
      case "leave":
        handleLeave(client);
        break;
    }
  });

  ws.on("close", () => {
    handleLeave(client);
    clients.delete(ws);
  });
});

function handleHello(client: Client, msg: ClientMessage & { type: "hello" }): void {
  // Apply client-provided profile
  client.displayName = sanitizeName(msg.deviceName);
  client.deviceType = validateDeviceType(msg.deviceType);
  client.avatar = validateAvatar(msg.avatar);

  // Join LAN room
  const lanRoom = roomManager.getOrCreateLanRoom(hashIP(getLanGroupKey(client.ip)));
  roomManager.addClientToRoom(lanRoom, client);

  const peers: PeerInfo[] = [];
  for (const peer of lanRoom.clients.values()) {
    if (peer.peerId !== client.peerId) {
      peers.push({
        peerId: peer.peerId,
        displayName: peer.displayName,
        deviceType: peer.deviceType,
        avatar: peer.avatar,
      });
      // Notify existing peers
      send(peer.ws, {
        type: "peer-joined",
        peerId: client.peerId,
        displayName: client.displayName,
        deviceType: client.deviceType,
        avatar: client.avatar,
      });
    }
  }
  console.log(`[hello] peer=${client.peerId} (${client.displayName}) room=${lanRoom.id} peers_in_room=${lanRoom.clients.size}`);

  // Join public room if code provided
  if (msg.joinCode) {
    const pubRoom = roomManager.getPublicRoomByCode(msg.joinCode);
    if (!pubRoom) {
      send(client.ws, { type: "error", code: "room-not-found", message: `Room "${msg.joinCode}" not found or expired` });
    } else {
      client.publicRoomId = pubRoom.id;
      roomManager.addClientToRoom(pubRoom, client);
      for (const peer of pubRoom.clients.values()) {
        if (peer.peerId !== client.peerId) {
          if (!peers.find((p) => p.peerId === peer.peerId)) {
            peers.push({
              peerId: peer.peerId,
              displayName: peer.displayName,
              deviceType: peer.deviceType,
              avatar: peer.avatar,
            });
          }
          send(peer.ws, {
            type: "peer-joined",
            peerId: client.peerId,
            displayName: client.displayName,
            deviceType: client.deviceType,
            avatar: client.avatar,
          });
        }
      }
    }
  }

  send(client.ws, {
    type: "welcome",
    peerId: client.peerId,
    roomId: client.lanRoomId,
    peers,
  });
}

function handleCreatePublicRoom(client: Client): void {
  try {
    const { code, roomId, expiresAt } = roomManager.createPublicRoom();
    client.publicRoomId = roomId;
    const room = roomManager.getPublicRoomByCode(code)!;
    roomManager.addClientToRoom(room, client);

    send(client.ws, {
      type: "public-room-created",
      code,
      url: `${BASE_URL}/${code}`,
      expiresAt: new Date(expiresAt).toISOString(),
    });
  } catch {
    send(client.ws, { type: "error", code: "rate-limit", message: "Failed to create room" });
  }
}

function handleSignal(client: Client, to: string, data: unknown): void {
  const target = roomManager.findPeerInClientRooms(client, to);
  if (target) {
    send(target.ws, { type: "signal", from: client.peerId, data });
  }
}

function handleLeave(client: Client): void {
  const lanRoom = roomManager.getOrCreateLanRoom(hashIP(getLanGroupKey(client.ip)));
  roomManager.removeClientFromRoom(lanRoom, client.peerId);
  for (const peer of lanRoom.clients.values()) {
    send(peer.ws, { type: "peer-left", peerId: client.peerId });
  }

  if (client.publicRoomId) {
    const pubRoom = roomManager.getRoomById(client.publicRoomId);
    if (pubRoom) {
      roomManager.removeClientFromRoom(pubRoom, client.peerId);
      for (const peer of pubRoom.clients.values()) {
        send(peer.ws, { type: "peer-left", peerId: client.peerId });
      }
    }
    client.publicRoomId = null;
  }
}

httpServer.listen(PORT, () => {
  console.log(`AnyDrop signaling server running on port ${PORT}`);
});