From ddff59a20ac86ddc300901c287fc834ce3f80eb7 Mon Sep 17 00:00:00 2001
From: William Casarin <jb55@jb55.com>
Date: Mon, 8 Jul 2024 14:12:38 -0700
Subject: [PATCH] ratelimit: switch to token-based rate limiting

allows for N actions per minute

Changelog-Changed: Switched to token-based ratelimiting
Closes: https://github.com/damus-io/noteguard/issues/2
---
 README.md                 |  4 +--
 noteguard.toml            |  4 +--
 src/filters/rate_limit.rs | 67 +++++++++++++++++++++++++--------------
 test/test-delayed         |  9 ++++++
 4 files changed, 57 insertions(+), 27 deletions(-)
 create mode 100755 test/test-delayed

diff --git a/README.md b/README.md
index b899651..5da4c47 100644
--- a/README.md
+++ b/README.md
@@ -18,7 +18,7 @@ The `pipeline` config specifies the order in which filters are run. When the fir
 pipeline = ["ratelimit"]
 
 [filters.ratelimit]
-delay_seconds = 1
+notes_per_minute = 8
 whitelist = ["127.0.0.1"]
 ```
 
@@ -34,7 +34,7 @@ The ratelimit filter limits the rate at which notes are written to the relay per
 
 Settings:
 
-- `delay_seconds`: the delay in seconds between accepted notes. 1 means only one note can be written per second. 2 means only 1 note can be written every 2 seconds, etc.
+- `notes_per_minute`: the number of notes per minute which are allowed to be written per ip.
 
 - `whitelist`: a list of IP4 or IP6 addresses that are allowed to bypass the ratelimit.
 
diff --git a/noteguard.toml b/noteguard.toml
index df49009..02bd9a7 100644
--- a/noteguard.toml
+++ b/noteguard.toml
@@ -2,5 +2,5 @@
 pipeline = ["ratelimit"]
 
 [filters.ratelimit]
-delay_seconds = 1
-whitelist = ["127.0.0.1"]
+posts_per_minute = 10
+whitelist = ["127.0.0.10"]
diff --git a/src/filters/rate_limit.rs b/src/filters/rate_limit.rs
index d4b398b..4279f24 100644
--- a/src/filters/rate_limit.rs
+++ b/src/filters/rate_limit.rs
@@ -3,20 +3,25 @@ use serde::Deserialize;
 use std::collections::HashMap;
 use std::time::{Duration, Instant};
 
-pub struct RateInfo {
-    pub last_note: Instant,
+pub struct Tokens {
+    pub tokens: i32,
+    pub last_post: Instant,
 }
 
 #[derive(Deserialize, Default)]
 pub struct RateLimit {
-    pub delay_seconds: u64,
+    pub posts_per_minute: i32,
     pub whitelist: Option<Vec<String>>,
 
     #[serde(skip)]
-    pub sources: HashMap<String, RateInfo>,
+    pub sources: HashMap<String, Tokens>,
 }
 
 impl NoteFilter for RateLimit {
+    fn name(&self) -> &'static str {
+        "ratelimit"
+    }
+
     fn filter_note(&mut self, msg: &InputMessage) -> OutputMessage {
         if let Some(whitelist) = &self.whitelist {
             if whitelist.contains(&msg.source_info) {
@@ -24,31 +29,47 @@ impl NoteFilter for RateLimit {
             }
         }
 
-        if self.sources.contains_key(&msg.source_info) {
-            let now = Instant::now();
-            let entry = self.sources.get_mut(&msg.source_info).expect("impossiburu");
-            if now - entry.last_note < Duration::from_secs(self.delay_seconds) {
-                return OutputMessage::new(
-                    msg.event.id.clone(),
-                    Action::Reject,
-                    Some("rate-limited: you are noting too fast".to_string()),
-                );
-            } else {
-                entry.last_note = Instant::now();
-                return OutputMessage::new(msg.event.id.clone(), Action::Accept, None);
-            }
-        } else {
+        if !self.sources.contains_key(&msg.source_info) {
             self.sources.insert(
                 msg.source_info.to_owned(),
-                RateInfo {
-                    last_note: Instant::now(),
+                Tokens {
+                    last_post: Instant::now(),
+                    tokens: self.posts_per_minute,
                 },
             );
             return OutputMessage::new(msg.event.id.clone(), Action::Accept, None);
         }
-    }
 
-    fn name(&self) -> &'static str {
-        "ratelimit"
+        let entry = self.sources.get_mut(&msg.source_info).expect("impossiburu");
+        let now = Instant::now();
+        let mut diff = now - entry.last_post;
+
+        let min = Duration::from_secs(60);
+        if diff > min {
+            diff = min;
+        }
+
+        let percent = (diff.as_secs() as f32) / 60.0;
+        let new_tokens = (percent * self.posts_per_minute as f32).floor() as i32;
+        entry.tokens += new_tokens - 1;
+
+        if entry.tokens <= 0 {
+            entry.tokens = 0;
+        }
+
+        if entry.tokens >= self.posts_per_minute {
+            entry.tokens = self.posts_per_minute - 1;
+        }
+
+        if entry.tokens == 0 {
+            return OutputMessage::new(
+                msg.event.id.clone(),
+                Action::Reject,
+                Some("rate-limited: you are noting too much".to_string()),
+            );
+        }
+
+        entry.last_post = now;
+        OutputMessage::new(msg.event.id.clone(), Action::Accept, None)
     }
 }
diff --git a/test/test-delayed b/test/test-delayed
new file mode 100755
index 0000000..845e1c6
--- /dev/null
+++ b/test/test-delayed
@@ -0,0 +1,9 @@
+#!/usr/bin/env sh
+
+while true
+do
+	echo '{"type": "new","receivedAt":12345,"sourceType":"IP4","sourceInfo": "127.0.0.2","event":{"id": "68421a122cef086512b2c5bd29ca6285ced8bd8e302e347e3c5d90466c860a76","pubkey": "16c21558762108afc34e4ff19e4ed51d9a48f79e0c34531efc423d21ab435e93","created_at": 1720408658,"kind": 1,"tags": [],"content": "hi","sig": "7b76471744ded2b720ca832cdc89e670f6093ce38aeef55a5c6a4e077883d7d80dda1e9051032fb1faa1c3c212c517e93ee42b3ceac8e8e9b04bad46a361de90"}}'
+
+	sleep 0.1
+done
+