malice

public

Star0

Watch0

1 commit 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, 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 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 => 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>,
        homeserver: String,
        username: String,
        password: String,
    ) -> Result<Arc<UserSession>, String> {
        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), sessions.clone()).await?;
                let _init = self.inner.init_lock.lock().await;
                return self
                    .start_user_session(sid.to_owned(), session, sessions)
                    .await;
            }
            self.apply_retention(&self.session_store(sid), Retain::Nothing)
                .await;
        }
        let sid = new_sid();
        let sessions = self.session_store(&sid);
        let session = login_with(
            PasswordLogin::new(homeserver, username, password)
                .with_device_display_name("malice web"),
            sessions.clone(),
        )
        .await?;
        let _init = self.inner.init_lock.lock().await;
        self.start_user_session(sid, session, sessions).await
    }

    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())
        }
        error => RestoreError::Stale(format!("restoring persisted session failed: {error}")),
    }
}

async fn login_with(
    login: PasswordLogin,
    sessions: Arc<dyn SessionStore>,
) -> Result<Session, String> {
    match Session::login_password(login, sessions).await {
        Ok(session) => Ok(session),
        Err(AccountError::StoreInUse) => Err("another engine is using this session".to_owned()),
        Err(AccountError::Unsaved { session, .. }) => session
            .save()
            .await
            .map_err(|_| "could not save the session".to_owned()),
        Err(error) => Err(format!(
            "login failed: {}",
            malice_core::login::describe_error(&error)
        )),
    }
}

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