From 80a0aae69deeeeb22871f4ee3606622304ca4596 Mon Sep 17 00:00:00 2001
From: Massimo Redaelli <massimo@typish.io>
Date: Wed, 8 Apr 2020 08:28:32 +0200
Subject: [PATCH] Support timing out wait on IDLE

---
 src/extensions/idle.rs | 46 +++++++++++++++++++++++++++++++-----------
 1 file changed, 34 insertions(+), 12 deletions(-)

diff --git a/src/extensions/idle.rs b/src/extensions/idle.rs
index 12aa514..451a222 100644
--- a/src/extensions/idle.rs
+++ b/src/extensions/idle.rs
@@ -31,6 +31,15 @@ pub struct Handle<'a, T: Read + Write> {
     done: bool,
 }
 
+/// The result of a wait on a [`Handle`]
+#[derive(Debug, PartialEq, Eq)]
+pub enum WaitOutcome {
+    /// The wait timed out
+    TimedOut,
+    /// The mailbox was modified
+    MailboxChanged,
+}
+
 /// Must be implemented for a transport in order for a `Session` using that transport to support
 /// operations with timeouts.
 ///
@@ -91,34 +100,37 @@ impl<'a, T: Read + Write + 'a> Handle<'a, T> {
     /// Internal helper that doesn't consume self.
     ///
     /// This is necessary so that we can keep using the inner `Session` in `wait_keepalive`.
-    fn wait_inner(&mut self) -> Result<()> {
+    fn wait_inner(&mut self, reconnect: bool) -> Result<WaitOutcome> {
         let mut v = Vec::new();
         loop {
-            match self.session.readline(&mut v).map(|_| ()) {
+            let result = match self.session.readline(&mut v).map(|_| ()) {
                 Err(Error::Io(ref e))
                     if e.kind() == io::ErrorKind::TimedOut
                         || e.kind() == io::ErrorKind::WouldBlock =>
                 {
-                    // we need to refresh the IDLE connection
-                    self.terminate()?;
-                    self.init()?;
-                    return self.wait_inner();
+                    if reconnect {
+                        self.terminate()?;
+                        self.init()?;
+                        return self.wait_inner(reconnect);
+                    }
+                    Ok(WaitOutcome::TimedOut)
                 }
-                r => r,
+                Ok(()) => Ok(WaitOutcome::MailboxChanged),
+                Err(r) => Err(r),
             }?;
 
             // Handle Dovecot's imap_idle_notify_interval message
             if v.eq_ignore_ascii_case(b"* OK Still here\r\n") {
                 v.clear();
             } else {
-                break Ok(());
+                break Ok(result);
             }
         }
     }
 
     /// Block until the selected mailbox changes.
     pub fn wait(mut self) -> Result<()> {
-        self.wait_inner()
+        self.wait_inner(true).map(|_| ())
     }
 }
 
@@ -147,16 +159,26 @@ impl<'a, T: SetReadTimeout + Read + Write + 'a> Handle<'a, T> {
         // This still allows a client to receive immediate mailbox updates even
         // though it need only "poll" at half hour intervals.
         let keepalive = self.keepalive;
-        self.wait_timeout(keepalive)
+        self.timed_wait(keepalive, true).map(|_| ())
     }
 
     /// Block until the selected mailbox changes, or until the given amount of time has expired.
-    pub fn wait_timeout(mut self, timeout: Duration) -> Result<()> {
+    #[deprecated(note = "use wait_with_timeout instead")]
+    pub fn wait_timeout(self, timeout: Duration) -> Result<()> {
+        self.wait_with_timeout(timeout).map(|_| ())
+    }
+
+    /// Block until the selected mailbox changes, or until the given amount of time has expired.
+    pub fn wait_with_timeout(self, timeout: Duration) -> Result<WaitOutcome> {
+        self.timed_wait(timeout, false)
+    }
+
+    fn timed_wait(mut self, timeout: Duration, reconnect: bool) -> Result<WaitOutcome> {
         self.session
             .stream
             .get_mut()
             .set_read_timeout(Some(timeout))?;
-        let res = self.wait_inner();
+        let res = self.wait_inner(reconnect);
         let _ = self.session.stream.get_mut().set_read_timeout(None).is_ok();
         res
     }