Implement master connection closure functionality (#11)

* Implement master connection closure functionality

... and update README with new examples. Add tests for connection closure in test-websocket-multiplexer.js.

https://linear.app/zetra/issue/HYX-774/ws-multiplexer-dynamic-upstream-switch

* Refactor master message handling and add validation for injection and close messages

Author: lovasoa

README.md

@@ -130,12 +130,13 @@ The multiplexer sends several types of messages to the master control:
 ```javascript
 {
   type: 'connection',
-  event: 'client-connected',  // or 'client-disconnected', 'upstream-connected', 'upstream-disconnected'
+  event: 'client-connected',  // or 'client-disconnected', 'upstream-connected', 'upstream-disconnected', 'connection-closed-by-master'
   connectionId: '/path',
   ip: '127.0.0.1',           // Only for client-connected
   headers: { ... },          // Only for client-connected
   code: 1000,                // Only for disconnection events
-  reason: 'Normal closure'   // Only for disconnection events
+  reason: 'Normal closure',  // Only for disconnection events
+  timestamp: '2024-01-01T12:00:00.000Z'  // Only for connection-closed-by-master
 }
 ```
 
@@ -277,6 +278,24 @@ master.send(JSON.stringify({
 }));
 ```
 
+### Example: Closing Connections
+
+You can cleanly close specific connections through the master control:
+
+```javascript
+// Close a specific connection (both client and upstream)
+master.send(JSON.stringify({
+  type: 'close',
+  connectionId: '/chat',
+  reason: 'Connection terminated by administrator'  // Optional reason
+}));
+```
+
+When a connection is closed via master control:
+- Both client and upstream connections are closed cleanly with code 1000
+- Message queues for the connection are cleared
+- A notification is sent to all root master connections with event type `connection-closed-by-master`
+
 ## Use Cases
 
 - Debugging WebSocket applications

test-websocket-multiplexer.js

@@ -620,6 +620,53 @@ class TestRunner {
     assert.equal(disconnectData.event, 'client-disconnected');
     assert.equal(disconnectData.connectionId, CONFIG.TEST_PATH);
 
+    // Test 8: Master connection close functionality
+    console.log('Test: Master connection close functionality');
+    
+    // Create a new client for the close test
+    const clientForClose = new WebSocketClient(
+      `ws://localhost:${this.config.PROXY_PORT}${this.config.TEST_PATH}`,
+      testHeaders
+    );
+    await clientForClose.connect();
+    
+    // Wait for connection to be established and consume the client-connected notification
+    await wait(500);
+    const clientConnectedMsg = await masterControl.receiveMessage();
+    const clientConnectedData = JSON.parse(clientConnectedMsg);
+    assert.equal(clientConnectedData.type, 'connection');
+    assert.equal(clientConnectedData.event, 'client-connected');
+    assert.equal(clientForClose.ws.readyState, WebSocket.OPEN, 'Connection should be open');
+
+    // Send close command via master
+    const closeReason = 'Test closure from master';
+    await masterControl.send(
+      JSON.stringify({
+        type: 'close',
+        connectionId: this.config.TEST_PATH,
+        reason: closeReason,
+      })
+    );
+    
+    // Verify the connection-closed-by-master notification
+    const closedByMasterMsg = await masterControl.receiveMessage();
+    const closedByMasterData = JSON.parse(closedByMasterMsg);
+    assert.equal(closedByMasterData.type, 'connection');
+    assert.equal(closedByMasterData.event, 'connection-closed-by-master');
+    assert.equal(closedByMasterData.connectionId, CONFIG.TEST_PATH);
+    assert.equal(closedByMasterData.reason, closeReason);
+    assert(closedByMasterData.timestamp);
+    
+    // Try to send a message - it should fail because the connection is closed
+    try {
+      await clientForClose.send('This should fail');
+      assert.fail('Message should not be sent, connection should be closed');
+    } catch (error) {
+      console.log('Connection properly closed, message send failed as expected: ' + error.message);
+    }
+
+    assert.equal(clientForClose.ws.readyState, WebSocket.CLOSED, 'Connection should be closed');
+
     // Cleanup
     masterControl.close();
     masterClientMonitor.close();

websocket-multiplex.js

@@ -165,7 +165,7 @@ function sendMessage(ws, message, source, target) {
 
 /**
  * Handles message injection from master to specified targets
- * @param {{ target: string, message: string | Buffer }} data - The message data containing target and message
+ * @param {InjectionMessage} data - The message data containing target and message
  */
 function handleMasterInjection(data) {
   if (data.target === 'all-clients') {
@@ -209,41 +209,150 @@ function handleMasterInjection(data) {
   }
 }
 
+/**
+ * Handles connection closure requests from master
+ * @param {CloseMessage} data - The close data containing connection ID and optional reason
+ */
+function handleMasterClose(data) {
+  const connectionId = data.connectionId;
+  const reason = data.reason || 'Closed by websocket-multiplex master control';
+  
+  logger.info(`Master requesting to close connection: ${connectionId}`);
+  
+  const client = connections.clients.get(connectionId);
+  const upstream = connections.upstreams.get(connectionId);
+  
+  if (!client && !upstream) {
+    logger.warn(`Connection ${connectionId} not found for closure`);
+    return;
+  }
+  
+  // Close client connection if it exists
+  if (client && client.ws.readyState === WebSocket.OPEN) {
+    logger.info(`Closing client connection for ${connectionId}`);
+    client.ws.close(1000, reason);
+  }
+  
+  // Close upstream connection if it exists
+  if (upstream && upstream.ws.readyState === WebSocket.OPEN) {
+    logger.info(`Closing upstream connection for ${connectionId}`);
+    upstream.ws.close(1000, reason);
+  }
+  
+  // Clean up message queue
+  connections.messageQueues.delete(connectionId);
+  
+  // Notify root masters about the forced closure
+  notifyRootMasters('connection', 'connection-closed-by-master', connectionId, {
+    reason,
+    timestamp: new Date().toISOString(),
+  });
+}
+
+/**
+ * @typedef {Object} InjectionMessage
+ * @property {'inject'} type - Message type
+ * @property {string} target - Target for injection ('all-clients', 'all-upstreams', 'client:...', 'upstream:...')
+ * @property {string | Buffer} message - Message to inject
+ */
+
+/**
+ * @typedef {Object} CloseMessage
+ * @property {'close'} type - Message type
+ * @property {string} connectionId - Connection ID to close
+ * @property {string} [reason] - Optional reason for closure
+ */
+
+/**
+ * Validates an injection message
+ * @param {any} data - Data to validate
+ * @returns {data is InjectionMessage} True if valid injection message
+ */
+function validateInjectionMessage(data) {
+  return (
+    data &&
+    typeof data === 'object' &&
+    data.type === 'inject' &&
+    typeof data.target === 'string' &&
+    data.target.length > 0 &&
+    (typeof data.message === 'string' || Buffer.isBuffer(data.message))
+  );
+}
+
+/**
+ * Validates a close message
+ * @param {any} data - Data to validate
+ * @returns {data is CloseMessage} True if valid close message
+ */
+function validateCloseMessage(data) {
+  return (
+    data &&
+    typeof data === 'object' &&
+    data.type === 'close' &&
+    typeof data.connectionId === 'string' &&
+    data.connectionId.length > 0
+  );
+}
+
+/**
+ * Handles messages from root master connections
+ * @param {string} message - The message received
+ * @throws {Error} When message is invalid
+ */
+function handleRootMasterMessage(message) {
+  const data = JSON.parse(message);
+  logger.debug(
+    `multiplexer <- master client: root ${JSON.stringify(data)}`
+  );
+
+  if (validateInjectionMessage(data)) return handleMasterInjection(data);
+  else if (validateCloseMessage(data)) return handleMasterClose(data);
+  else throw new Error(`Invalid master message: unsupported type '${data?.type}' or missing required fields. Message: ${JSON.stringify(data)}`);
+}
+
+/**
+ * Handles messages from master connection with type = 'client'
+ * @param {string} targetPath - The target client path
+ * @param {string} message - The message to forward
+ */
+function handleMasterMessageForClient(targetPath, message) {
+  const client = connections.clients.get(targetPath);
+  if (client?.connected) {
+    sendMessage(client.ws, message, 'master', `client:${targetPath}`);
+  }
+}
+
+/**
+ * Handles messages from master connection with type = 'upstream'
+ * @param {string} targetPath - The target upstream path
+ * @param {string} message - The message to forward
+ */
+function handleMasterMessageForUpstream(targetPath, message) {
+  const upstream = connections.upstreams.get(targetPath);
+  if (upstream?.connected) {
+    sendMessage(upstream.ws, message, 'master', `upstream:${targetPath}`);
+  }
+}
+
 /**
  * Handles messages received from the master connection
  * @param { MasterConnection } masterConnection - The master WebSocket connection
  * @param {string} message - The message received
  */
 function handleMasterMessage(masterConnection, message) {
-  const { path, type, targetPath, ws } = masterConnection;
-  if (type === 'root') {
-    try {
-      /** @type {{ type: string, target: string, message: string }} */
-      const data = JSON.parse(message);
-      const target = data.target;
-      const contents = data.message;
-      logger.debug(
-        `multiplexer <- master client: ${type} ${target} ${contents}`
-      );
-
-      if (data.type === 'inject') {
-        handleMasterInjection(data);
-      }
-    } catch (error) {
-      logger.error('Error processing master message:', error);
-    }
-  } else {
-    if (type === 'client') {
-      const client = connections.clients.get(targetPath);
-      if (client?.connected) {
-        sendMessage(client.ws, message, 'master', `client:${targetPath}`);
-      }
-    } else if (type === 'upstream') {
-      const upstream = connections.upstreams.get(targetPath);
-      if (upstream?.connected) {
-        sendMessage(upstream.ws, message, 'master', `upstream:${targetPath}`);
-      }
+  const { type, targetPath } = masterConnection;
+  
+  try {
+    switch (type) {
+      case 'root':
+        return handleRootMasterMessage(message);
+      case 'client':
+        return handleMasterMessageForClient(targetPath, message);
+      case 'upstream':
+        return handleMasterMessageForUpstream(targetPath, message);
     }
+  } catch (error) {
+    logger.error('Error processing master message:', error);
   }
 }