From 5c7eb95319aa54ef3b4b60d000d434824a666e18 Mon Sep 17 00:00:00 2001
From: Joey Harrington <jharrington@netflix.com>
Date: Wed, 9 Feb 2022 11:49:27 -0800
Subject: [PATCH] fix: use close event on response instead of socket (#1892)

---
 lib/plugins/bodyReader.js       |   2 +-
 test/plugins/bodyReader.test.js | 136 ++++++++++++++++++++++++++++----
 2 files changed, 123 insertions(+), 15 deletions(-)

diff --git a/lib/plugins/bodyReader.js b/lib/plugins/bodyReader.js
index 82c9d1016..60ff1f462 100644
--- a/lib/plugins/bodyReader.js
+++ b/lib/plugins/bodyReader.js
@@ -186,7 +186,7 @@ function bodyReader(options) {
         // add 'close and 'aborted' event handlers so that requests (and their
         // corresponding memory) don't leak if client stops sending data half
         // way through a POST request
-        req.socket.once('close', next);
+        res.once('close', next);
         req.once('aborted', next);
         req.resume();
     }
diff --git a/test/plugins/bodyReader.test.js b/test/plugins/bodyReader.test.js
index b799fef90..3e48a4298 100644
--- a/test/plugins/bodyReader.test.js
+++ b/test/plugins/bodyReader.test.js
@@ -18,25 +18,25 @@ var CLIENT;
 var PORT;
 
 describe('body reader', function() {
-    describe('gzip content encoding', function() {
-        beforeEach(function(done) {
-            SERVER = restify.createServer({
-                dtrace: helper.dtrace,
-                log: helper.getLog('server')
-            });
+    beforeEach(function(done) {
+        SERVER = restify.createServer({
+            dtrace: helper.dtrace,
+            log: helper.getLog('server')
+        });
 
-            SERVER.listen(0, '127.0.0.1', function() {
-                PORT = SERVER.address().port;
+        SERVER.listen(0, '127.0.0.1', function() {
+            PORT = SERVER.address().port;
 
-                done();
-            });
+            done();
         });
+    });
 
-        afterEach(function(done) {
-            CLIENT.close();
-            SERVER.close(done);
-        });
+    afterEach(function(done) {
+        CLIENT.close();
+        SERVER.close(done);
+    });
 
+    describe('gzip content encoding', function() {
         it('should parse gzip encoded content', function(done) {
             SERVER.use(restify.plugins.bodyParser());
 
@@ -187,4 +187,112 @@ describe('body reader', function() {
             req.write(postData);
         });
     });
+
+    it('should not add a listener for each call on same socket', done => {
+        SERVER.use(restify.plugins.bodyParser());
+
+        let serverReq, serverRes, serverReqSocket;
+        SERVER.post('/meals', function(req, res, next) {
+            serverReq = req;
+            serverRes = res;
+            serverReqSocket = req.socket;
+            res.send();
+            next();
+        });
+
+        CLIENT = restifyClients.createJsonClient({
+            url: 'http://127.0.0.1:' + PORT,
+            retry: false,
+            agent: new http.Agent({ keepAlive: true })
+        });
+
+        CLIENT.post('/meals', { breakfast: 'pancakes' }, (err, _, res) => {
+            assert.ifError(err);
+            assert.equal(res.statusCode, 200);
+
+            const firstReqSocket = serverReqSocket;
+            const numReqListeners = listenerCount(serverReq);
+            const numResListeners = listenerCount(serverRes);
+            const numReqSocketListeners = listenerCount(serverReq.socket);
+
+            // Without setImmediate, the second request will not reuse the socket.
+            setImmediate(() => {
+                CLIENT.post('/meals', { lunch: 'salad' }, (err2, __, res2) => {
+                    assert.ifError(err2);
+                    assert.equal(res2.statusCode, 200);
+                    assert.equal(
+                        serverReqSocket,
+                        firstReqSocket,
+                        'This test should issue two requests that share the ' +
+                            'same socket.'
+                    );
+                    // The number of listeners on each emitter should not have
+                    // increased since the first request.
+                    assert.equal(listenerCount(serverReq), numReqListeners);
+                    assert.equal(listenerCount(serverRes), numResListeners);
+                    assert.equal(
+                        listenerCount(serverReq.socket),
+                        numReqSocketListeners
+                    );
+                    done();
+                });
+            });
+        });
+    });
+
+    it('should call next for each successful request on same socket', done => {
+        let nextCallCount = 0;
+        SERVER.use(restify.plugins.bodyParser());
+        SERVER.use((req, res, next) => {
+            nextCallCount += 1;
+            next();
+        });
+
+        let serverReqSocket;
+        SERVER.post('/meals', function(req, res, next) {
+            res.send();
+            next();
+        });
+
+        CLIENT = restifyClients.createJsonClient({
+            url: 'http://127.0.0.1:' + PORT,
+            retry: false,
+            agent: new http.Agent({ keepAlive: true })
+        });
+
+        CLIENT.post('/meals', { breakfast: 'waffles' }, (err, _, res) => {
+            assert.ifError(err);
+            assert.equal(res.statusCode, 200);
+            const firstReqSocket = serverReqSocket;
+            assert.equal(nextCallCount, 1);
+
+            // Without setImmediate, the second request will not reuse the socket.
+            setImmediate(() => {
+                CLIENT.post('/meals', { lunch: 'candy' }, (err2, __, res2) => {
+                    assert.ifError(err2);
+                    assert.equal(res2.statusCode, 200);
+                    assert.equal(
+                        serverReqSocket,
+                        firstReqSocket,
+                        'This test should issue two requests that share the ' +
+                            'same socket.'
+                    );
+                    assert.equal(nextCallCount, 2);
+                    done();
+                });
+            });
+        });
+    });
 });
+
+/**
+ * @param {EventEmitter} emitter - An emitter
+ * @returns {number} - The total number of listeners across all events
+ */
+function listenerCount(emitter) {
+    let numListeners = 0;
+    for (const eventName of emitter.eventNames()) {
+        numListeners += emitter.listenerCount(eventName);
+    }
+    return numListeners;
+}