RateLimiter.php 2.78 KB
Newer Older
Qiang Xue committed
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24
<?php
/**
 * @link http://www.yiiframework.com/
 * @copyright Copyright (c) 2008 Yii Software LLC
 * @license http://www.yiiframework.com/license/
 */

namespace yii\rest;

use yii\base\Component;
use yii\web\Request;
use yii\web\Response;
use yii\web\TooManyRequestsHttpException;

/**
 * RateLimiter implements a rate limiting algorithm based on the [leaky bucket algorithm](http://en.wikipedia.org/wiki/Leaky_bucket).
 *
 * You may call [[check()]] to enforce rate limiting.
 *
 * @author Qiang Xue <qiang.xue@gmail.com>
 * @since 2.0
 */
class RateLimiter extends Component
{
25 26 27 28 29 30 31 32
    /**
     * @var boolean whether to include rate limit headers in the response
     */
    public $enableRateLimitHeaders = true;
    /**
     * @var string the message to be displayed when rate limit exceeds
     */
    public $errorMessage = 'Rate limit exceeded.';
Qiang Xue committed
33

34 35
    /**
     * Checks whether the rate limit exceeds.
36 37 38
     * @param RateLimitInterface $user the current user
     * @param Request $request
     * @param Response $response
39
     * @param \yii\base\Action $action the action to be executed
40 41 42 43 44 45 46 47 48
     * @throws TooManyRequestsHttpException if rate limit exceeds
     */
    public function check($user, $request, $response, $action)
    {
        $current = time();
        $params = [
            'request' => $request,
            'action' => $action,
        ];
Qiang Xue committed
49

50 51
        list ($limit, $window) = $user->getRateLimit($params);
        list ($allowance, $timestamp) = $user->loadAllowance($params);
Qiang Xue committed
52

53 54 55 56
        $allowance += (int) (($current - $timestamp) * $limit / $window);
        if ($allowance > $limit) {
            $allowance = $limit;
        }
Qiang Xue committed
57

58 59 60 61 62 63 64 65 66
        if ($allowance < 1) {
            $user->saveAllowance(0, $current, $params);
            $this->addRateLimitHeaders($response, $limit, 0, $window);
            throw new TooManyRequestsHttpException($this->errorMessage);
        } else {
            $user->saveAllowance($allowance - 1, $current, $params);
            $this->addRateLimitHeaders($response, $limit, 0, (int) (($limit - $allowance) * $window / $limit));
        }
    }
Qiang Xue committed
67

68 69 70
    /**
     * Adds the rate limit headers to the response
     * @param Response $response
71 72 73
     * @param integer $limit the maximum number of allowed requests during a period
     * @param integer $remaining the remaining number of allowed requests within the current period
     * @param integer $reset the number of seconds to wait before having maximum number of allowed requests again
74 75 76 77 78 79 80 81 82 83
     */
    protected function addRateLimitHeaders($response, $limit, $remaining, $reset)
    {
        if ($this->enableRateLimitHeaders) {
            $response->getHeaders()
                ->set('X-Rate-Limit-Limit', $limit)
                ->set('X-Rate-Limit-Remaining', $remaining)
                ->set('X-Rate-Limit-Reset', $reset);
        }
    }
Qiang Xue committed
84
}