fix: stabilize relay E2EE handshake

packages/relay/src/encrypted-channel.ts

@@ -87,19 +87,28 @@ export async function createDaemonChannel(
       reject(new Error("Handshake timeout"));
     }, 10000);
 
+    const bufferedMessages: Array<string | ArrayBuffer> = [];
+
     transport.onmessage = async (data) => {
       try {
-        if (typeof data !== "string") {
-          throw new Error("Expected string hello message");
-        }
+        const helloText =
+          typeof data === "string" ? data : new TextDecoder().decode(data);
 
-        const msg = JSON.parse(data) as HelloMessage;
+        const msg = JSON.parse(helloText) as HelloMessage;
         if (msg.type !== "hello" || !msg.key) {
           throw new Error("Invalid hello message");
         }
 
         clearTimeout(timeout);
 
+        // Buffer any subsequent messages that arrive while we're doing async
+        // WebCrypto work to derive the shared key. Without this, it's possible
+        // for the next message (already encrypted) to be misinterpreted as a
+        // second hello, causing the handshake to fail.
+        transport.onmessage = (next) => {
+          bufferedMessages.push(next);
+        };
+
         const clientPublicKey = await importPublicKey(msg.key);
         const sharedKey = await deriveSharedKey(
           daemonKeyPair.privateKey,
@@ -110,6 +119,10 @@ export async function createDaemonChannel(
         channel.setState("open");
         events.onopen?.();
 
+        for (const buffered of bufferedMessages) {
+          transport.onmessage?.(buffered);
+        }
+
         resolve(channel);
       } catch (error) {
         clearTimeout(timeout);

packages/relay/src/node-adapter.ts

@@ -72,15 +72,36 @@ export function createRelayServer(config: NodeRelayServerConfig): RelayServer {
     const sessionId = url.searchParams.get("session")!;
     const role = url.searchParams.get("role") as ConnectionRole;
 
-    const connection = wrapWebSocket(ws);
+    const connection = wrapWebSocket(ws, role);
     relay.addConnection(sessionId, role, connection);
 
-    ws.on("message", (data) => {
-      const message =
-        data instanceof Buffer
-          ? data.buffer.slice(data.byteOffset, data.byteOffset + data.byteLength)
-          : String(data);
-      relay.forward(sessionId, role, message as string | ArrayBuffer);
+    ws.on("message", (data, isBinary) => {
+      if (isBinary) {
+        const message =
+          data instanceof ArrayBuffer
+            ? data
+            : bufferToArrayBuffer(
+                data instanceof Buffer
+                  ? data
+                  : ArrayBuffer.isView(data)
+                    ? Buffer.from(data.buffer, data.byteOffset, data.byteLength)
+                    : Buffer.from(String(data), "utf8")
+              );
+        relay.forward(sessionId, role, message);
+        return;
+      }
+
+      const text =
+        typeof data === "string"
+          ? data
+          : data instanceof Buffer
+            ? data.toString("utf8")
+            : data instanceof ArrayBuffer
+              ? Buffer.from(data).toString("utf8")
+              : ArrayBuffer.isView(data)
+                ? Buffer.from(data.buffer, data.byteOffset, data.byteLength).toString("utf8")
+                : String(data);
+      relay.forward(sessionId, role, text);
     });
 
     ws.on("close", () => {
@@ -120,9 +141,9 @@ export function createRelayServer(config: NodeRelayServerConfig): RelayServer {
   };
 }
 
-function wrapWebSocket(ws: NodeWebSocket): RelayConnection {
+function wrapWebSocket(ws: NodeWebSocket, role: ConnectionRole): RelayConnection {
   return {
-    role: "server",
+    role,
     send: (data) => {
       if (ws.readyState === NodeWebSocket.OPEN) {
         ws.send(data);
@@ -133,3 +154,9 @@ function wrapWebSocket(ws: NodeWebSocket): RelayConnection {
     },
   };
 }
+
+function bufferToArrayBuffer(buffer: Buffer): ArrayBuffer {
+  const out = new Uint8Array(buffer.byteLength);
+  out.set(buffer);
+  return out.buffer;
+}