From a63a4bce90210774e46481ca9e29a076f6e733ae Mon Sep 17 00:00:00 2001
From: Artur K <fatal10110@gmail.com>
Date: Sat, 11 Sep 2021 14:55:09 +0300
Subject: [PATCH] http: limit requests per connection

Fixes: https://github.com/nodejs/node/issues/40071
PR-URL: https://github.com/nodejs/node/pull/40082
Reviewed-By: Matteo Collina <matteo.collina@gmail.com>
Reviewed-By: Robert Nagy <ronagy@icloud.com>
---
 doc/api/http.md                               |  16 +++
 lib/_http_outgoing.js                         |   5 +-
 lib/_http_server.js                           |  51 +++++---
 .../test-http-keep-alive-max-requests.js      | 116 ++++++++++++++++++
 ...t-http-keep-alive-pipeline-max-requests.js |  86 +++++++++++++
 5 files changed, 258 insertions(+), 16 deletions(-)
 create mode 100644 test/parallel/test-http-keep-alive-max-requests.js
 create mode 100644 test/parallel/test-http-keep-alive-pipeline-max-requests.js

diff --git a/doc/api/http.md b/doc/api/http.md
index 0cbfce0fbe67d9..f41d9f74e27718 100644
--- a/doc/api/http.md
+++ b/doc/api/http.md
@@ -1348,6 +1348,22 @@ By default, the Server does not timeout sockets. However, if a callback
 is assigned to the Server's `'timeout'` event, timeouts must be handled
 explicitly.
 
+### `server.maxRequestsPerSocket`
+<!-- YAML
+added: REPLACEME
+-->
+
+* {number} Requests per socket. **Default:** null (no limit)
+
+The maximum number of requests socket can handle
+before closing keep alive connection.
+
+A value of `null` will disable the limit.
+
+When limit is reach it will set `Connection` header value to `closed`,
+but will not actually close the connection, subsequent requests sent
+after the limit is reached will get `503 Service Unavailable` as a response.
+
 ### `server.timeout`
 <!-- YAML
 added: v0.9.12
diff --git a/lib/_http_outgoing.js b/lib/_http_outgoing.js
index 25fe8fb1ede73f..27e290a2b916cd 100644
--- a/lib/_http_outgoing.js
+++ b/lib/_http_outgoing.js
@@ -113,6 +113,7 @@ function OutgoingMessage() {
   this._last = false;
   this.chunkedEncoding = false;
   this.shouldKeepAlive = true;
+  this.maxRequestsOnConnectionReached = false;
   this._defaultKeepAlive = true;
   this.useChunkedEncodingByDefault = true;
   this.sendDate = false;
@@ -446,7 +447,9 @@ function _storeHeader(firstLine, headers) {
   } else if (!state.connection) {
     const shouldSendKeepAlive = this.shouldKeepAlive &&
         (state.contLen || this.useChunkedEncodingByDefault || this.agent);
-    if (shouldSendKeepAlive) {
+    if (shouldSendKeepAlive && this.maxRequestsOnConnectionReached) {
+      header += 'Connection: close\r\n';
+    } else if (shouldSendKeepAlive) {
       header += 'Connection: keep-alive\r\n';
       if (this._keepAliveTimeout && this._defaultKeepAlive) {
         const timeoutSeconds = MathFloor(this._keepAliveTimeout / 1000);
diff --git a/lib/_http_server.js b/lib/_http_server.js
index 3ef1b4ad5d4735..426f334bcbcb71 100644
--- a/lib/_http_server.js
+++ b/lib/_http_server.js
@@ -396,6 +396,7 @@ function Server(options, requestListener) {
   this.timeout = 0;
   this.keepAliveTimeout = 5000;
   this.maxHeadersCount = null;
+  this.maxRequestsPerSocket = null;
   this.headersTimeout = 60 * 1000; // 60 seconds
   this.requestTimeout = 0;
 }
@@ -487,6 +488,7 @@ function connectionListenerInternal(server, socket) {
     // need to pause TCP socket/HTTP parser, and wait until the data will be
     // sent to the client.
     outgoingData: 0,
+    requestsCount: 0,
     keepAliveTimeoutSet: false
   };
   state.onData = socketOnData.bind(undefined,
@@ -903,28 +905,47 @@ function parserOnIncoming(server, socket, state, req, keepAlive) {
          resOnFinish.bind(undefined,
                           req, res, socket, state, server));
 
-  if (req.headers.expect !== undefined &&
-      (req.httpVersionMajor === 1 && req.httpVersionMinor === 1)) {
-    if (RegExpPrototypeTest(continueExpression, req.headers.expect)) {
-      res._expect_continue = true;
+  let handled = false;
 
-      if (server.listenerCount('checkContinue') > 0) {
-        server.emit('checkContinue', req, res);
+  if (req.httpVersionMajor === 1 && req.httpVersionMinor === 1) {
+    if (typeof server.maxRequestsPerSocket === 'number') {
+      state.requestsCount++;
+      res.maxRequestsOnConnectionReached = (
+        server.maxRequestsPerSocket <= state.requestsCount);
+    }
+
+    if (typeof server.maxRequestsPerSocket === 'number' &&
+      (server.maxRequestsPerSocket < state.requestsCount)) {
+      handled = true;
+
+      res.writeHead(503);
+      res.end();
+    } else if (req.headers.expect !== undefined) {
+      handled = true;
+
+      if (RegExpPrototypeTest(continueExpression, req.headers.expect)) {
+        res._expect_continue = true;
+
+        if (server.listenerCount('checkContinue') > 0) {
+          server.emit('checkContinue', req, res);
+        } else {
+          res.writeContinue();
+          server.emit('request', req, res);
+        }
+      } else if (server.listenerCount('checkExpectation') > 0) {
+        server.emit('checkExpectation', req, res);
       } else {
-        res.writeContinue();
-        server.emit('request', req, res);
+        res.writeHead(417);
+        res.end();
       }
-    } else if (server.listenerCount('checkExpectation') > 0) {
-      server.emit('checkExpectation', req, res);
-    } else {
-      res.writeHead(417);
-      res.end();
     }
-  } else {
-    req.on('end', clearRequestTimeout);
+  }
 
+  if (!handled) {
+    req.on('end', clearRequestTimeout);
     server.emit('request', req, res);
   }
+
   return 0;  // No special treatment.
 }
 
diff --git a/test/parallel/test-http-keep-alive-max-requests.js b/test/parallel/test-http-keep-alive-max-requests.js
new file mode 100644
index 00000000000000..657b59ae6d9369
--- /dev/null
+++ b/test/parallel/test-http-keep-alive-max-requests.js
@@ -0,0 +1,116 @@
+'use strict';
+
+const common = require('../common');
+const net = require('net');
+const http = require('http');
+const assert = require('assert');
+
+const bodySent = 'This is my request';
+
+function assertResponse(headers, body, expectClosed) {
+  if (expectClosed) {
+    assert.match(headers, /Connection: close\r\n/m);
+    assert.strictEqual(headers.search(/Keep-Alive: timeout=5\r\n/m), -1);
+    assert.match(body, /Hello World!/m);
+  } else {
+    assert.match(headers, /Connection: keep-alive\r\n/m);
+    assert.match(headers, /Keep-Alive: timeout=5\r\n/m);
+    assert.match(body, /Hello World!/m);
+  }
+}
+
+function writeRequest(socket, withBody) {
+  if (withBody) {
+    socket.write('POST / HTTP/1.1\r\n');
+    socket.write('Connection: keep-alive\r\n');
+    socket.write('Content-Type: text/plain\r\n');
+    socket.write(`Content-Length: ${bodySent.length}\r\n\r\n`);
+    socket.write(`${bodySent}\r\n`);
+    socket.write('\r\n\r\n');
+  } else {
+    socket.write('GET / HTTP/1.1\r\n');
+    socket.write('Connection: keep-alive\r\n');
+    socket.write('\r\n\r\n');
+  }
+}
+
+const server = http.createServer((req, res) => {
+  let body = '';
+  req.on('data', (data) => {
+    body += data;
+  });
+
+  req.on('end', () => {
+    if (req.method === 'POST') {
+      assert.strictEqual(bodySent, body);
+    }
+    res.writeHead(200, { 'Content-Type': 'text/plain' });
+    res.write('Hello World!');
+    res.end();
+  });
+});
+
+function initialRequests(socket, numberOfRequests, cb) {
+  let buffer = '';
+
+  writeRequest(socket);
+
+  socket.on('data', (data) => {
+    buffer += data;
+
+    if (buffer.endsWith('\r\n\r\n')) {
+      if (--numberOfRequests === 0) {
+        socket.removeAllListeners('data');
+        cb();
+      } else {
+        const [headers, body] = buffer.trim().split('\r\n\r\n');
+        assertResponse(headers, body);
+        buffer = '';
+        writeRequest(socket, true);
+      }
+    }
+  });
+}
+
+
+server.maxRequestsPerSocket = 3;
+server.listen(0, common.mustCall((res) => {
+  const socket = new net.Socket();
+  const anotherSocket = new net.Socket();
+
+  socket.on('end', common.mustCall(() => {
+    server.close();
+  }));
+
+  socket.on('ready', common.mustCall(() => {
+    // Do 2 of 3 allowed requests and ensure they still alive
+    initialRequests(socket, 2, common.mustCall(() => {
+      anotherSocket.connect({ port: server.address().port });
+    }));
+  }));
+
+  anotherSocket.on('ready', common.mustCall(() => {
+    // Do another 2 requests with another socket
+    // enusre that this will not affect the first socket
+    initialRequests(anotherSocket, 2, common.mustCall(() => {
+      let buffer = '';
+
+      // Send the rest of the calls to the first socket
+      // and see connection is closed
+      socket.on('data', common.mustCall((data) => {
+        buffer += data;
+
+        if (buffer.endsWith('\r\n\r\n')) {
+          const [headers, body] = buffer.trim().split('\r\n\r\n');
+          assertResponse(headers, body, true);
+          anotherSocket.end();
+          socket.end();
+        }
+      }));
+
+      writeRequest(socket, true);
+    }));
+  }));
+
+  socket.connect({ port: server.address().port });
+}));
diff --git a/test/parallel/test-http-keep-alive-pipeline-max-requests.js b/test/parallel/test-http-keep-alive-pipeline-max-requests.js
new file mode 100644
index 00000000000000..9c5d46a57ce197
--- /dev/null
+++ b/test/parallel/test-http-keep-alive-pipeline-max-requests.js
@@ -0,0 +1,86 @@
+'use strict';
+
+const common = require('../common');
+const net = require('net');
+const http = require('http');
+const assert = require('assert');
+
+const bodySent = 'This is my request';
+
+function assertResponse(headers, body, expectClosed) {
+  if (expectClosed) {
+    assert.match(headers, /Connection: close\r\n/m);
+    assert.strictEqual(headers.search(/Keep-Alive: timeout=5\r\n/m), -1);
+    assert.match(body, /Hello World!/m);
+  } else {
+    assert.match(headers, /Connection: keep-alive\r\n/m);
+    assert.match(headers, /Keep-Alive: timeout=5\r\n/m);
+    assert.match(body, /Hello World!/m);
+  }
+}
+
+function writeRequest(socket) {
+  socket.write('POST / HTTP/1.1\r\n');
+  socket.write('Connection: keep-alive\r\n');
+  socket.write('Content-Type: text/plain\r\n');
+  socket.write(`Content-Length: ${bodySent.length}\r\n\r\n`);
+  socket.write(`${bodySent}\r\n`);
+  socket.write('\r\n\r\n');
+}
+
+const server = http.createServer((req, res) => {
+  let body = '';
+  req.on('data', (data) => {
+    body += data;
+  });
+
+  req.on('end', () => {
+    if (req.method === 'POST') {
+      assert.strictEqual(bodySent, body);
+    }
+
+    res.writeHead(200, { 'Content-Type': 'text/plain' });
+    res.write('Hello World!');
+    res.end();
+  });
+});
+
+server.maxRequestsPerSocket = 3;
+
+server.listen(0, common.mustCall((res) => {
+  const socket = new net.Socket();
+
+  socket.on('end', common.mustCall(() => {
+    server.close();
+  }));
+
+  socket.on('ready', common.mustCall(() => {
+    writeRequest(socket);
+    writeRequest(socket);
+    writeRequest(socket);
+    writeRequest(socket);
+  }));
+
+  let buffer = '';
+
+  socket.on('data', (data) => {
+    buffer += data;
+
+    const responseParts = buffer.trim().split('\r\n\r\n');
+
+    if (responseParts.length === 8) {
+      assertResponse(responseParts[0], responseParts[1]);
+      assertResponse(responseParts[2], responseParts[3]);
+      assertResponse(responseParts[4], responseParts[5], true);
+
+      assert.match(responseParts[6], /HTTP\/1\.1 503 Service Unavailable/m);
+      assert.match(responseParts[6], /Connection: close\r\n/m);
+      assert.strictEqual(responseParts[6].search(/Keep-Alive: timeout=5\r\n/m), -1);
+      assert.strictEqual(responseParts[7].search(/Hello World!/m), -1);
+
+      socket.end();
+    }
+  });
+
+  socket.connect({ port: server.address().port });
+}));