From 4ec136ae8257a1a9f40d83ad61c57954226e79e0 Mon Sep 17 00:00:00 2001
From: Justin Wind <justin.wind+git@gmail.com>
Date: Sat, 12 Mar 2022 14:03:12 -0800
Subject: [PATCH] add option to persist response body in context for HEAD
 requests

---
 lib/dingus.js      | 11 ++++++++---
 test/lib/dingus.js | 12 +++++++++++-
 2 files changed, 19 insertions(+), 4 deletions(-)

diff --git a/lib/dingus.js b/lib/dingus.js
index 11bdf99..1098ea7 100644
--- a/lib/dingus.js
+++ b/lib/dingus.js
@@ -226,11 +226,13 @@ class Dingus {
   /**
    * Intercept writes for head requests, do not send to client,
    * but send length, and make body available in context.
+   * N.B. If persisted, ctx.responseBody will be a raw buffer, be aware when logging.
    * @param {http.ClientRequest} req 
    * @param {http.ServerResponse} res 
    * @param {object} ctx 
+   * @param {Boolean} persistResponseBody
    */
-  static setHeadHandler(req, res, ctx) {
+  static setHeadHandler(req, res, ctx, persistResponseBody = false) {
     if (req.method === 'HEAD') {
       const origEnd = res.end.bind(res);
       const chunks = [];
@@ -240,8 +242,11 @@ class Dingus {
       };
       res.end = function (data, encoding, ...rest) {
         Dingus.pushBufChunk(chunks, data, encoding);
-        ctx.responseBody = Buffer.concat(chunks);
-        res.setHeader(Enum.Header.ContentLength, Buffer.byteLength(ctx.responseBody));
+        const responseBody = Buffer.concat(chunks);
+        res.setHeader(Enum.Header.ContentLength, Buffer.byteLength(responseBody));
+        if (persistResponseBody) {
+          ctx.responseBody = responseBody;
+        }
         return origEnd(undefined, encoding, ...rest);
       };
     }
diff --git a/test/lib/dingus.js b/test/lib/dingus.js
index 23b7e8c..b478a45 100644
--- a/test/lib/dingus.js
+++ b/test/lib/dingus.js
@@ -269,7 +269,7 @@ describe('Dingus', function () {
       };
       ctx = {};
     });
-    it('collects body without writing', function () {
+    it('collects response without writing', function () {
       Dingus.setHeadHandler(req, res, ctx);
       res.write(Buffer.from('foo'));
       res.write('baz');
@@ -277,6 +277,16 @@ describe('Dingus', function () {
       res.end('quux');
       assert(!origWrite.called);
       assert(origEnd.called);
+      assert.deepStrictEqual(ctx.responseBody, undefined);
+    });
+    it('collects response without writing, persists written data', function () {
+      Dingus.setHeadHandler(req, res, ctx, true);
+      res.write(Buffer.from('foo'));
+      res.write('baz');
+      res.write();
+      res.end('quux');
+      assert(!origWrite.called);
+      assert(origEnd.called);
       assert.deepStrictEqual(ctx.responseBody, Buffer.from('foobazquux'));
     });
     it('ignores non-head method', function () {
-- 
2.45.2