fix: does not show correct token time until 1 token is available

Fixes #5458

Author: kenjis

system/Throttle/Throttler.php

@@ -95,13 +95,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 = max(1, (int) $refresh);
+
             return true;
         }
 
@@ -110,15 +118,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 +127,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 = max(1, (int) ($refresh - $elapsed));
+
             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

@@ -45,6 +45,42 @@ public function testTokenTime()
         $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(100, $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(97, $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);