From 6bd90b7ce3520d06bda16304e2732de4402c1279 Mon Sep 17 00:00:00 2001
From: Roy Razon <roy@livecycle.io>
Date: Mon, 18 Sep 2023 13:45:25 +0300
Subject: [PATCH] fix: tunnel server injection

- non-html docs were not streaming
- responses with no content-type header were throwing
---
 tunnel-server/src/proxy/html-manipulation/index.ts | 12 ++++++++----
 tunnel-server/src/proxy/index.ts                   |  5 ++---
 2 files changed, 10 insertions(+), 7 deletions(-)

diff --git a/tunnel-server/src/proxy/html-manipulation/index.ts b/tunnel-server/src/proxy/html-manipulation/index.ts
index f1881373..83d7c178 100644
--- a/tunnel-server/src/proxy/html-manipulation/index.ts
+++ b/tunnel-server/src/proxy/html-manipulation/index.ts
@@ -28,22 +28,26 @@ export const injectScripts = async (
   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
-  if (!injectsStr) {
+  const contentTypeHeader = proxyRes.headers['content-type']
+
+  if (!injectsStr || !contentTypeHeader) {
+    proxyRes.pipe(res)
     return undefined
   }
 
   const {
     type: contentType,
     parameters: { charset: reqCharset },
-  } = parseContentType(proxyRes)
+  } = parseContentType(contentTypeHeader)
 
   if (contentType !== 'text/html') {
+    proxyRes.pipe(res)
     return undefined
   }
 
-  res.writeHead(proxyRes.statusCode as number, proxyRes.headers)
-
   const compress = compressionsForContentEncoding(proxyRes.headers['content-encoding'] || 'identity')
 
   const [input, output] = compress
diff --git a/tunnel-server/src/proxy/index.ts b/tunnel-server/src/proxy/index.ts
index 2d131e6b..1fdf1b92 100644
--- a/tunnel-server/src/proxy/index.ts
+++ b/tunnel-server/src/proxy/index.ts
@@ -139,8 +139,7 @@ export const proxy = ({
       ?.filter(({ pathRegex }) => !pathRegex || pathRegex.test(mutatedReq.url || ''))
       ?.map(({ src, defer, async }) => ({ src, defer, async }))
 
-    const shouldInject = Boolean(injects?.length)
-    if (shouldInject) {
+    if (injects?.length) {
       mutatedReq.headers[INJECT_SCRIPTS_HEADER] = JSON.stringify(injects)
     }
 
@@ -153,7 +152,7 @@ export const proxy = ({
         target: {
           socketPath: activeTunnel.target,
         },
-        selfHandleResponse: shouldInject,
+        selfHandleResponse: true, // handled by the injectScripts onProxyRes hook
       },
       err => errorHandler(log, err, req, res)
     )