fix: use flaps correctly

lib/Auth/MongoDBStorage.php

@@ -3,6 +3,7 @@
 namespace SimpleSAML\Module\privacyidea\Auth;
 
 use Exception;
+use MongoDB\BSON\UTCDateTime;
 use MongoDB\Client;
 use MongoDB\Collection;
 use BehEh\Flaps\StorageInterface;
@@ -33,7 +34,7 @@ class MongoDBStorage implements StorageInterface
     public function setValue($key, $value)
     {
         $this->collection->updateOne(
-            ['_id' => $key],
+            ['key' => $key],
             ['$set' => ['value' => $value]],
             ['upsert' => true]
         );
@@ -42,7 +43,7 @@ class MongoDBStorage implements StorageInterface
     public function incrementValue($key)
     {
         $result = $this->collection->findOneAndUpdate(
-            ['_id' => $key],
+            ['key' => $key],
             ['$inc' => ['value' => 1]],
             ['upsert' => true, 'returnDocument' => Collection::RETURN_DOCUMENT_AFTER]
         );
@@ -52,34 +53,50 @@ class MongoDBStorage implements StorageInterface
 
     public function getValue($key)
     {
-        $result = $this->collection->findOne(['_id' => $key]);
+        $result = $this->collection->findOne(['key' => $key]);
 
         $result = $this->collection->findOne(
-            ['$or' => ['$and' => ['_id' => $key, 'expireAt' =>
-                ['$gt' => gmdate('Y-m-d H:i:s')]]], ['$and' => ['_id' => $key, 'expire' => null]]]
+            ['$or' => [
+                [
+                    '$and' =>
+                        [
+                            ['key' => $key],
+                            ['expireAt' => ['$gt' => self::createMongoDate()]],
+                        ]
+                ],
+                [
+                    '$and' =>
+                    [
+                        ['key' => $key],
+                        ['expireAt' => null],
+                    ]
+                ]
+            ]]
         );
+
+        return $result ? $result['value'] : 0;
     }
 
     public function setTimestamp($key, $timestamp)
     {
         $this->collection->updateOne(
-            ['_id' => $key],
-            ['$set' => ['timestamp' => $timestamp]],
+            ['key' => $key],
+            ['$set' => ['timestamp' => new UTCDateTime($timestamp * 1000)]],
             ['upsert' => true]
         );
     }
 
     public function getTimestamp($key)
     {
-        $result = $this->collection->findOne(['_id' => $key]);
+        $result = $this->collection->findOne(['key' => $key]);
 
-        return $result ? $result['timestamp'] : 0;
+        return $result ? $result['timestamp']->toDateTime()->getTimestamp() : 0;
     }
 
     public function expire($key)
     {
-        $this->collection->deleteOne(['_id' => $key]);
-        $this->collection->deleteMany(['expireAt' => ['$lt' => gmdate('Y-m-d H:i:s')]]);
+        $this->collection->deleteOne(['key' => $key]);
+        $this->removeExpiredRecords();
     }
 
     /**
@@ -88,10 +105,22 @@ class MongoDBStorage implements StorageInterface
     public function expireIn($key, $seconds)
     {
         $this->collection->updateOne(
-            ['_id' => $key],
-            ['$set' => ['expireAt' => new DateTimeImmutable('+ ' . $seconds . ' seconds')]],
+            ['key' => $key],
+            ['$set' => ['expireAt' => self::createMongoDate($seconds)]],
             ['upsert' => true]
         );
-        $this->collection->deleteMany(['expireAt' => ['$lt' => gmdate('Y-m-d H:i:s')]]);
+        $this->removeExpiredRecords();
+    }
+
+    private function removeExpiredRecords()
+    {
+        $this->collection->deleteMany(['expireAt' => ['$lt' => self::createMongoDate()]]);
+    }
+
+    private static function createMongoDate($secondsFromNow = 0)
+    {
+        return new UTCDateTime(
+            (new \DateTimeImmutable('+' . $secondsFromNow . ' seconds'))->getTimestamp() * 1000
+        );
     }
 }

templates/LoginForm.php

@@ -314,4 +314,4 @@ if (!empty($this->data['links'])) {
 
 <?php
 $this->includeAtTemplateBase('includes/footer.php');
-?>
+?>
\ No newline at end of file

www/FormReceiver.php

@@ -3,6 +3,7 @@
 declare(strict_types=1);
 
 use BehEh\Flaps\Flaps;
+use BehEh\Flaps\Throttling\LeakyBucketStrategy;
 use BehEh\Flaps\Violation\PassiveViolationHandler;
 use SimpleSAML\Auth\Source;
 use SimpleSAML\Auth\State;
@@ -24,14 +25,16 @@ if (empty($stateId)) {
 $state = State::loadState($stateId, 'privacyidea:privacyidea');
 if (!empty($state['privacyidea:privacyidea']['rate_limiting'])) {
     $config = $state['privacyidea:privacyidea']['rate_limiting'];
-    $usernameStorage = new MongoDBStorage(
-        array_merge($config, ['collection_name' => $config['collection_prefix'] . 'username'])
+    $flapsStorage = new MongoDBStorage(
+        array_merge($config, ['collection_name' => $config['collection_prefix'] . 'flaps'])
     );
-    $ipStorage = new MongoDBStorage(array_merge($config, ['collection_name' => $config['collection_prefix'] . 'ip']));
-    $usernameFlaps = new Flaps($usernameStorage);
-    $ipFlaps = new Flaps($ipStorage);
-    $usernameFlaps->setViolationHandler(new PassiveViolationHandler());
-    $ipFlaps->setViolationHandler(new PassiveViolationHandler());
+    $flaps = new Flaps($flapsStorage);
+    $flaps->setDefaultViolationHandler(new PassiveViolationHandler());
+    $flapsUsername = $flaps->username;
+    $flapsUsername->pushThrottlingStrategy(new LeakyBucketStrategy(5, '60s'));
+    $flapsIP = $flaps->ip;
+    $flapsIP->pushThrottlingStrategy(new LeakyBucketStrategy(2, '1s'));
+
     $ip = $_SERVER['REMOTE_ADDR'];
 }
 
@@ -52,12 +55,12 @@ if (array_key_exists('username', $_REQUEST)) {
 
 if (!empty($state['privacyidea:privacyidea']['rate_limiting'])) {
     // Limit the requests based on UID and IP
-    if (!$usernameFlaps->login->limit($username)) {
+    if (!$flapsUsername->limit($username)) {
         Logger::warning('Rate limit exceeded for username ' . $username);
         $e = new \SimpleSAML\Error\Error('BADREQUEST', null, 429);
         throw new \SimpleSAML\Error\Exception('Rate limit exceeded', 429, $e);
     }
-    if (!$ipFlaps->login->limit($ip)) {
+    if (!$flapsIP->limit($ip)) {
         Logger::warning('Rate limit exceeded for IP address ' . $ip);
         $e = new \SimpleSAML\Error\Error('BADREQUEST', null, 429);
         throw new \SimpleSAML\Error\Exception('Rate limit exceeded', 429, $e);