malice

public

Star0

Watch0

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

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(Vec::new())).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;
}

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