Merge pull request #20868 from mozilla/fix-current-attempts-logic

bug(rate-limitting): bqWriter missing current attempts

Author: dschom

libs/accounts/rate-limit/src/lib/rate-limit.spec.ts

@@ -9,6 +9,8 @@ import { BlockRecord, Rule, TOO_MANY_ATTEMPTS } from './models';
 import { calculateRetryAfter, getKey } from './util';
 import { StatsD } from 'hot-shots';
 
+const MOCK_NOW = 1_700_000_000_000;
+
 describe('rate-limit', () => {
   let redis: Redis;
   let mockIncrement: jest.Mock;
@@ -857,6 +859,126 @@ describe('rate-limit', () => {
       );
     });
 
+    it('reports currentAttempts as the tallied count when no block is triggered', async () => {
+      mockIncr.mockResolvedValue(3); // under maxAttempts of 5
+
+      rateLimit = new RateLimit(
+        { rules: parseConfigRules(['test:ip:5:1 minute:5 minutes:block']) },
+        redis,
+        statsd,
+        { write: mockWrite }
+      );
+
+      await rateLimit.check('test', { ip: '1.2.3.4' });
+
+      expect(mockWrite).toHaveBeenCalledWith(
+        expect.objectContaining({
+          wasBlocked: false,
+          currentAttempts: 3,
+        })
+      );
+    });
+
+    it('reports currentAttempts as the tallied count when a new block is triggered', async () => {
+      mockIncr.mockResolvedValue(6); // exceeds maxAttempts of 5
+      redis.set = jest.fn().mockResolvedValue('OK');
+
+      rateLimit = new RateLimit(
+        { rules: parseConfigRules(['test:ip:5:1 minute:5 minutes:block']) },
+        redis,
+        statsd,
+        { write: mockWrite }
+      );
+
+      await rateLimit.check('test', { ip: '1.2.3.4' });
+
+      expect(mockWrite).toHaveBeenCalledWith(
+        expect.objectContaining({
+          wasBlocked: true,
+          currentAttempts: 6,
+        })
+      );
+    });
+
+    it('reports currentAttempts as maxAttempts when short-circuiting on a pre-existing block', async () => {
+      // A block already exists in Redis, so check() never tallies attempts.
+      // The stored record's own attempt count (99) must NOT be used; the
+      // reported value falls back to the rule's maxAttempts (5).
+      const existingBlock = {
+        action: 'test',
+        usedDefaultRule: false,
+        blockingOn: 'ip',
+        blockedValue: '1.2.3.4',
+        startTime: MOCK_NOW,
+        duration: 300,
+        attempts: 99,
+        reason: 'too-many-attempts',
+        policy: 'block',
+      };
+      // Bans are checked first via redis.get; return null for those, and the
+      // block record for the block key.
+      mockGet.mockImplementation((key: string) =>
+        Promise.resolve(
+          key.includes(':block:') ? JSON.stringify(existingBlock) : null
+        )
+      );
+
+      rateLimit = new RateLimit(
+        { rules: parseConfigRules(['test:ip:5:1 minute:5 minutes:block']) },
+        redis,
+        statsd,
+        { write: mockWrite }
+      );
+
+      await rateLimit.check('test', { ip: '1.2.3.4' });
+
+      expect(mockIncr).not.toHaveBeenCalled();
+      expect(mockWrite).toHaveBeenCalledWith(
+        expect.objectContaining({
+          wasBlocked: true,
+          currentAttempts: 5,
+        })
+      );
+    });
+
+    it('reports currentAttempts as undefined when a ban short-circuits the check', async () => {
+      // A ban has no associated rule (rules is empty), so the number of
+      // attempts cannot be determined.
+      const existingBan = {
+        action: 'test',
+        usedDefaultRule: false,
+        blockingOn: 'ip',
+        blockedValue: '1.2.3.4',
+        startTime: MOCK_NOW,
+        duration: 300,
+        attempts: 42,
+        reason: 'too-many-attempts',
+        policy: 'ban',
+      };
+      mockGet.mockImplementation((key: string) =>
+        Promise.resolve(
+          key.includes(':ban:') ? JSON.stringify(existingBan) : null
+        )
+      );
+
+      rateLimit = new RateLimit(
+        { rules: parseConfigRules(['test:ip:5:1 minute:5 minutes:ban']) },
+        redis,
+        statsd,
+        { write: mockWrite }
+      );
+
+      await rateLimit.check('test', { ip: '1.2.3.4' });
+
+      expect(mockIncr).not.toHaveBeenCalled();
+      expect(mockWrite).toHaveBeenCalledWith(
+        expect.objectContaining({
+          wasBlocked: true,
+          currentAttempts: undefined,
+        })
+      );
+    });
+
     it('calls bqWriter.write on skip() with wasSkipped true', () => {
       rateLimit = new RateLimit(
         { rules: {}, ignoreIPs: ['127.0.0.1'] },

libs/accounts/rate-limit/src/lib/rate-limit.ts

@@ -249,6 +249,7 @@ export class RateLimit {
     }
 
     const firstRule = rules[0];
+
     this.bqWriter.write({
       timestamp: now,
       action,
@@ -259,9 +260,24 @@ export class RateLimit {
       ruleMaxAttempts: firstRule?.rule.maxAttempts,
       ruleWindowSeconds: firstRule?.rule.windowDurationInSeconds,
       ruleBlockSeconds: firstRule?.rule.blockDurationInSeconds,
-      currentAttempts: attemptCounts?.get(
-        firstRule?.rule.blockingOn
-      ),
+      currentAttempts: (() => {
+        // When attempt counts were tallied on this check, they hold the real
+        // per-rule attempt total, so prefer them.
+        if (attemptCounts && firstRule) {
+          return attemptCounts.get(firstRule.rule.blockingOn);
+        }
+
+        // Without attempt counts we may be short-circuiting on a pre-existing
+        // block. In that case the user has, at minimum, reached maxAttempts, so
+        // report that as the best available approximation.
+        if (result != null && firstRule) {
+          return firstRule.rule.maxAttempts;
+        }
+
+        // Otherwise (e.g. a ban, where no rule applies) we can't say how many
+        // attempts occurred.
+        return undefined;
+      })(),
       wasBlocked: result != null,
       blockPolicy: result?.policy,
       blockDurationSeconds: result