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"));