feat(remote-control): compress textual tunnel responses with gzip

.changeset/tunnel-response-gzip.md

@@ -0,0 +1,5 @@
+---
+"@moonshot-ai/kimi-code": patch
+---
+
+Compress Remote Control tunnel responses with gzip.

packages/remote-control/src/remote-control.ts

@@ -2,6 +2,8 @@ import { hostname, platform } from 'node:os';
 import { join } from 'node:path';
 import { request as httpRequest, validateHeaderName, validateHeaderValue } from 'node:http';
 import { setTimeout as sleep } from 'node:timers/promises';
+import { promisify } from 'node:util';
+import { gzip } from 'node:zlib';
 
 import {
   createKimiDeviceId,
@@ -59,6 +61,14 @@ const BLOCKED_RESPONSE_HEADERS = new Set([
   'transfer-encoding',
   'upgrade',
 ]);
+const GZIP_MIN_BODY_BYTES = 1024;
+const GZIP_COMPRESSIBLE_TYPES = new Set([
+  'application/javascript',
+  'application/json',
+  'application/xml',
+  'image/svg+xml',
+]);
+const gzipAsync = promisify(gzip);
 
 interface RelayMessage {
   readonly type: string;
@@ -217,6 +227,27 @@ export function rewriteRemoteControlResponse(
   return body;
 }
 
+function acceptsGzipEncoding(headers: readonly [string, string][]): boolean {
+  let wildcard = false;
+  for (const [name, value] of headers) {
+    if (name.toLowerCase() !== 'accept-encoding') continue;
+    for (const token of value.split(',')) {
+      const [encoding, ...params] = token.trim().toLowerCase().split(';');
+      if (encoding !== 'gzip' && encoding !== '*') continue;
+      const quality = params.map((param) => param.trim()).find((param) => param.startsWith('q='));
+      const acceptable = quality === undefined || Number(quality.slice(2)) > 0;
+      if (encoding === 'gzip') return acceptable;
+      wildcard = wildcard || acceptable;
+    }
+  }
+  return wildcard;
+}
+
+function isGzipCompressibleType(contentType: string): boolean {
+  const mime = contentType.split(';', 1)[0]!.trim().toLowerCase();
+  return mime.startsWith('text/') || GZIP_COMPRESSIBLE_TYPES.has(mime);
+}
+
 export async function startRemoteControl(
   options: RemoteControlOptions,
 ): Promise<RemoteControlHandle> {
@@ -816,24 +847,48 @@ function requestLocalHttp(
         response.on('data', (chunk: Buffer | string) => chunks.push(Buffer.from(chunk)));
         response.once('error', reject);
         response.once('end', () => {
-          const contentType = response.headers['content-type'] ?? '';
-          const receivedBody = Buffer.concat(chunks);
-          const body =
-            response.headers['content-encoding'] === undefined
-              ? rewriteRemoteControlResponse(contentType, receivedBody, publicPrefix)
-              : receivedBody;
-          const rewritten = body !== receivedBody;
-          const headers = filterResponseHeaders(response.rawHeaders, rewritten);
-          if (rewritten) headers.push('Cache-Control', 'no-cache');
-          headers.push('Content-Length', String(body.length));
-          const statusCode = response.statusCode ?? 502;
-          const statusMessage = response.statusMessage ?? 'Bad Gateway';
-          resolve(
-            Buffer.concat([
+          void (async (): Promise<Buffer> => {
+            const contentType = response.headers['content-type'] ?? '';
+            const receivedBody = Buffer.concat(chunks);
+            let body =
+              response.headers['content-encoding'] === undefined
+                ? rewriteRemoteControlResponse(contentType, receivedBody, publicPrefix)
+                : receivedBody;
+            const rewritten = body !== receivedBody;
+            const headers = filterResponseHeaders(response.rawHeaders, rewritten);
+            if (rewritten) headers.push('Cache-Control', 'no-cache');
+            const negotiated =
+              response.headers['content-encoding'] === undefined &&
+              response.statusCode !== 206 &&
+              body.length >= GZIP_MIN_BODY_BYTES &&
+              isGzipCompressibleType(contentType);
+            if (negotiated) {
+              let varyCovers = false;
+              for (let index = 0; index < headers.length; index += 2) {
+                if (headers[index]!.toLowerCase() !== 'vary') continue;
+                const tokens = headers[index + 1]!
+                  .toLowerCase()
+                  .split(',')
+                  .map((token) => token.trim());
+                if (tokens.includes('*') || tokens.includes('accept-encoding')) varyCovers = true;
+              }
+              if (!varyCovers) headers.push('Vary', 'Accept-Encoding');
+            }
+            if (negotiated && acceptsGzipEncoding(parsed.headers)) {
+              body = await gzipAsync(body);
+              headers.push('Content-Encoding', 'gzip');
+              for (let index = headers.length - 2; index >= 0; index -= 2) {
+                if (headers[index]!.toLowerCase() === 'etag') headers.splice(index, 2);
+              }
+            }
+            headers.push('Content-Length', String(body.length));
+            const statusCode = response.statusCode ?? 502;
+            const statusMessage = response.statusMessage ?? 'Bad Gateway';
+            return Buffer.concat([
               Buffer.from(`HTTP/1.1 ${statusCode} ${statusMessage}\r\n${headerLines(headers)}\r\n\r\n`),
               body,
-            ]),
-          );
+            ]);
+          })().then(resolve, reject);
         });
       },
     );

packages/remote-control/test/remote-control.test.ts

@@ -4,6 +4,7 @@ import { existsSync, mkdtempSync, rmSync } from 'node:fs';
 import { mkdir, readFile, writeFile } from 'node:fs/promises';
 import { tmpdir } from 'node:os';
 import { join } from 'node:path';
+import { gunzipSync } from 'node:zlib';
 
 import {
   FileTokenStorage,
@@ -274,8 +275,38 @@ describe('Remote Control tunnel', () => {
     let localHttpRequest: IncomingMessage | undefined;
     let localWsRequest: IncomingMessage | undefined;
     const localWsServer = new WebSocketServer({ noServer: true });
+    const assetJs = `const boot = "/assets/boot.js";\n${'const chunk = "/assets/chunk.js";\n'.repeat(120)}`;
+    const assetPng = Buffer.alloc(4096, 7);
+    const assetSvg = `<svg xmlns="http://www.w3.org/2000/svg">${'<rect width="100" height="100"/>'.repeat(100)}</svg>`;
+    const assetText = 'chunk of text\n'.repeat(160);
     const localServer = createServer((request, response) => {
       localHttpRequest = request;
+      if (request.url === '/assets/index.js') {
+        response.writeHead(200, { 'Content-Type': 'text/javascript', ETag: '"v1"' });
+        response.end(assetJs);
+        return;
+      }
+      if (request.url === '/assets/logo.png') {
+        response.writeHead(200, { 'Content-Type': 'image/png' });
+        response.end(assetPng);
+        return;
+      }
+      if (request.url === '/assets/logo.svg') {
+        response.writeHead(200, {
+          'Content-Type': 'image/svg+xml',
+          'Cache-Control': 'public, max-age=31536000, immutable',
+        });
+        response.end(assetSvg);
+        return;
+      }
+      if (request.url === '/assets/partial.txt' && request.headers.range !== undefined) {
+        response.writeHead(206, {
+          'Content-Type': 'text/plain',
+          'Content-Range': 'bytes 0-2047/4096',
+        });
+        response.end(assetText);
+        return;
+      }
       response.writeHead(200, {
         'Content-Type': 'text/html',
         'Cache-Control': 'public, max-age=31536000, immutable',
@@ -352,7 +383,7 @@ describe('Remote Control tunnel', () => {
     expect(handle.url).toContain('?rc=1&from=kimi_code_cli');
 
     const rawRequest = Buffer.from(
-      'GET / HTTP/1.1\r\nHost: relay.test\r\nAuthorization: Bearer relay-token\r\nCookie: sid=1\r\nOrigin: https://relay.test\r\nConnection: X-Hop\r\nX-Hop: remove\r\nX-Keep: yes\r\n\r\n',
+      'GET / HTTP/1.1\r\nHost: relay.test\r\nAuthorization: Bearer relay-token\r\nCookie: sid=1\r\nOrigin: https://relay.test\r\nAccept-Encoding: gzip\r\nConnection: X-Hop\r\nX-Hop: remove\r\nX-Keep: yes\r\n\r\n',
     );
     const splitAt = Math.floor(rawRequest.length / 2);
     httpConnections[0]!.send(
@@ -384,6 +415,9 @@ describe('Remote Control tunnel', () => {
     expect(localHttpRequest?.headers['x-keep']).toBe('yes');
     expect(response).not.toContain('X-Remove');
     expect(response).not.toContain('immutable');
+    expect(response).not.toContain('Content-Encoding');
+    expect(response).not.toContain('Vary');
+    expect(localHttpRequest?.headers['accept-encoding']).toBeUndefined();
     expect(response).toContain('Cache-Control: no-cache');
     expect(response).toContain(`/coding-relay/devices/${handle.deviceId}/boot.js`);
 
@@ -400,6 +434,119 @@ describe('Remote Control tunnel', () => {
     await rotatedResponsePromise;
     await waitFor(() => localHttpRequest?.headers.authorization === 'Bearer rotated-server-token');
 
+    const gzipResponsePromise = nextJsonMessage(httpConnections[0]!);
+    httpConnections[0]!.send(
+      JSON.stringify({
+        request_id: 'request-3',
+        type: 'request',
+        is_last: true,
+        body_base64: Buffer.from(
+          'GET /assets/index.js HTTP/1.1\r\nHost: relay.test\r\nAccept-Encoding: br, gzip\r\n\r\n',
+        ).toString('base64'),
+      }),
+    );
+    const gzipResponse = Buffer.from(
+      (await gzipResponsePromise)['body_base64'] as string,
+      'base64',
+    );
+    const gzipSeparator = gzipResponse.indexOf('\r\n\r\n');
+    const gzipHead = gzipResponse.subarray(0, gzipSeparator).toString('latin1');
+    const gzipBody = gzipResponse.subarray(gzipSeparator + 4);
+    expect(gzipHead).toContain('HTTP/1.1 200 OK');
+    expect(gzipHead).toContain('Content-Encoding: gzip');
+    expect(gzipHead).toContain('Vary: Accept-Encoding');
+    expect(gzipHead).not.toContain('ETag');
+    expect(gzipHead).toContain(`Content-Length: ${gzipBody.length}`);
+    expect(gunzipSync(gzipBody).toString()).toBe(
+      assetJs.replaceAll('"/assets/', `"/coding-relay/devices/${handle.deviceId}/assets/`),
+    );
+
+    const binaryResponsePromise = nextJsonMessage(httpConnections[0]!);
+    httpConnections[0]!.send(
+      JSON.stringify({
+        request_id: 'request-4',
+        type: 'request',
+        is_last: true,
+        body_base64: Buffer.from(
+          'GET /assets/logo.png HTTP/1.1\r\nHost: relay.test\r\nAccept-Encoding: gzip\r\n\r\n',
+        ).toString('base64'),
+      }),
+    );
+    const binaryResponse = Buffer.from(
+      (await binaryResponsePromise)['body_base64'] as string,
+      'base64',
+    );
+    const binarySeparator = binaryResponse.indexOf('\r\n\r\n');
+    expect(binaryResponse.subarray(0, binarySeparator).toString('latin1')).not.toContain(
+      'Content-Encoding',
+    );
+    expect(binaryResponse.subarray(binarySeparator + 4).equals(assetPng)).toBe(true);
+
+    const excludedResponsePromise = nextJsonMessage(httpConnections[0]!);
+    httpConnections[0]!.send(
+      JSON.stringify({
+        request_id: 'request-5',
+        type: 'request',
+        is_last: true,
+        body_base64: Buffer.from(
+          'GET /assets/index.js HTTP/1.1\r\nHost: relay.test\r\nAccept-Encoding: gzip;q=0, *;q=1\r\n\r\n',
+        ).toString('base64'),
+      }),
+    );
+    const excludedResponse = Buffer.from(
+      (await excludedResponsePromise)['body_base64'] as string,
+      'base64',
+    );
+    const excludedSeparator = excludedResponse.indexOf('\r\n\r\n');
+    const excludedHead = excludedResponse.subarray(0, excludedSeparator).toString('latin1');
+    expect(excludedHead).not.toContain('Content-Encoding');
+    expect(excludedHead).toContain('Vary: Accept-Encoding');
+    expect(excludedHead).toContain('ETag: "v1"');
+    expect(excludedResponse.subarray(excludedSeparator + 4).toString()).toBe(
+      assetJs.replaceAll('"/assets/', `"/coding-relay/devices/${handle.deviceId}/assets/`),
+    );
+
+    const svgResponsePromise = nextJsonMessage(httpConnections[0]!);
+    httpConnections[0]!.send(
+      JSON.stringify({
+        request_id: 'request-6',
+        type: 'request',
+        is_last: true,
+        body_base64: Buffer.from(
+          'GET /assets/logo.svg HTTP/1.1\r\nHost: relay.test\r\nAccept-Encoding: gzip\r\n\r\n',
+        ).toString('base64'),
+      }),
+    );
+    const svgResponse = Buffer.from((await svgResponsePromise)['body_base64'] as string, 'base64');
+    const svgSeparator = svgResponse.indexOf('\r\n\r\n');
+    const svgHead = svgResponse.subarray(0, svgSeparator).toString('latin1');
+    expect(svgHead).toContain('Content-Encoding: gzip');
+    expect(svgHead).toContain('Vary: Accept-Encoding');
+    expect(svgHead).toContain('immutable');
+    expect(gunzipSync(svgResponse.subarray(svgSeparator + 4)).toString()).toBe(assetSvg);
+
+    const rangeResponsePromise = nextJsonMessage(httpConnections[0]!);
+    httpConnections[0]!.send(
+      JSON.stringify({
+        request_id: 'request-7',
+        type: 'request',
+        is_last: true,
+        body_base64: Buffer.from(
+          'GET /assets/partial.txt HTTP/1.1\r\nHost: relay.test\r\nAccept-Encoding: gzip\r\nRange: bytes=0-2047\r\n\r\n',
+        ).toString('base64'),
+      }),
+    );
+    const rangeResponse = Buffer.from(
+      (await rangeResponsePromise)['body_base64'] as string,
+      'base64',
+    );
+    const rangeSeparator = rangeResponse.indexOf('\r\n\r\n');
+    const rangeHead = rangeResponse.subarray(0, rangeSeparator).toString('latin1');
+    expect(rangeHead).toContain('206');
+    expect(rangeHead).toContain('Content-Range: bytes 0-2047/4096');
+    expect(rangeHead).not.toContain('Content-Encoding');
+    expect(rangeResponse.subarray(rangeSeparator + 4).toString()).toBe(assetText);
+
     managementConnections[0]!.send(
       JSON.stringify({
         type: 'open_ws',