arcway 0.1.0

package/server/ws/ws-router.js

@@ -0,0 +1,335 @@
+import { Server as SocketIOServer } from 'socket.io';
+import { randomUUID } from 'node:crypto';
+import { registerWsServer, unregisterWsServer } from './registry.js';
+import { parseCookies, resolveSession, flattenHeaders } from '../session/helpers.js';
+import { toErrorMessage } from '../helpers.js';
+
+class WsRouter {
+  _apiRouter;
+  _log;
+  _pingIntervalMs;
+  _io = null;
+  _clients = new Map();
+  _clientsBySocketId = new Map();
+  _path;
+  _enabled;
+  _sessionConfig;
+
+  constructor(config, { apiRouter, log, sessionConfig }) {
+    this._apiRouter = apiRouter;
+    this._log = log.extend({ logger: 'ws' });
+    this._pingIntervalMs = config?.pingIntervalMs ?? 3e4;
+    this._path = config?.path ?? '/ws';
+    this._sessionConfig = sessionConfig ?? null;
+
+    const wsRoutes = apiRouter.routes.filter((r) => r.wsEnabled);
+    this._enabled = wsRoutes.length > 0;
+
+    if (this._enabled) {
+      log.info(`WebSocket real-time: ${wsRoutes.length} route(s) on ${this._path}`);
+    }
+  }
+
+  attachToServer(httpServer) {
+    if (!this._enabled) return;
+
+    this._io = new SocketIOServer(httpServer, {
+      path: this._path,
+      serveClient: false,
+      transports: ['websocket'],
+      pingInterval: this._pingIntervalMs,
+      pingTimeout: this._pingIntervalMs * 2,
+    });
+
+    this._io.use(async (socket, next) => {
+      try {
+        const cookies = parseCookies({ headers: socket.handshake.headers });
+        const session = this._sessionConfig
+          ? await resolveSession(cookies, this._sessionConfig)
+          : null;
+
+        socket.data.session = session;
+        socket.data.cookies = cookies;
+        socket.data.headers = flattenHeaders(socket.handshake.headers);
+        next();
+      } catch (err) {
+        next(err);
+      }
+    });
+
+    this._io.on('connection', (socket) => this._onConnection(socket));
+
+    registerWsServer({
+      sendToSocket: (socketId, path, response) => this._sendToSocket(socketId, path, response),
+      broadcastToPath: (path, response) => this._broadcastToPath(path, response),
+    });
+  }
+
+  // Kept as no-op so WebServer's close() detection still works
+  handleUpgrade() {}
+
+  // No HTTP requests handled
+  handle() {
+    return false;
+  }
+
+  close() {
+    if (!this._enabled || !this._io) return;
+
+    unregisterWsServer();
+
+    const cleanupPromises = [];
+    for (const client of this._clients.values()) {
+      cleanupPromises.push(this._cleanupClient(client));
+    }
+
+    Promise.allSettled(cleanupPromises).then(() => {
+      this._io.close();
+    });
+  }
+
+  // ── Connection lifecycle ──
+
+  _onConnection(socket) {
+    const socketId = randomUUID();
+    const client = {
+      socket,
+      socketId,
+      subscriptions: new Map(),
+      session: socket.data.session,
+      headers: socket.data.headers,
+      cookies: socket.data.cookies,
+    };
+    this._clients.set(socket.id, client);
+    this._clientsBySocketId.set(socketId, client);
+
+    this._send(client, { path: this._path, data: { socketId }, status: 200 });
+
+    socket.on('msg', (data) => this._onMessage(client, data));
+    socket.on('disconnect', () => {
+      this._cleanupClient(client).catch((err) => {
+        this._log.error('WS client cleanup error on disconnect', { error: String(err) });
+      });
+    });
+    socket.on('error', (err) => {
+      this._log.error('WebSocket connection error', { error: String(err), socketId });
+    });
+  }
+
+  async _onMessage(client, parsed) {
+    if (!parsed || typeof parsed !== 'object') {
+      this._send(client, {
+        path: '',
+        error: { code: 'INVALID_MESSAGE', message: 'Invalid message' },
+      });
+      return;
+    }
+
+    if (!parsed.path || !parsed.method) {
+      this._send(client, {
+        path: '',
+        error: { code: 'INVALID_MESSAGE', message: 'Message must have "path" and "method" fields' },
+      });
+      return;
+    }
+
+    const method = parsed.method.toUpperCase();
+    if (method === 'SUBSCRIBE') {
+      await this._handleSubscribe(client, { ...parsed, method });
+    } else if (method === 'UNSUBSCRIBE') {
+      await this._handleUnsubscribe(client, { ...parsed, method });
+    } else if (['GET', 'POST', 'PUT', 'PATCH', 'DELETE'].includes(method)) {
+      await this._handleMethodCall(client, { ...parsed, method });
+    } else {
+      this._send(client, {
+        path: parsed.path,
+        error: { code: 'INVALID_METHOD', message: `Unsupported method: ${method}` },
+      });
+    }
+  }
+
+  // ── Subscribe / Unsubscribe ──
+
+  async _handleSubscribe(client, wsMsg) {
+    const msgPath = wsMsg.path;
+    if (client.subscriptions.has(msgPath)) {
+      this._send(client, {
+        path: msgPath,
+        error: { code: 'ALREADY_SUBSCRIBED', message: 'Already subscribed to this path' },
+      });
+      return;
+    }
+
+    const matched = this._apiRouter.findRoute('GET', msgPath);
+    if (!matched) {
+      this._send(client, {
+        path: msgPath,
+        error: { code: 'NOT_FOUND', message: `No GET route matches ${msgPath}` },
+        status: 404,
+      });
+      return;
+    }
+
+    const { route, params } = matched;
+    if (!route.wsEnabled || typeof route.config.ws !== 'function') {
+      this._send(client, {
+        path: msgPath,
+        error: {
+          code: 'WS_NOT_SUPPORTED',
+          message: 'This route does not support WebSocket subscriptions',
+        },
+        status: 400,
+      });
+      return;
+    }
+
+    const reqInfo = this._buildRequestInfo(client, { ...wsMsg, method: 'GET' }, params);
+    let cleanup;
+    try {
+      const ctx = this._apiRouter.buildCtx(reqInfo);
+      cleanup = await route.config.ws(ctx);
+    } catch (err) {
+      this._log.error(`WS subscribe error for ${msgPath}`, {
+        error: toErrorMessage(err),
+      });
+      this._send(client, {
+        path: msgPath,
+        error: { code: 'SUBSCRIBE_ERROR', message: 'Subscription setup failed' },
+        status: 500,
+      });
+      return;
+    }
+
+    client.subscriptions.set(msgPath, {
+      routePattern: route.pattern,
+      cleanup: cleanup ?? undefined,
+      req: reqInfo,
+    });
+
+    const response = await this._apiRouter.executeRoute(route, reqInfo);
+    this._send(client, { path: msgPath, ...response, id: wsMsg.id });
+  }
+
+  async _handleUnsubscribe(client, wsMsg) {
+    const msgPath = wsMsg.path;
+    const sub = client.subscriptions.get(msgPath);
+
+    if (!sub) {
+      this._send(client, {
+        path: msgPath,
+        error: { code: 'NOT_SUBSCRIBED', message: 'Not subscribed to this path' },
+      });
+      return;
+    }
+
+    if (sub.cleanup) {
+      try {
+        await sub.cleanup();
+      } catch (err) {
+        this._log.error(`WS unsubscribe cleanup error for ${msgPath}`, {
+          error: toErrorMessage(err),
+        });
+      }
+    }
+
+    client.subscriptions.delete(msgPath);
+  }
+
+  // ── Method calls (GET/POST/PUT/PATCH/DELETE) ──
+
+  async _handleMethodCall(client, wsMsg) {
+    const { path: msgPath, method } = wsMsg;
+    const matched = this._apiRouter.findRoute(method, msgPath);
+
+    if (!matched) {
+      this._send(client, {
+        path: msgPath,
+        error: { code: 'NOT_FOUND', message: `No ${method} route matches ${msgPath}` },
+        status: 404,
+        id: wsMsg.id,
+      });
+      return;
+    }
+
+    const { route, params } = matched;
+    const reqInfo = this._buildRequestInfo(client, wsMsg, params);
+    const response = await this._apiRouter.executeRoute(route, reqInfo);
+    this._send(client, { path: msgPath, ...response, id: wsMsg.id });
+
+    if (!response.error) {
+      const msg = { path: msgPath, data: response.data, status: response.status };
+      for (const c of this._clients.values()) {
+        if (c !== client && c.subscriptions.has(msgPath)) {
+          this._send(c, msg);
+        }
+      }
+    }
+  }
+
+  // ── Internals ──
+
+  _buildRequestInfo(client, wsMsg, params) {
+    return {
+      requestId: randomUUID(),
+      method: wsMsg.method,
+      path: wsMsg.path,
+      query: { ...(wsMsg.query ?? {}), ...params },
+      body: wsMsg.body,
+      headers: client.headers,
+      cookies: client.cookies,
+      session: client.session,
+      socketId: client.socketId,
+    };
+  }
+
+  _send(client, msg) {
+    if (client.socket.connected) {
+      client.socket.emit('msg', msg);
+    }
+  }
+
+  _sendToSocket(socketId, path, response) {
+    const client = this._clientsBySocketId.get(socketId);
+    if (!client) return;
+    this._send(client, {
+      path,
+      data: response.data,
+      error: response.error,
+      status: response.status,
+    });
+  }
+
+  _broadcastToPath(path, response) {
+    const msg = {
+      path,
+      data: response.data,
+      error: response.error,
+      status: response.status,
+    };
+    for (const client of this._clients.values()) {
+      if (client.subscriptions.has(path)) {
+        this._send(client, msg);
+      }
+    }
+  }
+
+  async _cleanupClient(client) {
+    for (const [path, sub] of client.subscriptions) {
+      if (sub.cleanup) {
+        try {
+          await sub.cleanup();
+        } catch (err) {
+          this._log.error(`WS cleanup error for ${path}`, {
+            error: toErrorMessage(err),
+          });
+        }
+      }
+    }
+
+    client.subscriptions.clear();
+    this._clients.delete(client.socket.id);
+    this._clientsBySocketId.delete(client.socketId);
+  }
+}
+
+export { WsRouter };