aerogramme/src/imap/command/selected.rs

use std::sync::Arc;
use std::num::NonZeroU64;

use anyhow::Result;
use imap_codec::imap_types::command::{Command, CommandBody, FetchModifier, StoreModifier};
use imap_codec::imap_types::core::Charset;
use imap_codec::imap_types::fetch::MacroOrMessageDataItemNames;
use imap_codec::imap_types::flag::{Flag, StoreResponse, StoreType};
use imap_codec::imap_types::mailbox::Mailbox as MailboxCodec;
use imap_codec::imap_types::response::{Code, CodeOther};
use imap_codec::imap_types::search::SearchKey;
use imap_codec::imap_types::sequence::SequenceSet;

use crate::imap::capability::{ClientCapability, ServerCapability};
use crate::imap::command::{anystate, authenticated, MailboxName};
use crate::imap::flow;
use crate::imap::mailbox_view::{MailboxView, UpdateParameters};
use crate::imap::response::Response;
use crate::imap::attributes::AttributesProxy;
use crate::mail::user::User;

pub struct SelectedContext<'a> {
    pub req: &'a Command<'static>,
    pub user: &'a Arc<User>,
    pub mailbox: &'a mut MailboxView,
    pub server_capabilities: &'a ServerCapability,
    pub client_capabilities: &'a mut ClientCapability,
}

pub async fn dispatch<'a>(
    ctx: SelectedContext<'a>,
) -> Result<(Response<'static>, flow::Transition)> {
    match &ctx.req.body {
        // Any State
        // noop is specific to this state
        CommandBody::Capability => {
            anystate::capability(ctx.req.tag.clone(), ctx.server_capabilities)
        }
        CommandBody::Logout => anystate::logout(),

        // Specific to this state (7 commands + NOOP)
        CommandBody::Close => ctx.close().await,
        CommandBody::Noop | CommandBody::Check => ctx.noop().await,
        CommandBody::Fetch {
            sequence_set,
            macro_or_item_names,
            modifiers,
            uid,
        } => ctx.fetch(sequence_set, macro_or_item_names, modifiers, uid).await,
        CommandBody::Search {
            charset,
            criteria,
            uid,
        } => ctx.search(charset, criteria, uid).await,
        CommandBody::Expunge => ctx.expunge().await,
        CommandBody::Store {
            sequence_set,
            kind,
            response,
            flags,
            modifiers,
            uid,
        } => ctx.store(sequence_set, kind, response, flags, modifiers, uid).await,
        CommandBody::Copy {
            sequence_set,
            mailbox,
            uid,
        } => ctx.copy(sequence_set, mailbox, uid).await,
        CommandBody::Move {
            sequence_set,
            mailbox,
            uid,
        } => ctx.r#move(sequence_set, mailbox, uid).await,

        // UNSELECT extension (rfc3691)
        CommandBody::Unselect => ctx.unselect().await,

        // In selected mode, we fallback to authenticated when needed
        _ => {
            authenticated::dispatch(authenticated::AuthenticatedContext {
                req: ctx.req,
                server_capabilities: ctx.server_capabilities,
                client_capabilities: ctx.client_capabilities,
                user: ctx.user,
            })
            .await
        }
    }
}

// --- PRIVATE ---

impl<'a> SelectedContext<'a> {
    async fn close(self) -> Result<(Response<'static>, flow::Transition)> {
        // We expunge messages,
        // but we don't send the untagged EXPUNGE responses
        let tag = self.req.tag.clone();
        self.expunge().await?;
        Ok((
            Response::build().tag(tag).message("CLOSE completed").ok()?,
            flow::Transition::Unselect,
        ))
    }

    async fn unselect(self) -> Result<(Response<'static>, flow::Transition)> {
        Ok((
            Response::build()
                .to_req(self.req)
                .message("UNSELECT completed")
                .ok()?,
            flow::Transition::Unselect,
        ))
    }

    pub async fn fetch(
        self,
        sequence_set: &SequenceSet,
        attributes: &'a MacroOrMessageDataItemNames<'static>,
        modifiers: &[FetchModifier],
        uid: &bool,
    ) -> Result<(Response<'static>, flow::Transition)> {
        let ap = AttributesProxy::new(attributes, modifiers, *uid);
        let mut changed_since: Option<NonZeroU64> = None;
        modifiers.iter().for_each(|m| match m {
            FetchModifier::ChangedSince(val) => {
                changed_since = Some(*val);
            },
        });

        match self.mailbox.fetch(sequence_set, &ap, changed_since, uid).await {
            Ok(resp) => {
                // Capabilities enabling logic only on successful command
                // (according to my understanding of the spec)
                self.client_capabilities.attributes_enable(&ap);
                self.client_capabilities.fetch_modifiers_enable(modifiers);

                // Response to the client
                Ok((
                    Response::build()
                        .to_req(self.req)
                        .message("FETCH completed")
                        .set_body(resp)
                        .ok()?,
                    flow::Transition::None,
                ))
            },
            Err(e) => Ok((
                Response::build()
                    .to_req(self.req)
                    .message(e.to_string())
                    .no()?,
                flow::Transition::None,
            )),
        }
    }

    pub async fn search(
        self,
        charset: &Option<Charset<'a>>,
        criteria: &SearchKey<'a>,
        uid: &bool,
    ) -> Result<(Response<'static>, flow::Transition)> {
        let (found, enable_condstore) = self.mailbox.search(charset, criteria, *uid).await?;
        if enable_condstore {
            self.client_capabilities.enable_condstore();
        }
        Ok((
            Response::build()
                .to_req(self.req)
                .set_body(found)
                .message("SEARCH completed")
                .ok()?,
            flow::Transition::None,
        ))
    }

    pub async fn noop(self) -> Result<(Response<'static>, flow::Transition)> {
        self.mailbox.internal.mailbox.force_sync().await?;

        let updates = self.mailbox.update(UpdateParameters::default()).await?;
        Ok((
            Response::build()
                .to_req(self.req)
                .message("NOOP completed.")
                .set_body(updates)
                .ok()?,
            flow::Transition::None,
        ))
    }

    async fn expunge(self) -> Result<(Response<'static>, flow::Transition)> {
        let tag = self.req.tag.clone();
        let data = self.mailbox.expunge().await?;

        Ok((
            Response::build()
                .tag(tag)
                .message("EXPUNGE completed")
                .set_body(data)
                .ok()?,
            flow::Transition::None,
        ))
    }

    async fn store(
        self,
        sequence_set: &SequenceSet,
        kind: &StoreType,
        response: &StoreResponse,
        flags: &[Flag<'a>],
        modifiers: &[StoreModifier],
        uid: &bool,
    ) -> Result<(Response<'static>, flow::Transition)> {
        let mut unchanged_since: Option<NonZeroU64> = None;
        modifiers.iter().for_each(|m| match m {
            StoreModifier::UnchangedSince(val) => {
                unchanged_since = Some(*val);
            },
        });

        let (data, modified) = self
            .mailbox
            .store(sequence_set, kind, response, flags, unchanged_since, uid)
            .await?;

        let mut ok_resp = Response::build()
                .to_req(self.req)
                .message("STORE completed")
                .set_body(data);


        match modified[..] {
            [] => (),
            [_head, ..] => {
                let modified_str = format!("MODIFIED {}", modified.into_iter().map(|x| x.to_string()).collect::<Vec<_>>().join(","));
                ok_resp = ok_resp.code(Code::Other(CodeOther::unvalidated(modified_str.into_bytes())));
            },
        };


        self.client_capabilities.store_modifiers_enable(modifiers);

        Ok((ok_resp.ok()?,
            flow::Transition::None,
        ))
    }

    async fn copy(
        self,
        sequence_set: &SequenceSet,
        mailbox: &MailboxCodec<'a>,
        uid: &bool,
    ) -> Result<(Response<'static>, flow::Transition)> {
        let name: &str = MailboxName(mailbox).try_into()?;

        let mb_opt = self.user.open_mailbox(&name).await?;
        let mb = match mb_opt {
            Some(mb) => mb,
            None => {
                return Ok((
                    Response::build()
                        .to_req(self.req)
                        .message("Destination mailbox does not exist")
                        .code(Code::TryCreate)
                        .no()?,
                    flow::Transition::None,
                ))
            }
        };

        let (uidval, uid_map) = self.mailbox.copy(sequence_set, mb, uid).await?;

        let copyuid_str = format!(
            "{} {} {}",
            uidval,
            uid_map
                .iter()
                .map(|(sid, _)| format!("{}", sid))
                .collect::<Vec<_>>()
                .join(","),
            uid_map
                .iter()
                .map(|(_, tuid)| format!("{}", tuid))
                .collect::<Vec<_>>()
                .join(",")
        );

        Ok((
            Response::build()
                .to_req(self.req)
                .message("COPY completed")
                .code(Code::Other(CodeOther::unvalidated(
                    format!("COPYUID {}", copyuid_str).into_bytes(),
                )))
                .ok()?,
            flow::Transition::None,
        ))
    }

    async fn r#move(
        self,
        sequence_set: &SequenceSet,
        mailbox: &MailboxCodec<'a>,
        uid: &bool,
    ) -> Result<(Response<'static>, flow::Transition)> {
        let name: &str = MailboxName(mailbox).try_into()?;

        let mb_opt = self.user.open_mailbox(&name).await?;
        let mb = match mb_opt {
            Some(mb) => mb,
            None => {
                return Ok((
                    Response::build()
                        .to_req(self.req)
                        .message("Destination mailbox does not exist")
                        .code(Code::TryCreate)
                        .no()?,
                    flow::Transition::None,
                ))
            }
        };

        let (uidval, uid_map, data) = self.mailbox.r#move(sequence_set, mb, uid).await?;

        // compute code
        let copyuid_str = format!(
            "{} {} {}",
            uidval,
            uid_map
                .iter()
                .map(|(sid, _)| format!("{}", sid))
                .collect::<Vec<_>>()
                .join(","),
            uid_map
                .iter()
                .map(|(_, tuid)| format!("{}", tuid))
                .collect::<Vec<_>>()
                .join(",")
        );

        Ok((
            Response::build()
                .to_req(self.req)
                .message("COPY completed")
                .code(Code::Other(CodeOther::unvalidated(
                    format!("COPYUID {}", copyuid_str).into_bytes(),
                )))
                .set_body(data)
                .ok()?,
            flow::Transition::None,
        ))
    }
}