From 4e316a74ec139f5d81db5d197d0d51735da9d67e Mon Sep 17 00:00:00 2001
From: Sam Adams <sam.adams@superpayments.com>
Date: Thu, 28 Nov 2024 07:36:49 +0000
Subject: [PATCH] fix: prevent possible slow resolution of rate limiter

We have found scenarios at scale where our requests lock up for 10+ seconds and wait for something in the rate limiter
while resolving a jwk endpoint.
We believe the intention is for the rate limiter to return immediately and not pause to wait for available tokens
(hence passing the fireImmediately arg to the RateLimiter constructor).

However, maybe because of a race condition at scale, it seems possible to end up trying
to resolve a token from the bucket even when there are no tokens left:
https://github.com/jhurliman/node-rate-limiter/blob/main/src/RateLimiter.ts#L83
Once we are on this path the rate limiter can trigger a wait for new 'tokens' become available:
https://github.com/jhurliman/node-rate-limiter/blob/main/src/TokenBucket.ts#L96

This fix uses the simpler `tryRemoveTokens` method which synchronously returns a boolean if tokens could be taken
(and so can't pause execution to wait for tokens to be available).
---
 src/wrappers/rateLimit.js | 28 ++++++++--------------------
 1 file changed, 8 insertions(+), 20 deletions(-)

diff --git a/src/wrappers/rateLimit.js b/src/wrappers/rateLimit.js
index b1837c4e..2b115280 100644
--- a/src/wrappers/rateLimit.js
+++ b/src/wrappers/rateLimit.js
@@ -9,26 +9,14 @@ function rateLimitWrapper(client, { jwksRequestsPerMinute = 10 }) {
   const limiter = new RateLimiter(jwksRequestsPerMinute, 'minute', true);
   logger(`Configured rate limiting to JWKS endpoint at ${jwksRequestsPerMinute}/minute`);
 
-  return async (kid) => await new Promise((resolve, reject) => {
-    limiter.removeTokens(1, async (err, remaining) => {
-      if (err) {
-        reject(err);
-      }
-
-      logger('Requests to the JWKS endpoint available for the next minute:', remaining);
-      if (remaining < 0) {
-        logger('Too many requests to the JWKS endpoint');
-        reject(new JwksRateLimitError('Too many requests to the JWKS endpoint'));
-      } else {
-        try {
-          const key = await getSigningKey(kid);
-          resolve(key);
-        } catch (error) {
-          reject(error);
-        }
-      }
-    });
-  });
+  return async (kid) => {
+    logger('Requests to the JWKS endpoint available for the next minute:', limiter.getTokensRemaining());
+    if (limiter.tryRemoveTokens(1)) {
+      return getSigningKey(kid);
+    }
+    logger('Too many requests to the JWKS endpoint');
+    throw new JwksRateLimitError('Too many requests to the JWKS endpoint');
+  };
 }
 
 module.exports.default = rateLimitWrapper;