@aegis-fluxion/core 0.4.0 → 0.7.0

package/dist/index.js

@@ -4,6 +4,8 @@ import WebSocket, { WebSocketServer } from 'ws';
 // src/index.ts
 var DEFAULT_CLOSE_CODE = 1e3;
 var DEFAULT_CLOSE_REASON = "";
+var POLICY_VIOLATION_CLOSE_CODE = 1008;
+var POLICY_VIOLATION_CLOSE_REASON = "Connection rejected by middleware.";
 var INTERNAL_HANDSHAKE_EVENT = "__handshake";
 var INTERNAL_RPC_REQUEST_EVENT = "__rpc:req";
 var INTERNAL_RPC_RESPONSE_EVENT = "__rpc:res";
@@ -15,6 +17,8 @@ var GCM_AUTH_TAG_LENGTH = 16;
 var ENCRYPTION_KEY_LENGTH = 32;
 var ENCRYPTED_PACKET_VERSION = 1;
 var ENCRYPTED_PACKET_PREFIX_LENGTH = 1 + GCM_IV_LENGTH + GCM_AUTH_TAG_LENGTH;
+var BINARY_PAYLOAD_MARKER = "__afxBinaryPayload";
+var BINARY_PAYLOAD_VERSION = 1;
 var DEFAULT_HEARTBEAT_INTERVAL_MS = 15e3;
 var DEFAULT_HEARTBEAT_TIMEOUT_MS = 15e3;
 var DEFAULT_RECONNECT_INITIAL_DELAY_MS = 250;
@@ -55,7 +59,103 @@ function rawDataToBuffer(rawData) {
   }
   return Buffer.from(rawData);
 }
-function serializeEnvelope(event, data) {
+function isBlobValue(value) {
+  return typeof Blob !== "undefined" && value instanceof Blob;
+}
+function isPlainObject(value) {
+  if (typeof value !== "object" || value === null) {
+    return false;
+  }
+  const prototype = Object.getPrototypeOf(value);
+  return prototype === Object.prototype || prototype === null;
+}
+function encodeBinaryPayload(kind, payloadBuffer, mimeType) {
+  const encodedPayload = {
+    [BINARY_PAYLOAD_MARKER]: BINARY_PAYLOAD_VERSION,
+    kind,
+    base64: payloadBuffer.toString("base64")
+  };
+  if (mimeType !== void 0 && mimeType.length > 0) {
+    encodedPayload.mimeType = mimeType;
+  }
+  return encodedPayload;
+}
+async function encodeEnvelopeData(value) {
+  if (Buffer.isBuffer(value)) {
+    return encodeBinaryPayload("buffer", value);
+  }
+  if (value instanceof Uint8Array) {
+    const typedArrayBuffer = Buffer.from(value.buffer, value.byteOffset, value.byteLength);
+    return encodeBinaryPayload("uint8array", typedArrayBuffer);
+  }
+  if (isBlobValue(value)) {
+    const blobBuffer = Buffer.from(await value.arrayBuffer());
+    return encodeBinaryPayload("blob", blobBuffer, value.type);
+  }
+  if (Array.isArray(value)) {
+    return Promise.all(value.map((item) => encodeEnvelopeData(item)));
+  }
+  if (isPlainObject(value)) {
+    const encodedEntries = await Promise.all(
+      Object.entries(value).map(async ([key, entryValue]) => {
+        return [key, await encodeEnvelopeData(entryValue)];
+      })
+    );
+    return Object.fromEntries(encodedEntries);
+  }
+  return value;
+}
+function isEncodedBinaryPayload(value) {
+  if (!isPlainObject(value)) {
+    return false;
+  }
+  if (value[BINARY_PAYLOAD_MARKER] !== BINARY_PAYLOAD_VERSION) {
+    return false;
+  }
+  if (value.kind !== "buffer" && value.kind !== "uint8array" && value.kind !== "blob") {
+    return false;
+  }
+  if (typeof value.base64 !== "string") {
+    return false;
+  }
+  if (value.mimeType !== void 0 && typeof value.mimeType !== "string") {
+    return false;
+  }
+  return true;
+}
+function decodeEnvelopeData(value) {
+  if (Array.isArray(value)) {
+    return value.map((item) => decodeEnvelopeData(item));
+  }
+  if (isEncodedBinaryPayload(value)) {
+    const binaryBuffer = Buffer.from(value.base64, "base64");
+    if (value.kind === "buffer") {
+      return binaryBuffer;
+    }
+    if (value.kind === "uint8array") {
+      return Uint8Array.from(binaryBuffer);
+    }
+    if (typeof Blob === "undefined") {
+      return binaryBuffer;
+    }
+    return new Blob([binaryBuffer], {
+      type: value.mimeType ?? ""
+    });
+  }
+  if (isPlainObject(value)) {
+    const decodedEntries = Object.entries(value).map(([key, entryValue]) => {
+      return [key, decodeEnvelopeData(entryValue)];
+    });
+    return Object.fromEntries(decodedEntries);
+  }
+  return value;
+}
+async function serializeEnvelope(event, data) {
+  const encodedData = await encodeEnvelopeData(data);
+  const envelope = { event, data: encodedData };
+  return JSON.stringify(envelope);
+}
+function serializePlainEnvelope(event, data) {
   const envelope = { event, data };
   return JSON.stringify(envelope);
 }
@@ -67,7 +167,7 @@ function parseEnvelope(rawData) {
   }
   return {
     event: parsed.event,
-    data: parsed.data
+    data: decodeEnvelopeData(parsed.data)
   };
 }
 function parseEnvelopeFromText(decodedPayload) {
@@ -77,7 +177,7 @@ function parseEnvelopeFromText(decodedPayload) {
   }
   return {
     event: parsed.event,
-    data: parsed.data
+    data: decodeEnvelopeData(parsed.data)
   };
 }
 function decodeCloseReason(reason) {
@@ -259,7 +359,9 @@ var SecureServer = class {
   disconnectHandlers = /* @__PURE__ */ new Set();
   readyHandlers = /* @__PURE__ */ new Set();
   errorHandlers = /* @__PURE__ */ new Set();
+  middlewareHandlers = [];
   handshakeStateBySocket = /* @__PURE__ */ new WeakMap();
+  middlewareMetadataBySocket = /* @__PURE__ */ new WeakMap();
   sharedSecretBySocket = /* @__PURE__ */ new WeakMap();
   encryptionKeyBySocket = /* @__PURE__ */ new WeakMap();
   pendingPayloadsBySocket = /* @__PURE__ */ new WeakMap();
@@ -348,6 +450,19 @@ var SecureServer = class {
     }
     return this;
   }
+  use(middleware) {
+    try {
+      if (typeof middleware !== "function") {
+        throw new Error("Server middleware must be a function.");
+      }
+      this.middlewareHandlers.push(middleware);
+    } catch (error) {
+      this.notifyError(
+        normalizeToError(error, "Failed to register server middleware.")
+      );
+    }
+    return this;
+  }
   emit(event, data) {
     try {
       if (isReservedEmitEvent(event)) {
@@ -355,7 +470,9 @@ var SecureServer = class {
       }
       const envelope = { event, data };
       for (const client of this.clientsById.values()) {
-        this.sendOrQueuePayload(client.socket, envelope);
+        void this.sendOrQueuePayload(client.socket, envelope).catch(() => {
+          return void 0;
+        });
       }
     } catch (error) {
       this.notifyError(normalizeToError(error, "Failed to emit server event."));
@@ -373,7 +490,9 @@ var SecureServer = class {
         throw new Error(`Client with id ${clientId} was not found.`);
       }
       if (!ackArgs.expectsAck) {
-        this.sendOrQueuePayload(client.socket, { event, data });
+        void this.sendOrQueuePayload(client.socket, { event, data }).catch(() => {
+          return void 0;
+        });
         return true;
       }
       const ackPromise = this.sendRpcRequest(
@@ -429,6 +548,7 @@ var SecureServer = class {
           client.socket,
           new Error("Server closed before ACK response was received.")
         );
+        this.middlewareMetadataBySocket.delete(client.socket);
         if (client.socket.readyState === WebSocket.OPEN || client.socket.readyState === WebSocket.CONNECTING) {
           client.socket.close(code, reason);
         }
@@ -491,6 +611,7 @@ var SecureServer = class {
         this.pendingRpcRequestsBySocket.delete(socket);
         this.handshakeStateBySocket.delete(socket);
         this.heartbeatStateBySocket.delete(socket);
+        this.middlewareMetadataBySocket.delete(socket);
         socket.terminate();
         continue;
       }
@@ -520,17 +641,46 @@ var SecureServer = class {
   }
   bindSocketServerEvents() {
     this.socketServer.on("connection", (socket, request) => {
-      this.handleConnection(socket, request);
+      void this.handleConnection(socket, request);
     });
     this.socketServer.on("error", (error) => {
       this.notifyError(normalizeToError(error, "WebSocket server encountered an error."));
     });
   }
-  handleConnection(socket, request) {
+  async handleConnection(socket, request) {
+    const connectionMetadata = /* @__PURE__ */ new Map();
+    this.middlewareMetadataBySocket.set(socket, connectionMetadata);
+    try {
+      await this.executeServerMiddleware({
+        phase: "connection",
+        socket,
+        request,
+        metadata: connectionMetadata
+      });
+    } catch (error) {
+      const normalizedError = normalizeToError(
+        error,
+        "Connection middleware rejected the incoming socket."
+      );
+      this.notifyError(normalizedError);
+      this.middlewareMetadataBySocket.delete(socket);
+      if (socket.readyState === WebSocket.OPEN || socket.readyState === WebSocket.CONNECTING) {
+        socket.close(
+          POLICY_VIOLATION_CLOSE_CODE,
+          normalizedError.message || POLICY_VIOLATION_CLOSE_REASON
+        );
+      }
+      return;
+    }
     try {
       const clientId = randomUUID();
       const handshakeState = this.createServerHandshakeState();
-      const client = this.createSecureServerClient(clientId, socket, request);
+      const client = this.createSecureServerClient(
+        clientId,
+        socket,
+        request,
+        connectionMetadata
+      );
       this.clientsById.set(clientId, client);
       this.clientIdBySocket.set(socket, clientId);
       this.handshakeStateBySocket.set(socket, handshakeState);
@@ -542,7 +692,7 @@ var SecureServer = class {
       });
       this.roomNamesByClientId.set(clientId, /* @__PURE__ */ new Set());
       socket.on("message", (rawData) => {
-        this.handleIncomingMessage(client, rawData);
+        void this.handleIncomingMessage(client, rawData);
       });
       socket.on("close", (code, reason) => {
         this.handleDisconnection(client, code, reason);
@@ -562,9 +712,10 @@ var SecureServer = class {
       this.notifyConnection(client);
     } catch (error) {
       this.notifyError(normalizeToError(error, "Failed to handle client connection."));
+      this.middlewareMetadataBySocket.delete(socket);
     }
   }
-  handleIncomingMessage(client, rawData) {
+  async handleIncomingMessage(client, rawData) {
     try {
       let envelope = null;
       try {
@@ -608,10 +759,16 @@ var SecureServer = class {
         return;
       }
       if (decryptedEnvelope.event === INTERNAL_RPC_REQUEST_EVENT) {
-        void this.handleRpcRequest(client, decryptedEnvelope.data);
+        await this.handleRpcRequest(client, decryptedEnvelope.data);
         return;
       }
-      this.dispatchCustomEvent(decryptedEnvelope.event, decryptedEnvelope.data, client);
+      const interceptedData = await this.applyMessageMiddleware(
+        "incoming",
+        client,
+        decryptedEnvelope.event,
+        decryptedEnvelope.data
+      );
+      this.dispatchCustomEvent(decryptedEnvelope.event, interceptedData, client);
     } catch (error) {
       this.notifyError(normalizeToError(error, "Failed to process incoming server message."));
     }
@@ -631,6 +788,7 @@ var SecureServer = class {
       );
       this.pendingRpcRequestsBySocket.delete(client.socket);
       this.heartbeatStateBySocket.delete(client.socket);
+      this.middlewareMetadataBySocket.delete(client.socket);
       const decodedReason = decodeCloseReason(reason);
       for (const handler of this.disconnectHandlers) {
         try {
@@ -676,6 +834,48 @@ var SecureServer = class {
       }
     }
   }
+  async executeServerMiddleware(context) {
+    if (this.middlewareHandlers.length === 0) {
+      return;
+    }
+    let currentIndex = -1;
+    const dispatch = async (index) => {
+      if (index <= currentIndex) {
+        throw new Error("Server middleware next() was called multiple times.");
+      }
+      currentIndex = index;
+      const middleware = this.middlewareHandlers[index];
+      if (!middleware) {
+        return;
+      }
+      await Promise.resolve(
+        middleware(context, async () => {
+          await dispatch(index + 1);
+        })
+      );
+    };
+    await dispatch(0);
+  }
+  async applyMessageMiddleware(phase, client, event, data) {
+    const metadata = this.middlewareMetadataBySocket.get(client.socket) ?? /* @__PURE__ */ new Map();
+    this.middlewareMetadataBySocket.set(client.socket, metadata);
+    const middlewareContext = {
+      phase,
+      client,
+      event,
+      data,
+      metadata
+    };
+    await this.executeServerMiddleware(middlewareContext);
+    return middlewareContext.data;
+  }
+  resolveClientBySocket(socket) {
+    const clientId = this.clientIdBySocket.get(socket);
+    if (!clientId) {
+      return null;
+    }
+    return this.clientsById.get(clientId) ?? null;
+  }
   sendRaw(socket, payload) {
     try {
       if (socket.readyState !== WebSocket.OPEN) {
@@ -686,22 +886,24 @@ var SecureServer = class {
       this.notifyError(normalizeToError(error, "Failed to send server payload."));
     }
   }
-  sendEncryptedEnvelope(socket, envelope) {
+  async sendEncryptedEnvelope(socket, envelope) {
+    if (socket.readyState !== WebSocket.OPEN) {
+      return;
+    }
+    const encryptionKey = this.encryptionKeyBySocket.get(socket);
+    if (!encryptionKey) {
+      const missingKeyError = new Error("Missing encryption key for connected socket.");
+      this.notifyError(missingKeyError);
+      throw missingKeyError;
+    }
     try {
-      if (socket.readyState !== WebSocket.OPEN) {
-        return;
-      }
-      const encryptionKey = this.encryptionKeyBySocket.get(socket);
-      if (!encryptionKey) {
-        throw new Error("Missing encryption key for connected socket.");
-      }
-      const encryptedPayload = encryptSerializedEnvelope(
-        serializeEnvelope(envelope.event, envelope.data),
-        encryptionKey
-      );
+      const serializedEnvelope = await serializeEnvelope(envelope.event, envelope.data);
+      const encryptedPayload = encryptSerializedEnvelope(serializedEnvelope, encryptionKey);
       socket.send(encryptedPayload);
     } catch (error) {
-      this.notifyError(normalizeToError(error, "Failed to send encrypted server payload."));
+      const normalizedError = normalizeToError(error, "Failed to send encrypted server payload.");
+      this.notifyError(normalizedError);
+      throw normalizedError;
     }
   }
   sendRpcRequest(socket, event, data, timeoutMs) {
@@ -722,13 +924,19 @@ var SecureServer = class {
         reject,
         timeoutHandle
       });
-      this.sendOrQueuePayload(socket, {
+      void this.sendOrQueuePayload(socket, {
         event: INTERNAL_RPC_REQUEST_EVENT,
         data: {
           id: requestId,
           event,
           data
         }
+      }).catch((error) => {
+        clearTimeout(timeoutHandle);
+        pendingRequests.delete(requestId);
+        reject(
+          normalizeToError(error, `Failed to dispatch ACK request for event "${event}".`)
+        );
       });
     });
   }
@@ -765,12 +973,18 @@ var SecureServer = class {
       return;
     }
     try {
+      const interceptedData = await this.applyMessageMiddleware(
+        "incoming",
+        client,
+        rpcRequestPayload.event,
+        rpcRequestPayload.data
+      );
       const ackResponse = await this.executeRpcRequestHandler(
         rpcRequestPayload.event,
-        rpcRequestPayload.data,
+        interceptedData,
         client
       );
-      this.sendEncryptedEnvelope(client.socket, {
+      await this.sendEncryptedEnvelope(client.socket, {
         event: INTERNAL_RPC_RESPONSE_EVENT,
         data: {
           id: rpcRequestPayload.id,
@@ -780,7 +994,7 @@ var SecureServer = class {
       });
     } catch (error) {
       const normalizedError = normalizeToError(error, "Server ACK request handler failed.");
-      this.sendEncryptedEnvelope(client.socket, {
+      await this.sendEncryptedEnvelope(client.socket, {
         event: INTERNAL_RPC_RESPONSE_EVENT,
         data: {
           id: rpcRequestPayload.id,
@@ -854,7 +1068,7 @@ var SecureServer = class {
   sendInternalHandshake(socket, localPublicKey) {
     this.sendRaw(
       socket,
-      serializeEnvelope(INTERNAL_HANDSHAKE_EVENT, {
+      serializePlainEnvelope(INTERNAL_HANDSHAKE_EVENT, {
         publicKey: localPublicKey
       })
     );
@@ -884,33 +1098,50 @@ var SecureServer = class {
   isClientHandshakeReady(socket) {
     return this.handshakeStateBySocket.get(socket)?.isReady ?? false;
   }
-  sendOrQueuePayload(socket, envelope) {
+  async sendOrQueuePayload(socket, envelope) {
+    let interceptedEnvelope = envelope;
+    if (!isReservedEmitEvent(envelope.event)) {
+      const targetClient = this.resolveClientBySocket(socket);
+      if (targetClient) {
+        const interceptedData = await this.applyMessageMiddleware(
+          "outgoing",
+          targetClient,
+          envelope.event,
+          envelope.data
+        );
+        interceptedEnvelope = {
+          event: envelope.event,
+          data: interceptedData
+        };
+      }
+    }
     if (!this.isClientHandshakeReady(socket)) {
-      this.queuePayload(socket, envelope);
+      this.queuePayload(socket, interceptedEnvelope);
       return;
     }
-    this.sendEncryptedEnvelope(socket, envelope);
+    await this.sendEncryptedEnvelope(socket, interceptedEnvelope);
   }
   queuePayload(socket, envelope) {
     const pendingPayloads = this.pendingPayloadsBySocket.get(socket) ?? [];
     pendingPayloads.push(envelope);
     this.pendingPayloadsBySocket.set(socket, pendingPayloads);
   }
-  flushQueuedPayloads(socket) {
+  async flushQueuedPayloads(socket) {
     const pendingPayloads = this.pendingPayloadsBySocket.get(socket);
     if (!pendingPayloads || pendingPayloads.length === 0) {
       return;
     }
     this.pendingPayloadsBySocket.delete(socket);
     for (const envelope of pendingPayloads) {
-      this.sendEncryptedEnvelope(socket, envelope);
+      await this.sendEncryptedEnvelope(socket, envelope);
     }
   }
-  createSecureServerClient(clientId, socket, request) {
+  createSecureServerClient(clientId, socket, request, metadata) {
     return {
       id: clientId,
       socket,
       request,
+      metadata,
       emit: (event, data, callbackOrOptions, maybeCallback) => {
         if (callbackOrOptions === void 0 && maybeCallback === void 0) {
           return this.emitTo(clientId, event, data);
@@ -1013,7 +1244,9 @@ var SecureServer = class {
       if (!client) {
         continue;
       }
-      this.sendOrQueuePayload(client.socket, envelope);
+      void this.sendOrQueuePayload(client.socket, envelope).catch(() => {
+        return void 0;
+      });
     }
   }
 };
@@ -1177,7 +1410,9 @@ var SecureClient = class {
         this.pendingPayloadQueue.push(envelope);
         return true;
       }
-      this.sendEncryptedEnvelope(envelope);
+      void this.sendEncryptedEnvelope(envelope).catch(() => {
+        return void 0;
+      });
       return true;
     } catch (error) {
       const normalizedError = normalizeToError(error, "Failed to emit client event.");
@@ -1423,22 +1658,26 @@ var SecureClient = class {
       }
     }
   }
-  sendEncryptedEnvelope(envelope) {
+  async sendEncryptedEnvelope(envelope) {
+    if (!this.socket || this.socket.readyState !== WebSocket.OPEN) {
+      const socketStateError = new Error("Client socket is not connected.");
+      this.notifyError(socketStateError);
+      throw socketStateError;
+    }
+    const encryptionKey = this.handshakeState?.encryptionKey;
+    if (!encryptionKey) {
+      const missingKeyError = new Error("Missing encryption key for client payload encryption.");
+      this.notifyError(missingKeyError);
+      throw missingKeyError;
+    }
     try {
-      if (!this.socket || this.socket.readyState !== WebSocket.OPEN) {
-        throw new Error("Client socket is not connected.");
-      }
-      const encryptionKey = this.handshakeState?.encryptionKey;
-      if (!encryptionKey) {
-        throw new Error("Missing encryption key for client payload encryption.");
-      }
-      const encryptedPayload = encryptSerializedEnvelope(
-        serializeEnvelope(envelope.event, envelope.data),
-        encryptionKey
-      );
+      const serializedEnvelope = await serializeEnvelope(envelope.event, envelope.data);
+      const encryptedPayload = encryptSerializedEnvelope(serializedEnvelope, encryptionKey);
       this.socket.send(encryptedPayload);
     } catch (error) {
-      this.notifyError(normalizeToError(error, "Failed to send encrypted client payload."));
+      const normalizedError = normalizeToError(error, "Failed to send encrypted client payload.");
+      this.notifyError(normalizedError);
+      throw normalizedError;
     }
   }
   sendRpcRequest(event, data, timeoutMs) {
@@ -1469,7 +1708,13 @@ var SecureClient = class {
         this.pendingPayloadQueue.push(rpcRequestEnvelope);
         return;
       }
-      this.sendEncryptedEnvelope(rpcRequestEnvelope);
+      void this.sendEncryptedEnvelope(rpcRequestEnvelope).catch((error) => {
+        clearTimeout(timeoutHandle);
+        this.pendingRpcRequests.delete(requestId);
+        reject(
+          normalizeToError(error, `Failed to dispatch ACK request for event "${event}".`)
+        );
+      });
     });
   }
   handleRpcResponse(data) {
@@ -1505,7 +1750,7 @@ var SecureClient = class {
         rpcRequestPayload.event,
         rpcRequestPayload.data
       );
-      this.sendEncryptedEnvelope({
+      await this.sendEncryptedEnvelope({
         event: INTERNAL_RPC_RESPONSE_EVENT,
         data: {
           id: rpcRequestPayload.id,
@@ -1515,7 +1760,7 @@ var SecureClient = class {
       });
     } catch (error) {
       const normalizedError = normalizeToError(error, "Client ACK request handler failed.");
-      this.sendEncryptedEnvelope({
+      await this.sendEncryptedEnvelope({
         event: INTERNAL_RPC_RESPONSE_EVENT,
         data: {
           id: rpcRequestPayload.id,
@@ -1560,7 +1805,7 @@ var SecureClient = class {
         throw new Error("Missing client handshake state.");
       }
       this.socket.send(
-        serializeEnvelope(INTERNAL_HANDSHAKE_EVENT, {
+        serializePlainEnvelope(INTERNAL_HANDSHAKE_EVENT, {
           publicKey: this.handshakeState.localPublicKey
         })
       );
@@ -1582,7 +1827,7 @@ var SecureClient = class {
       this.handshakeState.sharedSecret = sharedSecret;
       this.handshakeState.encryptionKey = deriveEncryptionKey(sharedSecret);
       this.handshakeState.isReady = true;
-      this.flushPendingPayloadQueue();
+      void this.flushPendingPayloadQueue();
       this.notifyReady();
     } catch (error) {
       this.notifyError(normalizeToError(error, "Failed to complete client handshake."));
@@ -1591,14 +1836,14 @@ var SecureClient = class {
   isHandshakeReady() {
     return this.handshakeState?.isReady ?? false;
   }
-  flushPendingPayloadQueue() {
+  async flushPendingPayloadQueue() {
     if (!this.socket || this.socket.readyState !== WebSocket.OPEN || !this.isHandshakeReady()) {
       return;
     }
     const pendingPayloads = this.pendingPayloadQueue;
     this.pendingPayloadQueue = [];
     for (const envelope of pendingPayloads) {
-      this.sendEncryptedEnvelope(envelope);
+      await this.sendEncryptedEnvelope(envelope);
     }
   }
 };