From 9926804f1b25bcd2a2164f58fd5d0ce1c61ac71e Mon Sep 17 00:00:00 2001
From: link2xt <link2xt@testrun.org>
Date: Sun, 10 Jul 2022 14:08:54 +0000
Subject: [PATCH] ratelimit: do not overflow leaky bucket

This way the time to wait until next message can
be sent is always limited.
---
 CHANGELOG.md     |  1 +
 src/ratelimit.rs | 17 ++++++-----------
 2 files changed, 7 insertions(+), 11 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 55555622d..342cdaef7 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -4,6 +4,7 @@
 
 ### Changes
 - handle drafts from mailto links in scanned QR #3492
+- do not overflow ratelimiter leaky bucket #3496
 
 ### Fixes
 - don't squash text parts of NDN into attachments #3497
diff --git a/src/ratelimit.rs b/src/ratelimit.rs
index 8ba987964..64edc1338 100644
--- a/src/ratelimit.rs
+++ b/src/ratelimit.rs
@@ -51,7 +51,7 @@ impl Ratelimit {
 
     /// Returns true if it is allowed to send a message.
     fn can_send_at(&self, now: SystemTime) -> bool {
-        self.current_value_at(now) <= self.quota
+        self.current_value_at(now) + 1.0 <= self.quota
     }
 
     /// Returns true if can send another message now.
@@ -62,7 +62,7 @@ impl Ratelimit {
     }
 
     fn send_at(&mut self, now: SystemTime) {
-        self.current_value = self.current_value_at(now) + 1.0;
+        self.current_value = f64::min(self.quota, self.current_value_at(now) + 1.0);
         self.last_update = now;
     }
 
@@ -77,10 +77,10 @@ impl Ratelimit {
 
     fn until_can_send_at(&self, now: SystemTime) -> Duration {
         let current_value = self.current_value_at(now);
-        if current_value <= self.quota {
+        if current_value + 1.0 <= self.quota {
             Duration::ZERO
         } else {
-            let requirement = current_value - self.quota;
+            let requirement = current_value + 1.0 - self.quota;
             let rate = self.quota / self.window.as_secs_f64();
             Duration::from_secs_f64(requirement / rate)
         }
@@ -109,8 +109,6 @@ mod tests {
         ratelimit.send_at(now);
         assert!(ratelimit.can_send_at(now));
         ratelimit.send_at(now);
-        assert!(ratelimit.can_send_at(now));
-        ratelimit.send_at(now);
 
         // Can't send more messages now.
         assert!(!ratelimit.can_send_at(now));
@@ -125,11 +123,8 @@ mod tests {
         // Send one more message anyway, over quota.
         ratelimit.send_at(now);
 
-        // Waiting 20 seconds is not enough.
-        let now = now + Duration::from_secs(20);
-        assert!(!ratelimit.can_send_at(now));
-
-        // Can send another message after 40 seconds.
+        // Always can send another message after 20 seconds,
+        // leaky bucket never overflows.
         let now = now + Duration::from_secs(20);
         assert!(ratelimit.can_send_at(now));