malice

public

Star0

Watch0

152 commits 2 branches 0 tags 0 releases4.4 MiB

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, push::HostedPush};

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>,
    push: Option<Arc<HostedPush>>,
}

#[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,
            }),
            push: None,
        }
    }

    pub fn with_push(self, push: HostedPush) -> Self {
        Self {
            push: Some(Arc::new(push)),
            ..self
        }
    }

    pub fn push(&self) -> Option<&HostedPush> {
        self.push.as_deref()
    }

    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 {
            let pushers = self
                .push()
                .map(|push| push.pusher_keys(sid))
                .unwrap_or_default();
            // answers once the server has let go of the device, and the engine stops itself
            let _ = timeout(Duration::from_secs(5), user.handle.logout(pushers)).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;
        if retain == Retain::Nothing
            && let Some(push) = self.push()
        {
            push.forget_session(sid);
        }
    }

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

fn new_sid() -> String {
    let mut bytes = [0u8; 32];
    OsRng.fill_bytes(&mut bytes);
    hex::encode(bytes)
}

/// Load the 32-byte master key, creating it with owner-only permissions when
/// absent. Secrets must never be world-readable.
pub fn load_or_create_master_key(data_dir: &Path) -> Result<[u8; 32], String> {
    create_private_dir(data_dir).map_err(|error| format!("creating data dir: {error}"))?;
    restrict_dir(data_dir);
    let path = data_dir.join("master.key");
    if path.exists() {
        let bytes = std::fs::read(&path).map_err(|error| format!("reading master key: {error}"))?;
        restrict_file(&path);
        if bytes.len() != 32 {
            return Err("master key has an unexpected length".to_owned());
        }
        let mut key = [0u8; 32];
        key.copy_from_slice(&bytes);
        return Ok(key);
    }
    let mut key = [0u8; 32];
    OsRng.fill_bytes(&mut key);
    let mut file =
        create_private_file(&path).map_err(|error| format!("creating master key: {error}"))?;
    if let Err(error) = std::io::Write::write_all(&mut file, &key).and_then(|_| file.sync_all()) {
        let _ = std::fs::remove_file(&path);
        return Err(format!("writing master key: {error}"));
    }
    Ok(key)
}

fn create_private_dir(path: &Path) -> std::io::Result<()> {
    let mut builder = std::fs::DirBuilder::new();
    builder.recursive(true);
    #[cfg(unix)]
    {
        use std::os::unix::fs::DirBuilderExt;
        builder.mode(0o700);
    }
    builder.create(path)
}

fn create_private_file(path: &Path) -> std::io::Result<std::fs::File> {
    let mut options = std::fs::OpenOptions::new();
    options.write(true).create_new(true);
    #[cfg(unix)]
    {
        use std::os::unix::fs::OpenOptionsExt;
        options.mode(0o600);
    }
    options.open(path)
}

pub fn ensure_dirs(data_dir: &Path) {
    for dir in [
        data_dir.to_path_buf(),
        data_dir.join("sessions"),
        data_dir.join("stores"),
    ] {
        let _ = create_private_dir(&dir);
        restrict_dir(&dir);
    }
}

#[cfg(unix)]
fn restrict_dir(path: &Path) {
    use std::os::unix::fs::PermissionsExt;
    let _ = std::fs::set_permissions(path, std::fs::Permissions::from_mode(0o700));
}

#[cfg(not(unix))]
fn restrict_dir(_path: &Path) {}

#[cfg(unix)]
fn restrict_file(path: &Path) {
    use std::os::unix::fs::PermissionsExt;
    let _ = std::fs::set_permissions(path, std::fs::Permissions::from_mode(0o600));
}

#[cfg(not(unix))]
fn restrict_file(_path: &Path) {}

#[cfg(test)]
mod tests {
    use super::*;

    fn state() -> (AppState, tempfile::TempDir) {
        let dir = tempfile::tempdir().expect("temp dir");
        (
            AppState::new(dir.path().to_owned(), [7; 32], false, false),
            dir,
        )
    }

    #[tokio::test]
    async fn sid_locks_only_serialise_the_same_sid() {
        let (state, _dir) = state();
        let held = state.lock_sid("alpha").await;
        let other = timeout(Duration::from_millis(50), state.lock_sid("beta")).await;
        assert!(other.is_ok(), "a different sid must not wait");
        let same = timeout(Duration::from_millis(50), state.lock_sid("alpha")).await;
        assert!(same.is_err(), "the same sid must wait for the holder");
        drop(held);
        let again = timeout(Duration::from_millis(50), state.lock_sid("alpha")).await;
        assert!(again.is_ok(), "the sid is free once the holder drops");
    }

    #[tokio::test]
    async fn sid_locks_forget_released_sids() {
        let (state, _dir) = state();
        drop(state.lock_sid("alpha").await);
        drop(state.lock_sid("beta").await);
        let live = state.inner.sid_locks.lock().unwrap().len();
        assert!(live <= 1, "released locks are pruned on the next lookup");
    }
}

1	use std::{
2	collections::HashMap,
3	path::{Path, PathBuf},
4	sync::{Arc, PoisonError, Weak},
5	time::Duration,
6	};
7
8	use futures_util::StreamExt;
9	use macltok::{
10	AccountError, ClientHandle, EncryptedSessionStore, EngineConfig, EngineError, ErrorKind,
11	Mirror, MirrorHandle, PasswordLogin, Restore, Retain, Session, SessionStatus, SessionStore,
12	SoftLogout, StateChange, StoreConfig, Termination,
13	};
14	use rand_core::{OsRng, RngCore};
15	use sha2::{Digest, Sha256};
16	use tokio::{
17	sync::{Mutex, OwnedMutexGuard},
18	task::JoinHandle,
19	time::timeout,
20	};
21
22	use malice_core::{
23	error::{CoreError, Failure},
24	login::login_error,
25	};
26
27	use crate::{csrf::valid_sid, push::HostedPush};
28
29	const AUTH_TIMEOUT: Duration = Duration::from_secs(20);
30
31	pub struct UserSession {
32	pub sid: String,
33	pub user_id: String,
34	pub handle: ClientHandle,
35	pub mirror: MirrorHandle,
36	sessions: Arc<dyn SessionStore>,
37	engine: Mutex<Option<JoinHandle<Result<Termination, EngineError>>>>,
38	mirror_task: Mutex<Option<JoinHandle<()>>>,
39	}
40
41	struct Inner {
42	sessions: tokio::sync::RwLock<HashMap<String, Arc<UserSession>>>,
43	/// Serialises restores, logins and removals of the same sid so two engines
44	/// never open the same SQLite store at once, without one slow homeserver
45	/// holding up every other user.
46	sid_locks: std::sync::Mutex<HashMap<String, Weak<Mutex<()>>>>,
47	data_dir: PathBuf,
48	master_key: [u8; 32],
49	cookie_secure: bool,
50	allow_http_homeserver: bool,
51	}
52
53	#[derive(Clone)]
54	pub struct AppState {
55	inner: Arc<Inner>,
56	push: Option<Arc<HostedPush>>,
57	}
58
59	#[derive(Debug)]
60	pub enum RestoreError {
61	/// No cookie session and nothing on disk: plain guest.
62	NotFound,
63	/// Credential on disk is unusable: caller should clear the cookie.
64	Stale(String),
65	/// Credential looks fine but the engine would not start.
66	Unavailable(String),
67	/// The server logged this device out; it can log back in as the same user.
68	SoftLoggedOut { user_id: String, homeserver: String },
69	}
70
71	impl AppState {
72	pub fn new(
73	data_dir: PathBuf,
74	master_key: [u8; 32],
75	cookie_secure: bool,
76	allow_http_homeserver: bool,
77	) -> Self {
78	Self {
79	inner: Arc::new(Inner {
80	sessions: tokio::sync::RwLock::new(HashMap::new()),
81	sid_locks: std::sync::Mutex::new(HashMap::new()),
82	data_dir,
83	master_key,
84	cookie_secure,
85	allow_http_homeserver,
86	}),
87	push: None,
88	}
89	}
90
91	pub fn with_push(self, push: HostedPush) -> Self {
92	Self {
93	push: Some(Arc::new(push)),
94	..self
95	}
96	}
97
98	pub fn push(&self) -> Option<&HostedPush> {
99	self.push.as_deref()
100	}
101
102	pub fn allow_http_homeserver(&self) -> bool {
103	self.inner.allow_http_homeserver
104	}
105
106	pub fn cookie_secure(&self) -> bool {
107	self.inner.cookie_secure
108	}
109
110	async fn lock_sid(&self, sid: &str) -> OwnedMutexGuard<()> {
111	let lock = {
112	let mut locks = self
113	.inner
114	.sid_locks
115	.lock()
116	.unwrap_or_else(PoisonError::into_inner);
117	locks.retain(\|_, lock\| lock.strong_count() > 0);
118	if let Some(lock) = locks.get(sid).and_then(Weak::upgrade) {
119	lock
120	} else {
121	let lock = Arc::new(Mutex::new(()));
122	locks.insert(sid.to_owned(), Arc::downgrade(&lock));
123	lock
124	}
125	};
126	lock.lock_owned().await
127	}
128
129	fn sessions_dir(&self) -> PathBuf {
130	self.inner.data_dir.join("sessions")
131	}
132
133	fn session_file(&self, sid: &str) -> PathBuf {
134	self.sessions_dir().join(format!("{sid}.bin"))
135	}
136
137	fn store_dir(&self, sid: &str) -> PathBuf {
138	self.inner.data_dir.join("stores").join(sid)
139	}
140
141	fn store_passphrase(&self, sid: &str) -> String {
142	let mut hasher = Sha256::new();
143	hasher.update(self.inner.master_key);
144	hasher.update(b"malice-store-v1:");
145	hasher.update(sid.as_bytes());
146	hex::encode(hasher.finalize())
147	}
148
149	fn store_config(&self, sid: &str) -> StoreConfig {
150	StoreConfig::sqlite(self.store_dir(sid), self.store_passphrase(sid))
151	}
152
153	/// A fresh handle onto a sid's record and store. Cheap: nothing touches
154	/// disk until the handle is actually used.
155	fn session_store(&self, sid: &str) -> Arc<dyn SessionStore> {
156	Arc::new(EncryptedSessionStore::file(
157	self.session_file(sid),
158	self.store_config(sid),
159	))
160	}
161
162	async fn apply_retention(&self, sessions: &Arc<dyn SessionStore>, retain: Retain) {
163	let result = match retain {
164	Retain::RecordAndStore => Ok(()),
165	Retain::Store => sessions.soft_log_out().await,
166	Retain::Nothing => sessions.discard().await,
167	};
168	if let Err(error) = result {
169	tracing::warn!(%error, "applying session retention failed");
170	}
171	}
172
173	/// Resolve a cookie sid to a live session, restoring from disk on demand.
174	/// The shared map is never held across network or engine work: the fast
175	/// path clones an `Arc` under a read lock, and restores serialise on a
176	/// short init lock with a double-checked insert.
177	pub async fn get_or_restore(&self, sid: &str) -> Result<Arc<UserSession>, RestoreError> {
178	if !valid_sid(sid) {
179	return Err(RestoreError::NotFound);
180	}
181	if let Some(user) = self.inner.sessions.read().await.get(sid).cloned() {
182	return Ok(user);
183	}
184	let _init = self.lock_sid(sid).await;
185	if let Some(user) = self.inner.sessions.read().await.get(sid).cloned() {
186	return Ok(user);
187	}
188	let sessions = self.session_store(sid);
189	let restore = Session::restore(sessions.clone()).await.map_err(\|error\| {
190	tracing::warn!("restoring persisted session failed: {error}");
191	restore_error(error)
192	})?;
193	match restore {
194	None \| Some(Restore::OAuthPending) => Err(RestoreError::NotFound),
195	Some(Restore::SoftLoggedOut(soft)) => Err(RestoreError::SoftLoggedOut {
196	user_id: soft.user_id.to_string(),
197	homeserver: soft.homeserver,
198	}),
199	Some(Restore::Resumed(session)) => self
200	.start_user_session(sid.to_owned(), session, sessions)
201	.await
202	.map_err(\|error\| {
203	tracing::warn!("starting restored session failed: {error}");
204	RestoreError::Unavailable(error)
205	}),
206	}
207	}
208
209	/// Restore just far enough to tell whether `sid`'s record is a soft
210	/// logout, without disturbing anything else.
211	async fn soft_logged_out(&self, sid: &str) -> Option<SoftLogout> {
212	match Session::restore(self.session_store(sid)).await {
213	Ok(Some(Restore::SoftLoggedOut(soft))) => Some(soft),
214	_ => None,
215	}
216	}
217
218	/// Password login: creates a fresh opaque sid and starts the engine,
219	/// which saves the record itself. The password never leaves this call.
220	/// When the cookie sid names the same user's soft-logged-out device, logs
221	/// back in as that device against the same sid instead, keeping its keys.
222	pub async fn login(
223	&self,
224	cookie_sid: Option<&str>,
225	login: PasswordLogin,
226	username: &str,
227	password: String,
228	) -> Result<Arc<UserSession>, CoreError> {
229	if let Some(sid) = cookie_sid.filter(\|sid\| valid_sid(sid))
230	&& let Some(soft) = self.soft_logged_out(sid).await
231	{
232	if username_matches(username, &soft.user_id) {
233	let sessions = self.session_store(sid);
234	let session = login_with(
235	soft.password_login(password.clone()),
236	&password,
237	sessions.clone(),
238	)
239	.await?;
240	let _init = self.lock_sid(sid).await;
241	return self
242	.start_user_session(sid.to_owned(), session, sessions)
243	.await
244	.map_err(\|error\| CoreError::new(Failure::Upstream, error));
245	}
246	self.apply_retention(&self.session_store(sid), Retain::Nothing)
247	.await;
248	}
249	let sid = new_sid();
250	let sessions = self.session_store(&sid);
251	let session = login_with(login, &password, sessions.clone()).await?;
252	let _init = self.lock_sid(&sid).await;
253	self.start_user_session(sid, session, sessions)
254	.await
255	.map_err(\|error\| CoreError::new(Failure::Upstream, error))
256	}
257
258	async fn start_user_session(
259	&self,
260	sid: String,
261	session: Session,
262	sessions: Arc<dyn SessionStore>,
263	) -> Result<Arc<UserSession>, String> {
264	let (engine, handle, events) = session
265	.build(EngineConfig::default())
266	.await
267	.map_err(\|error\| format!("starting sync engine: {error}"))?;
268	let engine_task = tokio::spawn(engine.run());
269	let (mirror, mirror_handle) = Mirror::new(handle.clone(), events);
270	let mirror_task = tokio::spawn(mirror.run());
271	let Some(user_id) = wait_authenticated(&mirror_handle).await else {
272	engine_task.abort();
273	let _ = engine_task.await;
274	mirror_task.abort();
275	let _ = mirror_task.await;
276	return Err("sync engine stopped before authenticating".to_owned());
277	};
278	let user = Arc::new(UserSession {
279	sid: sid.clone(),
280	user_id,
281	handle,
282	mirror: mirror_handle,
283	sessions,
284	engine: Mutex::new(Some(engine_task)),
285	mirror_task: Mutex::new(Some(mirror_task)),
286	});
287	{
288	let mut sessions_map = self.inner.sessions.write().await;
289	sessions_map.insert(sid, user.clone());
290	}
291	let supervisor_state = self.clone();
292	let supervisor_user = user.clone();
293	tokio::spawn(async move {
294	supervise(supervisor_state, supervisor_user).await;
295	});
296	Ok(user)
297	}
298
299	/// Logout: best-effort server-side logout, then drop all local state and
300	/// delete persisted secrets. Never holds the map lock across the network
301	/// call. Forces `Retain::Nothing` regardless of what the engine's own
302	/// termination said, so an explicit logout always ends in a full discard
303	/// even if it raced the supervisor over a server-initiated soft logout.
304	pub async fn logout(&self, sid: &str) {
305	let user = { self.inner.sessions.read().await.get(sid).cloned() };
306	if let Some(user) = user {
307	let pushers = self
308	.push()
309	.map(\|push\| push.pusher_keys(sid))
310	.unwrap_or_default();
311	// answers once the server has let go of the device, and the engine stops itself
312	let _ = timeout(Duration::from_secs(5), user.handle.logout(pushers)).await;
313	}
314	self.remove_session(sid, Retain::Nothing).await;
315	}
316
317	/// Forget an in-memory session, abort its engine and mirror tasks, and
318	/// apply the given retention to its persisted secrets. Called both for
319	/// explicit logout (always `Retain::Nothing`) and from the supervisor
320	/// once an engine's `Termination` says what its record and store are
321	/// still worth. Retention always runs, even if the session was already
322	/// gone from the map, so a racing logout's `Nothing` still lands.
323	pub async fn remove_session(&self, sid: &str, retain: Retain) {
324	// A restore must not reopen the store between removing the map entry
325	// and deleting its persisted credentials.
326	let _init = self.lock_sid(sid).await;
327	let user = self.inner.sessions.write().await.remove(sid);
328	let sessions = match &user {
329	Some(user) => user.sessions.clone(),
330	None => self.session_store(sid),
331	};
332	if let Some(user) = &user {
333	if let Some(task) = user.engine.lock().await.take() {
334	task.abort();
335	}
336	if let Some(task) = user.mirror_task.lock().await.take() {
337	task.abort();
338	}
339	}
340	self.apply_retention(&sessions, retain).await;
341	if retain == Retain::Nothing
342	&& let Some(push) = self.push()
343	{
344	push.forget_session(sid);
345	}
346	}
347
348	/// Graceful shutdown: ask every engine to stop, keep persisted records so
349	/// the next boot restores them.
350	pub async fn shutdown_gracefully(&self) {
351	let users: Vec<Arc<UserSession>> =
352	self.inner.sessions.read().await.values().cloned().collect();
353	for user in &users {
354	// the supervisor keeps waiting, so a shutdown can answer once `Stopped` is out
355	let handle = user.handle.clone();
356	tokio::spawn(async move {
357	let _ = timeout(Duration::from_secs(5), handle.shutdown()).await;
358	});
359	}
360	// Give engines a moment to emit `Stopped`; the supervisor then drops
361	// each in-memory session while leaving persisted records behind.
362	timeout(Duration::from_secs(5), async {
363	loop {
364	if self.inner.sessions.read().await.is_empty() {
365	break;
366	}
367	tokio::time::sleep(Duration::from_millis(50)).await;
368	}
369	})
370	.await
371	.ok();
372	for user in &users {
373	if let Some(task) = user.engine.lock().await.take() {
374	task.abort();
375	}
376	if let Some(task) = user.mirror_task.lock().await.take() {
377	task.abort();
378	}
379	}
380	self.inner.sessions.write().await.clear();
381	}
382	}
383
384	fn restore_error(error: AccountError) -> RestoreError {
385	match error {
386	AccountError::StoreInUse => {
387	RestoreError::Unavailable("another connection is using this session".to_owned())
388	}
389	AccountError::Operation {
390	kind: ErrorKind::Network \| ErrorKind::Server \| ErrorKind::RateLimited,
391	..
392	} => RestoreError::Unavailable(format!("restoring persisted session failed: {error}")),
393	error @ (AccountError::Operation { .. }
394	\| AccountError::MissingSession
395	\| AccountError::UnsupportedAuthentication
396	\| AccountError::OAuthCallbackAlreadySubmitted
397	\| AccountError::NoPendingOAuth
398	\| AccountError::InvalidHomeserver(_)
399	\| AccountError::AlreadyLoggedIn(_)
400	\| AccountError::StoreDeletePending
401	\| AccountError::Unsaved { .. }
402	\| AccountError::Engine(_)) => {
403	RestoreError::Stale(format!("restoring persisted session failed: {error}"))
404	}
405	}
406	}
407
408	async fn login_with(
409	login: PasswordLogin,
410	password: &str,
411	sessions: Arc<dyn SessionStore>,
412	) -> Result<Session, CoreError> {
413	match Session::login_password(login, sessions).await {
414	Ok(session) => Ok(session),
415	Err(AccountError::Unsaved { session, .. }) => session
416	.save()
417	.await
418	.map_err(\|_\| CoreError::new(Failure::Upstream, "could not save the session")),
419	Err(error) => Err(login_error(&error, password)),
420	}
421	}
422
423	fn username_matches(submitted: &str, user_id: &ruma::OwnedUserId) -> bool {
424	let submitted = submitted.trim();
425	submitted == user_id.as_str() \|\| submitted == user_id.localpart()
426	}
427
428	/// Wait for the mirror to learn the user id. Fails once the subscription
429	/// ends or the session leaves `Active` before that happens.
430	async fn wait_authenticated(mirror: &MirrorHandle) -> Option<String> {
431	timeout(AUTH_TIMEOUT, async {
432	let mut updates = mirror.subscribe();
433	loop {
434	let update = updates.next().await?;
435	match update.change {
436	StateChange::State(state) => {
437	if let Some(user_id) = &state.user_id {
438	return Some(user_id.to_string());
439	}
440	if state.session != SessionStatus::Active {
441	return None;
442	}
443	}
444	StateChange::Authenticated(user_id) => return Some(user_id.to_string()),
445	StateChange::Session(status) if status != SessionStatus::Active => return None,
446	StateChange::Session(_)
447	\| StateChange::SessionSaveFailed(_)
448	\| StateChange::SessionSaved
449	\| StateChange::Connection(_)
450	\| StateChange::Encryption(_)
451	\| StateChange::Verification(_)
452	\| StateChange::VerificationTarget(_)
453	\| StateChange::Mentions(_)
454	\| StateChange::NotificationSettings
455	\| StateChange::Rooms(_)
456	\| StateChange::ImagePacks(_)
457	\| StateChange::Timeline { .. }
458	\| StateChange::History { .. }
459	\| StateChange::HistoryRepair(_)
460	\| StateChange::RecentHistory { .. }
461	\| StateChange::ThreadList { .. }
462	\| StateChange::EventContextClosed { .. }
463	\| StateChange::SearchClosed { .. }
464	\| StateChange::Typing(_)
465	\| StateChange::Upload { .. }
466	\| StateChange::UploadFinished { .. }
467	\| StateChange::Failed(_) => {}
468	}
469	}
470	})
471	.await
472	.ok()
473	.flatten()
474	}
475
476	/// Awaits a session's engine, then its mirror, removes it from the map, and
477	/// applies the retention its `Termination` earns. Taking the engine
478	/// `JoinHandle` is the single point two racing teardowns (this task and an
479	/// explicit logout) arbitrate over: whichever finds it empty does nothing
480	/// further, so the engine and mirror are never aborted twice.
481	async fn supervise(state: AppState, user: Arc<UserSession>) {
482	let Some(engine_task) = user.engine.lock().await.take() else {
483	return;
484	};
485	let retain = match engine_task.await {
486	Ok(Ok(termination)) => termination.retain(),
487	Ok(Err(error)) => {
488	tracing::warn!(%error, "sync engine ended with an error; keeping its session");
489	Retain::RecordAndStore
490	}
491	Err(error) => {
492	tracing::warn!(%error, "sync engine task did not finish cleanly; keeping its session");
493	Retain::RecordAndStore
494	}
495	};
496	if let Some(mirror_task) = user.mirror_task.lock().await.take() {
497	let _ = mirror_task.await;
498	}
499	state.remove_session(&user.sid, retain).await;
500	}
501
502	fn new_sid() -> String {
503	let mut bytes = [0u8; 32];
504	OsRng.fill_bytes(&mut bytes);
505	hex::encode(bytes)
506	}
507
508	/// Load the 32-byte master key, creating it with owner-only permissions when
509	/// absent. Secrets must never be world-readable.
510	pub fn load_or_create_master_key(data_dir: &Path) -> Result<[u8; 32], String> {
511	create_private_dir(data_dir).map_err(\|error\| format!("creating data dir: {error}"))?;
512	restrict_dir(data_dir);
513	let path = data_dir.join("master.key");
514	if path.exists() {
515	let bytes = std::fs::read(&path).map_err(\|error\| format!("reading master key: {error}"))?;
516	restrict_file(&path);
517	if bytes.len() != 32 {
518	return Err("master key has an unexpected length".to_owned());
519	}
520	let mut key = [0u8; 32];
521	key.copy_from_slice(&bytes);
522	return Ok(key);
523	}
524	let mut key = [0u8; 32];
525	OsRng.fill_bytes(&mut key);
526	let mut file =
527	create_private_file(&path).map_err(\|error\| format!("creating master key: {error}"))?;
528	if let Err(error) = std::io::Write::write_all(&mut file, &key).and_then(\|_\| file.sync_all()) {
529	let _ = std::fs::remove_file(&path);
530	return Err(format!("writing master key: {error}"));
531	}
532	Ok(key)
533	}
534
535	fn create_private_dir(path: &Path) -> std::io::Result<()> {
536	let mut builder = std::fs::DirBuilder::new();
537	builder.recursive(true);
538	#[cfg(unix)]
539	{
540	use std::os::unix::fs::DirBuilderExt;
541	builder.mode(0o700);
542	}
543	builder.create(path)
544	}
545
546	fn create_private_file(path: &Path) -> std::io::Result<std::fs::File> {
547	let mut options = std::fs::OpenOptions::new();
548	options.write(true).create_new(true);
549	#[cfg(unix)]
550	{
551	use std::os::unix::fs::OpenOptionsExt;
552	options.mode(0o600);
553	}
554	options.open(path)
555	}
556
557	pub fn ensure_dirs(data_dir: &Path) {
558	for dir in [
559	data_dir.to_path_buf(),
560	data_dir.join("sessions"),
561	data_dir.join("stores"),
562	] {
563	let _ = create_private_dir(&dir);
564	restrict_dir(&dir);
565	}
566	}
567
568	#[cfg(unix)]
569	fn restrict_dir(path: &Path) {
570	use std::os::unix::fs::PermissionsExt;
571	let _ = std::fs::set_permissions(path, std::fs::Permissions::from_mode(0o700));
572	}
573
574	#[cfg(not(unix))]
575	fn restrict_dir(_path: &Path) {}
576
577	#[cfg(unix)]
578	fn restrict_file(path: &Path) {
579	use std::os::unix::fs::PermissionsExt;
580	let _ = std::fs::set_permissions(path, std::fs::Permissions::from_mode(0o600));
581	}
582
583	#[cfg(not(unix))]
584	fn restrict_file(_path: &Path) {}
585
586	#[cfg(test)]
587	mod tests {
588	use super::*;
589
590	fn state() -> (AppState, tempfile::TempDir) {
591	let dir = tempfile::tempdir().expect("temp dir");
592	(
593	AppState::new(dir.path().to_owned(), [7; 32], false, false),
594	dir,
595	)
596	}
597
598	#[tokio::test]
599	async fn sid_locks_only_serialise_the_same_sid() {
600	let (state, _dir) = state();
601	let held = state.lock_sid("alpha").await;
602	let other = timeout(Duration::from_millis(50), state.lock_sid("beta")).await;
603	assert!(other.is_ok(), "a different sid must not wait");
604	let same = timeout(Duration::from_millis(50), state.lock_sid("alpha")).await;
605	assert!(same.is_err(), "the same sid must wait for the holder");
606	drop(held);
607	let again = timeout(Duration::from_millis(50), state.lock_sid("alpha")).await;
608	assert!(again.is_ok(), "the sid is free once the holder drops");
609	}
610
611	#[tokio::test]
612	async fn sid_locks_forget_released_sids() {
613	let (state, _dir) = state();
614	drop(state.lock_sid("alpha").await);
615	drop(state.lock_sid("beta").await);
616	let live = state.inner.sid_locks.lock().unwrap().len();
617	assert!(live <= 1, "released locks are pruned on the next lookup");
618	}
619	}