symfony · fabpot · Sep 30, 2020 · Sep 27, 2020
diff --git a/src/Symfony/Component/RateLimiter/CompoundLimiter.php b/src/Symfony/Component/RateLimiter/CompoundLimiter.php
@@ -25,17 +25,28 @@ final class CompoundLimiter implements LimiterInterface
      */
     public function __construct(array $limiters)
     {
+        if (!$limiters) {
+            throw new \LogicException(sprintf('"%s::%s()" require at least one limiter.', self::class, __METHOD__));
+        }
         $this->limiters = $limiters;
     }
 
-    public function consume(int $tokens = 1): bool
+    public function consume(int $tokens = 1): Limit
     {
-        $allow = true;
+        $minimalLimit = null;
         foreach ($this->limiters as $limiter) {
-            $allow = $limiter->consume($tokens) && $allow;
+            $limit = $limiter->consume($tokens);
+
+            if (0 === $limit->getRemainingTokens()) {
+                return $limit;
+            }
+
+            if (null === $minimalLimit || $limit->getRemainingTokens() < $minimalLimit->getRemainingTokens()) {
+                $minimalLimit = $limit;
+            }
         }
 
-        return $allow;
+        return $minimalLimit;
     }
 
     public function reset(): void

diff --git a/src/Symfony/Component/RateLimiter/FixedWindowLimiter.php b/src/Symfony/Component/RateLimiter/FixedWindowLimiter.php
@@ -43,7 +43,7 @@ public function __construct(string $id, int $limit, \DateInterval $interval, Sto
     /**
      * {@inheritdoc}
      */
-    public function consume(int $tokens = 1): bool
+    public function consume(int $tokens = 1): Limit
     {
         $this->lock->acquire(true);
 
@@ -54,17 +54,28 @@ public function consume(int $tokens = 1): bool
             }
 
             $hitCount = $window->getHitCount();
-            $availableTokens = $this->limit - $hitCount;
+            $availableTokens = $this->getAvailableTokens($hitCount);
+            $windowStart = \DateTimeImmutable::createFromFormat('U', time());
             if ($availableTokens < $tokens) {
-                return false;
+                return new Limit($availableTokens, $this->getRetryAfter($windowStart), false);
             }
 
             $window->add($tokens);
             $this->storage->save($window);
 
-            return true;
+            return new Limit($this->getAvailableTokens($window->getHitCount()), $this->getRetryAfter($windowStart), true);
         } finally {
             $this->lock->release();
         }
     }
+
+    public function getAvailableTokens(int $hitCount): int
+    {
+        return $this->limit - $hitCount;
+    }
+
+    private function getRetryAfter(\DateTimeImmutable $windowStart): \DateTimeImmutable
+    {
+        return $windowStart->add(new \DateInterval(sprintf('PT%sS', $this->interval)));
+    }
 }
diff --git a/src/Symfony/Component/RateLimiter/Limit.php b/src/Symfony/Component/RateLimiter/Limit.php
@@ -0,0 +1,46 @@
+<?php
+
+/*
+ * This file is part of the Symfony package.
+ *
+ * (c) Fabien Potencier <fabien@symfony.com>
+ *
+ * For the full copyright and license information, please view the LICENSE
+ * file that was distributed with this source code.
+ */
+
+namespace Symfony\Component\RateLimiter;
+
+/**
+ * @author Valentin Silvestre <vsilvestre.pro@gmail.com>
+ *
+ * @experimental in 5.2
+ */
+class Limit
+{
+    private $availableTokens;
+    private $retryAfter;
+    private $accepted;
+
+    public function __construct(int $availableTokens, \DateTimeImmutable $retryAfter, bool $accepted)
+    {
+        $this->availableTokens = $availableTokens;
+        $this->retryAfter = $retryAfter;
+        $this->accepted = $accepted;
+    }
+
+    public function isAccepted(): bool
+    {
+        return $this->accepted;
+    }
+
+    public function getRetryAfter(): \DateTimeImmutable
+    {
+        return $this->retryAfter;
+    }
+
+    public function getRemainingTokens(): int
+    {
+        return $this->availableTokens;
+    }
+}
diff --git a/src/Symfony/Component/RateLimiter/LimiterInterface.php b/src/Symfony/Component/RateLimiter/LimiterInterface.php
@@ -24,7 +24,7 @@ interface LimiterInterface
      *
      * @param int $tokens the number of tokens required
      */
-    public function consume(int $tokens = 1): bool;
+    public function consume(int $tokens = 1): Limit;
 
     /**
      * Resets the limit.

diff --git a/src/Symfony/Component/RateLimiter/NoLimiter.php b/src/Symfony/Component/RateLimiter/NoLimiter.php
@@ -23,9 +23,9 @@
  */
 final class NoLimiter implements LimiterInterface
 {
-    public function consume(int $tokens = 1): bool
+    public function consume(int $tokens = 1): Limit
     {
-        return true;
+        return new Limit(\INF, new \DateTimeImmutable(), true, 'no_limit');
     }
 
     public function reset(): void

diff --git a/src/Symfony/Component/RateLimiter/README.md b/src/Symfony/Component/RateLimiter/README.md
@@ -32,7 +32,7 @@ $limiter->reserve(1)->wait();
 // ... execute the code
 
 // only claims 1 token if it's free at this moment (useful if you plan to skip this process)
-if ($limiter->consume(1)) {
+if ($limiter->consume(1)->isAccepted()) {
    // ... execute the code
 }
 ```

diff --git a/src/Symfony/Component/RateLimiter/Rate.php b/src/Symfony/Component/RateLimiter/Rate.php
@@ -73,6 +73,16 @@ public function calculateTimeForTokens(int $tokens): int
         return TimeUtil::dateIntervalToSeconds($this->refillTime) * $cyclesRequired;
     }
 
+    /**
+     * Calculates the next moment of token availability.
+     *
+     * @return \DateTimeImmutable the next moment a token will be available
+     */
+    public function calculateNextTokenAvailability(): \DateTimeImmutable
+    {
+        return (new \DateTimeImmutable())->add($this->refillTime);
+    }
+
     /**
      * Calculates the number of new free tokens during $duration.
      *

diff --git a/src/Symfony/Component/RateLimiter/Tests/CompoundLimiterTest.php b/src/Symfony/Component/RateLimiter/Tests/CompoundLimiterTest.php
@@ -38,19 +38,20 @@ public function testConsume()
         $limiter3 = $this->createLimiter(12, new \DateInterval('PT30S'));
         $limiter = new CompoundLimiter([$limiter1, $limiter2, $limiter3]);
 
-        $this->assertFalse($limiter->consume(5), 'Limiter 1 reached the limit');
+        // Reach limiter 1 limit, verify that limiter2 available tokens reduced by 5 and fetch successfully limiter 1
+        $this->assertEquals(3, $limiter->consume(5)->getRemainingTokens(), 'Limiter 1 reached the limit');
         sleep(1); // reset limiter1's window
-        $limiter->consume(2);
+        $this->assertTrue($limiter->consume(2)->isAccepted());
 
-        $this->assertTrue($limiter->consume());
-        $this->assertFalse($limiter->consume(), 'Limiter 2 reached the limit');
+        // Reach limiter 2 limit, verify that limiter2 available tokens reduced by 5 and and fetch successfully
+        $this->assertEquals(0, $limiter->consume()->getRemainingTokens(), 'Limiter 2 has no remaining tokens left');
         sleep(9); // reset limiter2's window
+        $this->assertTrue($limiter->consume(3)->isAccepted());
 
-        $this->assertTrue($limiter->consume(3));
-        $this->assertFalse($limiter->consume(), 'Limiter 3 reached the limit');
+        // Reach limiter 3 limit, verify that limiter2 available tokens reduced by 5 and fetch successfully
+        $this->assertEquals(0, $limiter->consume()->getRemainingTokens(), 'Limiter 3 reached the limit');
         sleep(20); // reset limiter3's window
-
-        $this->assertTrue($limiter->consume());
+        $this->assertTrue($limiter->consume()->isAccepted());
     }
 
     private function createLimiter(int $limit, \DateInterval $interval): FixedWindowLimiter

diff --git a/src/Symfony/Component/RateLimiter/Tests/FixedWindowLimiterTest.php b/src/Symfony/Component/RateLimiter/Tests/FixedWindowLimiterTest.php
@@ -40,8 +40,10 @@ public function testConsume()
             sleep(5);
         }
 
-        $this->assertTrue($limiter->consume());
-        $this->assertFalse($limiter->consume());
+        $limit = $limiter->consume();
+        $this->assertTrue($limit->isAccepted());
+        $limit = $limiter->consume();
+        $this->assertFalse($limit->isAccepted());
     }
 
     public function testConsumeOutsideInterval()
@@ -55,7 +57,9 @@ public function testConsumeOutsideInterval()
         $limiter->consume(9);
         // ...try bursting again at the start of the next window
         sleep(10);
-        $this->assertTrue($limiter->consume(10));
+        $limit = $limiter->consume(10);
+        $this->assertEquals(0, $limit->getRemainingTokens());
+        $this->assertEquals(time() + 60, $limit->getRetryAfter()->getTimestamp());
     }
 
     private function createLimiter(): FixedWindowLimiter

diff --git a/src/Symfony/Component/RateLimiter/Tests/TokenBucketLimiterTest.php b/src/Symfony/Component/RateLimiter/Tests/TokenBucketLimiterTest.php
@@ -69,13 +69,21 @@ public function testReserveMaxWaitingTime()
 
     public function testConsume()
     {
-        $limiter = $this->createLimiter();
+        $rate = Rate::perSecond(10);
+        $limiter = $this->createLimiter(10, $rate);
 
         // enough free tokens
-        $this->assertTrue($limiter->consume(5));
+        $limit = $limiter->consume(5);
+        $this->assertTrue($limit->isAccepted());
+        $this->assertEquals(5, $limit->getRemainingTokens());
+        $this->assertEqualsWithDelta(time(), $limit->getRetryAfter()->getTimestamp(), 1);
         // there are only 5 available free tokens left now
-        $this->assertFalse($limiter->consume(10));
-        $this->assertTrue($limiter->consume(5));
+        $limit = $limiter->consume(10);
+        $this->assertEquals(5, $limit->getRemainingTokens());
+
+        $limit = $limiter->consume(5);
+        $this->assertEquals(0, $limit->getRemainingTokens());
+        $this->assertEqualsWithDelta(time(), $limit->getRetryAfter()->getTimestamp(), 1);
     }
 
     private function createLimiter($initialTokens = 10, Rate $rate = null)

diff --git a/src/Symfony/Component/RateLimiter/TokenBucketLimiter.php b/src/Symfony/Component/RateLimiter/TokenBucketLimiter.php
@@ -103,14 +103,20 @@ public function reserve(int $tokens = 1, ?float $maxTime = null): Reservation
     /**
      * {@inheritdoc}
      */
-    public function consume(int $tokens = 1): bool
+    public function consume(int $tokens = 1): Limit
     {
+        $bucket = $this->storage->fetch($this->id);
+        if (null === $bucket) {
+            $bucket = new TokenBucket($this->id, $this->maxBurst, $this->rate);
+        }
+        $now = microtime(true);
+
         try {
             $this->reserve($tokens, 0);
 
-            return true;
+            return new Limit($bucket->getAvailableTokens($now) - $tokens, $this->rate->calculateNextTokenAvailability(), true);
         } catch (MaxWaitDurationExceededException $e) {
-            return false;
+            return new Limit($bucket->getAvailableTokens($now), $this->rate->calculateNextTokenAvailability(), false);
         }
     }
 }
diff --git a/src/Symfony/Component/Security/Http/EventListener/LoginThrottlingListener.php b/src/Symfony/Component/Security/Http/EventListener/LoginThrottlingListener.php
@@ -48,7 +48,7 @@ public function checkPassport(CheckPassportEvent $event): void
         $limiterKey = $this->createLimiterKey($username, $request);
 
         $limiter = $this->limiter->create($limiterKey);
-        if (!$limiter->consume()) {
+        if (!$limiter->consume()->isAccepted()) {
             throw new TooManyLoginAttemptsAuthenticationException();
         }
     }