From cd685c1b2b28d2427357a58a8d21e7d31d5f2ec6 Mon Sep 17 00:00:00 2001
From: Phillip Thelen <phillip@habitica.com>
Date: Wed, 29 May 2024 18:34:40 +0200
Subject: [PATCH] Fix liveliness probes being rate limited (#15236)

* Do not rate limit any liveliness probes

* update example config
---
 config.json.example                           |  1 +
 test/api/unit/middlewares/rateLimiter.test.js | 61 +++++++++++++++++++
 website/server/middlewares/rateLimiter.js     |  2 +
 3 files changed, 64 insertions(+)

diff --git a/config.json.example b/config.json.example
index 62c3439cd2..876f4e7c88 100644
--- a/config.json.example
+++ b/config.json.example
@@ -84,6 +84,7 @@
   "BLOCKED_IPS": "",
   "LOG_AMPLITUDE_EVENTS": "false",
   "RATE_LIMITER_ENABLED": "false",
+  "LIVELINESS_PROBE_KEY": "",
   "REDIS_HOST": "aaabbbcccdddeeefff",
   "REDIS_PORT": "1234",
   "REDIS_PASSWORD": "12345678",
diff --git a/test/api/unit/middlewares/rateLimiter.test.js b/test/api/unit/middlewares/rateLimiter.test.js
index 7b98773f48..4fa40ab343 100644
--- a/test/api/unit/middlewares/rateLimiter.test.js
+++ b/test/api/unit/middlewares/rateLimiter.test.js
@@ -87,6 +87,67 @@ describe('rateLimiter middleware', () => {
     expect(logger.error).to.have.been.calledWithMatch(Error, 'Rate Limiter Error');
   });
 
+  it('does not throw when LIVELINESS_PROBE_KEY is correct', async () => {
+    nconfGetStub.withArgs('RATE_LIMITER_ENABLED').returns('true');
+    nconfGetStub.withArgs('LIVELINESS_PROBE_KEY').returns('abc');
+    const attachRateLimiter = requireAgain(pathToRateLimiter).default;
+
+    req.query.liveliness = 'abc';
+    await attachRateLimiter(req, res, next);
+
+    expect(next).to.have.been.calledOnce;
+    const calledWith = next.getCall(0).args;
+    expect(typeof calledWith[0] === 'undefined').to.equal(true);
+    expect(res.set).to.not.have.been.called;
+  });
+
+  it('limits when LIVELINESS_PROBE_KEY is incorrect', async () => {
+    nconfGetStub.withArgs('RATE_LIMITER_ENABLED').returns('true');
+    nconfGetStub.withArgs('LIVELINESS_PROBE_KEY').returns('abc');
+    const attachRateLimiter = requireAgain(pathToRateLimiter).default;
+
+    req.query.liveliness = 'das';
+    await attachRateLimiter(req, res, next);
+
+    expect(next).to.have.been.calledOnce;
+    expect(res.set).to.have.been.calledWithMatch({
+      'X-RateLimit-Limit': 30,
+      'X-RateLimit-Remaining': 29,
+      'X-RateLimit-Reset': sinon.match(Date),
+    });
+  });
+
+  it('limits when LIVELINESS_PROBE_KEY is not set', async () => {
+    nconfGetStub.withArgs('RATE_LIMITER_ENABLED').returns('true');
+    nconfGetStub.withArgs('LIVELINESS_PROBE_KEY').returns(undefined);
+    const attachRateLimiter = requireAgain(pathToRateLimiter).default;
+
+    await attachRateLimiter(req, res, next);
+
+    expect(next).to.have.been.calledOnce;
+    expect(res.set).to.have.been.calledWithMatch({
+      'X-RateLimit-Limit': 30,
+      'X-RateLimit-Remaining': 29,
+      'X-RateLimit-Reset': sinon.match(Date),
+    });
+  });
+
+  it('throws when LIVELINESS_PROBE_KEY is blank', async () => {
+    nconfGetStub.withArgs('RATE_LIMITER_ENABLED').returns('true');
+    nconfGetStub.withArgs('LIVELINESS_PROBE_KEY').returns('');
+    const attachRateLimiter = requireAgain(pathToRateLimiter).default;
+
+    req.query.liveliness = '';
+    await attachRateLimiter(req, res, next);
+
+    expect(next).to.have.been.calledOnce;
+    expect(res.set).to.have.been.calledWithMatch({
+      'X-RateLimit-Limit': 30,
+      'X-RateLimit-Remaining': 29,
+      'X-RateLimit-Reset': sinon.match(Date),
+    });
+  });
+
   it('throws when there are no available points remaining', async () => {
     nconfGetStub.withArgs('RATE_LIMITER_ENABLED').returns('true');
     const attachRateLimiter = requireAgain(pathToRateLimiter).default;
diff --git a/website/server/middlewares/rateLimiter.js b/website/server/middlewares/rateLimiter.js
index b6706db11f..e62f85ff92 100644
--- a/website/server/middlewares/rateLimiter.js
+++ b/website/server/middlewares/rateLimiter.js
@@ -21,6 +21,7 @@ const RATE_LIMITER_ENABLED = nconf.get('RATE_LIMITER_ENABLED') === 'true';
 const REDIS_HOST = nconf.get('REDIS_HOST');
 const REDIS_PASSWORD = nconf.get('REDIS_PASSWORD');
 const REDIS_PORT = nconf.get('REDIS_PORT');
+const LIVELINESS_PROBE_KEY = nconf.get('LIVELINESS_PROBE_KEY');
 
 let redisClient;
 let rateLimiter;
@@ -71,6 +72,7 @@ function setResponseHeaders (res, rateLimiterRes) {
 
 export default function rateLimiterMiddleware (req, res, next) {
   if (!RATE_LIMITER_ENABLED) return next();
+  if (LIVELINESS_PROBE_KEY && req.query.liveliness === LIVELINESS_PROBE_KEY) return next();
 
   const userId = req.header('x-api-user');