1use 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#[derive(Clone, Debug, Error)]
68pub enum ImapGreetingGetError {
69 #[error("IMAP greeting failed: BYE {0}")]
71 Bye(String),
72 #[error("IMAP greeting failed: reached unexpected EOF on stream")]
74 Eof,
75 #[error("IMAP greeting failed: decode error")]
77 DecodingFailure(Secret<Box<[u8]>>),
78 #[error("IMAP greeting failed: parse error: message is poisoned")]
81 MessageIsPoisoned(Secret<Box<[u8]>>),
82 #[error("IMAP greeting failed: parse error: message is too long")]
85 MessageTooLong(Secret<Box<[u8]>>),
86 #[error(transparent)]
88 Capability(#[from] ImapCapabilityGetError),
89}
90
91#[derive(Debug)]
93pub struct ImapGreetingOk {
94 pub capability: Vec<Capability<'static>>,
97 pub pre_authenticated: bool,
100}
101
102#[derive(Clone, Debug, Default, Eq, PartialEq)]
104pub struct ImapGreetingGetOptions {
105 pub ensure_capabilities: bool,
107}
108
109pub 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 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 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}