redweb 0.8.0 → 0.9.0

package/src/ws/BaseHandler.js

@@ -18,47 +18,47 @@ class BaseHandler {
      * @param {WebSocket & {sendJson: (message: Object) => void, broadcast: (message: Object) => void}} socket - The WebSocket connection that sent the message.
      * @param {any} message - The incoming message in parsed JSON.
      */
-    async handleMessage(socket, message) {
-        const validationResult = await this.validateMessage(message, socket);
-        if (validationResult === false) {
-            throw new Error('Invalid message');
-        }
-        return this.onMessage(socket, message);
-    }
-
-    validateMessage() {
-        return true;
-    }
-
-    /**
-     * Handles an incoming binary message.
-     * @param {WebSocket & {sendJson: (message: Object) => void}} socket - The WebSocket connection that sent the message.
-     * @param {Buffer} buffer - The incoming binary message.
-     */
-    async handleBinaryMessage(socket, buffer) {
-        return this.onBinaryMessage(socket, buffer);
-    }
-
-    /**
-     * Method to be overriden to process messages.
-     * @param {WebSocket} socket - The WebSocket connection that sent the message.
-     * @param {any} message - The incoming message in parsed JSON.
+    async handleMessage(socket, message) {
+        const validationResult = await this.validateMessage(message, socket);
+        if (validationResult === false) {
+            throw new Error('Invalid message');
+        }
+        return this.onMessage(socket, message);
+    }
+
+    validateMessage() {
+        return true;
+    }
+
+    /**
+     * Handles an incoming binary message.
+     * @param {WebSocket & {sendJson: (message: Object) => void}} socket - The WebSocket connection that sent the message.
+     * @param {Buffer} buffer - The incoming binary message.
      */
-    onMessage(socket, message) {
-        throw new Error('onMessage must be implemented by the handler.');
-    }
-
-    /**
-     * Method to be overriden to process binary messages.
-     * @param {WebSocket & {sendJson: (message: Object) => void}} socket - The WebSocket connection that sent the message.
-     * @param {Buffer} buffer - The incoming binary message.
-     */
-    onBinaryMessage(socket, buffer) {
-        socket.sendJson({ error: 'Binary messages are not supported by this handler' });
-    }
-
-    onInitialContact(socket) {
-
+    async handleBinaryMessage(socket, buffer) {
+        return this.onBinaryMessage(socket, buffer);
+    }
+
+    /**
+     * Method to be overriden to process messages.
+     * @param {WebSocket} socket - The WebSocket connection that sent the message.
+     * @param {any} message - The incoming message in parsed JSON.
+     */
+    onMessage(socket, message) {
+        throw new Error('onMessage must be implemented by the handler.');
+    }
+
+    /**
+     * Method to be overriden to process binary messages.
+     * @param {WebSocket & {sendJson: (message: Object) => void}} socket - The WebSocket connection that sent the message.
+     * @param {Buffer} buffer - The incoming binary message.
+     */
+    onBinaryMessage(socket, buffer) {
+        socket.sendJson({ error: 'Binary messages are not supported by this handler' });
+    }
+
+    onInitialContact(socket) {
+
     }
 }
 

package/src/ws/BaseSocketServer.js

@@ -1,31 +1,33 @@
 /**
  * @typedef {Object} SocketServerOptions
- * @property {import('http').Server} [server]                    HTTP server to bind to
- * @property {number}              [port=3000]                  Port to listen on
- * @property {boolean}             [listen=true]                Whether owned servers should automatically start listening
- * @property {Array<new () => import('./SocketRoute').SocketRoute>} [routes]
- */
-
-const DefaultRoute = require('./DefaultRoute');
-const {
-  listenServer,
-  closeServer,
-  settleTasks,
-  throwCleanupErrors,
-  validateListenerOptions,
-} = require('../serverLifecycle');
-
-const SOCKET_OPTIONS = {
-  port: 3000,
-  bind: '0.0.0.0',
-  ssl:  null,
-  listen: true,
-  routes: [],
-  fallbackToRoot: false,
-  closeServerOnShutdown: undefined,
-  logger: console,
-  listenCallback: undefined,
-};
+ * @property {import('http').Server} [server]                    HTTP server to bind to
+ * @property {number}              [port=3000]                  Port to listen on
+ * @property {boolean}             [listen=true]                Whether owned servers should automatically start listening
+ * @property {Array<new () => import('./SocketRoute').SocketRoute>} [routes]
+ */
+
+const DefaultRoute = require('./DefaultRoute');
+const { PLACEMENT_REDIRECT, ADMISSION_SETTLEMENT } = require('./AdmissionPolicy');
+const { PROTOCOL_REJECTION } = require('./ProtocolPolicy');
+const {
+  listenServer,
+  closeServer,
+  settleTasks,
+  throwCleanupErrors,
+  validateListenerOptions,
+} = require('../serverLifecycle');
+
+const SOCKET_OPTIONS = {
+  port: 3000,
+  bind: '0.0.0.0',
+  ssl:  null,
+  listen: true,
+  routes: [],
+  fallbackToRoot: false,
+  closeServerOnShutdown: undefined,
+  logger: console,
+  listenCallback: undefined,
+};
 
 /**
  * Base WebSocket server
@@ -35,54 +37,62 @@ class BaseSocketServer {
    * @param {import('http').Server} server
    * @param {SocketServerOptions}  [options]
    */
-  constructor(server, options = {}, ownsServer = false, name = 'SocketServer') {
-    if (!server || typeof server.on !== 'function') throw new TypeError('A Node HTTP(S) server is required.');
-    Object.assign(this, { ...SOCKET_OPTIONS, ...options });
-    validateListenerOptions(this);
-    if (!Array.isArray(this.routes)) throw new TypeError('`routes` must be an array.');
-    this.server = server;
-    this.ownsServer = ownsServer;
-    this.closeServerOnShutdown = options.closeServerOnShutdown ?? ownsServer;
+  constructor(server, options = {}, ownsServer = false, name = 'SocketServer') {
+    if (!server || typeof server.on !== 'function') throw new TypeError('A Node HTTP(S) server is required.');
+    Object.assign(this, { ...SOCKET_OPTIONS, ...options });
+    validateListenerOptions(this);
+    if (!Array.isArray(this.routes)) throw new TypeError('`routes` must be an array.');
+    this.server = server;
+    this.ownsServer = ownsServer;
+    this.closeServerOnShutdown = options.closeServerOnShutdown ?? ownsServer;
+    this.draining = false;
+    this.pendingUpgrades = new Map();
+    this.rawConnections = this.closeServerOnShutdown ? new Set() : null;
+    this._connectionHandler = this.rawConnections ? socket => {
+      this.rawConnections.add(socket);
+      socket.once('close', () => this.rawConnections.delete(socket));
+    } : null;
+    if (this._connectionHandler) this.server.on('connection', this._connectionHandler);
 
     /* ─── ROUTE INITIALISATION ─────────────────────────── */
-    const RouteClasses = options.routes?.length ? [...options.routes] : [DefaultRoute];
-    this.routes = [];
-    try {
-      for (const RouteClass of RouteClasses) {
-        const route = new RouteClass(server, { logger: this.logger });
-        if (this.routes.some(existing => existing.path === route.path)) {
-          this.disposeRoutes([route]);
-          throw new Error('WebSocket route paths must be unique.');
-        }
-        this.routes.push(route);
-      }
-    } catch (error) {
-      this.disposeRoutes(this.routes);
-      throw error;
-    }
-
-    this._upgradeHandler = this.handleUpgrade.bind(this);
-    this.server.on('upgrade', this._upgradeHandler);
-
-    const shouldListen = (ownsServer && this.listen !== false) || (!ownsServer && options.listen === true);
-    if (shouldListen) {
-      listenServer(this.server, {
-        port: this.port,
-        bind: this.bind,
-        callback: this.listenCallback,
-        logger: this.logger,
-        name,
-      });
-    }
-  }
-
-  disposeRoutes(routes) {
-    routes.forEach(route => {
-      Promise.resolve()
-        .then(() => route.shutdown?.())
-        .catch(error => this.logger?.error?.('Error shutting down route:', error));
-    });
-  }
+    const RouteClasses = options.routes?.length ? [...options.routes] : [DefaultRoute];
+    this.routes = [];
+    try {
+      for (const RouteClass of RouteClasses) {
+        const route = new RouteClass(server, { logger: this.logger });
+        if (this.routes.some(existing => existing.path === route.path)) {
+          this.disposeRoutes([route]);
+          throw new Error('WebSocket route paths must be unique.');
+        }
+        this.routes.push(route);
+      }
+    } catch (error) {
+      this.disposeRoutes(this.routes);
+      throw error;
+    }
+
+    this._upgradeHandler = this.handleUpgrade.bind(this);
+    this.server.on('upgrade', this._upgradeHandler);
+
+    const shouldListen = (ownsServer && this.listen !== false) || (!ownsServer && options.listen === true);
+    if (shouldListen) {
+      listenServer(this.server, {
+        port: this.port,
+        bind: this.bind,
+        callback: this.listenCallback,
+        logger: this.logger,
+        name,
+      });
+    }
+  }
+
+  disposeRoutes(routes) {
+    routes.forEach(route => {
+      Promise.resolve()
+        .then(() => route.shutdown?.())
+        .catch(error => this.logger?.error?.('Error shutting down route:', error));
+    });
+  }
 
   handleUpgrade(req, sock, head) {
     // Some websocket clients (e.g., certain UE plugins) are finicky about the
@@ -97,50 +107,135 @@ class BaseSocketServer {
     })();
 
     const route =
-      this.routes.find(r => r.path === path) ||
-      (this.fallbackToRoot ? this.routes.find(r => r.path === '/') : undefined);
+      this.routes.find(r => r.path === path) ||
+      (this.fallbackToRoot ? this.routes.find(r => r.path === '/') : undefined);
 
     if (!route) return sock.destroy();
+    if (this.draining) return this.rejectUpgrade(sock, 503, 'Service Unavailable');
+    if (route.isReady?.() === false) return this.rejectUpgrade(sock, 503, 'Service Unavailable');
+
+    if (route.admissionPolicy || route.protocolPolicy || route.transportPolicy && route.transportPolicy.maxConnections !== Infinity) {
+      let reservation;
+      try {
+        reservation = route.reserveUpgrade(req);
+      } catch (error) {
+        this.logger?.error?.('WebSocket admission reservation failed:', error);
+      }
+      if (!reservation) return this.rejectUpgrade(sock, 503, 'Service Unavailable');
+      const controller = new AbortController();
+      this.pendingUpgrades.set(sock, { controller, route, reservation });
+      void Promise.resolve()
+        .then(() => route.authorizeUpgrade(req, sock, controller.signal))
+        .then(accepted => {
+          if (sock.destroyed) return;
+          if (this.draining) return this.rejectUpgrade(sock, 503, 'Service Unavailable');
+          if (!accepted) {
+            const redirect = req[PLACEMENT_REDIRECT];
+            if (redirect) return this.rejectUpgrade(sock, 307, 'Temporary Redirect', { Location: redirect });
+            const rejection = req[PROTOCOL_REJECTION];
+            return rejection
+              ? this.rejectUpgrade(sock, rejection.statusCode, rejection.statusText, rejection.headers)
+              : this.rejectUpgrade(sock, 401, 'Unauthorized');
+          }
+          this.completeUpgrade(route, req, sock, head);
+        })
+        .catch(error => {
+          this.logger?.error?.('WebSocket admission failed:', error);
+          if (!sock.destroyed) this.rejectUpgrade(sock, 401, 'Unauthorized');
+        })
+        .finally(async () => {
+          await req[ADMISSION_SETTLEMENT];
+          this.pendingUpgrades.delete(sock);
+          route.releaseUpgrade(reservation);
+        });
+      return;
+    }
+
+    this.completeUpgrade(route, req, sock, head);
+  }
 
+  completeUpgrade(route, req, sock, head) {
     route.server.handleUpgrade(req, sock, head, (s, r) =>
       route.server.emit('connection', s, r)
     );
   }
 
+  rejectUpgrade(socket, statusCode, statusText, headers = {}) {
+    try {
+      const extraHeaders = Object.entries(headers).map(([name, value]) => `${name}: ${value}\r\n`).join('');
+      socket.end?.(
+        `HTTP/1.1 ${statusCode} ${statusText}\r\n` +
+        extraHeaders +
+        'Connection: close\r\n' +
+        'Content-Length: 0\r\n\r\n'
+      );
+    } catch {
+      socket.destroy?.();
+    }
+  }
+
   /**
    * Dynamically attach a new route at runtime
    * @param {new () => import('./SocketRoute').SocketRoute} RouteClass
    */
-  addRoute(RouteClass) {
-    const route = new RouteClass(this.server, { logger: this.logger });
-    if (this.routes.some(existing => existing.path === route.path)) {
-      this.disposeRoutes([route]);
-      throw new Error(`A WebSocket route already exists at ${route.path}.`);
-    }
-    this.routes.push(route);
-    return route;
+  addRoute(RouteClass) {
+    const route = new RouteClass(this.server, { logger: this.logger });
+    if (this.routes.some(existing => existing.path === route.path)) {
+      this.disposeRoutes([route]);
+      throw new Error(`A WebSocket route already exists at ${route.path}.`);
+    }
+    this.routes.push(route);
+    return route;
   }
 
   /**
    * Gracefully tear down all routes (and their services)
    */
-  shutdown() {
-    if (!this._shutdownPromise) this._shutdownPromise = this.performShutdown();
-    return this._shutdownPromise;
-  }
-
-  async performShutdown() {
-    this.server.off?.('upgrade', this._upgradeHandler);
-    const errors = await settleTasks(this.routes.map(route => () => route.shutdown?.()));
-    if (this.closeServerOnShutdown && this.server.listening) {
-      try {
-        await closeServer(this.server);
-      } catch (error) {
-        errors.push(error);
-      }
-    }
-    throwCleanupErrors(errors, 'One or more WebSocket server cleanup operations failed.');
-  }
-}
+  shutdown() {
+    if (!this._shutdownPromise) this._shutdownPromise = this.performShutdown();
+    return this._shutdownPromise;
+  }
+
+  isReady() {
+    return !this.draining && this.routes.every(route => route.isReady?.() !== false);
+  }
+
+  beginDrain() {
+    if (this.draining) return false;
+    this.draining = true;
+    this.pendingUpgrades.forEach(({ controller }) => controller.abort());
+    this.routes.forEach(route => route.beginDrain?.());
+    return true;
+  }
+
+  async performShutdown() {
+    this.beginDrain();
+    this.server.off?.('upgrade', this._upgradeHandler);
+    const deadline = Date.now() + Math.max(0, ...this.routes.map(route => route.shutdownTimeoutMs));
+    const errors = await settleTasks(this.routes.map(route => () => route.shutdown?.()));
+    if (this.closeServerOnShutdown && this.server.listening) {
+      try {
+        await this.closeOwnedServer(Math.max(0, deadline - Date.now()));
+      } catch (error) {
+        errors.push(error);
+      }
+    }
+    if (this._connectionHandler) this.server.off?.('connection', this._connectionHandler);
+    throwCleanupErrors(errors, 'One or more WebSocket server cleanup operations failed.');
+  }
+
+  closeOwnedServer(timeoutMs) {
+    let timer;
+    const closing = closeServer(this.server);
+    const timeout = new Promise(resolve => {
+      timer = setTimeout(() => {
+        this.rawConnections?.forEach(socket => socket.destroy?.());
+        resolve();
+      }, timeoutMs);
+      timer.unref?.();
+    });
+    return Promise.race([closing, timeout]).finally(() => clearTimeout(timer));
+  }
+}
 
 module.exports = { BaseSocketServer, SOCKET_OPTIONS };

package/src/ws/DefaultHandler.js

@@ -5,9 +5,9 @@ class DefaultHandler extends BaseHandler {
         super("DefaultHandler");
     }
 
-    onMessage(socket, message) {
-        socket.sendJson({ message: `I got your message of ${JSON.stringify(message)}` });
-    }
-}
+    onMessage(socket, message) {
+        socket.sendJson({ message: `I got your message of ${JSON.stringify(message)}` });
+    }
+}
 
-module.exports = DefaultHandler;
+module.exports = DefaultHandler;

package/src/ws/DefaultRoute.js

@@ -2,14 +2,14 @@ const SocketRoute  = require("./SocketRoute");
 const DefaultHandler = require('./DefaultHandler');
 
 class DefaultRoute extends SocketRoute {
-    constructor(server, options = {}) {
-        super({
-            server,
-            path: "/",
-            handlers: [DefaultHandler],
-            logger: options.logger,
-        })
+    constructor(server, options = {}) {
+        super({
+            server,
+            path: "/",
+            handlers: [DefaultHandler],
+            logger: options.logger,
+        })
     }
 }
 
-module.exports = DefaultRoute;
+module.exports = DefaultRoute;