perf: Use governor for rate limiting (#20)

Fixes: #19

Signed-off-by: Jagoda Ślązak <jslazak@jslazak.com>

filtermail/src/config.rs

@@ -1,6 +1,7 @@
 //! Configuration file handling for filtermail.
 
 use serde::{Deserialize, Deserializer};
+use std::num::NonZeroU32;
 use std::path::{Path, PathBuf};
 
 /// Chatmail configuration subset used by filtermail.
@@ -16,7 +17,7 @@ pub struct Config {
     pub postfix_reinject_port_incoming: u16,
     #[serde(default = "Config::default_max_message_size")]
     pub max_message_size: usize,
-    pub max_user_send_per_minute: usize,
+    pub max_user_send_per_minute: NonZeroU32,
     #[serde(default, deserialize_with = "deserialize_sequence")]
     pub passthrough_senders: Vec<String>,
     #[serde(default, deserialize_with = "deserialize_sequence")]

filtermail/src/main.rs

@@ -30,7 +30,6 @@ pub(crate) mod inbound;
 pub(crate) mod message;
 pub(crate) mod openpgp;
 pub(crate) mod outbound;
-pub(crate) mod rate_limiter;
 pub(crate) mod smtp_server;
 pub(crate) mod utils;
 

filtermail/src/outbound.rs

@@ -3,11 +3,11 @@
 use crate::ENCRYPTION_NEEDED_523;
 use crate::config::Config;
 use crate::message::{check_encrypted, is_securejoin, recipient_matches_passthrough};
-use crate::rate_limiter::SendRateLimiter;
 pub use crate::smtp_server::Envelope;
 use crate::smtp_server::SmtpHandler;
 use crate::utils::{extract_address, format_smtp_error};
 use async_trait::async_trait;
+use governor::{DefaultKeyedRateLimiter, Quota, RateLimiter};
 use lettre::{AsyncSmtpTransport, AsyncTransport, Tokio1Executor};
 use mailparse::{MailHeaderMap, parse_mail};
 use std::sync::{Arc, Mutex};
@@ -15,14 +15,15 @@ use std::sync::{Arc, Mutex};
 /// Handler for outgoing SMTP messages.
 pub struct OutgoingBeforeQueueHandler {
     config: Arc<Config>,
-    send_rate_limiter: Arc<Mutex<SendRateLimiter>>,
+    send_rate_limiter: Arc<Mutex<DefaultKeyedRateLimiter<String>>>,
 }
 
 impl OutgoingBeforeQueueHandler {
     pub fn new(config: Config) -> Self {
+        let quota = Quota::per_minute(config.max_user_send_per_minute);
         Self {
             config: Arc::new(config),
-            send_rate_limiter: Arc::new(Mutex::new(SendRateLimiter::default())),
+            send_rate_limiter: Arc::new(Mutex::new(RateLimiter::keyed(quota))),
         }
     }
 }
@@ -37,13 +38,25 @@ impl SmtpHandler for OutgoingBeforeQueueHandler {
             return Err(format!("500 Invalid from address <{}>", address));
         }
 
-        let max_sent = self.config.max_user_send_per_minute;
-        let mut limiter = self.send_rate_limiter.lock().unwrap();
-        if !limiter.is_sending_allowed(address, max_sent) {
-            log::debug!("Rate limit exceeded for {address}");
-            return Err(format!("450 4.7.1: Too much mail from {address}"));
+        let Ok(limiter) = self.send_rate_limiter.lock() else {
+            log::error!("send_rate_limiter lock panicked!");
+            return Err("451 Temporary server error".to_string());
+        };
+        if let Err(e) = limiter.check_key(&address.to_string()) {
+            // "<example@example.org> rate limited until: ..."
+            log::debug!("<{address}> {e}");
+            return Err(format!("450 4.7.1: Too much mail from <{address}>, {e}"));
         }
 
+        // Cleanup
+        //
+        // This is only called after a successful check,
+        // so a spam of mails from the same user will not cause calling this repeatedly.
+        // In the future, in case of higher traffic this can be further optimized by e.g. calling it
+        // every N messages or in a separate task every N minutes.
+        // Time complexity is O(n) where n is the number of unique senders in the last minute.
+        limiter.retain_recent();
+
         Ok(())
     }
 

filtermail/src/rate_limiter.rs

@@ -1,35 +0,0 @@
-//! Module for rate limiting.
-
-use std::collections::HashMap;
-use std::time::{Duration, SystemTime};
-
-const ONE_MINUTE: Duration = Duration::from_secs(60);
-
-/// A rate limiter tracking send timestamps per address.
-#[derive(Debug, Default)]
-pub struct SendRateLimiter {
-    address_timestamps: HashMap<String, Vec<SystemTime>>,
-}
-
-impl SendRateLimiter {
-    pub fn is_sending_allowed(&mut self, mail_from: &str, max_send_per_minute: usize) -> bool {
-        self.address_timestamps.retain(|_, timestamps| {
-            timestamps
-                .last()
-                .map(|t| t.elapsed().unwrap_or_default() <= ONE_MINUTE)
-                .unwrap_or(false)
-        });
-
-        let last = self
-            .address_timestamps
-            .entry(mail_from.to_string())
-            .or_default();
-        last.retain(|&send_time| send_time.elapsed().unwrap_or_default() <= ONE_MINUTE);
-        if last.len() <= max_send_per_minute {
-            last.push(SystemTime::now());
-            true
-        } else {
-            false
-        }
-    }
-}