@dxos/edge-client 0.8.2-main.36232bc → 0.8.2-main.5ca3450

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@dxos/edge-client",
-  "version": "0.8.2-main.36232bc",
+  "version": "0.8.2-main.5ca3450",
   "description": "EDGE Client",
   "homepage": "https://dxos.org",
   "bugs": "https://github.com/dxos/dxos/issues",
@@ -14,6 +14,11 @@
       "browser": "./dist/lib/browser/index.mjs",
       "node": "./dist/lib/node-esm/index.mjs"
     },
+    "./muxer": {
+      "types": "./dist/types/src/edge-ws-muxer.d.ts",
+      "browser": "./dist/lib/browser/edge-ws-muxer.mjs",
+      "node": "./dist/lib/node-esm/edge-ws-muxer.mjs"
+    },
     "./testing": {
       "types": "./dist/types/src/testing/index.d.ts",
       "browser": "./dist/lib/browser/testing/index.mjs",
@@ -36,21 +41,21 @@
   "dependencies": {
     "isomorphic-ws": "^5.0.0",
     "ws": "^8.14.2",
-    "@dxos/async": "0.8.2-main.36232bc",
-    "@dxos/context": "0.8.2-main.36232bc",
-    "@dxos/credentials": "0.8.2-main.36232bc",
-    "@dxos/debug": "0.8.2-main.36232bc",
-    "@dxos/keyring": "0.8.2-main.36232bc",
-    "@dxos/keys": "0.8.2-main.36232bc",
-    "@dxos/invariant": "0.8.2-main.36232bc",
-    "@dxos/crypto": "0.8.2-main.36232bc",
-    "@dxos/log": "0.8.2-main.36232bc",
-    "@dxos/node-std": "0.8.2-main.36232bc",
-    "@dxos/util": "0.8.2-main.36232bc",
-    "@dxos/protocols": "0.8.2-main.36232bc"
+    "@dxos/context": "0.8.2-main.5ca3450",
+    "@dxos/async": "0.8.2-main.5ca3450",
+    "@dxos/credentials": "0.8.2-main.5ca3450",
+    "@dxos/debug": "0.8.2-main.5ca3450",
+    "@dxos/invariant": "0.8.2-main.5ca3450",
+    "@dxos/keys": "0.8.2-main.5ca3450",
+    "@dxos/crypto": "0.8.2-main.5ca3450",
+    "@dxos/keyring": "0.8.2-main.5ca3450",
+    "@dxos/node-std": "0.8.2-main.5ca3450",
+    "@dxos/protocols": "0.8.2-main.5ca3450",
+    "@dxos/log": "0.8.2-main.5ca3450",
+    "@dxos/util": "0.8.2-main.5ca3450"
   },
   "devDependencies": {
-    "@dxos/test-utils": "0.8.2-main.36232bc"
+    "@dxos/test-utils": "0.8.2-main.5ca3450"
   },
   "publishConfig": {
     "access": "public"

package/src/edge-ws-connection.ts

@@ -8,23 +8,18 @@ import { scheduleTask, scheduleTaskInterval } from '@dxos/async';
 import { Context, Resource } from '@dxos/context';
 import { invariant } from '@dxos/invariant';
 import { log, logInfo } from '@dxos/log';
+import { EdgeWebsocketProtocol } from '@dxos/protocols';
 import { buf } from '@dxos/protocols/buf';
 import { MessageSchema, type Message } from '@dxos/protocols/buf/dxos/edge/messenger_pb';
 
 import { protocol } from './defs';
 import { type EdgeIdentity } from './edge-identity';
-import { CLOUDFLARE_MESSAGE_LENGTH_LIMIT, WebSocketMuxer } from './edge-ws-muxer';
+import { CLOUDFLARE_MESSAGE_MAX_BYTES, WebSocketMuxer } from './edge-ws-muxer';
 import { toUint8Array } from './protocol';
 
 const SIGNAL_KEEPALIVE_INTERVAL = 4_000;
 const SIGNAL_KEEPALIVE_TIMEOUT = 12_000;
 
-const EDGE_WEBSOCKET_PROTOCOL_V0 = 'edge-ws-v0';
-/**
- * Supports message segmentation and muxing.
- */
-export const EDGE_WEBSOCKET_PROTOCOL_V1 = 'edge-ws-v1';
-
 export type EdgeWsConnectionCallbacks = {
   onConnected: () => void;
   onMessage: (message: Message) => void;
@@ -57,9 +52,9 @@ export class EdgeWsConnection extends Resource {
     invariant(this._ws);
     invariant(this._wsMuxer);
     log('sending...', { peerKey: this._identity.peerKey, payload: protocol.getPayloadType(message) });
-    if (this._ws?.protocol.includes(EDGE_WEBSOCKET_PROTOCOL_V0)) {
+    if (this._ws?.protocol.includes(EdgeWebsocketProtocol.V0)) {
       const binary = buf.toBinary(MessageSchema, message);
-      if (binary.length > CLOUDFLARE_MESSAGE_LENGTH_LIMIT) {
+      if (binary.length > CLOUDFLARE_MESSAGE_MAX_BYTES) {
         log.error('Message dropped because it was too large (>1MB).', {
           byteLength: binary.byteLength,
           serviceId: message.serviceId,
@@ -74,7 +69,7 @@ export class EdgeWsConnection extends Resource {
   }
 
   protected override async _open() {
-    const baseProtocols = [EDGE_WEBSOCKET_PROTOCOL_V0, EDGE_WEBSOCKET_PROTOCOL_V1];
+    const baseProtocols = [...Object.values(EdgeWebsocketProtocol)];
     this._ws = new WebSocket(
       this._connectionInfo.url.toString(),
       this._connectionInfo.protocolHeader
@@ -125,7 +120,7 @@ export class EdgeWsConnection extends Resource {
         return;
       }
 
-      const message = this._ws?.protocol?.includes(EDGE_WEBSOCKET_PROTOCOL_V0)
+      const message = this._ws?.protocol?.includes(EdgeWebsocketProtocol.V0)
         ? buf.fromBinary(MessageSchema, bytes)
         : muxer.receiveData(bytes);
 

package/src/edge-ws-muxer.test.ts

@@ -0,0 +1,55 @@
+//
+// Copyright 2024 DXOS.org
+//
+
+import { describe, expect, test } from 'vitest';
+
+import { bufWkt } from '@dxos/protocols/buf';
+import { TextMessageSchema } from '@dxos/protocols/buf/dxos/edge/messenger_pb';
+
+import { protocol } from './defs';
+import type { EdgeIdentity } from './edge-identity';
+import { WebSocketMuxer } from './edge-ws-muxer';
+
+const MAX_CHUNK_LENGTH = 16;
+
+describe('WebSocketMuxerTest', () => {
+  test('basic message reassembly', async () => {
+    const { muxer: muxer1, sentMessages } = await createMuxer();
+    const { muxer: muxer2 } = await createMuxer();
+    const content = 'A'.repeat(MAX_CHUNK_LENGTH);
+    await muxer1.send(textMessage(content));
+
+    expect(sentMessages.length).toBeGreaterThan(1);
+    for (const chunk of sentMessages.slice(0, -1)) {
+      expect(muxer2.receiveData(chunk)).toBeUndefined();
+    }
+
+    const reassembledMessage = muxer2.receiveData(sentMessages.at(-1)!)!;
+    expect(reassembledMessage).toBeDefined();
+
+    const decoded = bufWkt.anyUnpack(reassembledMessage.payload!, TextMessageSchema);
+    expect(decoded?.message).toStrictEqual(content);
+  });
+});
+
+const textMessage = (message: string, source?: EdgeIdentity) =>
+  protocol.createMessage(TextMessageSchema, {
+    source: source && { peerKey: source.peerKey, identityKey: source.identityKey },
+    serviceId: 'test-service',
+    payload: { message },
+  });
+
+const createMuxer = async () => {
+  const sentMessages: Uint8Array[] = [];
+  const muxer = new WebSocketMuxer(
+    {
+      readyState: 1,
+      send: (message: string) => {
+        sentMessages.push(Buffer.from(message));
+      },
+    },
+    { maxChunkLength: MAX_CHUNK_LENGTH },
+  );
+  return { muxer, sentMessages };
+};

package/src/edge-ws-muxer.ts

@@ -1,7 +1,6 @@
 //
 // Copyright 2025 DXOS.org
 //
-import WebSocket from 'isomorphic-ws';
 
 import { Trigger } from '@dxos/async';
 import { log } from '@dxos/log';
@@ -25,22 +24,30 @@ const FLAG_SEGMENT_SEQ = 1;
 const FLAG_SEGMENT_SEQ_TERMINATED = 1 << 1;
 
 /**
- * 1MB websocket message limit: https://developers.cloudflare.com/durable-objects/platform/limits/
+ * https://developers.cloudflare.com/durable-objects/platform/limits/
  */
-export const CLOUDFLARE_MESSAGE_LENGTH_LIMIT = 1024 * 1024;
+export const CLOUDFLARE_MESSAGE_MAX_BYTES = 1000 * 1000; // 1MB
+export const CLOUDFLARE_RPC_MAX_BYTES = 32 * 1000 * 1000; // 32MB
 
 const MAX_CHUNK_LENGTH = 16384;
-const MAX_BUFFERED_AMOUNT = CLOUDFLARE_MESSAGE_LENGTH_LIMIT;
+const MAX_BUFFERED_AMOUNT = CLOUDFLARE_MESSAGE_MAX_BYTES;
 const BUFFER_FULL_BACKOFF_TIMEOUT = 100;
 
 export class WebSocketMuxer {
-  private readonly _incomingMessageAccumulator = new Map<number, Buffer[]>();
-  private readonly _outgoingMessageChunks = new Map<number, MessageChunk[]>();
-  private readonly _serviceToChannel = new Map<string, number>();
+  private readonly _inMessageAccumulator = new Map<number, Buffer[]>();
+  private readonly _outMessageChunks = new Map<number, MessageChunk[]>();
+  private readonly _outMessageChannelByService = new Map<string, number>();
 
   private _sendTimeout: any | undefined;
 
-  constructor(private readonly _ws: WebSocket) {}
+  private readonly _maxChunkLength: number;
+
+  constructor(
+    private readonly _ws: WebSocketCompat,
+    config?: { maxChunkLength: number },
+  ) {
+    this._maxChunkLength = config?.maxChunkLength ?? MAX_CHUNK_LENGTH;
+  }
 
   /**
    * Resolves when all the message chunks get enqueued for sending.
@@ -48,16 +55,20 @@ export class WebSocketMuxer {
   public async send(message: Message): Promise<void> {
     const binary = buf.toBinary(MessageSchema, message);
     const channelId = this._resolveChannel(message);
-    if (channelId == null && binary.length > CLOUDFLARE_MESSAGE_LENGTH_LIMIT) {
-      log.error('Large message dropped because channel resolution failed.', {
+    if (
+      (channelId == null && binary.byteLength > CLOUDFLARE_MESSAGE_MAX_BYTES) ||
+      binary.byteLength > CLOUDFLARE_RPC_MAX_BYTES
+    ) {
+      log.error('Large message dropped', {
         byteLength: binary.byteLength,
         serviceId: message.serviceId,
         payload: protocol.getPayloadType(message),
+        channelId,
       });
       return;
     }
 
-    if (channelId == null || binary.length < MAX_CHUNK_LENGTH) {
+    if (channelId == null || binary.length < this._maxChunkLength) {
       const flags = Buffer.from([0]);
       this._ws.send(Buffer.concat([flags, binary]));
       return;
@@ -65,23 +76,23 @@ export class WebSocketMuxer {
 
     const terminatorSentTrigger = new Trigger();
     const messageChunks: MessageChunk[] = [];
-    for (let i = 0; i < binary.length; i += MAX_CHUNK_LENGTH) {
-      const chunk = binary.slice(i, i + MAX_CHUNK_LENGTH);
-      const isLastChunk = i + MAX_CHUNK_LENGTH < binary.length;
+    for (let i = 0; i < binary.length; i += this._maxChunkLength) {
+      const chunk = binary.slice(i, i + this._maxChunkLength);
+      const isLastChunk = i + this._maxChunkLength >= binary.length;
       if (isLastChunk) {
         const flags = Buffer.from([FLAG_SEGMENT_SEQ | FLAG_SEGMENT_SEQ_TERMINATED, channelId]);
         messageChunks.push({ payload: Buffer.concat([flags, chunk]), trigger: terminatorSentTrigger });
       } else {
-        const flags = Buffer.from([FLAG_SEGMENT_SEQ]);
+        const flags = Buffer.from([FLAG_SEGMENT_SEQ, channelId]);
         messageChunks.push({ payload: Buffer.concat([flags, chunk]) });
       }
     }
 
-    const queuedMessages = this._outgoingMessageChunks.get(channelId);
+    const queuedMessages = this._outMessageChunks.get(channelId);
     if (queuedMessages) {
       queuedMessages.push(...messageChunks);
     } else {
-      this._outgoingMessageChunks.set(channelId, messageChunks);
+      this._outMessageChunks.set(channelId, messageChunks);
     }
 
     this._sendChunkedMessages();
@@ -95,12 +106,12 @@ export class WebSocketMuxer {
     }
 
     const [flags, channelId, ...payload] = data;
-    let chunkAccumulator = this._incomingMessageAccumulator.get(channelId);
+    let chunkAccumulator = this._inMessageAccumulator.get(channelId);
     if (chunkAccumulator) {
       chunkAccumulator.push(Buffer.from(payload));
     } else {
       chunkAccumulator = [Buffer.from(payload)];
-      this._incomingMessageAccumulator.set(channelId, chunkAccumulator);
+      this._inMessageAccumulator.set(channelId, chunkAccumulator);
     }
 
     if ((flags & FLAG_SEGMENT_SEQ_TERMINATED) === 0) {
@@ -108,7 +119,7 @@ export class WebSocketMuxer {
     }
 
     const message = buf.fromBinary(MessageSchema, Buffer.concat(chunkAccumulator));
-    this._incomingMessageAccumulator.delete(channelId);
+    this._inMessageAccumulator.delete(channelId);
     return message;
   }
 
@@ -117,12 +128,12 @@ export class WebSocketMuxer {
       clearTimeout(this._sendTimeout);
       this._sendTimeout = undefined;
     }
-    for (const channelChunks of this._outgoingMessageChunks.values()) {
+    for (const channelChunks of this._outMessageChunks.values()) {
       channelChunks.forEach((chunk) => chunk.trigger?.wake());
     }
-    this._outgoingMessageChunks.clear();
-    this._incomingMessageAccumulator.clear();
-    this._serviceToChannel.clear();
+    this._outMessageChunks.clear();
+    this._inMessageAccumulator.clear();
+    this._outMessageChannelByService.clear();
   }
 
   private _sendChunkedMessages() {
@@ -139,10 +150,12 @@ export class WebSocketMuxer {
 
       let timeout = 0;
       const emptyChannels: number[] = [];
-      for (const [channelId, messages] of this._outgoingMessageChunks.entries()) {
-        if (this._ws.bufferedAmount + MAX_CHUNK_LENGTH > MAX_BUFFERED_AMOUNT) {
-          timeout = BUFFER_FULL_BACKOFF_TIMEOUT;
-          break;
+      for (const [channelId, messages] of this._outMessageChunks.entries()) {
+        if (this._ws.bufferedAmount != null) {
+          if (this._ws.bufferedAmount + MAX_CHUNK_LENGTH > MAX_BUFFERED_AMOUNT) {
+            timeout = BUFFER_FULL_BACKOFF_TIMEOUT;
+            break;
+          }
         }
 
         const nextMessage = messages.shift();
@@ -154,9 +167,9 @@ export class WebSocketMuxer {
         }
       }
 
-      emptyChannels.forEach((channelId) => this._outgoingMessageChunks.delete(channelId));
+      emptyChannels.forEach((channelId) => this._outMessageChunks.delete(channelId));
 
-      if (this._outgoingMessageChunks.size > 0) {
+      if (this._outMessageChunks.size > 0) {
         this._sendTimeout = setTimeout(send, timeout);
       } else {
         this._sendTimeout = undefined;
@@ -169,15 +182,24 @@ export class WebSocketMuxer {
     if (!message.serviceId) {
       return undefined;
     }
-    let id = this._serviceToChannel.get(message.serviceId);
+    let id = this._outMessageChannelByService.get(message.serviceId);
     if (!id) {
-      id = this._serviceToChannel.size + 1;
-      this._serviceToChannel.set(message.serviceId, id);
+      id = this._outMessageChannelByService.size + 1;
+      this._outMessageChannelByService.set(message.serviceId, id);
     }
     return id;
   }
 }
 
+type WebSocketCompat = {
+  readonly readyState: number;
+  /**
+   * Not available in workerd.
+   */
+  bufferedAmount?: number;
+  send(message: (ArrayBuffer | ArrayBufferView) | string): void;
+};
+
 type MessageChunk = {
   payload: Buffer;
   /**
@@ -185,3 +207,11 @@ type MessageChunk = {
    */
   trigger?: Trigger;
 };
+
+/**
+ * To avoid using isomorphic-ws on edge.
+ */
+enum WebSocket {
+  CLOSING = 2,
+  CLOSED = 3,
+}

package/src/index.ts

@@ -11,3 +11,4 @@ export * from './errors';
 export * from './auth';
 export * from './edge-http-client';
 export * from './edge-identity';
+export * from './edge-ws-muxer';

package/src/testing/test-utils.ts

@@ -6,11 +6,11 @@ import WebSocket from 'isomorphic-ws';
 
 import { Trigger } from '@dxos/async';
 import { log } from '@dxos/log';
+import { EdgeWebsocketProtocol } from '@dxos/protocols';
 import { buf } from '@dxos/protocols/buf';
 import { MessageSchema, TextMessageSchema, type Message } from '@dxos/protocols/buf/dxos/edge/messenger_pb';
 
 import { protocol } from '../defs';
-import { EDGE_WEBSOCKET_PROTOCOL_V1 } from '../edge-ws-connection';
 import { WebSocketMuxer } from '../edge-ws-muxer';
 import { toUint8Array } from '../protocol';
 
@@ -26,7 +26,7 @@ export const createTestEdgeWsServer = async (port = DEFAULT_PORT, params?: TestE
   const wsServer = new WebSocket.Server({
     port,
     verifyClient: createConnectionDelayHandler(params),
-    handleProtocols: () => [EDGE_WEBSOCKET_PROTOCOL_V1],
+    handleProtocols: () => [EdgeWebsocketProtocol.V1],
   });
 
   let connection: { ws: WebSocket; muxer: WebSocketMuxer } | undefined;