macltok

use std::sync::Arc;

#[cfg(feature = "sqlite")]

use std::{collections::BTreeSet, path::PathBuf};

use ruma::TransactionId;

use crate::{

    AccountError, EncryptedSessionStore, OAuthLogin, PasswordLogin, PendingOAuthLogin, Restore,

    Session, SessionRecord, SessionStore, SoftLogout, StoreConfig, account::homeserver_input,

    engine::StoreLock,

};

/// every account one owner has, a store each under one root and one passphrase. whoever holds

/// the set can resume any soft-logged-out device in it with just its password, and the device's

/// store has its e2ee keys, so a set belongs to one person and never to a whole machine

pub struct Accounts {

    root: Root,

    passphrase: String,

}

enum Root {

    #[cfg(feature = "sqlite")]

    Directory(PathBuf),

    #[cfg(feature = "indexeddb")]

    Prefix(String),

}

/// one account's record and store

#[derive(Clone)]

pub struct Account {

    id: String,

    sessions: Arc<EncryptedSessionStore>,

}

impl Account {

    /// stable for the account's life, and safe in a cookie or a url

    #[must_use]

    pub fn id(&self) -> &str {

        &self.id

    }

    /// for `Termination::retain`. a discarded account drops out of `Accounts::list`

    #[must_use]

    pub fn sessions(&self) -> Arc<dyn SessionStore> {

        self.sessions.clone()

    }

    pub async fn restore(&self) -> Result<Option<Restore>, AccountError> {

        Session::restore(self.sessions()).await

    }

}

impl std::fmt::Debug for Account {

    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {

        formatter

            .debug_struct("Account")

            .field("id", &self.id)

            .finish_non_exhaustive()

    }

}

enum Found {

    Nobody,

    SoftLoggedOut(Account, SoftLogout),

}

impl Accounts {

    /// each account is a store directory under `root` with its record beside it

    #[cfg(feature = "sqlite")]

    #[must_use]

    pub fn sqlite(root: impl Into<PathBuf>, passphrase: impl Into<String>) -> Self {

        Self {

            root: Root::Directory(root.into()),

            passphrase: passphrase.into(),

        }

    }

    /// each account is a store named `<prefix>::<id>`, listed in `<prefix>::macltok-accounts`

    #[cfg(feature = "indexeddb")]

    #[must_use]

    pub fn indexeddb(prefix: impl Into<String>, passphrase: impl Into<String>) -> Self {

        Self {

            root: Root::Prefix(prefix.into()),

            passphrase: passphrase.into(),

        }

    }

    /// `None` for anything `Account::id` couldn't have been

    #[must_use]

    pub fn account(&self, id: &str) -> Option<Account> {

        let valid =

            !id.is_empty() && id.len() <= 64 && id.bytes().all(|b| b.is_ascii_alphanumeric());

        valid.then(|| self.open(id.to_owned()))

    }

    /// every account with a record. a store left without one, by a login that died halfway or a

    /// discard that couldn't finish, is deleted on the way past

    pub async fn list(&self) -> Result<Vec<Account>, AccountError> {

        let mut accounts = Vec::new();

        for account in self.slots().await? {

            match account.sessions.load().await {

                Ok(None) => self.sweep(&account).await,

                Ok(Some(_)) | Err(_) => accounts.push(account),

            }

        }

        Ok(accounts)

    }

    /// a soft-logged-out account the login names is resumed as its own device, keeping its keys.

    /// anyone else gets a new account

    pub async fn login_password(

        &self,

        login: PasswordLogin,

    ) -> Result<(Account, Session), AccountError> {

        homeserver_input(&login.homeserver, login.allow_insecure_loopback)?;

        let found = self

            .find(|record| record_names(record, &login.homeserver, &login.username))

            .await?;

        match found {

            Found::SoftLoggedOut(account, soft) => {

                let store_lock = account.sessions.store().lock().await?;

                let login = PasswordLogin {

                    device_display_name: login.device_display_name,

                    ..soft.password_login(login.password)

                };

                let session =

                    Session::login_password_locked(login, account.sessions(), store_lock).await?;

                Ok((account, session))

            }

            Found::Nobody => {

                let (account, store_lock) = self.fresh().await?;

                let login = PasswordLogin {

                    device_id: None,

                    ..login

                };

                match Session::login_password_locked(login, account.sessions(), store_lock).await {

                    Ok(session) => Ok((account, session)),

                    Err(error) => Err(self.abandon(&account, error).await),

                }

            }

        }

    }

    /// resumes the soft-logged-out account whose device `login` names, as

    /// `SoftLogout::oauth_login` sets it. without one it's a new account, since which account an

    /// OAuth login is for isn't known until it finishes

    pub async fn begin_oauth(

        &self,

        login: OAuthLogin,

    ) -> Result<(Account, PendingOAuthLogin), AccountError> {

        homeserver_input(&login.homeserver, login.allow_insecure_loopback)?;

        let device_id = login.device_id.clone();

        let found = self

            .find(|record| {

                record

                    .soft_logout()

                    .is_some_and(|soft| device_id.as_ref() == Some(&soft.device_id))

            })

            .await?;

        match found {

            Found::SoftLoggedOut(account, soft) => {

                let store_lock = account.sessions.store().lock().await?;

                let login = OAuthLogin {

                    homeserver: soft.homeserver,

                    client_id: login.client_id.or(soft.client_id),

                    allow_insecure_loopback: true,

                    ..login

                };

                let pending =

                    Session::begin_oauth_locked(login, account.sessions(), store_lock).await?;

                Ok((account, pending))

            }

            Found::Nobody => {

                let (account, store_lock) = self.fresh().await?;

                let login = OAuthLogin {

                    device_id: None,

                    ..login

                };

                match Session::begin_oauth_locked(login, account.sessions(), store_lock).await {

                    Ok(pending) => Ok((account, pending)),

                    Err(error) => Err(self.abandon(&account, error).await),

                }

            }

        }

    }

    /// a live login for the same person is refused rather than given a second device

    async fn find(&self, resumes: impl Fn(&SessionRecord) -> bool) -> Result<Found, AccountError> {

        for account in self.slots().await? {

            let Ok(Some(record)) = account.sessions.load().await else {

                continue;

            };

            if !resumes(&record) {

                continue;

            }

            return match record.soft_logout() {

                Some(soft) => Ok(Found::SoftLoggedOut(account, soft)),

                None => Err(AccountError::AlreadyLoggedIn(record.user_id().clone())),

            };

        }

        Ok(Found::Nobody)

    }

    /// locked before anything else can see it, so `list` never sweeps a login still under way

    async fn fresh(&self) -> Result<(Account, StoreLock), AccountError> {

        let account = self.open(TransactionId::new().to_string());

        let store_lock = account.sessions.store().lock().await?;

        #[cfg(feature = "indexeddb")]

        {

            let Root::Prefix(prefix) = &self.root;

            index::add(prefix, &account.id).await?;

        }

        Ok((account, store_lock))

    }

    /// a failed first login leaves nothing behind. one whose record didn't save has a live device

    /// on the server, so that stays for `UnsavedSession::save`

    async fn abandon(&self, account: &Account, error: AccountError) -> AccountError {

        if !matches!(error, AccountError::Unsaved { .. }) {

            self.sweep(account).await;

        }

        error

    }

    /// best effort, since whatever it leaves is found and swept again by the next `list`

    async fn sweep(&self, account: &Account) {

        let Ok(_store_lock) = account.sessions.store().lock().await else {

            return;

        };

        if !matches!(account.sessions.load().await, Ok(None)) {

            return;

        }

        if account.sessions.discard().await.is_err() {

            return;

        }

        #[cfg(feature = "indexeddb")]

        {

            let Root::Prefix(prefix) = &self.root;

            let _ = index::remove(prefix, &account.id).await;

        }

    }

    fn open(&self, id: String) -> Account {

        let sessions = match &self.root {

            #[cfg(feature = "sqlite")]

            Root::Directory(root) => EncryptedSessionStore::file(

                root.join(format!("{id}.session")),

                StoreConfig::sqlite(root.join(&id), self.passphrase.clone()),

            ),

            #[cfg(feature = "indexeddb")]

            Root::Prefix(prefix) => EncryptedSessionStore::indexeddb(StoreConfig::indexeddb(

                format!("{prefix}::{id}"),

                self.passphrase.clone(),

            )),

        };

        Account {

            id,

            sessions: Arc::new(sessions),

        }

    }

    async fn slots(&self) -> Result<Vec<Account>, AccountError> {

        let ids = match &self.root {

            #[cfg(feature = "sqlite")]

            Root::Directory(root) => directory_slots(root)?,

            #[cfg(feature = "indexeddb")]

            Root::Prefix(prefix) => index::list(prefix).await?,

        };

        Ok(ids.iter().filter_map(|id| self.account(id)).collect())

    }

}

/// a localpart or a full user id names the same account, but only on the same server. the server

/// is compared by name as well as by url, so a delegated `example.org` still finds the account

/// discovery took to `matrix.example.org`

fn record_names(record: &SessionRecord, homeserver: &str, username: &str) -> bool {

    let user_id = record.user_id();

    let username = username.trim();

    let user = if username.starts_with('@') {

        username.eq_ignore_ascii_case(user_id.as_str())

    } else {

        username.eq_ignore_ascii_case(user_id.localpart())

    };

    let typed = authority(homeserver);

    user && (typed.eq_ignore_ascii_case(user_id.server_name().as_str())

        || typed.eq_ignore_ascii_case(authority(record.homeserver())))

}

fn authority(homeserver: &str) -> &str {

    let homeserver = homeserver.trim();

    let rest = homeserver

        .split_once("://")

        .map_or(homeserver, |(_, rest)| rest);

    rest.split('/').next().unwrap_or(rest)

}

#[cfg(feature = "sqlite")]

fn directory_slots(root: &std::path::Path) -> Result<BTreeSet<String>, AccountError> {

    use std::io::ErrorKind;

    let entries = match std::fs::read_dir(root) {

        Ok(entries) => entries,

        Err(error) if error.kind() == ErrorKind::NotFound => return Ok(BTreeSet::new()),

        Err(error) => return Err(crate::account::account_error("listing the accounts", error)),

    };

    let mut ids = BTreeSet::new();

    for entry in entries {

        let entry =

            entry.map_err(|error| crate::account::account_error("listing the accounts", error))?;

        let name = entry.file_name();

        let Some(name) = name.to_str() else {

            continue;

        };

        if entry.path().is_dir() {

            ids.insert(name.to_owned());

        } else if let Some(id) = name.strip_suffix(".session") {

            ids.insert(id.to_owned());

        }

    }

    Ok(ids)

}

/// a key per account, so two tabs adding at once can't lose one to the other

#[cfg(feature = "indexeddb")]

mod index {

    use std::collections::BTreeSet;

    use indexed_db_futures::{

        Build,

        database::Database,

        error::{Error, OpenDbError},

        prelude::*,

        transaction::{Transaction, TransactionMode},

    };

    use crate::{AccountError, account::account_error};

    const ACCOUNTS: &str = "accounts";

    fn name(prefix: &str) -> String {

        format!("{prefix}::macltok-accounts")

    }

    async fn open(prefix: &str) -> Result<Database, OpenDbError> {

        Database::open(&name(prefix))

            .with_version(1u32)

            .with_on_upgrade_needed(|_, transaction: &Transaction<'_>| -> Result<(), Error> {

                transaction.db().create_object_store(ACCOUNTS).build()?;

                Ok(())

            })

            .build()?

            .await

    }

    pub(super) async fn list(prefix: &str) -> Result<BTreeSet<String>, AccountError> {

        let failed = |error| account_error("listing the accounts", error);

        let database = open(prefix).await.map_err(failed)?;

        let ids = keys(&database).await;

        database.close();

        ids.map_err(|error| account_error("listing the accounts", error))

    }

    pub(super) async fn add(prefix: &str, id: &str) -> Result<(), AccountError> {

        let failed = |error| account_error("adding the account", error);

        let database = open(prefix).await.map_err(failed)?;

        let added = write(&database, |store| {

            store.put(true).with_key(id).build()?;

            Ok(())

        })

        .await;

        database.close();

        added.map_err(|error| account_error("adding the account", error))

    }

    pub(super) async fn remove(prefix: &str, id: &str) -> Result<(), AccountError> {

        let failed = |error| account_error("removing the account", error);

        let database = open(prefix).await.map_err(failed)?;

        let removed = write(&database, |store| {

            store.delete(id).build()?;

            Ok(())

        })

        .await;

        database.close();

        removed.map_err(|error| account_error("removing the account", error))

    }

    async fn keys(database: &Database) -> Result<BTreeSet<String>, Error> {

        let transaction = database

            .transaction(ACCOUNTS)

            .with_mode(TransactionMode::Readonly)

            .build()?;

        let store = transaction.object_store(ACCOUNTS)?;

        store.get_all_keys::<String>().await?.collect()

    }

    async fn write(

        database: &Database,

        change: impl FnOnce(&indexed_db_futures::object_store::ObjectStore<'_>) -> Result<(), Error>,

    ) -> Result<(), Error> {

        let transaction = database

            .transaction(ACCOUNTS)

            .with_mode(TransactionMode::Readwrite)

            .build()?;

        change(&transaction.object_store(ACCOUNTS)?)?;

        transaction.commit().await

    }

}