Skip to main content

io_imap/rfc3501/
capability.rs

1//! IMAP CAPABILITY coroutine returning the advertised capability list.
2//!
3//! # Example
4//!
5//! ```rust,no_run
6//! use std::{
7//!     io::{Read, Write},
8//!     net::TcpStream,
9//! };
10//!
11//! use io_imap::{
12//!     codec::fragmentizer::Fragmentizer,
13//!     coroutine::{ImapCoroutine, ImapCoroutineState, ImapYield},
14//!     rfc3501::capability::ImapCapabilityGet,
15//! };
16//!
17//! // Ready stream needed (TCP-connected, TLS-negociated, IMAP-authenticated)
18//! let mut stream = TcpStream::connect("localhost:143").unwrap();
19//!
20//! let mut fragmentizer = Fragmentizer::new(50 * 1024 * 1024);
21//! let mut buf = [0u8; 4096];
22//!
23//! let mut coroutine = ImapCapabilityGet::new();
24//! let mut arg = None;
25//!
26//! let capability = loop {
27//!     match coroutine.resume(&mut fragmentizer, arg.take()) {
28//!         ImapCoroutineState::Yielded(ImapYield::WantsWrite(bytes)) => {
29//!             stream.write_all(&bytes).unwrap();
30//!         }
31//!         ImapCoroutineState::Yielded(ImapYield::WantsRead) => {
32//!             let n = stream.read(&mut buf).unwrap();
33//!             arg = Some(&buf[..n]);
34//!         }
35//!         ImapCoroutineState::Complete(Ok(capability)) => break capability,
36//!         ImapCoroutineState::Complete(Err(err)) => panic!("{err}"),
37//!     }
38//! };
39//!
40//! println!("{capability:?}");
41//! ```
42
43use core::fmt;
44
45use alloc::{string::String, string::ToString, vec::Vec};
46
47use imap_codec::{
48    CommandCodec,
49    fragmentizer::Fragmentizer,
50    imap_types::{
51        command::{Command, CommandBody},
52        core::TagGenerator,
53        response::{Capability, Code, Data, StatusBody, StatusKind, Tagged},
54    },
55};
56use log::trace;
57use thiserror::Error;
58
59use crate::{coroutine::*, imap_try, send::*};
60
61/// Failure causes during the IMAP CAPABILITY flow.
62#[derive(Clone, Debug, Error)]
63pub enum ImapCapabilityGetError {
64    #[error("IMAP CAPABILITY failed: NO {0}")]
65    No(String),
66    #[error("IMAP CAPABILITY failed: BAD {0}")]
67    Bad(String),
68    #[error("IMAP CAPABILITY failed: BYE {0}")]
69    Bye(String),
70
71    #[error("IMAP CAPABILITY failed: server did not return a tagged response")]
72    MissingTagged,
73    #[error("IMAP CAPABILITY failed: server did not advertise any capability")]
74    MissingCapability,
75
76    #[error("IMAP CAPABILITY failed: {0}")]
77    Send(#[from] SendImapCommandError),
78}
79
80/// I/O-free IMAP CAPABILITY coroutine.
81pub struct ImapCapabilityGet {
82    state: State,
83}
84
85impl ImapCapabilityGet {
86    pub fn new() -> Self {
87        let command = Command {
88            tag: TagGenerator::new().generate(),
89            body: CommandBody::Capability,
90        };
91
92        trace!("send IMAP command {command:?}");
93
94        let state = State::Send(SendImapCommand::new(CommandCodec::new(), command));
95
96        Self { state }
97    }
98}
99
100impl Default for ImapCapabilityGet {
101    fn default() -> Self {
102        Self::new()
103    }
104}
105
106impl ImapCoroutine for ImapCapabilityGet {
107    type Yield = ImapYield;
108    type Return = Result<Vec<Capability<'static>>, ImapCapabilityGetError>;
109
110    fn resume(
111        &mut self,
112        fragmentizer: &mut Fragmentizer,
113        arg: Option<&[u8]>,
114    ) -> ImapCoroutineState<Self::Yield, Self::Return> {
115        loop {
116            trace!("capability: {}", self.state);
117
118            match &mut self.state {
119                State::Send(send) => {
120                    let out = imap_try!(send, fragmentizer, arg);
121
122                    if let Some(bye) = out.bye {
123                        let err = ImapCapabilityGetError::Bye(bye.text.to_string());
124                        return ImapCoroutineState::Complete(Err(err));
125                    }
126
127                    let Some(Tagged { body, .. }) = out.tagged else {
128                        let err = ImapCapabilityGetError::MissingTagged;
129                        return ImapCoroutineState::Complete(Err(err));
130                    };
131
132                    let code = match body.kind {
133                        StatusKind::Ok => body.code,
134                        StatusKind::No => {
135                            let err = ImapCapabilityGetError::No(body.text.to_string());
136                            return ImapCoroutineState::Complete(Err(err));
137                        }
138                        StatusKind::Bad => {
139                            let err = ImapCapabilityGetError::Bad(body.text.to_string());
140                            return ImapCoroutineState::Complete(Err(err));
141                        }
142                    };
143
144                    let mut new_capability = None;
145
146                    if let Some(Code::Capability(capability)) = code {
147                        new_capability.replace(capability);
148                    }
149
150                    for data in out.data {
151                        if let Data::Capability(capability) = data {
152                            new_capability.replace(capability);
153                        }
154                    }
155
156                    for StatusBody { code, .. } in out.untagged {
157                        if let Some(Code::Capability(capability)) = code {
158                            new_capability.replace(capability);
159                        }
160                    }
161
162                    let Some(capability) = new_capability else {
163                        let err = ImapCapabilityGetError::MissingCapability;
164                        return ImapCoroutineState::Complete(Err(err));
165                    };
166
167                    return ImapCoroutineState::Complete(Ok(capability.into_iter().collect()));
168                }
169            }
170        }
171    }
172}
173
174enum State {
175    Send(SendImapCommand<CommandCodec>),
176}
177
178impl fmt::Display for State {
179    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
180        match self {
181            Self::Send(_) => f.write_str("send capability"),
182        }
183    }
184}
185
186#[cfg(test)]
187mod tests {
188    use core::str;
189
190    use alloc::borrow::ToOwned;
191
192    use super::*;
193
194    #[test]
195    fn data_capability_returns_capabilities() {
196        let mut cap = ImapCapabilityGet::new();
197        let mut frag = Fragmentizer::new(50 * 1024 * 1024);
198
199        let bytes = expect_wants_write(&mut cap, &mut frag, None);
200        let line = str::from_utf8(&bytes).expect("utf8 command");
201        let tag = first_word(line).to_owned();
202        assert!(line.trim_end().ends_with("CAPABILITY"));
203
204        expect_wants_read(&mut cap, &mut frag);
205
206        let reply =
207            format!("* CAPABILITY IMAP4REV1 STARTTLS IDLE\r\n{tag} OK CAPABILITY completed\r\n");
208        let caps = expect_complete_ok(&mut cap, &mut frag, reply.as_bytes());
209        assert_eq!(3, caps.len());
210        assert!(caps.contains(&Capability::Imap4Rev1));
211        assert!(caps.contains(&Capability::StartTls));
212        assert!(caps.contains(&Capability::Idle));
213    }
214
215    #[test]
216    fn tagged_code_capability_returns_capabilities() {
217        let mut cap = ImapCapabilityGet::new();
218        let mut frag = Fragmentizer::new(50 * 1024 * 1024);
219
220        let bytes = expect_wants_write(&mut cap, &mut frag, None);
221        let tag = first_word(str::from_utf8(&bytes).expect("utf8 command")).to_owned();
222
223        expect_wants_read(&mut cap, &mut frag);
224
225        let reply = format!("{tag} OK [CAPABILITY IMAP4REV1 IDLE] done\r\n");
226        let caps = expect_complete_ok(&mut cap, &mut frag, reply.as_bytes());
227        assert_eq!(2, caps.len());
228    }
229
230    #[test]
231    fn no_capability_returns_missing_error() {
232        let mut cap = ImapCapabilityGet::new();
233        let mut frag = Fragmentizer::new(50 * 1024 * 1024);
234
235        let bytes = expect_wants_write(&mut cap, &mut frag, None);
236        let tag = first_word(str::from_utf8(&bytes).expect("utf8 command")).to_owned();
237
238        expect_wants_read(&mut cap, &mut frag);
239
240        let reply = format!("{tag} OK CAPABILITY completed\r\n");
241        let err = expect_complete_err(&mut cap, &mut frag, reply.as_bytes());
242        assert!(matches!(err, ImapCapabilityGetError::MissingCapability));
243    }
244
245    #[test]
246    fn tagged_no_returns_no_error() {
247        let mut cap = ImapCapabilityGet::new();
248        let mut frag = Fragmentizer::new(50 * 1024 * 1024);
249
250        let bytes = expect_wants_write(&mut cap, &mut frag, None);
251        let tag = first_word(str::from_utf8(&bytes).expect("utf8 command")).to_owned();
252
253        expect_wants_read(&mut cap, &mut frag);
254
255        let reply = format!("{tag} NO server is sulking\r\n");
256        let err = expect_complete_err(&mut cap, &mut frag, reply.as_bytes());
257        let ImapCapabilityGetError::No(text) = err else {
258            panic!("expected ImapCapabilityGetError::No, got {err:?}");
259        };
260        assert_eq!(text, "server is sulking");
261    }
262
263    #[test]
264    fn bye_returns_bye_error() {
265        let mut cap = ImapCapabilityGet::new();
266        let mut frag = Fragmentizer::new(50 * 1024 * 1024);
267
268        let _ = expect_wants_write(&mut cap, &mut frag, None);
269        expect_wants_read(&mut cap, &mut frag);
270
271        let err = expect_complete_err(&mut cap, &mut frag, b"* BYE going down\r\n");
272        let ImapCapabilityGetError::Bye(text) = err else {
273            panic!("expected ImapCapabilityGetError::Bye, got {err:?}");
274        };
275        assert_eq!(text, "going down");
276    }
277
278    // --- utils
279
280    fn expect_wants_write(
281        cor: &mut ImapCapabilityGet,
282        frag: &mut Fragmentizer,
283        arg: Option<&[u8]>,
284    ) -> Vec<u8> {
285        match cor.resume(frag, arg) {
286            ImapCoroutineState::Yielded(ImapYield::WantsWrite(bytes)) => bytes,
287            state => panic!("expected WantsWrite, got {state:?}"),
288        }
289    }
290
291    fn expect_wants_read(cor: &mut ImapCapabilityGet, frag: &mut Fragmentizer) {
292        match cor.resume(frag, None) {
293            ImapCoroutineState::Yielded(ImapYield::WantsRead) => {}
294            state => panic!("expected WantsRead, got {state:?}"),
295        }
296    }
297
298    fn expect_complete_ok(
299        cor: &mut ImapCapabilityGet,
300        frag: &mut Fragmentizer,
301        reply: &[u8],
302    ) -> Vec<Capability<'static>> {
303        match cor.resume(frag, Some(reply)) {
304            ImapCoroutineState::Complete(Ok(value)) => value,
305            state => panic!("expected Complete(Ok), got {state:?}"),
306        }
307    }
308
309    fn expect_complete_err(
310        cor: &mut ImapCapabilityGet,
311        frag: &mut Fragmentizer,
312        reply: &[u8],
313    ) -> ImapCapabilityGetError {
314        match cor.resume(frag, Some(reply)) {
315            ImapCoroutineState::Complete(Err(err)) => err,
316            state => panic!("expected Complete(Err), got {state:?}"),
317        }
318    }
319
320    fn first_word(line: &str) -> &str {
321        line.split_whitespace()
322            .next()
323            .expect("first whitespace-separated token")
324    }
325}