From ccec26ffa712a2c313bda15f4583989c65e72eb9 Mon Sep 17 00:00:00 2001
From: link2xt <link2xt@testrun.org>
Date: Thu, 21 Dec 2023 01:08:43 +0000
Subject: [PATCH] fix(imap): limit the rate of LOGIN attempts rather than
 connection attempts

As ratelimit was introduced to avoid reconnecting immediately after disconnecting
in case of bugs in IMAP protocol handling,
connection attempts should only be counted when IMAP is actually used,
i.e. when the first command (LOGIN) is sent.
---
 src/imap.rs | 29 ++++++++++++++++++++++++-----
 1 file changed, 24 insertions(+), 5 deletions(-)

diff --git a/src/imap.rs b/src/imap.rs
index 526224ef9..02a74c46a 100644
--- a/src/imap.rs
+++ b/src/imap.rs
@@ -94,6 +94,20 @@ pub struct Imap {
     login_failed_once: bool,
 
     pub(crate) connectivity: ConnectivityStore,
+
+    /// Rate limit for IMAP connection usage attempts.
+    ///
+    /// Rate limit is checked before connecting
+    /// and updated right before login attempt.
+    /// It does not limit the number of connection attempts
+    /// if the network is bad as only successful connections are counted.
+    ///
+    /// Main purpose of this rate limit is
+    /// to prevent busy loop in case
+    /// connection gets dropped over and over due to IMAP bug,
+    /// e.g. the server returning invalid response to SELECT command
+    /// immediately after logging in or returning an error in response to LOGIN command
+    /// due to internal server error.
     ratelimit: RwLock<Ratelimit>,
 }
 
@@ -256,7 +270,7 @@ impl Imap {
             session: None,
             login_failed_once: false,
             connectivity: Default::default(),
-            // 1 connection per minute + a burst of 2.
+            // 1 login per minute + a burst of 2.
             ratelimit: RwLock::new(Ratelimit::new(Duration::new(120, 0), 2.0)),
         };
 
@@ -306,6 +320,12 @@ impl Imap {
         }
 
         self.connectivity.set_connecting(context).await;
+
+        // Check rate limit before trying to connect
+        // to avoid connecting and not using the connection
+        // in case we are currently ratelimited.
+        // Otherwise connection may become unusable due to NAT forgetting about it
+        // before we attempt to actually login.
         let ratelimit_duration = self.ratelimit.read().await.until_can_send();
         if !ratelimit_duration.is_zero() {
             warn!(
@@ -316,10 +336,7 @@ impl Imap {
             tokio::time::sleep(ratelimit_duration).await;
         }
 
-        let oauth2 = self.config.lp.oauth2;
-
         info!(context, "Connecting to IMAP server");
-        self.ratelimit.write().await.send();
         let connection_res: Result<Client> = if self.config.lp.security == Socket::Starttls
             || self.config.lp.security == Socket::Plain
         {
@@ -369,11 +386,13 @@ impl Imap {
                 Client::connect_secure(context, imap_server, imap_port, config.strict_tls).await
             }
         };
-
         let client = connection_res?;
+        self.ratelimit.write().await.send();
+
         let config = &self.config;
         let imap_user: &str = config.lp.user.as_ref();
         let imap_pw: &str = config.lp.password.as_ref();
+        let oauth2 = self.config.lp.oauth2;
 
         let login_res = if oauth2 {
             info!(context, "Logging into IMAP server with OAuth 2");