From a7ecd0ed0823e657cf25b7c99a8ea49ce1127889 Mon Sep 17 00:00:00 2001
From: Kevin Bond <kevinbond@gmail.com>
Date: Wed, 30 Sep 2020 09:50:50 -0400
Subject: [PATCH] [RateLimiter] add Limit::ensureAccepted() and
 RateLimitExceededException

---
 .../Exception/RateLimitExceededException.php  | 46 +++++++++++++++++++
 src/Symfony/Component/RateLimiter/Limit.php   | 14 ++++++
 .../Component/RateLimiter/Tests/LimitTest.php | 43 +++++++++++++++++
 3 files changed, 103 insertions(+)
 create mode 100644 src/Symfony/Component/RateLimiter/Exception/RateLimitExceededException.php
 create mode 100644 src/Symfony/Component/RateLimiter/Tests/LimitTest.php

diff --git a/src/Symfony/Component/RateLimiter/Exception/RateLimitExceededException.php b/src/Symfony/Component/RateLimiter/Exception/RateLimitExceededException.php
new file mode 100644
index 0000000000..f96516e653
--- /dev/null
+++ b/src/Symfony/Component/RateLimiter/Exception/RateLimitExceededException.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\Exception;
+
+use Symfony\Component\RateLimiter\Limit;
+
+/**
+ * @author Kevin Bond <kevinbond@gmail.com>
+ *
+ * @experimental in 5.2
+ */
+class RateLimitExceededException extends \RuntimeException
+{
+    private $limit;
+
+    public function __construct(Limit $limit, $code = 0, \Throwable $previous = null)
+    {
+        parent::__construct('Rate Limit Exceeded', $code, $previous);
+
+        $this->limit = $limit;
+    }
+
+    public function getLimit(): Limit
+    {
+        return $this->limit;
+    }
+
+    public function getRetryAfter(): \DateTimeImmutable
+    {
+        return $this->limit->getRetryAfter();
+    }
+
+    public function getRemainingTokens(): int
+    {
+        return $this->limit->getRemainingTokens();
+    }
+}
diff --git a/src/Symfony/Component/RateLimiter/Limit.php b/src/Symfony/Component/RateLimiter/Limit.php
index 6042b6987b..fb59ceb127 100644
--- a/src/Symfony/Component/RateLimiter/Limit.php
+++ b/src/Symfony/Component/RateLimiter/Limit.php
@@ -11,6 +11,8 @@
 
 namespace Symfony\Component\RateLimiter;
 
+use Symfony\Component\RateLimiter\Exception\RateLimitExceededException;
+
 /**
  * @author Valentin Silvestre <vsilvestre.pro@gmail.com>
  *
@@ -34,6 +36,18 @@ class Limit
         return $this->accepted;
     }
 
+    /**
+     * @throws RateLimitExceededException if not accepted
+     */
+    public function ensureAccepted(): self
+    {
+        if (!$this->accepted) {
+            throw new RateLimitExceededException($this);
+        }
+
+        return $this;
+    }
+
     public function getRetryAfter(): \DateTimeImmutable
     {
         return $this->retryAfter;
diff --git a/src/Symfony/Component/RateLimiter/Tests/LimitTest.php b/src/Symfony/Component/RateLimiter/Tests/LimitTest.php
new file mode 100644
index 0000000000..a986581114
--- /dev/null
+++ b/src/Symfony/Component/RateLimiter/Tests/LimitTest.php
@@ -0,0 +1,43 @@
+<?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\Tests;
+
+use PHPUnit\Framework\TestCase;
+use Symfony\Component\RateLimiter\Exception\RateLimitExceededException;
+use Symfony\Component\RateLimiter\Limit;
+
+class LimitTest extends TestCase
+{
+    public function testEnsureAcceptedDoesNotThrowExceptionIfAccepted()
+    {
+        $limit = new Limit(10, new \DateTimeImmutable(), true);
+
+        $this->assertSame($limit, $limit->ensureAccepted());
+    }
+
+    public function testEnsureAcceptedThrowsRateLimitExceptionIfNotAccepted()
+    {
+        $limit = new Limit(10, $retryAfter = new \DateTimeImmutable(), false);
+
+        try {
+            $limit->ensureAccepted();
+        } catch (RateLimitExceededException $exception) {
+            $this->assertSame($limit, $exception->getLimit());
+            $this->assertSame(10, $exception->getRemainingTokens());
+            $this->assertSame($retryAfter, $exception->getRetryAfter());
+
+            return;
+        }
+
+        $this->fail('RateLimitExceededException not thrown.');
+    }
+}