From 328a174f733b64a81bc3d7fe38da7cd9f0dc4ebc Mon Sep 17 00:00:00 2001
From: Frederik Bolding <frederik.bolding@gmail.com>
Date: Thu, 10 Mar 2022 18:04:05 +0100
Subject: [PATCH] Fix RateLimitController aggressive timeout issue (#716)

---
 src/ratelimit/RateLimitController.test.ts | 24 +++++++++++++++++++++++
 src/ratelimit/RateLimitController.ts      | 16 ++++++++-------
 2 files changed, 33 insertions(+), 7 deletions(-)

diff --git a/src/ratelimit/RateLimitController.test.ts b/src/ratelimit/RateLimitController.test.ts
index d5abc2da1b7..1ece067307c 100644
--- a/src/ratelimit/RateLimitController.test.ts
+++ b/src/ratelimit/RateLimitController.test.ts
@@ -147,4 +147,28 @@ describe('RateLimitController', () => {
       message,
     );
   });
+
+  it('timeout is only applied once per window', async () => {
+    const messenger = getRestrictedMessenger();
+    const controller = new RateLimitController({
+      implementations,
+      messenger,
+      rateLimitCount: 2,
+    });
+    expect(
+      await controller.call(origin, 'showNativeNotification', origin, message),
+    ).toBeUndefined();
+    jest.advanceTimersByTime(2500);
+    expect(
+      await controller.call(origin, 'showNativeNotification', origin, message),
+    ).toBeUndefined();
+    expect(controller.state.requests.showNativeNotification[origin]).toBe(2);
+    jest.advanceTimersByTime(2500);
+    expect(controller.state.requests.showNativeNotification[origin]).toBe(0);
+    expect(
+      await controller.call(origin, 'showNativeNotification', origin, message),
+    ).toBeUndefined();
+    jest.advanceTimersByTime(2500);
+    expect(controller.state.requests.showNativeNotification[origin]).toBe(1);
+  });
 });
diff --git a/src/ratelimit/RateLimitController.ts b/src/ratelimit/RateLimitController.ts
index 6eff3dede10..8d0340a2aaa 100644
--- a/src/ratelimit/RateLimitController.ts
+++ b/src/ratelimit/RateLimitController.ts
@@ -169,13 +169,15 @@ export class RateLimitController<
    */
   private recordRequest(api: keyof RateLimitedApis, origin: string) {
     this.update((state) => {
-      (state as any).requests[api][origin] =
-        ((state as any).requests[api][origin] ?? 0) + 1;
-
-      setTimeout(
-        () => this.resetRequestCount(api, origin),
-        this.rateLimitTimeout,
-      );
+      const previous = (state as any).requests[api][origin] ?? 0;
+      (state as any).requests[api][origin] = previous + 1;
+
+      if (previous === 0) {
+        setTimeout(
+          () => this.resetRequestCount(api, origin),
+          this.rateLimitTimeout,
+        );
+      }
     });
   }