malice

use std::{

    collections::HashMap,

    path::{Path, PathBuf},

    sync::{Arc, PoisonError, Weak},

    time::Duration,

};

use futures_util::StreamExt;

use macltok::{

    AccountError, ClientHandle, EncryptedSessionStore, EngineConfig, EngineError, ErrorKind,

    Mirror, MirrorHandle, PasswordLogin, Restore, Retain, Session, SessionStatus, SessionStore,

    SoftLogout, StateChange, StoreConfig, Termination,

};

use rand_core::{OsRng, RngCore};

use sha2::{Digest, Sha256};

use tokio::{

    sync::{Mutex, OwnedMutexGuard},

    task::JoinHandle,

    time::timeout,

};

use malice_core::{

    error::{CoreError, Failure},

    login::login_error,

};

use crate::csrf::valid_sid;

const AUTH_TIMEOUT: Duration = Duration::from_secs(20);

pub struct UserSession {

    pub sid: String,

    pub user_id: String,

    pub handle: ClientHandle,

    pub mirror: MirrorHandle,

    sessions: Arc<dyn SessionStore>,

    engine: Mutex<Option<JoinHandle<Result<Termination, EngineError>>>>,

    mirror_task: Mutex<Option<JoinHandle<()>>>,

}

struct Inner {

    sessions: tokio::sync::RwLock<HashMap<String, Arc<UserSession>>>,

    /// Serialises restores, logins and removals of the same sid so two engines

    /// never open the same SQLite store at once, without one slow homeserver

    /// holding up every other user.

    sid_locks: std::sync::Mutex<HashMap<String, Weak<Mutex<()>>>>,

    data_dir: PathBuf,

    master_key: [u8; 32],

    cookie_secure: bool,

    allow_http_homeserver: bool,

}

#[derive(Clone)]

pub struct AppState {

    inner: Arc<Inner>,

}

#[derive(Debug)]

pub enum RestoreError {

    /// No cookie session and nothing on disk: plain guest.

    NotFound,

    /// Credential on disk is unusable: caller should clear the cookie.

    Stale(String),

    /// Credential looks fine but the engine would not start.

    Unavailable(String),

    /// The server logged this device out; it can log back in as the same user.

    SoftLoggedOut { user_id: String, homeserver: String },

}

impl AppState {

    pub fn new(

        data_dir: PathBuf,

        master_key: [u8; 32],

        cookie_secure: bool,

        allow_http_homeserver: bool,

    ) -> Self {

        Self {

            inner: Arc::new(Inner {

                sessions: tokio::sync::RwLock::new(HashMap::new()),

                sid_locks: std::sync::Mutex::new(HashMap::new()),

                data_dir,

                master_key,

                cookie_secure,

                allow_http_homeserver,

            }),

        }

    }

    pub fn allow_http_homeserver(&self) -> bool {

        self.inner.allow_http_homeserver

    }

    pub fn cookie_secure(&self) -> bool {

        self.inner.cookie_secure

    }

    async fn lock_sid(&self, sid: &str) -> OwnedMutexGuard<()> {

        let lock = {

            let mut locks = self

                .inner

                .sid_locks

                .lock()

                .unwrap_or_else(PoisonError::into_inner);

            locks.retain(|_, lock| lock.strong_count() > 0);

            if let Some(lock) = locks.get(sid).and_then(Weak::upgrade) {

                lock

            } else {

                let lock = Arc::new(Mutex::new(()));

                locks.insert(sid.to_owned(), Arc::downgrade(&lock));

                lock

            }

        };

        lock.lock_owned().await

    }

    fn sessions_dir(&self) -> PathBuf {

        self.inner.data_dir.join("sessions")

    }

    fn session_file(&self, sid: &str) -> PathBuf {

        self.sessions_dir().join(format!("{sid}.bin"))

    }

    fn store_dir(&self, sid: &str) -> PathBuf {

        self.inner.data_dir.join("stores").join(sid)

    }

    fn store_passphrase(&self, sid: &str) -> String {

        let mut hasher = Sha256::new();

        hasher.update(self.inner.master_key);

        hasher.update(b"malice-store-v1:");

        hasher.update(sid.as_bytes());

        hex::encode(hasher.finalize())

    }

    fn store_config(&self, sid: &str) -> StoreConfig {

        StoreConfig::sqlite(self.store_dir(sid), self.store_passphrase(sid))

    }

    /// A fresh handle onto a sid's record and store. Cheap: nothing touches

    /// disk until the handle is actually used.

    fn session_store(&self, sid: &str) -> Arc<dyn SessionStore> {

        Arc::new(EncryptedSessionStore::file(

            self.session_file(sid),

            self.store_config(sid),

        ))

    }

    async fn apply_retention(&self, sessions: &Arc<dyn SessionStore>, retain: Retain) {

        let result = match retain {

            Retain::RecordAndStore => Ok(()),

            Retain::Store => sessions.soft_log_out().await,

            Retain::Nothing => sessions.discard().await,

        };

        if let Err(error) = result {

            tracing::warn!(%error, "applying session retention failed");

        }

    }

    /// Resolve a cookie sid to a live session, restoring from disk on demand.

    /// The shared map is never held across network or engine work: the fast

    /// path clones an `Arc` under a read lock, and restores serialise on a

    /// short init lock with a double-checked insert.

    pub async fn get_or_restore(&self, sid: &str) -> Result<Arc<UserSession>, RestoreError> {

        if !valid_sid(sid) {

            return Err(RestoreError::NotFound);

        }

        if let Some(user) = self.inner.sessions.read().await.get(sid).cloned() {

            return Ok(user);

        }

        let _init = self.lock_sid(sid).await;

        if let Some(user) = self.inner.sessions.read().await.get(sid).cloned() {

            return Ok(user);

        }

        let sessions = self.session_store(sid);

        let restore = Session::restore(sessions.clone()).await.map_err(|error| {

            tracing::warn!("restoring persisted session failed: {error}");

            restore_error(error)

        })?;

        match restore {

            None | Some(Restore::OAuthPending) => Err(RestoreError::NotFound),

            Some(Restore::SoftLoggedOut(soft)) => Err(RestoreError::SoftLoggedOut {

                user_id: soft.user_id.to_string(),

                homeserver: soft.homeserver,

            }),

            Some(Restore::Resumed(session)) => self

                .start_user_session(sid.to_owned(), session, sessions)

                .await

                .map_err(|error| {

                    tracing::warn!("starting restored session failed: {error}");

                    RestoreError::Unavailable(error)

                }),

        }

    }

    /// Restore just far enough to tell whether `sid`'s record is a soft

    /// logout, without disturbing anything else.

    async fn soft_logged_out(&self, sid: &str) -> Option<SoftLogout> {

        match Session::restore(self.session_store(sid)).await {

            Ok(Some(Restore::SoftLoggedOut(soft))) => Some(soft),

            _ => None,

        }

    }

    /// Password login: creates a fresh opaque sid and starts the engine,

    /// which saves the record itself. The password never leaves this call.

    /// When the cookie sid names the same user's soft-logged-out device, logs

    /// back in as that device against the same sid instead, keeping its keys.

    pub async fn login(

        &self,

        cookie_sid: Option<&str>,

        login: PasswordLogin,

        username: &str,

        password: String,

    ) -> Result<Arc<UserSession>, CoreError> {

        if let Some(sid) = cookie_sid.filter(|sid| valid_sid(sid))

            && let Some(soft) = self.soft_logged_out(sid).await

        {

            if username_matches(username, &soft.user_id) {

                let sessions = self.session_store(sid);

                let session = login_with(

                    soft.password_login(password.clone()),

                    &password,

                    sessions.clone(),

                )

                .await?;

                let _init = self.lock_sid(sid).await;

                return self

                    .start_user_session(sid.to_owned(), session, sessions)

                    .await

                    .map_err(|error| CoreError::new(Failure::Upstream, error));

            }

            self.apply_retention(&self.session_store(sid), Retain::Nothing)

                .await;

        }

        let sid = new_sid();

        let sessions = self.session_store(&sid);

        let session = login_with(login, &password, sessions.clone()).await?;

        let _init = self.lock_sid(&sid).await;

        self.start_user_session(sid, session, sessions)

            .await

            .map_err(|error| CoreError::new(Failure::Upstream, error))

    }

    async fn start_user_session(

        &self,

        sid: String,

        session: Session,

        sessions: Arc<dyn SessionStore>,

    ) -> Result<Arc<UserSession>, String> {

        let (engine, handle, events) = session

            .build(EngineConfig::default())

            .await

            .map_err(|error| format!("starting sync engine: {error}"))?;

        let engine_task = tokio::spawn(engine.run());

        let (mirror, mirror_handle) = Mirror::new(handle.clone(), events);

        let mirror_task = tokio::spawn(mirror.run());

        let Some(user_id) = wait_authenticated(&mirror_handle).await else {

            engine_task.abort();

            let _ = engine_task.await;

            mirror_task.abort();

            let _ = mirror_task.await;

            return Err("sync engine stopped before authenticating".to_owned());

        };

        let user = Arc::new(UserSession {

            sid: sid.clone(),

            user_id,

            handle,

            mirror: mirror_handle,

            sessions,

            engine: Mutex::new(Some(engine_task)),

            mirror_task: Mutex::new(Some(mirror_task)),

        });

        {

            let mut sessions_map = self.inner.sessions.write().await;

            sessions_map.insert(sid, user.clone());

        }

        let supervisor_state = self.clone();

        let supervisor_user = user.clone();

        tokio::spawn(async move {

            supervise(supervisor_state, supervisor_user).await;

        });

        Ok(user)

    }

    /// Logout: best-effort server-side logout, then drop all local state and

    /// delete persisted secrets. Never holds the map lock across the network

    /// call. Forces `Retain::Nothing` regardless of what the engine's own

    /// termination said, so an explicit logout always ends in a full discard

    /// even if it raced the supervisor over a server-initiated soft logout.

    pub async fn logout(&self, sid: &str) {

        let user = { self.inner.sessions.read().await.get(sid).cloned() };

        if let Some(user) = user {

            // answers once the server has let go of the device, and the engine stops itself

            let _ = timeout(Duration::from_secs(5), user.handle.logout()).await;

        }

        self.remove_session(sid, Retain::Nothing).await;

    }

    /// Forget an in-memory session, abort its engine and mirror tasks, and

    /// apply the given retention to its persisted secrets. Called both for

    /// explicit logout (always `Retain::Nothing`) and from the supervisor

    /// once an engine's `Termination` says what its record and store are

    /// still worth. Retention always runs, even if the session was already

    /// gone from the map, so a racing logout's `Nothing` still lands.

    pub async fn remove_session(&self, sid: &str, retain: Retain) {

        // A restore must not reopen the store between removing the map entry

        // and deleting its persisted credentials.

        let _init = self.lock_sid(sid).await;

        let user = self.inner.sessions.write().await.remove(sid);

        let sessions = match &user {

            Some(user) => user.sessions.clone(),

            None => self.session_store(sid),

        };

        if let Some(user) = &user {

            if let Some(task) = user.engine.lock().await.take() {

                task.abort();

            }

            if let Some(task) = user.mirror_task.lock().await.take() {

                task.abort();

            }

        }

        self.apply_retention(&sessions, retain).await;

    }

    /// Graceful shutdown: ask every engine to stop, keep persisted records so

    /// the next boot restores them.

    pub async fn shutdown_gracefully(&self) {

        let users: Vec<Arc<UserSession>> =

            self.inner.sessions.read().await.values().cloned().collect();

        for user in &users {

            // the supervisor keeps waiting, so a shutdown can answer once `Stopped` is out

            let handle = user.handle.clone();

            tokio::spawn(async move {

                let _ = timeout(Duration::from_secs(5), handle.shutdown()).await;

            });

        }

        // Give engines a moment to emit `Stopped`; the supervisor then drops

        // each in-memory session while leaving persisted records behind.

        timeout(Duration::from_secs(5), async {

            loop {

                if self.inner.sessions.read().await.is_empty() {

                    break;

                }

                tokio::time::sleep(Duration::from_millis(50)).await;

            }

        })

        .await

        .ok();

        for user in &users {

            if let Some(task) = user.engine.lock().await.take() {

                task.abort();

            }

            if let Some(task) = user.mirror_task.lock().await.take() {

                task.abort();

            }

        }

        self.inner.sessions.write().await.clear();

    }

}

fn restore_error(error: AccountError) -> RestoreError {

    match error {

        AccountError::StoreInUse => {

            RestoreError::Unavailable("another connection is using this session".to_owned())

        }

        AccountError::Operation {

            kind: ErrorKind::Network | ErrorKind::Server | ErrorKind::RateLimited,

            ..

        } => RestoreError::Unavailable(format!("restoring persisted session failed: {error}")),

        error @ (AccountError::Operation { .. }

        | AccountError::MissingSession

        | AccountError::UnsupportedAuthentication

        | AccountError::OAuthCallbackAlreadySubmitted

        | AccountError::NoPendingOAuth

        | AccountError::InvalidHomeserver(_)

        | AccountError::AlreadyLoggedIn(_)

        | AccountError::StoreDeletePending

        | AccountError::Unsaved { .. }

        | AccountError::Engine(_)) => {

            RestoreError::Stale(format!("restoring persisted session failed: {error}"))

        }

    }

}

async fn login_with(

    login: PasswordLogin,

    password: &str,

    sessions: Arc<dyn SessionStore>,

) -> Result<Session, CoreError> {

    match Session::login_password(login, sessions).await {

        Ok(session) => Ok(session),

        Err(AccountError::Unsaved { session, .. }) => session

            .save()

            .await

            .map_err(|_| CoreError::new(Failure::Upstream, "could not save the session")),

        Err(error) => Err(login_error(&error, password)),

    }

}

fn username_matches(submitted: &str, user_id: &ruma::OwnedUserId) -> bool {

    let submitted = submitted.trim();

    submitted == user_id.as_str() || submitted == user_id.localpart()

}

/// Wait for the mirror to learn the user id. Fails once the subscription

/// ends or the session leaves `Active` before that happens.

async fn wait_authenticated(mirror: &MirrorHandle) -> Option<String> {

    timeout(AUTH_TIMEOUT, async {

        let mut updates = mirror.subscribe();

        loop {

            let update = updates.next().await?;

            match update.change {

                StateChange::State(state) => {

                    if let Some(user_id) = &state.user_id {

                        return Some(user_id.to_string());

                    }

                    if state.session != SessionStatus::Active {

                        return None;

                    }

                }

                StateChange::Authenticated(user_id) => return Some(user_id.to_string()),

                StateChange::Session(status) if status != SessionStatus::Active => return None,

                StateChange::Session(_)

                | StateChange::SessionSaveFailed(_)

                | StateChange::SessionSaved

                | StateChange::Connection(_)

                | StateChange::Encryption(_)

                | StateChange::Verification(_)

                | StateChange::VerificationTarget(_)

                | StateChange::Mentions(_)

                | StateChange::NotificationSettings

                | StateChange::Rooms(_)

                | StateChange::ImagePacks(_)

                | StateChange::Timeline { .. }

                | StateChange::History { .. }

                | StateChange::HistoryRepair(_)

                | StateChange::RecentHistory { .. }

                | StateChange::ThreadList { .. }

                | StateChange::EventContextClosed { .. }

                | StateChange::SearchClosed { .. }

                | StateChange::Typing(_)

                | StateChange::Upload { .. }

                | StateChange::UploadFinished { .. }

                | StateChange::Failed(_) => {}

            }

        }

    })

    .await

    .ok()

    .flatten()

}

/// Awaits a session's engine, then its mirror, removes it from the map, and

/// applies the retention its `Termination` earns. Taking the engine

/// `JoinHandle` is the single point two racing teardowns (this task and an

/// explicit logout) arbitrate over: whichever finds it empty does nothing

/// further, so the engine and mirror are never aborted twice.

async fn supervise(state: AppState, user: Arc<UserSession>) {

    let Some(engine_task) = user.engine.lock().await.take() else {

        return;

    };

    let retain = match engine_task.await {

        Ok(Ok(termination)) => termination.retain(),

        Ok(Err(error)) => {

            tracing::warn!(%error, "sync engine ended with an error; keeping its session");

            Retain::RecordAndStore

        }

        Err(error) => {

            tracing::warn!(%error, "sync engine task did not finish cleanly; keeping its session");

            Retain::RecordAndStore

        }

    };

    if let Some(mirror_task) = user.mirror_task.lock().await.take() {

        let _ = mirror_task.await;

    }

    state.remove_session(&user.sid, retain).await;