aerogramme/src/mail/user.rs

use std::collections::{HashMap, BTreeMap};
use std::sync::{Arc, Weak};

use anyhow::{Result, bail};
use lazy_static::lazy_static;
use serde::{Serialize, Deserialize};
use k2v_client::{K2vClient, CausalityToken, K2vValue};
use tokio::sync::watch;

use crate::cryptoblob::{seal_serialize, open_deserialize};
use crate::login::{Credentials, StorageCredentials};
use crate::mail::mailbox::Mailbox;
use crate::mail::unique_ident::{UniqueIdent, gen_ident};
use crate::mail::uidindex::ImapUidvalidity;
use crate::time::now_msec;

const MAILBOX_HIERARCHY_DELIMITER: &str = "/";

/// INBOX is the only mailbox that must always exist.
/// It is created automatically when the account is created.
/// IMAP allows the user to rename INBOX to something else,
/// in this case all messages from INBOX are moved to a mailbox
/// with the new name and the INBOX mailbox still exists and is empty.
/// In our implementation, we indeed move the underlying mailbox
/// to the new name (i.e. the new name has the same id as the previous
/// INBOX), and we create a new empty mailbox for INBOX.
const INBOX: &str = "INBOX";

const MAILBOX_LIST_PK: &str = "mailboxes";
const MAILBOX_LIST_SK: &str = "list";

pub struct User {
    pub username: String,
    pub creds: Credentials,
    pub k2v: K2vClient,
    pub mailboxes: std::sync::Mutex<HashMap<UniqueIdent, Weak<Mailbox>>>,

    tx_inbox_id: watch::Sender<Option<(UniqueIdent, ImapUidvalidity)>>,
}

impl User {
    pub async fn new(username: String, creds: Credentials) -> Result<Arc<Self>> {
        let cache_key = (username.clone(), creds.storage.clone());

        {
            let cache = USER_CACHE.lock().unwrap();
            if let Some(u) = cache.get(&cache_key).and_then(Weak::upgrade) {
                return Ok(u);
            }
        }

        let user = Self::open(username, creds).await?;

        let mut cache = USER_CACHE.lock().unwrap();
        if let Some(concurrent_user) = cache.get(&cache_key).and_then(Weak::upgrade) {
            drop(user);
            Ok(concurrent_user)
        } else {
            cache.insert(cache_key, Arc::downgrade(&user));
            Ok(user)
        }
    }

    /// Lists user's available mailboxes
    pub async fn list_mailboxes(&self) -> Result<Vec<String>> {
        let (list, _ct) = self.load_mailbox_list().await?;
        Ok(list.into_iter().map(|(k, _)| k).collect())
    }

    /// Opens an existing mailbox given its IMAP name.
    pub async fn open_mailbox(&self, name: &str) -> Result<Option<Arc<Mailbox>>> {
        let (list, _ct) = self.load_mailbox_list().await?;
        match list.get(name) {
            Some(MailboxListEntry { id_lww: (_, Some(mbid)), uidvalidity }) =>
                self.open_mailbox_by_id(*mbid, *uidvalidity).await,
            _ =>
                bail!("Mailbox does not exist: {}", name),
        }
    }

    /// Creates a new mailbox in the user's IMAP namespace.
    pub fn create_mailbox(&self, _name: &str) -> Result<()> {
        unimplemented!()
    }

    /// Deletes a mailbox in the user's IMAP namespace.
    pub fn delete_mailbox(&self, _name: &str) -> Result<()> {
        unimplemented!()
    }

    /// Renames a mailbox in the user's IMAP namespace.
    pub fn rename_mailbox(&self, _old_name: &str, _new_name: &str) -> Result<()> {
        unimplemented!()
    }

    // ---- Internal user & mailbox management ----
    
    async fn open(username: String, creds: Credentials) -> Result<Arc<Self>> {
        let k2v = creds.k2v_client()?;

        let (tx_inbox_id, rx_inbox_id) = watch::channel(None);

        let user = Arc::new(Self {
            username,
            creds,
            k2v,
            tx_inbox_id,
            mailboxes: std::sync::Mutex::new(HashMap::new()),
        });

        // Ensure INBOX exists (done inside load_mailbox_list)
        user.load_mailbox_list().await?;

        Ok(user)
    }

    async fn open_mailbox_by_id(&self, id: UniqueIdent, min_uidvalidity: ImapUidvalidity) -> Result<Option<Arc<Mailbox>>> {
        {
            let cache = self.mailboxes.lock().unwrap();
            if let Some(mb) = cache.get(&id).and_then(Weak::upgrade) {
                return Ok(Some(mb));
            }
        }

        let mb = Arc::new(Mailbox::open(&self.creds, id, min_uidvalidity).await?);

        let mut cache = self.mailboxes.lock().unwrap();
        if let Some(concurrent_mb) = cache.get(&id).and_then(Weak::upgrade) {
            drop(mb); // we worked for nothing but at least we didn't starve someone else
            Ok(Some(concurrent_mb))
        } else {
            cache.insert(id, Arc::downgrade(&mb));
            Ok(Some(mb))
        }
    }

    // ---- Mailbox list management ----

    async fn load_mailbox_list(&self) -> Result<(MailboxList, Option<CausalityToken>)> {
        let cv = match self.k2v.read_item(MAILBOX_LIST_PK, MAILBOX_LIST_SK).await {
            Err(k2v_client::Error::NotFound) => return Ok((BTreeMap::new(), None)),
            Err(e) => return Err(e.into()),
            Ok(cv) => cv,
        };

        let mut list = BTreeMap::new();
        for v in cv.value {
            if let K2vValue::Value(vbytes) = v {
                let list2 = open_deserialize::<MailboxList>(&vbytes, &self.creds.keys.master)?;
                list = merge_mailbox_lists(list, list2);
            }
        }

        // If INBOX doesn't exist, create a new mailbox with that name
        // and save new mailbox list.
        // Also, ensure that the mpsc::watch that keeps track of the
        // inbox id is up-to-date.
        let (inbox_id, inbox_uidvalidity) = match list.get_mut(INBOX) {
            None => {
                let (id, uidvalidity) = (gen_ident(), ImapUidvalidity::new(1).unwrap());
                list.insert(INBOX.into(), MailboxListEntry {
                    id_lww: (now_msec(), Some(id)),
                    uidvalidity,
                });
                self.save_mailbox_list(&list, Some(cv.causality.clone())).await?;
                (id, uidvalidity)
            }
            Some(MailboxListEntry { id_lww: id_lww @ (_, None), uidvalidity }) => {
                let id = gen_ident();
                id_lww.0 = std::cmp::max(id_lww.0 + 1, now_msec());
                id_lww.1 = Some(id);
                *uidvalidity = ImapUidvalidity::new(uidvalidity.get() + 1).unwrap();
                let uidvalidity = *uidvalidity;
                self.save_mailbox_list(&list, Some(cv.causality.clone())).await?;
                (id, uidvalidity)
            }
            Some(MailboxListEntry { id_lww: (_, Some(id)), uidvalidity }) => {
                (*id, *uidvalidity)
            }
        };
        self.tx_inbox_id.send(Some((inbox_id, inbox_uidvalidity))).unwrap();

        Ok((list, Some(cv.causality)))
    }

    async fn save_mailbox_list(&self, list: &MailboxList, ct: Option<CausalityToken>) -> Result<()> {
        let list_blob = seal_serialize(list, &self.creds.keys.master)?;
        self.k2v.insert_item(MAILBOX_LIST_PK, MAILBOX_LIST_SK, list_blob, ct).await?;
        Ok(())
    }
}

// ---- User's mailbox list (serialized in K2V) ----

type MailboxList = BTreeMap<String, MailboxListEntry>;

#[derive(Serialize, Deserialize, Clone, Copy)]
struct MailboxListEntry {
    id_lww: (u64, Option<UniqueIdent>),
    uidvalidity: ImapUidvalidity,
}

impl MailboxListEntry {
    fn merge(&mut self, other: &Self) {
        // Simple CRDT merge rule
        if other.id_lww.0 > self.id_lww.0
        || (other.id_lww.0 == self.id_lww.0 && other.id_lww.1 > self.id_lww.1) {
            self.id_lww = other.id_lww;
        }
        self.uidvalidity = std::cmp::max(self.uidvalidity, other.uidvalidity);
    }
}

fn merge_mailbox_lists(mut list1: MailboxList, list2: MailboxList) -> MailboxList {
    for (k, v) in list2.into_iter() {
        if let Some(e) = list1.get_mut(&k) {
            e.merge(&v);
        } else {
            list1.insert(k, v);
        }
    }
    list1
}

// ---- User cache ----

lazy_static! {
    static ref USER_CACHE: std::sync::Mutex<HashMap<(String, StorageCredentials), Weak<User>>> = std::sync::Mutex::new(HashMap::new());
}