Allow custom services to respond with Content-Encoding header

Previously, if the response from a custom service included a
`Content-Encoding` header, we would respond with decoded content to
`workerd` without removing the `Content-Encoding` header. This caused
`workerd` to try to decode already decoded content, resulting in an
error. This was particularly a problem with `wrangler pages dev`, as
upstream dev servers could very easily return gzipped content.

This change updates the `writeResponse()` function with code from
Miniflare 2, to re-encode the body before responding.

Note, due to a fixed but unreleased bug in `undici`
(nodejs/undici#2159), defining a custom service that directly proxies
somewhere else with `fetch()` like `wrangler pages dev` will fail, if
`Content-Encoding` contains multiple encodings.

Closes cloudflare/workers-sdk#3408

packages/miniflare/src/index.ts

@@ -5,8 +5,9 @@ import http from "http";
 import net from "net";
 import os from "os";
 import path from "path";
-import { Duplex } from "stream";
+import { Duplex, Transform, Writable } from "stream";
 import { ReadableStream } from "stream/web";
+import zlib from "zlib";
 import type { RequestInitCfProperties } from "@cloudflare/workers-types/experimental";
 import exitHook from "exit-hook";
 import stoppable from "stoppable";
@@ -277,15 +278,68 @@ const restrictedWebSocketUpgradeHeaders = [
   "sec-websocket-accept",
 ];
 
+export function _transformsForContentEncoding(encoding?: string): Transform[] {
+  const encoders: Transform[] = [];
+  if (!encoding) return encoders;
+
+  // Reverse of https://github.com/nodejs/undici/blob/48d9578f431cbbd6e74f77455ba92184f57096cf/lib/fetch/index.js#L1660
+  const codings = encoding
+    .toLowerCase()
+    .split(",")
+    .map((x) => x.trim());
+  for (const coding of codings) {
+    if (/(x-)?gzip/.test(coding)) {
+      encoders.push(zlib.createGzip());
+    } else if (/(x-)?deflate/.test(coding)) {
+      encoders.push(zlib.createDeflate());
+    } else if (coding === "br") {
+      encoders.push(zlib.createBrotliCompress());
+    } else {
+      // Unknown encoding, don't do any encoding at all
+      encoders.length = 0;
+      break;
+    }
+  }
+  return encoders;
+}
+
 async function writeResponse(response: Response, res: http.ServerResponse) {
   const headers = Object.fromEntries(response.headers);
+
+  // If a `Content-Encoding` header is set, we'll need to encode the body
+  // (likely only set by custom service bindings)
+  const encoding = headers["content-encoding"]?.toString();
+  const encoders = _transformsForContentEncoding(encoding);
+  if (encoders.length > 0) {
+    // `Content-Length` if set, will be wrong as it's for the decoded length
+    delete headers["content-length"];
+  }
+
   res.writeHead(response.status, response.statusText, headers);
+
+  // `initialStream` is the stream we'll write the response to. It
+  // should end up as the first encoder, piping to the next encoder,
+  // and finally piping to the response:
+  //
+  // encoders[0] (initialStream) -> encoders[1] -> res
+  //
+  // Not using `pipeline(passThrough, ...encoders, res)` here as that
+  // gives a premature close error with server sent events. This also
+  // avoids creating an extra stream even when we're not encoding.
+  let initialStream: Writable = res;
+  for (let i = encoders.length - 1; i >= 0; i--) {
+    encoders[i].pipe(initialStream);
+    initialStream = encoders[i];
+  }
+
+  // Response body may be null if empty
   if (response.body) {
     for await (const chunk of response.body) {
-      if (chunk) res.write(chunk);
+      if (chunk) initialStream.write(chunk);
     }
   }
-  res.end();
+
+  initialStream.end();
 }
 
 function safeReadableStreamFrom(iterable: AsyncIterable<Uint8Array>) {

packages/miniflare/test/index.spec.ts

@@ -1,13 +1,15 @@
 import assert from "assert";
 import http from "http";
 import { AddressInfo } from "net";
+import { Writable } from "stream";
 import test from "ava";
 import {
   DeferredPromise,
   MessageEvent,
   Miniflare,
   MiniflareCoreError,
   MiniflareOptions,
+  _transformsForContentEncoding,
   fetch,
 } from "miniflare";
 import {
@@ -106,6 +108,50 @@ test("Miniflare: routes to multiple workers with fallback", async (t) => {
   t.is(await res.text(), "a");
 });
 
+test("Miniflare: custom service using Content-Encoding header", async (t) => {
+  const testBody = "x".repeat(100);
+  const { http } = await useServer(t, (req, res) => {
+    const testEncoding = req.headers["x-test-encoding"]?.toString();
+    const encoders = _transformsForContentEncoding(testEncoding);
+    let initialStream: Writable = res;
+    for (let i = encoders.length - 1; i >= 0; i--) {
+      encoders[i].pipe(initialStream);
+      initialStream = encoders[i];
+    }
+    res.writeHead(200, { "Content-Encoding": testEncoding });
+    initialStream.write(testBody);
+    initialStream.end();
+  });
+  const mf = new Miniflare({
+    script: `addEventListener("fetch", (event) => {
+      event.respondWith(CUSTOM.fetch(event.request));
+    })`,
+    serviceBindings: {
+      CUSTOM(request) {
+        return fetch(http, request);
+      },
+    },
+  });
+  t.teardown(() => mf.dispose());
+
+  const test = async (encoding: string) => {
+    const res = await mf.dispatchFetch("http://localhost", {
+      headers: { "X-Test-Encoding": encoding },
+    });
+    t.is(res.headers.get("Content-Encoding"), encoding);
+    t.is(await res.text(), testBody, encoding);
+  };
+
+  await test("gzip");
+  await test("deflate");
+  await test("br");
+  // `undici`'s `fetch()` is currently broken when `Content-Encoding` specifies
+  // multiple encodings. Once https://github.com/nodejs/undici/pull/2159 is
+  // released, we can re-enable this test.
+  // TODO(soon): re-enable this test
+  // await test("deflate, gzip");
+});
+
 test("Miniflare: web socket kitchen sink", async (t) => {
   // Create deferred promises for asserting asynchronous event results
   const clientEventPromise = new DeferredPromise<MessageEvent>();