Non-blocking mechanism to transform request/response body streams #192

Activity

1. avarayr commented on May 13, 2025

   

   avarayr

   on May 13, 2025

   Last edited by avarayr

   More actions

   Thank you for this @pimterry and thank you for this amazing well-tested library.

   I ended up forking and hacking up a solution that fit my usecase, here's how I did it.

   request-handlers.ts

   // definitions
   export interface PassThroughResponseStream extends Omit<PassThroughResponse, "body" | "rawHeaders"> {
     body: Readable;
   }
   export interface CompletedRequest extends Request {
       body: CompletedBody;
       rawTrailers: RawTrailers;
       trailers: Trailers;
   }
   public readonly onStreamingResponse?: (
       res: PassThroughResponseStream,
       req: CompletedRequest,
   ) => MaybePromise<CallbackResponseStreamResult | void> | void;
   
   
   // request-handlers.ts
   
   /**
    * Check if we need to transform the response
    */
   if (this.onStreamingResponse) {
     let resTransformStream: CallbackResponseStreamResult | void;
     let serverHeaders = rawHeadersToObject(serverRawHeaders);
   
     let reqHeader = rawHeadersToObjectPreservingCase(rawHeaders);
     resTransformStream = await this.onStreamingResponse(
       {
         id: clientReq.id,
         statusCode: serverStatusCode,
         statusMessage: serverRes.statusMessage,
         headers: serverHeaders,
         body: decodeBodyStream(serverRes, serverHeaders), // decode the body for convenience
       },
       {
         id: clientReq.id,
         protocol: protocol?.replace(":", "") ?? "",
         method: method,
         httpVersion: serverRes.httpVersion,
         url: clientReq.url, // original, unmodified URL
         modifiedReqUrl: clientReq.url !== reqUrl ? reqUrl : undefined, // modified URL
         destination: {
           hostname: hostname || "localhost",
           port: effectivePort,
         },
         path: path ?? "",
         headers: reqHeader,
         rawHeaders: rawHeaders,
         timingEvents: clientReq.timingEvents,
         tags: clientReq.tags,
         body: buildBodyReader(
           reqBodyOverride ? Buffer.from(reqBodyOverride.buffer) : await clientReq.body.asDecodedBuffer(),
           reqHeader,
         ),
         rawTrailers: clientReq.rawTrailers ?? [],
         trailers: rawHeadersToObject(clientReq.rawTrailers ?? []),
       },
     );
   
     if (!resTransformStream) {
       // what are we supposed to do here?
       serverRes.pipe(clientRes);
       resolve();
       return;
     }
   
     if (resTransformStream === "close" || resTransformStream === "reset") {
       // If you kill the connection, we need to fire an upstream event separately here, since
       // this means the body won't be delivered in normal response events.
       if (options.emitEventCallback) {
         options.emitEventCallback!("passthrough-response-body", {
           overridden: true,
           rawBody: undefined,
         });
       }
   
       if (resTransformStream === "close") {
         (clientReq as any).socket.end();
       } else if (resTransformStream === "reset") {
         requireSocketResetSupport();
         resetOrDestroy(clientReq);
       }
   
       throw new AbortError(
         `Connection ${resTransformStream === "close" ? "closed" : "reset"} intentionally by rule`,
         `E_RULE_BRES_${resTransformStream.toUpperCase()}`,
       );
     }
   
     validateCustomHeaders(serverHeaders, resTransformStream?.headers);
   
     serverStatusCode = resTransformStream?.statusCode || serverStatusCode;
     serverStatusMessage = resTransformStream?.statusMessage || serverStatusMessage;
   
     serverHeaders = resTransformStream?.headers || serverHeaders;
   
     // delete the content-encoding to avoid needless CPU usage
     // there's no need to encode again if the downstream client is perfectly happy with raw bytes
     delete serverHeaders["content-encoding"];
   
     /**
      * Check if the .body is a buffer
      * If it's a buffer, we have the entire response ready to go, we can calculate the content-length
      * And pass it down immediately
      */
     if (Buffer.isBuffer(resTransformStream.body)) {
       serverHeaders["content-length"] = getContentLengthAfterModification(
         resTransformStream.body,
         serverRes.headers,
         resTransformStream?.headers,
         method === "HEAD", // HEAD responses are allowed mismatched content-length
       );
   
       // rebuild raw headers after content length modification
       serverRawHeaders = objectHeadersToRaw(serverHeaders);
   
       writeHead(
         clientRes,
         serverStatusCode,
         serverStatusMessage,
         serverRawHeaders.filter(([key, value]) => {
           if (key === ":status") return false;
           if (!validateHeader(key, value)) {
             console.warn(`Not forwarding invalid header: "${key}: ${value}"`);
             // Nothing else we can do in this case regardless - setHeaders will
             // throw within Node if we try to set this value.
             return false;
           }
           return true;
         }),
       );
   
       // Return the override data to the client:
       clientRes.end(resTransformStream.body);
       // Dump the real response data, in case that body wasn't read yet:
       serverRes.resume();
       resolve();
       return;
     } else {
       /**
        * Body is a stream | undefined
        * Write the modified headers and pipe the stream to client
        *
        * Because we're potentially transforming the body,
        * we're unsure of the content-length, so we don't set it
        */
       delete serverHeaders["content-length"];
   
       // get the raw headers
       serverRawHeaders = objectHeadersToRaw(serverHeaders);
   
       // send the headers first
       writeHead(
         clientRes,
         serverStatusCode,
         serverStatusMessage,
         serverRawHeaders.filter(([key, value]) => {
           if (key === ":status") return false;
           if (!validateHeader(key, value)) {
             console.warn(`Not forwarding invalid header: "${key}: ${value}"`);
             // Nothing else we can do in this case regardless - setHeaders will
             // throw within Node if we try to set this value.
             return false;
           }
           return true;
         }),
       );
   
       resTransformStream.body.pipe(clientRes);
       resTransformStream.body.once("end", resolve);
       return;
     }
   }

   Couple important notes to consider:

   1. This comes before calling beforeResponse and transformResponse, most of the body of the code is essentially copy-pasted from beforeResponse, with several modifications to avoid buffering response

   2. Since we're dealing with transform streams, I'm not sure how we can reliably calculate the resulting Content-Length... I think we can't. Therefore I'm just deleting the Content-Length (since the HTTP RFC seems to allow omitting it, and clients will just read the response until connection closes.

   3. I'm deleting the Content-Encoding, thus the downstream client will receive decoded content. This might be a questionable decision, but in my use case, I need to achieve low latency, and re-encoding seems to be unnecessary.
2. avarayr commented on May 13, 2025

   

   avarayr

   on May 13, 2025

   Last edited by avarayr

   More actions

   I also had to write a custom decodeBodyStream function because I couldn't find support for it in the @httptoolkit/http-encoding

   request-utils.ts

   export function decodeBodyStream(stream: stream.Readable, headers: Headers | RawHeaders) {
     const contentEncoding = getHeaderValue(headers, "content-encoding");
   
     if (!contentEncoding) return stream;
   
     let encodingLayers: string[] = [];
     if (contentEncoding) {
       const headerValue = Array.isArray(contentEncoding) ? contentEncoding.join(",") : String(contentEncoding);
       encodingLayers = headerValue.split(",").map((e) => e.trim().toLowerCase());
     }
   
     if (encodingLayers.length > 0) {
       // Apply decompression layers in order
       for (const enc of encodingLayers) {
         if (enc === "gzip" || enc === "x-gzip") {
           stream = stream.pipe(zlib.createGunzip());
         } else if (enc === "deflate") {
           stream = stream.pipe(zlib.createInflate());
         } else if (enc === "br" || enc === "brotli") {
           stream = stream.pipe(zlib.createBrotliDecompress());
         } else if (enc === "zstd") {
           if (typeof zlib.createZstdDecompress === "function") {
             stream = stream.pipe(zlib.createZstdDecompress());
           } else {
             // i had a weird edge case where zstd was not available in a particular electron version. not sure what happened there
             console.warn("Skipping Zstd streaming decompression: zlib.createZstdDecompress not available");
           }
         }
       }
     }
   
     return stream;
   }
   
   export function encodeBodyStream(stream: stream.Readable, headers: Headers | RawHeaders) {
     const contentEncoding = getHeaderValue(headers, "content-encoding");
   
     if (!contentEncoding) return stream;
   
     let encodingLayers: string[] = [];
     if (contentEncoding) {
       const headerValue = Array.isArray(contentEncoding) ? contentEncoding.join(",") : String(contentEncoding);
       encodingLayers = headerValue.split(",").map((e) => e.trim().toLowerCase());
     }
   
     if (encodingLayers.length > 0) {
       // Apply decompression layers in order
       for (const enc of encodingLayers) {
         if (enc === "gzip" || enc === "x-gzip") {
           stream = stream.pipe(zlib.createGzip());
         } else if (enc === "deflate") {
           stream = stream.pipe(zlib.createDeflate());
         } else if (enc === "br" || enc === "brotli") {
           stream = stream.pipe(zlib.createBrotliCompress());
         } else if (enc === "zstd") {
           if (typeof zlib.createZstdCompress === "function") {
             stream = stream.pipe(zlib.createZstdCompress());
           } else {
             console.warn("Skipping Zstd streaming compression: zlib.createZstdCompress not available");
           }
         }
       }
     }
   
     return stream;
   }