From 3f0890e3d372b354d28a50b3340b3bd054db3033 Mon Sep 17 00:00:00 2001
From: Roy Razon <roy@livecycle.io>
Date: Tue, 19 Sep 2023 09:40:42 +0300
Subject: [PATCH] tunnel server: fix encoding

when injecting, responses were not compressed.
also fixed wrong chunked response.
---
 .../src/proxy/html-manipulation/index.ts      | 40 +++++++++++++------
 .../html-manipulation/inject-transform.ts     |  2 +-
 2 files changed, 28 insertions(+), 14 deletions(-)

diff --git a/tunnel-server/src/proxy/html-manipulation/index.ts b/tunnel-server/src/proxy/html-manipulation/index.ts
index 6409a44d..584038b3 100644
--- a/tunnel-server/src/proxy/html-manipulation/index.ts
+++ b/tunnel-server/src/proxy/html-manipulation/index.ts
@@ -3,24 +3,40 @@ import zlib from 'zlib'
 import stream from 'stream'
 import { parse as parseContentType } from 'content-type'
 import iconv from 'iconv-lite'
+import { inspect } from 'node:util'
 import { INJECT_SCRIPTS_HEADER } from '../common'
 import { InjectHtmlScriptTransform } from './inject-transform'
 import { ScriptInjection } from '../../tunnel-store'
 
-const compressionsForContentEncoding = (contentEncoding: string) => {
+const compressionsForContentEncoding = (
+  contentEncoding: string | undefined,
+): [stream.Transform, stream.Transform] | undefined => {
+  if (!contentEncoding || contentEncoding === 'identity') {
+    return undefined
+  }
   if (contentEncoding === 'gzip') {
-    return [zlib.createGunzip(), zlib.createGzip()] as const
+    return [zlib.createGunzip(), zlib.createGzip()]
   }
   if (contentEncoding === 'deflate') {
-    return [zlib.createInflate(), zlib.createDeflate()] as const
+    return [zlib.createInflate(), zlib.createDeflate()]
   }
   if (contentEncoding === 'br') {
-    return [zlib.createBrotliDecompress(), zlib.createBrotliCompress()] as const
+    return [zlib.createBrotliDecompress(), zlib.createBrotliCompress()]
   }
-  if (contentEncoding === 'identity') {
-    return undefined
+  throw new Error(`unsupported content encoding: ${inspect(contentEncoding)}`)
+}
+
+const streamsForContentEncoding = (
+  contentEncoding: string | undefined,
+  input: stream.Readable,
+  output: stream.Writable,
+): [stream.Readable, stream.Writable] => {
+  const compress = compressionsForContentEncoding(contentEncoding)
+  if (!compress) {
+    return [input, output]
   }
-  throw new Error(`unsupported content encoding: "${contentEncoding}"`)
+  compress[1].pipe(output)
+  return [input.pipe(compress[0]), compress[1]]
 }
 
 export const injectScripts = (
@@ -28,12 +44,11 @@ export const injectScripts = (
   req: Pick<IncomingMessage, 'headers'>,
   res: stream.Writable & Pick<ServerResponse<IncomingMessage>, 'writeHead'>,
 ) => {
-  res.writeHead(proxyRes.statusCode as number, proxyRes.headers)
-
   const injectsStr = req.headers[INJECT_SCRIPTS_HEADER] as string | undefined
   const contentTypeHeader = proxyRes.headers['content-type']
 
   if (!injectsStr || !contentTypeHeader) {
+    res.writeHead(proxyRes.statusCode as number, proxyRes.headers)
     proxyRes.pipe(res)
     return undefined
   }
@@ -44,15 +59,14 @@ export const injectScripts = (
   } = parseContentType(contentTypeHeader)
 
   if (contentType !== 'text/html') {
+    res.writeHead(proxyRes.statusCode as number, proxyRes.headers)
     proxyRes.pipe(res)
     return undefined
   }
 
-  const compress = compressionsForContentEncoding(proxyRes.headers['content-encoding'] || 'identity')
+  res.writeHead(proxyRes.statusCode as number, { ...proxyRes.headers, 'transfer-encoding': '' })
 
-  const [input, output] = compress
-    ? [proxyRes.pipe(compress[0]), res.pipe(compress[1])]
-    : [proxyRes, res]
+  const [input, output] = streamsForContentEncoding(proxyRes.headers['content-encoding'], proxyRes, res)
 
   const injects = JSON.parse(injectsStr) as Omit<ScriptInjection, 'pathRegex'>[]
   const transform = new InjectHtmlScriptTransform(injects)
diff --git a/tunnel-server/src/proxy/html-manipulation/inject-transform.ts b/tunnel-server/src/proxy/html-manipulation/inject-transform.ts
index e12af5ff..3941d1a7 100644
--- a/tunnel-server/src/proxy/html-manipulation/inject-transform.ts
+++ b/tunnel-server/src/proxy/html-manipulation/inject-transform.ts
@@ -85,7 +85,7 @@ export class InjectHtmlScriptTransform extends stream.Transform {
   override _transform(chunk: string, _encoding: BufferEncoding | 'buffer', callback: stream.TransformCallback): void {
     if (typeof chunk !== 'string') {
       // chunk must be string rather than Buffer so htmlDetector offsets would be in character units, not bytes
-      throw new Error(`Invalid chunk, expected string, received ${typeof chunk}: ${chunk}`)
+      throw new Error(`Invalid chunk, expected string, received ${Buffer.isBuffer(chunk) ? 'Buffer' : typeof chunk}: ${chunk}`)
     }
 
     if (this.injected) {