macltok

#[cfg(feature = "sqlite")]

use std::path::PathBuf;

#[cfg(not(target_family = "wasm"))]

use std::sync::atomic::{AtomicBool, Ordering};

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

use matrix_sdk::{

    Client, ClientBuilder, SendOutsideWasm, SessionChange, SessionMeta, SyncOutsideWasm,

    ThreadingSupport,

    authentication::{

        AuthSession,

        matrix::MatrixSession,

        oauth::{

            ClientId, ClientRegistrationData, CsrfToken, OAuthSession, PendingAuthorization,

            UserSession,

            registration::{ApplicationType, ClientMetadata, Localized, OAuthGrantType},

        },

    },

    cross_process_lock::CrossProcessLockConfig,

    encryption::{BackupDownloadStrategy, EncryptionSettings},

    ruma::{MilliSecondsSinceUnixEpoch, OwnedDeviceId, OwnedUserId, serde::Raw},

};

use serde::{Deserialize, Serialize};

use thiserror::Error;

use tokio::sync::{broadcast, mpsc};

use url::Url;

use crate::{

    ClientHandle, Engine, EngineConfig, EngineError, ErrorKind, EventStream, ServerError,

    SessionStore,

    engine::{Persistence, StoreLock},

    error::SourceError,

};

/// a logged in client and the `SessionStore` its record lives in. the engine it builds

/// saves every token rotation there, unless told `without_persistence`

pub struct Session {

    client: Client,

    tracking: Tracking,

    store_lock: StoreLock,

}

/// taken as soon as the client exists, since the sdk can refresh before there's an engine

struct Tracking {

    store: Arc<dyn SessionStore>,

    session_changes: broadcast::Receiver<SessionChange>,

    persistence: Option<Persistence>,

    #[cfg(not(target_family = "wasm"))]

    inline: Arc<AtomicBool>,

}

#[non_exhaustive]

pub struct PasswordLogin {

    pub homeserver: String,

    pub username: String,

    pub password: String,

    pub device_display_name: Option<String>,

    /// set after a soft logout so the store's keys still belong to us

    pub device_id: Option<OwnedDeviceId>,

    /// let plain http through, for loopback hosts only

    pub allow_insecure_loopback: bool,

}

impl PasswordLogin {

    #[must_use]

    pub fn new(

        homeserver: impl Into<String>,

        username: impl Into<String>,

        password: impl Into<String>,

    ) -> Self {

        Self {

            homeserver: homeserver.into(),

            username: username.into(),

            password: password.into(),

            device_display_name: None,

            device_id: None,

            allow_insecure_loopback: false,

        }

    }

    #[must_use]

    pub fn with_device_display_name(mut self, name: impl Into<String>) -> Self {

        self.device_display_name = Some(name.into());

        self

    }

    #[must_use]

    pub fn with_device_id(mut self, device_id: OwnedDeviceId) -> Self {

        self.device_id = Some(device_id);

        self

    }

    #[must_use]

    pub fn allow_insecure_loopback(mut self, allow: bool) -> Self {

        self.allow_insecure_loopback = allow;

        self

    }

}

#[non_exhaustive]

pub struct OAuthLogin {

    pub homeserver: String,

    pub redirect_uri: String,

    pub client_uri: String,

    pub client_name: Option<String>,

    pub client_id: Option<String>,

    pub device_id: Option<OwnedDeviceId>,

    pub application_type: OAuthApplicationType,

    /// let plain http through, for loopback hosts only

    pub allow_insecure_loopback: bool,

}

impl OAuthLogin {

    #[must_use]

    pub fn new(

        homeserver: impl Into<String>,

        redirect_uri: impl Into<String>,

        client_uri: impl Into<String>,

        application_type: OAuthApplicationType,

    ) -> Self {

        Self {

            homeserver: homeserver.into(),

            redirect_uri: redirect_uri.into(),

            client_uri: client_uri.into(),

            client_name: None,

            client_id: None,

            device_id: None,

            application_type,

            allow_insecure_loopback: false,

        }

    }

    #[must_use]

    pub fn with_client_name(mut self, name: impl Into<String>) -> Self {

        self.client_name = Some(name.into());

        self

    }

    /// a statically registered client id, skipping dynamic registration

    #[must_use]

    pub fn with_client_id(mut self, client_id: impl Into<String>) -> Self {

        self.client_id = Some(client_id.into());

        self

    }

    #[must_use]

    pub fn with_device_id(mut self, device_id: OwnedDeviceId) -> Self {

        self.device_id = Some(device_id);

        self

    }

    #[must_use]

    pub fn allow_insecure_loopback(mut self, allow: bool) -> Self {

        self.allow_insecure_loopback = allow;

        self

    }

}

#[derive(Clone, Copy, Debug, Eq, PartialEq)]

#[non_exhaustive]

pub enum OAuthApplicationType {

    Native,

    Web,

}

pub struct PendingOAuthLogin {

    client: Client,

    sessions: Arc<dyn SessionStore>,

    tracking: Option<(Tracking, StoreLock)>,

    authorization_url: String,

    state: Option<CsrfToken>,

    previous: Option<SessionRecord>,

}

impl PendingOAuthLogin {

    #[must_use]

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

        &self.authorization_url

    }

    #[must_use]

    pub fn registered_client_id(&self) -> String {

        self.client

            .oauth()

            .client_id()

            .expect("OAuth client was registered before authorization")

            .as_str()

            .to_owned()

    }

    #[must_use]

    pub const fn is_pending(&self) -> bool {

        self.state.is_some()

    }

    pub async fn finish(&mut self, callback_url: &str) -> Result<Session, AccountError> {

        let callback_url = Url::parse(callback_url)

            .map_err(|error| account_error("parsing the OAuth callback URL", error))?;

        let state = self

            .state

            .take()

            .ok_or(AccountError::OAuthCallbackAlreadySubmitted)?;

        let oauth = self.client.oauth();

        if let Err(error) = oauth.finish_login(callback_url.into()).await {

            oauth.abort_login(&state).await;

            if let Some(previous) = self.previous.take() {

                put_back(self.sessions.as_ref(), previous).await;

            }

            return Err(account_error("OAuth login", error));

        }

        let (tracking, store_lock) = self

            .tracking

            .take()

            .ok_or(AccountError::OAuthCallbackAlreadySubmitted)?;

        Session::saved(

            self.client.clone(),

            self.sessions.clone(),

            tracking,

            store_lock,

        )

        .await

    }

    /// puts back what the store held before, or deletes the store if it held nothing

    pub async fn abort(self) {

        let Self {

            client,

            sessions,

            tracking,

            state,

            previous,

            ..

        } = self;

        let Some(state) = state else {

            return;

        };

        client.oauth().abort_login(&state).await;

        drop(client);

        match previous {

            Some(previous) => put_back(sessions.as_ref(), previous).await,

            None => {

                if let Err(error) = sessions.discard().await {

                    tracing::warn!(%error, "deleting an abandoned OAuth login's store failed");

                }

            }

        }

        drop(tracking);

    }

}

async fn put_back(sessions: &dyn SessionStore, previous: SessionRecord) {

    if let Err(error) = sessions.save(&previous).await {

        tracing::warn!(%error, "putting back the record an OAuth login replaced failed");

    }

}

pub struct StoreConfig {

    location: StoreLocation,

    passphrase: String,

}

enum StoreLocation {

    #[cfg(feature = "sqlite")]

    Sqlite(PathBuf),

    #[cfg(feature = "indexeddb")]

    IndexedDb(String),

}

impl StoreConfig {

    #[cfg(feature = "sqlite")]

    #[must_use]

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

        Self {

            location: StoreLocation::Sqlite(path.into()),

            passphrase: passphrase.into(),

        }

    }

    #[cfg(feature = "indexeddb")]

    #[must_use]

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

        Self {

            location: StoreLocation::IndexedDb(name.into()),

            passphrase: passphrase.into(),

        }

    }

    /// delete the store's databases, once nothing has them open

    pub async fn delete(&self) -> Result<(), AccountError> {

        match &self.location {

            #[cfg(feature = "sqlite")]

            StoreLocation::Sqlite(path) => crate::session_store::delete_sqlite_store(path)

                .map_err(|error| account_error("deleting the Matrix store", error)),

            #[cfg(feature = "indexeddb")]

            StoreLocation::IndexedDb(name) => crate::session_store::delete_indexeddb_store(name)

                .await

                .map_err(|error| error.into_account_error("deleting the Matrix store")),

        }

    }

    /// taken before the client opens the store, and held until the engine and every

    /// `ClientHandle` are gone

    pub(crate) async fn lock(&self) -> Result<StoreLock, AccountError> {

        match &self.location {

            #[cfg(feature = "sqlite")]

            StoreLocation::Sqlite(path) => StoreLock::file(path),

            #[cfg(feature = "indexeddb")]

            StoreLocation::IndexedDb(name) => StoreLock::web(name).await,

        }

    }

    pub(crate) fn passphrase(&self) -> &str {

        &self.passphrase

    }

    #[cfg(feature = "sqlite")]

    pub(crate) fn mentions_location(&self) -> PathBuf {

        let StoreLocation::Sqlite(path) = &self.location;

        path.join("macltok-mentions.sqlite3")

    }

    #[cfg(feature = "indexeddb")]

    pub(crate) fn mentions_location(&self) -> String {

        format!("{}::macltok-mentions", self.indexeddb_name())

    }

    #[cfg(feature = "indexeddb")]

    pub(crate) fn indexeddb_name(&self) -> &str {

        let StoreLocation::IndexedDb(name) = &self.location;

        name

    }

}

/// Serializable authentication secrets. A `SessionStore` keeps it, and the engine

/// replaces it there on every token rotation.

#[derive(Clone, Deserialize, Serialize)]

pub struct SessionRecord {

    homeserver: String,

    authentication: AuthenticationRecord,

}

#[derive(Clone, Deserialize, Serialize)]

#[serde(tag = "type", content = "session", rename_all = "snake_case")]

enum AuthenticationRecord {

    Matrix(MatrixSession),

    OAuth {

        client_id: ClientId,

        user: UserSession,

    },

    SoftLoggedOut {

        meta: SessionMeta,

        client_id: Option<ClientId>,

    },

    /// an oauth login waiting on its redirect. `authorization` holds the pkce verifier

    PendingOAuth {

        client_id: ClientId,

        state: CsrfToken,

        authorization: Box<PendingAuthorization>,

        started_ms: u64,

        /// what the store held before, put back if the login never finishes

        previous: Option<Box<SessionRecord>>,

    },

}

/// authorization codes live for minutes, so a login this old was abandoned

const OAUTH_PENDING_FOR_MS: u64 = 60 * 60 * 1000;

/// the authorization server answers in the query, or the fragment when asked to

pub(crate) fn callback_state(callback: &Url) -> Option<String> {

    let query = callback.query_pairs();

    let fragment = url::form_urlencoded::parse(callback.fragment().unwrap_or_default().as_bytes());

    query

        .chain(fragment)

        .find(|(key, _)| key == "state")

        .map(|(_, state)| state.into_owned())

}

fn now_ms() -> u64 {

    MilliSecondsSinceUnixEpoch::now().get().into()

}

impl SessionRecord {

    #[must_use]

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

        &self.homeserver

    }

    /// a pending oauth login is nobody yet, unless it resumes someone

    pub(crate) fn user_id(&self) -> Option<&OwnedUserId> {

        match &self.authentication {

            AuthenticationRecord::Matrix(session) => Some(&session.meta.user_id),

            AuthenticationRecord::OAuth { user, .. } => Some(&user.meta.user_id),

            AuthenticationRecord::SoftLoggedOut { meta, .. } => Some(&meta.user_id),

            AuthenticationRecord::PendingOAuth { previous, .. } => previous.as_ref()?.user_id(),

        }

    }

    /// what a login should resume. a pending oauth login stands in for the record before it

    pub(crate) fn resumable(self) -> Option<Self> {

        match self.authentication {

            AuthenticationRecord::PendingOAuth { previous, .. } => {

                previous.and_then(|previous| previous.resumable())

            }

            _ => Some(self),

        }

    }

    pub(crate) fn pending_oauth_state(&self) -> Option<&str> {

        match &self.authentication {

            AuthenticationRecord::PendingOAuth { state, .. } => Some(state.secret()),

            _ => None,

        }

    }

    pub(crate) fn oauth_expired(&self) -> bool {

        let AuthenticationRecord::PendingOAuth { started_ms, .. } = &self.authentication else {

            return false;

        };

        now_ms().saturating_sub(*started_ms) > OAUTH_PENDING_FOR_MS

    }

    pub(crate) fn soft_logout(&self) -> Option<SoftLogout> {

        let (meta, client_id) = match &self.authentication {

            AuthenticationRecord::SoftLoggedOut { meta, client_id } => (meta, client_id),

            AuthenticationRecord::PendingOAuth { previous, .. } => {

                return previous.as_ref()?.soft_logout();

            }

            _ => return None,

        };

        Some(SoftLogout {

            homeserver: self.homeserver.clone(),

            user_id: meta.user_id.clone(),

            device_id: meta.device_id.clone(),

            client_id: client_id

                .as_ref()

                .map(|client_id| client_id.as_str().to_owned()),

        })

    }

    pub(crate) fn soft_logged_out(self) -> Self {

        let authentication = match self.authentication {

            AuthenticationRecord::Matrix(session) => AuthenticationRecord::SoftLoggedOut {

                meta: session.meta,

                client_id: None,

            },

            AuthenticationRecord::OAuth { client_id, user } => {

                AuthenticationRecord::SoftLoggedOut {

                    meta: user.meta,

                    client_id: Some(client_id),

                }

            }

            soft @ AuthenticationRecord::SoftLoggedOut { .. } => soft,

            pending @ AuthenticationRecord::PendingOAuth { .. } => pending,

        };

        Self {

            homeserver: self.homeserver,

            authentication,

        }

    }

}

pub enum Restore {

    Resumed(Session),

    SoftLoggedOut(SoftLogout),

    /// an oauth login is waiting on its redirect. `Session::finish_oauth` takes the callback

    OAuthPending,

}

/// who to log in again as. the store keeps this device's e2ee keys, so a login

/// as any other device would orphan them

#[derive(Clone, Debug)]

#[non_exhaustive]

pub struct SoftLogout {

    pub homeserver: String,

    pub user_id: OwnedUserId,

    pub device_id: OwnedDeviceId,

    /// set when the login was OAuth, for `OAuthLogin::with_client_id`

    pub client_id: Option<String>,

}

impl SoftLogout {

    #[must_use]

    pub fn password_login(&self, password: impl Into<String>) -> PasswordLogin {

        PasswordLogin::new(self.homeserver.clone(), self.user_id.as_str(), password)

            .with_device_id(self.device_id.clone())

            .allow_insecure_loopback(true)

    }