diff --git a/.changeset/tunnel-response-gzip.md b/.changeset/tunnel-response-gzip.md new file mode 100644 index 00000000000..593c163d637 --- /dev/null +++ b/.changeset/tunnel-response-gzip.md @@ -0,0 +1,5 @@ +--- +"@moonshot-ai/kimi-code": patch +--- + +Compress Remote Control tunnel responses with gzip. diff --git a/packages/remote-control/src/remote-control.ts b/packages/remote-control/src/remote-control.ts index 88a9a3b1d19..1eaef02e205 100644 --- a/packages/remote-control/src/remote-control.ts +++ b/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 { @@ -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 => { + 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); }); }, ); diff --git a/packages/remote-control/test/remote-control.test.ts b/packages/remote-control/test/remote-control.test.ts index a3a1d66ac49..e1da0659088 100644 --- a/packages/remote-control/test/remote-control.test.ts +++ b/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 = `${''.repeat(100)}`; + 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',