replica-bridge 0.3.5 → 0.3.6

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "replica-bridge",
-  "version": "0.3.5",
+  "version": "0.3.6",
   "description": "Lets the ChatReplica iPhone app use Codex on this computer through a relay, end-to-end encrypted.",
   "type": "module",
   "bin": {
@@ -11,7 +11,7 @@
   },
   "dependencies": {
     "qrcode": "1.5.4",
-    "ws": "8.18.3"
+    "ws": "8.22.0"
   },
   "license": "UNLICENSED"
 }

package/src/bridge/service.js

@@ -10,7 +10,7 @@ import { Duplex } from "node:stream";
 import { WebSocket } from "ws";
                                   
 import {
-  CloseCode, FrameCipher, MAX_FRAME_BYTES, RELAY_REPLY_LIMIT, StreamBuffer, b64url, fromB64url, generateKeyPair, hostSessionKeys,
+  CloseCode, FrameCipher, ChunkAssembler, MAX_FRAME_BYTES, REASSEMBLED_MAX_BYTES, RELAY_REPLY_LIMIT, StreamBuffer, encodeLine, b64url, fromB64url, generateKeyPair, hostSessionKeys,
 } from "../protocol.js";
                                                                             
 import { relaySocketURL } from "../http.js";
@@ -50,6 +50,12 @@ import { annotateWriterError } from "./writer.js";
                                   
                                                                                                                        
                          
+                                                                                                     
+                     
+                                                  
+                         
+                                                                                                                            
+                              
  
 
                 
@@ -125,6 +131,14 @@ export class BridgeService {
     });
   }
 
+  /**
+   * Waits while the relay socket has more than a few chunks queued, so a long message goes out at the network's pace
+   * instead of tens of MiB piling up in memory ahead of everything else on this computer's socket.
+   */
+          async drained() {
+    while (this.relay?.readyState === WebSocket.OPEN && this.relay.bufferedAmount > RELAY_SEND_QUEUE_BYTES) await new Promise((resolve) => setTimeout(resolve, 10));
+  }
+
           sendEnvelope(envelope              ) {
     if (this.relay?.readyState === WebSocket.OPEN) this.relay.send(JSON.stringify(envelope));
   }
@@ -191,12 +205,20 @@ export class BridgeService {
     const message = conn.cipher.openFrame(frame);
     if (message.t === "resume") {
       if (conn.stream) throw new Error("second resume");
-      await this.resume(conn, message.stream, message.ack);
+      await this.resume(conn, message.stream, message.ack, message.seg === 1);
     } else if (message.t === "msg") {
       const stream = conn.stream;
       if (!stream) throw new Error("message before resume");
       if (!stream.in.accept(message.seq)) return;
-      if (stream.appServer.readyState === WebSocket.OPEN) stream.appServer.send(message.line);
+      if (stream.chunks.busy) throw new Error("message in the middle of a chunked one");
+      this.sendToAppServer(stream, message.line);
+      this.scheduleAck(stream);
+    } else if (message.t === "chunk") {
+      const stream = conn.stream;
+      if (!stream) throw new Error("message before resume");
+      if (!stream.in.accept(message.seq)) return;
+      const line = stream.chunks.take(message);
+      if (line !== null) this.sendToAppServer(stream, line);
       this.scheduleAck(stream);
     } else if (message.t === "ack") {
       conn.stream?.out.ack(message.seq);
@@ -209,7 +231,7 @@ export class BridgeService {
 
   // MARK: - Streams
 
-          async resume(conn      , streamId               , peerAck        ) {
+          async resume(conn      , streamId               , peerAck        , segmented         ) {
     let stream = streamId ? this.streams.get(streamId) : undefined;
     const resumed = !!stream && stream.deviceId === conn.deviceId;
     if (!resumed) {
@@ -219,31 +241,40 @@ export class BridgeService {
     if (s.conn && s.conn !== conn) this.closeConn(s.conn, "resumed elsewhere");
     if (s.expiry) { clearTimeout(s.expiry); s.expiry = null; }
     s.conn = conn;
+    s.segmented = segmented;
     conn.stream = s;
-    this.sendInner(conn, { t: "ready", stream: s.id, resumed, ack: s.in.received });
+    this.sendInner(conn, { t: "ready", stream: s.id, resumed, ack: s.in.received, seg: 1 });
     // A new stream's sequence numbers start over, whatever the phone acknowledged on its old one.
     const after = resumed ? peerAck : 0;
     s.out.ack(after);
-    for (const item of s.out.pending(after)) this.sendInner(conn, { t: "msg", seq: item.seq, line: item.line });
+    for (const message of s.out.pending(after)) this.sendInner(conn, message);
     this.log(`${resumed ? "resumed" : "opened"} stream ${s.id} for ${conn.deviceId}`);
   }
 
           async openStream(deviceId        )                  {
-    const appServer = await this.connectAppServer();
-    const stream         = { id: crypto.randomBytes(12).toString("base64url"), deviceId, appServer, out: new StreamBuffer(), in: new StreamBuffer(), conn: null, expiry: null, ackTimer: null, relayed: Promise.resolve() };
+    const { ws: appServer, frameBytes } = await this.connectAppServer();
+    const stream         = {
+      id: crypto.randomBytes(12).toString("base64url"), deviceId, appServer, out: new StreamBuffer(STREAM_BUFFER_BYTES), in: new StreamBuffer(STREAM_BUFFER_BYTES),
+      conn: null, expiry: null, ackTimer: null, relayed: Promise.resolve(), segmented: false, chunks: new ChunkAssembler(), appServerFrameBytes: frameBytes,
+    };
     this.streams.set(stream.id, stream);
     appServer.on("message", (data         ) => {
       const text = data.toString();
-      const fitted = fitLine(text);
+      const segmented = stream.segmented;
+      const fitted = fitLine(text, segmented ? REASSEMBLED_MAX_BYTES : RELAY_REPLY_LIMIT);
       if (fitted === null) return this.log(`stream ${stream.id}: dropped a ${text.length}-byte app-server notification too large for the relay`);
       if (fitted !== text) this.log(`stream ${stream.id}: answered a ${text.length}-byte app-server reply too large for the relay with an error`);
       const annotated = annotateWriterError(fitted, this.options.findWriter);
       stream.relayed = stream.relayed.then(async () => {
         const line = await annotated;
-        if (!this.streams.has(stream.id)) return;
-        let item;
-        try { item = stream.out.push(line); } catch { return this.endStream(stream, "the phone stayed away too long and the buffer filled up"); }
-        if (stream.conn) this.sendInner(stream.conn, { t: "msg", seq: item.seq, line: item.line });
+        const parts = segmented ? await encodeLine(line) : [line];
+        for (const part of parts) {
+          if (!this.streams.has(stream.id)) return;
+          let message;
+          try { message = stream.out.push(part); } catch { return this.endStream(stream, "the phone stayed away too long and the buffer filled up"); }
+          if (stream.conn) this.sendInner(stream.conn, message);
+          if (parts.length > 1) await this.drained();
+        }
       });
     });
     appServer.on("close", () => this.endStream(stream, "codex app-server closed the connection"));
@@ -282,11 +313,22 @@ export class BridgeService {
 
   // MARK: - codex app-server
 
+  /** Sends a line to the app-server, as WebSocket fragments when it is longer than the app-server takes in one frame. */
+          sendToAppServer(stream        , line        ) {
+    const ws = stream.appServer;
+    if (ws.readyState !== WebSocket.OPEN) return;
+    const bytes = Buffer.from(line, "utf8");
+    if (bytes.length <= stream.appServerFrameBytes) return ws.send(line);
+    for (let at = 0; at < bytes.length; at += stream.appServerFrameBytes) {
+      ws.send(bytes.subarray(at, at + stream.appServerFrameBytes), { binary: false, fin: at + stream.appServerFrameBytes >= bytes.length });
+    }
+  }
+
   /**
    * One `codex app-server proxy` per stream, with WebSocket over its stdin and stdout. The proxy only relays bytes to
    * the app-server's control socket, wherever Codex keeps it; it exits at once when no app-server is running.
    */
-          openProxy()                     {
+          openProxy()                                                 {
     return new Promise((resolve, reject) => {
       const child               = spawnCommand(this.options.codexPath ?? "codex", ["app-server", "proxy"], { stdio: ["pipe", "pipe", "pipe"] });
       let stderr = "";
@@ -297,13 +339,18 @@ export class BridgeService {
       // Read app-server messages well past what one relay frame holds, so fitLine answers a big reply with an error
       // rather than the oversized message closing the app-server connection.
       const ws = new WebSocket("ws://localhost/", { createConnection: () => pipe         , perMessageDeflate: false, maxPayload: APP_SERVER_MAX_MESSAGE_BYTES });
+      let frameBytes = APP_SERVER_FRAME_BYTES;
+      ws.once("upgrade", (res) => {
+        const advertised = Number(res.headers["x-codex-websocket-max-unfragmented-message-bytes"]);
+        if (Number.isInteger(advertised) && advertised >= 64 * 1024) frameBytes = advertised;
+      });
       const fail = (why        ) => { ws.removeAllListeners("open"); killCommand(child); reject(new Error(why)); };
       child.once("error", (error) => fail(`could not run codex: ${error.message}`));
       child.once("exit", (code) => { if (ws.readyState !== WebSocket.OPEN) fail(`codex app-server proxy exited (${code})${stderr ? `: ${stderr.trim()}` : ""}`); });
       ws.once("error", (error) => { if (ws.readyState !== WebSocket.OPEN) fail(error.message); });
       ws.once("open", () => {
         ws.on("close", () => killCommand(child));
-        resolve(ws);
+        resolve({ ws, frameBytes });
       });
     });
   }
@@ -318,7 +365,7 @@ export class BridgeService {
     });
   }
 
-          async connectAppServer()                     {
+          async connectAppServer()                                                 {
     try {
       return await this.openProxy();
     } catch (error) {
@@ -349,16 +396,24 @@ export class BridgeService {
 /** Largest app-server message the bridge reads before refusing it. */
 const APP_SERVER_MAX_MESSAGE_BYTES = 256 * 1024 * 1024;
 
+/** What the bridge keeps per direction for a phone that is away: room for a REASSEMBLED_MAX_BYTES line deflate could not shrink. */
+const STREAM_BUFFER_BYTES = 128 * 1024 * 1024;
+/** Chunks queued on the relay socket before the next one waits. */
+const RELAY_SEND_QUEUE_BYTES = 1024 * 1024;
+/** What codex-cli 0.163.0 sends in x-codex-websocket-max-unfragmented-message-bytes, for an app-server that sends none. */
+const APP_SERVER_FRAME_BYTES = 16 * 1024 * 1024;
+
 /**
- * An app-server line as it can cross the relay: unchanged when it fits RELAY_REPLY_LIMIT, a JSON-RPC error with the
- * same id when it is a reply that does not (`error.data.replicaTooLarge: {bytes, limit}`, PROTOCOL.md), and null for a
- * notification that does not, which has no one to answer.
+ * An app-server line as it can cross the relay: unchanged when it fits `limit` (RELAY_REPLY_LIMIT, or
+ * REASSEMBLED_MAX_BYTES for a phone that takes chunks), a JSON-RPC error with the same id when it is a reply that does
+ * not (`error.data.replicaTooLarge: {bytes, limit}`, PROTOCOL.md), and null for a notification that does not, which
+ * has no one to answer.
  */
-export function fitLine(line        )                {
+export function fitLine(line        , limit = RELAY_REPLY_LIMIT)                {
   const bytes = Buffer.byteLength(line, "utf8");
-  if (bytes <= RELAY_REPLY_LIMIT) return line;
+  if (bytes <= limit) return line;
   let id         ;
   try { id = (JSON.parse(line)                    ).id; } catch { return null; }
   if (id === undefined || id === null) return null;
-  return JSON.stringify({ id, error: { code: -32010, message: `the reply is ${bytes} bytes, more than the relay carries (${RELAY_REPLY_LIMIT})`, data: { replicaTooLarge: { bytes, limit: RELAY_REPLY_LIMIT } } } });
+  return JSON.stringify({ id, error: { code: -32010, message: `the reply is ${bytes} bytes, more than the relay carries (${limit})`, data: { replicaTooLarge: { bytes, limit } } } });
 }

package/src/protocol.js

@@ -2,6 +2,8 @@
 // Everything here is plain Node: X25519, HKDF-SHA256, ChaCha20-Poly1305 and HMAC-SHA256 match CryptoKit's
 // Curve25519.KeyAgreement, HKDF<SHA256>, ChaChaPoly and HMAC<SHA256>, so the Swift side can follow the same steps.
 import crypto from "node:crypto";
+import zlib from "node:zlib";
+import { promisify } from "node:util";
 
 export const PROTOCOL = "chatreplica-relay-v1";
 export const LINK_SCHEME = "chatreplica:";
@@ -15,6 +17,17 @@ export const MAX_FRAME_BYTES = 16 * 1024 * 1024;
  * of a big file) come back as a JSON-RPC error instead of a frame the relay would refuse by closing the connection.
  */
 export const RELAY_REPLY_LIMIT = 11 * 1024 * 1024;
+/**
+ * Long messages, the way Codex remote-control splits them (openai/codex app-server-transport remote_control/segment.rs):
+ * when both sides say `seg` (the phone in resume, the bridge in ready), a line over SEGMENT_BYTES goes as `chunk`
+ * messages of at most SEGMENT_BYTES each, at most SEGMENT_COUNT_MAX of them, up to REASSEMBLED_MAX_BYTES; the receiver
+ * checks every chunk's segment, count and size. Unlike remote-control the payload is raw deflate of the line (`z`),
+ * since the relay's frames are encrypted and nothing on the way can compress them. A chunk travels sealed and
+ * base64url encoded in about 180 KiB of frame, so pings still get through while a long message is on its way.
+ */
+export const SEGMENT_BYTES = 100 * 1024;
+export const SEGMENT_COUNT_MAX = 1024;
+export const REASSEMBLED_MAX_BYTES = 100 * 1024 * 1024;
 
 export const b64url = (data            )         => Buffer.from(data).toString("base64url");
 export const fromB64url = (text        )         => Buffer.from(text, "base64url");
@@ -208,12 +221,78 @@ export class FrameCipher {
 
 /** What is inside a data frame. */
                           
-                                                       
-                                                                 
-                                           
+                                                                
+                                                                          
+              
+                
                              
                                     
 
+                                                                 
+/** One segment of a long line: `data` is base64 of bytes [segment * chunk size, …) of the payload, `size` its length. */
+                                                                                                                             
+                                                        
+/** What a StreamBuffer keeps until the peer acknowledges it. */
+                                                 
+
+/** The payload cut into chunks: SEGMENT_BYTES each, or larger only to stay within SEGMENT_COUNT_MAX. */
+export function chunkParts(payload        , z         , segmentBytes = SEGMENT_BYTES)              {
+  const size = payload.length;
+  const chunkSize = Math.max(segmentBytes, Math.ceil(size / SEGMENT_COUNT_MAX));
+  const segments = Math.max(1, Math.ceil(size / chunkSize));
+  return Array.from({ length: segments }, (_, segment) => ({
+    segment, segments, size, ...(z ? { z: 1          } : {}), data: payload.subarray(segment * chunkSize, (segment + 1) * chunkSize).toString("base64"),
+  }));
+}
+
+const deflateRaw = promisify(zlib.deflateRaw);
+
+/** A line as a `seg` peer takes it: as it is up to SEGMENT_BYTES, otherwise deflated and in chunks. */
+export async function encodeLine(line        )                                  {
+  const bytes = Buffer.from(line, "utf8");
+  if (bytes.length <= SEGMENT_BYTES) return [line];
+  return chunkParts(await deflateRaw(bytes), true);
+}
+
+/** encodeLine without yielding, for the test phone. */
+export function encodeLineSync(line        )                         {
+  const bytes = Buffer.from(line, "utf8");
+  if (bytes.length <= SEGMENT_BYTES) return [line];
+  return chunkParts(zlib.deflateRawSync(bytes), true);
+}
+
+/**
+ * Puts one stream's chunks back together. Chunks arrive in seq order, so a line's segments are consecutive; anything
+ * else (a skipped segment, a changed count or size, more bytes than announced, a `msg` in the middle) is a protocol
+ * error, since the peer is authenticated and only a bug would send it.
+ */
+export class ChunkAssembler {
+          parts           = [];
+          received = 0;
+          first                   = null;
+
+  get busy() { return this.first !== null; }
+
+  take(message           )                {
+    const { segment, segments, size } = message;
+    if (!Number.isInteger(segments) || segments < 1 || segments > SEGMENT_COUNT_MAX || !Number.isInteger(size) || size < 1 || size > REASSEMBLED_MAX_BYTES) {
+      throw new Error("chunk count or size out of range");
+    }
+    const first = this.first ?? message;
+    if (segment !== this.parts.length || first.segments !== segments || first.size !== size || first.z !== message.z) throw new Error(`chunk ${segment} out of order`);
+    const data = Buffer.from(message.data, "base64");
+    this.received += data.length;
+    if (data.length === 0 || this.received > size) throw new Error("chunk carries more than its size");
+    this.first = first;
+    this.parts.push(data);
+    if (segment < segments - 1) return null;
+    if (this.received !== size) throw new Error("chunks carry less than their size");
+    const payload = Buffer.concat(this.parts);
+    this.parts = []; this.received = 0; this.first = null;
+    return (message.z ? zlib.inflateRawSync(payload, { maxOutputLength: REASSEMBLED_MAX_BYTES }) : payload).toString("utf8");
+  }
+}
+
 /** Relay ↔ host multiplexing: one host WebSocket carries every phone connection. */
                           
                                                   
@@ -236,26 +315,28 @@ export const CloseCode = {
 export class StreamBuffer {
   nextSeq = 1;
   received = 0;
-          unacked                                                = [];
+          unacked                                        = [];
           bytes = 0;
 
                    maxBytes        ;
   constructor(maxBytes = 32 * 1024 * 1024) { this.maxBytes = maxBytes; }
 
-  push(line        )                                {
-    const item = { seq: this.nextSeq++, line, size: Buffer.byteLength(line) };
+  push(part                    )           {
+    const seq = this.nextSeq++;
+    const message           = typeof part === "string" ? { t: "msg", seq, line: part } : { t: "chunk", seq, ...part };
+    const item = { message, size: Buffer.byteLength(message.t === "msg" ? message.line : message.data) };
     this.unacked.push(item);
     this.bytes += item.size;
     if (this.bytes > this.maxBytes) throw new Error("unacknowledged buffer full");
-    return item;
+    return message;
   }
 
   ack(seq        ) {
-    while (this.unacked.length && this.unacked[0].seq <= seq) this.bytes -= this.unacked.shift() .size;
+    while (this.unacked.length && this.unacked[0].message.seq <= seq) this.bytes -= this.unacked.shift() .size;
   }
 
-  pending(after        )                                  {
-    return this.unacked.filter((item) => item.seq > after);
+  pending(after        )             {
+    return this.unacked.filter((item) => item.message.seq > after).map((item) => item.message);
   }
 
   /** True when the line is new; duplicates from a resend are dropped. Gaps cannot happen on one ordered socket. */