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

@@ -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.
 

noteguard.toml

@@ -2,5 +2,5 @@
 pipeline = ["ratelimit"]
 
 [filters.ratelimit]
-delay_seconds = 1
+posts_per_minute = 10
-whitelist = ["127.0.0.1"]
+whitelist = ["127.0.0.10"]

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 struct Tokens {
-    pub last_note: Instant,
+    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) {
+        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 {
             self.sources.insert(
                 msg.source_info.to_owned(),
-                RateInfo {
+                Tokens {
-                    last_note: Instant::now(),
+                    last_post: Instant::now(),
+                    tokens: self.posts_per_minute,
                 },
             );
             return OutputMessage::new(msg.event.id.clone(), Action::Accept, None);
         }
+
+        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;
         }
 
-    fn name(&self) -> &'static str {
+        let percent = (diff.as_secs() as f32) / 60.0;
-        "ratelimit"
+        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)
     }
 }

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
+