From 8bda777b07f9508afc1ebd51027063fbe365d921 Mon Sep 17 00:00:00 2001 From: Gabe A Date: Tue, 4 Aug 2026 10:30:57 -0400 Subject: [PATCH 1/4] add bounded uidonly protocol guard --- crates/core/src/imap/mod.rs | 3 + crates/core/src/imap/uidonly.rs | 1042 +++++++++++++++++++++++++ crates/core/src/imap/uidonly_tests.rs | 517 ++++++++++++ crates/core/src/logger/mod.rs | 7 +- 4 files changed, 1568 insertions(+), 1 deletion(-) create mode 100644 crates/core/src/imap/uidonly.rs create mode 100644 crates/core/src/imap/uidonly_tests.rs diff --git a/crates/core/src/imap/mod.rs b/crates/core/src/imap/mod.rs index 0866f1a1..161c760c 100644 --- a/crates/core/src/imap/mod.rs +++ b/crates/core/src/imap/mod.rs @@ -26,5 +26,8 @@ pub mod session; pub mod stats; #[cfg(test)] mod tests; +pub(crate) mod uidonly; +#[cfg(test)] +mod uidonly_tests; #[cfg(test)] pub mod mock_server; diff --git a/crates/core/src/imap/uidonly.rs b/crates/core/src/imap/uidonly.rs new file mode 100644 index 00000000..610ec88c --- /dev/null +++ b/crates/core/src/imap/uidonly.rs @@ -0,0 +1,1042 @@ +// +// Copyright (c) 2025-2026 rustmailer.com (https://rustmailer.com) +// +// This file is part of the Bichon Email Archiving Project +// +// This program is free software: you can redistribute it and/or modify +// it under the terms of the GNU Affero General Public License as published by +// the Free Software Foundation, either version 3 of the License, or +// (at your option) any later version. +// +// This program is distributed in the hope that it will be useful, +// but WITHOUT ANY WARRANTY; without even the implied warranty of +// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +// GNU Affero General Public License for more details. +// +// You should have received a copy of the GNU Affero General Public License +// along with this program. If not, see . + +// The routing layer lands later in the stack and consumes this module's API. +#![allow(dead_code)] + +//! Bounded RFC 9586 compatibility for the released `async-imap` parser. +//! +//! The adapter is installed around the final transport before `Client` is +//! constructed. It is transparent through authentication, observes the exact +//! `ENABLE UIDONLY` exchange, then translates only top-level `UIDFETCH` atoms +//! to `FETCH`. Literal bytes are never searched or rewritten. + +use crate::imap::session::SessionStream; +use std::collections::VecDeque; +use std::future::Future; +use std::io; +use std::pin::Pin; +use std::sync::{Arc, Mutex}; +use std::task::{Context, Poll}; +use std::time::Duration; +use tokio::io::{AsyncRead, AsyncWrite, ReadBuf}; +use tokio::time::Sleep; + +const CHUNK: usize = 8 * 1024; + +/// Hard limits applied before bytes reach `imap-proto`. +#[derive(Clone, Debug)] +pub(crate) struct UidOnlyLimits { + pub max_control_line_bytes: usize, + pub max_literal_bytes: usize, + pub max_response_bytes: usize, + pub max_command_literal_bytes: usize, + pub max_command_response_bytes: usize, + pub max_command_responses: usize, + pub max_command_runtime: Duration, +} + +impl Default for UidOnlyLimits { + fn default() -> Self { + Self { + max_control_line_bytes: 64 * 1024, + max_literal_bytes: 25 * 1024 * 1024, + max_response_bytes: 26 * 1024 * 1024, + max_command_literal_bytes: 100 * 1024 * 1024, + max_command_response_bytes: 128 * 1024 * 1024, + max_command_responses: 2_048, + max_command_runtime: Duration::from_secs(5 * 60), + } + } +} + +impl UidOnlyLimits { + fn validate(&self) -> io::Result<()> { + if self.max_control_line_bytes < 2 + || self.max_literal_bytes == 0 + || self.max_response_bytes < self.max_literal_bytes + || self.max_command_literal_bytes < self.max_literal_bytes + || self.max_command_response_bytes < self.max_response_bytes + || self.max_command_responses < 2 + || self.max_command_runtime.is_zero() + { + return Err(invalid("invalid UIDONLY limits")); + } + Ok(()) + } +} + +#[derive(Clone, Debug)] +pub(crate) struct UidOnlyHandle { + health: Arc>, +} + +impl UidOnlyHandle { + pub fn ensure_active(&self) -> io::Result<()> { + let health = self.health.lock().expect("UIDONLY health poisoned"); + if let Some(reason) = &health.poison { + return Err(invalid(reason.clone())); + } + if !health.active { + return Err(invalid("UIDONLY activation was not confirmed")); + } + Ok(()) + } + + pub fn poison_reason(&self) -> Option { + self.health + .lock() + .expect("UIDONLY health poisoned") + .poison + .clone() + } + + pub fn literal_bytes_received(&self) -> u64 { + self.health + .lock() + .expect("UIDONLY health poisoned") + .literal_bytes + } + + /// Sets the pre-read literal ceiling for the next exact-message fetch. + pub fn arm_next_fetch_literal_limit(&self, limit: usize) -> io::Result<()> { + let mut health = self.health.lock().expect("UIDONLY health poisoned"); + if limit == 0 + || !health.active + || health.command_in_flight + || health.next_literal_limit.is_some() + { + return Err(invalid( + "cannot arm UIDONLY literal limit in the current state", + )); + } + health.next_literal_limit = Some(limit); + Ok(()) + } +} + +#[derive(Debug, Default)] +struct Health { + active: bool, + poison: Option, + literal_bytes: u64, + command_in_flight: bool, + next_literal_limit: Option, +} + +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +enum Mode { + PassThrough, + Enabling, + Active, + Poisoned, +} + +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +enum CommandKind { + Enable, + Examine, + Inventory, + ExactFetch, + Logout, +} + +#[derive(Debug)] +struct PendingCommand { + tag: Vec, + kind: CommandKind, + saw_enabled: bool, + response_bytes: usize, + response_count: usize, + literal_bytes: usize, + literal_limit: usize, + deadline: Pin>, +} + +/// Literal-aware stream bridge. Use [`wrap`] for Bichon's erased transport. +#[derive(Debug)] +pub(crate) struct UidOnlyStream { + inner: T, + limits: UidOnlyLimits, + health: Arc>, + mode: Mode, + pending: Option, + input: VecDeque, + output: VecDeque, + control_line: Vec, + command_line: Vec, + outgoing: VecDeque, + literal_remaining: usize, + response_bytes: usize, + response_uidfetch: bool, + first_line: bool, + in_response: bool, + eof: bool, +} + +impl UidOnlyStream { + pub fn new(inner: T, limits: UidOnlyLimits) -> io::Result<(Self, UidOnlyHandle)> { + limits.validate()?; + let health = Arc::new(Mutex::new(Health::default())); + let handle = UidOnlyHandle { + health: Arc::clone(&health), + }; + Ok(( + Self { + inner, + limits, + health, + mode: Mode::PassThrough, + pending: None, + input: VecDeque::new(), + output: VecDeque::new(), + control_line: Vec::new(), + command_line: Vec::new(), + outgoing: VecDeque::new(), + literal_remaining: 0, + response_bytes: 0, + response_uidfetch: false, + first_line: true, + in_response: false, + eof: false, + }, + handle, + )) + } + + fn fail(&mut self, kind: io::ErrorKind, reason: impl Into) -> io::Error { + let reason = reason.into(); + self.mode = Mode::Poisoned; + let mut health = self.health.lock().expect("UIDONLY health poisoned"); + health.active = false; + health.command_in_flight = false; + health.next_literal_limit = None; + health.poison = Some(reason.clone()); + io::Error::new(kind, reason) + } + + fn check_deadline(&mut self, cx: &mut Context<'_>) -> io::Result<()> { + let elapsed = self + .pending + .as_mut() + .is_some_and(|pending| pending.deadline.as_mut().poll(cx).is_ready()); + if elapsed { + return Err(self.fail(io::ErrorKind::TimedOut, "UIDONLY command timed out")); + } + Ok(()) + } + + fn start_command(&mut self, tag: &[u8], kind: CommandKind) -> io::Result<()> { + if self.pending.is_some() { + return Err(self.fail( + io::ErrorKind::InvalidInput, + "UIDONLY permits only one command in flight", + )); + } + let literal_limit = if kind == CommandKind::ExactFetch { + let limit = self + .health + .lock() + .expect("UIDONLY health poisoned") + .next_literal_limit + .take(); + let Some(limit) = limit else { + return Err(self.fail( + io::ErrorKind::InvalidInput, + "exact UIDONLY fetch was not armed with a literal limit", + )); + }; + limit.min(self.limits.max_command_literal_bytes) + } else { + 0 + }; + self.health + .lock() + .expect("UIDONLY health poisoned") + .command_in_flight = true; + self.pending = Some(PendingCommand { + tag: tag.to_vec(), + kind, + saw_enabled: false, + response_bytes: 0, + response_count: 0, + literal_bytes: 0, + literal_limit, + deadline: Box::pin(tokio::time::sleep(self.limits.max_command_runtime)), + }); + Ok(()) + } + + fn validate_command(&mut self, line: &[u8]) -> io::Result<()> { + let line = line + .strip_suffix(b"\r\n") + .ok_or_else(|| invalid("outbound IMAP line is not CRLF terminated"))?; + let Some(space) = line.iter().position(|byte| *byte == b' ') else { + if self.mode == Mode::PassThrough { + return Ok(()); + } + return Err(self.fail(io::ErrorKind::InvalidInput, "untagged UIDONLY command")); + }; + let (tag, command) = (&line[..space], &line[space + 1..]); + if tag.is_empty() || tag == b"*" || tag == b"+" { + if self.mode == Mode::PassThrough { + return Ok(()); + } + return Err(self.fail(io::ErrorKind::InvalidInput, "invalid UIDONLY command tag")); + } + + match self.mode { + Mode::PassThrough if command.eq_ignore_ascii_case(b"ENABLE UIDONLY") => { + self.start_command(tag, CommandKind::Enable)?; + self.mode = Mode::Enabling; + } + Mode::PassThrough => {} + Mode::Enabling => { + return Err(self.fail( + io::ErrorKind::InvalidInput, + "command sent before ENABLE UIDONLY completed", + )); + } + Mode::Active => { + let kind = if starts_ci(command, b"EXAMINE ") && command.len() > 8 { + CommandKind::Examine + } else if command.eq_ignore_ascii_case(b"LOGOUT") { + CommandKind::Logout + } else if let Some(kind) = uid_fetch_command_kind(command) { + kind + } else { + return Err(self.fail( + io::ErrorKind::InvalidInput, + "command is not allowed after UIDONLY activation", + )); + }; + self.start_command(tag, kind)?; + } + Mode::Poisoned => return Err(invalid("UIDONLY stream is poisoned")), + } + Ok(()) + } + + fn start_response(&mut self) -> io::Result<()> { + self.in_response = true; + self.first_line = true; + self.response_uidfetch = false; + self.response_bytes = 0; + if let Some(command) = self.pending.as_mut() { + command.response_count = command + .response_count + .checked_add(1) + .ok_or_else(|| invalid("UIDONLY response count overflow"))?; + if command.response_count > self.limits.max_command_responses { + return Err(self.fail( + io::ErrorKind::InvalidData, + "UIDONLY command response count exceeded", + )); + } + } + Ok(()) + } + + fn add_wire_bytes(&mut self, count: usize, literal: bool) -> io::Result<()> { + self.response_bytes = self + .response_bytes + .checked_add(count) + .ok_or_else(|| invalid("UIDONLY response byte count overflow"))?; + if self.response_bytes > self.limits.max_response_bytes { + return Err(self.fail(io::ErrorKind::InvalidData, "UIDONLY response too large")); + } + if let Some(command) = self.pending.as_mut() { + command.response_bytes = command + .response_bytes + .checked_add(count) + .ok_or_else(|| invalid("UIDONLY command byte count overflow"))?; + if command.response_bytes > self.limits.max_command_response_bytes { + return Err(self.fail( + io::ErrorKind::InvalidData, + "UIDONLY command response bytes exceeded", + )); + } + } + if literal { + let mut health = self.health.lock().expect("UIDONLY health poisoned"); + health.literal_bytes = health + .literal_bytes + .checked_add(count as u64) + .ok_or_else(|| invalid("UIDONLY literal meter overflow"))?; + } + Ok(()) + } + + fn reserve_literal(&mut self, length: usize) -> io::Result<()> { + if length > self.limits.max_literal_bytes { + return Err(self.fail(io::ErrorKind::InvalidData, "UIDONLY literal too large")); + } + if self.response_bytes.saturating_add(length) > self.limits.max_response_bytes { + return Err(self.fail( + io::ErrorKind::InvalidData, + "UIDONLY literal exceeds response budget", + )); + } + if let Some(command) = self.pending.as_mut() { + if command.response_bytes.saturating_add(length) + > self.limits.max_command_response_bytes + || command.literal_bytes.saturating_add(length) > command.literal_limit + { + return Err(self.fail( + io::ErrorKind::InvalidData, + "UIDONLY literal exceeds command budget", + )); + } + command.literal_bytes += length; + } else if self.mode == Mode::Active { + return Err(self.fail( + io::ErrorKind::InvalidData, + "UIDONLY literal arrived outside a command", + )); + } + self.literal_remaining = length; + Ok(()) + } + + fn classify_first_line(&mut self, line: &mut Vec) -> io::Result<()> { + if self.mode == Mode::Enabling + && numeric_atom(line).is_some_and(|atom| atom.eq_ignore_ascii_case(b"UIDFETCH")) + { + return Err(self.fail( + io::ErrorKind::InvalidData, + "UIDFETCH arrived before UIDONLY activation", + )); + } + if self.mode != Mode::Active { + return Ok(()); + } + if untagged_atom(line).is_some_and(|atom| atom.eq_ignore_ascii_case(b"VANISHED")) { + return Err(self.fail( + io::ErrorKind::InvalidData, + "VANISHED aborts UIDONLY acquisition", + )); + } + if contains_ci(line, b"[UNSEEN ") { + if self + .pending + .as_ref() + .is_some_and(|command| command.kind == CommandKind::Examine) + && starts_ci(line, b"* OK [UNSEEN ") + { + // Cyrus 3.12 advertises UIDONLY but still emits this optional + // sequence-number response code on EXAMINE. Do not expose the + // unusable sequence number to the parser or acquisition logic. + *line = b"* OK UIDONLY ignored UNSEEN response code\r\n".to_vec(); + } else { + return Err(self.fail( + io::ErrorKind::InvalidData, + "sequence-bearing UNSEEN is forbidden in UIDONLY mode", + )); + } + } + if contains_ci(line, b"[UIDNOTSTICKY") { + return Err(self.fail( + io::ErrorKind::InvalidData, + "UIDNOTSTICKY is forbidden in UIDONLY mode", + )); + } + if self.pending.as_ref().is_some_and(|command| { + matches!( + command.kind, + CommandKind::Inventory | CommandKind::ExactFetch + ) + }) && line.starts_with(b"* ") + && !numeric_atom(line).is_some_and(|atom| atom.eq_ignore_ascii_case(b"UIDFETCH")) + && !count_atom(line).is_some_and(|atom| { + atom.eq_ignore_ascii_case(b"EXISTS") || atom.eq_ignore_ascii_case(b"RECENT") + }) + && !untagged_atom(line).is_some_and(|atom| atom.eq_ignore_ascii_case(b"OK")) + { + let reason = if starts_ci(line, b"* NO [MESSAGELIMIT ") { + "server reported MESSAGELIMIT during UIDONLY fetch" + } else { + "unexpected untagged response during UIDONLY fetch" + }; + return Err(self.fail(io::ErrorKind::InvalidData, reason)); + } + match numeric_atom_with_range(line) { + Some((atom, start, end)) if atom.eq_ignore_ascii_case(b"UIDFETCH") => { + if is_ignorable_uidfetch_notification(line)? { + *line = b"* OK UIDONLY ignored flag notification\r\n".to_vec(); + return Ok(()); + } + let kind = self.pending.as_ref().map(|command| command.kind); + validate_uidfetch_shape(line, kind) + .map_err(|error| self.fail(io::ErrorKind::InvalidData, error.to_string()))?; + let mut rewritten = Vec::with_capacity(line.len() - 3); + rewritten.extend_from_slice(&line[..start]); + rewritten.extend_from_slice(b"FETCH"); + rewritten.extend_from_slice(&line[end..]); + *line = rewritten; + self.response_uidfetch = true; + } + Some((atom, _, _)) if atom.eq_ignore_ascii_case(b"FETCH") => { + return Err(self.fail( + io::ErrorKind::InvalidData, + "raw FETCH is forbidden after UIDONLY activation", + )); + } + Some((atom, _, _)) if atom.eq_ignore_ascii_case(b"EXPUNGE") => { + return Err(self.fail( + io::ErrorKind::InvalidData, + "sequence EXPUNGE is forbidden in UIDONLY mode", + )); + } + _ => {} + } + Ok(()) + } + + fn finish_response(&mut self, line: &[u8]) -> io::Result<()> { + let completion = tagged_completion(line); + if self.mode == Mode::Enabling { + if line.eq_ignore_ascii_case(b"* ENABLED UIDONLY\r\n") { + if let Some(command) = self.pending.as_mut() { + command.saw_enabled = true; + } + } else if untagged_atom(line).is_some_and(|atom| atom.eq_ignore_ascii_case(b"ENABLED")) + { + return Err(self.fail( + io::ErrorKind::InvalidData, + "server returned a non-exact ENABLED UIDONLY response", + )); + } + } + if self.mode != Mode::PassThrough { + if let Some((tag, status)) = completion { + let Some(command) = self.pending.as_ref() else { + return Err(self.fail( + io::ErrorKind::InvalidData, + "tagged completion arrived without a command", + )); + }; + if tag != command.tag { + return Err(self.fail( + io::ErrorKind::InvalidData, + "tagged completion did not match the active command", + )); + } + if !status.eq_ignore_ascii_case(b"OK") { + return Err(self.fail( + io::ErrorKind::InvalidData, + "UIDONLY command did not complete OK", + )); + } + if command.kind == CommandKind::Enable && !command.saw_enabled { + return Err(self.fail( + io::ErrorKind::InvalidData, + "ENABLE completed without exact ENABLED UIDONLY", + )); + } + let enabled = command.kind == CommandKind::Enable; + self.pending = None; + self.health + .lock() + .expect("UIDONLY health poisoned") + .command_in_flight = false; + if enabled { + self.mode = Mode::Active; + self.health.lock().expect("UIDONLY health poisoned").active = true; + } + } + } + self.in_response = false; + self.first_line = true; + self.response_uidfetch = false; + self.response_bytes = 0; + Ok(()) + } + + fn process_line(&mut self) -> io::Result<()> { + let mut line = std::mem::take(&mut self.control_line); + if self.first_line { + self.classify_first_line(&mut line)?; + self.first_line = false; + } else if self.response_uidfetch && line != b")\r\n" { + return Err(self.fail( + io::ErrorKind::InvalidData, + "only closing ')' may follow an exact UIDFETCH literal", + )); + } + let bare = &line[..line.len() - 2]; + let literal = literal_length(bare)?; + if self.mode == Mode::Enabling && literal.is_some() { + return Err(self.fail( + io::ErrorKind::InvalidData, + "literal response is forbidden during UIDONLY activation", + )); + } + if self.mode == Mode::Active && literal.is_some() && !self.response_uidfetch { + return Err(self.fail( + io::ErrorKind::InvalidData, + "literal outside UIDFETCH is forbidden in UIDONLY mode", + )); + } + if let Some(length) = literal { + self.reserve_literal(length)?; + self.output.extend(line); + } else { + self.finish_response(&line)?; + self.output.extend(line); + } + Ok(()) + } + + fn process_input(&mut self) -> io::Result { + if !self.output.is_empty() { + return Ok(true); + } + if !self.in_response && !self.input.is_empty() { + self.start_response()?; + } + if self.literal_remaining > 0 { + let count = self.literal_remaining.min(self.input.len()).min(CHUNK); + if count == 0 { + return Ok(false); + } + self.add_wire_bytes(count, true)?; + self.input.make_contiguous(); + let (bytes, _) = self.input.as_slices(); + self.output.extend(&bytes[..count]); + self.input.drain(..count); + self.literal_remaining -= count; + return Ok(true); + } + while let Some(byte) = self.input.pop_front() { + self.control_line.push(byte); + self.add_wire_bytes(1, false)?; + if self.control_line.len() > self.limits.max_control_line_bytes { + return Err(self.fail(io::ErrorKind::InvalidData, "UIDONLY control line too long")); + } + if byte == b'\n' { + if !self.control_line.ends_with(b"\r\n") { + return Err(self.fail(io::ErrorKind::InvalidData, "bare LF in IMAP response")); + } + self.process_line()?; + return Ok(true); + } + } + Ok(false) + } + + fn copy_output(&mut self, destination: &mut ReadBuf<'_>) { + self.output.make_contiguous(); + let (bytes, _) = self.output.as_slices(); + let count = destination.remaining().min(bytes.len()); + destination.put_slice(&bytes[..count]); + self.output.drain(..count); + } + + fn drain_outgoing(&mut self, cx: &mut Context<'_>) -> Poll> + where + T: AsyncWrite + Unpin, + { + while !self.outgoing.is_empty() { + let written = { + let (head, _) = self.outgoing.as_slices(); + match Pin::new(&mut self.inner).poll_write(cx, head) { + Poll::Pending => return Poll::Pending, + Poll::Ready(Err(error)) => return Poll::Ready(Err(error)), + Poll::Ready(Ok(0)) => { + return Poll::Ready(Err(io::Error::new( + io::ErrorKind::WriteZero, + "failed to write UIDONLY command", + ))) + } + Poll::Ready(Ok(written)) => written, + } + }; + self.outgoing.drain(..written); + } + Poll::Ready(Ok(())) + } +} + +impl AsyncRead for UidOnlyStream { + fn poll_read( + self: Pin<&mut Self>, + cx: &mut Context<'_>, + destination: &mut ReadBuf<'_>, + ) -> Poll> { + let this = self.get_mut(); + if destination.remaining() == 0 { + return Poll::Ready(Ok(())); + } + if let Err(error) = this.check_deadline(cx) { + return Poll::Ready(Err(error)); + } + loop { + if !this.output.is_empty() { + this.copy_output(destination); + return Poll::Ready(Ok(())); + } + match this.process_input() { + Ok(true) => continue, + Ok(false) => {} + Err(error) => return Poll::Ready(Err(error)), + } + if this.eof { + if this.in_response || !this.control_line.is_empty() || this.literal_remaining > 0 { + let error = this.fail(io::ErrorKind::UnexpectedEof, "truncated IMAP response"); + return Poll::Ready(Err(error)); + } + if this.pending.is_some() { + let error = this.fail( + io::ErrorKind::UnexpectedEof, + "connection closed before UIDONLY command completion", + ); + return Poll::Ready(Err(error)); + } + return Poll::Ready(Ok(())); + } + let mut bytes = [0_u8; CHUNK]; + let mut buffer = ReadBuf::new(&mut bytes); + match Pin::new(&mut this.inner).poll_read(cx, &mut buffer) { + Poll::Pending => return Poll::Pending, + Poll::Ready(Err(error)) => return Poll::Ready(Err(error)), + Poll::Ready(Ok(())) if buffer.filled().is_empty() => this.eof = true, + Poll::Ready(Ok(())) => this.input.extend(buffer.filled()), + } + } + } +} + +impl AsyncWrite for UidOnlyStream { + fn poll_write( + self: Pin<&mut Self>, + _cx: &mut Context<'_>, + bytes: &[u8], + ) -> Poll> { + let this = self.get_mut(); + if this.mode == Mode::Poisoned { + return Poll::Ready(Err(invalid("UIDONLY stream is poisoned"))); + } + for byte in bytes { + this.command_line.push(*byte); + if this.command_line.len() > this.limits.max_control_line_bytes { + let error = this.fail(io::ErrorKind::InvalidInput, "outbound IMAP line too long"); + return Poll::Ready(Err(error)); + } + if *byte == b'\n' { + let line = std::mem::take(&mut this.command_line); + if let Err(error) = this.validate_command(&line) { + return Poll::Ready(Err(error)); + } + this.outgoing.extend(line); + } + } + Poll::Ready(Ok(bytes.len())) + } + + fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + let this = self.get_mut(); + if let Err(error) = this.check_deadline(cx) { + return Poll::Ready(Err(error)); + } + if !this.command_line.is_empty() { + return Poll::Ready(Err(this.fail( + io::ErrorKind::InvalidInput, + "flush attempted with a partial IMAP command", + ))); + } + match this.drain_outgoing(cx) { + Poll::Ready(Ok(())) => Pin::new(&mut this.inner).poll_flush(cx), + other => other, + } + } + + fn poll_shutdown(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + match self.as_mut().poll_flush(cx) { + Poll::Ready(Ok(())) => Pin::new(&mut self.get_mut().inner).poll_shutdown(cx), + other => other, + } + } +} + +impl SessionStream for UidOnlyStream {} + +/// Installs the bridge around Bichon's final (already TLS-upgraded) transport. +pub(crate) fn wrap( + stream: Box, + limits: UidOnlyLimits, +) -> io::Result<(Box, UidOnlyHandle)> { + let (stream, handle) = UidOnlyStream::new(stream, limits)?; + Ok((Box::new(stream), handle)) +} + +/// Injection-safe arguments for ascending RFC 9394 inventory paging. +pub(crate) fn inventory_args( + cursor: u32, + high: u32, + page_size: u32, +) -> io::Result<(String, String)> { + if cursor == 0 || high == 0 || cursor > high || page_size == 0 { + return Err(invalid("invalid UIDONLY inventory bounds")); + } + Ok(( + format!("{cursor}:{high}"), + format!("(UID RFC822.SIZE) (PARTIAL 1:{page_size})"), + )) +} + +/// Injection-safe arguments for one complete, literal-backed raw message. +pub(crate) fn exact_args(uid: u32) -> io::Result<(String, &'static str)> { + if uid == 0 { + return Err(invalid("UID 0 is invalid")); + } + Ok((uid.to_string(), "(UID RFC822.SIZE BODY.PEEK[])")) +} + +fn uid_fetch_command_kind(command: &[u8]) -> Option { + let Some(rest) = strip_prefix_ci(command, b"UID FETCH ") else { + return None; + }; + let Some(space) = rest.iter().position(|byte| *byte == b' ') else { + return None; + }; + let (set, query) = (&rest[..space], &rest[space + 1..]); + if let Some(colon) = set.iter().position(|byte| *byte == b':') { + let Some(low) = parse_nonzero_u32(&set[..colon]) else { + return None; + }; + let Some(high) = parse_nonzero_u32(&set[colon + 1..]) else { + return None; + }; + if low > high { + return None; + } + query + .strip_prefix(b"(UID RFC822.SIZE) (PARTIAL 1:") + .and_then(|tail| tail.strip_suffix(b")")) + .and_then(parse_nonzero_u32) + .map(|_| CommandKind::Inventory) + } else { + (parse_nonzero_u32(set).is_some() && query == b"(UID RFC822.SIZE BODY.PEEK[])") + .then_some(CommandKind::ExactFetch) + } +} + +fn validate_uidfetch_shape(line: &[u8], command: Option) -> io::Result<()> { + let bare = line + .strip_suffix(b"\r\n") + .ok_or_else(|| invalid("invalid UIDFETCH response"))?; + let rest = bare + .strip_prefix(b"* ") + .ok_or_else(|| invalid("invalid UIDFETCH response"))?; + let number_end = rest + .iter() + .position(|byte| *byte == b' ') + .ok_or_else(|| invalid("invalid UIDFETCH response"))?; + let leading_uid = parse_nonzero_u32(&rest[..number_end]) + .ok_or_else(|| invalid("invalid leading UIDFETCH UID"))?; + let attributes = strip_prefix_ci(&rest[number_end + 1..], b"UIDFETCH ") + .and_then(|value| value.strip_prefix(b"(")) + .ok_or_else(|| invalid("invalid UIDFETCH attribute list"))?; + + let (attributes, exact) = match command { + Some(CommandKind::Inventory) => ( + attributes + .strip_suffix(b")") + .ok_or_else(|| invalid("inventory UIDFETCH must not contain a literal"))?, + false, + ), + Some(CommandKind::ExactFetch) => { + literal_length(bare)? + .ok_or_else(|| invalid("exact UIDFETCH BODY[] must be literal-backed"))?; + let marker = attributes + .iter() + .rposition(|byte| *byte == b'{') + .ok_or_else(|| invalid("exact UIDFETCH is missing its literal marker"))?; + (&attributes[..marker], true) + } + _ => return Err(invalid("UIDFETCH arrived outside a UID fetch command")), + }; + + let mut uid = None; + let mut size = None; + let mut body = false; + let mut tokens = attributes + .split(|b| b.is_ascii_whitespace()) + .filter(|t| !t.is_empty()); + while let Some(token) = tokens.next() { + if token.eq_ignore_ascii_case(b"UID") && uid.is_none() { + uid = tokens.next().and_then(parse_nonzero_u32); + if uid.is_none() { + return Err(invalid("UIDFETCH has an invalid or duplicate UID")); + } + } else if token.eq_ignore_ascii_case(b"RFC822.SIZE") && size.is_none() { + size = tokens.next().and_then(parse_u32); + if size.is_none() { + return Err(invalid("UIDFETCH has an invalid or duplicate RFC822.SIZE")); + } + } else if exact && token.eq_ignore_ascii_case(b"BODY[]") && !body { + body = true; + if tokens.next().is_some() { + return Err(invalid("BODY[] must be the final exact UIDFETCH attribute")); + } + } else { + return Err(invalid("unexpected or duplicate UIDFETCH attribute")); + } + } + let uid = uid.ok_or_else(|| invalid("UIDFETCH omitted UID"))?; + size.ok_or_else(|| invalid("UIDFETCH omitted RFC822.SIZE"))?; + if uid != leading_uid { + return Err(invalid("leading UIDFETCH UID disagrees with UID attribute")); + } + if body != exact { + return Err(invalid("UIDFETCH body did not match command")); + } + Ok(()) +} + +fn is_ignorable_uidfetch_notification(line: &[u8]) -> io::Result { + let bare = line + .strip_suffix(b"\r\n") + .ok_or_else(|| invalid("invalid UIDFETCH notification"))?; + if literal_length(bare)?.is_some() { + return Ok(false); + } + let Some((atom, _, end)) = numeric_atom_with_range(line) else { + return Ok(false); + }; + if !atom.eq_ignore_ascii_case(b"UIDFETCH") { + return Ok(false); + } + let attributes = bare[end..] + .strip_prefix(b" (") + .and_then(|value| value.strip_suffix(b")")) + .ok_or_else(|| invalid("invalid UIDFETCH notification attributes"))?; + Ok(contains_ci(attributes, b"FLAGS") + && !contains_ci(attributes, b"RFC822.SIZE") + && !contains_ci(attributes, b"BODY")) +} + +fn literal_length(line: &[u8]) -> io::Result> { + if !line.ends_with(b"}") { + return Ok(None); + } + let Some(open) = line.iter().rposition(|byte| *byte == b'{') else { + return Err(invalid("invalid IMAP literal marker")); + }; + if open > 0 && line[open - 1] == b'~' { + return Err(invalid("literal8 is unsupported")); + } + let digits = &line[open + 1..line.len() - 1]; + if digits.is_empty() || !digits.iter().all(u8::is_ascii_digit) { + return Err(invalid("invalid IMAP literal marker")); + } + let length = std::str::from_utf8(digits) + .ok() + .and_then(|value| value.parse::().ok()) + .ok_or_else(|| invalid("IMAP literal length overflow"))?; + Ok(Some(length)) +} + +fn numeric_atom_with_range(line: &[u8]) -> Option<(&[u8], usize, usize)> { + numbered_atom_with_range(line, false) +} + +fn numbered_atom_with_range(line: &[u8], allow_zero: bool) -> Option<(&[u8], usize, usize)> { + let bare = line.strip_suffix(b"\r\n")?; + let rest = bare.strip_prefix(b"* ")?; + let number_end = rest.iter().position(|byte| *byte == b' ')?; + if allow_zero { + parse_u32(&rest[..number_end])?; + } else { + parse_nonzero_u32(&rest[..number_end])?; + } + let atom_start = 2 + number_end + 1; + let tail = &bare[atom_start..]; + let atom_end = tail + .iter() + .position(|byte| *byte == b' ' || *byte == b'\r' || *byte == b'\n') + .unwrap_or(tail.len()); + Some((&tail[..atom_end], atom_start, atom_start + atom_end)) +} + +fn numeric_atom(line: &[u8]) -> Option<&[u8]> { + numeric_atom_with_range(line).map(|value| value.0) +} + +fn count_atom(line: &[u8]) -> Option<&[u8]> { + numbered_atom_with_range(line, true).map(|value| value.0) +} + +fn untagged_atom(line: &[u8]) -> Option<&[u8]> { + let bare = line.strip_suffix(b"\r\n")?; + let rest = bare.strip_prefix(b"* ")?; + let end = rest + .iter() + .position(|byte| *byte == b' ') + .unwrap_or(rest.len()); + Some(&rest[..end]) +} + +fn tagged_completion(line: &[u8]) -> Option<(&[u8], &[u8])> { + let bare = line.strip_suffix(b"\r\n")?; + if bare.starts_with(b"* ") || bare.starts_with(b"+ ") { + return None; + } + let mut fields = bare.split(|byte| *byte == b' '); + let tag = fields.next()?; + let status = fields.next()?; + if [b"OK".as_slice(), b"NO".as_slice(), b"BAD".as_slice()] + .iter() + .any(|candidate| status.eq_ignore_ascii_case(candidate)) + { + Some((tag, status)) + } else { + None + } +} + +fn parse_nonzero_u32(bytes: &[u8]) -> Option { + parse_u32(bytes).filter(|number| *number != 0) +} + +fn parse_u32(bytes: &[u8]) -> Option { + if bytes.is_empty() || !bytes.iter().all(u8::is_ascii_digit) { + return None; + } + std::str::from_utf8(bytes).ok()?.parse().ok() +} + +fn starts_ci(value: &[u8], prefix: &[u8]) -> bool { + value + .get(..prefix.len()) + .is_some_and(|head| head.eq_ignore_ascii_case(prefix)) +} + +fn strip_prefix_ci<'a>(value: &'a [u8], prefix: &[u8]) -> Option<&'a [u8]> { + starts_ci(value, prefix).then(|| &value[prefix.len()..]) +} + +fn contains_ci(value: &[u8], needle: &[u8]) -> bool { + value + .windows(needle.len()) + .any(|window| window.eq_ignore_ascii_case(needle)) +} + +fn invalid(reason: impl Into) -> io::Error { + io::Error::new(io::ErrorKind::InvalidData, reason.into()) +} diff --git a/crates/core/src/imap/uidonly_tests.rs b/crates/core/src/imap/uidonly_tests.rs new file mode 100644 index 00000000..37e4248b --- /dev/null +++ b/crates/core/src/imap/uidonly_tests.rs @@ -0,0 +1,517 @@ +// +// Copyright (c) 2025-2026 rustmailer.com (https://rustmailer.com) +// +// This file is part of the Bichon Email Archiving Project +// +// This program is free software: you can redistribute it and/or modify +// it under the terms of the GNU Affero General Public License as published by +// the Free Software Foundation, either version 3 of the License, or +// (at your option) any later version. +// +// This program is distributed in the hope that it will be useful, +// but WITHOUT ANY WARRANTY; without even the implied warranty of +// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +// GNU Affero General Public License for more details. +// +// You should have received a copy of the GNU Affero General Public License +// along with this program. If not, see . + +use super::uidonly::{exact_args, inventory_args, UidOnlyHandle, UidOnlyLimits, UidOnlyStream}; +use futures::TryStreamExt; +use std::collections::VecDeque; +use std::io; +use std::pin::Pin; +use std::sync::{Arc, Mutex}; +use std::task::{Context, Poll}; +use std::time::Duration; +use tokio::io::{AsyncRead, AsyncWrite, ReadBuf}; + +#[derive(Debug)] +struct ScriptIo { + input: VecDeque, + max_chunk: usize, + pending_at_eof: bool, + writes: Arc>>, +} + +impl ScriptIo { + fn new(input: impl Into>, max_chunk: usize) -> (Self, Arc>>) { + let writes = Arc::new(Mutex::new(Vec::new())); + ( + Self { + input: input.into().into(), + max_chunk, + pending_at_eof: false, + writes: Arc::clone(&writes), + }, + writes, + ) + } + + fn pending(input: impl Into>, max_chunk: usize) -> Self { + let (mut io, _) = Self::new(input, max_chunk); + io.pending_at_eof = true; + io + } +} + +impl AsyncRead for ScriptIo { + fn poll_read( + mut self: Pin<&mut Self>, + _cx: &mut Context<'_>, + destination: &mut ReadBuf<'_>, + ) -> Poll> { + let count = destination + .remaining() + .min(self.max_chunk) + .min(self.input.len()); + if count == 0 && self.pending_at_eof { + return Poll::Pending; + } + for _ in 0..count { + destination.put_slice(&[self.input.pop_front().expect("input length checked")]); + } + Poll::Ready(Ok(())) + } +} + +impl AsyncWrite for ScriptIo { + fn poll_write( + self: Pin<&mut Self>, + _cx: &mut Context<'_>, + bytes: &[u8], + ) -> Poll> { + self.writes.lock().expect("writes poisoned").extend(bytes); + Poll::Ready(Ok(bytes.len())) + } + + fn poll_flush(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll> { + Poll::Ready(Ok(())) + } + + fn poll_shutdown(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll> { + Poll::Ready(Ok(())) + } +} + +type Session = async_imap::Session>; + +async fn authenticated( + after_login: &[u8], + limits: UidOnlyLimits, + chunk: usize, +) -> (Session, UidOnlyHandle, Arc>>) { + let mut transcript = b"* OK synthetic ready\r\nA0001 OK LOGIN completed\r\n".to_vec(); + transcript.extend_from_slice(after_login); + let (io, writes) = ScriptIo::new(transcript, chunk); + let (stream, handle) = UidOnlyStream::new(io, limits).expect("valid limits"); + let mut client = async_imap::Client::new(stream); + client + .read_response() + .await + .expect("greeting read") + .expect("greeting present"); + let session = client.login("synthetic", "redacted").await.expect("login"); + (session, handle, writes) +} + +async fn enabled( + after_enable: &[u8], + limits: UidOnlyLimits, + chunk: usize, +) -> (Session, UidOnlyHandle) { + let mut transcript = b"* ENABLED UIDONLY\r\nA0002 OK ENABLE completed\r\n".to_vec(); + transcript.extend_from_slice(after_enable); + let (mut session, handle, _) = authenticated(&transcript, limits, chunk).await; + session + .run_command_and_check_ok("ENABLE UIDONLY") + .await + .expect("enable"); + handle.ensure_active().expect("adapter active"); + (session, handle) +} + +async fn rejected_fetch(response: &[u8], exact: bool) { + let (mut session, handle) = enabled(response, UidOnlyLimits::default(), 1).await; + if exact { + handle + .arm_next_fetch_literal_limit(16) + .expect("arm body bound"); + } + let stream = if exact { + session + .uid_fetch("7", "(UID RFC822.SIZE BODY.PEEK[])") + .await + } else { + session + .uid_fetch("1:7", "(UID RFC822.SIZE) (PARTIAL 1:7)") + .await + } + .expect("command"); + assert!(stream.try_collect::>().await.is_err()); + assert!(handle.poison_reason().is_some()); +} + +#[tokio::test] +async fn public_async_imap_session_parses_fragmented_uidfetch_without_touching_literal() { + let body = b"From: sender@example.invalid\r\n\r\n* 9 UIDFETCH {999}\r\nsynthetic"; + let mut response = format!( + "* 2 EXISTS\r\n* OK [UNSEEN 1] optional sequence metadata\r\n* OK [UIDVALIDITY 7] epoch\r\n* OK [UIDNEXT 50] next\r\nA0003 OK [READ-ONLY] EXAMINE completed\r\n\ + * 3 EXISTS\r\n* OK [UIDNEXT 51] arrival\r\n* 50 UIDFETCH (FLAGS (\\Seen))\r\n\ + * 42 UIDFETCH (UID 42 RFC822.SIZE {} BODY[] {{{}}}\r\n", + body.len(), + body.len() + ) + .into_bytes(); + response.extend_from_slice(body); + response.extend_from_slice(b")\r\nA0004 OK FETCH completed\r\n"); + + let (mut session, handle) = enabled(&response, UidOnlyLimits::default(), 1).await; + let mailbox = session.examine("Synthetic").await.expect("examine"); + assert_eq!(mailbox.exists, 2); + assert_eq!(mailbox.uid_validity, Some(7)); + assert_eq!(mailbox.uid_next, Some(50)); + + let (set, query) = exact_args(42).expect("safe UID"); + handle + .arm_next_fetch_literal_limit(body.len()) + .expect("arm body bound"); + let fetched: Vec<_> = session + .uid_fetch(set, query) + .await + .expect("command") + .try_collect() + .await + .expect("fetch response"); + assert_eq!(fetched.len(), 1); + assert_eq!(fetched[0].message, 42); + assert_eq!(fetched[0].uid, Some(42)); + assert_eq!(fetched[0].body(), Some(body.as_slice())); + assert_eq!(handle.literal_bytes_received(), body.len() as u64); +} + +#[tokio::test] +async fn activation_requires_exact_enabled_and_matching_ok() { + for response in [ + b"A0002 OK ENABLE completed\r\n".as_slice(), + b"* ENABLED UIDONLY QRESYNC\r\nA0002 OK ENABLE completed\r\n".as_slice(), + b"* ENABLED UIDONLY\r\nA9999 OK ENABLE completed\r\n".as_slice(), + b"* ENABLED UIDONLY\r\nA0002 NO unavailable\r\n".as_slice(), + ] { + let (mut session, handle, _) = authenticated(response, UidOnlyLimits::default(), 2).await; + assert!(session + .run_command_and_check_ok("ENABLE UIDONLY") + .await + .is_err()); + assert!(handle.ensure_active().is_err()); + assert!(handle.poison_reason().is_some()); + } +} + +#[tokio::test] +async fn forbidden_post_activation_responses_fail_closed() { + for response in [ + b"* 7 FETCH (UID 7)\r\nA0003 OK FETCH completed\r\n".as_slice(), + b"* 7 EXPUNGE\r\nA0003 OK FETCH completed\r\n".as_slice(), + b"* VANISHED 7\r\nA0003 OK FETCH completed\r\n".as_slice(), + b"* OK [UNSEEN 7] sequence\r\nA0003 OK FETCH completed\r\n".as_slice(), + b"* OK [UIDNOTSTICKY] invalid\r\nA0003 OK FETCH completed\r\n".as_slice(), + b"* NO [MESSAGELIMIT 1000] bounded\r\nA0003 OK FETCH completed\r\n".as_slice(), + b"* SEARCH 7\r\nA0003 OK FETCH completed\r\n".as_slice(), + b"* ESEARCH UID ALL 7\r\nA0003 OK FETCH completed\r\n".as_slice(), + ] { + rejected_fetch(response, true).await; + } +} + +#[tokio::test] +async fn body_must_be_full_and_literal_backed() { + for response in [ + b"* 7 UIDFETCH (UID 7 RFC822.SIZE 1 BODY[] \"x\")\r\nA0003 OK FETCH completed\r\n" + .as_slice(), + b"* 7 UIDFETCH (UID 7 RFC822.SIZE 1 BODY[]<0> {1}\r\nx)\r\nA0003 OK FETCH completed\r\n" + .as_slice(), + ] { + rejected_fetch(response, true).await; + } +} + +#[tokio::test] +async fn uidfetch_shape_accepts_requested_fields_in_either_order() { + let response = b"* 42 UIDFETCH (RFC822.SIZE 1 UID 42)\r\nA0003 OK inventory\r\n\ + * 42 UIDFETCH (RFC822.SIZE 1 UID 42 BODY[] {1}\r\nx)\r\nA0004 OK exact\r\n"; + let (mut session, handle) = enabled(response, UidOnlyLimits::default(), 1).await; + let inventory: Vec<_> = session + .uid_fetch("1:42", "(UID RFC822.SIZE) (PARTIAL 1:42)") + .await + .expect("inventory command") + .try_collect() + .await + .expect("inventory response"); + assert_eq!(inventory.len(), 1); + assert_eq!(inventory[0].message, 42); + assert_eq!(inventory[0].uid, Some(42)); + assert_eq!(inventory[0].size, Some(1)); + + handle + .arm_next_fetch_literal_limit(1) + .expect("arm body bound"); + let exact: Vec<_> = session + .uid_fetch("42", "(UID RFC822.SIZE BODY.PEEK[])") + .await + .expect("exact command") + .try_collect() + .await + .expect("exact response"); + assert_eq!(exact.len(), 1); + assert_eq!(exact[0].message, 42); + assert_eq!(exact[0].uid, Some(42)); + assert_eq!(exact[0].size, Some(1)); + assert_eq!(exact[0].body(), Some(b"x".as_slice())); +} + +#[tokio::test] +async fn uidfetch_shape_rejects_duplicates_mismatch_missing_and_extras() { + let cases: &[(&[u8], bool)] = &[ + ( + b"* 7 UIDFETCH (UID 7 UID 7 RFC822.SIZE 1)\r\nA0003 OK done\r\n", + false, + ), + ( + b"* 7 UIDFETCH (UID 7 RFC822.SIZE 1 RFC822.SIZE 1)\r\nA0003 OK done\r\n", + false, + ), + ( + b"* 7 UIDFETCH (UID 8 RFC822.SIZE 1)\r\nA0003 OK done\r\n", + false, + ), + ( + b"* 7 UIDFETCH (UID 7 FLAGS () RFC822.SIZE 1)\r\nA0003 OK done\r\n", + false, + ), + (b"* 7 UIDFETCH (UID 7)\r\nA0003 OK done\r\n", false), + ( + b"* 7 UIDFETCH (UID 7 RFC822.SIZE 1 BODY[] {1}\r\nx)\r\nA0003 OK done\r\n", + false, + ), + ( + b"* 7 UIDFETCH (UID 7 UID 7 RFC822.SIZE 1 BODY[] {1}\r\nx)\r\nA0003 OK done\r\n", + true, + ), + ( + b"* 7 UIDFETCH (UID 7 RFC822.SIZE 1 RFC822.SIZE 1 BODY[] {1}\r\nx)\r\nA0003 OK done\r\n", + true, + ), + ( + b"* 7 UIDFETCH (UID 8 RFC822.SIZE 1 BODY[] {1}\r\nx)\r\nA0003 OK done\r\n", + true, + ), + ( + b"* 7 UIDFETCH (UID 7 BODY[] {1}\r\nx)\r\nA0003 OK done\r\n", + true, + ), + ( + b"* 7 UIDFETCH (UID 7 RFC822.SIZE 1 FLAGS () BODY[] {1}\r\nx)\r\nA0003 OK done\r\n", + true, + ), + ( + b"* 7 UIDFETCH (UID 7 RFC822.SIZE 1 BODY[] BODY[] {1}\r\nx)\r\nA0003 OK done\r\n", + true, + ), + ( + b"* 7 UIDFETCH (UID 7 RFC822.SIZE 1 BODY[] {1}\r\nx FLAGS ())\r\nA0003 OK done\r\n", + true, + ), + ]; + for (response, exact) in cases { + rejected_fetch(response, *exact).await; + } +} + +#[tokio::test] +async fn tagged_no_is_not_mistaken_for_an_empty_success() { + rejected_fetch(b"A0003 NO [LIMIT] synthetic failure\r\n", false).await; +} + +#[tokio::test] +async fn zero_exists_and_recent_are_valid_fetch_notifications() { + let (mut session, handle) = enabled( + b"* 0 EXISTS\r\n* 0 RECENT\r\nA0003 OK FETCH completed\r\n", + UidOnlyLimits::default(), + 1, + ) + .await; + let stream = session + .uid_fetch("1:9", "(UID RFC822.SIZE) (PARTIAL 1:9)") + .await + .expect("command"); + assert!(stream.try_collect::>().await.unwrap().is_empty()); + handle.ensure_active().expect("adapter remains healthy"); +} + +#[tokio::test] +async fn literal_limits_and_second_literals_reject_before_excess_body_bytes() { + let mut limits = UidOnlyLimits::default(); + limits.max_literal_bytes = 4; + limits.max_response_bytes = 128; + limits.max_command_literal_bytes = 5; + limits.max_command_response_bytes = 512; + + let response = + b"* 7 UIDFETCH (UID 7 RFC822.SIZE 1 BODY[] {6}\r\nabcdef)\r\nA0003 OK FETCH completed\r\n"; + let (mut session, handle) = enabled(response, limits.clone(), 1).await; + handle + .arm_next_fetch_literal_limit(5) + .expect("arm body bound"); + let stream = session + .uid_fetch("7", "(UID RFC822.SIZE BODY.PEEK[])") + .await + .expect("command"); + assert!(stream.try_collect::>().await.is_err()); + assert_eq!(handle.literal_bytes_received(), 0); + + let response = b"* 7 UIDFETCH (UID 7 RFC822.SIZE 6 BODY[] {3}\r\nabc BODY[HEADER] {3}\r\ndef)\r\nA0003 OK FETCH completed\r\n"; + let (mut session, handle) = enabled(response, limits, 1).await; + handle + .arm_next_fetch_literal_limit(5) + .expect("arm body bound"); + let stream = session + .uid_fetch("7", "(UID RFC822.SIZE BODY.PEEK[])") + .await + .expect("command"); + assert!(stream.try_collect::>().await.is_err()); + assert_eq!(handle.literal_bytes_received(), 3); + assert!(handle + .poison_reason() + .expect("poison") + .contains("only closing")); +} + +#[tokio::test] +async fn command_runtime_is_bounded_while_the_server_is_silent() { + let transcript = b"* OK synthetic ready\r\nA0001 OK LOGIN completed\r\n\ + * ENABLED UIDONLY\r\nA0002 OK ENABLE completed\r\n"; + let mut limits = UidOnlyLimits::default(); + limits.max_command_runtime = Duration::from_millis(10); + let (stream, handle) = + UidOnlyStream::new(ScriptIo::pending(transcript, 2), limits).expect("adapter"); + let mut client = async_imap::Client::new(stream); + client + .read_response() + .await + .expect("greeting") + .expect("greeting"); + let mut session = client.login("synthetic", "redacted").await.expect("login"); + session + .run_command_and_check_ok("ENABLE UIDONLY") + .await + .expect("enable"); + handle + .arm_next_fetch_literal_limit(16) + .expect("arm body bound"); + let stream = session + .uid_fetch("7", "(UID RFC822.SIZE BODY.PEEK[])") + .await + .expect("command"); + let result = tokio::time::timeout(Duration::from_secs(1), stream.try_collect::>()) + .await + .expect("adapter deadline fired"); + assert!(result.is_err()); + assert!(handle + .poison_reason() + .expect("poison") + .contains("timed out")); +} + +#[tokio::test] +async fn response_count_is_bounded() { + let mut limits = UidOnlyLimits::default(); + limits.max_command_responses = 2; + let response = b"* 1 UIDFETCH (UID 1 RFC822.SIZE 1)\r\n\ + * 2 UIDFETCH (UID 2 RFC822.SIZE 1)\r\n\ + A0003 OK FETCH completed\r\n"; + let (mut session, handle) = enabled(response, limits, 4).await; + let stream = session + .uid_fetch("1:2", "(UID RFC822.SIZE) (PARTIAL 1:2)") + .await + .expect("command"); + assert!(stream.try_collect::>().await.is_err()); + assert!(handle + .poison_reason() + .expect("poison") + .contains("response count")); +} + +#[tokio::test] +async fn control_line_and_whole_response_bytes_are_bounded() { + let mut limits = UidOnlyLimits::default(); + limits.max_control_line_bytes = 64; + let response = b"* 7 UIDFETCH (UID 7 RFC822.SIZE 123456789 FLAGS (\\Seen \\Answered \\Flagged) BODY[] {1}\r\nx)\r\nA0003 OK done\r\n"; + let (mut session, handle) = enabled(response, limits, 2).await; + handle + .arm_next_fetch_literal_limit(1) + .expect("arm body bound"); + let stream = session + .uid_fetch("7", "(UID RFC822.SIZE BODY.PEEK[])") + .await + .expect("command"); + assert!(stream.try_collect::>().await.is_err()); + assert!(handle + .poison_reason() + .expect("poison") + .contains("control line")); + + let mut limits = UidOnlyLimits::default(); + limits.max_literal_bytes = 4; + limits.max_response_bytes = 36; + limits.max_command_literal_bytes = 4; + limits.max_command_response_bytes = 128; + let response = b"* 7 UIDFETCH (UID 7 RFC822.SIZE 4 BODY[] {4}\r\nabcd)\r\nA0003 OK done\r\n"; + let (mut session, handle) = enabled(response, limits, 2).await; + handle + .arm_next_fetch_literal_limit(4) + .expect("arm body bound"); + let stream = session + .uid_fetch("7", "(UID RFC822.SIZE BODY.PEEK[])") + .await + .expect("command"); + assert!(stream.try_collect::>().await.is_err()); + let reason = handle.poison_reason().expect("poison"); + assert!(reason.contains("response"), "{reason}"); +} + +#[tokio::test] +async fn one_in_flight_and_read_only_command_allowlist_are_enforced_before_write() { + let (mut session, handle) = enabled(b"", UidOnlyLimits::default(), 8).await; + handle + .arm_next_fetch_literal_limit(16) + .expect("arm body bound"); + session + .run_command("UID FETCH 1 (UID RFC822.SIZE BODY.PEEK[])") + .await + .expect("first command"); + assert!(session.run_command("LOGOUT").await.is_err()); + assert!(handle.poison_reason().is_some()); + + let (mut session, handle) = enabled(b"", UidOnlyLimits::default(), 8).await; + assert!(session + .run_command("UID STORE 1 +FLAGS (\\Deleted)") + .await + .is_err()); + assert!(handle.poison_reason().is_some()); +} + +#[test] +fn numeric_query_builders_reject_invalid_bounds() { + assert_eq!( + inventory_args(10, 20, 100).expect("valid inventory"), + ( + "10:20".to_string(), + "(UID RFC822.SIZE) (PARTIAL 1:100)".to_string() + ) + ); + assert!(inventory_args(0, 20, 100).is_err()); + assert!(inventory_args(20, 10, 100).is_err()); + assert!(inventory_args(10, 20, 0).is_err()); + assert!(exact_args(0).is_err()); +} diff --git a/crates/core/src/logger/mod.rs b/crates/core/src/logger/mod.rs index ed3ec3ed..64cdd356 100644 --- a/crates/core/src/logger/mod.rs +++ b/crates/core/src/logger/mod.rs @@ -37,7 +37,12 @@ impl FormatTime for LocalTimer { pub fn initialize_logging() { let level = validate_log_level(&SETTINGS.bichon_log_level); if matches!(level, Level::DEBUG) || matches!(level, Level::TRACE) { - LogTracer::init().unwrap(); + // async-imap trace events include raw commands and parser buffers, + // which may contain credentials or complete message literals. + LogTracer::builder() + .ignore_crate("async_imap") + .init() + .unwrap(); } if SETTINGS.bichon_log_to_file { setup_file_logger(level).unwrap(); From 9c6904373e84dd343fc6b10aa75e024529fd0277 Mon Sep 17 00:00:00 2001 From: Gabe A Date: Tue, 4 Aug 2026 10:33:45 -0400 Subject: [PATCH 2/4] add canonical uid-scoped storage primitives --- crates/blob/src/bucket.rs | 21 + crates/blob/src/engine.rs | 5 +- crates/core/src/account/migration.rs | 5 +- crates/core/src/envelope/extractor.rs | 586 +++++++++++++++++--- crates/core/src/mailbox/delete.rs | 7 +- crates/core/src/message/delete.rs | 8 +- crates/core/src/store/blob.rs | 220 +++++++- crates/core/src/store/tantivy/attachment.rs | 25 + crates/core/src/store/tantivy/dedup.rs | 47 +- crates/core/src/store/tantivy/envelope.rs | 129 ++++- 10 files changed, 957 insertions(+), 96 deletions(-) diff --git a/crates/blob/src/bucket.rs b/crates/blob/src/bucket.rs index d7432a25..4b4c82bc 100644 --- a/crates/blob/src/bucket.rs +++ b/crates/blob/src/bucket.rs @@ -163,6 +163,27 @@ impl IndexStore { self.get(key).map(|r| r.is_some()) } + /// Check several keys in one read transaction. + pub fn exists_batch(&self, keys: &[[u8; 32]]) -> Result> { + let txn = self + .db + .begin_read() + .map_err(|e| crate::error::Error::IndexDb(format!("read txn: {}", e)))?; + let table = txn + .open_table(INDEX_TABLE) + .map_err(|e| crate::error::Error::IndexDb(format!("open table: {}", e)))?; + keys.iter() + .map(|key| { + let record = table + .get(key) + .map_err(|e| crate::error::Error::IndexDb(format!("get: {}", e)))? + .map(|guard| IndexRecord::decode(&guard.value().0)) + .transpose()?; + Ok(record.is_some_and(|record| !record.is_tombstone())) + }) + .collect() + } + /// Insert or update a record for a key. Committed in a single write txn. pub fn insert(&self, record: &IndexRecord) -> Result<()> { let txn = self diff --git a/crates/blob/src/engine.rs b/crates/blob/src/engine.rs index 5656b5f5..04c31b4d 100644 --- a/crates/blob/src/engine.rs +++ b/crates/blob/src/engine.rs @@ -320,6 +320,10 @@ impl Engine { self.shared.index_store.exists(key) } + pub fn exists_batch(&self, keys: &[[u8; 32]]) -> Result> { + self.shared.index_store.exists_batch(keys) + } + // ── Batch delete ───────────────────────────────────────────────────── pub fn delete_batch(&self, keys: &[[u8; 32]]) -> Result<()> { @@ -375,7 +379,6 @@ impl Engine { let mut records: Vec = Vec::with_capacity(entries.len()); let mut ends: Vec<(u32, u64)> = Vec::with_capacity(entries.len()); - for (key, value, codec) in entries { if value.len() > crate::types::MAX_VALUE_SIZE { return Err(Error::ValueTooLarge { size: value.len() }); diff --git a/crates/core/src/account/migration.rs b/crates/core/src/account/migration.rs index ed0bcfae..531f701e 100644 --- a/crates/core/src/account/migration.rs +++ b/crates/core/src/account/migration.rs @@ -39,7 +39,7 @@ use crate::{ id, oauth2::token::OAuth2AccessToken, raise_error, - store::tantivy::{attachment::ATTACHMENT_MANAGER, envelope::ENVELOPE_MANAGER}, + store::tantivy::envelope::ENVELOPE_MANAGER, users::{payload::UserUpdateRequest, role::DEFAULT_ACCOUNT_MANAGER_ROLE_ID, UserModel}, utc_now, }; @@ -502,9 +502,6 @@ impl Account { ENVELOPE_MANAGER .delete_account_envelopes(account.id) .await?; - ATTACHMENT_MANAGER - .delete_account_attachments(account.id) - .await?; Self::delete_account(account)?; info!("Sequential cleanup completed for account: {}", account.id); Ok(()) diff --git a/crates/core/src/envelope/extractor.rs b/crates/core/src/envelope/extractor.rs index 88195a72..700868d7 100644 --- a/crates/core/src/envelope/extractor.rs +++ b/crates/core/src/envelope/extractor.rs @@ -25,10 +25,11 @@ use crate::error::code::ErrorCode; use crate::error::BichonResult; use crate::imap::executor::ImapExecutor; use crate::message::content::AttachmentInfo; -use crate::store::blob::{DetachedEmail, BLOB_MANAGER}; +use crate::store::blob::{DetachedEmail, UidOnlyBlob, BLOB_MANAGER}; use crate::store::tantivy::attachment::ATTACHMENT_MANAGER; +use crate::store::tantivy::dedup::UIDONLY_SHARD_BIT; use crate::store::tantivy::dedup_cache::DEDUP_CACHE; -use crate::store::tantivy::envelope::ENVELOPE_MANAGER; +use crate::store::tantivy::envelope::{CANONICAL_STORAGE_GATE, ENVELOPE_MANAGER}; use crate::store::tantivy::model::{AttachmentModel, EnvelopeWithAttachments}; use crate::utils::html::extract_text; use crate::utils::{compute_content_hash, hex_hash}; @@ -37,11 +38,82 @@ use crate::{raise_error, utc_now}; use async_imap::types::Fetch; use bytes::Bytes; use mail_parser::{Address, HeaderName, Message, MessageParser, MimeHeaders}; -use tantivy::TantivyDocument; use tantivy::schema::Facet; +use tantivy::TantivyDocument; use tracing::error; use uuid::Uuid; +pub(crate) const UIDONLY_PROJECTION_BATCH_MESSAGES: usize = 32; +pub(crate) const UIDONLY_PROJECTION_BATCH_BYTES: usize = 64 * 1024 * 1024; + +pub(crate) struct UidOnlyMessage { + pub body: Vec, + pub uid: u32, +} + +struct PreparedEnvelope { + envelope_id: String, + content_hash: String, + detached_email: DetachedEmail, + attachment_docs: Vec, + envelope_doc: TantivyDocument, +} + +pub(crate) fn uidonly_envelope_id( + account_id: u64, + mailbox_id: u64, + uid_validity: u32, + uid: u32, + source_scope: &str, +) -> String { + let mut hasher = blake3::Hasher::new(); + hasher.update(b"bichon-uidonly-envelope-v1\0"); + hasher.update(&account_id.to_be_bytes()); + hasher.update(&mailbox_id.to_be_bytes()); + hasher.update(&uid_validity.to_be_bytes()); + hasher.update(&uid.to_be_bytes()); + hasher.update(source_scope.as_bytes()); + format!("uidonly-{}", hasher.finalize().to_hex()) +} + +pub(crate) fn verify_uidonly_projections( + account_id: u64, + mailbox_id: u64, + uid_validity: u32, + uids: &[u32], + source_scope: &str, +) -> BichonResult> { + let identities: Vec<_> = uids + .iter() + .map(|uid| { + ( + *uid, + uidonly_envelope_id(account_id, mailbox_id, uid_validity, *uid, source_scope), + ) + }) + .collect(); + let envelope_ids: Vec<_> = identities.iter().map(|(_, id)| id.clone()).collect(); + let markers = ENVELOPE_MANAGER.get_uidonly_projection_markers(account_id, &envelope_ids)?; + let candidates: Vec<_> = identities + .iter() + .enumerate() + .filter_map(|(index, (uid, envelope_id))| { + let marker = markers.get(envelope_id)?; + (marker.mailbox_id == mailbox_id + && marker.uid == *uid + && marker.shard_id == UIDONLY_SHARD_BIT) + .then(|| (index, marker.content_hash.clone())) + }) + .collect(); + let hashes: Vec<_> = candidates.iter().map(|(_, hash)| hash.clone()).collect(); + let exists = BLOB_MANAGER.has_uidonly_exact_batch(&hashes)?; + let mut verified = vec![false; uids.len()]; + for ((index, _), exists) in candidates.into_iter().zip(exists) { + verified[index] = exists; + } + Ok(verified) +} + pub async fn extract_envelope_and_store_it( fetch: Fetch, account_id: u64, @@ -64,7 +136,9 @@ pub async fn extract_envelope_and_store_it( } }; let size = fetch.size.unwrap_or(body.len() as u32); - extract_envelope_core(body, uid, size, internal_date, account_id, mailbox_id).await + extract_envelope_core(body, uid, size, internal_date, account_id, mailbox_id) + .await + .map(|_| ()) } pub async fn extract_envelope_from_eml( @@ -72,7 +146,9 @@ pub async fn extract_envelope_from_eml( account_id: u64, mailbox_id: u64, ) -> BichonResult<()> { - extract_envelope_core(body, 0, body.len() as u32, 0, account_id, mailbox_id).await + extract_envelope_core(body, 0, body.len() as u32, 0, account_id, mailbox_id) + .await + .map(|_| ()) } pub async fn extract_envelope_from_smtp( @@ -89,6 +165,91 @@ pub async fn extract_envelope_from_smtp( mailbox_id, ) .await + .map(|_| ()) +} + +pub(crate) async fn project_uidonly_messages( + messages: Vec, + account_id: u64, + mailbox_id: u64, + uid_validity: u32, + source_scope: &str, +) -> BichonResult<()> { + let total_bytes = messages.iter().try_fold(0usize, |total, message| { + total.checked_add(message.body.len()).ok_or_else(|| { + raise_error!( + "UIDONLY projection batch byte count overflow".into(), + ErrorCode::PayloadTooLarge + ) + }) + })?; + if messages.len() > UIDONLY_PROJECTION_BATCH_MESSAGES + || (messages.len() > 1 && total_bytes > UIDONLY_PROJECTION_BATCH_BYTES) + { + return Err(raise_error!( + "UIDONLY projection batch exceeds its memory bound".into(), + ErrorCode::PayloadTooLarge + )); + } + + let mut blobs = Vec::new(); + let mut attachment_batches = Vec::new(); + let mut envelope_documents = Vec::new(); + for message in messages { + let envelope_id = uidonly_envelope_id( + account_id, + mailbox_id, + uid_validity, + message.uid, + source_scope, + ); + let size = u32::try_from(message.body.len()).map_err(|_| { + raise_error!( + "UIDONLY message is too large to project".into(), + ErrorCode::PayloadTooLarge + ) + })?; + match prepare_envelope_core( + &message.body, + message.uid, + size, + 0, + account_id, + mailbox_id, + Some((envelope_id, UIDONLY_SHARD_BIT)), + ) + .await? + { + None => { + return Err(raise_error!( + "UIDONLY message was unexpectedly filtered".into(), + ErrorCode::InternalError + )) + } + Some(prepared) => { + let envelope_id = prepared.envelope_id.clone(); + blobs.push(UidOnlyBlob { + content_hash: prepared.content_hash, + raw: message.body, + attachments: prepared.detached_email.attachments.unwrap_or_default(), + }); + attachment_batches.push((envelope_id.clone(), prepared.attachment_docs)); + envelope_documents.push((envelope_id, prepared.envelope_doc)); + } + } + } + + if !blobs.is_empty() { + let _canonical_guard = CANONICAL_STORAGE_GATE.lock().await; + BLOB_MANAGER.store_uidonly_exact_batch(blobs).await?; + ATTACHMENT_MANAGER + .replace_uidonly_batches(attachment_batches) + .await?; + ENVELOPE_MANAGER + .replace_uidonly_documents(envelope_documents) + .await?; + } + Ok(()) } async fn extract_envelope_core( @@ -99,44 +260,95 @@ async fn extract_envelope_core( account_id: u64, mailbox_id: u64, ) -> BichonResult<()> { + match prepare_envelope_core(body, uid, size, internal_date, account_id, mailbox_id, None) + .await? + { + None => Ok(()), + Some(prepared) => { + BLOB_MANAGER.queue(prepared.detached_email).await; + ENVELOPE_MANAGER.queue(prepared.envelope_doc).await; + DEDUP_CACHE.insert(account_id, mailbox_id, &prepared.content_hash); + for doc in prepared.attachment_docs { + ATTACHMENT_MANAGER.queue(doc).await; + } + Ok(()) + } + } +} + +async fn prepare_envelope_core( + body: &[u8], + uid: u32, + size: u32, + internal_date: i64, + account_id: u64, + mailbox_id: u64, + uidonly_identity: Option<(String, u64)>, +) -> BichonResult> { //The content hash of the original raw EML let email_content_hash = compute_content_hash(body); - if DEDUP_CACHE.contains(account_id, mailbox_id, &email_content_hash) { + if uidonly_identity.is_none() + && DEDUP_CACHE.contains(account_id, mailbox_id, &email_content_hash) + { tracing::debug!("Duplicate email detected"); //println!("Duplicate email detected"); - return Ok(()); + return Ok(None); } - let message: Message<'_> = MessageParser::new().parse(body).ok_or_else(|| { - raise_error!( - "Email header parse result is not available".into(), - ErrorCode::InternalError - ) - })?; + let is_uidonly = uidonly_identity.is_some(); + let message: Message<'_> = match MessageParser::new().parse(body) { + Some(message) => message, + None if is_uidonly => { + return prepare_unparseable_uidonly( + body, + uid, + internal_date, + account_id, + mailbox_id, + uidonly_identity.expect("UIDONLY identity checked above"), + ) + .await; + } + None => { + return Err(raise_error!( + "Email header parse result is not available".into(), + ErrorCode::InternalError + )); + } + }; - if let Ok(account) = AccountModel::get(account_id) { - if let Some(ref rules) = account.archive_rules { - let sender = message.from().and_then(|addr| { - AddrVec::from(addr).0.into_iter().next().and_then(|a| a.address) - }); - let subject = message.subject().map(|s| s.to_string()); - - let is_spam = !rules.spam_headers.is_empty() - && rules.spam_headers.iter().any(|h| { - message - .header_raw(h.clone()) - .map(|v| matches!(v.trim().to_lowercase().as_str(), "yes" | "true")) - .unwrap_or(false) + // UIDONLY is a complete-mailbox export path. Its caller rejects enabled + // archive rules, so do not let a mid-run configuration change silently + // turn later UIDs into filtered receipts. + if !is_uidonly { + if let Ok(account) = AccountModel::get(account_id) { + if let Some(ref rules) = account.archive_rules { + let sender = message.from().and_then(|addr| { + AddrVec::from(addr) + .0 + .into_iter() + .next() + .and_then(|a| a.address) }); - - if !rules.should_archive(sender.as_deref(), subject.as_deref(), size, is_spam) { - tracing::debug!( - account_id, - uid, - sender = sender.as_deref().unwrap_or("?"), - subject = subject.as_deref().unwrap_or("?"), - "Email filtered out by archive rules" - ); - return Ok(()); + let subject = message.subject().map(|s| s.to_string()); + + let is_spam = !rules.spam_headers.is_empty() + && rules.spam_headers.iter().any(|h| { + message + .header_raw(h.clone()) + .map(|v| matches!(v.trim().to_lowercase().as_str(), "yes" | "true")) + .unwrap_or(false) + }); + + if !rules.should_archive(sender.as_deref(), subject.as_deref(), size, is_spam) { + tracing::debug!( + account_id, + uid, + sender = sender.as_deref().unwrap_or("?"), + subject = subject.as_deref().unwrap_or("?"), + "Email filtered out by archive rules" + ); + return Ok(None); + } } } } @@ -150,7 +362,11 @@ async fn extract_envelope_core( String::new() }; - let text = text.split_whitespace().collect::>().join(" "); + let text = if is_uidonly { + normalize_whitespace_bounded(&text, 16 * 1024 * 1024) + } else { + text.split_whitespace().collect::>().join(" ") + }; let preview = if text.chars().count() > preview_limit { text.chars().take(preview_limit).collect::() + "..." @@ -202,10 +418,29 @@ async fn extract_envelope_core( .and_then(|add| add.address) .unwrap_or_else(|| "unknown".to_string()); let attachment_count = message.attachment_count(); - let attachments = detach_and_store_attachments(body, &message, &email_content_hash, account_id, mailbox_id).await; + let (attachments, detached_email) = prepare_detached_attachments( + body, + &message, + &email_content_hash, + account_id, + mailbox_id, + !is_uidonly, + ) + .await; - let envelope_id = Uuid::new_v4().to_string(); + let (envelope_id, shard_id) = + uidonly_identity.unwrap_or_else(|| (Uuid::new_v4().to_string(), 0)); let now = utc_now!(); + let stored_size = if is_uidonly { + u32::try_from(body.len()).map_err(|_| { + raise_error!( + "UIDONLY message size exceeds Bichon's envelope size field".into(), + ErrorCode::PayloadTooLarge + ) + })? + } else { + size + }; let mut final_tags = Vec::new(); @@ -274,7 +509,7 @@ async fn extract_envelope_core( .collect(); let envelope = Envelope { - id: envelope_id, + id: envelope_id.clone(), message_id, account_id, mailbox_id, @@ -288,7 +523,7 @@ async fn extract_envelope_core( date, internal_date, ingest_at: now, - size, + size: stored_size, thread_id, attachment_count, regular_attachment_count: attachment_docs.len(), @@ -303,7 +538,7 @@ async fn extract_envelope_core( envelope, attachments: Some(attachments), }; - let doc = ea.to_document(&body_text, 0)?; + let doc = ea.to_document(&body_text, shard_id)?; tracing::debug!( "[account {}][mailbox {}] extract: uid={} msg_id={} content_hash={}", account_id, @@ -312,12 +547,80 @@ async fn extract_envelope_core( &ea.envelope.message_id, &ea.envelope.content_hash, ); - ENVELOPE_MANAGER.queue(doc).await; - DEDUP_CACHE.insert(account_id, mailbox_id, &email_content_hash); - for doc in attachment_docs { - ATTACHMENT_MANAGER.queue(doc).await; + Ok(Some(PreparedEnvelope { + envelope_id, + content_hash: email_content_hash, + detached_email, + attachment_docs, + envelope_doc: doc, + })) +} + +fn normalize_whitespace_bounded(input: &str, max_bytes: usize) -> String { + let mut output = String::with_capacity(input.len().min(max_bytes)); + for word in input.split_whitespace() { + let separator = usize::from(!output.is_empty()); + if output + .len() + .checked_add(separator) + .and_then(|len| len.checked_add(word.len())) + .is_none_or(|len| len > max_bytes) + { + break; + } + if separator == 1 { + output.push(' '); + } + output.push_str(word); } - Ok(()) + output +} + +async fn prepare_unparseable_uidonly( + body: &[u8], + uid: u32, + internal_date: i64, + account_id: u64, + mailbox_id: u64, + identity: (String, u64), +) -> BichonResult> { + let (envelope_id, shard_id) = identity; + let size = u32::try_from(body.len()).map_err(|_| { + raise_error!( + "UIDONLY message size exceeds Bichon's envelope size field".into(), + ErrorCode::PayloadTooLarge + ) + })?; + let content_hash = compute_content_hash(body); + let envelope = Envelope { + id: envelope_id.clone(), + message_id: format!("<{}@uidonly.bichon.invalid>", &envelope_id), + account_id, + mailbox_id, + uid, + from: "unknown".into(), + internal_date, + ingest_at: utc_now!(), + size, + thread_id: hex_hash(&envelope_id), + content_hash: content_hash.clone(), + ..Envelope::default() + }; + let doc = EnvelopeWithAttachments { + envelope, + attachments: Some(Vec::new()), + } + .to_document("", shard_id)?; + Ok(Some(PreparedEnvelope { + envelope_id, + content_hash: content_hash.clone(), + detached_email: DetachedEmail { + email: (content_hash, Bytes::new()), + attachments: Some(Vec::new()), + }, + attachment_docs: Vec::new(), + envelope_doc: doc, + })) } pub fn extract_envelope_from_nested_message( @@ -426,13 +729,14 @@ pub fn extract_references(message: &Message<'_>) -> Option> { } } -pub async fn detach_and_store_attachments( +async fn prepare_detached_attachments( original_body: &[u8], message: &Message<'_>, eml_content_hash: &str, account_id: u64, mailbox_id: u64, -) -> Vec { + legacy_storage: bool, +) -> (Vec, DetachedEmail) { let rules = if account_id > 0 { AccountModel::get(account_id) .ok() @@ -451,7 +755,10 @@ pub async fn detach_and_store_attachments( .and_then(|addr| AddrVec::from(addr).0.into_iter().next()) .and_then(|add| add.address); - let mut stripped_eml = original_body.to_vec(); + let mut stripped_eml = legacy_storage + .then(|| original_body.to_vec()) + .unwrap_or_default(); + let shared_body = (!legacy_storage).then(|| Bytes::copy_from_slice(original_body)); let mut attachment_infos = Vec::new(); // Step 1: Collect and sort attachment ranges in reverse to maintain offset integrity let mut ranges: Vec<_> = message @@ -492,15 +799,19 @@ pub async fn detach_and_store_attachments( if range_valid { let raw_bytes = &original_body[raw_start..raw_end]; // The actual content stored in the blob is the raw undecoded data. - attachments.push((content_hash.clone(), Bytes::copy_from_slice(raw_bytes))); - - // Replace raw attachment content with a hash-based placeholder - let placeholder = format!("<>", &content_hash); - stripped_eml.splice(raw_start..raw_end, placeholder.as_bytes().iter().cloned()); + let bytes = shared_body + .as_ref() + .map(|body| body.slice(raw_start..raw_end)) + .unwrap_or_else(|| Bytes::copy_from_slice(raw_bytes)); + attachments.push((content_hash.clone(), bytes)); + if legacy_storage { + // Replace raw attachment content with a hash-based placeholder + let placeholder = format!("<>", &content_hash); + stripped_eml.splice(raw_start..raw_end, placeholder.as_bytes().iter().cloned()); + } } else { - // Invalid range: store a zero-length blob so the consistency - // check passes; reattachment will log a warning for the missing - // blob data but won't panic. + // Exact raw remains canonical when a malformed attachment range + // cannot be sliced. Keep the legacy zero-length compatibility blob. attachments.push((content_hash.clone(), Bytes::new())); } @@ -541,7 +852,8 @@ pub async fn detach_and_store_attachments( if !inline || !has_cid { let decoded_len = att.contents().len(); - if should_extract + if legacy_storage + && should_extract && decoded_len <= crate::ext::text_extractor::MAX_EXTRACT_BYTES && crate::ext::text_extractor::should_try_extract(&file_type, &ext) { @@ -597,15 +909,33 @@ pub async fn detach_and_store_attachments( } } } - // Step 4: Store the final stripped EML content - BLOB_MANAGER - .queue(DetachedEmail { + ( + attachment_infos, + DetachedEmail { email: (eml_content_hash.to_string(), Bytes::from(stripped_eml)), attachments: Some(attachments), - }) - .await; + }, + ) +} - attachment_infos +pub async fn detach_and_store_attachments( + original_body: &[u8], + message: &Message<'_>, + eml_content_hash: &str, + account_id: u64, + mailbox_id: u64, +) -> Vec { + let (infos, detached) = prepare_detached_attachments( + original_body, + message, + eml_content_hash, + account_id, + mailbox_id, + true, + ) + .await; + BLOB_MANAGER.queue(detached).await; + infos } pub fn reattach_eml_content( @@ -626,7 +956,7 @@ pub fn reattach_eml_content( })?; let restored_eml = BLOB_MANAGER - .get_email(&e.envelope.content_hash)? + .get_canonical_email(&e.envelope.content_hash)? .ok_or_else(|| { raise_error!( format!( @@ -637,7 +967,11 @@ pub fn reattach_eml_content( ) })?; - if !e.envelope.has_any_attachments() { + // UIDONLY stores the complete raw message, while legacy entries store a + // detached representation that still needs attachment substitution. + if compute_content_hash(&restored_eml) == e.envelope.content_hash + || !e.envelope.has_any_attachments() + { return Ok((e.envelope, restored_eml)); } @@ -647,7 +981,7 @@ pub fn reattach_eml_content( return Err(raise_error!( format!( "Consistency check failed: envelope.attachment_count ({}) does not match attachments.len ({})", - e.envelope.attachment_count, + e.envelope.attachment_count, actual_count ), ErrorCode::InternalError @@ -716,7 +1050,10 @@ pub async fn reattach_eml_content_self_healing( .envelope; // Fast path: the content blob is present, reuse the regular reattach logic. - if BLOB_MANAGER.get_email(&envelope.content_hash)?.is_some() { + if BLOB_MANAGER + .get_canonical_email(&envelope.content_hash)? + .is_some() + { return reattach_eml_content(account_id, envelope_id); } @@ -802,6 +1139,75 @@ async fn recover_message_blob(envelope: &Envelope) -> BichonResult { mod test { use html2text::config; + #[test] + fn uidonly_identity_is_stable_and_source_scoped() { + let first = super::uidonly_envelope_id(7, 11, 13, 17, "imap.example:993\0alice"); + assert_eq!( + first, + super::uidonly_envelope_id(7, 11, 13, 17, "imap.example:993\0alice") + ); + assert_ne!( + first, + super::uidonly_envelope_id(7, 11, 13, 17, "imap.example:993\0bob") + ); + assert_ne!( + first, + super::uidonly_envelope_id(7, 11, 13, 17, "export.example:993\0alice") + ); + } + + #[test] + fn uidonly_search_text_normalization_is_bounded() { + assert_eq!( + super::normalize_whitespace_bounded(" a\n b c ", 5), + "a b c" + ); + assert_eq!( + super::normalize_whitespace_bounded("alpha beta", 5), + "alpha" + ); + } + + #[tokio::test] + async fn uidonly_projection_batch_is_verified_and_idempotent() { + let raw = b"From: sender@example.invalid\r\nTo: archive@example.invalid\r\n\ + Subject: fixture\r\n\r\nbody\r\n"; + let messages = || { + [17, 18] + .into_iter() + .map(|uid| super::UidOnlyMessage { + body: raw.to_vec(), + uid, + }) + .collect() + }; + super::project_uidonly_messages(messages(), 7, 11, 13, "source-a") + .await + .unwrap(); + super::project_uidonly_messages(messages(), 7, 11, 13, "source-a") + .await + .unwrap(); + assert_eq!( + super::verify_uidonly_projections(7, 11, 13, &[17, 18, 19], "source-a").unwrap(), + [true, true, false] + ); + } + + #[tokio::test] + async fn uidonly_projection_batch_count_is_bounded() { + let messages = (0..=super::UIDONLY_PROJECTION_BATCH_MESSAGES) + .map(|uid| super::UidOnlyMessage { + body: Vec::new(), + uid: uid as u32, + }) + .collect(); + assert!( + super::project_uidonly_messages(messages, 7, 11, 13, "source-a") + .await + .is_err() + ); + } + #[test] fn test_various_html_with_overflow_enabled() { let cases = [ @@ -899,4 +1305,46 @@ mod test { // in reattach_eml_content doesn't fail later. assert_eq!(infos.len(), 1); } + + #[tokio::test] + async fn uidonly_nested_attachments_use_bounded_shared_slices() { + let raw = b"From: outer@example.invalid\r\n\ +MIME-Version: 1.0\r\n\ +Content-Type: multipart/mixed; boundary=outer\r\n\r\n\ +--outer\r\n\ +Content-Type: message/rfc822\r\n\ +Content-Disposition: attachment; filename=forwarded.eml\r\n\r\n\ +From: inner@example.invalid\r\n\ +MIME-Version: 1.0\r\n\ +Content-Type: multipart/mixed; boundary=inner\r\n\r\n\ +--inner\r\n\ +Content-Type: application/octet-stream\r\n\ +Content-Disposition: attachment; filename=data.bin\r\n\ +Content-Transfer-Encoding: base64\r\n\r\n\ +YWJjZA==\r\n\ +--inner--\r\n\ +--outer\r\n\ +Content-Type: application/octet-stream\r\n\ +Content-Disposition: attachment; filename=outer.bin\r\n\ +Content-Transfer-Encoding: base64\r\n\r\n\ +ZWZnaA==\r\n\ +--outer--\r\n"; + let message = super::MessageParser::new().parse(raw).unwrap(); + let ranges: Vec<_> = message + .attachments() + .map(|part| (part.raw_body_offset(), part.raw_end_offset())) + .collect(); + assert!(ranges.len() >= 2); + + let (_, detached) = super::prepare_detached_attachments( + raw, + &message, + &super::compute_content_hash(raw), + 0, + 0, + false, + ) + .await; + assert_eq!(detached.attachments.unwrap().len(), ranges.len()); + } } diff --git a/crates/core/src/mailbox/delete.rs b/crates/core/src/mailbox/delete.rs index 574323cf..c5922113 100644 --- a/crates/core/src/mailbox/delete.rs +++ b/crates/core/src/mailbox/delete.rs @@ -19,7 +19,7 @@ use crate::{ cache::imap::mailbox::MailBox, error::BichonResult, - store::tantivy::{attachment::ATTACHMENT_MANAGER, envelope::ENVELOPE_MANAGER}, + store::tantivy::envelope::ENVELOPE_MANAGER, }; pub async fn delete_mailbox_impl(account_id: u64, mailbox_id: u64) -> BichonResult<()> { @@ -45,10 +45,7 @@ pub async fn delete_mailbox_impl(account_id: u64, mailbox_id: u64) -> BichonResu } ENVELOPE_MANAGER - .delete_mailbox_envelopes(account_id, ids_to_delete.clone()) - .await?; - ATTACHMENT_MANAGER - .delete_mailbox_attachments(account_id, ids_to_delete.clone()) + .delete_mailbox_envelopes(account_id, ids_to_delete) .await?; Ok(()) } diff --git a/crates/core/src/message/delete.rs b/crates/core/src/message/delete.rs index 20185d0f..1ef8cb7d 100644 --- a/crates/core/src/message/delete.rs +++ b/crates/core/src/message/delete.rs @@ -17,15 +17,9 @@ // along with this program. If not, see . use crate::error::BichonResult; -use crate::store::tantivy::attachment::ATTACHMENT_MANAGER; use crate::store::tantivy::envelope::ENVELOPE_MANAGER; use std::collections::HashMap; pub async fn delete_messages_impl(request: HashMap>) -> BichonResult<()> { - ENVELOPE_MANAGER - .delete_envelopes_multi_account(request.clone()) - .await?; - ATTACHMENT_MANAGER - .delete_attachments_multi_account(request) - .await + ENVELOPE_MANAGER.delete_envelopes_multi_account(request).await } diff --git a/crates/core/src/store/blob.rs b/crates/core/src/store/blob.rs index 145b4959..40b05b46 100644 --- a/crates/core/src/store/blob.rs +++ b/crates/core/src/store/blob.rs @@ -22,11 +22,17 @@ use crate::{ envelope::extractor::reattach_eml_content_self_healing, error::{code::ErrorCode, BichonResult}, settings::dir::DATA_DIR_MANAGER, + utils::compute_content_hash, }; use bichon_blob::{Codec, Config, Engine}; use bytes::Bytes; -use std::{io::Cursor, sync::Arc, sync::LazyLock}; +use std::{ + collections::{hash_map::Entry, HashMap, HashSet}, + io::Cursor, + sync::Arc, + sync::LazyLock, +}; use tokio::{ sync::{mpsc, Mutex}, task::{self, JoinHandle}, @@ -34,11 +40,24 @@ use tokio::{ pub static BLOB_MANAGER: LazyLock = LazyLock::new(BlobManager::new); +pub(crate) fn uidonly_exact_raw_blob_key(content_hash: &str) -> String { + let mut hasher = blake3::Hasher::new(); + hasher.update(b"bichon-uidonly-exact-raw-v1\0"); + hasher.update(content_hash.as_bytes()); + hex::encode(hasher.finalize().as_bytes()) +} + pub struct DetachedEmail { pub email: (String, Bytes), pub attachments: Option>, } +pub(crate) struct UidOnlyBlob { + pub content_hash: String, + pub raw: Vec, + pub attachments: Vec<(String, Bytes)>, +} + pub struct BlobManager { sender: mpsc::Sender, engine: Arc, @@ -56,6 +75,104 @@ fn hex_to_key(hex: &str) -> BichonResult<[u8; 32]> { Ok(key) } +fn insert_uidonly_blob( + values: &mut HashMap<[u8; 32], Vec>, + order: &mut Vec<[u8; 32]>, + key: [u8; 32], + value: Vec, + reject_mismatch: bool, +) -> BichonResult<()> { + match values.entry(key) { + Entry::Vacant(entry) => { + order.push(key); + entry.insert(value); + } + Entry::Occupied(entry) if reject_mismatch && entry.get() != &value => { + return Err(raise_error!( + "one UIDONLY blob key mapped to different bytes in a batch".into(), + ErrorCode::InternalError + )); + } + Entry::Occupied(_) => {} + } + Ok(()) +} + +fn store_uidonly_exact_batch_inner(engine: &Engine, blobs: Vec) -> BichonResult<()> { + let mut values = HashMap::new(); + let mut order = Vec::new(); + let mut exact_keys = HashSet::new(); + let mut exact_readbacks = Vec::new(); + + for blob in blobs { + if compute_content_hash(&blob.raw) != blob.content_hash { + return Err(raise_error!( + "UIDONLY raw bytes do not match their content hash".into(), + ErrorCode::InternalError + )); + } + let exact_key = hex_to_key(&uidonly_exact_raw_blob_key(&blob.content_hash))?; + if exact_keys.insert(exact_key) { + exact_readbacks.push((exact_key, blob.content_hash.clone(), blob.raw.len())); + } + insert_uidonly_blob(&mut values, &mut order, exact_key, blob.raw, true)?; + for (hash, bytes) in blob.attachments { + // Attachment keys predate UIDONLY and are based on decoded bytes, + // while their stored MIME slices may differ in transfer encoding. + // Preserve the first value, matching the legacy dedup behavior. + insert_uidonly_blob( + &mut values, + &mut order, + hex_to_key(&hash)?, + bytes.to_vec(), + false, + )?; + } + } + + let existing = engine + .exists_batch(&order) + .map_err(|e| raise_error!(format!("{e:#?}"), ErrorCode::InternalError))?; + let entries: Vec<_> = order + .iter() + .copied() + .zip(existing) + .filter(|(_, exists)| !exists) + .map(|key| { + let key = key.0; + ( + key, + values.remove(&key).expect("ordered UIDONLY blob"), + Codec::Lz4, + ) + }) + .collect(); + engine + .put_batch(&entries) + .map_err(|e| raise_error!(format!("{e:#?}"), ErrorCode::InternalError))?; + drop(entries); + drop(values); + + for (key, content_hash, size) in exact_readbacks { + let stored = engine + .get(&key) + .map_err(|e| raise_error!(format!("{e:#?}"), ErrorCode::InternalError))? + .ok_or_else(|| { + raise_error!( + "UIDONLY exact raw blob missing after durable write".into(), + ErrorCode::InternalError + ) + })?; + if stored.len() != size || compute_content_hash(&stored) != content_hash { + return Err(raise_error!( + "UIDONLY exact raw blob failed readback verification".into(), + ErrorCode::InternalError + )); + } + } + Ok(()) +} + impl BlobManager { pub async fn shutdown(&self) { let mut guard = self.handle.lock().await; @@ -195,10 +312,57 @@ impl BlobManager { } } + /// Durably stores a batch of exact RFC822 messages, then reads every unique + /// raw blob back before any search-index completion marker may be written. + pub(crate) async fn store_uidonly_exact_batch( + &self, + blobs: Vec, + ) -> BichonResult<()> { + let engine = Arc::clone(&self.engine); + tokio::task::spawn_blocking(move || store_uidonly_exact_batch_inner(&engine, blobs)) + .await + .map_err(|e| raise_error!(format!("{e:#?}"), ErrorCode::InternalError))? + } + pub fn get_email(&self, content_hash: &str) -> BichonResult> { self.get(content_hash) } + pub(crate) fn get_uidonly_exact(&self, content_hash: &str) -> BichonResult> { + let Some(raw) = self.get(&uidonly_exact_raw_blob_key(content_hash))? else { + return Ok(None); + }; + if compute_content_hash(&raw) != content_hash { + return Err(raise_error!( + "UIDONLY exact raw blob digest mismatch".into(), + ErrorCode::InternalError + )); + } + Ok(Some(raw)) + } + + /// Fast restart check. This is a receipt only when the caller has already + /// found a valid envelope marker committed after the initial full readback. + pub(crate) fn has_uidonly_exact_batch( + &self, + content_hashes: &[String], + ) -> BichonResult> { + let keys: Result, _> = content_hashes + .iter() + .map(|hash| hex_to_key(&uidonly_exact_raw_blob_key(hash))) + .collect(); + self.engine + .exists_batch(&keys?) + .map_err(|e| raise_error!(format!("{e:#?}"), ErrorCode::InternalError)) + } + + pub(crate) fn get_canonical_email(&self, content_hash: &str) -> BichonResult> { + match self.get_uidonly_exact(content_hash)? { + Some(raw) => Ok(Some(raw)), + None => self.get(content_hash), + } + } + pub fn get_attachment(&self, content_hash: &str) -> BichonResult> { self.get(content_hash) } @@ -222,10 +386,11 @@ impl BlobManager { I2: IntoIterator, I2::Item: AsRef, { - let mut keys: Vec<[u8; 32]> = email_content_hashes - .into_iter() - .map(|h| hex_to_key(h.as_ref())) - .collect::>()?; + let mut keys = Vec::new(); + for hash in email_content_hashes { + keys.push(hex_to_key(hash.as_ref())?); + keys.push(hex_to_key(&uidonly_exact_raw_blob_key(hash.as_ref()))?); + } for h in attachment_content_hashes { keys.push(hex_to_key(h.as_ref())?); @@ -241,6 +406,51 @@ impl BlobManager { } } +#[cfg(test)] +mod tests { + use super::{store_uidonly_exact_batch_inner, uidonly_exact_raw_blob_key, UidOnlyBlob}; + use crate::utils::compute_content_hash; + use bichon_blob::{Codec, Config, Engine}; + use bytes::Bytes; + + #[test] + fn exact_raw_write_is_durable_verified_and_idempotent() { + let dir = std::env::temp_dir().join(format!("bichon-uidonly-{}", uuid::Uuid::new_v4())); + let engine = Engine::open(&dir, Config::default()).unwrap(); + let raw = b"From: sender@example.invalid\r\n\r\nbody\r\n"; + let hash = compute_content_hash(raw); + let exact_hash = uidonly_exact_raw_blob_key(&hash); + assert_ne!(exact_hash, hash); + assert_eq!(exact_hash, uidonly_exact_raw_blob_key(&hash)); + let attachment = Bytes::from_static(b"attachment"); + let attachment_hash = compute_content_hash(&attachment); + let key = super::hex_to_key(&hash).unwrap(); + engine + .put(key, b"legacy detached value", Codec::Lz4) + .unwrap(); + let blob = || UidOnlyBlob { + content_hash: hash.clone(), + raw: raw.to_vec(), + attachments: vec![(attachment_hash.clone(), attachment.clone())], + }; + store_uidonly_exact_batch_inner(&engine, vec![blob(), blob()]).unwrap(); + store_uidonly_exact_batch_inner(&engine, vec![blob()]).unwrap(); + assert_eq!(engine.get(&key).unwrap().unwrap(), b"legacy detached value"); + let exact_key = super::hex_to_key(&uidonly_exact_raw_blob_key(&hash)).unwrap(); + assert_eq!(engine.get(&exact_key).unwrap().unwrap(), raw); + let attachment_key = super::hex_to_key(&attachment_hash).unwrap(); + assert_eq!(engine.get(&attachment_key).unwrap().unwrap(), attachment); + assert_eq!( + engine + .exists_batch(&[key, exact_key, attachment_key]) + .unwrap(), + vec![true, true, true] + ); + engine.shutdown().unwrap(); + std::fs::remove_dir_all(dir).unwrap(); + } +} + /// Returns a reader over the raw EML for an indexed message. /// /// If the message's content blob is missing from the blob store, it is fetched diff --git a/crates/core/src/store/tantivy/attachment.rs b/crates/core/src/store/tantivy/attachment.rs index dd727e14..9d6b2cac 100644 --- a/crates/core/src/store/tantivy/attachment.rs +++ b/crates/core/src/store/tantivy/attachment.rs @@ -209,6 +209,31 @@ impl IndexManager { let _ = self.sender.send(doc).await; } + /// Idempotently replaces a UIDONLY projection batch in one commit, before + /// the corresponding envelope completion markers are published. + pub(crate) async fn replace_uidonly_batches( + &self, + batches: Vec<(String, Vec)>, + ) -> BichonResult<()> { + let mut writer = self.index_writer.lock().await; + let f_envelope_id = SchemaTools::attachment_fields().f_envelope_id; + let mut operations = Vec::new(); + for (envelope_id, docs) in batches { + operations.push(UserOperation::Delete(Term::from_field_text( + f_envelope_id, + &envelope_id, + ))); + operations.extend(docs.into_iter().map(UserOperation::Add)); + } + writer + .run(operations) + .map_err(|e| raise_error!(format!("{e:#?}"), ErrorCode::InternalError))?; + writer + .commit() + .map_err(|e| raise_error!(format!("{e:#?}"), ErrorCode::InternalError))?; + Ok(()) + } + fn open_or_create_index(index_dir: &PathBuf) -> Index { let need_create = !index_dir.exists() || index_dir diff --git a/crates/core/src/store/tantivy/dedup.rs b/crates/core/src/store/tantivy/dedup.rs index ecd9d698..8621cba9 100644 --- a/crates/core/src/store/tantivy/dedup.rs +++ b/crates/core/src/store/tantivy/dedup.rs @@ -11,7 +11,7 @@ use crate::raise_error; use crate::store::tantivy::attachment::ATTACHMENT_MANAGER; use crate::store::tantivy::envelope::ENVELOPE_MANAGER; use crate::store::tantivy::fields::{ - F_ACCOUNT_ID, F_CONTENT_HASH, F_ID, F_INGEST_AT, F_MAILBOX_ID, + F_ACCOUNT_ID, F_CONTENT_HASH, F_ID, F_INGEST_AT, F_MAILBOX_ID, F_SHARD_ID, }; use crate::store::tantivy::schema::SchemaTools; @@ -35,6 +35,12 @@ struct DedupEntry { /// Value = all documents sharing that key, to be reduced to exactly one. type DedupMap = HashMap<(u64, String), Vec>; +pub(crate) const UIDONLY_SHARD_BIT: u64 = 1 << 63; + +pub(crate) fn is_uidonly_shard(shard_id: u64) -> bool { + shard_id & UIDONLY_SHARD_BIT != 0 +} + // ─── Public entry point ─────────────────────────────────────────────────────── /// Background deduplication task. @@ -178,6 +184,10 @@ fn dedup_account( .fast_fields() .i64(F_INGEST_AT) .map_err(|e| raise_error!(format!("{:#?}", e), ErrorCode::InternalError))?; + let shard_col = segment_reader + .fast_fields() + .u64(F_SHARD_ID) + .map_err(|e| raise_error!(format!("{:#?}", e), ErrorCode::InternalError))?; // content_hash and f_id are text fields with FAST; stored as dictionary-encoded strings let hash_col = segment_reader .fast_fields() @@ -199,6 +209,10 @@ fn dedup_account( if account_col.values.get_val(doc_id) != account_id { continue; } + // UIDONLY records represent server UIDs, not unique bodies. + if is_uidonly_shard(shard_col.values.get_val(doc_id)) { + continue; + } let mailbox_id = mailbox_col.values.get_val(doc_id); let ingest_at = ingest_col.values.get_val(doc_id); @@ -408,6 +422,20 @@ mod tests { mailbox: u64, hash: &str, ingest_at: i64, + ) { + add_email_with_shard(f, w, id, account, mailbox, hash, ingest_at, 0); + } + + #[allow(clippy::too_many_arguments)] + fn add_email_with_shard( + f: &EmailFields, + w: &mut IndexWriter, + id: &str, + account: u64, + mailbox: u64, + hash: &str, + ingest_at: i64, + shard_id: u64, ) { let mut doc = TantivyDocument::new(); doc.add_text(f.f_id, id); @@ -415,6 +443,7 @@ mod tests { doc.add_u64(f.f_mailbox_id, mailbox); doc.add_text(f.f_content_hash, hash); doc.add_i64(f.f_ingest_at, ingest_at); + doc.add_u64(f.f_shard_id, shard_id); w.add_document(doc).unwrap(); } @@ -517,6 +546,22 @@ mod tests { .await; } + #[tokio::test] + async fn dedup_preserves_distinct_uidonly_records_with_identical_raw() { + Harness::run( + "uidonly-identical-raw", + |ef, ew, af, aw| { + add_email_with_shard(ef, ew, "epoch-7-uid-1", 1, 2, "same", 1, UIDONLY_SHARD_BIT); + add_email_with_shard(ef, ew, "epoch-7-uid-2", 1, 2, "same", 2, UIDONLY_SHARD_BIT); + add_attachment(af, aw, "att-1", "epoch-7-uid-1", 1, 2); + add_attachment(af, aw, "att-2", "epoch-7-uid-2", 1, 2); + }, + &["epoch-7-uid-1", "epoch-7-uid-2"], + &["att-1", "att-2"], + ) + .await; + } + #[tokio::test] async fn dedup_keeps_latest_among_many_duplicates() { Harness::run( diff --git a/crates/core/src/store/tantivy/envelope.rs b/crates/core/src/store/tantivy/envelope.rs index 9854fbf7..c76917ee 100644 --- a/crates/core/src/store/tantivy/envelope.rs +++ b/crates/core/src/store/tantivy/envelope.rs @@ -68,7 +68,10 @@ use tantivy::{ }, collector::{Count, DocSetCollector, FacetCollector, TopDocs}, indexer::{LogMergePolicy, UserOperation}, - query::{AllQuery, BooleanQuery, EmptyQuery, Occur, Query, QueryParser, RangeQuery, TermQuery}, + query::{ + AllQuery, BooleanQuery, EmptyQuery, Occur, Query, QueryParser, RangeQuery, TermQuery, + TermSetQuery, + }, schema::{IndexRecordOption, Value}, DocAddress, Index, IndexReader, IndexWriter, Order, TantivyDocument, Term, }; @@ -80,6 +83,7 @@ use tokio::{ use tracing::{info, warn}; pub static ENVELOPE_MANAGER: LazyLock = LazyLock::new(IndexManager::new); +pub(crate) static CANONICAL_STORAGE_GATE: Mutex<()> = Mutex::const_new(()); pub struct IndexManager { index: Arc, @@ -89,6 +93,14 @@ pub struct IndexManager { handle: Mutex>>, } +#[derive(Clone, Debug, Eq, PartialEq)] +pub(crate) struct UidOnlyProjectionMarker { + pub mailbox_id: u64, + pub uid: u32, + pub content_hash: String, + pub shard_id: u64, +} + impl IndexManager { pub(crate) fn index_writer(&self) -> &Arc> { &self.index_writer @@ -229,6 +241,104 @@ impl IndexManager { } } + /// Replaces and commits a batch of final UIDONLY completion markers. Raw + /// bytes and attachment documents must already be committed. + pub(crate) async fn replace_uidonly_documents( + &self, + documents: Vec<(String, TantivyDocument)>, + ) -> BichonResult<()> { + let mut writer = self.index_writer.lock().await; + let f_id = SchemaTools::email_fields().f_id; + let operations: Vec<_> = documents + .into_iter() + .flat_map(|(envelope_id, doc)| { + [ + UserOperation::Delete(Term::from_field_text(f_id, &envelope_id)), + UserOperation::Add(doc), + ] + }) + .collect(); + writer + .run(operations) + .map_err(|e| raise_error!(format!("{e:#?}"), ErrorCode::InternalError))?; + writer + .commit() + .map_err(|e| raise_error!(format!("{e:#?}"), ErrorCode::InternalError))?; + Ok(()) + } + + pub(crate) fn get_uidonly_projection_markers( + &self, + account_id: u64, + envelope_ids: &[String], + ) -> BichonResult> { + if envelope_ids.is_empty() { + return Ok(HashMap::new()); + } + let searcher = self.create_searcher()?; + let f = SchemaTools::email_fields(); + let query = BooleanQuery::new(vec![ + ( + Occur::Must, + self.account_query(account_id) as Box, + ), + ( + Occur::Must, + Box::new(TermSetQuery::new( + envelope_ids + .iter() + .map(|id| Term::from_field_text(f.f_id, id)), + )), + ), + ]); + let docs = searcher + .search(&query, &DocSetCollector) + .map_err(|e| raise_error!(format!("{e:#?}"), ErrorCode::InternalError))?; + let mut markers = HashMap::with_capacity(docs.len()); + let mut seen_ids = HashSet::with_capacity(docs.len()); + for address in docs { + let doc = searcher + .doc::(address) + .map_err(|e| raise_error!(format!("{e:#?}"), ErrorCode::InternalError))?; + let Some(envelope_id) = doc + .get_first(f.f_id) + .and_then(|v| v.as_str()) + .map(ToOwned::to_owned) + else { + continue; + }; + if !seen_ids.insert(envelope_id.clone()) { + return Err(raise_error!( + "multiple documents exist for one deterministic UIDONLY envelope id".into(), + ErrorCode::Incompatible + )); + } + let marker = ( + doc.get_first(f.f_mailbox_id).and_then(|v| v.as_u64()), + doc.get_first(f.f_uid) + .and_then(|v| v.as_u64()) + .and_then(|v| u32::try_from(v).ok()), + doc.get_first(f.f_content_hash) + .and_then(|v| v.as_str()) + .map(ToOwned::to_owned), + doc.get_first(f.f_shard_id).and_then(|v| v.as_u64()), + ); + if let (Some(mailbox_id), Some(uid), Some(content_hash), Some(shard_id)) = marker + { + markers.insert( + envelope_id, + UidOnlyProjectionMarker { + mailbox_id, + uid, + content_hash, + shard_id, + }, + ); + } + } + Ok(markers) + } + fn open_or_create_index(index_dir: &PathBuf) -> Index { let need_create = !index_dir.exists() || index_dir @@ -870,6 +980,7 @@ impl IndexManager { } pub async fn delete_account_envelopes(&self, account_id: u64) -> BichonResult<()> { + let _canonical_guard = CANONICAL_STORAGE_GATE.lock().await; let query = self.account_query(account_id); let (eml_content_hashes, attachments_content_hashes) = self.collect_content_hashes(query)?; @@ -908,6 +1019,7 @@ impl IndexManager { if mailbox_ids.is_empty() { return Ok(()); } + let _canonical_guard = CANONICAL_STORAGE_GATE.lock().await; let mut eml_content_hashes: HashSet = HashSet::new(); let mut attachments_content_hashes: HashSet = HashSet::new(); @@ -933,6 +1045,10 @@ impl IndexManager { .commit() .map_err(|e| raise_error!(format!("{:#?}", e), ErrorCode::InternalError))?; + ATTACHMENT_MANAGER + .delete_mailbox_attachments(account_id, mailbox_ids.clone()) + .await?; + if !eml_content_hashes.is_empty() || !attachments_content_hashes.is_empty() { self.cleanup_unused_content( &mut writer, @@ -1061,6 +1177,7 @@ impl IndexManager { tracing::warn!("delete_envelopes_multi_account: deletes is empty, nothing to delete"); return Ok(()); } + let _canonical_guard = CANONICAL_STORAGE_GATE.lock().await; let mut eml_content_hash_triples: HashSet<(u64, u64, String)> = HashSet::new(); let mut attachments_content_hashes: HashSet = HashSet::new(); @@ -1087,13 +1204,13 @@ impl IndexManager { let mut writer = self.index_writer.lock().await; - for (account_id, envelope_ids) in deletes { + for (account_id, envelope_ids) in &deletes { let unique_ids: HashSet<&String> = envelope_ids.iter().collect(); if unique_ids.is_empty() { continue; } for eid in unique_ids { - let query = self.envelope_query(account_id, eid); + let query = self.envelope_query(*account_id, eid); writer .delete_query(query) .map_err(|e| raise_error!(format!("{:#?}", e), ErrorCode::InternalError))?; @@ -1103,6 +1220,10 @@ impl IndexManager { .commit() .map_err(|e| raise_error!(format!("{:#?}", e), ErrorCode::InternalError))?; + ATTACHMENT_MANAGER + .delete_attachments_multi_account(deletes) + .await?; + if !eml_content_hash_triples.is_empty() || !attachments_content_hashes.is_empty() { let eml_content_hashes: HashSet = eml_content_hash_triples .iter() @@ -1335,7 +1456,7 @@ impl IndexManager { // blob store, referenced by f_content_hash. if let Some(hash_val) = old_doc.get_first(f.f_content_hash) { if let Some(content_hash) = hash_val.as_str() { - match BLOB_MANAGER.get_email(content_hash) { + match BLOB_MANAGER.get_canonical_email(content_hash) { Ok(Some(eml_bytes)) => { if let Some(message) = MessageParser::new().parse(&eml_bytes) { let text = message From d6efe7bf82f964b82e778cc44ba4e08729e8106c Mon Sep 17 00:00:00 2001 From: Gabe A Date: Tue, 4 Aug 2026 10:36:16 -0400 Subject: [PATCH 3/4] route limited mailboxes through uidonly --- crates/admin/src/meta.rs | 2 + crates/core/src/account/migration.rs | 34 + crates/core/src/account/payload.rs | 5 + crates/core/src/account/view.rs | 2 + crates/core/src/cache/imap/download/flow.rs | 262 ++- crates/core/src/cache/imap/download/mod.rs | 11 +- .../core/src/cache/imap/download/rebuild.rs | 91 +- crates/core/src/cache/imap/mailbox.rs | 4 + crates/core/src/imap/client.rs | 17 +- crates/core/src/imap/manager.rs | 55 +- crates/core/src/imap/mock_server.rs | 16 +- crates/core/src/imap/mod.rs | 1 + crates/core/src/imap/uidonly_acquisition.rs | 1823 +++++++++++++++++ crates/core/src/import/mod.rs | 2 + crates/smtp/src/server.rs | 1 + web/src/api/account/api.ts | 3 +- web/src/features/accounts/account-new.tsx | 2 + .../accounts/account-settings-page.tsx | 2 + .../components/__tests__/schema-edge.test.ts | 1 + .../components/__tests__/schema.test.ts | 1 + .../features/accounts/components/schema.ts | 1 + .../accounts/components/tab-server.tsx | 16 + web/src/locales/en.json | 4 +- 23 files changed, 2252 insertions(+), 104 deletions(-) create mode 100644 crates/core/src/imap/uidonly_acquisition.rs diff --git a/crates/admin/src/meta.rs b/crates/admin/src/meta.rs index dfd91181..4f1f7e44 100644 --- a/crates/admin/src/meta.rs +++ b/crates/admin/src/meta.rs @@ -256,6 +256,7 @@ impl From for AccountModel { updated_at: value.updated_at, created_by: value.created_by, use_dangerous: value.use_dangerous, + uidonly_enabled: false, pgp_key: value.pgp_key, imap_quota_window: None, imap_quota_bytes: None, @@ -666,6 +667,7 @@ impl From for bichon_core::cache::imap::mailbox::MailBox { uid_next: value.uid_next, uid_validity: value.uid_validity, highest_uid: None, + uidonly_source_scope: None, } } } diff --git a/crates/core/src/account/migration.rs b/crates/core/src/account/migration.rs index 531f701e..5c31020d 100644 --- a/crates/core/src/account/migration.rs +++ b/crates/core/src/account/migration.rs @@ -301,6 +301,9 @@ pub struct Account { pub updated_at: i64, pub created_by: u64, //user id pub use_dangerous: bool, + /// UIDONLY acquisition is opt-in because it changes remote message identity. + #[serde(default)] + pub uidonly_enabled: bool, pub pgp_key: Option, pub imap_quota_bytes: Option, pub imap_quota_window: Option, @@ -345,6 +348,7 @@ impl Account { created_at: utc_now!(), updated_at: utc_now!(), use_dangerous: request.use_dangerous, + uidonly_enabled: request.uidonly_enabled.unwrap_or_default(), pgp_key: request.pgp_key, created_by: user_id, download_batch_size: request.download_batch_size, @@ -668,6 +672,10 @@ impl Account { new.use_dangerous = use_dangerous; } + if let Some(uidonly_enabled) = request.uidonly_enabled { + new.uidonly_enabled = uidonly_enabled; + } + if let Some(pgp_key) = request.pgp_key { new.pgp_key = Some(pgp_key); } @@ -704,6 +712,32 @@ impl Account { mod tests { use super::*; + #[test] + fn uidonly_setting_defaults_off_and_only_changes_explicitly() { + let account = Account::new(1, AccountCreateRequest::default()).unwrap(); + assert!(!account.uidonly_enabled); + let update = |old: &Account, value| { + Account::apply_update_fields( + old, + AccountUpdateRequest { + uidonly_enabled: value, + ..Default::default() + }, + ) + .unwrap() + }; + let enabled = update(&account, Some(true)); + assert!(enabled.uidonly_enabled); + assert!(update(&enabled, None).uidonly_enabled); + assert!(!update(&enabled, Some(false)).uidonly_enabled); + + let mut legacy = serde_json::to_value(enabled).unwrap(); + legacy.as_object_mut().unwrap().remove("uidonly_enabled"); + assert!(!serde_json::from_value::(legacy) + .unwrap() + .uidonly_enabled); + } + // ── FilterRule ─────────────────────────────────────────────────── #[test] diff --git a/crates/core/src/account/payload.rs b/crates/core/src/account/payload.rs index a0c55943..6b1fd929 100644 --- a/crates/core/src/account/payload.rs +++ b/crates/core/src/account/payload.rs @@ -52,6 +52,9 @@ pub struct AccountCreateRequest { pub download_batch_size: Option, pub max_email_size_bytes: Option, pub use_dangerous: bool, + /// Explicitly enables UIDONLY acquisition when required by the server. + /// Omitted values remain disabled for backward compatibility. + pub uidonly_enabled: Option, pub pgp_key: Option, pub imap_quota_bytes: Option, pub imap_quota_window: Option, @@ -188,6 +191,8 @@ pub struct AccountUpdateRequest { pub download_batch_size: Option, pub max_email_size_bytes: Option, pub use_dangerous: Option, + /// Enables or disables UIDONLY acquisition for this account. + pub uidonly_enabled: Option, pub pgp_key: Option, pub imap_quota_bytes: Option, diff --git a/crates/core/src/account/view.rs b/crates/core/src/account/view.rs index d800f0b0..f585e29c 100644 --- a/crates/core/src/account/view.rs +++ b/crates/core/src/account/view.rs @@ -52,6 +52,7 @@ pub struct AccountResp { pub created_user_name: String, pub created_user_email: String, pub use_dangerous: bool, + pub uidonly_enabled: bool, pub pgp_key: Option, pub imap_quota_bytes: Option, pub imap_quota_window: Option, @@ -90,6 +91,7 @@ impl AccountResp { .map(|u| u.email.clone()) .unwrap_or_else(|| "N/A".to_string()), use_dangerous: account.use_dangerous, + uidonly_enabled: account.uidonly_enabled, pgp_key: account.pgp_key, imap_quota_bytes: account.imap_quota_bytes, imap_quota_window: account.imap_quota_window, diff --git a/crates/core/src/cache/imap/download/flow.rs b/crates/core/src/cache/imap/download/flow.rs index d41f53ea..23f6f0c1 100644 --- a/crates/core/src/cache/imap/download/flow.rs +++ b/crates/core/src/cache/imap/download/flow.rs @@ -36,6 +36,10 @@ use crate::{ compress_uid_list, generate_uid_sequence_hashset, slow_server_message, ImapExecutor, DEFAULT_BATCH_SIZE, }, + imap::uidonly_acquisition::{ + account_requires_uidonly, connect_and_acquire_or_legacy, + mailbox_has_uidonly_proof, AcquisitionRoute, + }, store::tantivy::envelope::ENVELOPE_MANAGER, }, }; @@ -58,6 +62,12 @@ pub async fn fetch_and_save_by_date( direction: FetchDirection, token: CancellationToken, ) -> BichonResult> { + if account_requires_uidonly(account) || mailbox.uidonly_source_scope.is_some() { + return Err(raise_error!( + "Date-scoped acquisition is unsafe on a UIDONLY-limited server".into(), + ErrorCode::Incompatible + )); + } let account_id = account.id; let mut session = match ImapExecutor::create_connection(account_id).await { Ok(session) => session, @@ -258,24 +268,61 @@ pub async fn fetch_and_save_by_date( } /// Fetches all messages from a mailbox. -/// Returns `Ok(Some(max_uid))` with the highest UID stored, or `Ok(None)` if empty. -/// -/// The mailbox is enumerated via `UID SEARCH ALL` first, then downloaded in UID -/// batches. Unlike sequence-number paging, UIDs stay stable while the download -/// runs (new arrivals only get larger UIDs), so no message is silently skipped -/// when the server changes mid-download. +/// Returns a checkpoint only after the complete mailbox run succeeds. +/// The legacy route enumerates `UID SEARCH ALL` first, so new arrivals cannot +/// shift messages between batches while the download is running. pub async fn fetch_and_save_full_mailbox( account: &AccountModel, mailbox: &MailBox, + force_uidonly: bool, token: CancellationToken, -) -> BichonResult> { +) -> BichonResult { let mailbox_id = mailbox.id; let account_id = account.id; - let mut session = match ImapExecutor::create_connection(account_id).await { - Ok(session) => session, + // Classify the actual acquisition connection. Cached capabilities route + // known-limited accounts early, but can never authorize a legacy fetch. + let route = connect_and_acquire_or_legacy( + account, + mailbox, + force_uidonly, + token.clone(), + |progress| { + DownloadState::update_folder_progress( + account_id, + mailbox.name.clone(), + progress.planned, + progress.resolved, + FolderStatus::Downloading, + None, + ) + }, + ) + .await; + let mut session = match route { + Ok(AcquisitionRoute::Acquired { + report, + source_scope, + }) => { + DownloadState::update_folder_progress( + account_id, + mailbox.name.clone(), + report.inventoried, + report.archived, + FolderStatus::Success, + None, + )?; + let mut updated = mailbox.clone(); + updated.highest_uid = report.checkpoint; + updated.uid_validity = Some(report.uid_validity); + updated.uid_next = Some(report.uid_next); + updated.exists = report.exists; + updated.uidonly_source_scope = Some(source_scope); + return Ok(updated); + } + Ok(AcquisitionRoute::Legacy(session)) => session, Err(e) => { - let err_msg = format!("Connection failed for this folder: {:#?}", e); + let err_msg = format!("Full mailbox acquisition failed: {:#?}", e); DownloadState::update_folder_progress( account_id, mailbox.name.clone(), @@ -323,7 +370,10 @@ pub async fn fetch_and_save_full_mailbox( None, )?; session.logout().await.ok(); - return Ok(None); + let mut updated = mailbox.clone(); + updated.highest_uid = None; + updated.uidonly_source_scope = None; + return Ok(updated); } let max_uid = *uid_list.last().unwrap(); @@ -337,7 +387,6 @@ pub async fn fetch_and_save_full_mailbox( ); let mut current_processed = 0u64; - let mut has_error_or_cancel = false; for (index, batch) in uid_batches.into_iter().enumerate() { if token.is_cancelled() { @@ -354,8 +403,11 @@ pub async fn fetch_and_save_full_mailbox( FolderStatus::Cancelled, None, )?; - has_error_or_cancel = true; - break; + session.logout().await.ok(); + return Err(raise_error!( + "Full mailbox download cancelled".into(), + ErrorCode::InternalError + )); } // Heartbeat: keeps `current_folder` fresh so the web UI can show @@ -452,24 +504,25 @@ pub async fn fetch_and_save_full_mailbox( FolderStatus::Failed, Some(err_msg), )?; - has_error_or_cancel = true; - break; + session.logout().await.ok(); + return Err(e); } } } - if !has_error_or_cancel { - DownloadState::update_folder_progress( - account_id, - mailbox.name.clone(), - planned, - current_processed, - FolderStatus::Success, - None, - )?; - } + DownloadState::update_folder_progress( + account_id, + mailbox.name.clone(), + planned, + current_processed, + FolderStatus::Success, + None, + )?; session.logout().await.ok(); - Ok(max_uid.into()) + let mut updated = mailbox.clone(); + updated.highest_uid = Some(max_uid); + updated.uidonly_source_scope = None; + Ok(updated) } /// Generates a synthetic UIDVALIDITY for IMAP servers that don't provide it. @@ -759,14 +812,37 @@ async fn reconcile_uid_validity_change( Ok(max_uid) } +fn mailbox_source_changed(account: &AccountModel, mailbox: &MailBox) -> bool { + mailbox.uidonly_source_scope.is_some() && !mailbox_has_uidonly_proof(account, mailbox) +} + +fn intersecting_source_changed(account: &AccountModel, mailboxes: &[(MailBox, MailBox)]) -> bool { + mailboxes + .iter() + .any(|(local, _)| mailbox_source_changed(account, local)) +} + pub async fn reconcile_mailboxes( account: &AccountModel, remote_mailboxes: &[MailBox], local_mailboxes: &[MailBox], token: CancellationToken, ) -> BichonResult<()> { - let start_time = Instant::now(); let existing_mailboxes = find_intersecting_mailboxes(local_mailboxes, remote_mailboxes); + let source_changed = intersecting_source_changed(account, &existing_mailboxes); + let known_limited = local_mailboxes + .iter() + .any(|mailbox| mailbox_has_uidonly_proof(account, mailbox)) + || account_requires_uidonly(account); + if (known_limited || source_changed) + && (account.date_since.is_some() || account.date_before.is_some()) + { + return Err(raise_error!( + "Date-scoped acquisition is unsafe on a UIDONLY-limited server".into(), + ErrorCode::Incompatible + )); + } + let start_time = Instant::now(); let account_id = account.id; if !existing_mailboxes.is_empty() { let mut mailboxes_to_update = Vec::with_capacity(existing_mailboxes.len()); @@ -777,6 +853,7 @@ pub async fn reconcile_mailboxes( )?; for (local_mailbox, remote_mailbox) in &existing_mailboxes { + let source_changed = mailbox_source_changed(account, local_mailbox); if token.is_cancelled() { DownloadState::update_session_status( account.id, @@ -786,6 +863,30 @@ pub async fn reconcile_mailboxes( break; } + if account.date_since.is_none() + && account.date_before.is_none() + && (known_limited || source_changed) + { + let can_resume = local_mailbox.uid_validity == remote_mailbox.uid_validity + && !source_changed + && mailbox_has_uidonly_proof(account, local_mailbox); + let mut acquisition_mailbox = (*remote_mailbox).clone(); + if can_resume { + acquisition_mailbox.highest_uid = local_mailbox.highest_uid; + acquisition_mailbox.uidonly_source_scope = + local_mailbox.uidonly_source_scope.clone(); + } + let result = fetch_and_save_full_mailbox( + account, + &acquisition_mailbox, + known_limited, + token.clone(), + ) + .await?; + mailboxes_to_update.push(result); + continue; + } + // Handle missing UIDVALIDITY from non-compliant IMAP servers // (e.g., Tencent Enterprise Mail, etc.) let remote_uid_validity = match remote_mailbox.uid_validity { @@ -885,7 +986,11 @@ pub async fn reconcile_mailboxes( )?; break; } - if mailbox.exists > 0 { + if mailbox.exists > 0 + || (account.date_since.is_none() + && account.date_before.is_none() + && (known_limited || account_requires_uidonly(account))) + { let account = account.clone(); let mailbox = mailbox.clone(); @@ -901,41 +1006,50 @@ pub async fn reconcile_mailboxes( }; let result = match &account.date_since { - Some(date_since) => { - rebuild_mailbox_cache_by_date( + Some(date_since) => rebuild_mailbox_cache_by_date( + &account, + mailbox.id, + &date_since.since_date()?, + &mailbox, + FetchDirection::Since, + token.clone(), + ) + .await + .map(|highest_uid| { + let mut updated = mailbox.clone(); + updated.highest_uid = highest_uid; + updated + }), + None => match &account.date_before { + Some(r) => rebuild_mailbox_cache_by_date( &account, mailbox.id, - &date_since.since_date()?, + &r.calculate_date()?, &mailbox, - FetchDirection::Since, + FetchDirection::Before, token.clone(), ) .await - } - None => match &account.date_before { - Some(r) => { - rebuild_mailbox_cache_by_date( + .map(|highest_uid| { + let mut updated = mailbox.clone(); + updated.highest_uid = highest_uid; + updated + }), + None => { + rebuild_mailbox_cache( &account, - mailbox.id, - &r.calculate_date()?, &mailbox, - FetchDirection::Before, + &mailbox, + known_limited, token.clone(), ) .await } - None => { - rebuild_mailbox_cache(&account, &mailbox, &mailbox, token.clone()).await - } }, }; match result { - Ok(new_highest_uid) => { - let mut updated = mailbox.clone(); - updated.highest_uid = new_highest_uid; - MailBox::batch_upsert(&[updated])?; - } + Ok(updated) => MailBox::batch_upsert(&[updated])?, Err(err) => { has_error = true; tracing::error!("Folder sync task failed: {:#?}", err); @@ -1003,8 +1117,14 @@ async fn perform_incremental_sync( .await? } None => { - fetch_and_save_full_mailbox(account, remote_mailbox, token) - .await? + fetch_and_save_full_mailbox( + account, + remote_mailbox, + false, + token, + ) + .await? + .highest_uid } }, }; @@ -1063,6 +1183,48 @@ mod tests { // Pure unit tests (no network) // ============================================================ + #[test] + fn stale_nonremote_uidonly_scope_does_not_force_current_mailbox_rescan() { + let local = [ + MailBox { + name: "INBOX".into(), + ..Default::default() + }, + MailBox { + name: "Unsubscribed".into(), + uidonly_source_scope: Some("old-source".into()), + ..Default::default() + }, + ]; + let remote = [MailBox { + name: "INBOX".into(), + ..Default::default() + }]; + let current = find_intersecting_mailboxes(&local, &remote); + assert!(mailbox_source_changed(&AccountModel::default(), &local[1])); + assert!(!intersecting_source_changed( + &AccountModel::default(), + ¤t + )); + } + + #[tokio::test] + async fn known_limited_server_rejects_date_scoped_fallback_before_connecting() { + let account = AccountModel { + capabilities: Some(vec!["UIDONLY".into(), "MESSAGELIMIT=10000".into()]), + ..Default::default() + }; + assert!(fetch_and_save_by_date( + &account, + "2026-01-01", + &MailBox::default(), + FetchDirection::Since, + CancellationToken::new(), + ) + .await + .is_err()); + } + #[test] fn test_generate_synthetic_uidvalidity_deterministic() { let a = generate_synthetic_uidvalidity("INBOX"); diff --git a/crates/core/src/cache/imap/download/mod.rs b/crates/core/src/cache/imap/download/mod.rs index c4653f96..03942218 100644 --- a/crates/core/src/cache/imap/download/mod.rs +++ b/crates/core/src/cache/imap/download/mod.rs @@ -78,11 +78,14 @@ pub async fn process_imap_download( } }; session.logout().await.ok(); + // The discovery connection refreshes cached capabilities. Route the + // download from that fresh snapshot, not the stale task input. + let account = AccountModel::get(account_id)?; if matches!(download_task, DownloadTask::FullFetch) { let result = match &account.date_since { Some(date_since) => { rebuild_cache_by_date( - account, + &account, &remote_mailboxes, &date_since.since_date()?, FetchDirection::Since, @@ -93,7 +96,7 @@ pub async fn process_imap_download( None => match &account.date_before { Some(r) => { rebuild_cache_by_date( - account, + &account, &remote_mailboxes, &r.calculate_date()?, FetchDirection::Before, @@ -101,7 +104,7 @@ pub async fn process_imap_download( ) .await } - None => rebuild_cache(account, &remote_mailboxes, token).await, + None => rebuild_cache(&account, &remote_mailboxes, token).await, }, }; match result { @@ -122,7 +125,7 @@ pub async fn process_imap_download( } let local_mailboxes = MailBox::list_all(account_id)?; - match reconcile_mailboxes(account, &remote_mailboxes, &local_mailboxes, token).await { + match reconcile_mailboxes(&account, &remote_mailboxes, &local_mailboxes, token).await { Ok(_) => DownloadState::update_session_status(account_id, DownloadStatus::Success, None)?, Err(e) => { let err_msg = format!("Email Download interrupted: {:#?}", e); diff --git a/crates/core/src/cache/imap/download/rebuild.rs b/crates/core/src/cache/imap/download/rebuild.rs index 96d31d9e..50d19cab 100644 --- a/crates/core/src/cache/imap/download/rebuild.rs +++ b/crates/core/src/cache/imap/download/rebuild.rs @@ -23,19 +23,28 @@ use crate::{ }, cache::{ imap::{ - download::flow::{fetch_and_save_by_date, fetch_and_save_full_mailbox, FetchDirection}, + download::flow::{ + fetch_and_save_by_date, fetch_and_save_full_mailbox, FetchDirection, + }, mailbox::MailBox, }, SEMAPHORE, }, error::{code::ErrorCode, BichonResult}, + imap::uidonly_acquisition::account_requires_uidonly, raise_error, - store::tantivy::{attachment::ATTACHMENT_MANAGER, envelope::ENVELOPE_MANAGER}, + store::tantivy::envelope::ENVELOPE_MANAGER, }; use tokio_util::sync::CancellationToken; use tracing::{error, info}; +fn can_trust_empty_mailbox(account: &AccountModel, mailbox: &MailBox) -> bool { + mailbox.exists == 0 + && mailbox.uidonly_source_scope.is_none() + && !account_requires_uidonly(account) +} + pub async fn rebuild_cache( account: &AccountModel, remote_mailboxes: &[MailBox], @@ -59,7 +68,7 @@ pub async fn rebuild_cache( )?; break; } - if mailbox.exists == 0 { + if can_trust_empty_mailbox(account, mailbox) { info!( "Account {}: Mailbox '{}' on the remote server has no emails. Skipping fetch for this mailbox.", account.id, &mailbox.name @@ -88,12 +97,8 @@ pub async fn rebuild_cache( } }; - match fetch_and_save_full_mailbox(&account, &mailbox, token.clone()).await { - Ok(new_highest_uid) => { - let mut updated = mailbox.clone(); - updated.highest_uid = new_highest_uid; - MailBox::batch_upsert(&[updated])?; - } + match fetch_and_save_full_mailbox(&account, &mailbox, false, token.clone()).await { + Ok(updated) => MailBox::batch_upsert(&[updated])?, Err(err) => { has_error = true; tracing::error!("Folder sync task failed: {:#?}", err); @@ -114,6 +119,29 @@ pub async fn rebuild_cache( Ok(()) } +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn only_ordinary_empty_mailboxes_use_the_fast_path() { + let mailbox = MailBox::default(); + assert!(can_trust_empty_mailbox(&AccountModel::default(), &mailbox)); + let limited = AccountModel { + capabilities: Some(vec!["UIDONLY".into(), "MESSAGELIMIT=10000".into()]), + ..Default::default() + }; + assert!(!can_trust_empty_mailbox(&limited, &mailbox)); + assert!(!can_trust_empty_mailbox( + &AccountModel::default(), + &MailBox { + uidonly_source_scope: Some("proof".into()), + ..Default::default() + } + )); + } +} + pub async fn rebuild_cache_by_date( account: &AccountModel, remote_mailboxes: &[MailBox], @@ -121,6 +149,16 @@ pub async fn rebuild_cache_by_date( direction: FetchDirection, token: CancellationToken, ) -> BichonResult<()> { + if account_requires_uidonly(account) + || remote_mailboxes + .iter() + .any(|mailbox| mailbox.uidonly_source_scope.is_some()) + { + return Err(raise_error!( + "Date-scoped acquisition is unsafe on a UIDONLY-limited server".into(), + ErrorCode::Incompatible + )); + } MailBox::batch_insert(remote_mailboxes)?; DownloadState::init_folder_details( account.id, @@ -203,33 +241,13 @@ pub async fn rebuild_mailbox_cache( account: &AccountModel, local_mailbox: &MailBox, remote_mailbox: &MailBox, + force_uidonly: bool, token: CancellationToken, -) -> BichonResult> { +) -> BichonResult { ENVELOPE_MANAGER .delete_mailbox_envelopes(account.id, vec![local_mailbox.id]) .await?; - ATTACHMENT_MANAGER - .delete_mailbox_attachments(account.id, vec![local_mailbox.id]) - .await?; - if remote_mailbox.exists == 0 { - info!( - "Account {}: Mailbox '{}' has no emails on the remote server. The mailbox is empty, no envelopes to fetch.", - account.id, - &local_mailbox.name - ); - DownloadState::update_folder_progress( - account.id, - remote_mailbox.name.clone(), - 0, - 0, - FolderStatus::Success, - None, - )?; - return Ok(None); - } - - let result = fetch_and_save_full_mailbox(account, remote_mailbox, token).await?; - Ok(result) + fetch_and_save_full_mailbox(account, remote_mailbox, force_uidonly, token).await } pub async fn rebuild_mailbox_cache_by_date( @@ -240,12 +258,15 @@ pub async fn rebuild_mailbox_cache_by_date( direction: FetchDirection, token: CancellationToken, ) -> BichonResult> { + if account_requires_uidonly(account) || remote.uidonly_source_scope.is_some() { + return Err(raise_error!( + "Date-scoped acquisition is unsafe on a UIDONLY-limited server".into(), + ErrorCode::Incompatible + )); + } ENVELOPE_MANAGER .delete_mailbox_envelopes(account.id, vec![local_mailbox_id]) .await?; - ATTACHMENT_MANAGER - .delete_mailbox_attachments(account.id, vec![local_mailbox_id]) - .await?; if remote.exists == 0 { info!( "Account {}: Mailbox '{}' has no emails on the remote server. The mailbox is empty, no envelopes to fetch.", diff --git a/crates/core/src/cache/imap/mailbox.rs b/crates/core/src/cache/imap/mailbox.rs index db11e4e2..69485c18 100644 --- a/crates/core/src/cache/imap/mailbox.rs +++ b/crates/core/src/cache/imap/mailbox.rs @@ -60,6 +60,10 @@ pub struct MailBox { /// Used for incremental sync: next fetch starts from `highest_uid + 1`. /// If `None`, a fallback query against the Tantivy index will be performed once. pub highest_uid: Option, + /// Source fingerprint proving that `highest_uid` came from a complete + /// UIDONLY traversal rather than a provider-limited legacy view. + #[serde(default)] + pub uidonly_source_scope: Option, } impl MemDbModel for MailBox { diff --git a/crates/core/src/imap/client.rs b/crates/core/src/imap/client.rs index 85cbceea..00ec2311 100644 --- a/crates/core/src/imap/client.rs +++ b/crates/core/src/imap/client.rs @@ -21,10 +21,11 @@ use crate::error::code::ErrorCode; use crate::error::BichonResult; use crate::imap::session::SessionStream; use crate::imap::stats::StatsWrapper; +use crate::imap::uidonly::{self, UidOnlyHandle, UidOnlyLimits}; +use crate::raise_error; use crate::utils::net::establish_tcp_connection_with_timeout; use crate::utils::net::establish_tls_connection; use crate::utils::tls::establish_tls_stream; -use crate::raise_error; use async_imap::Client as ImapClient; use async_imap::Session as ImapSession; use std::net::SocketAddr; @@ -87,6 +88,20 @@ impl Client { } } + /// Installs the UIDONLY protocol guard around the final transport. This is + /// called after STARTTLS has finished, but before authentication consumes + /// the client, so the same authenticated connection can be enabled later. + pub(crate) fn with_uidonly(self, limits: UidOnlyLimits) -> BichonResult<(Self, UidOnlyHandle)> { + let stream = self.inner.into_inner(); + let (stream, handle) = uidonly::wrap(stream, limits).map_err(|_| { + raise_error!( + "Invalid UIDONLY transport limits".into(), + ErrorCode::InvalidParameter + ) + })?; + Ok((Self::new(stream), handle)) + } + pub(crate) async fn login( self, username: &str, diff --git a/crates/core/src/imap/manager.rs b/crates/core/src/imap/manager.rs index e048edd0..64cec5fc 100644 --- a/crates/core/src/imap/manager.rs +++ b/crates/core/src/imap/manager.rs @@ -24,6 +24,7 @@ use crate::imap::capabilities::{capability_to_string, check_capabilities, fetch_ use crate::imap::client::Client; use crate::imap::oauth2::OAuth2; use crate::imap::session::SessionStream; +use crate::imap::uidonly::{UidOnlyHandle, UidOnlyLimits}; use crate::oauth2::token::OAuth2AccessToken; use crate::{bichon_version, decrypt, raise_error}; use async_imap::Session; @@ -31,6 +32,12 @@ use tracing::{error, warn}; pub struct ImapConnectionManager; +pub(crate) struct UidOnlyConnection { + pub session: Session>, + pub handle: UidOnlyHandle, + pub capabilities: Vec, +} + impl ImapConnectionManager { async fn create_client(account: &AccountModel) -> BichonResult { assert_eq!(account.account_type, AccountType::IMAP); @@ -93,13 +100,11 @@ impl ImapConnectionManager { } } - pub async fn build(account_id: u64) -> BichonResult>> { - let account = AccountModel::get(account_id)?; + async fn create_client_with_retry(account: &AccountModel) -> BichonResult { let account_email = account.email.clone(); - let mut client = None; for attempt in 0..3u32 { - match Self::create_client(&account).await { + match Self::create_client(account).await { Ok(c) => { client = Some(c); break; @@ -123,14 +128,19 @@ impl ImapConnectionManager { } } - let client = client.ok_or_else(|| { + client.ok_or_else(|| { raise_error!( format!("Failed to create IMAP {}'s client after 3 attempts", account_email), ErrorCode::NetworkError ) - })?; + }) + } - let mut session = match Self::authenticate(client, &account).await { + async fn initialize_session( + account: &AccountModel, + client: Client, + ) -> BichonResult<(Session>, Vec)> { + let mut session = match Self::authenticate(client, account).await { Ok(session) => session, Err(error) => { error!("Failed to authenticate IMAP session: {:#?}", error); @@ -138,10 +148,10 @@ impl ImapConnectionManager { } }; - match fetch_capabilities(&mut session).await { + let to_save = match fetch_capabilities(&mut session).await { Ok(capabilities) => { let to_save: Vec = capabilities.iter().map(capability_to_string).collect(); - AccountModel::update_capabilities(account_id, to_save)?; + AccountModel::update_capabilities(account.id, to_save.clone())?; if let Err(error) = check_capabilities(&capabilities) { error!("Failed to check IMAP capabilities: {:#?}", error); return Err(error); @@ -159,13 +169,38 @@ impl ImapConnectionManager { warn!("IMAP ID command failed (ignored): {:#?}", e); } } + to_save } Err(error) => { error!("Failed to fetch IMAP capabilities: {:#?}", error); return Err(error); } - } + }; + + Ok((session, to_save)) + } + pub async fn build(account_id: u64) -> BichonResult>> { + let account = AccountModel::get(account_id)?; + let client = Self::create_client_with_retry(&account).await?; + let (session, _) = Self::initialize_session(&account, client).await?; Ok(session) } + + /// Builds one capability-probed connection with the UIDONLY guard already + /// installed. The returned session is still in ordinary mode until an + /// exact `ENABLE UIDONLY` exchange succeeds. + pub(crate) async fn build_uidonly( + account: &AccountModel, + limits: UidOnlyLimits, + ) -> BichonResult { + let client = Self::create_client_with_retry(account).await?; + let (client, handle) = client.with_uidonly(limits)?; + let (session, capabilities) = Self::initialize_session(account, client).await?; + Ok(UidOnlyConnection { + session, + handle, + capabilities, + }) + } } diff --git a/crates/core/src/imap/mock_server.rs b/crates/core/src/imap/mock_server.rs index c37c3455..5cd5e141 100644 --- a/crates/core/src/imap/mock_server.rs +++ b/crates/core/src/imap/mock_server.rs @@ -38,7 +38,7 @@ //! ``` use std::net::SocketAddr; -use std::sync::Arc; +use std::sync::{Arc, Mutex}; use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader}; use tokio::net::{TcpListener, TcpStream}; @@ -47,6 +47,7 @@ type Response = Vec; pub struct MockImapServer { greeting: Vec, script: Vec<(String, Response)>, + commands: Arc>>, } impl MockImapServer { @@ -54,6 +55,7 @@ impl MockImapServer { Self { greeting: b"* OK Mock IMAP server ready\r\n".to_vec(), script: Vec::new(), + commands: Arc::new(Mutex::new(Vec::new())), } } @@ -77,6 +79,7 @@ impl MockImapServer { let addr = listener.local_addr().expect("local_addr"); let server = Arc::new(self); + let commands = Arc::clone(&server.commands); tokio::spawn(async move { loop { @@ -92,7 +95,7 @@ impl MockImapServer { } }); - MockImapServerHandle { addr } + MockImapServerHandle { addr, commands } } async fn handle_connection(&self, mut stream: TcpStream) { @@ -112,6 +115,10 @@ impl MockImapServer { Ok(_) => {} Err(_) => break, } + self.commands + .lock() + .expect("commands poisoned") + .push(line.trim_end_matches(['\r', '\n']).to_string()); let tag = extract_tag(&line).unwrap_or("A0"); let matched = self.find_match(&line); @@ -151,6 +158,7 @@ impl Default for MockImapServer { /// is dropped. pub struct MockImapServerHandle { addr: SocketAddr, + commands: Arc>>, } impl MockImapServerHandle { @@ -161,6 +169,10 @@ impl MockImapServerHandle { pub fn port(&self) -> u16 { self.addr.port() } + + pub fn commands(&self) -> Vec { + self.commands.lock().expect("commands poisoned").clone() + } } fn extract_tag(line: &str) -> Option<&str> { diff --git a/crates/core/src/imap/mod.rs b/crates/core/src/imap/mod.rs index 161c760c..3d623f09 100644 --- a/crates/core/src/imap/mod.rs +++ b/crates/core/src/imap/mod.rs @@ -27,6 +27,7 @@ pub mod stats; #[cfg(test)] mod tests; pub(crate) mod uidonly; +pub(crate) mod uidonly_acquisition; #[cfg(test)] mod uidonly_tests; #[cfg(test)] diff --git a/crates/core/src/imap/uidonly_acquisition.rs b/crates/core/src/imap/uidonly_acquisition.rs new file mode 100644 index 00000000..a376c66e --- /dev/null +++ b/crates/core/src/imap/uidonly_acquisition.rs @@ -0,0 +1,1823 @@ +// +// Copyright (c) 2025-2026 rustmailer.com (https://rustmailer.com) +// +// This file is part of the Bichon Email Archiving Project +// +// This program is free software: you can redistribute it and/or modify +// it under the terms of the GNU Affero General Public License as published by +// the Free Software Foundation, either version 3 of the License, or +// (at your option) any later version. +// +// This program is distributed in the hope that it will be useful, +// but WITHOUT ANY WARRANTY; without even the implied warranty of +// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +// GNU Affero General Public License for more details. +// +// You should have received a copy of the GNU Affero General Public License +// along with this program. If not, see . + +//! Streaming, fail-closed acquisition for RFC 9586 UIDONLY mailboxes. +//! +//! There is deliberately no secondary ledger. The canonical envelope plus +//! verified exact-raw blob is the per-UID receipt. Transient disconnects retry +//! the same operation after re-proving the fixed mailbox epoch; a failed run +//! leaves the mailbox checkpoint unchanged. + +use crate::account::entity::{AuthType, Encryption}; +use crate::account::migration::AccountModel; +use crate::cache::imap::mailbox::MailBox; +use crate::envelope::extractor::{ + project_uidonly_messages, verify_uidonly_projections, UidOnlyMessage, + UIDONLY_PROJECTION_BATCH_BYTES, UIDONLY_PROJECTION_BATCH_MESSAGES, +}; +use crate::error::code::ErrorCode; +use crate::error::BichonResult; +use crate::imap::executor::DEFAULT_MAX_EMAIL_SIZE; +use crate::imap::manager::ImapConnectionManager; +use crate::imap::session::SessionStream; +use crate::imap::uidonly::{exact_args, inventory_args, UidOnlyHandle, UidOnlyLimits}; +use crate::raise_error; +use async_imap::Session; +use futures::{FutureExt, TryStreamExt}; +use std::future::Future; +use std::panic::AssertUnwindSafe; +use std::time::Duration; +use tokio_util::sync::CancellationToken; + +const UIDONLY_DEFAULT_PAGE_SIZE: u32 = 1_000; +const MAX_UIDONLY_RECONNECTS: u32 = 3; + +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub(crate) struct AcquisitionLimits { + pub max_literal_bytes: u64, + pub max_operation_runtime: Duration, + pub page_size: u32, +} + +impl AcquisitionLimits { + pub(crate) fn bounded(max_literal_bytes: u64) -> Self { + Self { + max_literal_bytes, + // Bound each network or durable-storage operation, not the whole + // archive: a valid large mailbox may need to run for days. + max_operation_runtime: Duration::from_secs(10 * 60), + page_size: 1_000, + } + } + + fn validate(self) -> BichonResult { + if self.max_literal_bytes == 0 + || self.max_operation_runtime.is_zero() + || self.page_size == 0 + { + return Err(raise_error!( + "UIDONLY acquisition limits must be nonzero".into(), + ErrorCode::InvalidParameter + )); + } + Ok(self) + } +} + +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub(crate) struct MailboxSnapshot { + pub exists: u32, + pub uid_validity: u32, + pub uid_next: u32, +} + +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub(crate) struct InventoryItem { + pub uid: u32, +} + +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub(crate) struct AcquisitionProgress { + pub planned: u64, + pub resolved: u64, + pub downloaded: u64, +} + +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub(crate) struct AcquisitionReport { + pub uid_validity: u32, + pub uid_next: u32, + pub exists: u32, + pub checkpoint: Option, + pub inventoried: u64, + pub archived: u64, +} + +#[allow(async_fn_in_trait)] +pub(crate) trait UidOnlyTransport { + async fn snapshot(&mut self, mailbox: &str) -> BichonResult; + + async fn inventory_page( + &mut self, + cursor: u32, + high: u32, + page_size: u32, + ) -> BichonResult>; + + /// Fetch exactly one full message. `literal_budget` is a pre-read ceiling, + /// not a post-read accounting hint. + async fn fetch_exact(&mut self, uid: u32, literal_budget: u64) -> BichonResult; + + async fn reconnect(&mut self, _page_size: u32) -> BichonResult<()> { + Err(raise_error!( + "UIDONLY transport cannot reconnect".into(), + ErrorCode::NetworkError + )) + } +} + +#[allow(async_fn_in_trait)] +pub(crate) trait CanonicalArchive { + fn resume_after(&self) -> Option { + None + } + + fn begin_epoch(&mut self, _uid_validity: u32) -> BichonResult<()> { + Ok(()) + } + + async fn verify_many(&mut self, uids: &[u32]) -> BichonResult>; + + /// A successful result means every stored message has passed the durable + /// raw readback and final envelope-marker commit barrier. + async fn project_many(&mut self, messages: Vec) -> BichonResult<()>; +} + +async fn bounded(future: F, runtime: Duration, token: &CancellationToken) -> BichonResult +where + F: Future>, +{ + tokio::select! { + _ = token.cancelled() => Err(raise_error!( + "UIDONLY acquisition cancelled".into(), + ErrorCode::InternalError + )), + result = tokio::time::timeout(runtime, future) => result.map_err(|_| raise_error!( + "UIDONLY operation runtime ceiling exceeded".into(), + ErrorCode::RequestTimeout + ))?, + } +} + +fn retryable_transport_error(code: ErrorCode) -> bool { + matches!( + code, + ErrorCode::NetworkError | ErrorCode::ConnectionTimeout | ErrorCode::RequestTimeout + ) +} + +async fn recover_transport( + transport: &mut T, + mailbox: &str, + initial: MailboxSnapshot, + limits: AcquisitionLimits, + token: &CancellationToken, + reconnects: &mut u32, +) -> BichonResult<()> { + let mut last_error = raise_error!( + "UIDONLY transport recovery failed".into(), + ErrorCode::NetworkError + ); + while *reconnects < MAX_UIDONLY_RECONNECTS { + *reconnects += 1; + let delay = if cfg!(test) { + Duration::ZERO + } else { + Duration::from_secs(1 << (*reconnects - 1)) + }; + let reconnected = bounded( + async { + tokio::time::sleep(delay).await; + transport.reconnect(limits.page_size).await + }, + limits.max_operation_runtime, + token, + ) + .await; + if let Err(error) = reconnected { + if retryable_transport_error(error.code()) { + last_error = error; + continue; + } + return Err(error); + } + + match bounded( + transport.snapshot(mailbox), + limits.max_operation_runtime, + token, + ) + .await + { + Ok(snapshot) + if snapshot.uid_validity == initial.uid_validity + && snapshot.uid_next >= initial.uid_next => + { + return Ok(()) + } + Ok(_) => { + return Err(raise_error!( + "UIDONLY mailbox epoch changed while reconnecting".into(), + ErrorCode::Incompatible + )) + } + Err(error) if retryable_transport_error(error.code()) => last_error = error, + Err(error) => return Err(error), + } + } + Err(last_error) +} + +async fn inventory_with_reconnect( + transport: &mut T, + mailbox: &str, + initial: MailboxSnapshot, + cursor: u32, + high: u32, + limits: AcquisitionLimits, + token: &CancellationToken, +) -> BichonResult> { + let mut reconnects = 0; + loop { + match bounded( + transport.inventory_page(cursor, high, limits.page_size), + limits.max_operation_runtime, + token, + ) + .await + { + Err(error) if retryable_transport_error(error.code()) => { + recover_transport(transport, mailbox, initial, limits, token, &mut reconnects) + .await?; + } + result => return result, + } + } +} + +async fn fetch_with_reconnect( + transport: &mut T, + mailbox: &str, + initial: MailboxSnapshot, + uid: u32, + limits: AcquisitionLimits, + token: &CancellationToken, +) -> BichonResult { + let mut reconnects = 0; + loop { + match bounded( + transport.fetch_exact(uid, limits.max_literal_bytes), + limits.max_operation_runtime, + token, + ) + .await + { + Err(error) if retryable_transport_error(error.code()) => { + recover_transport(transport, mailbox, initial, limits, token, &mut reconnects) + .await?; + } + result => return result, + } + } +} + +async fn snapshot_with_reconnect( + transport: &mut T, + mailbox: &str, + initial: MailboxSnapshot, + limits: AcquisitionLimits, + token: &CancellationToken, +) -> BichonResult { + let mut reconnects = 0; + loop { + match bounded( + transport.snapshot(mailbox), + limits.max_operation_runtime, + token, + ) + .await + { + Err(error) if retryable_transport_error(error.code()) => { + recover_transport(transport, mailbox, initial, limits, token, &mut reconnects) + .await?; + } + result => return result, + } + } +} + +async fn flush_projection_batch( + archive: &mut A, + pending: &mut Vec, + limits: AcquisitionLimits, + token: &CancellationToken, +) -> BichonResult { + if pending.is_empty() { + return Ok(0); + } + let count = pending.len() as u64; + bounded( + archive.project_many(std::mem::take(pending)), + limits.max_operation_runtime, + token, + ) + .await?; + Ok(count) +} + +/// Reconciles one immutable UID range. Any ambiguity is an error, so callers +/// must persist `checkpoint` only from a returned report. +pub(crate) async fn run_acquisition( + transport: &mut T, + archive: &mut A, + mailbox: &str, + expected_uid_validity: Option, + limits: AcquisitionLimits, + token: CancellationToken, + mut progress: P, +) -> BichonResult +where + T: UidOnlyTransport, + A: CanonicalArchive, + P: FnMut(AcquisitionProgress) -> BichonResult<()>, +{ + let limits = limits.validate()?; + let snapshot = bounded( + transport.snapshot(mailbox), + limits.max_operation_runtime, + &token, + ) + .await?; + if snapshot.uid_validity == 0 || snapshot.uid_next == 0 { + return Err(raise_error!( + "UIDONLY EXAMINE omitted a valid UIDVALIDITY or UIDNEXT".into(), + ErrorCode::ImapUnexpectedResult + )); + } + if expected_uid_validity.is_some_and(|expected| expected != snapshot.uid_validity) { + return Err(raise_error!( + "UIDVALIDITY changed between mailbox discovery and UIDONLY EXAMINE".into(), + ErrorCode::Incompatible + )); + } + archive.begin_epoch(snapshot.uid_validity)?; + let planned = u64::from(snapshot.exists); + let high = snapshot.uid_next - 1; + let resume_after = archive.resume_after(); + if resume_after.is_some_and(|checkpoint| checkpoint > high) { + return Err(raise_error!( + "UIDONLY UIDNEXT moved behind the proven checkpoint".into(), + ErrorCode::Incompatible + )); + } + let full_snapshot = resume_after.is_none(); + let mut cursor = resume_after.map_or(1, |uid| uid.saturating_add(1)); + let mut inventoried = 0_u64; + let mut archived = 0_u64; + let mut downloaded = 0_u64; + let mut pending = Vec::with_capacity(UIDONLY_PROJECTION_BATCH_MESSAGES); + let mut pending_bytes = 0_u64; + + progress(AcquisitionProgress { + planned, + resolved: 0, + downloaded: 0, + })?; + + while high > 0 && cursor <= high { + let page = + inventory_with_reconnect(transport, mailbox, snapshot, cursor, high, limits, &token) + .await?; + if page.len() > limits.page_size as usize { + return Err(raise_error!( + "UIDONLY inventory page exceeded its requested bound".into(), + ErrorCode::ImapUnexpectedResult + )); + } + if page.is_empty() { + break; + } + + let mut previous = cursor - 1; + for item in &page { + if token.is_cancelled() { + return Err(raise_error!( + "UIDONLY acquisition cancelled".into(), + ErrorCode::InternalError + )); + } + if item.uid < cursor || item.uid > high || item.uid <= previous { + return Err(raise_error!( + "UIDONLY inventory was duplicate, unordered, or outside the fixed range".into(), + ErrorCode::ImapUnexpectedResult + )); + } + previous = item.uid; + inventoried = inventoried.checked_add(1).ok_or_else(|| { + raise_error!( + "UIDONLY inventory count overflow".into(), + ErrorCode::PayloadTooLarge + ) + })?; + } + let uids = page.iter().map(|item| item.uid).collect::>(); + let verified = bounded( + archive.verify_many(&uids), + limits.max_operation_runtime, + &token, + ) + .await?; + if verified.len() != page.len() { + return Err(raise_error!( + "UIDONLY receipt lookup returned the wrong result count".into(), + ErrorCode::InternalError + )); + } + + for (item, is_verified) in page.into_iter().zip(verified) { + if is_verified { + archived += 1; + continue; + } + + let message = + fetch_with_reconnect(transport, mailbox, snapshot, item.uid, limits, &token) + .await?; + if message.uid != item.uid { + return Err(raise_error!( + "UIDONLY exact fetch returned the wrong UID".into(), + ErrorCode::ImapUnexpectedResult + )); + } + let actual = message.body.len() as u64; + if actual > limits.max_literal_bytes { + return Err(raise_error!( + "UIDONLY exact body exceeded its pre-read budget".into(), + ErrorCode::PayloadTooLarge + )); + } + if !pending.is_empty() + && (pending.len() >= UIDONLY_PROJECTION_BATCH_MESSAGES + || pending_bytes.saturating_add(actual) > UIDONLY_PROJECTION_BATCH_BYTES as u64) + { + let stored = flush_projection_batch(archive, &mut pending, limits, &token).await?; + archived += stored; + downloaded += stored; + pending_bytes = 0; + } + pending_bytes += actual; + pending.push(message); + } + + let stored = flush_projection_batch(archive, &mut pending, limits, &token).await?; + archived += stored; + downloaded += stored; + pending_bytes = 0; + + cursor = previous.checked_add(1).unwrap_or(u32::MAX); + progress(AcquisitionProgress { + planned, + resolved: inventoried, + downloaded, + })?; + if previous == u32::MAX { + break; + } + } + + if (full_snapshot && inventoried != planned) || archived != inventoried { + return Err(raise_error!( + "UIDONLY inventory did not reconcile the EXAMINE message count".into(), + ErrorCode::ImapUnexpectedResult + )); + } + let final_snapshot = + snapshot_with_reconnect(transport, mailbox, snapshot, limits, &token).await?; + if final_snapshot.uid_validity != snapshot.uid_validity + || final_snapshot.uid_next < snapshot.uid_next + { + return Err(raise_error!( + "UIDONLY mailbox epoch changed during acquisition".into(), + ErrorCode::Incompatible + )); + } + Ok(AcquisitionReport { + uid_validity: snapshot.uid_validity, + uid_next: snapshot.uid_next, + exists: snapshot.exists, + checkpoint: (high > 0).then_some(high), + inventoried, + archived, + }) +} + +pub(crate) enum AcquisitionRoute { + Acquired { + report: AcquisitionReport, + source_scope: String, + }, + Legacy(Session>), +} + +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +enum CapabilityRoute { + Legacy, + UidOnly { message_limit: u32 }, +} + +fn is_capability(value: &str, expected: &str) -> bool { + value.trim().eq_ignore_ascii_case(expected) +} + +fn is_message_limit_marker(value: &str) -> bool { + value + .trim() + .get(.."MESSAGELIMIT".len()) + .is_some_and(|head| head.eq_ignore_ascii_case("MESSAGELIMIT")) +} + +fn classify_capabilities(capabilities: &[String]) -> BichonResult { + let uidonly = capabilities + .iter() + .filter(|value| is_capability(value, "UIDONLY")) + .count(); + let partial = capabilities + .iter() + .filter(|value| is_capability(value, "PARTIAL")) + .count(); + let limits: Vec<_> = capabilities + .iter() + .filter(|value| is_message_limit_marker(value)) + .collect(); + + // PARTIAL is a standalone RFC 9394 extension. It is not by itself + // evidence that ordinary mailbox views are limited. + if uidonly == 0 && limits.is_empty() { + return Ok(CapabilityRoute::Legacy); + } + if uidonly != 1 || partial != 1 || limits.len() != 1 { + return Err(raise_error!( + "Server advertised an incomplete or ambiguous UIDONLY capability set".into(), + ErrorCode::Incompatible + )); + } + let value = limits[0].trim(); + let prefix = "MESSAGELIMIT="; + let number = value + .get(prefix.len()..) + .filter(|_| { + value + .get(..prefix.len()) + .is_some_and(|head| head.eq_ignore_ascii_case(prefix)) + }) + .and_then(|value| value.parse::().ok()) + .filter(|value| *value > 0) + .ok_or_else(|| { + raise_error!( + "Server advertised an invalid UIDONLY MESSAGELIMIT".into(), + ErrorCode::Incompatible + ) + })?; + Ok(CapabilityRoute::UidOnly { + message_limit: number, + }) +} + +fn require_uidonly_opt_in( + account: &AccountModel, + capability_route: CapabilityRoute, +) -> BichonResult<()> { + if matches!(capability_route, CapabilityRoute::UidOnly { .. }) && !account.uidonly_enabled { + return Err(raise_error!( + "Server requires UIDONLY acquisition; enable uidonly_enabled for this account after reviewing the storage transition".into(), + ErrorCode::Incompatible + )); + } + Ok(()) +} + +/// Cached capabilities only decide whether legacy reconciliation may be +/// bypassed. The acquisition connection always reclassifies fresh capabilities. +pub(crate) fn account_requires_uidonly(account: &AccountModel) -> bool { + account.capabilities.as_ref().is_some_and(|capabilities| { + capabilities + .iter() + .any(|value| is_capability(value, "UIDONLY") || is_message_limit_marker(value)) + }) +} + +pub(crate) fn mailbox_has_uidonly_proof(account: &AccountModel, mailbox: &MailBox) -> bool { + source_scope(account) + .ok() + .as_deref() + .is_some_and(|scope| mailbox.uidonly_source_scope.as_deref() == Some(scope)) +} + +fn source_scope(account: &AccountModel) -> BichonResult { + let imap = account.imap.as_ref().ok_or_else(|| { + raise_error!( + "UIDONLY account has no IMAP configuration".into(), + ErrorCode::MissingConfiguration + ) + })?; + let host = imap.host.trim().trim_end_matches('.').to_ascii_lowercase(); + let principal = account.login_name.as_ref().unwrap_or(&account.email); + if host.is_empty() || principal.is_empty() || imap.port == 0 { + return Err(raise_error!( + "UIDONLY source identity is incomplete".into(), + ErrorCode::MissingConfiguration + )); + } + let mut hasher = blake3::Hasher::new(); + hasher.update(b"bichon-uidonly-source-v2\0"); + for field in [host.as_bytes(), principal.as_bytes()] { + hasher.update(&(field.len() as u64).to_be_bytes()); + hasher.update(field); + } + hasher.update(&imap.port.to_be_bytes()); + Ok(hasher.finalize().to_hex().to_string()) +} + +fn connection_scope(account: &AccountModel) -> BichonResult { + let imap = account.imap.as_ref().ok_or_else(|| { + raise_error!( + "UIDONLY account has no IMAP configuration".into(), + ErrorCode::MissingConfiguration + ) + })?; + let encryption = match imap.encryption { + Encryption::Ssl => 1_u8, + Encryption::StartTls => 2, + Encryption::None => 3, + }; + let mut hasher = blake3::Hasher::new(); + hasher.update(b"bichon-uidonly-connection-v1\0"); + hasher.update(source_scope(account)?.as_bytes()); + hasher.update(&[encryption, u8::from(account.use_dangerous)]); + match imap.use_proxy { + Some(proxy) => { + hasher.update(&[1]); + hasher.update(&proxy.to_be_bytes()); + } + None => { + hasher.update(&[0]); + } + } + hasher.update(&[match imap.auth.auth_type { + AuthType::Password => 1, + AuthType::OAuth2 => 2, + }]); + Ok(hasher.finalize().to_hex().to_string()) +} + +fn protocol_limits(max_literal_bytes: u64) -> BichonResult { + if max_literal_bytes == 0 || max_literal_bytes > u64::from(u32::MAX) { + return Err(raise_error!( + "UIDONLY message-size limit is outside the supported range".into(), + ErrorCode::InvalidParameter + )); + } + let literal = usize::try_from(max_literal_bytes).map_err(|_| { + raise_error!( + "UIDONLY message-size limit is not representable".into(), + ErrorCode::InvalidParameter + ) + })?; + let response = literal.checked_add(128 * 1024).ok_or_else(|| { + raise_error!( + "UIDONLY response-size limit overflow".into(), + ErrorCode::InvalidParameter + ) + })?; + let command = response.checked_add(128 * 1024).ok_or_else(|| { + raise_error!( + "UIDONLY command-size limit overflow".into(), + ErrorCode::InvalidParameter + ) + })?; + Ok(UidOnlyLimits { + max_control_line_bytes: 64 * 1024, + max_literal_bytes: literal, + max_response_bytes: response, + max_command_literal_bytes: literal, + max_command_response_bytes: command, + // A 1,000-message inventory page plus bounded unsolicited mailbox + // updates must still fit without making the response count unbounded. + max_command_responses: 2_048, + max_command_runtime: Duration::from_secs(5 * 60), + }) +} + +fn protocol_probe_literal(max_literal_bytes: u64) -> u64 { + let representable = u64::from(u32::MAX).min(usize::MAX.saturating_sub(256 * 1024) as u64); + max_literal_bytes.clamp(1, representable) +} + +fn imap_error_code(error: &async_imap::error::Error) -> ErrorCode { + match error { + async_imap::error::Error::ConnectionLost => ErrorCode::NetworkError, + async_imap::error::Error::Io(error) + if matches!( + error.kind(), + std::io::ErrorKind::BrokenPipe + | std::io::ErrorKind::ConnectionReset + | std::io::ErrorKind::ConnectionAborted + | std::io::ErrorKind::NotConnected + | std::io::ErrorKind::UnexpectedEof + | std::io::ErrorKind::TimedOut + | std::io::ErrorKind::WriteZero + ) => + { + ErrorCode::NetworkError + } + _ => ErrorCode::ImapUnexpectedResult, + } +} + +struct SessionUidOnlyTransport { + session: Session>, + handle: UidOnlyHandle, + account: AccountModel, + limits: UidOnlyLimits, + connection_scope: String, +} + +impl SessionUidOnlyTransport { + fn healthy_handle(handle: &UidOnlyHandle) -> BichonResult<()> { + if let Some(reason) = handle.poison_reason() { + return Err(raise_error!(reason, ErrorCode::ImapUnexpectedResult)); + } + handle.ensure_active().map_err(|_| { + raise_error!( + "UIDONLY transport is not active".into(), + ErrorCode::ImapUnexpectedResult + ) + }) + } + + fn healthy(&self) -> BichonResult<()> { + Self::healthy_handle(&self.handle) + } + + async fn enable( + session: &mut Session>, + handle: &UidOnlyHandle, + ) -> BichonResult<()> { + session + .run_command_and_check_ok("ENABLE UIDONLY") + .await + .map_err(|error| { + raise_error!( + "Server did not enable UIDONLY".into(), + imap_error_code(&error) + ) + })?; + Self::healthy_handle(handle) + } + + async fn collect_fetches( + &mut self, + set: String, + query: impl AsRef, + ) -> BichonResult> { + let query = query.as_ref().to_string(); + let operation = async { + let stream = self.session.uid_fetch(set, query).await.map_err(|error| { + raise_error!( + "UIDONLY fetch command failed".into(), + imap_error_code(&error) + ) + })?; + stream.try_collect::>().await.map_err(|error| { + raise_error!( + "UIDONLY fetch response failed".into(), + imap_error_code(&error) + ) + }) + }; + let result = AssertUnwindSafe(operation) + .catch_unwind() + .await + .map_err(|_| { + raise_error!( + "UIDONLY parser rejected a malformed response".into(), + ErrorCode::ImapUnexpectedResult + ) + })??; + self.healthy()?; + Ok(result) + } +} + +impl UidOnlyTransport for SessionUidOnlyTransport { + async fn snapshot(&mut self, mailbox: &str) -> BichonResult { + self.healthy()?; + let mailbox = AssertUnwindSafe(self.session.examine(mailbox)) + .catch_unwind() + .await + .map_err(|_| { + raise_error!( + "UIDONLY parser rejected EXAMINE".into(), + ErrorCode::ImapUnexpectedResult + ) + })? + .map_err(|error| { + raise_error!("UIDONLY EXAMINE failed".into(), imap_error_code(&error)) + })?; + self.healthy()?; + Ok(MailboxSnapshot { + exists: mailbox.exists, + uid_validity: mailbox.uid_validity.ok_or_else(|| { + raise_error!( + "UIDONLY EXAMINE omitted UIDVALIDITY".into(), + ErrorCode::ImapUnexpectedResult + ) + })?, + uid_next: mailbox.uid_next.ok_or_else(|| { + raise_error!( + "UIDONLY EXAMINE omitted UIDNEXT".into(), + ErrorCode::ImapUnexpectedResult + ) + })?, + }) + } + + async fn inventory_page( + &mut self, + cursor: u32, + high: u32, + page_size: u32, + ) -> BichonResult> { + let (set, query) = inventory_args(cursor, high, page_size).map_err(|_| { + raise_error!( + "Invalid UIDONLY inventory bounds".into(), + ErrorCode::InvalidParameter + ) + })?; + self.collect_fetches(set, query) + .await? + .into_iter() + .map(|fetch| { + let uid = fetch.uid.ok_or_else(|| { + raise_error!( + "UIDONLY inventory omitted UID".into(), + ErrorCode::ImapUnexpectedResult + ) + })?; + if fetch.message != uid { + return Err(raise_error!( + "UIDONLY inventory leading UID did not match UID data item".into(), + ErrorCode::ImapUnexpectedResult + )); + } + fetch.size.ok_or_else(|| { + raise_error!( + "UIDONLY inventory omitted RFC822.SIZE".into(), + ErrorCode::ImapUnexpectedResult + ) + })?; + Ok(InventoryItem { uid }) + }) + .collect() + } + + async fn fetch_exact(&mut self, uid: u32, literal_budget: u64) -> BichonResult { + let limit = usize::try_from(literal_budget).map_err(|_| { + raise_error!( + "UIDONLY literal budget is not representable".into(), + ErrorCode::InvalidParameter + ) + })?; + self.handle + .arm_next_fetch_literal_limit(limit) + .map_err(|_| { + raise_error!( + "UIDONLY exact fetch could not arm its pre-read limit".into(), + ErrorCode::ImapUnexpectedResult + ) + })?; + let before = self.handle.literal_bytes_received(); + let (set, query) = exact_args(uid).map_err(|_| { + raise_error!( + "Invalid UIDONLY exact UID".into(), + ErrorCode::InvalidParameter + ) + })?; + let mut fetched = self.collect_fetches(set, query).await?; + if fetched.len() != 1 { + return Err(raise_error!( + "UIDONLY exact fetch did not return exactly one message".into(), + ErrorCode::ImapUnexpectedResult + )); + } + let fetch = fetched.pop().expect("length checked"); + if fetch.message != uid || fetch.uid != Some(uid) || fetch.size.is_none() { + return Err(raise_error!( + "UIDONLY exact fetch returned mismatched metadata".into(), + ErrorCode::ImapUnexpectedResult + )); + } + let raw = fetch.body().ok_or_else(|| { + raise_error!( + "UIDONLY exact fetch omitted its full literal body".into(), + ErrorCode::ImapUnexpectedResult + ) + })?; + let after = self.handle.literal_bytes_received(); + if after.checked_sub(before) != Some(raw.len() as u64) { + return Err(raise_error!( + "UIDONLY exact body did not match literal accounting".into(), + ErrorCode::ImapUnexpectedResult + )); + } + Ok(UidOnlyMessage { + uid, + body: raw.to_vec(), + }) + } + + async fn reconnect(&mut self, page_size: u32) -> BichonResult<()> { + let current = AccountModel::get(self.account.id)?; + if connection_scope(¤t)? != self.connection_scope { + return Err(raise_error!( + "UIDONLY account connection changed while reconnecting".into(), + ErrorCode::Incompatible + )); + } + let connection = + ImapConnectionManager::build_uidonly(&self.account, self.limits.clone()).await?; + let CapabilityRoute::UidOnly { message_limit } = + classify_capabilities(&connection.capabilities)? + else { + return Err(raise_error!( + "UIDONLY reconnect lost required capabilities".into(), + ErrorCode::Incompatible + )); + }; + if message_limit < page_size { + return Err(raise_error!( + "UIDONLY reconnect reduced MESSAGELIMIT below the fixed page size".into(), + ErrorCode::Incompatible + )); + } + let mut session = connection.session; + Self::enable(&mut session, &connection.handle).await?; + self.session = session; + self.handle = connection.handle; + Ok(()) + } +} + +struct BichonCanonicalArchive { + account_id: u64, + mailbox_id: u64, + source_scope: String, + uid_validity: Option, + resume_after: Option, +} + +impl CanonicalArchive for BichonCanonicalArchive { + fn resume_after(&self) -> Option { + self.resume_after + } + + fn begin_epoch(&mut self, uid_validity: u32) -> BichonResult<()> { + if self + .uid_validity + .is_some_and(|existing| existing != uid_validity) + { + return Err(raise_error!( + "UIDONLY archive epoch changed".into(), + ErrorCode::Incompatible + )); + } + self.uid_validity = Some(uid_validity); + Ok(()) + } + + async fn verify_many(&mut self, uids: &[u32]) -> BichonResult> { + verify_uidonly_projections( + self.account_id, + self.mailbox_id, + self.uid_validity.expect("begin_epoch is called first"), + uids, + &self.source_scope, + ) + } + + async fn project_many(&mut self, messages: Vec) -> BichonResult<()> { + project_uidonly_messages( + messages, + self.account_id, + self.mailbox_id, + self.uid_validity.expect("begin_epoch is called first"), + &self.source_scope, + ) + .await + } +} + +pub(crate) async fn connect_and_acquire_or_legacy

( + account: &AccountModel, + mailbox: &MailBox, + force_uidonly: bool, + token: CancellationToken, + progress: P, +) -> BichonResult +where + P: FnMut(AcquisitionProgress) -> BichonResult<()>, +{ + let configured_scope = source_scope(account)?; + let known_limited = force_uidonly + || account_requires_uidonly(account) + || mailbox.uidonly_source_scope.as_deref() == Some(&configured_scope); + if known_limited + && account + .archive_rules + .as_ref() + .is_some_and(|rules| rules.enabled) + { + return Err(raise_error!( + "UIDONLY acquisition does not yet support enabled archive rules".into(), + ErrorCode::Incompatible + )); + } + let max_literal = account + .max_email_size_bytes + .unwrap_or(DEFAULT_MAX_EMAIL_SIZE); + let wire_limits = protocol_limits(protocol_probe_literal(max_literal))?; + let connection = ImapConnectionManager::build_uidonly(account, wire_limits.clone()).await?; + let capability_route = classify_capabilities(&connection.capabilities)?; + require_uidonly_opt_in(account, capability_route)?; + if capability_route == CapabilityRoute::Legacy { + if known_limited { + return Err(raise_error!( + "Known limited server omitted UIDONLY capabilities on the acquisition connection" + .into(), + ErrorCode::Incompatible + )); + } + return Ok(AcquisitionRoute::Legacy(connection.session)); + } + let CapabilityRoute::UidOnly { message_limit } = capability_route else { + unreachable!("legacy returned above") + }; + let wire_limits = protocol_limits(max_literal)?; + if account + .archive_rules + .as_ref() + .is_some_and(|rules| rules.enabled) + { + return Err(raise_error!( + "UIDONLY acquisition does not yet support enabled archive rules".into(), + ErrorCode::Incompatible + )); + } + + let frozen_connection = connection_scope(account)?; + let mut transport = SessionUidOnlyTransport { + session: connection.session, + handle: connection.handle, + account: account.clone(), + limits: wire_limits, + connection_scope: frozen_connection.clone(), + }; + SessionUidOnlyTransport::enable(&mut transport.session, &transport.handle).await?; + + let frozen_scope = configured_scope; + let resume_after = (mailbox.uidonly_source_scope.as_deref() == Some(&frozen_scope)) + .then_some(mailbox.highest_uid) + .flatten(); + let mut archive = BichonCanonicalArchive { + account_id: account.id, + mailbox_id: mailbox.id, + source_scope: frozen_scope.clone(), + uid_validity: None, + resume_after, + }; + let mut limits = AcquisitionLimits::bounded(max_literal); + limits.page_size = UIDONLY_DEFAULT_PAGE_SIZE.min(message_limit).max(1); + let result = run_acquisition( + &mut transport, + &mut archive, + &mailbox.encoded_name(), + mailbox.uid_validity, + limits, + token, + progress, + ) + .await; + let report = match result { + Ok(report) => report, + Err(error) => return Err(error), + }; + + let current = AccountModel::get(account.id)?; + if source_scope(¤t)? != frozen_scope + || connection_scope(¤t)? != frozen_connection + || current + .archive_rules + .as_ref() + .is_some_and(|rules| rules.enabled) + { + return Err(raise_error!( + "UIDONLY account source, connection policy, or archive rules changed during acquisition" + .into(), + ErrorCode::Incompatible + )); + } + transport.session.logout().await.ok(); + Ok(AcquisitionRoute::Acquired { + report, + source_scope: frozen_scope, + }) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::account::entity::{AuthConfig, ImapConfig}; + use crate::imap::client::Client; + use crate::imap::mock_server::{examine_response, MockImapServer, MockImapServerHandle}; + use std::collections::{BTreeMap, BTreeSet, VecDeque}; + + fn limits() -> AcquisitionLimits { + AcquisitionLimits { + max_literal_bytes: 64 * 1024, + max_operation_runtime: Duration::from_secs(20), + page_size: 2, + } + } + + #[test] + fn capability_routing_ignores_standalone_partial_and_rejects_partial_uidonly_sets() { + assert_eq!( + classify_capabilities(&["IMAP4rev1".into(), "PARTIAL".into()]).unwrap(), + CapabilityRoute::Legacy + ); + assert_eq!( + classify_capabilities(&[ + "IMAP4rev1".into(), + "uidonly".into(), + "partial".into(), + "messagelimit=10000".into(), + ]) + .unwrap(), + CapabilityRoute::UidOnly { + message_limit: 10_000 + } + ); + for capabilities in [ + vec!["UIDONLY".into(), "PARTIAL".into()], + vec!["MESSAGELIMIT=1000".into(), "PARTIAL".into()], + vec!["UIDONLY".into(), "PARTIAL".into(), "MESSAGELIMIT=0".into()], + ] { + assert!(classify_capabilities(&capabilities).is_err()); + } + let cached = |capability: &str| AccountModel { + capabilities: Some(vec![capability.into()]), + ..Default::default() + }; + assert!(!account_requires_uidonly(&cached("PARTIAL"))); + assert!(account_requires_uidonly(&cached("UIDONLY"))); + assert!(account_requires_uidonly(&cached("MESSAGELIMIT=10000"))); + } + + #[test] + fn uidonly_route_requires_explicit_account_opt_in() { + let route = CapabilityRoute::UidOnly { + message_limit: 10_000, + }; + let error = require_uidonly_opt_in(&AccountModel::default(), route).unwrap_err(); + assert_eq!(error.code(), ErrorCode::Incompatible); + assert!(error.to_string().contains("enable uidonly_enabled")); + + let enabled = AccountModel { + uidonly_enabled: true, + ..Default::default() + }; + assert!(require_uidonly_opt_in(&enabled, route).is_ok()); + assert!(require_uidonly_opt_in(&AccountModel::default(), CapabilityRoute::Legacy).is_ok()); + } + + #[test] + fn capability_probe_accepts_limits_that_only_uidonly_rejects() { + for configured in [0, u64::from(u32::MAX) + 1] { + assert!(protocol_limits(protocol_probe_literal(configured)).is_ok()); + assert!(protocol_limits(configured).is_err()); + } + } + + #[test] + fn source_identity_is_stable_across_connection_policy_changes() { + let account = AccountModel { + email: "alice@example.invalid".into(), + imap: Some(ImapConfig { + host: "IMAP.EXAMPLE.INVALID.".into(), + port: 993, + encryption: Encryption::Ssl, + auth: AuthConfig { + auth_type: AuthType::Password, + password: None, + }, + use_proxy: None, + }), + ..Default::default() + }; + let source = source_scope(&account).unwrap(); + let connection = connection_scope(&account).unwrap(); + for changed in [ + AccountModel { + use_dangerous: true, + ..account.clone() + }, + AccountModel { + imap: account.imap.clone().map(|mut imap| { + imap.use_proxy = Some(7); + imap + }), + ..account.clone() + }, + AccountModel { + imap: account.imap.clone().map(|mut imap| { + imap.auth.auth_type = AuthType::OAuth2; + imap + }), + ..account.clone() + }, + ] { + assert_eq!(source_scope(&changed).unwrap(), source); + assert_ne!(connection_scope(&changed).unwrap(), connection); + } + let mut other_principal = account; + other_principal.login_name = Some("other@example.invalid".into()); + assert_ne!(source_scope(&other_principal).unwrap(), source); + } + + struct FakeTransport { + snapshot: MailboxSnapshot, + pages: VecDeque>, + messages: BTreeMap>, + fetched: Vec, + expected_cursor: Option, + } + + impl FakeTransport { + fn sparse() -> Self { + let body = |uid| { + format!("From: sender@invalid\r\nMessage-ID: <{uid}@invalid>\r\n\r\nbody") + .into_bytes() + }; + Self { + snapshot: MailboxSnapshot { + exists: 3, + uid_validity: 77, + uid_next: 51, + }, + pages: VecDeque::from([ + vec![InventoryItem { uid: 2 }, InventoryItem { uid: 30 }], + vec![InventoryItem { uid: 50 }], + vec![], + ]), + messages: BTreeMap::from([(2, body(2)), (30, body(30)), (50, body(50))]), + fetched: Vec::new(), + expected_cursor: None, + } + } + } + + impl UidOnlyTransport for FakeTransport { + async fn snapshot(&mut self, _mailbox: &str) -> BichonResult { + Ok(self.snapshot) + } + + async fn inventory_page( + &mut self, + cursor: u32, + _high: u32, + _page_size: u32, + ) -> BichonResult> { + if let Some(expected) = self.expected_cursor.take() { + assert_eq!(cursor, expected); + } + Ok(self.pages.pop_front().unwrap_or_default()) + } + + async fn fetch_exact( + &mut self, + uid: u32, + literal_budget: u64, + ) -> BichonResult { + self.fetched.push(uid); + let raw = self.messages.get(&uid).cloned().ok_or_else(|| { + raise_error!( + "synthetic missing body".into(), + ErrorCode::ImapUnexpectedResult + ) + })?; + if raw.len() as u64 > literal_budget { + return Err(raise_error!( + "synthetic literal budget".into(), + ErrorCode::PayloadTooLarge + )); + } + Ok(UidOnlyMessage { uid, body: raw }) + } + } + + struct FlakyTransport { + inner: FakeTransport, + inventory_failures: u32, + fetch_failures: u32, + reconnects: u32, + inventory_attempts: u32, + fetch_attempts: Vec, + snapshot_after_reconnect: Option, + } + + impl FlakyTransport { + fn sparse() -> Self { + Self { + inner: FakeTransport::sparse(), + inventory_failures: 0, + fetch_failures: 0, + reconnects: 0, + inventory_attempts: 0, + fetch_attempts: Vec::new(), + snapshot_after_reconnect: None, + } + } + } + + impl UidOnlyTransport for FlakyTransport { + async fn snapshot(&mut self, mailbox: &str) -> BichonResult { + self.inner.snapshot(mailbox).await + } + + async fn inventory_page( + &mut self, + cursor: u32, + high: u32, + page_size: u32, + ) -> BichonResult> { + self.inventory_attempts += 1; + if self.inventory_failures > 0 { + self.inventory_failures -= 1; + return Err(raise_error!( + "synthetic disconnect".into(), + ErrorCode::NetworkError + )); + } + self.inner.inventory_page(cursor, high, page_size).await + } + + async fn fetch_exact( + &mut self, + uid: u32, + literal_budget: u64, + ) -> BichonResult { + self.fetch_attempts.push(uid); + if self.fetch_failures > 0 { + self.fetch_failures -= 1; + return Err(raise_error!( + "synthetic disconnect".into(), + ErrorCode::NetworkError + )); + } + self.inner.fetch_exact(uid, literal_budget).await + } + + async fn reconnect(&mut self, _page_size: u32) -> BichonResult<()> { + self.reconnects += 1; + if let Some(snapshot) = self.snapshot_after_reconnect.take() { + self.inner.snapshot = snapshot; + } + Ok(()) + } + } + + #[derive(Default)] + struct FakeArchive { + verified: BTreeSet, + projected: Vec, + fail_project: bool, + resume_after: Option, + } + + impl CanonicalArchive for FakeArchive { + fn resume_after(&self) -> Option { + self.resume_after + } + + async fn verify_many(&mut self, uids: &[u32]) -> BichonResult> { + Ok(uids.iter().map(|uid| self.verified.contains(uid)).collect()) + } + + async fn project_many(&mut self, messages: Vec) -> BichonResult<()> { + if self.fail_project { + return Err(raise_error!( + "synthetic projection failure".into(), + ErrorCode::InternalError + )); + } + for message in messages { + self.projected.push(message.uid); + self.verified.insert(message.uid); + } + Ok(()) + } + } + + async fn run_test( + transport: &mut T, + archive: &mut A, + ) -> BichonResult { + run_acquisition( + transport, + archive, + "synthetic", + Some(77), + limits(), + CancellationToken::new(), + |_| Ok(()), + ) + .await + } + + #[tokio::test] + async fn sparse_snapshot_streams_and_checkpoints_fixed_high() { + let mut transport = FakeTransport::sparse(); + let mut archive = FakeArchive::default(); + let report = run_test(&mut transport, &mut archive).await.unwrap(); + assert_eq!(report.checkpoint, Some(50)); + assert_eq!((report.inventoried, report.archived), (3, 3)); + assert_eq!(transport.fetched, [2, 30, 50]); + } + + #[tokio::test] + async fn reconnect_retries_the_same_inventory_cursor() { + let mut transport = FlakyTransport::sparse(); + transport.inventory_failures = 1; + let report = run_test(&mut transport, &mut FakeArchive::default()) + .await + .unwrap(); + + assert_eq!(report.checkpoint, Some(50)); + assert_eq!(transport.reconnects, 1); + assert_eq!(transport.inventory_attempts, 3); + } + + #[tokio::test] + async fn reconnect_retries_one_exact_uid_without_double_projection() { + let mut transport = FlakyTransport::sparse(); + transport.fetch_failures = 1; + let mut archive = FakeArchive::default(); + let report = run_test(&mut transport, &mut archive).await.unwrap(); + + assert_eq!(report.archived, 3); + assert_eq!(transport.reconnects, 1); + assert_eq!(transport.fetch_attempts, [2, 2, 30, 50]); + assert_eq!(archive.projected, [2, 30, 50]); + } + + #[tokio::test] + async fn reconnect_rejects_changed_uidvalidity_or_lower_uidnext() { + for changed in [ + MailboxSnapshot { + uid_validity: 78, + ..FakeTransport::sparse().snapshot + }, + MailboxSnapshot { + uid_next: 50, + ..FakeTransport::sparse().snapshot + }, + ] { + let mut transport = FlakyTransport::sparse(); + transport.inventory_failures = 1; + transport.snapshot_after_reconnect = Some(changed); + let error = run_test(&mut transport, &mut FakeArchive::default()) + .await + .unwrap_err(); + + assert_eq!(error.code(), ErrorCode::Incompatible); + assert!(transport.inner.fetched.is_empty()); + } + } + + #[tokio::test] + async fn persistent_disconnect_uses_only_three_reconnects() { + let mut transport = FlakyTransport::sparse(); + transport.inventory_failures = 4; + let error = run_test(&mut transport, &mut FakeArchive::default()) + .await + .unwrap_err(); + + assert_eq!(error.code(), ErrorCode::NetworkError); + assert_eq!(transport.reconnects, 3); + assert_eq!(transport.inventory_attempts, 4); + } + + #[tokio::test] + async fn proven_checkpoint_scans_only_new_uids() { + let mut transport = FakeTransport::sparse(); + transport.snapshot.exists = 4; + transport.snapshot.uid_next = 60; + transport.pages = VecDeque::from([vec![InventoryItem { uid: 55 }]]); + transport + .messages + .insert(55, b"Subject: new\r\n\r\nbody".to_vec()); + transport.expected_cursor = Some(51); + let report = run_test( + &mut transport, + &mut FakeArchive { + resume_after: Some(50), + ..Default::default() + }, + ) + .await + .unwrap(); + assert_eq!((report.inventoried, report.archived), (1, 1)); + assert_eq!(report.checkpoint, Some(59)); + assert_eq!(transport.fetched, [55]); + } + + #[tokio::test] + async fn uidnext_behind_proven_checkpoint_fails_closed() { + let mut transport = FakeTransport::sparse(); + let error = run_test( + &mut transport, + &mut FakeArchive { + resume_after: Some(51), + ..Default::default() + }, + ) + .await + .unwrap_err(); + assert_eq!(error.code(), ErrorCode::Incompatible); + assert!(transport.fetched.is_empty()); + } + + #[tokio::test] + async fn restart_skips_verified_receipts_without_a_secondary_ledger() { + let mut transport = FakeTransport::sparse(); + let mut archive = FakeArchive { + verified: BTreeSet::from([2, 30]), + ..Default::default() + }; + let report = run_test(&mut transport, &mut archive).await.unwrap(); + assert_eq!(report.archived, 3); + assert_eq!(transport.fetched, [50]); + } + + #[tokio::test] + async fn short_inventory_never_returns_a_checkpoint() { + let mut transport = FakeTransport::sparse(); + transport.snapshot.exists = 4; + let error = run_test(&mut transport, &mut FakeArchive::default()) + .await + .unwrap_err(); + assert_eq!(error.code(), ErrorCode::ImapUnexpectedResult); + } + + #[tokio::test] + async fn cancellation_and_projection_failure_abort_the_run() { + let token = CancellationToken::new(); + let cancel = token.clone(); + let mut transport = FakeTransport::sparse(); + let error = run_acquisition( + &mut transport, + &mut FakeArchive::default(), + "synthetic", + Some(77), + limits(), + token, + move |progress| { + if progress.resolved > 0 { + cancel.cancel(); + } + Ok(()) + }, + ) + .await + .unwrap_err(); + assert!(error.to_string().contains("cancelled")); + + let mut transport = FakeTransport::sparse(); + let error = run_test( + &mut transport, + &mut FakeArchive { + fail_project: true, + ..Default::default() + }, + ) + .await + .unwrap_err(); + assert!(error.to_string().contains("synthetic projection failure")); + } + + struct LogicalMillion { + cursor: u32, + } + + impl UidOnlyTransport for LogicalMillion { + async fn snapshot(&mut self, _mailbox: &str) -> BichonResult { + Ok(MailboxSnapshot { + exists: 1_000_000, + uid_validity: 9, + uid_next: 1_000_001, + }) + } + + async fn inventory_page( + &mut self, + cursor: u32, + high: u32, + page_size: u32, + ) -> BichonResult> { + assert_eq!(cursor, self.cursor); + let end = high.min(cursor + page_size - 1); + self.cursor = end + 1; + Ok((cursor..=end).map(|uid| InventoryItem { uid }).collect()) + } + + async fn fetch_exact( + &mut self, + _uid: u32, + _literal_budget: u64, + ) -> BichonResult { + unreachable!("every logical receipt verifies") + } + } + + struct AllVerified; + + impl CanonicalArchive for AllVerified { + async fn verify_many(&mut self, uids: &[u32]) -> BichonResult> { + Ok(vec![true; uids.len()]) + } + + async fn project_many(&mut self, _messages: Vec) -> BichonResult<()> { + unreachable!("every logical receipt verifies") + } + } + + #[tokio::test] + async fn million_message_inventory_is_page_bounded() { + let mut transport = LogicalMillion { cursor: 1 }; + let mut bounded = limits(); + bounded.page_size = 1_000; + bounded.max_operation_runtime = Duration::from_secs(30); + let report = run_acquisition( + &mut transport, + &mut AllVerified, + "synthetic", + Some(9), + bounded, + CancellationToken::new(), + |_| Ok(()), + ) + .await + .unwrap(); + assert_eq!(report.inventoried, 1_000_000); + assert_eq!(report.checkpoint, Some(1_000_000)); + } + + fn inventory_response(entries: &[(u32, u32)]) -> Vec { + let mut response = + b"* 3 EXISTS\r\n* OK [UIDNEXT 10] unchanged\r\n* 9 UIDFETCH (FLAGS (\\Seen))\r\n" + .to_vec(); + for (uid, size) in entries { + response.extend_from_slice( + format!("* {uid} UIDFETCH (UID {uid} RFC822.SIZE {size})\r\n").as_bytes(), + ); + } + response.extend_from_slice(b"{TAG} OK inventory completed\r\n"); + response + } + + fn exact_response(uid: u32, reported_size: u32, raw: &[u8]) -> Vec { + let mut response = format!( + "* {uid} UIDFETCH (UID {uid} RFC822.SIZE {reported_size} BODY[] {{{}}}\r\n", + raw.len() + ) + .into_bytes(); + response.extend_from_slice(raw); + response.extend_from_slice(b")\r\n{TAG} OK exact fetch completed\r\n"); + response + } + + async fn connect_fake_transport(server: &MockImapServerHandle) -> SessionUidOnlyTransport { + let account = AccountModel { + email: "synthetic-user".into(), + imap: Some(ImapConfig { + host: server.host(), + port: server.port(), + encryption: Encryption::None, + auth: AuthConfig { + auth_type: AuthType::Password, + password: None, + }, + use_proxy: None, + }), + ..Default::default() + }; + let limits = protocol_limits(64 * 1024).expect("bounded protocol limits"); + let client = Client::connection( + &server.host(), + &Encryption::None, + server.port(), + None, + false, + ) + .await + .expect("localhost connection"); + let (client, handle) = client + .with_uidonly(limits.clone()) + .expect("install UIDONLY guard"); + let mut session = client + .login("synthetic-user", "synthetic-secret") + .await + .expect("synthetic login"); + let capabilities = session.capabilities().await.expect("capabilities"); + assert!(capabilities.has_str("UIDONLY")); + assert!(capabilities.has_str("PARTIAL")); + assert!(capabilities.has_str("MESSAGELIMIT=2")); + session + .run_command_and_check_ok("ENABLE UIDONLY") + .await + .expect("enable UIDONLY"); + handle.ensure_active().expect("UIDONLY confirmed"); + SessionUidOnlyTransport { + session, + handle, + connection_scope: connection_scope(&account).unwrap(), + account, + limits, + } + } + + #[tokio::test] + async fn tcp_fake_yahoo_uidonly_pages_sparse_uids_and_checkpoints_after_verification() { + let messages = [ + (2, b"".as_slice()), + (7, b"Subject: seven\r\n\r\nseven".as_slice()), + (9, b"Subject: nine\r\n\r\nnine".as_slice()), + ]; + let server = MockImapServer::new() + .greeting("* OK synthetic Yahoo-like IMAP ready\r\n") + .respond("LOGIN", "{TAG} OK LOGIN completed\r\n") + .respond( + "CAPABILITY", + "* CAPABILITY IMAP4rev1 ENABLE UIDONLY PARTIAL MESSAGELIMIT=2\r\n{TAG} OK CAPABILITY completed\r\n", + ) + .respond( + "ENABLE UIDONLY", + "* ENABLED UIDONLY\r\n{TAG} OK ENABLE completed\r\n", + ) + .respond("EXAMINE", examine_response("Synthetic", 3, 77, 10)) + .respond("UID FETCH 1:9", inventory_response(&[(2, 999), (7, 1)])) + .respond("UID FETCH 8:9", inventory_response(&[(9, 0)])) + // The reported sizes are deliberately advisory and wrong. The + // literal byte count is the acquisition/accounting authority. + .respond("UID FETCH 2 (", exact_response(2, 1, messages[0].1)) + .respond("UID FETCH 7 (", exact_response(7, 2, messages[1].1)) + .respond("UID FETCH 9 (", exact_response(9, 3, messages[2].1)) + .start() + .await; + let mut transport = connect_fake_transport(&server).await; + let mut archive = FakeArchive::default(); + let report = run_acquisition( + &mut transport, + &mut archive, + "Synthetic", + Some(77), + limits(), + CancellationToken::new(), + |_| Ok(()), + ) + .await + .expect("complete fixed snapshot"); + + assert_eq!(report.checkpoint, Some(9)); + assert_eq!((report.inventoried, report.archived), (3, 3)); + assert_eq!(archive.projected, [2, 7, 9]); + + let commands = server.commands(); + let commands = commands.join("\n"); + assert_eq!(commands.matches(" EXAMINE ").count(), 2); + assert!(commands.contains(" UID FETCH 1:9 (UID RFC822.SIZE) (PARTIAL 1:2)")); + assert!(commands.contains(" UID FETCH 8:9 (UID RFC822.SIZE) (PARTIAL 1:2)")); + assert_eq!(commands.matches("BODY.PEEK[]").count(), 3); + let commands = commands.to_ascii_uppercase(); + assert!( + ![" STORE ", " MOVE ", " COPY ", " DELETE ", " EXPUNGE", " CLOSE",] + .iter() + .any(|forbidden| commands.contains(forbidden)) + ); + } +} diff --git a/crates/core/src/import/mod.rs b/crates/core/src/import/mod.rs index bd8db2ae..917c399b 100644 --- a/crates/core/src/import/mod.rs +++ b/crates/core/src/import/mod.rs @@ -122,6 +122,7 @@ impl ImportEmls { uid_next: None, uid_validity: None, highest_uid: None, + uidonly_source_scope: None, }; let mailbox_id = mailbox.id; // Upsert the mailbox, creating it if it doesn't exist @@ -413,6 +414,7 @@ pub(super) fn resolve_mailbox(account: &AccountModel, folder: &str) -> BichonRes uid_next: None, uid_validity: None, highest_uid: None, + uidonly_source_scope: None, }; let mailbox_id = mailbox.id; MailBox::batch_upsert(&[mailbox])?; diff --git a/crates/smtp/src/server.rs b/crates/smtp/src/server.rs index ad38ff48..4f3b06b8 100644 --- a/crates/smtp/src/server.rs +++ b/crates/smtp/src/server.rs @@ -656,6 +656,7 @@ async fn parse_email(data: &[u8], session: &Session) -> BichonResult<()> { uid_next: None, uid_validity: None, highest_uid: None, + uidonly_source_scope: None, }; if let Err(e) = MailBox::batch_upsert(&[mailbox]) { diff --git a/web/src/api/account/api.ts b/web/src/api/account/api.ts index ae1604c1..6c1dea48 100644 --- a/web/src/api/account/api.ts +++ b/web/src/api/account/api.ts @@ -154,6 +154,7 @@ export interface AccountModel { created_at: number; updated_at: number; use_dangerous: boolean; + uidonly_enabled: boolean; pgp_key?: string; imap_quota_window?: QuotaWindow; imap_quota_bytes?: number; @@ -225,4 +226,4 @@ export const autoconfig = async (email: string) => { export const access_assign = async (data: Record) => { const response = await axiosInstance.post("api/v1/accounts/access/assignments", data); return response.data; -}; \ No newline at end of file +}; diff --git a/web/src/features/accounts/account-new.tsx b/web/src/features/accounts/account-new.tsx index 70e6cdc6..b90e394f 100644 --- a/web/src/features/accounts/account-new.tsx +++ b/web/src/features/accounts/account-new.tsx @@ -51,6 +51,7 @@ const defaultValues: AccountFormValues = { }, enabled: true, use_dangerous: false, + uidonly_enabled: false, date_since: undefined, date_before: undefined, download_interval_min: 60, @@ -126,6 +127,7 @@ export function AccountNewPage() { }, enabled: data.enabled, use_dangerous: data.use_dangerous, + uidonly_enabled: data.uidonly_enabled, date_since: data.date_since, date_before: data.date_before, download_interval_min: data.download_interval_min, diff --git a/web/src/features/accounts/account-settings-page.tsx b/web/src/features/accounts/account-settings-page.tsx index 1bcb8896..556935e7 100644 --- a/web/src/features/accounts/account-settings-page.tsx +++ b/web/src/features/accounts/account-settings-page.tsx @@ -57,6 +57,7 @@ function mapAccountToFormValues(account: AccountModel): AccountFormValues { imap, enabled: account.enabled, use_dangerous: account.use_dangerous, + uidonly_enabled: account.uidonly_enabled, date_since: account.date_since ?? undefined, date_before: account.date_before ?? undefined, download_interval_min: account.download_interval_min ?? 60, @@ -150,6 +151,7 @@ export function AccountSettingsPage({ accountId }: AccountSettingsPageProps) { }, enabled: data.enabled, use_dangerous: data.use_dangerous, + uidonly_enabled: data.uidonly_enabled, date_since: data.date_since, date_before: data.date_before, download_interval_min: data.download_interval_min, diff --git a/web/src/features/accounts/components/__tests__/schema-edge.test.ts b/web/src/features/accounts/components/__tests__/schema-edge.test.ts index c4568640..d1f186e9 100644 --- a/web/src/features/accounts/components/__tests__/schema-edge.test.ts +++ b/web/src/features/accounts/components/__tests__/schema-edge.test.ts @@ -16,6 +16,7 @@ const baseData = { }, enabled: true, use_dangerous: false, + uidonly_enabled: false, download_interval_min: 60, download_batch_size: 30, auto_download_new_mailboxes: true, diff --git a/web/src/features/accounts/components/__tests__/schema.test.ts b/web/src/features/accounts/components/__tests__/schema.test.ts index 8110f976..9195a137 100644 --- a/web/src/features/accounts/components/__tests__/schema.test.ts +++ b/web/src/features/accounts/components/__tests__/schema.test.ts @@ -16,6 +16,7 @@ const validAccountData = { }, enabled: true, use_dangerous: false, + uidonly_enabled: false, download_interval_min: 60, download_batch_size: 30, auto_download_new_mailboxes: true, diff --git a/web/src/features/accounts/components/schema.ts b/web/src/features/accounts/components/schema.ts index 3baa3200..8ad7ed9f 100644 --- a/web/src/features/accounts/components/schema.ts +++ b/web/src/features/accounts/components/schema.ts @@ -90,6 +90,7 @@ export const getAccountSchema = (isEdit: boolean, t: (key: string) => string) => imap: getImapConfigSchema(isEdit, t), enabled: z.boolean(), use_dangerous: z.boolean(), + uidonly_enabled: z.boolean(), date_since: dateSelectionSchema(t).optional(), date_before: relativeDateSchema(t).optional(), download_interval_min: z diff --git a/web/src/features/accounts/components/tab-server.tsx b/web/src/features/accounts/components/tab-server.tsx index 9e9c91a0..931231bd 100644 --- a/web/src/features/accounts/components/tab-server.tsx +++ b/web/src/features/accounts/components/tab-server.tsx @@ -209,6 +209,22 @@ export function TabServer({ isEdit }: TabServerProps) { )} /> + + ( + + + + +

+ {t('accounts.uidonlyEnabled')} + {t('accounts.uidonlyEnabledDescription')} +
+ + )} + /> ); } diff --git a/web/src/locales/en.json b/web/src/locales/en.json index ffa6f04e..2b64dd25 100644 --- a/web/src/locales/en.json +++ b/web/src/locales/en.json @@ -380,6 +380,8 @@ "updateFailed": "Update failed, please try again later", "updateTheEmailAccountHere": "Update the email account here. ", "updatedAt": "Updated At", + "uidonlyEnabled": "Enable UID-only acquisition", + "uidonlyEnabledDescription": "Required for large Yahoo mailboxes. This changes remote message identity for this account; disabling it later stops synchronization instead of falling back to the limited mailbox view.", "useDangerous": "Trust Any TLS Certificate", "useDangerousDescription": "Enable this option only if you are connecting to an IMAP server with a public or self-signed certificate that may not be recognized by your system. Using this setting bypasses standard certificate validation, which can expose you to man-in-the-middle attacks. Only enable if you understand the risks.", "useNoProxy": "No Proxy", @@ -1833,4 +1835,4 @@ "singleRequestBatchSizeTooLarge": "Batch size must be at most 200", "singleRequestBatchSizeTooSmall": "Batch size must be at least 10" } -} \ No newline at end of file +} From 5e8779670b18bd6d31a805a81f76b27e6081f95b Mon Sep 17 00:00:00 2001 From: Gabe A Date: Wed, 5 Aug 2026 13:25:53 -0400 Subject: [PATCH 4/4] batch UIDONLY body fetches --- crates/core/src/imap/uidonly.rs | 43 +- crates/core/src/imap/uidonly_acquisition.rs | 670 +++++++++++++++++--- crates/core/src/imap/uidonly_tests.rs | 49 +- 3 files changed, 649 insertions(+), 113 deletions(-) diff --git a/crates/core/src/imap/uidonly.rs b/crates/core/src/imap/uidonly.rs index 610ec88c..bb67b2a5 100644 --- a/crates/core/src/imap/uidonly.rs +++ b/crates/core/src/imap/uidonly.rs @@ -799,12 +799,43 @@ pub(crate) fn inventory_args( )) } -/// Injection-safe arguments for one complete, literal-backed raw message. -pub(crate) fn exact_args(uid: u32) -> io::Result<(String, &'static str)> { - if uid == 0 { - return Err(invalid("UID 0 is invalid")); +/// Injection-safe arguments for one bounded batch of complete raw messages. +/// +/// The caller supplies an already bounded, strictly increasing UID list. A +/// comma-only sequence set keeps exact-body commands distinguishable from the +/// range-plus-PARTIAL inventory command admitted below. +pub(crate) fn exact_args(uids: &[u32]) -> io::Result<(String, &'static str)> { + if uids.is_empty() + || uids.iter().any(|uid| *uid == 0) + || uids.windows(2).any(|window| window[0] >= window[1]) + { + return Err(invalid( + "exact UIDONLY UIDs must be nonzero, unique, and increasing", + )); + } + Ok(( + uids.iter() + .map(u32::to_string) + .collect::>() + .join(","), + "(UID RFC822.SIZE BODY.PEEK[])", + )) +} + +fn exact_uid_set(set: &[u8]) -> bool { + let mut previous = 0_u32; + let mut count = 0_usize; + for value in set.split(|byte| *byte == b',') { + let Some(uid) = parse_nonzero_u32(value) else { + return false; + }; + if uid <= previous { + return false; + } + previous = uid; + count += 1; } - Ok((uid.to_string(), "(UID RFC822.SIZE BODY.PEEK[])")) + count > 0 } fn uid_fetch_command_kind(command: &[u8]) -> Option { @@ -831,7 +862,7 @@ fn uid_fetch_command_kind(command: &[u8]) -> Option { .and_then(parse_nonzero_u32) .map(|_| CommandKind::Inventory) } else { - (parse_nonzero_u32(set).is_some() && query == b"(UID RFC822.SIZE BODY.PEEK[])") + (exact_uid_set(set) && query == b"(UID RFC822.SIZE BODY.PEEK[])") .then_some(CommandKind::ExactFetch) } } diff --git a/crates/core/src/imap/uidonly_acquisition.rs b/crates/core/src/imap/uidonly_acquisition.rs index a376c66e..4cd828e4 100644 --- a/crates/core/src/imap/uidonly_acquisition.rs +++ b/crates/core/src/imap/uidonly_acquisition.rs @@ -32,19 +32,22 @@ use crate::envelope::extractor::{ }; use crate::error::code::ErrorCode; use crate::error::BichonResult; -use crate::imap::executor::DEFAULT_MAX_EMAIL_SIZE; +use crate::imap::executor::{DEFAULT_BATCH_SIZE, DEFAULT_MAX_EMAIL_SIZE}; use crate::imap::manager::ImapConnectionManager; use crate::imap::session::SessionStream; use crate::imap::uidonly::{exact_args, inventory_args, UidOnlyHandle, UidOnlyLimits}; use crate::raise_error; use async_imap::Session; use futures::{FutureExt, TryStreamExt}; +use std::collections::{BTreeMap, VecDeque}; use std::future::Future; use std::panic::AssertUnwindSafe; use std::time::Duration; use tokio_util::sync::CancellationToken; const UIDONLY_DEFAULT_PAGE_SIZE: u32 = 1_000; +const UIDONLY_MAX_FETCH_BATCH_MESSAGES: u32 = 32; +const UIDONLY_FETCH_BATCH_BYTES: u64 = UIDONLY_PROJECTION_BATCH_BYTES as u64; const MAX_UIDONLY_RECONNECTS: u32 = 3; #[derive(Clone, Copy, Debug, Eq, PartialEq)] @@ -52,6 +55,8 @@ pub(crate) struct AcquisitionLimits { pub max_literal_bytes: u64, pub max_operation_runtime: Duration, pub page_size: u32, + pub fetch_batch_size: u32, + pub fetch_batch_bytes: u64, } impl AcquisitionLimits { @@ -62,6 +67,8 @@ impl AcquisitionLimits { // archive: a valid large mailbox may need to run for days. max_operation_runtime: Duration::from_secs(10 * 60), page_size: 1_000, + fetch_batch_size: DEFAULT_BATCH_SIZE.min(UIDONLY_MAX_FETCH_BATCH_MESSAGES), + fetch_batch_bytes: UIDONLY_FETCH_BATCH_BYTES, } } @@ -69,6 +76,10 @@ impl AcquisitionLimits { if self.max_literal_bytes == 0 || self.max_operation_runtime.is_zero() || self.page_size == 0 + || self.fetch_batch_size == 0 + || self.fetch_batch_size > UIDONLY_MAX_FETCH_BATCH_MESSAGES + || self.fetch_batch_bytes == 0 + || self.fetch_batch_bytes > UIDONLY_FETCH_BATCH_BYTES { return Err(raise_error!( "UIDONLY acquisition limits must be nonzero".into(), @@ -89,6 +100,7 @@ pub(crate) struct MailboxSnapshot { #[derive(Clone, Copy, Debug, Eq, PartialEq)] pub(crate) struct InventoryItem { pub uid: u32, + pub rfc822_size: u64, } #[derive(Clone, Copy, Debug, Eq, PartialEq)] @@ -119,9 +131,14 @@ pub(crate) trait UidOnlyTransport { page_size: u32, ) -> BichonResult>; - /// Fetch exactly one full message. `literal_budget` is a pre-read ceiling, - /// not a post-read accounting hint. - async fn fetch_exact(&mut self, uid: u32, literal_budget: u64) -> BichonResult; + /// Fetch one exact, bounded UID set. The count and aggregate byte budget + /// are pre-read ceilings, not post-read accounting hints. + async fn fetch_many( + &mut self, + items: &[InventoryItem], + literal_budget: u64, + command_budget: u64, + ) -> BichonResult>; async fn reconnect(&mut self, _page_size: u32) -> BichonResult<()> { Err(raise_error!( @@ -260,18 +277,20 @@ async fn inventory_with_reconnect( } } -async fn fetch_with_reconnect( +async fn fetch_batch_with_reconnect( transport: &mut T, mailbox: &str, initial: MailboxSnapshot, - uid: u32, + items: Vec, limits: AcquisitionLimits, token: &CancellationToken, -) -> BichonResult { +) -> BichonResult> { let mut reconnects = 0; - loop { + let mut pending = VecDeque::from([items]); + let mut messages = Vec::new(); + while let Some(batch) = pending.pop_front() { match bounded( - transport.fetch_exact(uid, limits.max_literal_bytes), + transport.fetch_many(&batch, limits.max_literal_bytes, limits.fetch_batch_bytes), limits.max_operation_runtime, token, ) @@ -280,10 +299,53 @@ async fn fetch_with_reconnect( Err(error) if retryable_transport_error(error.code()) => { recover_transport(transport, mailbox, initial, limits, token, &mut reconnects) .await?; + pending.push_front(batch); } - result => return result, + Err(error) if error.code() == ErrorCode::PayloadTooLarge && batch.len() > 1 => { + recover_transport(transport, mailbox, initial, limits, token, &mut reconnects) + .await?; + let midpoint = batch.len() / 2; + pending.push_front(batch[midpoint..].to_vec()); + pending.push_front(batch[..midpoint].to_vec()); + } + Err(error) => return Err(error), + Ok(batch_messages) => messages.extend(batch_messages), } } + Ok(messages) +} + +fn planned_fetch_batches( + items: Vec, + count_limit: u32, + byte_limit: u64, +) -> Vec> { + let mut batches = Vec::new(); + let mut batch = Vec::new(); + let mut bytes = 0_u64; + for item in items { + let declared = item.rfc822_size.max(1); + if !batch.is_empty() + && (batch.len() >= count_limit as usize || bytes.saturating_add(declared) > byte_limit) + { + batches.push(std::mem::take(&mut batch)); + bytes = 0; + } + bytes = bytes.saturating_add(declared); + batch.push(item); + } + if !batch.is_empty() { + batches.push(batch); + } + batches +} + +fn effective_fetch_batch_size(requested: Option, message_limit: u32) -> u32 { + requested + .unwrap_or(DEFAULT_BATCH_SIZE) + .min(message_limit) + .min(UIDONLY_MAX_FETCH_BATCH_MESSAGES) + .max(1) } async fn snapshot_with_reconnect( @@ -439,39 +501,63 @@ where )); } + let mut unresolved = Vec::new(); for (item, is_verified) in page.into_iter().zip(verified) { if is_verified { archived += 1; - continue; + } else { + unresolved.push(item); } + } - let message = - fetch_with_reconnect(transport, mailbox, snapshot, item.uid, limits, &token) + for batch in planned_fetch_batches( + unresolved, + limits.fetch_batch_size, + limits.fetch_batch_bytes, + ) { + let requested = batch.clone(); + let messages = + fetch_batch_with_reconnect(transport, mailbox, snapshot, batch, limits, &token) .await?; - if message.uid != item.uid { + if messages.len() != requested.len() { return Err(raise_error!( - "UIDONLY exact fetch returned the wrong UID".into(), + "UIDONLY body batch returned the wrong result count".into(), ErrorCode::ImapUnexpectedResult )); } - let actual = message.body.len() as u64; - if actual > limits.max_literal_bytes { - return Err(raise_error!( - "UIDONLY exact body exceeded its pre-read budget".into(), - ErrorCode::PayloadTooLarge - )); - } - if !pending.is_empty() - && (pending.len() >= UIDONLY_PROJECTION_BATCH_MESSAGES - || pending_bytes.saturating_add(actual) > UIDONLY_PROJECTION_BATCH_BYTES as u64) - { - let stored = flush_projection_batch(archive, &mut pending, limits, &token).await?; - archived += stored; - downloaded += stored; - pending_bytes = 0; + for (item, message) in requested.into_iter().zip(messages) { + if message.uid != item.uid { + return Err(raise_error!( + "UIDONLY exact fetch returned the wrong UID".into(), + ErrorCode::ImapUnexpectedResult + )); + } + let actual = message.body.len() as u64; + if actual > limits.max_literal_bytes { + return Err(raise_error!( + "UIDONLY exact body exceeded its pre-read budget".into(), + ErrorCode::PayloadTooLarge + )); + } + if !pending.is_empty() + && (pending.len() >= UIDONLY_PROJECTION_BATCH_MESSAGES + || pending_bytes.saturating_add(actual) + > UIDONLY_PROJECTION_BATCH_BYTES as u64) + { + let stored = + flush_projection_batch(archive, &mut pending, limits, &token).await?; + archived += stored; + downloaded += stored; + pending_bytes = 0; + progress(AcquisitionProgress { + planned, + resolved: archived, + downloaded, + })?; + } + pending_bytes += actual; + pending.push(message); } - pending_bytes += actual; - pending.push(message); } let stored = flush_projection_batch(archive, &mut pending, limits, &token).await?; @@ -482,7 +568,7 @@ where cursor = previous.checked_add(1).unwrap_or(u32::MAX); progress(AcquisitionProgress { planned, - resolved: inventoried, + resolved: archived, downloaded, })?; if previous == u32::MAX { @@ -688,7 +774,8 @@ fn protocol_limits(max_literal_bytes: u64) -> BichonResult { ErrorCode::InvalidParameter ) })?; - let response = literal.checked_add(128 * 1024).ok_or_else(|| { + let command_literal = literal.max(UIDONLY_FETCH_BATCH_BYTES as usize); + let response = command_literal.checked_add(128 * 1024).ok_or_else(|| { raise_error!( "UIDONLY response-size limit overflow".into(), ErrorCode::InvalidParameter @@ -704,7 +791,7 @@ fn protocol_limits(max_literal_bytes: u64) -> BichonResult { max_control_line_bytes: 64 * 1024, max_literal_bytes: literal, max_response_bytes: response, - max_command_literal_bytes: literal, + max_command_literal_bytes: command_literal, max_command_response_bytes: command, // A 1,000-message inventory page plus bounded unsolicited mailbox // updates must still fit without making the response count unbounded. @@ -875,21 +962,29 @@ impl UidOnlyTransport for SessionUidOnlyTransport { ErrorCode::ImapUnexpectedResult )); } - fetch.size.ok_or_else(|| { + let rfc822_size = fetch.size.ok_or_else(|| { raise_error!( "UIDONLY inventory omitted RFC822.SIZE".into(), ErrorCode::ImapUnexpectedResult ) })?; - Ok(InventoryItem { uid }) + Ok(InventoryItem { + uid, + rfc822_size: u64::from(rfc822_size), + }) }) .collect() } - async fn fetch_exact(&mut self, uid: u32, literal_budget: u64) -> BichonResult { - let limit = usize::try_from(literal_budget).map_err(|_| { + async fn fetch_many( + &mut self, + items: &[InventoryItem], + literal_budget: u64, + command_budget: u64, + ) -> BichonResult> { + let limit = usize::try_from(command_budget).map_err(|_| { raise_error!( - "UIDONLY literal budget is not representable".into(), + "UIDONLY command budget is not representable".into(), ErrorCode::InvalidParameter ) })?; @@ -902,43 +997,99 @@ impl UidOnlyTransport for SessionUidOnlyTransport { ) })?; let before = self.handle.literal_bytes_received(); - let (set, query) = exact_args(uid).map_err(|_| { + let uids = items.iter().map(|item| item.uid).collect::>(); + let (set, query) = exact_args(&uids).map_err(|_| { raise_error!( - "Invalid UIDONLY exact UID".into(), + "Invalid UIDONLY exact UID set".into(), ErrorCode::InvalidParameter ) })?; - let mut fetched = self.collect_fetches(set, query).await?; - if fetched.len() != 1 { + let fetched = self.collect_fetches(set, query).await?; + if fetched.len() != items.len() { return Err(raise_error!( - "UIDONLY exact fetch did not return exactly one message".into(), + "UIDONLY exact fetch returned the wrong result count".into(), ErrorCode::ImapUnexpectedResult )); } - let fetch = fetched.pop().expect("length checked"); - if fetch.message != uid || fetch.uid != Some(uid) || fetch.size.is_none() { - return Err(raise_error!( - "UIDONLY exact fetch returned mismatched metadata".into(), - ErrorCode::ImapUnexpectedResult - )); + + let requested = items + .iter() + .map(|item| (item.uid, *item)) + .collect::>(); + let mut messages = BTreeMap::new(); + let mut actual_total = 0_u64; + for fetch in fetched { + let uid = fetch.uid.ok_or_else(|| { + raise_error!( + "UIDONLY exact fetch omitted UID".into(), + ErrorCode::ImapUnexpectedResult + ) + })?; + if !requested.contains_key(&uid) || fetch.message != uid || fetch.size.is_none() { + return Err(raise_error!( + "UIDONLY exact fetch returned mismatched metadata".into(), + ErrorCode::ImapUnexpectedResult + )); + } + let raw = fetch.body().ok_or_else(|| { + raise_error!( + "UIDONLY exact fetch omitted its full literal body".into(), + ErrorCode::ImapUnexpectedResult + ) + })?; + let actual = raw.len() as u64; + if actual > literal_budget { + return Err(raise_error!( + "UIDONLY exact body exceeded its per-literal budget".into(), + ErrorCode::PayloadTooLarge + )); + } + actual_total = actual_total.checked_add(actual).ok_or_else(|| { + raise_error!( + "UIDONLY exact batch byte count overflow".into(), + ErrorCode::PayloadTooLarge + ) + })?; + if actual_total > command_budget { + return Err(raise_error!( + "UIDONLY exact batch exceeded its command budget".into(), + ErrorCode::PayloadTooLarge + )); + } + if messages + .insert( + uid, + UidOnlyMessage { + uid, + body: raw.to_vec(), + }, + ) + .is_some() + { + return Err(raise_error!( + "UIDONLY exact fetch returned a duplicate UID".into(), + ErrorCode::ImapUnexpectedResult + )); + } } - let raw = fetch.body().ok_or_else(|| { - raise_error!( - "UIDONLY exact fetch omitted its full literal body".into(), - ErrorCode::ImapUnexpectedResult - ) - })?; let after = self.handle.literal_bytes_received(); - if after.checked_sub(before) != Some(raw.len() as u64) { + if after.checked_sub(before) != Some(actual_total) { return Err(raise_error!( - "UIDONLY exact body did not match literal accounting".into(), + "UIDONLY exact batch did not match literal accounting".into(), ErrorCode::ImapUnexpectedResult )); } - Ok(UidOnlyMessage { - uid, - body: raw.to_vec(), - }) + items + .iter() + .map(|item| { + messages.remove(&item.uid).ok_or_else(|| { + raise_error!( + "UIDONLY exact batch omitted a requested UID".into(), + ErrorCode::ImapUnexpectedResult + ) + }) + }) + .collect() } async fn reconnect(&mut self, page_size: u32) -> BichonResult<()> { @@ -1102,6 +1253,8 @@ where }; let mut limits = AcquisitionLimits::bounded(max_literal); limits.page_size = UIDONLY_DEFAULT_PAGE_SIZE.min(message_limit).max(1); + limits.fetch_batch_size = + effective_fetch_batch_size(account.download_batch_size, message_limit); let result = run_acquisition( &mut transport, &mut archive, @@ -1147,11 +1300,15 @@ mod tests { use std::collections::{BTreeMap, BTreeSet, VecDeque}; fn limits() -> AcquisitionLimits { - AcquisitionLimits { - max_literal_bytes: 64 * 1024, - max_operation_runtime: Duration::from_secs(20), - page_size: 2, - } + let mut limits = AcquisitionLimits::bounded(64 * 1024); + limits.max_operation_runtime = Duration::from_secs(20); + limits.page_size = 2; + limits.fetch_batch_size = 1; + limits + } + + fn inventory_item(uid: u32, rfc822_size: u64) -> InventoryItem { + InventoryItem { uid, rfc822_size } } #[test] @@ -1213,6 +1370,28 @@ mod tests { } } + #[test] + fn fetch_batch_limits_honor_account_server_and_byte_bounds() { + assert_eq!(effective_fetch_batch_size(None, 1_000), 30); + assert_eq!(effective_fetch_batch_size(Some(8), 1_000), 8); + assert_eq!(effective_fetch_batch_size(Some(100), 1_000), 32); + assert_eq!(effective_fetch_batch_size(Some(30), 4), 4); + + let items = vec![ + inventory_item(1, 40), + inventory_item(2, 40), + inventory_item(3, 10), + ]; + let batches = planned_fetch_batches(items, 2, 64); + assert_eq!( + batches + .iter() + .map(|batch| batch.iter().map(|item| item.uid).collect::>()) + .collect::>(), + [vec![1], vec![2, 3]] + ); + } + #[test] fn source_identity_is_stable_across_connection_policy_changes() { let account = AccountModel { @@ -1264,6 +1443,7 @@ mod tests { pages: VecDeque>, messages: BTreeMap>, fetched: Vec, + fetch_batches: Vec>, expected_cursor: Option, } @@ -1280,12 +1460,13 @@ mod tests { uid_next: 51, }, pages: VecDeque::from([ - vec![InventoryItem { uid: 2 }, InventoryItem { uid: 30 }], - vec![InventoryItem { uid: 50 }], + vec![inventory_item(2, 64), inventory_item(30, 64)], + vec![inventory_item(50, 64)], vec![], ]), messages: BTreeMap::from([(2, body(2)), (30, body(30)), (50, body(50))]), fetched: Vec::new(), + fetch_batches: Vec::new(), expected_cursor: None, } } @@ -1308,25 +1489,37 @@ mod tests { Ok(self.pages.pop_front().unwrap_or_default()) } - async fn fetch_exact( + async fn fetch_many( &mut self, - uid: u32, + items: &[InventoryItem], literal_budget: u64, - ) -> BichonResult { - self.fetched.push(uid); - let raw = self.messages.get(&uid).cloned().ok_or_else(|| { - raise_error!( - "synthetic missing body".into(), - ErrorCode::ImapUnexpectedResult - ) - })?; - if raw.len() as u64 > literal_budget { - return Err(raise_error!( - "synthetic literal budget".into(), - ErrorCode::PayloadTooLarge - )); + command_budget: u64, + ) -> BichonResult> { + self.fetch_batches + .push(items.iter().map(|item| item.uid).collect()); + let mut total = 0_u64; + let mut messages = Vec::new(); + for item in items { + self.fetched.push(item.uid); + let raw = self.messages.get(&item.uid).cloned().ok_or_else(|| { + raise_error!( + "synthetic missing body".into(), + ErrorCode::ImapUnexpectedResult + ) + })?; + total += raw.len() as u64; + if raw.len() as u64 > literal_budget || total > command_budget { + return Err(raise_error!( + "synthetic literal budget".into(), + ErrorCode::PayloadTooLarge + )); + } + messages.push(UidOnlyMessage { + uid: item.uid, + body: raw, + }); } - Ok(UidOnlyMessage { uid, body: raw }) + Ok(messages) } } @@ -1334,10 +1527,12 @@ mod tests { inner: FakeTransport, inventory_failures: u32, fetch_failures: u32, + aggregate_overflow_once: bool, reconnects: u32, inventory_attempts: u32, fetch_attempts: Vec, snapshot_after_reconnect: Option, + message_limit_after_reconnect: Option, } impl FlakyTransport { @@ -1346,10 +1541,12 @@ mod tests { inner: FakeTransport::sparse(), inventory_failures: 0, fetch_failures: 0, + aggregate_overflow_once: false, reconnects: 0, inventory_attempts: 0, fetch_attempts: Vec::new(), snapshot_after_reconnect: None, + message_limit_after_reconnect: None, } } } @@ -1376,12 +1573,21 @@ mod tests { self.inner.inventory_page(cursor, high, page_size).await } - async fn fetch_exact( + async fn fetch_many( &mut self, - uid: u32, + items: &[InventoryItem], literal_budget: u64, - ) -> BichonResult { - self.fetch_attempts.push(uid); + command_budget: u64, + ) -> BichonResult> { + self.fetch_attempts + .extend(items.iter().map(|item| item.uid)); + if self.aggregate_overflow_once && items.len() > 1 { + self.aggregate_overflow_once = false; + return Err(raise_error!( + "synthetic aggregate overflow".into(), + ErrorCode::PayloadTooLarge + )); + } if self.fetch_failures > 0 { self.fetch_failures -= 1; return Err(raise_error!( @@ -1389,11 +1595,23 @@ mod tests { ErrorCode::NetworkError )); } - self.inner.fetch_exact(uid, literal_budget).await + self.inner + .fetch_many(items, literal_budget, command_budget) + .await } - async fn reconnect(&mut self, _page_size: u32) -> BichonResult<()> { + async fn reconnect(&mut self, page_size: u32) -> BichonResult<()> { self.reconnects += 1; + if self + .message_limit_after_reconnect + .take() + .is_some_and(|message_limit| message_limit < page_size) + { + return Err(raise_error!( + "synthetic reconnect reduced MESSAGELIMIT".into(), + ErrorCode::Incompatible + )); + } if let Some(snapshot) = self.snapshot_after_reconnect.take() { self.inner.snapshot = snapshot; } @@ -1459,6 +1677,49 @@ mod tests { assert_eq!(transport.fetched, [2, 30, 50]); } + #[tokio::test] + async fn body_batches_are_exact_and_progress_follows_projection_commits() { + let count = 40_u32; + let mut transport = FakeTransport { + snapshot: MailboxSnapshot { + exists: count, + uid_validity: 77, + uid_next: count + 1, + }, + pages: VecDeque::from([(1..=count) + .map(|uid| inventory_item(uid, 1)) + .collect::>()]), + messages: (1..=count).map(|uid| (uid, vec![b'x'])).collect(), + fetched: Vec::new(), + fetch_batches: Vec::new(), + expected_cursor: None, + }; + let mut bounded = limits(); + bounded.page_size = count; + bounded.fetch_batch_size = 30; + let mut progress = Vec::new(); + let report = run_acquisition( + &mut transport, + &mut FakeArchive::default(), + "synthetic", + Some(77), + bounded, + CancellationToken::new(), + |state| { + progress.push(state); + Ok(()) + }, + ) + .await + .unwrap(); + + assert_eq!(report.archived, u64::from(count)); + assert_eq!(transport.fetch_batches[0], (1..=30).collect::>()); + assert_eq!(transport.fetch_batches[1], (31..=40).collect::>()); + assert!(progress.iter().any(|state| state.downloaded == 32)); + assert_eq!(progress.last().unwrap().downloaded, u64::from(count)); + } + #[tokio::test] async fn reconnect_retries_the_same_inventory_cursor() { let mut transport = FlakyTransport::sparse(); @@ -1522,12 +1783,51 @@ mod tests { assert_eq!(transport.inventory_attempts, 4); } + #[tokio::test] + async fn aggregate_overflow_reconnects_and_splits_without_skipping() { + let mut transport = FlakyTransport::sparse(); + transport.aggregate_overflow_once = true; + let mut archive = FakeArchive::default(); + let mut bounded = limits(); + bounded.fetch_batch_size = 2; + let report = run_acquisition( + &mut transport, + &mut archive, + "synthetic", + Some(77), + bounded, + CancellationToken::new(), + |_| Ok(()), + ) + .await + .unwrap(); + + assert_eq!(report.archived, 3); + assert_eq!(transport.reconnects, 1); + assert_eq!(archive.projected, [2, 30, 50]); + } + + #[tokio::test] + async fn reconnect_rejects_reduced_messagelimit_without_advancing() { + let mut transport = FlakyTransport::sparse(); + transport.inventory_failures = 1; + transport.message_limit_after_reconnect = Some(1); + let mut archive = FakeArchive::default(); + + let error = run_test(&mut transport, &mut archive).await.unwrap_err(); + + assert_eq!(error.code(), ErrorCode::Incompatible); + assert_eq!(transport.reconnects, 1); + assert!(transport.inner.fetched.is_empty()); + assert!(archive.projected.is_empty()); + } + #[tokio::test] async fn proven_checkpoint_scans_only_new_uids() { let mut transport = FakeTransport::sparse(); transport.snapshot.exists = 4; transport.snapshot.uid_next = 60; - transport.pages = VecDeque::from([vec![InventoryItem { uid: 55 }]]); + transport.pages = VecDeque::from([vec![inventory_item(55, 64)]]); transport .messages .insert(55, b"Subject: new\r\n\r\nbody".to_vec()); @@ -1642,14 +1942,15 @@ mod tests { assert_eq!(cursor, self.cursor); let end = high.min(cursor + page_size - 1); self.cursor = end + 1; - Ok((cursor..=end).map(|uid| InventoryItem { uid }).collect()) + Ok((cursor..=end).map(|uid| inventory_item(uid, 64)).collect()) } - async fn fetch_exact( + async fn fetch_many( &mut self, - _uid: u32, + _items: &[InventoryItem], _literal_budget: u64, - ) -> BichonResult { + _command_budget: u64, + ) -> BichonResult> { unreachable!("every logical receipt verifies") } } @@ -1711,9 +2012,29 @@ mod tests { response } - async fn connect_fake_transport(server: &MockImapServerHandle) -> SessionUidOnlyTransport { + fn exact_batch_response(entries: &[(u32, u32, &[u8])]) -> Vec { + let mut response = Vec::new(); + for (uid, reported_size, raw) in entries { + response.extend_from_slice( + format!( + "* {uid} UIDFETCH (UID {uid} RFC822.SIZE {reported_size} BODY[] {{{}}}\r\n", + raw.len() + ) + .as_bytes(), + ); + response.extend_from_slice(raw); + response.extend_from_slice(b")\r\n"); + } + response.extend_from_slice(b"{TAG} OK exact fetch completed\r\n"); + response + } + + async fn connect_fake_transport_with_message_limit( + server: &MockImapServerHandle, + expected_message_limit: u32, + ) -> SessionUidOnlyTransport { let account = AccountModel { - email: "synthetic-user".into(), + email: "synthetic-user@example.invalid".into(), imap: Some(ImapConfig { host: server.host(), port: server.port(), @@ -1740,13 +2061,13 @@ mod tests { .with_uidonly(limits.clone()) .expect("install UIDONLY guard"); let mut session = client - .login("synthetic-user", "synthetic-secret") + .login("synthetic-user@example.invalid", "synthetic-secret") .await .expect("synthetic login"); let capabilities = session.capabilities().await.expect("capabilities"); assert!(capabilities.has_str("UIDONLY")); assert!(capabilities.has_str("PARTIAL")); - assert!(capabilities.has_str("MESSAGELIMIT=2")); + assert!(capabilities.has_str(&format!("MESSAGELIMIT={expected_message_limit}"))); session .run_command_and_check_ok("ENABLE UIDONLY") .await @@ -1761,6 +2082,10 @@ mod tests { } } + async fn connect_fake_transport(server: &MockImapServerHandle) -> SessionUidOnlyTransport { + connect_fake_transport_with_message_limit(server, 2).await + } + #[tokio::test] async fn tcp_fake_yahoo_uidonly_pages_sparse_uids_and_checkpoints_after_verification() { let messages = [ @@ -1784,19 +2109,26 @@ mod tests { .respond("UID FETCH 8:9", inventory_response(&[(9, 0)])) // The reported sizes are deliberately advisory and wrong. The // literal byte count is the acquisition/accounting authority. - .respond("UID FETCH 2 (", exact_response(2, 1, messages[0].1)) - .respond("UID FETCH 7 (", exact_response(7, 2, messages[1].1)) + .respond( + "UID FETCH 2,7 (", + exact_batch_response(&[ + (7, 2, messages[1].1), + (2, 1, messages[0].1), + ]), + ) .respond("UID FETCH 9 (", exact_response(9, 3, messages[2].1)) .start() .await; let mut transport = connect_fake_transport(&server).await; let mut archive = FakeArchive::default(); + let mut bounded = limits(); + bounded.fetch_batch_size = 2; let report = run_acquisition( &mut transport, &mut archive, "Synthetic", Some(77), - limits(), + bounded, CancellationToken::new(), |_| Ok(()), ) @@ -1812,7 +2144,99 @@ mod tests { assert_eq!(commands.matches(" EXAMINE ").count(), 2); assert!(commands.contains(" UID FETCH 1:9 (UID RFC822.SIZE) (PARTIAL 1:2)")); assert!(commands.contains(" UID FETCH 8:9 (UID RFC822.SIZE) (PARTIAL 1:2)")); - assert_eq!(commands.matches("BODY.PEEK[]").count(), 3); + assert_eq!(commands.matches("BODY.PEEK[]").count(), 2); + let commands = commands.to_ascii_uppercase(); + assert!( + ![" STORE ", " MOVE ", " COPY ", " DELETE ", " EXPUNGE", " CLOSE",] + .iter() + .any(|forbidden| commands.contains(forbidden)) + ); + } + + #[tokio::test] + async fn tcp_fake_yahoo_inventory_crosses_10000_before_one_30_body_command() { + const MESSAGE_COUNT: u32 = 10_030; + const PAGE_SIZE: u32 = 1_000; + const FIRST_BODY_UID: u32 = 10_001; + + let mut server = MockImapServer::new() + .greeting("* OK synthetic Yahoo-like IMAP ready\r\n") + .respond("LOGIN", "{TAG} OK LOGIN completed\r\n") + .respond( + "CAPABILITY", + "* CAPABILITY IMAP4rev1 ENABLE UIDONLY PARTIAL MESSAGELIMIT=30\r\n{TAG} OK CAPABILITY completed\r\n", + ) + .respond( + "ENABLE UIDONLY", + "* ENABLED UIDONLY\r\n{TAG} OK ENABLE completed\r\n", + ) + .respond( + "EXAMINE", + examine_response("Synthetic", MESSAGE_COUNT, 77, MESSAGE_COUNT + 1), + ); + + let mut cursor = 1_u32; + while cursor <= MESSAGE_COUNT { + let end = MESSAGE_COUNT.min(cursor + PAGE_SIZE - 1); + let entries = (cursor..=end).map(|uid| (uid, 1_u32)).collect::>(); + server = server.respond( + format!("UID FETCH {cursor}:{MESSAGE_COUNT} ("), + inventory_response(&entries), + ); + cursor = end + 1; + } + + let requested_uids = (FIRST_BODY_UID..=MESSAGE_COUNT).collect::>(); + let body_command = format!( + "UID FETCH {} (", + requested_uids + .iter() + .map(u32::to_string) + .collect::>() + .join(",") + ); + let bodies = requested_uids + .iter() + .rev() + .map(|uid| (*uid, 1_u32, b"x".as_slice())) + .collect::>(); + let server = server + .respond(&body_command, exact_batch_response(&bodies)) + .start() + .await; + + let mut transport = connect_fake_transport_with_message_limit(&server, 30).await; + let mut archive = FakeArchive { + verified: (1..FIRST_BODY_UID).collect(), + ..Default::default() + }; + let mut bounded = limits(); + bounded.page_size = PAGE_SIZE; + bounded.fetch_batch_size = 30; + let report = run_acquisition( + &mut transport, + &mut archive, + "Synthetic", + Some(77), + bounded, + CancellationToken::new(), + |_| Ok(()), + ) + .await + .expect("inventory and body retrieval must cross the 10,000-message boundary"); + + assert_eq!(report.inventoried, u64::from(MESSAGE_COUNT)); + assert_eq!(report.archived, u64::from(MESSAGE_COUNT)); + assert_eq!(report.checkpoint, Some(MESSAGE_COUNT)); + assert_eq!(archive.projected, requested_uids); + + let commands = server.commands().join("\n"); + assert_eq!( + commands.matches("RFC822.SIZE) (PARTIAL 1:1000)").count(), + 11 + ); + assert_eq!(commands.matches("BODY.PEEK[]").count(), 1); + assert!(commands.contains(&body_command)); let commands = commands.to_ascii_uppercase(); assert!( ![" STORE ", " MOVE ", " COPY ", " DELETE ", " EXPUNGE", " CLOSE",] @@ -1820,4 +2244,40 @@ mod tests { .any(|forbidden| commands.contains(forbidden)) ); } + + #[tokio::test] + async fn session_body_batch_rejects_missing_duplicate_and_extraneous_uids() { + let one = b"x".as_slice(); + for response in [ + exact_batch_response(&[(7, 1, one)]), + exact_batch_response(&[(7, 1, one), (7, 1, one)]), + exact_batch_response(&[(7, 1, one), (99, 1, one)]), + ] { + let server = MockImapServer::new() + .greeting("* OK synthetic IMAP ready\r\n") + .respond("LOGIN", "{TAG} OK LOGIN completed\r\n") + .respond( + "CAPABILITY", + "* CAPABILITY IMAP4rev1 ENABLE UIDONLY PARTIAL MESSAGELIMIT=2\r\n{TAG} OK CAPABILITY completed\r\n", + ) + .respond( + "ENABLE UIDONLY", + "* ENABLED UIDONLY\r\n{TAG} OK ENABLE completed\r\n", + ) + .respond("UID FETCH 7,42 (", response) + .start() + .await; + let mut transport = connect_fake_transport(&server).await; + let error = transport + .fetch_many( + &[inventory_item(7, 1), inventory_item(42, 1)], + 64 * 1024, + 64 * 1024, + ) + .await + .err() + .expect("malformed batch must fail"); + assert_eq!(error.code(), ErrorCode::ImapUnexpectedResult); + } + } } diff --git a/crates/core/src/imap/uidonly_tests.rs b/crates/core/src/imap/uidonly_tests.rs index 37e4248b..4afb066c 100644 --- a/crates/core/src/imap/uidonly_tests.rs +++ b/crates/core/src/imap/uidonly_tests.rs @@ -172,7 +172,7 @@ async fn public_async_imap_session_parses_fragmented_uidfetch_without_touching_l assert_eq!(mailbox.uid_validity, Some(7)); assert_eq!(mailbox.uid_next, Some(50)); - let (set, query) = exact_args(42).expect("safe UID"); + let (set, query) = exact_args(&[42]).expect("safe UID"); handle .arm_next_fetch_literal_limit(body.len()) .expect("arm body bound"); @@ -236,6 +236,23 @@ async fn body_must_be_full_and_literal_backed() { } } +#[tokio::test] +async fn disconnect_mid_literal_fails_closed_without_accepting_partial_body() { + let response = b"* 7 UIDFETCH (UID 7 RFC822.SIZE 4 BODY[] {4}\r\nab"; + let (mut session, handle) = enabled(response, UidOnlyLimits::default(), 1).await; + handle + .arm_next_fetch_literal_limit(4) + .expect("arm body bound"); + + let stream = session + .uid_fetch("7", "(UID RFC822.SIZE BODY.PEEK[])") + .await + .expect("command"); + + assert!(stream.try_collect::>().await.is_err()); + assert!(handle.poison_reason().is_some()); +} + #[tokio::test] async fn uidfetch_shape_accepts_requested_fields_in_either_order() { let response = b"* 42 UIDFETCH (RFC822.SIZE 1 UID 42)\r\nA0003 OK inventory\r\n\ @@ -270,6 +287,30 @@ async fn uidfetch_shape_accepts_requested_fields_in_either_order() { assert_eq!(exact[0].body(), Some(b"x".as_slice())); } +#[tokio::test] +async fn exact_multi_uid_command_accepts_out_of_order_responses_and_accounts_all_literals() { + let response = b"* 42 UIDFETCH (UID 42 RFC822.SIZE 1 BODY[] {1}\r\ny)\r\n\ + * 7 UIDFETCH (UID 7 RFC822.SIZE 1 BODY[] {1}\r\nx)\r\n\ + A0003 OK exact\r\n"; + let (mut session, handle) = enabled(response, UidOnlyLimits::default(), 1).await; + let (set, query) = exact_args(&[7, 42]).expect("safe exact set"); + handle + .arm_next_fetch_literal_limit(2) + .expect("arm aggregate body bound"); + let fetched: Vec<_> = session + .uid_fetch(set, query) + .await + .expect("exact batch command") + .try_collect() + .await + .expect("exact batch response"); + + assert_eq!(fetched.len(), 2); + assert_eq!(fetched[0].uid, Some(42)); + assert_eq!(fetched[1].uid, Some(7)); + assert_eq!(handle.literal_bytes_received(), 2); +} + #[tokio::test] async fn uidfetch_shape_rejects_duplicates_mismatch_missing_and_extras() { let cases: &[(&[u8], bool)] = &[ @@ -513,5 +554,9 @@ fn numeric_query_builders_reject_invalid_bounds() { assert!(inventory_args(0, 20, 100).is_err()); assert!(inventory_args(20, 10, 100).is_err()); assert!(inventory_args(10, 20, 0).is_err()); - assert!(exact_args(0).is_err()); + assert_eq!(exact_args(&[2, 7, 42]).expect("safe UID batch").0, "2,7,42"); + assert!(exact_args(&[]).is_err()); + assert!(exact_args(&[0]).is_err()); + assert!(exact_args(&[2, 2]).is_err()); + assert!(exact_args(&[7, 2]).is_err()); }