From 820171bf8f9785499cc12ad4c98362002819c9b9 Mon Sep 17 00:00:00 2001
From: link2xt <link2xt@testrun.org>
Date: Mon, 13 Nov 2023 15:00:47 +0000
Subject: [PATCH] feat(imap): detect NOTIFY extension

---
 src/config.rs            |  6 ++++++
 src/context.rs           |  2 ++
 src/imap.rs              | 28 ++++++++++++++++++++++++++--
 src/imap/capabilities.rs |  4 ++++
 src/imap/client.rs       |  1 +
 src/imap/session.rs      |  8 ++++++++
 6 files changed, 47 insertions(+), 2 deletions(-)

diff --git a/src/config.rs b/src/config.rs
index 53fed8338..d56acb985 100644
--- a/src/config.rs
+++ b/src/config.rs
@@ -296,6 +296,12 @@ pub enum Config {
     #[strum(props(default = "0"))]
     DisableIdle,
 
+    /// Whether to avoid using IMAP NOTIFY even if the server supports it.
+    ///
+    /// This is a developer option for testing prefetch without NOTIFY.
+    #[strum(props(default = "0"))]
+    DisableNotify,
+
     /// Defines the max. size (in bytes) of messages downloaded automatically.
     /// 0 = no limit.
     #[strum(props(default = "0"))]
diff --git a/src/context.rs b/src/context.rs
index 1e04c618b..30125911e 100644
--- a/src/context.rs
+++ b/src/context.rs
@@ -595,6 +595,7 @@ impl Context {
         let bcc_self = self.get_config_int(Config::BccSelf).await?;
         let sync_msgs = self.get_config_int(Config::SyncMsgs).await?;
         let disable_idle = self.get_config_bool(Config::DisableIdle).await?;
+        let disable_notify = self.get_config_bool(Config::DisableNotify).await?;
 
         let prv_key_cnt = self.sql.count("SELECT COUNT(*) FROM keypairs;", ()).await?;
 
@@ -708,6 +709,7 @@ impl Context {
         res.insert("bcc_self", bcc_self.to_string());
         res.insert("sync_msgs", sync_msgs.to_string());
         res.insert("disable_idle", disable_idle.to_string());
+        res.insert("disable_notify", disable_notify.to_string());
         res.insert("private_key_count", prv_key_cnt.to_string());
         res.insert("public_key_count", pub_key_cnt.to_string());
         res.insert("fingerprint", fingerprint_str);
diff --git a/src/imap.rs b/src/imap.rs
index 5a6479313..fbda1e410 100644
--- a/src/imap.rs
+++ b/src/imap.rs
@@ -448,6 +448,22 @@ impl Imap {
         self.session = None;
     }
 
+    /// Tries to setup NOTIFY.
+    pub async fn setup_notify(&mut self, context: &Context) -> Result<()> {
+        let session = self
+            .session
+            .as_mut()
+            .context("no IMAP connection established")?;
+        if session.can_notify() && !session.notify_set {
+            let cmd = format!("NOTIFY SET (selected (Messagenew {PREFETCH_FLAGS} messageexpunge))");
+            session.run_command_and_check_ok(cmd).await?;
+            info!(context, "Enabled NOTIFY");
+            session.notify_set = true;
+        }
+
+        Ok(())
+    }
+
     /// FETCH-MOVE-DELETE iteration.
     ///
     /// Prefetches headers and downloads new message from the folder, moves messages away from the
@@ -720,7 +736,9 @@ impl Imap {
                 .await
                 .context("prefetch_existing_msgs")?
         } else {
-            self.prefetch(old_uid_next).await.context("prefetch")?
+            self.prefetch(context, old_uid_next)
+                .await
+                .context("prefetch")?
         };
         let read_cnt = msgs.len();
 
@@ -1329,7 +1347,13 @@ impl Imap {
 
     /// Prefetch all messages greater than or equal to `uid_next`. Returns a list of fetch results
     /// in the order of ascending delivery time to the server (INTERNALDATE).
-    async fn prefetch(&mut self, uid_next: u32) -> Result<Vec<(u32, async_imap::types::Fetch)>> {
+    async fn prefetch(
+        &mut self,
+        context: &Context,
+        uid_next: u32,
+    ) -> Result<Vec<(u32, async_imap::types::Fetch)>> {
+        self.setup_notify(context).await?;
+
         let session = self
             .session
             .as_mut()
diff --git a/src/imap/capabilities.rs b/src/imap/capabilities.rs
index 14385c088..ebcb93afe 100644
--- a/src/imap/capabilities.rs
+++ b/src/imap/capabilities.rs
@@ -9,6 +9,10 @@ pub(crate) struct Capabilities {
     /// <https://tools.ietf.org/html/rfc2177>
     pub can_idle: bool,
 
+    /// True if the server has NOTIFY capability as defined in
+    /// <https://tools.ietf.org/html/rfc5465>
+    pub can_notify: bool,
+
     /// True if the server has MOVE capability as defined in
     /// <https://tools.ietf.org/html/rfc6851>
     pub can_move: bool,
diff --git a/src/imap/client.rs b/src/imap/client.rs
index b052a5471..3a077863d 100644
--- a/src/imap/client.rs
+++ b/src/imap/client.rs
@@ -56,6 +56,7 @@ async fn determine_capabilities(
     };
     let capabilities = Capabilities {
         can_idle: caps.has_str("IDLE"),
+        can_notify: caps.has_str("NOTIFY"),
         can_move: caps.has_str("MOVE"),
         can_check_quota: caps.has_str("QUOTA"),
         can_condstore: caps.has_str("CONDSTORE"),
diff --git a/src/imap/session.rs b/src/imap/session.rs
index ced292591..5021a756d 100644
--- a/src/imap/session.rs
+++ b/src/imap/session.rs
@@ -19,6 +19,9 @@ pub(crate) struct Session {
     pub selected_mailbox: Option<Mailbox>,
 
     pub selected_folder_needs_expunge: bool,
+
+    /// True if NOTIFY SET command was executed in this session.
+    pub notify_set: bool,
 }
 
 impl Deref for Session {
@@ -46,6 +49,7 @@ impl Session {
             selected_folder: None,
             selected_mailbox: None,
             selected_folder_needs_expunge: false,
+            notify_set: false,
         }
     }
 
@@ -53,6 +57,10 @@ impl Session {
         self.capabilities.can_idle
     }
 
+    pub fn can_notify(&self) -> bool {
+        self.capabilities.can_notify
+    }
+
     pub fn can_move(&self) -> bool {
         self.capabilities.can_move
     }