Merge pull request #5470 from kenjis/fix-Throttle

fix: Throttler does not show correct token time

Author: kenjis

system/Throttle/Throttler.php

@@ -79,10 +79,9 @@ public function getTokenTime(): int
      *
      * Example:
      *
-     *  if (! $throttler->check($request->ipAddress(), 60, MINUTE))
-     * {
+     *  if (! $throttler->check($request->ipAddress(), 60, MINUTE)) {
      *      die('You submitted over 60 requests within a minute.');
-     * }
+     *  }
      *
      * @param string $key      The name to use as the "bucket" name.
      * @param int    $capacity The number of requests the "bucket" can hold
@@ -95,13 +94,21 @@ public function check(string $key, int $capacity, int $seconds, int $cost = 1):
     {
         $tokenName = $this->prefix . $key;
 
+        // Number of tokens to add back per second
+        $rate = $capacity / $seconds;
+        // Number of seconds to get one token
+        $refresh = 1 / $rate;
+
         // Check to see if the bucket has even been created yet.
         if (($tokens = $this->cache->get($tokenName)) === null) {
             // If it hasn't been created, then we'll set it to the maximum
             // capacity - 1, and save it to the cache.
-            $this->cache->save($tokenName, $capacity - $cost, $seconds);
+            $tokens = $capacity - $cost;
+            $this->cache->save($tokenName, $tokens, $seconds);
             $this->cache->save($tokenName . 'Time', $this->time(), $seconds);
 
+            $this->tokenTime = 0;
+
             return true;
         }
 
@@ -110,15 +117,6 @@ public function check(string $key, int $capacity, int $seconds, int $cost = 1):
         $throttleTime = $this->cache->get($tokenName . 'Time');
         $elapsed      = $this->time() - $throttleTime;
 
-        // Number of tokens to add back per second
-        $rate = $capacity / $seconds;
-
-        // How many seconds till a new token is available.
-        // We must have a minimum wait of 1 second for a new token.
-        // Primarily stored to allow devs to report back to users.
-        $newTokenAvailable = (1 / $rate) - $elapsed;
-        $this->tokenTime   = max(1, $newTokenAvailable);
-
         // Add tokens based up on number per second that
         // should be refilled, then checked against capacity
         // to be sure the bucket didn't overflow.
@@ -128,12 +126,21 @@ public function check(string $key, int $capacity, int $seconds, int $cost = 1):
         // If $tokens >= 1, then we are safe to perform the action, but
         // we need to decrement the number of available tokens.
         if ($tokens >= 1) {
-            $this->cache->save($tokenName, $tokens - $cost, $seconds);
+            $tokens = $tokens - $cost;
+            $this->cache->save($tokenName, $tokens, $seconds);
             $this->cache->save($tokenName . 'Time', $this->time(), $seconds);
 
+            $this->tokenTime = 0;
+
             return true;
         }
 
+        // How many seconds till a new token is available.
+        // We must have a minimum wait of 1 second for a new token.
+        // Primarily stored to allow devs to report back to users.
+        $newTokenAvailable = (int) ($refresh - $elapsed - $refresh * $tokens);
+        $this->tokenTime   = max(1, $newTokenAvailable);
+
         return false;
     }
 

tests/system/Throttle/ThrottleTest.php

@@ -36,15 +36,51 @@ public function testTokenTime()
         // set $rate
         $rate = 1;    // allow 1 request per minute
 
-        // first check just creates a bucket, so tokenTime should be 0
+        // When the first check you have a token, so tokenTime should be 0
         $throttler->check('127.0.0.1', $rate, MINUTE);
         $this->assertSame(0, $throttler->getTokenTime());
 
-        // additional check affects tokenTime, so tokenTime should be 1 or greater
+        // When additional check you don't have one token, so tokenTime should be 1 or greater
         $throttler->check('127.0.0.1', $rate, MINUTE);
         $this->assertGreaterThanOrEqual(1, $throttler->getTokenTime());
     }
 
+    /**
+     * @see https://github.com/codeigniter4/CodeIgniter4/issues/5458
+     */
+    public function testTokenTimeCalculation()
+    {
+        $time = 1639441295;
+
+        $throttler = new Throttler($this->cache);
+        $throttler->setTestTime($time);
+
+        $capacity = 2;
+        $seconds  = 200;
+
+        // refresh = 200 / 2 = 100 seconds
+        // refresh rate = 2 / 200 = 0.01 token per second
+
+        // token should be 2
+        $this->assertTrue($throttler->check('test', $capacity, $seconds));
+        // token should be 2 - 1 = 1
+        $this->assertSame(0, $throttler->getTokenTime(), 'Wrong token time');
+
+        // do nothing for 3 seconds
+        $throttler = $throttler->setTestTime($time + 3);
+
+        // token should be 1 + 3 * 0.01 = 1.03
+        $this->assertTrue($throttler->check('test', $capacity, $seconds));
+        // token should be 1.03 - 1 = 0.03
+        $this->assertSame(0, $throttler->getTokenTime(), 'Wrong token time');
+
+        $this->assertFalse($throttler->check('test', $capacity, $seconds));
+        // token should still be 0.03 because check failed
+
+        // expect remaining time: (1 - 0.03) * 100 = 97
+        $this->assertSame(97, $throttler->getTokenTime(), 'Wrong token time');
+    }
+
     public function testIPSavesBucket()
     {
         $throttler = new Throttler($this->cache);