[Miniflare 3] Allow custom services to respond with Content-Encoding and multiple Set-Cookie headers

packages/miniflare/package.json

@@ -35,6 +35,7 @@
     "glob-to-regexp": "^0.4.1",
     "http-cache-semantics": "^4.1.0",
     "kleur": "^4.1.5",
+    "set-cookie-parser": "^2.6.0",
     "source-map-support": "0.5.21",
     "stoppable": "^1.1.0",
     "undici": "^5.13.0",
@@ -50,6 +51,7 @@
     "@types/estree": "^1.0.0",
     "@types/glob-to-regexp": "^0.4.1",
     "@types/http-cache-semantics": "^4.0.1",
+    "@types/set-cookie-parser": "^2.4.2",
     "@types/source-map-support": "^0.5.6",
     "@types/stoppable": "^1.1.1",
     "@types/ws": "^8.5.3",

packages/miniflare/src/index.ts

@@ -5,10 +5,12 @@ 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 { splitCookiesString } from "set-cookie-parser";
 import stoppable from "stoppable";
 import { WebSocketServer } from "ws";
 import { z } from "zod";
@@ -277,15 +279,78 @@ 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);
+  // Convert headers into Node-friendly format
+  const headers: http.OutgoingHttpHeaders = {};
+  for (const entry of response.headers) {
+    const key = entry[0].toLowerCase();
+    const value = entry[1];
+    if (key === "set-cookie") {
+      headers[key] = splitCookiesString(value);
+    } else {
+      headers[key] = value;
+    }
+  }
+
+  // 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,82 @@ 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: custom service using Set-Cookie header", async (t) => {
+  const testCookies = [
+    "key1=value1; Max-Age=3600",
+    "key2=value2; Domain=example.com; Secure",
+  ];
+  const { http } = await useServer(t, (req, res) => {
+    res.writeHead(200, { "Set-Cookie": testCookies });
+    res.end();
+  });
+  const mf = new Miniflare({
+    modules: true,
+    script: `export default {
+      async fetch(request, env, ctx) {
+        const res = await env.CUSTOM.fetch(request);
+        return Response.json(res.headers.getSetCookie());
+      }
+    }`,
+    serviceBindings: {
+      CUSTOM(request) {
+        return fetch(http, request);
+      },
+    },
+    // Enable `Headers#getSetCookie()`:
+    // https://github.com/cloudflare/workerd/blob/14b54764609c263ea36ab862bb8bf512f9b1387b/src/workerd/io/compatibility-date.capnp#L273-L278
+    compatibilityDate: "2023-03-01",
+  });
+  t.teardown(() => mf.dispose());
+
+  const res = await mf.dispatchFetch("http://localhost");
+  t.deepEqual(await res.json(), testCookies);
+});
+
 test("Miniflare: web socket kitchen sink", async (t) => {
   // Create deferred promises for asserting asynchronous event results
   const clientEventPromise = new DeferredPromise<MessageEvent>();