From c49b9c82a633f8c6ef7379cc0a2dd276598b35cb Mon Sep 17 00:00:00 2001
From: Eran Hammer <eran@hammer.io>
Date: Wed, 18 Oct 2017 00:16:27 -0700
Subject: [PATCH] Expose compressor to response streams. Closes #3599

---
 API.md           |  8 ++++++-
 lib/transmit.js  |  9 +++++++-
 test/transmit.js | 60 ++++++++++++++++++++++++++++++++++++++++--------
 3 files changed, 66 insertions(+), 11 deletions(-)

diff --git a/API.md b/API.md
index 0408d4e32..0750740fc 100755
--- a/API.md
+++ b/API.md
@@ -3323,7 +3323,13 @@ The return value must be one of:
 - `Stream` object
     - must be compatible with the "streams2" API and not be in `objectMode`.
     - if the stream object has a `statusCode` property, that status code will be used as
-      the default response code.
+      the default response code based on the [`passThrough`](#response.settings.passThrough)
+      option.
+    - if the stream object has a `headers` property, the headers will be included in the response
+      based on the [`passThrough`](#response.settings.passThrough) option.
+    - if the stream object has a function property `setCompressor(compressor)` and the response
+      passes through a compressor, a reference to the compressor stream will be passed to the
+      response stream via this method.
 - any object or array
     - must not include circular references.
 - a toolkit signal:
diff --git a/lib/transmit.js b/lib/transmit.js
index 10b598b03..05cba5572 100755
--- a/lib/transmit.js
+++ b/lib/transmit.js
@@ -222,7 +222,14 @@ internals.encoding = function (response, encoding) {
 
     delete response.headers['content-length'];
     response._header('content-encoding', encoding);
-    return request._core.compression.encoder(request, encoding);
+    const compressor = request._core.compression.encoder(request, encoding);
+    if (response.variety === 'stream' &&
+        typeof response._payload.setCompressor === 'function') {
+
+        response._payload.setCompressor(compressor);
+    }
+
+    return compressor;
 };
 
 
diff --git a/test/transmit.js b/test/transmit.js
index f7270bb25..1cbb99759 100755
--- a/test/transmit.js
+++ b/test/transmit.js
@@ -371,15 +371,8 @@ describe('transmission', () => {
             expect(res.headers['content-encoding']).to.equal('gzip');
             expect(res.headers.vary).to.equal('accept-encoding');
 
-            await new Promise((resolve) => {
-
-                Zlib.unzip(res.rawPayload, (err, result) => {
-
-                    expect(err).to.not.exist();
-                    expect(result.toString()).to.equal('/**/docall({"first":"1","last":"2"});');
-                    resolve();
-                });
-            });
+            const uncompressed = await internals.uncompress('unzip', res.rawPayload);
+            expect(uncompressed.toString()).to.equal('/**/docall({"first":"1","last":"2"});');
         });
 
         it('returns an JSONP response when response is a buffer', async () => {
@@ -1772,6 +1765,49 @@ describe('transmission', () => {
         });
     });
 
+    describe('encoding()', () => {
+
+        it('passes compressor to stream', async () => {
+
+            const handler = (request, h) => {
+
+                const TestStream = class extends Stream.Readable {
+
+                    _read(size) {
+
+                        if (this.isDone) {
+                            return;
+                        }
+                        this.isDone = true;
+
+                        this.push('some payload');
+                        this._compressor.flush();
+
+                        setTimeout(() => {
+
+                            this.push(' and some other payload');
+                            this.push(null);
+                        }, 10);
+                    }
+
+                    setCompressor(compressor) {
+
+                        this._compressor = compressor;
+                    }
+                };
+
+                return h.response(new TestStream()).type('text/html');
+            };
+
+            const server = Hapi.server({ compression: { minBytes: 1 } });
+            server.route({ method: 'GET', path: '/', handler });
+
+            const res = await server.inject({ url: '/', headers: { 'accept-encoding': 'gzip' } });
+            const uncompressed = await internals.uncompress('unzip', res.rawPayload);
+            expect(uncompressed.toString()).to.equal('some payload and some other payload');
+        });
+    });
+
     describe('writeHead()', () => {
 
         it('set custom statusMessage', async () => {
@@ -1841,3 +1877,9 @@ internals.compress = function (encoder, value) {
 
     return new Promise((resolve) => Zlib[encoder](value, (ignoreErr, compressed) => resolve(compressed)));
 };
+
+
+internals.uncompress = function (decoder, value) {
+
+    return new Promise((resolve) => Zlib[decoder](value, (ignoreErr, uncompressed) => resolve(uncompressed)));
+};