Fail2ban and IP address blocking support (closes #164)

crates/imap/src/core/client.rs

@@ -28,7 +28,6 @@ use imap_proto::{
     Command, ResponseCode, StatusResponse,
 };
 use jmap::auth::rate_limit::AuthenticatedLimiter;
-use parking_lot::Mutex;
 use utils::listener::{
     limiter::{ConcurrencyLimiter, RateLimiter},
     SessionStream,
@@ -231,9 +230,8 @@ impl<T: SessionStream> Session<T> {
             if !data
                 .imap
                 .get_authenticated_limiter(data.account_id)
-                .lock()
                 .request_limiter
-                .is_allowed()
+                .is_allowed(&self.imap.rate_requests)
             {
                 return Err(StatusResponse::no("Too many requests")
                     .with_tag(request.tag)
@@ -392,19 +390,16 @@ impl<T: SessionStream> State<T> {
 }
 
 impl IMAP {
-    pub fn get_authenticated_limiter(&self, account_id: u32) -> Arc<Mutex<AuthenticatedLimiter>> {
+    pub fn get_authenticated_limiter(&self, account_id: u32) -> Arc<AuthenticatedLimiter> {
         self.rate_limiter
             .get(&account_id)
             .map(|limiter| limiter.clone())
             .unwrap_or_else(|| {
-                let limiter = Arc::new(Mutex::new(AuthenticatedLimiter {
-                    request_limiter: RateLimiter::new(
-                        self.rate_requests.requests,
-                        self.rate_requests.period,
-                    ),
+                let limiter = Arc::new(AuthenticatedLimiter {
+                    request_limiter: RateLimiter::new(&self.rate_requests),
                     concurrent_requests: ConcurrencyLimiter::new(self.rate_concurrent),
                     concurrent_uploads: ConcurrencyLimiter::new(self.rate_concurrent),
-                }));
+                });
                 self.rate_limiter.insert(account_id, limiter.clone());
                 limiter
             })

crates/imap/src/core/mod.rs

@@ -23,6 +23,7 @@
 
 use std::{
     collections::BTreeMap,
+    net::IpAddr,
     sync::{atomic::AtomicU32, Arc},
     time::Duration,
 };
@@ -35,10 +36,7 @@ use imap_proto::{
     Command, ResponseCode, StatusResponse,
 };
 use jmap::{
-    auth::{
-        rate_limit::{AuthenticatedLimiter, RemoteAddress},
-        AccessToken,
-    },
+    auth::{rate_limit::AuthenticatedLimiter, AccessToken},
     JMAP,
 };
 use store::roaring::RoaringBitmap;
@@ -82,7 +80,7 @@ pub struct IMAP {
     pub greeting_plain: Vec<u8>,
     pub greeting_tls: Vec<u8>,
 
-    pub rate_limiter: DashMap<u32, Arc<parking_lot::Mutex<AuthenticatedLimiter>>>,
+    pub rate_limiter: DashMap<u32, Arc<AuthenticatedLimiter>>,
     pub rate_requests: Rate,
     pub rate_concurrent: u64,
 }
@@ -100,7 +98,7 @@ pub struct Session<T: SessionStream> {
     pub stream_rx: ReadHalf<T>,
     pub stream_tx: Arc<tokio::sync::Mutex<WriteHalf<T>>>,
     pub in_flight: InFlight,
-    pub remote_addr: RemoteAddress,
+    pub remote_addr: IpAddr,
     pub span: tracing::Span,
 }
 

crates/imap/src/core/session.rs

@@ -24,7 +24,6 @@
 use std::{borrow::Cow, sync::Arc};
 
 use imap_proto::{protocol::ProtocolVersion, receiver::Receiver};
-use jmap::auth::rate_limit::RemoteAddress;
 use tokio::io::{AsyncReadExt, AsyncWriteExt};
 use tokio_rustls::server::TlsStream;
 use utils::listener::{stream::NullIo, SessionManager, SessionStream};
@@ -139,7 +138,7 @@ impl<T: SessionStream> Session<T> {
             instance: session.instance,
             span: session.span,
             in_flight: session.in_flight,
-            remote_addr: RemoteAddress::IpAddress(session.remote_ip),
+            remote_addr: session.remote_ip,
             stream_rx,
             stream_tx: Arc::new(tokio::sync::Mutex::new(stream_tx)),
         })

crates/imap/src/op/authenticate.rs

@@ -23,6 +23,7 @@
 
 use std::sync::Arc;
 
+use directory::AuthResult;
 use imap_proto::{
     protocol::{authenticate::Mechanism, capability::Capability},
     receiver::{self, Request},
@@ -116,9 +117,15 @@ impl<T: SessionStream> Session<T> {
         // Authenticate
         let access_token = match credentials {
             Credentials::Plain { username, secret } | Credentials::XOauth2 { username, secret } => {
-                self.jmap
-                    .authenticate_plain(&username, &secret, &self.remote_addr)
+                match self
+                    .jmap
+                    .authenticate_plain(&username, &secret, self.remote_addr)
                     .await
+                {
+                    AuthResult::Success(token) => Some(token),
+                    AuthResult::Failure => None,
+                    AuthResult::Banned => return Err(()),
+                }
             }
             Credentials::OAuthBearer { token } => {
                 match self
@@ -145,7 +152,6 @@ impl<T: SessionStream> Session<T> {
             let in_flight = self
                 .imap
                 .get_authenticated_limiter(access_token.primary_id())
-                .lock()
                 .concurrent_requests
                 .is_allowed();
             if let Some(in_flight) = in_flight {