Add response lifecycle tracking and checks (#4257)

* Add response lifecycle tracking and checks. Closes #4256

* Cleanup

Author: kanongil

lib/response.js

@@ -55,10 +55,11 @@ exports = module.exports = internals.Response = class {
 
         this._events = null;
         this._payload = null;                       // Readable stream
-        this._error = null;                         // The boom object when created from an error (used for logging)
+        this._error = options.error || null;        // The boom object when created from an error (used for logging)
         this._contentType = null;                   // Used if no explicit content-type is set and type is known
         this._takeover = false;
         this._statusCode = false;                   // true when code() called
+        this._state = this._error ? 'prepare' : 'init'; // One of 'init', 'prepare', 'marshall', 'close'
 
         this._processors = {
             marshal: options.marshal,
@@ -495,6 +496,10 @@ exports = module.exports = internals.Response = class {
 
     _prepare() {
 
+        Hoek.assert(this._state === 'init');
+
+        this._state = 'prepare';
+
         this._passThrough();
 
         if (!this._processors.prepare) {
@@ -558,10 +563,14 @@ exports = module.exports = internals.Response = class {
 
     async _marshal() {
 
-        let source = this.source;
+        Hoek.assert(this._state === 'prepare');
+
+        this._state = 'marshall';
 
         // Processor marshal
 
+        let source = this.source;
+
         if (this._processors.marshal) {
             try {
                 source = await this._processors.marshal(this);
@@ -638,6 +647,12 @@ exports = module.exports = internals.Response = class {
 
     _close() {
 
+        if (this._state === 'close') {
+            return;
+        }
+
+        this._state = 'close';
+
         if (this._processors.close) {
             try {
                 this._processors.close(this);

lib/toolkit.js

@@ -101,7 +101,7 @@ exports.Manager = class {
 
         if (typeof response !== 'symbol') {
             response = request._core.Response.wrap(response, request);
-            if (!response.isBoom) {
+            if (!response.isBoom && response._state === 'init') {
                 response = await response._prepare();
             }
         }

lib/transmit.js

@@ -38,7 +38,9 @@ exports.send = async function (request) {
 internals.marshal = async function (response) {
 
     for (const func of response.request._route._marshalCycle) {
-        await func(response);
+        if (response._state !== 'close') {
+            await func(response);
+        }
     }
 };
 
@@ -72,8 +74,7 @@ internals.fail = async function (request, boom) {
 internals.error = function (request, boom) {
 
     const error = boom.output;
-    const response = new request._core.Response(error.payload, request);
-    response._error = boom;
+    const response = new request._core.Response(error.payload, request, { error: boom });
     response.code(error.statusCode);
     response.headers = Hoek.clone(error.headers);               // Prevent source from being modified
     return response;

test/response.js

@@ -1290,6 +1290,37 @@ describe('Response', () => {
             const res = await server.inject('/');
             expect(res.statusCode).to.equal(500);
         });
+
+        it('is only called once for returned responses', async () => {
+
+            let calls = 0;
+            const pre = (request, h) => {
+
+                const prepare = (response) => {
+
+                    ++calls;
+                    return response;
+                };
+
+                return request.generateResponse(null, { prepare });
+            };
+
+            const server = Hapi.server();
+            server.route({
+                method: 'GET',
+                path: '/',
+                options: {
+                    pre: [
+                        { method: pre, assign: 'p' }
+                    ],
+                    handler: (request) => request.preResponses.p
+                }
+            });
+
+            const res = await server.inject('/');
+            expect(res.statusCode).to.equal(204);
+            expect(calls).to.equal(1);
+        });
     });
 
     describe('_tap()', () => {