macltok

public
Star0
Watch0
26 commits 3 branches 0 tags 0 releases3.0 MiB
#[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, UserSession,
            registration::{ApplicationType, ClientMetadata, Localized, OAuthGrantType},
        },
    },
    cross_process_lock::CrossProcessLockConfig,
    encryption::{BackupDownloadStrategy, EncryptionSettings},
    ruma::{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, 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 {
    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>,
}

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,
        }
    }

    #[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
    }
}

#[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,
}

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,
        }
    }

    #[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
    }
}

#[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>,
}

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;
            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
    }

    pub async fn abort(mut self) {
        if let Some(state) = self.state.take() {
            self.client.oauth().abort_login(&state).await;
        }
    }
}

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 = "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>,
    },
}

impl SessionRecord {
    #[must_use]
    pub fn homeserver(&self) -> &str {
        &self.homeserver
    }

    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,
        };
        Self {
            homeserver: self.homeserver,
            authentication,
        }
    }
}

pub enum Restore {
    Resumed(Session),
    SoftLoggedOut(SoftLogout),
}

/// 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())
    }
}

#[derive(Debug, Error)]
pub enum AccountError {
    #[error("{context} failed")]
    Operation {
        context: &'static str,
        #[source]
        source: Box<SourceError>,
    },
    #[error("the homeserver returned no authenticated session")]
    MissingSession,
    #[error("the SDK returned an unsupported authentication session")]
    UnsupportedAuthentication,
    #[error("the OAuth callback was already submitted")]
    OAuthCallbackAlreadySubmitted,
    /// an IndexedDB delete is waiting on the store's other connections. it goes through
    /// once they close, and opening the store again waits behind it
    #[error("the store is still open elsewhere, it will be deleted once that closes")]
    StoreDeletePending,
    /// another engine, in this process or another, has the store. only one runs per store
    #[error("another engine is using this store")]
    StoreInUse,
    /// the login worked but its first record didn't save. the server already has the
    /// device, so retry the save rather than log in again
    #[error("saving the new login's record failed")]
    Unsaved {
        session: Box<UnsavedSession>,
        #[source]
        source: Box<AccountError>,
    },
    #[error(transparent)]
    Engine(#[from] EngineError),
}

/// a live login whose record never reached its `SessionStore`
pub struct UnsavedSession {
    session: Session,
    sessions: Arc<dyn SessionStore>,
}

impl UnsavedSession {
    pub async fn save(self) -> Result<Session, AccountError> {
        let saved = match session_record(&self.session.client) {
            Ok(record) => self.sessions.save(&record).await,
            Err(error) => Err(error),
        };
        match saved {
            Ok(()) => Ok(self.session),
            Err(source) => Err(AccountError::Unsaved {
                session: Box::new(self),
                source: Box::new(source),
            }),
        }
    }

    /// the record to keep somewhere else, if the store won't take it
    pub fn record(&self) -> Result<SessionRecord, AccountError> {
        session_record(&self.session.client)
    }
}

impl std::fmt::Debug for UnsavedSession {
    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        formatter
            .debug_struct("UnsavedSession")
            .finish_non_exhaustive()
    }
}

impl Session {
    pub async fn login_password(
        login: PasswordLogin,
        sessions: Arc<dyn SessionStore>,
    ) -> Result<Self, AccountError> {
        let store_lock = sessions.store().lock().await?;
        let client = matrix_client_builder(sessions.store())
            .server_name_or_homeserver_url(login.homeserver.trim())
            .build()
            .await
            .map_err(|error| account_error("building the Matrix client", error))?;
        let tracking = Tracking::start(&client, sessions.clone())?;
        let mut request = client
            .matrix_auth()
            .login_username(login.username.trim(), &login.password)
            .request_refresh_token();
        if let Some(name) = login.device_display_name.as_deref() {
            request = request.initial_device_display_name(name);
        }
        if let Some(device_id) = login.device_id.as_deref() {
            request = request.device_id(device_id.as_str());
        }
        request
            .send()
            .await
            .map_err(|error| account_error("password login", error))?;
        Self::saved(client, sessions, tracking, store_lock).await
    }

    /// the pending login lives only in this process, so the one that begins it has to finish it
    pub async fn begin_oauth(
        login: OAuthLogin,
        sessions: Arc<dyn SessionStore>,
    ) -> Result<PendingOAuthLogin, AccountError> {
        let redirect_uri = Url::parse(login.redirect_uri.trim())
            .map_err(|error| account_error("parsing the OAuth redirect URI", error))?;
        let client_uri = Url::parse(login.client_uri.trim())
            .map_err(|error| account_error("parsing the OAuth client URI", error))?;
        let store_lock = sessions.store().lock().await?;
        let client = matrix_client_builder(sessions.store())
            .server_name_or_homeserver_url(login.homeserver.trim())
            .build()
            .await
            .map_err(|error| account_error("building the Matrix client", error))?;
        let tracking = Tracking::start(&client, sessions.clone())?;
        let application_type = match login.application_type {
            OAuthApplicationType::Native => ApplicationType::Native,
            OAuthApplicationType::Web => ApplicationType::Web,
        };
        let mut metadata = ClientMetadata::new(
            application_type,
            vec![OAuthGrantType::AuthorizationCode {
                redirect_uris: vec![redirect_uri.clone()],
            }],
            Localized::new(client_uri, []),
        );
        metadata.client_name = login.client_name.map(|name| Localized::new(name, []));
        let metadata = Raw::new(&metadata)
            .map_err(|error| account_error("serializing OAuth client metadata", error))?;
        let mut registration = ClientRegistrationData::new(metadata);
        if let Some(client_id) = login.client_id {
            registration.static_registrations = Some(HashMap::from([(
                client.homeserver(),
                ClientId::new(client_id),
            )]));
        }
        let authorization = client
            .oauth()
            .login(redirect_uri, login.device_id, Some(registration), None)
            .build()
            .await
            .map_err(|error| account_error("starting OAuth login", error))?;
        Ok(PendingOAuthLogin {
            client,
            sessions,
            tracking: Some((tracking, store_lock)),
            authorization_url: authorization.url.to_string(),
            state: Some(authorization.state),
        })
    }

    /// Restore the record `sessions` holds, or `None` when it holds none.
    /// The caller must also handle rejection and explicit logout.
    pub async fn restore(sessions: Arc<dyn SessionStore>) -> Result<Option<Restore>, AccountError> {
        let Some(record) = sessions.load().await? else {
            return Ok(None);
        };
        let authentication = match record.authentication {
            AuthenticationRecord::Matrix(session) => AuthSession::Matrix(session),
            AuthenticationRecord::OAuth { client_id, user } => {
                OAuthSession { client_id, user }.into()
            }
            AuthenticationRecord::SoftLoggedOut { meta, client_id } => {
                return Ok(Some(Restore::SoftLoggedOut(SoftLogout {
                    homeserver: record.homeserver,
                    user_id: meta.user_id,
                    device_id: meta.device_id,
                    client_id: client_id.map(|client_id| client_id.as_str().to_owned()),
                })));
            }
        };
        let store_lock = sessions.store().lock().await?;
        let client = matrix_client_builder(sessions.store())
            .homeserver_url(&record.homeserver)
            .build()
            .await
            .map_err(|error| account_error("reopening the Matrix store", error))?;
        let tracking = Tracking::start(&client, sessions)?;
        client
            .restore_session(authentication)
            .await
            .map_err(|error| account_error("restoring the authenticated session", error))?;
        Ok(Some(Restore::Resumed(Self {
            client,
            tracking,
            store_lock,
        })))
    }

    /// leave the record to the host, which then saves `ClientHandle::session_record()`
    /// on every `SessionTokensChanged`
    #[must_use]
    pub fn without_persistence(mut self) -> Self {
        self.tracking.persistence = None;
        #[cfg(not(target_family = "wasm"))]
        self.tracking.inline.store(false, Ordering::Release);
        self
    }

    pub async fn build(
        self,
        config: EngineConfig,
    ) -> Result<(Engine, ClientHandle, EventStream), AccountError> {
        Engine::with_persistence(
            self.client,
            config,
            self.tracking.session_changes,
            self.tracking.persistence,
            Some(self.store_lock),
        )
        .await
        .map_err(AccountError::from)
    }

    async fn saved(
        client: Client,
        sessions: Arc<dyn SessionStore>,
        tracking: Tracking,
        store_lock: StoreLock,
    ) -> Result<Self, AccountError> {
        UnsavedSession {
            session: Self {
                client,
                tracking,
                store_lock,
            },
            sessions,
        }
        .save()
        .await
    }
}

impl Tracking {
    fn start(client: &Client, sessions: Arc<dyn SessionStore>) -> Result<Self, AccountError> {
        let session_changes = client.subscribe_to_session_changes();
        #[cfg_attr(target_family = "wasm", allow(unused_variables))]
        let (results, inline_results) = mpsc::unbounded_channel();
        #[cfg(not(target_family = "wasm"))]
        let inline = save_inline(client, sessions.clone(), results)?;
        Ok(Self {
            session_changes,
            persistence: Some(Persistence {
                save: Box::new(move |client| {
                    let sessions = sessions.clone();
                    Box::pin(async move {
                        let record = session_record(&client).map_err(save_failed)?;
                        sessions.save(&record).await.map_err(save_failed)
                    })
                }),
                inline: inline_results,
            }),
            #[cfg(not(target_family = "wasm"))]
            inline,
        })
    }
}

/// the sdk calls this inside its refresh, after it has the new tokens and before it uses
/// them, so a store that writes there and then can't lose a refresh token to a crash
#[cfg(not(target_family = "wasm"))]
fn save_inline(
    client: &Client,
    sessions: Arc<dyn SessionStore>,
    results: mpsc::UnboundedSender<Result<(), crate::Error>>,
) -> Result<Arc<AtomicBool>, AccountError> {
    let enabled = Arc::new(AtomicBool::new(true));
    let saving = enabled.clone();
    client
        .set_session_callbacks(
            // only asked for with the cross process refresh lock, which we never enable
            Box::new(|client| {
                client
                    .session_tokens()
                    .ok_or_else(|| AccountError::MissingSession.into())
            }),
            Box::new(move |client| {
                if !saving.load(Ordering::Acquire) {
                    return Ok(());
                }
                let record = session_record(&client)?;
                if let Some(result) = sessions.save_blocking(&record) {
                    let _ = results.send(result.map_err(save_failed));
                }
                Ok(())
            }),
        )
        .map_err(|error| account_error("registering the session callbacks", error))?;
    Ok(enabled)
}

fn save_failed(error: AccountError) -> crate::Error {
    let message = format!("saving the session record failed, {error}");
    crate::Error::with_source(ErrorKind::Store, message, error)
}

pub(crate) fn session_record(client: &Client) -> Result<SessionRecord, AccountError> {
    let authentication = match client.session().ok_or(AccountError::MissingSession)? {
        AuthSession::Matrix(session) => AuthenticationRecord::Matrix(session),
        AuthSession::OAuth(session) => {
            let OAuthSession { client_id, user } = *session;
            AuthenticationRecord::OAuth { client_id, user }
        }
        _ => return Err(AccountError::UnsupportedAuthentication),
    };
    Ok(SessionRecord {
        homeserver: client.homeserver().to_string(),
        authentication,
    })
}

fn encryption_settings() -> EncryptionSettings {
    EncryptionSettings {
        auto_enable_cross_signing: true,
        backup_download_strategy: BackupDownloadStrategy::OneShot,
        auto_enable_backups: true,
    }
}

/// `StoreLock` already keeps every other client off the store, and the sdk's lease would cost
/// a store write per locked section and a task renewing it every 50ms
fn matrix_client_builder(store: &StoreConfig) -> ClientBuilder {
    let builder = Client::builder()
        .with_encryption_settings(encryption_settings())
        .with_threading_support(ThreadingSupport::Enabled {
            with_subscriptions: false,
        })
        .handle_refresh_tokens()
        .cross_process_store_config(CrossProcessLockConfig::SingleProcess);
    match &store.location {
        #[cfg(feature = "sqlite")]
        StoreLocation::Sqlite(path) => builder.sqlite_store(path, Some(&store.passphrase)),
        #[cfg(feature = "indexeddb")]
        StoreLocation::IndexedDb(name) => builder.indexeddb_store(name, Some(&store.passphrase)),
    }
}

pub(crate) fn account_error(
    context: &'static str,
    source: impl std::error::Error + SendOutsideWasm + SyncOutsideWasm + 'static,
) -> AccountError {
    AccountError::Operation {
        context,
        source: Box::new(source),
    }
}