io_imap/rfc3501/
fetch_stream_batch.rs

1//! IMAP batched FETCH body-stream coroutine: fetches the bodies of a whole
2//! sequence set in **one** command and streams each straight to a per-message
3//! sink, so N bodies cost one round trip instead of N.
4//!
5//! Sends `UID FETCH <set> (UID BODY.PEEK[])` (peek so syncing does not set
6//! `\Seen`; `UID` so every returned body is self-identifying). The response is a
7//! run of `* <seq> FETCH (UID <uid> BODY[] {len}\r\n<body>)` items followed by
8//! the tagged status. The coroutine, per message, parses the header line for its
9//! UID, announces it ([`MessageStart`](ImapMessageFetchStreamBatchYield::MessageStart)),
10//! streams the body ([`BodyChunk`](ImapMessageFetchStreamBatchYield::BodyChunk) /
11//! [`WantsStream`](ImapMessageFetchStreamBatchYield::WantsStream)), then closes it
12//! ([`MessageEnd`](ImapMessageFetchStreamBatchYield::MessageEnd)) and moves to the
13//! next — looping until the tagged response.
14//!
15//! The caller is expected to open a fresh sink at `MessageStart` and commit it at
16//! `MessageEnd`; the body of the message in between is routed to that sink. A UID
17//! requested but absent on the server simply never appears (fewer messages than
18//! requested). If a body's FETCH line carries no parseable `UID` (a server that
19//! ordered `UID` *after* the body literal — not seen in practice), the coroutine
20//! fails with [`UidMissing`](ImapMessageFetchStreamBatchError::UidMissing) so the
21//! caller can fall back to per-message fetches rather than misroute a body.
22
23use core::{fmt, num::NonZeroU32};
24
25use alloc::{string::String, string::ToString, vec, vec::Vec};
26
27use imap_codec::{
28    CommandCodec, ResponseCodec,
29    encode::Encoder,
30    fragmentizer::{FragmentInfo, Fragmentizer},
31    imap_types::{
32        command::{Command, CommandBody},
33        core::TagGenerator,
34        fetch::{MacroOrMessageDataItemNames, MessageDataItemName},
35        response::{Response, Status, StatusKind},
36        sequence::SequenceSet,
37    },
38};
39use log::{debug, trace};
40use thiserror::Error;
41
42use crate::coroutine::*;
43
44/// Failure causes during the batched FETCH body-stream flow.
45#[derive(Clone, Debug, Error)]
46pub enum ImapMessageFetchStreamBatchError {
47    /// The server rejected the command with a NO response.
48    #[error("IMAP batched FETCH failed: NO {0}")]
49    No(String),
50    /// The server rejected the command with a BAD response.
51    #[error("IMAP batched FETCH failed: BAD {0}")]
52    Bad(String),
53    /// The server closed the session with an untagged BYE.
54    #[error("IMAP batched FETCH failed: BYE {0}")]
55    Bye(String),
56    /// The exchange ended without a tagged response from the server.
57    #[error("IMAP batched FETCH failed: server did not return a tagged response")]
58    MissingTagged,
59    /// The socket reached EOF before a message's declared body octets were all
60    /// streamed.
61    #[error("IMAP batched FETCH failed: stream ended before the declared body length")]
62    ShortBody,
63    /// A body's FETCH line carried no parseable `UID`, so the body could not be
64    /// attributed to a message. The caller should fall back to per-message
65    /// fetches rather than risk misrouting.
66    #[error("IMAP batched FETCH failed: FETCH body line without a parseable UID")]
67    UidMissing,
68}
69
70/// Yield variants from the batched FETCH body-stream coroutine.
71#[derive(Debug)]
72pub enum ImapMessageFetchStreamBatchYield {
73    /// The caller reads from its stream and resumes with the bytes.
74    WantsRead,
75    /// The caller writes the given bytes to its stream and resumes.
76    WantsWrite(Vec<u8>),
77    /// A new message's body is beginning; the caller opens a sink for `uid`. The
78    /// following `BodyChunk` / `WantsStream` octets belong to it, until
79    /// `MessageEnd`.
80    MessageStart {
81        /// The message's UID, parsed from its FETCH line.
82        uid: u32,
83    },
84    /// Body octets the coroutine already read past the header line; the caller
85    /// writes them to the current message's sink.
86    BodyChunk(Vec<u8>),
87    /// Read exactly `len` octets off the socket straight into the current
88    /// message's sink; resume with `None` on success or `Some(&[])` if the socket
89    /// ran short.
90    WantsStream {
91        /// Number of body octets left to stream for the current message.
92        len: u32,
93    },
94    /// The current message's body is complete; the caller commits its sink.
95    MessageEnd,
96}
97
98/// I/O-free IMAP batched FETCH coroutine streaming many message bodies from one
99/// command.
100pub struct ImapMessageFetchStreamBatch {
101    state: State,
102    command: Option<Vec<u8>>,
103    pending: Vec<u8>,
104    remaining: u32,
105    stream_pending: bool,
106    codec: ResponseCodec,
107}
108
109impl ImapMessageFetchStreamBatch {
110    /// Builds a coroutine streaming the `BODY.PEEK[]` of every message in
111    /// `sequence_set`; when `uid` is `true`, sends `UID FETCH`.
112    pub fn new(sequence_set: SequenceSet, uid: bool) -> Self {
113        let command = Command {
114            tag: TagGenerator::new().generate(),
115            body: CommandBody::Fetch {
116                sequence_set,
117                macro_or_item_names: MacroOrMessageDataItemNames::MessageDataItemNames(vec![
118                    // NOTE: UID first so it lands on the header line, ahead of the
119                    // body literal, and can be parsed before the body streams.
120                    MessageDataItemName::Uid,
121                    MessageDataItemName::BodyExt {
122                        section: None,
123                        partial: None,
124                        peek: true,
125                    },
126                ]),
127                uid,
128                modifiers: Vec::new(),
129            },
130        };
131
132        trace!("send IMAP command {command:?}");
133
134        let command = CommandCodec::new().encode(&command).dump();
135
136        Self {
137            state: State::SendCommand,
138            command: Some(command),
139            pending: Vec::new(),
140            remaining: 0,
141            stream_pending: false,
142            codec: ResponseCodec::new(),
143        }
144    }
145}
146
147impl ImapCoroutine for ImapMessageFetchStreamBatch {
148    type Yield = ImapMessageFetchStreamBatchYield;
149    type Return = Result<(), ImapMessageFetchStreamBatchError>;
150
151    fn resume(
152        &mut self,
153        fragmentizer: &mut Fragmentizer,
154        mut arg: Option<&[u8]>,
155    ) -> ImapCoroutineState<Self::Yield, Self::Return> {
156        loop {
157            match self.state {
158                State::SendCommand => {
159                    let command = self.command.take().expect("command sent once");
160                    self.state = State::NextItem;
161                    debug!("{}", self.state);
162                    return ImapCoroutineState::Yielded(
163                        ImapMessageFetchStreamBatchYield::WantsWrite(command),
164                    );
165                }
166                // Between messages: parse lines until the next FETCH body header
167                // (a new message) or the tagged status (done).
168                State::NextItem => {
169                    if let Some(bytes) = arg.take() {
170                        if bytes.is_empty() {
171                            let err = ImapMessageFetchStreamBatchError::MissingTagged;
172                            return ImapCoroutineState::Complete(Err(err));
173                        }
174                        self.pending.extend_from_slice(bytes);
175                    }
176
177                    loop {
178                        let Some(nl) = self.pending.iter().position(|&b| b == b'\n') else {
179                            return ImapCoroutineState::Yielded(
180                                ImapMessageFetchStreamBatchYield::WantsRead,
181                            );
182                        };
183
184                        let line: Vec<u8> = self.pending.drain(..=nl).collect();
185                        fragmentizer.enqueue_bytes(&line);
186
187                        match fragmentizer.progress() {
188                            // NOTE: a FETCH line announcing a body literal: parse
189                            // its UID, start the message, stream the body next.
190                            Some(FragmentInfo::Line {
191                                announcement: Some(announcement),
192                                ..
193                            }) => {
194                                let Some(uid) = parse_uid(&line) else {
195                                    return ImapCoroutineState::Complete(Err(
196                                        ImapMessageFetchStreamBatchError::UidMissing,
197                                    ));
198                                };
199                                self.remaining = announcement.length;
200                                self.state = State::Stream;
201                                debug!("{}", self.state);
202                                return ImapCoroutineState::Yielded(
203                                    ImapMessageFetchStreamBatchYield::MessageStart { uid },
204                                );
205                            }
206                            // NOTE: a complete line without literal: the tagged
207                            // status (done), a BYE, or an untagged line we skip
208                            // (the literal-closing `)`, stray untagged data).
209                            Some(FragmentInfo::Line {
210                                announcement: None, ..
211                            }) => {
212                                if let Some(result) = self.decode_terminal(fragmentizer) {
213                                    return result;
214                                }
215                            }
216                            _ => {}
217                        }
218                    }
219                }
220                State::Stream => {
221                    if self.remaining == 0 {
222                        // NOTE: drop the bypassed literal, close the message, and
223                        // resume line parsing for the next item / tagged status.
224                        fragmentizer.skip_message();
225                        self.state = State::NextItem;
226                        debug!("{}", self.state);
227                        return ImapCoroutineState::Yielded(
228                            ImapMessageFetchStreamBatchYield::MessageEnd,
229                        );
230                    }
231
232                    if !self.pending.is_empty() {
233                        let take = (self.remaining as usize).min(self.pending.len());
234                        let chunk: Vec<u8> = self.pending.drain(..take).collect();
235                        self.remaining -= take as u32;
236                        return ImapCoroutineState::Yielded(
237                            ImapMessageFetchStreamBatchYield::BodyChunk(chunk),
238                        );
239                    }
240
241                    if self.stream_pending {
242                        self.stream_pending = false;
243                        if matches!(arg.take(), Some(&[])) {
244                            let err = ImapMessageFetchStreamBatchError::ShortBody;
245                            return ImapCoroutineState::Complete(Err(err));
246                        }
247                        self.remaining = 0;
248                        continue;
249                    }
250
251                    self.stream_pending = true;
252                    return ImapCoroutineState::Yielded(
253                        ImapMessageFetchStreamBatchYield::WantsStream {
254                            len: self.remaining,
255                        },
256                    );
257                }
258            }
259        }
260    }
261}
262
263impl ImapMessageFetchStreamBatch {
264    /// Decodes the completed line in `fragmentizer`. Returns `Some` for a terminal
265    /// tagged status or BYE; `None` for undecodable or untagged lines to skip (the
266    /// literal-closing `)`, stray untagged data).
267    fn decode_terminal(
268        &self,
269        fragmentizer: &Fragmentizer,
270    ) -> Option<
271        ImapCoroutineState<
272            ImapMessageFetchStreamBatchYield,
273            Result<(), ImapMessageFetchStreamBatchError>,
274        >,
275    > {
276        match fragmentizer.decode_message(&self.codec) {
277            Ok(Response::Status(Status::Tagged(tagged))) => {
278                let text = tagged.body.text.to_string();
279                let result = match tagged.body.kind {
280                    StatusKind::Ok => Ok(()),
281                    StatusKind::No => Err(ImapMessageFetchStreamBatchError::No(text)),
282                    StatusKind::Bad => Err(ImapMessageFetchStreamBatchError::Bad(text)),
283                };
284                Some(ImapCoroutineState::Complete(result))
285            }
286            Ok(Response::Status(Status::Bye(bye))) => {
287                let err = ImapMessageFetchStreamBatchError::Bye(bye.text.to_string());
288                Some(ImapCoroutineState::Complete(Err(err)))
289            }
290            _ => None,
291        }
292    }
293}
294
295/// Extracts the `UID` value from a FETCH response line, e.g.
296/// `* 12 FETCH (UID 34 BODY[] {1234}`. Scans for a `UID` token (word-boundary
297/// left, whitespace + digits right); `None` when absent or unparseable.
298fn parse_uid(line: &[u8]) -> Option<u32> {
299    let mut i = 0;
300    while i + 3 <= line.len() {
301        let is_uid = line[i..i + 3].eq_ignore_ascii_case(b"UID");
302        let boundary_left = i == 0 || !line[i - 1].is_ascii_alphanumeric();
303        if is_uid && boundary_left {
304            let mut j = i + 3;
305            let mut saw_space = false;
306            while j < line.len() && line[j] == b' ' {
307                j += 1;
308                saw_space = true;
309            }
310            let start = j;
311            while j < line.len() && line[j].is_ascii_digit() {
312                j += 1;
313            }
314            if saw_space && j > start {
315                return core::str::from_utf8(&line[start..j]).ok()?.parse().ok();
316            }
317        }
318        i += 1;
319    }
320    None
321}
322
323/// A convenience over [`ImapMessageFetchStreamBatch::new`] taking a single UID —
324/// unused by the batch driver but handy in tests.
325#[allow(dead_code)]
326fn single(uid: NonZeroU32) -> ImapMessageFetchStreamBatch {
327    ImapMessageFetchStreamBatch::new(SequenceSet::from(uid), true)
328}
329
330#[derive(Clone, Copy)]
331enum State {
332    SendCommand,
333    NextItem,
334    Stream,
335}
336
337impl fmt::Display for State {
338    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
339        match self {
340            Self::SendCommand => f.write_str("send batched fetch command"),
341            Self::NextItem => f.write_str("parse next fetch item"),
342            Self::Stream => f.write_str("stream body"),
343        }
344    }
345}
346
347#[cfg(test)]
348mod tests {
349    use core::str;
350
351    use alloc::{borrow::ToOwned, format, vec::Vec};
352
353    use super::*;
354
355    /// Drives the coroutine over a single canned reply, collecting `(uid, body)`
356    /// pairs. Panics on any error.
357    fn run_ok(cmd_set: &str, reply_after_tag: impl Fn(&str) -> String) -> Vec<(u32, Vec<u8>)> {
358        let set: SequenceSet = cmd_set.try_into().unwrap();
359        let mut cor = ImapMessageFetchStreamBatch::new(set, true);
360        let mut frag = Fragmentizer::new(50 * 1024 * 1024);
361
362        // First resume writes the command; capture its tag.
363        let cmd = match cor.resume(&mut frag, None) {
364            ImapCoroutineState::Yielded(ImapMessageFetchStreamBatchYield::WantsWrite(b)) => b,
365            s => panic!("expected WantsWrite, got {s:?}"),
366        };
367        let tag = str::from_utf8(&cmd)
368            .unwrap()
369            .split_whitespace()
370            .next()
371            .unwrap()
372            .to_owned();
373        let reply = reply_after_tag(&tag);
374
375        let mut out: Vec<(u32, Vec<u8>)> = Vec::new();
376        let mut cur_uid: Option<u32> = None;
377        let mut cur_body: Vec<u8> = Vec::new();
378        let mut fed = false;
379        let mut arg: Option<&[u8]> = None;
380        let reply_bytes = reply.as_bytes();
381
382        loop {
383            match cor.resume(&mut frag, arg.take()) {
384                ImapCoroutineState::Complete(Ok(())) => break,
385                ImapCoroutineState::Complete(Err(e)) => panic!("unexpected error: {e:?}"),
386                ImapCoroutineState::Yielded(ImapMessageFetchStreamBatchYield::WantsRead) => {
387                    // Feed the whole reply on the first read, EOF-empty after.
388                    arg = if !fed {
389                        fed = true;
390                        Some(reply_bytes)
391                    } else {
392                        Some(&[])
393                    };
394                }
395                ImapCoroutineState::Yielded(ImapMessageFetchStreamBatchYield::WantsWrite(_)) => {}
396                ImapCoroutineState::Yielded(ImapMessageFetchStreamBatchYield::MessageStart {
397                    uid,
398                }) => {
399                    cur_uid = Some(uid);
400                    cur_body.clear();
401                }
402                ImapCoroutineState::Yielded(ImapMessageFetchStreamBatchYield::BodyChunk(b)) => {
403                    cur_body.extend_from_slice(&b);
404                }
405                ImapCoroutineState::Yielded(ImapMessageFetchStreamBatchYield::WantsStream {
406                    ..
407                }) => {
408                    // The reply is fed in one go, so bodies always arrive as
409                    // BodyChunk from pending; a WantsStream here means the test's
410                    // canned reply was too fragmented — signal short.
411                    arg = Some(&[]);
412                }
413                ImapCoroutineState::Yielded(ImapMessageFetchStreamBatchYield::MessageEnd) => {
414                    out.push((cur_uid.take().unwrap(), core::mem::take(&mut cur_body)));
415                }
416            }
417        }
418        out
419    }
420
421    #[test]
422    fn command_requests_uid_and_body_peek() {
423        let set: SequenceSet = "1,2,3".try_into().unwrap();
424        let mut cor = ImapMessageFetchStreamBatch::new(set, true);
425        let mut frag = Fragmentizer::new(50 * 1024 * 1024);
426        let cmd = match cor.resume(&mut frag, None) {
427            ImapCoroutineState::Yielded(ImapMessageFetchStreamBatchYield::WantsWrite(b)) => b,
428            s => panic!("expected WantsWrite, got {s:?}"),
429        };
430        let line = str::from_utf8(&cmd).unwrap();
431        // imap-codec collapses `1,2,3` to the range `1:3`.
432        assert!(line.contains("UID FETCH 1:3 (UID BODY.PEEK[])"), "{line}");
433    }
434
435    #[test]
436    fn streams_two_bodies_routed_by_uid() {
437        let bodies = run_ok("10,11", |tag| {
438            format!(
439                "* 1 FETCH (UID 10 BODY[] {{5}}\r\nhello)\r\n\
440                 * 2 FETCH (UID 11 BODY[] {{5}}\r\nworld)\r\n\
441                 {tag} OK FETCH completed\r\n"
442            )
443        });
444        assert_eq!(bodies.len(), 2);
445        assert_eq!(bodies[0], (10, b"hello".to_vec()));
446        assert_eq!(bodies[1], (11, b"world".to_vec()));
447    }
448
449    #[test]
450    fn routes_by_uid_not_by_position() {
451        // Server returns them in a different order than requested: routing must
452        // follow the UID on each line, not arrival order.
453        let bodies = run_ok("10,11", |tag| {
454            format!(
455                "* 2 FETCH (UID 11 BODY[] {{3}}\r\nBBB)\r\n\
456                 * 1 FETCH (UID 10 BODY[] {{3}}\r\nAAA)\r\n\
457                 {tag} OK done\r\n"
458            )
459        });
460        assert_eq!(bodies, vec![(11, b"BBB".to_vec()), (10, b"AAA".to_vec())]);
461    }
462
463    #[test]
464    fn skips_interleaved_untagged_and_missing_uids() {
465        // A requested UID with no data simply never appears; an interleaved
466        // untagged EXPUNGE between items is skipped.
467        let bodies = run_ok("10,11,12", |tag| {
468            format!(
469                "* 1 FETCH (UID 10 BODY[] {{2}}\r\nhi)\r\n\
470                 * 3 EXPUNGE\r\n\
471                 * 4 FETCH (UID 12 BODY[] {{2}}\r\nyo)\r\n\
472                 {tag} OK done\r\n"
473            )
474        });
475        assert_eq!(bodies, vec![(10, b"hi".to_vec()), (12, b"yo".to_vec())]);
476    }
477
478    #[test]
479    fn empty_result_completes_clean() {
480        let bodies = run_ok("99", |tag| format!("{tag} OK nothing\r\n"));
481        assert!(bodies.is_empty());
482    }
483
484    #[test]
485    fn tagged_no_is_an_error() {
486        let set: SequenceSet = "1".try_into().unwrap();
487        let mut cor = ImapMessageFetchStreamBatch::new(set, true);
488        let mut frag = Fragmentizer::new(50 * 1024 * 1024);
489        let cmd = match cor.resume(&mut frag, None) {
490            ImapCoroutineState::Yielded(ImapMessageFetchStreamBatchYield::WantsWrite(b)) => b,
491            s => panic!("{s:?}"),
492        };
493        let tag = str::from_utf8(&cmd)
494            .unwrap()
495            .split_whitespace()
496            .next()
497            .unwrap()
498            .to_owned();
499        // WantsRead, then feed a NO.
500        assert!(matches!(
501            cor.resume(&mut frag, None),
502            ImapCoroutineState::Yielded(ImapMessageFetchStreamBatchYield::WantsRead)
503        ));
504        let reply = format!("{tag} NO mailbox gone\r\n");
505        match cor.resume(&mut frag, Some(reply.as_bytes())) {
506            ImapCoroutineState::Complete(Err(ImapMessageFetchStreamBatchError::No(t))) => {
507                assert_eq!(t, "mailbox gone")
508            }
509            s => panic!("expected No error, got {s:?}"),
510        }
511    }
512
513    #[test]
514    fn body_line_without_uid_errs_for_fallback() {
515        // A body FETCH line whose UID we cannot parse must error (so the caller
516        // falls back to per-message) rather than misroute the body.
517        let set: SequenceSet = "1".try_into().unwrap();
518        let mut cor = ImapMessageFetchStreamBatch::new(set, true);
519        let mut frag = Fragmentizer::new(50 * 1024 * 1024);
520        let _ = cor.resume(&mut frag, None); // WantsWrite
521        let _ = cor.resume(&mut frag, None); // WantsRead
522        let reply = "* 1 FETCH (BODY[] {3}\r\nxxx)\r\nA1 OK done\r\n";
523        match cor.resume(&mut frag, Some(reply.as_bytes())) {
524            ImapCoroutineState::Complete(Err(ImapMessageFetchStreamBatchError::UidMissing)) => {}
525            s => panic!("expected UidMissing, got {s:?}"),
526        }
527    }
528
529    #[test]
530    fn parse_uid_finds_the_token() {
531        assert_eq!(parse_uid(b"* 12 FETCH (UID 34 BODY[] {5}\r\n"), Some(34));
532        assert_eq!(
533            parse_uid(b"* 1 FETCH (FLAGS (\\Seen) UID 7 BODY[] {2}\r\n"),
534            Some(7)
535        );
536        assert_eq!(parse_uid(b"* 1 FETCH (BODY[] {2}\r\n"), None);
537        // A "UID"-like substring in a word must not match (word boundary).
538        assert_eq!(parse_uid(b"* 1 FETCH (XUID 9 BODY[] {2}\r\n"), None);
539    }
540}
io_imap/rfc3501/fetch_stream_batch.rs

io_imap/rfc3501/
fetch_stream_batch.rs
