fix: ensure ratelimit expiry is set every time (#556)

firebase · May 18, 2024 · 09cb208 · 09cb208
1 parent 4bdb0a6
commit 09cb208
Show file tree

Hide file tree

Showing 2 changed files with 63 additions and 6 deletions.
diff --git a/src/CachedKeySet.php b/src/CachedKeySet.php
@@ -212,15 +212,21 @@ private function rateLimitExceeded(): bool
         }
 
         $cacheItem = $this->cache->getItem($this->rateLimitCacheKey);
-        if (!$cacheItem->isHit()) {
-            $cacheItem->expiresAfter(60); // # of calls are cached each minute
+
+        $cacheItemData = [];
+        if ($cacheItem->isHit() && \is_array($data = $cacheItem->get())) {
+            $cacheItemData = $data;
         }
 
-        $callsPerMinute = (int) $cacheItem->get();
+        $callsPerMinute = $cacheItemData['callsPerMinute'] ?? 0;
+        $expiry = $cacheItemData['expiry'] ?? new \DateTime('+60 seconds', new \DateTimeZone('UTC'));
+
         if (++$callsPerMinute > $this->maxCallsPerMinute) {
             return true;
         }
-        $cacheItem->set($callsPerMinute);
+
+        $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
@@ -344,7 +344,7 @@ public function testRateLimit()
         $cachedKeySet = new CachedKeySet(
             $this->testJwksUri,
             $this->getMockHttpClient($this->testJwks1, $shouldBeCalledTimes),
-            $factory = $this->getMockHttpFactory($shouldBeCalledTimes),
+            $this->getMockHttpFactory($shouldBeCalledTimes),
             new TestMemoryCacheItemPool(),
             10,  // expires after seconds
             true // enable rate limiting
@@ -358,6 +358,54 @@ 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),
+            $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');
+        $cacheItem->set([
+            'expiry' => new \DateTime('+1 second', new \DateTimeZone('UTC')),
+            'callsPerMinute' => 0,
+        ]);
+        $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
      */
@@ -466,7 +514,10 @@ final class TestMemoryCacheItemPool implements CacheItemPoolInterface
 
     public function getItem($key): CacheItemInterface
     {
-        return current($this->getItems([$key]));
+        $item = current($this->getItems([$key]));
+        $item->expiresAt(null); // mimic symfony cache behavior
+
+        return $item;
     }
 
     public function getItems(array $keys = []): iterable