malice

public

Star0

Watch0

68 commits 2 branches 0 tags 0 releases4.4 MiB

use std::{
    collections::HashMap,
    path::{Path, PathBuf},
    sync::Arc,
    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, 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 concurrent restores of the same sid so two engines never
    /// open the same SQLite store at once.
    init_lock: 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()),
                init_lock: Mutex::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
    }

    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.inner.init_lock.lock().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.inner.init_lock.lock().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.inner.init_lock.lock().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.inner.init_lock.lock().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::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) {}

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