review comment, add test

firebase · May 18, 2024 · ec55a4a · ec55a4a
1 parent f80d7a1
commit ec55a4a
Show file tree

Hide file tree

Showing 2 changed files with 56 additions and 8 deletions.
diff --git a/src/CachedKeySet.php b/src/CachedKeySet.php
@@ -213,19 +213,20 @@ private function rateLimitExceeded(): bool
 
         $cacheItem = $this->cache->getItem($this->rateLimitCacheKey);
 
-        if (!$cacheItem->isHit() || !\is_array($cacheItemData = $cacheItem->get())) {
-            $cacheItemData = [
-                'callsPerMinute' => 0,
-                'expiry' => new \DateTime('+60 seconds', new \DateTimeZone('UTC')),
-            ];
+        $cacheItemData = [];
+        if ($cacheItem->isHit() && \is_array($data = $cacheItem->get())) {
+            $cacheItemData = $data;
         }
 
-        if (++$cacheItemData['callsPerMinute'] > $this->maxCallsPerMinute) {
+        $callsPerMinute = $cacheItemData['callsPerMinute'] ?? 0;
+        $expiry = $cacheItemData['expiry'] ?? new \DateTime('+60 seconds', new \DateTimeZone('UTC'));
+
+        if (++$callsPerMinute > $this->maxCallsPerMinute) {
             return true;
         }
 
-        $cacheItem->set($cacheItemData);
-        $cacheItem->expiresAt($cacheItemData['expiry']);
+        $cacheItem->set(['expiry' => $expiry, 'callsPerMinute' => $callsPerMinute]);
+        $cacheItem->expiresAt($expiry);
         $this->cache->save($cacheItem);
         return false;
     }

diff --git a/tests/CachedKeySetTest.php b/tests/CachedKeySetTest.php
@@ -358,6 +358,53 @@ public function testRateLimit()
         $this->assertFalse(isset($cachedKeySet[$invalidKid]));
     }
 
+    public function testRateLimitWithExpiresAfter()
+    {
+        // We request the key 17 times, HTTP should only be called 15 times
+        $shouldBeCalledTimes = 10;
+        $cachedTimes = 2;
+        $afterExpirationTimes = 5;
+
+        $totalHttpTimes = $shouldBeCalledTimes + $afterExpirationTimes;
+
+        $cachePool = new TestMemoryCacheItemPool();
+
+        // Instantiate the cached key set
+        $cachedKeySet = new CachedKeySet(
+            $this->testJwksUri,
+            $this->getMockHttpClient($this->testJwks1, $totalHttpTimes),
+            $factory = $this->getMockHttpFactory($totalHttpTimes),
+            $cachePool,
+            10,   // expires after seconds
+            true // enable rate limiting
+        );
+
+        // Set the rate limit cache to expire after 1 second
+        $cacheItem = $cachePool->getItem('jwksratelimitjwkshttpsjwk.uri');
+        $cacheItemData = $cacheItem->get();
+        $cacheItemData['expiry'] = new \DateTime('+1 seconds', new \DateTimeZone('UTC'));
+        $cacheItem->set($cacheItemData);
+        $cacheItem->expiresAfter(1);
+        $cachePool->save($cacheItem);
+
+        $invalidKid = 'invalidkey';
+        for ($i = 0; $i < $shouldBeCalledTimes; $i++) {
+            $this->assertFalse(isset($cachedKeySet[$invalidKid]));
+        }
+
+        // The next calls do not call HTTP
+        for ($i = 0; $i < $cachedTimes; $i++) {
+            $this->assertFalse(isset($cachedKeySet[$invalidKid]));
+        }
+
+        sleep(1); // wait for cache to expire
+
+        // These calls DO call HTTP because the cache has expired
+        for ($i = 0; $i < $afterExpirationTimes; $i++) {
+            $this->assertFalse(isset($cachedKeySet[$invalidKid]));
+        }
+    }
+
     /**
      * @dataProvider provideFullIntegration
      */