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";
import {
  initPush,
  getVapidPublicKey,
  storeSubscription,
  notifyOfflineDevices,
  getOfflineSubscribers,
  wakeDevice,
} from "./push.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;

const connectionCounts = new Map<string, { count: number; windowStart: number }>();

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

function normalizeIP(ip: string): string {
  return ip.startsWith("::ffff:") ? ip.slice(7) : ip;
}

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.");
}

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"}`);

// Initialize Web Push
initPush();

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;
}

function sanitizeName(name: unknown): string {
  if (typeof name !== "string") return "Appareil";
  return name.trim().slice(0, 30) || "Appareil";
}

function validateDeviceType(dt: unknown): DeviceType {
  const valid: DeviceType[] = ["phone", "tablet", "laptop", "desktop"];
  return typeof dt === "string" && valid.includes(dt as DeviceType) ? (dt as DeviceType) : "laptop";
}

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;
}

// ── Pair codes (temporary mapping: short code → groupId, 5 min TTL) ──

import { customAlphabet } from "nanoid";
const generatePairCode = customAlphabet("ABCDEFGHJKLMNPQRSTUVWXYZ23456789", 6);
const pairCodes = new Map<string, { groupId: string; expiresAt: number }>();

setInterval(() => {
  const now = Date.now();
  for (const [code, entry] of pairCodes) {
    if (now > entry.expiresAt) pairCodes.delete(code);
  }
}, 60_000);

// ── HTTP server ──

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));

  const client: Client = {
    ws,
    peerId,
    deviceId: null,
    displayName: "Appareil",
    deviceType: "laptop",
    ip,
    lanGroupKey: groupKey,
    lanRoomId,
    groupId: null,
    groupRoomId: null,
    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 "subscribe-push":
        handleSubscribePush(client, msg);
        break;
      case "wake-peer":
        handleWakePeer(client, msg);
        break;
      case "create-pair-code":
        handleCreatePairCode(client, msg);
        break;
      case "resolve-pair-code":
        handleResolvePairCode(client, msg);
        break;
      case "leave":
        handleLeave(client);
        break;
    }
  });

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

function getLocalSubnet(localIP: unknown): string | null {
  if (typeof localIP !== "string") return null;
  const parts = localIP.split(".");
  if (parts.length !== 4) return null;
  if (!isPrivateIP(localIP)) return null;
  return parts.slice(0, 3).join(".");
}

// Maps local subnet → LAN room hash, so devices on the same local subnet share a room
const subnetToRoomHash = new Map<string, string>();

function handleHello(client: Client, msg: ClientMessage & { type: "hello" }): void {
  client.displayName = sanitizeName(msg.deviceName);
  client.deviceType = validateDeviceType(msg.deviceType);
  client.avatar = validateAvatar(msg.avatar);
  client.deviceId = typeof msg.deviceId === "string" ? msg.deviceId : null;

  const publicGroupKey = getLanGroupKey(client.ip);
  const localSubnet = getLocalSubnet(msg.localIP);

  // Determine which room hash to use:
  // - If the client provides a localIP subnet, check if another device on the same subnet
  //   is already in a room (possibly with a different public IP due to Private Relay / VPN).
  //   If so, join that room. Otherwise, register this subnet → room mapping.
  // - Fall back to public IP grouping.
  let roomHash: string;
  if (localSubnet) {
    const existingHash = subnetToRoomHash.get(localSubnet);
    if (existingHash) {
      roomHash = existingHash;
    } else {
      roomHash = hashIP(publicGroupKey);
      subnetToRoomHash.set(localSubnet, roomHash);
    }
    console.log(`[hello] localIP=${msg.localIP} subnet=${localSubnet} → roomHash=${roomHash}`);
  } else {
    roomHash = hashIP(publicGroupKey);
  }

  const lanRoom = roomManager.getOrCreateLanRoom(roomHash);
  client.lanRoomId = lanRoom.id;
  client.lanGroupKey = localSubnet || publicGroupKey;
  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,
      });
      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 device group room if paired
  if (typeof msg.groupId === "string" && msg.groupId.length > 0 && msg.groupId.length <= 64) {
    client.groupId = msg.groupId;
    const groupRoom = roomManager.getOrCreateGroupRoom(msg.groupId);
    client.groupRoomId = groupRoom.id;
    roomManager.addClientToRoom(groupRoom, client);
    for (const peer of groupRoom.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,
        });
      }
    }
    console.log(`[hello] peer=${client.peerId} joined group=${msg.groupId} peers_in_group=${groupRoom.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,
          });
        }
      }
    }
  }

  // Collect online device IDs
  const onlineDeviceIds = new Set<string>();
  for (const peer of lanRoom.clients.values()) {
    if (peer.deviceId) onlineDeviceIds.add(peer.deviceId);
  }

  // Add offline push-subscribed peers to the list (excluding self)
  const offlineSubs = getOfflineSubscribers(client.lanGroupKey, onlineDeviceIds, client.deviceId ?? undefined);
  for (const sub of offlineSubs) {
    peers.push({
      peerId: `offline:${sub.deviceId}`,
      displayName: sub.displayName,
      deviceType: sub.deviceType,
      online: false,
      deviceId: sub.deviceId,
    });
  }

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

  // Notify offline devices via push (skip self)
  notifyOfflineDevices(client.lanGroupKey, onlineDeviceIds, client.deviceId, client.displayName, client.deviceType);
}

function handleSubscribePush(client: Client, msg: ClientMessage & { type: "subscribe-push" }): void {
  if (!msg.deviceId || !msg.subscription) return;
  storeSubscription(msg.deviceId, client.displayName, client.deviceType, msg.subscription as any, client.lanGroupKey);
}

function handleWakePeer(client: Client, msg: ClientMessage & { type: "wake-peer" }): void {
  if (!msg.deviceId) return;
  wakeDevice(client.lanGroupKey, msg.deviceId, client.displayName, client.deviceType);
}

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 handleCreatePairCode(client: Client, msg: ClientMessage & { type: "create-pair-code" }): void {
  if (typeof msg.groupId !== "string" || msg.groupId.length === 0) return;

  // Check if this groupId already has an active code
  for (const [code, entry] of pairCodes) {
    if (entry.groupId === msg.groupId && Date.now() < entry.expiresAt) {
      send(client.ws, { type: "pair-code-created", code });
      return;
    }
  }

  const code = generatePairCode();
  pairCodes.set(code, { groupId: msg.groupId, expiresAt: Date.now() + 5 * 60 * 1000 });
  console.log(`[pair] created code ${code} for group ${msg.groupId.slice(0, 8)}...`);
  send(client.ws, { type: "pair-code-created", code });
}

function handleResolvePairCode(client: Client, msg: ClientMessage & { type: "resolve-pair-code" }): void {
  if (typeof msg.code !== "string" || msg.code.length !== 6) {
    send(client.ws, { type: "error", code: "pair-code-not-found", message: "Code d'appairage invalide" });
    return;
  }
  const entry = pairCodes.get(msg.code.toUpperCase());
  if (!entry || Date.now() > entry.expiresAt) {
    send(client.ws, { type: "error", code: "pair-code-not-found", message: "Code d'appairage invalide ou expiré" });
    return;
  }
  pairCodes.delete(msg.code.toUpperCase());
  console.log(`[pair] resolved code ${msg.code} → group ${entry.groupId.slice(0, 8)}...`);
  send(client.ws, { type: "pair-code-resolved", groupId: entry.groupId });
}

function handleLeave(client: Client): void {
  const lanRoom = roomManager.getRoomById(client.lanRoomId);
  if (!lanRoom) return;
  roomManager.removeClientFromRoom(lanRoom, client.peerId);
  for (const peer of lanRoom.clients.values()) {
    send(peer.ws, { type: "peer-left", peerId: client.peerId });
  }

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

  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}`);
});