Rate limit per role (#3)

* prepare configuration * rate limit per user role * fix tests * add rate limit sort * use username for authenticated user as RateLimitExceededException message * fix test * fix service * add role based rate limit docs
IndraGunawan · Jul 1, 2017 · 4e19f93 · 4e19f93
1 parent 185c3a5
commit 4e19f93
Show file tree

Hide file tree

Showing 16 changed files with 550 additions and 66 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -4,4 +4,7 @@ CHANGELOG
 master
 ------
 
-* deprecated the `indragunawan_api_rate_limit.storage` configuration key.
+* Feature. Using PSR/Cache instead of DoctrineCache (#2).
+* Feature. Role based rate limit. (#3)
+* Deprecated the `indragunawan_api_rate_limit.storage` configuration key.
+* Deprecated the `indragunawan_api_rate_limit.throttle.period` and `indragunawan_api_rate_limit.throttle.limit` configuration key.
diff --git a/DependencyInjection/Configuration.php b/DependencyInjection/Configuration.php
@@ -51,10 +51,49 @@ public function getConfigTreeBuilder()
                     ->end()
                 ->end()
                 ->arrayNode('throttle')
+                    ->beforeNormalization()
+                        ->ifTrue(function ($v) { return is_array($v) && (isset($v['limit']) || isset($v['period'])); })
+                        ->then(function ($v) {
+                            $v['default'] = [];
+                            if (isset($v['limit'])) {
+                                @trigger_error('The indragunawan_api_rate_limit.throttle.limit configuration key is deprecated since version v0.2.0 and will be removed in v0.3.0. Use the indragunawan_api_rate_limit.throttle.default.limit configuration key instead.', E_USER_DEPRECATED);
+
+                                $v['default']['limit'] = $v['limit'];
+                            }
+
+                            if (isset($v['period'])) {
+                                @trigger_error('The indragunawan_api_rate_limit.throttle.period configuration key is deprecated since version v0.2.0 and will be removed in v0.3.0. Use the indragunawan_api_rate_limit.throttle.default.period configuration key instead.', E_USER_DEPRECATED);
+
+                                $v['default']['period'] = $v['period'];
+                            }
+
+                            return $v;
+                        })
+                    ->end()
                     ->addDefaultsIfNotSet()
                     ->children()
-                        ->integerNode('limit')->defaultValue(60)->end()
+                        ->integerNode('limit')->min(1)->defaultValue(60)->end()
                         ->integerNode('period')->min(1)->defaultValue(60)->end()
+                        ->arrayNode('default')
+                            ->addDefaultsIfNotSet()
+                            ->children()
+                                ->integerNode('limit')->min(1)->defaultValue(60)->end()
+                                ->integerNode('period')->min(1)->defaultValue(60)->end()
+                            ->end()
+                        ->end()
+                        ->arrayNode('roles')
+                            ->useAttributeAsKey('name')
+                            ->prototype('array')
+                                ->children()
+                                    ->integerNode('limit')->isRequired()->min(1)->end()
+                                    ->integerNode('period')->isRequired()->min(1)->end()
+                                ->end()
+                            ->end()
+                        ->end()
+                        ->enumNode('sort')
+                            ->values(['first-match', 'rate-limit-asc', 'rate-limit-desc'])
+                            ->defaultValue('rate-limit-desc')
+                        ->end()
                     ->end()
                 ->end()
                 ->arrayNode('exception')

diff --git a/DependencyInjection/IndragunawanApiRateLimitExtension.php b/DependencyInjection/IndragunawanApiRateLimitExtension.php
@@ -64,8 +64,18 @@ private function registerServiceConfig(ContainerBuilder $container, array $confi
             $cache = new Definition(FilesystemAdapter::class, ['api_rate_limit', 0, $container->getParameter('kernel.cache_dir')]);
         }
 
+        if ('rate-limit-asc' === $config['throttle']['sort']) {
+            uasort($config['throttle']['roles'], function (array $a, array $b) {
+                return ($a['limit'] / $a['period']) <=> ($b['limit'] / $b['period']);
+            });
+        } elseif ('rate-limit-desc' === $config['throttle']['sort']) {
+            uasort($config['throttle']['roles'], function (array $a, array $b) {
+                return ($b['limit'] / $b['period']) <=> ($a['limit'] / $a['period']);
+            });
+        }
+
         $container->getDefinition('indragunawan_api_rate_limit.service.rate_limit_handler')
             ->replaceArgument(0, $cache)
-            ->replaceArgument(1, $config['throttle']);
+            ->replaceArgument(3, $config['throttle']);
     }
 }
diff --git a/EventListener/RateLimitListener.php b/EventListener/RateLimitListener.php
@@ -15,6 +15,7 @@
 use Indragunawan\ApiRateLimitBundle\Service\RateLimitHandler;
 use Symfony\Component\HttpFoundation\Request;
 use Symfony\Component\HttpKernel\Event\GetResponseEvent;
+use Symfony\Component\Security\Core\Authentication\Token\Storage\TokenStorageInterface;
 
 /**
  * @author Indra Gunawan <hello@indra.my.id>
@@ -36,11 +37,17 @@ class RateLimitListener
      */
     private $exceptionConfig;
 
-    public function __construct(bool $enabled, RateLimitHandler $rateLimitHandler, array $exceptionConfig)
+    /**
+     * @var TokenStorageInterface
+     */
+    private $tokenStorage;
+
+    public function __construct(bool $enabled, RateLimitHandler $rateLimitHandler, array $exceptionConfig, TokenStorageInterface $tokenStorage)
     {
         $this->enabled = $enabled;
         $this->rateLimitHandler = $rateLimitHandler;
         $this->exceptionConfig = $exceptionConfig;
+        $this->tokenStorage = $tokenStorage;
     }
 
     public function onKernelRequest(GetResponseEvent $event)
@@ -81,7 +88,14 @@ protected function createRateLimitExceededException(Request $request)
     {
         $config = $this->exceptionConfig;
         $class = $config['custom_exception'] ?? RateLimitExceededException::class;
+        $username = null;
+
+        if (null !== $token = $this->tokenStorage->getToken()) {
+            if (is_object($token->getUser())) {
+                $username = $token->getUsername();
+            }
+        }
 
-        return new $class($config['status_code'], $config['message'], $request->getClientIp());
+        return new $class($config['status_code'], $config['message'], $request->getClientIp(), $username);
     }
 }
diff --git a/Exception/RateLimitExceededException.php b/Exception/RateLimitExceededException.php
@@ -18,9 +18,9 @@
  */
 class RateLimitExceededException extends HttpException
 {
-    public function __construct(int $statusCode, string $message, string $ip)
+    public function __construct(int $statusCode, string $message, string $ip, string $username = null)
     {
-        $message = sprintf($message, $ip);
+        $message = sprintf($message, $username ?: $ip);
 
         parent::__construct($statusCode, $message);
     }

diff --git a/README.md b/README.md
@@ -18,11 +18,11 @@ Documentation
     * [Disable rate limit per resource](Resources/doc/usage.md#disable-rate-limit-per-resource)
     * [Custom Exception](Resources/doc/usage.md#custom-exception)
     * [Custom Cache](Resources/doc/usage.md#custom-cache)
+    * [Role based rate limit](Resources/doc/usage.md#role-based-rate-limit)
 
 Todo
 ----
 
-* Rate limit per user ROLE
 
 License
 -------

diff --git a/Resources/config/event_listener.xml b/Resources/config/event_listener.xml
@@ -13,13 +13,16 @@
         <service id="indragunawan_api_rate_limit.event_listener.header_modification" class="%indragunawan_api_rate_limit.event_listener.header_modification.class%">
             <argument/>
             <argument/>
+
             <tag name="kernel.event_listener" event="kernel.response" method="onKernelResponse" priority="0" />
         </service>
 
         <service id="indragunawan_api_rate_limit.event_listener.rate_limit" class="%indragunawan_api_rate_limit.event_listener.rate_limit.class%">
             <argument/>
             <argument type="service" id="indragunawan_api_rate_limit.service.rate_limit_handler" />
             <argument/>
+            <argument type="service" id="security.token_storage" />
+
             <tag name="kernel.event_listener" event="kernel.request" method="onKernelRequest" priority="0" />
         </service>
     </services>

diff --git a/Resources/config/services.xml b/Resources/config/services.xml
@@ -15,6 +15,8 @@
     <services>
         <service id="indragunawan_api_rate_limit.service.rate_limit_handler" class="%indragunawan_api_rate_limit.service.rate_limit_handler.class%" public="false">
             <argument/>
+            <argument type="service" id="security.token_storage" />
+            <argument type="service" id="security.authorization_checker" />
             <argument/>
         </service>
     </services>

diff --git a/Resources/doc/configuration.md b/Resources/doc/configuration.md
@@ -22,10 +22,13 @@ indragunawan_api_rate_limit:
             remaining: X-RateLimit-Remaining
             reset: X-RateLimit-Reset
 
-    # Limit the request per period per IP address
+    # Limit the request per period per IP address / user
     throttle:
-        limit: 60 # max attempts per period
-        period: 60 # in seconds
+        default:
+            limit: 60 # max attempts per period
+            period: 60 # in seconds
+        roles: {}
+        sort: 'rate-limit-desc' # available value 'first-match', 'rate-limit-asc', 'rate-limit-desc'. default value rate-limit-desc
 
     # Exception thrown when rate limit exceeded
     exception:

diff --git a/Resources/doc/usage.md b/Resources/doc/usage.md
@@ -99,4 +99,29 @@ indragunawan_api_rate_limit:
 
 ---
 
+Role based rate limit
+---------------------
+
+By default, default throttle applies to all users (anonymous / authenticated user). If you want to apply rate limiting to all users who belong to the specified role, you can use user role based throttle configuration. If a user does not belong to your defined roles, the throttle will fall back to default.
+
+The `indragunawan_api_rate_limit.throttle.sort` configuration key is to sort only's the user role base throttle. The default value is `rate-limit-desc` that sorts high-low the request per second, other values are `rate-limit-asc` sorts low-high the request per second and `first-match` which uses your defined role throttle order.
+
+```yml
+indragunawan_api_rate_limit:
+    throttle:
+        default:
+            limit: 60 # max attempts per period
+            period: 60 # in seconds
+        roles:
+            ROLE_USER:
+                limit: 100
+                period: 60
+            ROLE_ADMIN:
+                limit: 1000
+                period: 60
+        sort: 'rate-limit-desc' # available value 'first-match', 'rate-limit-asc', 'rate-limit-desc'. default value 'rate-limit-desc'
+```
+
+---
+
 [Return to the index.](../../README.md)
diff --git a/Service/RateLimitHandler.php b/Service/RateLimitHandler.php
@@ -15,6 +15,9 @@
 use Indragunawan\ApiRateLimitBundle\Annotation\ApiRateLimit;
 use Psr\Cache\CacheItemPoolInterface;
 use Symfony\Component\HttpFoundation\Request;
+use Symfony\Component\Security\Core\Authentication\Token\Storage\TokenStorageInterface;
+use Symfony\Component\Security\Core\Authorization\AuthorizationCheckerInterface;
+use Symfony\Component\Security\Core\Exception\AuthenticationCredentialsNotFoundException;
 
 /**
  * @author Indra Gunawan <hello@indra.my.id>
@@ -26,6 +29,16 @@ class RateLimitHandler
      */
     private $cacheItemPool;
 
+    /**
+     * @var TokenStorageInterface
+     */
+    private $tokenStorage;
+
+    /**
+     * @var AuthorizationCheckerInterface
+     */
+    private $authorizationChecker;
+
     /**
      * @var array
      */
@@ -56,27 +69,26 @@ class RateLimitHandler
      */
     private $rateLimitExceeded = false;
 
-    public function __construct(CacheItemPoolInterface $cacheItemPool, array $throttleConfig)
-    {
+    public function __construct(
+        CacheItemPoolInterface $cacheItemPool,
+        TokenStorageInterface $tokenStorage,
+        AuthorizationCheckerInterface $authorizationChecker,
+        array $throttleConfig
+    ) {
         $this->cacheItemPool = $cacheItemPool;
+        $this->tokenStorage = $tokenStorage;
+        $this->authorizationChecker = $authorizationChecker;
         $this->throttleConfig = $throttleConfig;
     }
 
-    public function handle(Request $request)
+    public function isEnabled()
     {
-        $key = $this->getKey($request);
-        $limit = $this->throttleConfig['limit'];
-        $period = $this->throttleConfig['period'];
-
-        $annotationReader = new AnnotationReader();
-        $annotation = $annotationReader->getClassAnnotation(new \ReflectionClass($request->attributes->get('_api_resource_class')), ApiRateLimit::class);
-        if (null !== $annotation) {
-            $this->enabled = $annotation->enabled;
-        }
+        return $this->enabled;
+    }
 
-        if ($this->enabled) {
-            $this->decreaseRateLimitRemaining($key, $limit, $period);
-        }
+    public function isRateLimitExceeded()
+    {
+        return $this->rateLimitExceeded;
     }
 
     public function getRateLimitInfo(): array
@@ -88,14 +100,31 @@ public function getRateLimitInfo(): array
         ];
     }
 
-    protected function getKey(Request $request): string
+    public static function generateCacheKey(string $ip, string $userName = null, string $userRole = null): string
     {
-        return sprintf('api_rate_limit$%s', sha1($request->getClientIp()));
+        return sprintf('_api_rate_limit_metadata$%s', sha1($userName && $userRole ? sprintf('%s$%s', $userRole, $userName) : $ip));
     }
 
-    protected function decreaseRateLimitRemaining(string $key, int $limit, int $period, int $cost = 1)
+    public function handle(Request $request)
     {
+        list($key, $limit, $period) = $this->getThrottle($request);
+
+        $annotationReader = new AnnotationReader();
+        $annotation = $annotationReader->getClassAnnotation(new \ReflectionClass($request->attributes->get('_api_resource_class')), ApiRateLimit::class);
+        if (null !== $annotation) {
+            $this->enabled = $annotation->enabled;
+        }
+
+        if ($this->enabled) {
+            $this->decreaseRateLimitRemaining($key, $limit, $period);
+        }
+    }
+
+    protected function decreaseRateLimitRemaining(string $key, int $limit, int $period)
+    {
+        $cost = 1;
         $currentTime = gmdate('U');
+
         $rateLimitInfo = $this->cacheItemPool->getItem($key);
         $rateLimit = $rateLimitInfo->get();
         if ($rateLimitInfo->isHit() && $currentTime <= $rateLimit['reset']) {
@@ -135,13 +164,28 @@ protected function decreaseRateLimitRemaining(string $key, int $limit, int $peri
         $this->reset = $reset;
     }
 
-    public function isEnabled()
+    private function getThrottle(Request $request)
     {
-        return $this->enabled;
-    }
+        $userName = null;
+        $userRole = null;
+        $limit = $this->throttleConfig['default']['limit'];
+        $period = $this->throttleConfig['default']['period'];
+
+        foreach ($this->throttleConfig['roles'] as $role => $throttle) {
+            try {
+                if ($this->authorizationChecker->isGranted($role)) {
+                    $userName = $this->tokenStorage->getToken()->getUsername();
+                    $userRole = $role;
+                    $limit = $throttle['limit'];
+                    $period = $throttle['period'];
+
+                    break;
+                }
+            } catch (AuthenticationCredentialsNotFoundException $e) {
+                // do nothing
+            }
+        }
 
-    public function isRateLimitExceeded()
-    {
-        return $this->rateLimitExceeded;
+        return [self::generateCacheKey($request->getClientIp(), $userName, $userRole), $limit, $period];
     }
 }
diff --git a/Tests/DependencyInjection/ConfigurationTest.php b/Tests/DependencyInjection/ConfigurationTest.php
@@ -102,7 +102,6 @@ public function testExceptionClassNotSubclass()
 
     public function testValidExceptionClass()
     {
-        // var_dump(\Indragunawan\ApiRateLimitBundle\Exception\RateLimitExceededException::class); die();
         $config = $this->processor->processConfiguration(
             $this->configuration,
             [
@@ -116,4 +115,21 @@ public function testValidExceptionClass()
 
         $this->assertSame(ValidRateLimitExceededException::class, $config['exception']['custom_exception']);
     }
+
+    public function testDeprecatedConfiguration()
+    {
+        $config = $this->processor->processConfiguration(
+            $this->configuration,
+            [
+                [
+                    'throttle' => [
+                        'limit' => 10,
+                        'period' => 10,
+                    ],
+                ],
+            ]
+        );
+
+        $this->assertSame(['limit' => 10, 'period' => 10], $config['throttle']['default']);
+    }
 }