From 6bfee4638d42a83078f2ecf2bab98e0451187053 Mon Sep 17 00:00:00 2001 From: John Coffey Date: Thu, 29 Oct 2026 11:48:59 -0700 Subject: [PATCH] IMAP: check each pipelined command when it runs Every command in a read buffer was checked against the session state before any of them ran. A SELECT pipelined with the FETCH or STORE after it got "BAD No mailbox is selected" for those, though RFC 9051 allows pipelining them. The other way round was worse: with a mailbox selected, UNSELECT (or CLOSE, or UNAUTHENTICATE) followed in the same write by FETCH let the FETCH through, and its handler then found no mailbox and panicked on unreachable!(), dropping the connection. Each command is now checked just before it runs, in the state the commands before it left. Commands batched with the one before them (FETCH, STATUS, CREATE, DELETE) are each checked, and counted against the rate limit, on their own; one that isn't allowed is answered and left out, and a refusal that ends the session still ends it. STARTTLS still returns before anything pipelined after it runs. The pipelining test now sends SELECT with STORE and FETCH in one write, and expects BAD for a FETCH pipelined after UNSELECT. On the old code it fails (the session panicked); with this change the imap suite passes. --- crates/imap/src/core/client.rs | 118 ++++++++++++++++++++------------- tests/src/imap/pipelining.rs | 33 ++++++--- 2 files changed, 98 insertions(+), 53 deletions(-) diff --git a/crates/imap/src/core/client.rs b/crates/imap/src/core/client.rs index 435d4d61..bc8a5f7e 100644 --- a/crates/imap/src/core/client.rs +++ b/crates/imap/src/core/client.rs @@ -36,18 +36,10 @@ impl Session { 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 Session { 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 Session { .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 Session { .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 Session { 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>>, - mut grouped_requests: Vec>, -) -> Vec> { - 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 Session { + // 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>>, + first: Request, + ) -> Option>> { + 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) -> trc::Result> { let state = &self.state; // Rate limit request diff --git a/tests/src/imap/pipelining.rs b/tests/src/imap/pipelining.rs index b408286f..160cc30c 100644 --- a/tests/src/imap/pipelining.rs +++ b/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::>().join(" ")) + .collect::>(); + 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")); -- 2.54.0