crates/imap/src/core/client.rs

@@ -36,18 +36,10 @@ impl<T: SessionStream> Session<T> {
 
         loop {
             match self.receiver.parse(&mut bytes) {
-                Ok(request) => match self.is_allowed(request).await {
-                    Ok(request) => {
-                        has_expunge |=
-                            matches!(request.command, Command::Expunge(_) | Command::Close);
-                        requests.push(request);
-                    }
-                    Err(err) => {
-                        if !self.write_error(err).await {
-                            return SessionResult::Close;
-                        }
-                    }
-                },
+                Ok(request) => {
+                    has_expunge |= matches!(request.command, Command::Expunge(_) | Command::Close);
+                    requests.push(request);
+                }
                 Err(receiver::Error::NeedsMoreData) => {
                     break;
                 }
@@ -95,6 +87,18 @@ impl<T: SessionStream> Session<T> {
 
         let mut requests = requests.into_iter().peekable();
         while let Some(request) = requests.next() {
+            // inbuxa: each command is checked against the state it runs in,
+            // not the one the buffer was read in, so a SELECT pipelined with
+            // the FETCH or STORE after it lets them run (RFC 9051 5.5)
+            let request = match self.is_allowed(request).await {
+                Ok(request) => request,
+                Err(err) => {
+                    if !self.write_error(err).await {
+                        return SessionResult::Close;
+                    }
+                    continue;
+                }
+            };
             // inbuxa: ST-6: commands that only read may be served by a read
             // replica; any write they make still goes to the primary
             let replica_accounts = match (&request.command, &self.state) {
@@ -131,22 +135,31 @@ impl<T: SessionStream> Session<T> {
                         .handle_select(request)
                         .await
                         .map(|_| SessionResult::Continue),
-                    Command::Create => self
-                        .handle_create(group_requests(&mut requests, vec![request]))
-                        .await
-                        .map(|_| SessionResult::Continue),
-                    Command::Delete => self
-                        .handle_delete(group_requests(&mut requests, vec![request]))
-                        .await
-                        .map(|_| SessionResult::Continue),
+                    Command::Create => match self.group_allowed(&mut requests, request).await {
+                        Some(requests) => self
+                            .handle_create(requests)
+                            .await
+                            .map(|_| SessionResult::Continue),
+                        None => Ok(SessionResult::Close),
+                    },
+                    Command::Delete => match self.group_allowed(&mut requests, request).await {
+                        Some(requests) => self
+                            .handle_delete(requests)
+                            .await
+                            .map(|_| SessionResult::Continue),
+                        None => Ok(SessionResult::Close),
+                    },
                     Command::Rename => self
                         .handle_rename(request)
                         .await
                         .map(|_| SessionResult::Continue),
-                    Command::Status => self
-                        .handle_status(group_requests(&mut requests, vec![request]))
-                        .await
-                        .map(|_| SessionResult::Continue),
+                    Command::Status => match self.group_allowed(&mut requests, request).await {
+                        Some(requests) => self
+                            .handle_status(requests)
+                            .await
+                            .map(|_| SessionResult::Continue),
+                        None => Ok(SessionResult::Close),
+                    },
                     Command::Append => self
                         .handle_append(request)
                         .await
@@ -167,10 +180,13 @@ impl<T: SessionStream> Session<T> {
                         .handle_search(request, false, is_uid)
                         .await
                         .map(|_| SessionResult::Continue),
-                    Command::Fetch(_) => self
-                        .handle_fetch(group_requests(&mut requests, vec![request]))
-                        .await
-                        .map(|_| SessionResult::Continue),
+                    Command::Fetch(_) => match self.group_allowed(&mut requests, request).await {
+                        Some(requests) => self
+                            .handle_fetch(requests)
+                            .await
+                            .map(|_| SessionResult::Continue),
+                        None => Ok(SessionResult::Close),
+                    },
                     Command::Store(is_uid) => self
                         .handle_store(request, is_uid, !has_expunge)
                         .await
@@ -290,7 +306,11 @@ impl<T: SessionStream> Session<T> {
             let result = if replica_accounts.is_empty() {
                 dispatch.await
             } else {
-                store::backend::scaleout::replica::replica_read(replica_accounts, dispatch).await
+                Box::pin(store::backend::scaleout::replica::replica_read(
+                    replica_accounts,
+                    dispatch,
+                ))
+                .await
             };
 
             drop(inline_permit);
@@ -349,23 +369,31 @@ fn takes_own_permit(command: &Command) -> bool {
     )
 }
 
-pub fn group_requests(
-    requests: &mut Peekable<IntoIter<Request<Command>>>,
-    mut grouped_requests: Vec<Request<Command>>,
-) -> Vec<Request<Command>> {
-    let last_command = grouped_requests.last().unwrap().command;
-    loop {
-        match requests.peek() {
-            Some(request) if request.command == last_command => {
-                grouped_requests.push(requests.next().unwrap());
-            }
-            _ => break,
-        }
-    }
-    grouped_requests
-}
-
 impl<T: SessionStream> Session<T> {
+    // inbuxa: the commands of one kind that follow a command, run with it as
+    // one batch; each is checked (and counted against the rate limit) on its
+    // own, and one that isn't allowed is answered and left out. None when
+    // that answer ends the session.
+    async fn group_allowed(
+        &self,
+        requests: &mut Peekable<IntoIter<Request<Command>>>,
+        first: Request<Command>,
+    ) -> Option<Vec<Request<Command>>> {
+        let command = first.command;
+        let mut grouped = vec![first];
+        while let Some(request) = requests.next_if(|request| request.command == command) {
+            match self.is_allowed(request).await {
+                Ok(request) => grouped.push(request),
+                Err(err) => {
+                    if !self.write_error(err).await {
+                        return None;
+                    }
+                }
+            }
+        }
+        Some(grouped)
+    }
+
     async fn is_allowed(&self, request: Request<Command>) -> trc::Result<Request<Command>> {
         let state = &self.state;
         // Rate limit request

tests/src/imap/pipelining.rs

@@ -82,15 +82,13 @@ pub async fn test(test: &TestServer) {
         ));
     }
 
-    // STORE, then FETCH of the same messages, then DELETE, in one write.
-    // SELECT goes first on its own: the commands after it are checked
-    // against the session state as they are read.
-    tx.write_all(b"q1 SELECT \"Pipelined\"\r\n").await.unwrap();
-    let result = read_tags(&mut rx, &["q1"]).await;
-    result.check_in_order(&mut problems, "SELECT", &["q1"]);
-    let tags = ["q2", "q3", "q4", "q5"];
+    // SELECT, then STORE and FETCH of the same messages, then UNSELECT and
+    // DELETE, in one write: each command is checked against the state it
+    // runs in, so the ones after SELECT find a mailbox selected
+    let tags = ["q1", "q2", "q3", "q4", "q5"];
     tx.write_all(
         concat!(
+            "q1 SELECT \"Pipelined\"\r\n",
             "q2 STORE 1:* +FLAGS.SILENT (\\Flagged)\r\n",
             "q3 FETCH 1:* (FLAGS)\r\n",
             "q4 UNSELECT\r\n",
@@ -101,7 +99,7 @@ pub async fn test(test: &TestServer) {
     .await
     .unwrap();
     let result = read_tags(&mut rx, &tags).await;
-    result.check_in_order(&mut problems, "STORE/FETCH/UNSELECT/DELETE", &tags);
+    result.check_in_order(&mut problems, "SELECT/STORE/FETCH/UNSELECT/DELETE", &tags);
     let fetched = result.untagged("FETCH (");
     if fetched.len() != MESSAGES || fetched.iter().any(|line| !line.contains("\\Flagged")) {
         problems.push(format!(
@@ -109,6 +107,25 @@ pub async fn test(test: &TestServer) {
         ));
     }
 
+    // and the other way round: a FETCH after UNSELECT in the same write
+    // finds no mailbox selected
+    let tags = ["r1", "r2", "r3"];
+    tx.write_all(b"r1 SELECT \"INBOX\"\r\nr2 UNSELECT\r\nr3 FETCH 1:* (FLAGS)\r\n")
+        .await
+        .unwrap();
+    let result = read_tags(&mut rx, &tags).await;
+    let outcome = result
+        .tagged
+        .iter()
+        .map(|line| line.splitn(3, ' ').take(2).collect::<Vec<_>>().join(" "))
+        .collect::<Vec<_>>();
+    if outcome != ["r1 OK", "r2 OK", "r3 BAD"] {
+        problems.push(format!(
+            "SELECT/UNSELECT/FETCH: expected OK, OK, BAD in order: {:?}",
+            result.tagged
+        ));
+    }
+
     tx.write_all(b"z LOGOUT\r\n").await.unwrap();
 
     assert!(problems.is_empty(), "{}", problems.join("\n"));