Skip to main content

io_imap/rfc3501/
greeting.rs

1//! IMAP server greeting reader; optionally forces a CAPABILITY round-trip
2//! if the greeting carries none.
3//!
4//! # Example
5//!
6//! ```rust,no_run
7//! use std::{
8//!     io::{Read, Write},
9//!     net::TcpStream,
10//! };
11//!
12//! use io_imap::{
13//!     codec::fragmentizer::Fragmentizer,
14//!     coroutine::{ImapCoroutine, ImapCoroutineState, ImapYield},
15//!     rfc3501::greeting::{ImapGreetingGet, ImapGreetingGetOptions},
16//! };
17//!
18//! // Ready stream needed (TCP-connected, TLS-negotiated)
19//! let mut stream = TcpStream::connect("localhost:143").unwrap();
20//!
21//! let mut fragmentizer = Fragmentizer::new(50 * 1024 * 1024);
22//! let mut buf = [0u8; 4096];
23//!
24//! let opts = ImapGreetingGetOptions {
25//!     ensure_capabilities: true,
26//! };
27//! let mut coroutine = ImapGreetingGet::new(opts);
28//! let mut arg = None;
29//!
30//! let greeting = loop {
31//!     match coroutine.resume(&mut fragmentizer, arg.take()) {
32//!         ImapCoroutineState::Yielded(ImapYield::WantsWrite(bytes)) => {
33//!             stream.write_all(&bytes).unwrap();
34//!         }
35//!         ImapCoroutineState::Yielded(ImapYield::WantsRead) => {
36//!             let n = stream.read(&mut buf).unwrap();
37//!             arg = Some(&buf[..n]);
38//!         }
39//!         ImapCoroutineState::Complete(Ok(greeting)) => break greeting,
40//!         ImapCoroutineState::Complete(Err(err)) => panic!("{err}"),
41//!     }
42//! };
43//!
44//! println!("{greeting:?}");
45//! ```
46
47use core::{fmt, mem};
48
49use alloc::{boxed::Box, string::String, string::ToString, vec::Vec};
50
51use imap_codec::{
52    GreetingCodec,
53    fragmentizer::{DecodeMessageError, FragmentInfo, Fragmentizer},
54    imap_types::{
55        IntoStatic,
56        response::{Capability, Code, GreetingKind},
57        secret::Secret,
58        utils::escape_byte_string,
59    },
60};
61use log::{debug, trace};
62use thiserror::Error;
63
64use crate::{coroutine::*, imap_try, rfc3501::capability::*};
65
66/// Failure causes while reading the IMAP greeting.
67#[derive(Clone, Debug, Error)]
68pub enum ImapGreetingGetError {
69    /// The server refused the connection with a BYE greeting.
70    #[error("IMAP greeting failed: BYE {0}")]
71    Bye(String),
72    /// The stream reached EOF before a complete greeting arrived.
73    #[error("IMAP greeting failed: reached unexpected EOF on stream")]
74    Eof,
75    /// The greeting line could not be decoded; carries the raw bytes.
76    #[error("IMAP greeting failed: decode error")]
77    DecodingFailure(Secret<Box<[u8]>>),
78    /// The `Fragmentizer` poisoned the message after a framing error;
79    /// carries the raw bytes.
80    #[error("IMAP greeting failed: parse error: message is poisoned")]
81    MessageIsPoisoned(Secret<Box<[u8]>>),
82    /// The greeting exceeded the `Fragmentizer` size limit; carries the
83    /// raw bytes.
84    #[error("IMAP greeting failed: parse error: message is too long")]
85    MessageTooLong(Secret<Box<[u8]>>),
86    /// The follow-up CAPABILITY round-trip failed.
87    #[error(transparent)]
88    Capability(#[from] ImapCapabilityGetError),
89}
90
91/// Decoded greeting outcome.
92#[derive(Debug)]
93pub struct ImapGreetingOk {
94    /// Capabilities advertised in the greeting (or fetched right after,
95    /// see [`ImapGreetingGetOptions::ensure_capabilities`]).
96    pub capability: Vec<Capability<'static>>,
97    /// Whether the greeting was `PREAUTH`: the session is already
98    /// authenticated and LOGIN must be skipped.
99    pub pre_authenticated: bool,
100}
101
102/// Options for [`ImapGreetingGet::new`].
103#[derive(Clone, Debug, Default, Eq, PartialEq)]
104pub struct ImapGreetingGetOptions {
105    /// Fetch capabilities explicitly when the greeting carries none.
106    pub ensure_capabilities: bool,
107}
108
109/// I/O-free IMAP greeting-read coroutine.
110pub struct ImapGreetingGet {
111    codec: GreetingCodec,
112    state: State,
113    wants_read: bool,
114    observed: Vec<Capability<'static>>,
115    pre_authenticated: bool,
116    opts: ImapGreetingGetOptions,
117}
118
119impl ImapGreetingGet {
120    /// Builds a coroutine reading the server greeting; depending on
121    /// `opts`, it follows up with CAPABILITY when the greeting has none.
122    pub fn new(opts: ImapGreetingGetOptions) -> Self {
123        Self {
124            codec: GreetingCodec::new(),
125            state: State::Read,
126            wants_read: false,
127            observed: Vec::new(),
128            pre_authenticated: false,
129            opts,
130        }
131    }
132}
133
134impl ImapCoroutine for ImapGreetingGet {
135    type Yield = ImapYield;
136    type Return = Result<ImapGreetingOk, ImapGreetingGetError>;
137
138    fn resume(
139        &mut self,
140        fragmentizer: &mut Fragmentizer,
141        mut arg: Option<&[u8]>,
142    ) -> ImapCoroutineState<Self::Yield, Self::Return> {
143        loop {
144            if mem::take(&mut self.wants_read) {
145                return ImapCoroutineState::Yielded(ImapYield::WantsRead);
146            }
147
148            match &mut self.state {
149                State::Read => match arg.take() {
150                    Some(&[]) => {
151                        return ImapCoroutineState::Complete(Err(ImapGreetingGetError::Eof));
152                    }
153                    Some(bytes) => {
154                        trace!("read bytes: {}", escape_byte_string(bytes));
155                        fragmentizer.enqueue_bytes(bytes);
156                        self.state = State::Deserialize;
157                    }
158                    None => {
159                        self.wants_read = true;
160                    }
161                },
162                State::Deserialize => match fragmentizer.progress() {
163                    Some(info @ FragmentInfo::Line { .. }) => {
164                        let bytes = fragmentizer.fragment_bytes(info);
165                        trace!("read greeting line: {}", escape_byte_string(bytes));
166
167                        if !fragmentizer.is_message_complete() {
168                            continue;
169                        }
170
171                        match fragmentizer.decode_message(&self.codec) {
172                            Ok(greeting) if greeting.kind == GreetingKind::Bye => {
173                                let err = ImapGreetingGetError::Bye(greeting.text.to_string());
174                                return ImapCoroutineState::Complete(Err(err));
175                            }
176                            Ok(greeting) => {
177                                self.pre_authenticated = greeting.kind == GreetingKind::PreAuth;
178
179                                if let Some(Code::Capability(capability)) = greeting.code {
180                                    self.observed = capability.into_static().into_iter().collect();
181                                }
182
183                                if self.opts.ensure_capabilities && self.observed.is_empty() {
184                                    self.state = State::Capability(Box::default());
185                                    debug!("{}", self.state);
186                                    continue;
187                                }
188
189                                return ImapCoroutineState::Complete(Ok(ImapGreetingOk {
190                                    capability: mem::take(&mut self.observed),
191                                    pre_authenticated: self.pre_authenticated,
192                                }));
193                            }
194                            Err(err) => {
195                                let bytes = fragmentizer.message_bytes();
196                                let bytes = Secret::new(bytes.into());
197                                let err = match err {
198                                    DecodeMessageError::DecodingFailure(_)
199                                    | DecodeMessageError::DecodingRemainder { .. } => {
200                                        ImapGreetingGetError::DecodingFailure(bytes)
201                                    }
202                                    DecodeMessageError::MessageTooLong { .. } => {
203                                        ImapGreetingGetError::MessageTooLong(bytes)
204                                    }
205                                    DecodeMessageError::MessagePoisoned { .. } => {
206                                        ImapGreetingGetError::MessageIsPoisoned(bytes)
207                                    }
208                                };
209                                return ImapCoroutineState::Complete(Err(err));
210                            }
211                        }
212                    }
213                    // NOTE: greetings never carry literals.
214                    Some(FragmentInfo::Literal { .. }) => unreachable!(),
215                    None => {
216                        self.state = State::Read;
217                    }
218                },
219                State::Capability(capability) => {
220                    let caps = imap_try!(capability.as_mut(), fragmentizer, arg.take());
221                    return ImapCoroutineState::Complete(Ok(ImapGreetingOk {
222                        capability: caps,
223                        pre_authenticated: self.pre_authenticated,
224                    }));
225                }
226            }
227        }
228    }
229}
230
231enum State {
232    Read,
233    Deserialize,
234    Capability(Box<ImapCapabilityGet>),
235}
236
237impl fmt::Display for State {
238    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
239        match self {
240            Self::Read => f.write_str("read greeting"),
241            Self::Deserialize => f.write_str("decode greeting"),
242            Self::Capability(_) => f.write_str("fetch capabilities"),
243        }
244    }
245}
246
247#[cfg(test)]
248mod tests {
249    use alloc::vec::Vec;
250
251    use crate::rfc3501::greeting::*;
252
253    #[test]
254    fn ok_with_inline_capability_returns_ok() {
255        let mut greeting = ImapGreetingGet::new(ImapGreetingGetOptions {
256            ensure_capabilities: true,
257        });
258        let mut frag = Fragmentizer::new(50 * 1024 * 1024);
259
260        expect_wants_read(&mut greeting, &mut frag);
261
262        let reply = b"* OK [CAPABILITY IMAP4REV1 IDLE] hello\r\n";
263        let ok = expect_complete_ok(&mut greeting, &mut frag, reply);
264        assert!(!ok.pre_authenticated);
265        assert_eq!(2, ok.capability.len());
266    }
267
268    #[test]
269    fn ok_without_inline_capability_triggers_extra_round_trip() {
270        let mut greeting = ImapGreetingGet::new(ImapGreetingGetOptions {
271            ensure_capabilities: true,
272        });
273        let mut frag = Fragmentizer::new(50 * 1024 * 1024);
274
275        expect_wants_read(&mut greeting, &mut frag);
276        expect_wants_write_after(&mut greeting, &mut frag, b"* OK hello\r\n");
277    }
278
279    #[test]
280    fn preauth_sets_flag() {
281        let mut greeting = ImapGreetingGet::new(ImapGreetingGetOptions::default());
282        let mut frag = Fragmentizer::new(50 * 1024 * 1024);
283
284        expect_wants_read(&mut greeting, &mut frag);
285
286        let reply = b"* PREAUTH [CAPABILITY IMAP4REV1] welcome\r\n";
287        let ok = expect_complete_ok(&mut greeting, &mut frag, reply);
288        assert!(ok.pre_authenticated);
289    }
290
291    #[test]
292    fn bye_returns_bye_error() {
293        let mut greeting = ImapGreetingGet::new(ImapGreetingGetOptions::default());
294        let mut frag = Fragmentizer::new(50 * 1024 * 1024);
295
296        expect_wants_read(&mut greeting, &mut frag);
297
298        let err = expect_complete_err(&mut greeting, &mut frag, b"* BYE service unavailable\r\n");
299        let ImapGreetingGetError::Bye(text) = err else {
300            panic!("expected ImapGreetingGetError::Bye, got {err:?}");
301        };
302        assert_eq!(text, "service unavailable");
303    }
304
305    #[test]
306    fn eof_returns_eof_error() {
307        let mut greeting = ImapGreetingGet::new(ImapGreetingGetOptions::default());
308        let mut frag = Fragmentizer::new(50 * 1024 * 1024);
309
310        expect_wants_read(&mut greeting, &mut frag);
311
312        let err = expect_complete_err(&mut greeting, &mut frag, b"");
313        assert!(matches!(err, ImapGreetingGetError::Eof));
314    }
315
316    fn expect_wants_read(cor: &mut ImapGreetingGet, frag: &mut Fragmentizer) {
317        match cor.resume(frag, None) {
318            ImapCoroutineState::Yielded(ImapYield::WantsRead) => {}
319            state => panic!("expected WantsRead, got {state:?}"),
320        }
321    }
322
323    fn expect_wants_write_after(
324        cor: &mut ImapGreetingGet,
325        frag: &mut Fragmentizer,
326        arg: &[u8],
327    ) -> Vec<u8> {
328        match cor.resume(frag, Some(arg)) {
329            ImapCoroutineState::Yielded(ImapYield::WantsWrite(bytes)) => bytes,
330            state => panic!("expected WantsWrite, got {state:?}"),
331        }
332    }
333
334    fn expect_complete_ok(
335        cor: &mut ImapGreetingGet,
336        frag: &mut Fragmentizer,
337        reply: &[u8],
338    ) -> ImapGreetingOk {
339        match cor.resume(frag, Some(reply)) {
340            ImapCoroutineState::Complete(Ok(value)) => value,
341            state => panic!("expected Complete(Ok), got {state:?}"),
342        }
343    }
344
345    fn expect_complete_err(
346        cor: &mut ImapGreetingGet,
347        frag: &mut Fragmentizer,
348        reply: &[u8],
349    ) -> ImapGreetingGetError {
350        match cor.resume(frag, Some(reply)) {
351            ImapCoroutineState::Complete(Err(err)) => err,
352            state => panic!("expected Complete(Err), got {state:?}"),
353        }
354    }
355}