1use 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#[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
80pub 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 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}