Fix RateLimitController aggressive timeout issue (#716)

MetaMask · Mar 10, 2022 · 6ee6e41 · 6ee6e41
1 parent a799f42
commit 6ee6e41
Show file tree

Hide file tree

Showing 2 changed files with 33 additions and 7 deletions.
diff --git 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
@@ -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,
+        );
+      }
     });
   }