macltok

public

Star0

Watch0

6 commits 3 branches 0 tags 0 releases3.0 MiB

use std::{
    collections::HashMap,
    path::{Path, PathBuf},
};

use matrix_sdk::{
    Client, ClientBuilder, ThreadingSupport,
    authentication::{
        AuthSession,
        matrix::MatrixSession,
        oauth::{
            ClientId, ClientRegistrationData, CsrfToken, OAuthSession, UserSession,
            registration::{ApplicationType, ClientMetadata, Localized, OAuthGrantType},
        },
    },
    encryption::{BackupDownloadStrategy, EncryptionSettings},
    ruma::{OwnedDeviceId, serde::Raw},
};
use serde::{Deserialize, Serialize};
use thiserror::Error;
use url::Url;

use crate::{ClientHandle, Engine, EngineConfig, EngineError, EventStream};

pub struct Session {
    client: Client,
    record: SessionRecord,
}

#[non_exhaustive]
pub struct PasswordLogin {
    pub homeserver: String,
    pub username: String,
    pub password: String,
    pub device_display_name: Option<String>,
    /// set after a soft logout so the store's keys still belong to us
    pub device_id: Option<OwnedDeviceId>,
}

impl PasswordLogin {
    #[must_use]
    pub fn new(
        homeserver: impl Into<String>,
        username: impl Into<String>,
        password: impl Into<String>,
    ) -> Self {
        Self {
            homeserver: homeserver.into(),
            username: username.into(),
            password: password.into(),
            device_display_name: None,
            device_id: None,
        }
    }

    #[must_use]
    pub fn with_device_display_name(mut self, name: impl Into<String>) -> Self {
        self.device_display_name = Some(name.into());
        self
    }

    #[must_use]
    pub fn with_device_id(mut self, device_id: OwnedDeviceId) -> Self {
        self.device_id = Some(device_id);
        self
    }
}

#[non_exhaustive]
pub struct OAuthLogin {
    pub homeserver: String,
    pub redirect_uri: String,
    pub client_uri: String,
    pub client_name: Option<String>,
    pub client_id: Option<String>,
    pub device_id: Option<OwnedDeviceId>,
    pub application_type: OAuthApplicationType,
}

impl OAuthLogin {
    #[must_use]
    pub fn new(
        homeserver: impl Into<String>,
        redirect_uri: impl Into<String>,
        client_uri: impl Into<String>,
        application_type: OAuthApplicationType,
    ) -> Self {
        Self {
            homeserver: homeserver.into(),
            redirect_uri: redirect_uri.into(),
            client_uri: client_uri.into(),
            client_name: None,
            client_id: None,
            device_id: None,
            application_type,
        }
    }

    #[must_use]
    pub fn with_client_name(mut self, name: impl Into<String>) -> Self {
        self.client_name = Some(name.into());
        self
    }

    /// a statically registered client id, skipping dynamic registration
    #[must_use]
    pub fn with_client_id(mut self, client_id: impl Into<String>) -> Self {
        self.client_id = Some(client_id.into());
        self
    }

    #[must_use]
    pub fn with_device_id(mut self, device_id: OwnedDeviceId) -> Self {
        self.device_id = Some(device_id);
        self
    }
}

#[derive(Clone, Copy, Debug, Eq, PartialEq)]
#[non_exhaustive]
pub enum OAuthApplicationType {
    Native,
    Web,
}

pub struct PendingOAuthLogin {
    client: Client,
    authorization_url: String,
    state: Option<CsrfToken>,
}

impl PendingOAuthLogin {
    #[must_use]
    pub fn authorization_url(&self) -> &str {
        &self.authorization_url
    }

    #[must_use]
    pub fn registered_client_id(&self) -> String {
        self.client
            .oauth()
            .client_id()
            .expect("OAuth client was registered before authorization")
            .as_str()
            .to_owned()
    }

    #[must_use]
    pub const fn is_pending(&self) -> bool {
        self.state.is_some()
    }

    pub async fn finish(&mut self, callback_url: &str) -> Result<Session, AccountError> {
        let callback_url = Url::parse(callback_url)
            .map_err(|error| account_error("parsing the OAuth callback URL", error))?;
        let state = self
            .state
            .take()
            .ok_or(AccountError::OAuthCallbackAlreadySubmitted)?;
        let oauth = self.client.oauth();
        if let Err(error) = oauth.finish_login(callback_url.into()).await {
            oauth.abort_login(&state).await;
            return Err(account_error("OAuth login", error));
        }
        Session::from_client(self.client.clone())
    }

    pub async fn abort(mut self) {
        if let Some(state) = self.state.take() {
            self.client.oauth().abort_login(&state).await;
        }
    }
}

pub struct StoreConfig {
    path: PathBuf,
    passphrase: String,
}

impl StoreConfig {
    #[must_use]
    pub fn new(path: impl Into<PathBuf>, passphrase: impl Into<String>) -> Self {
        Self {
            path: path.into(),
            passphrase: passphrase.into(),
        }
    }

    #[must_use]
    pub fn path(&self) -> &Path {
        &self.path
    }
}

/// Serializable authentication secrets. Store this separately from the
/// encrypted SQLite store in protected, preferably encrypted storage.
/// A saved record must be replaced when `Event::SessionTokensChanged` arrives.
#[derive(Clone, Deserialize, Serialize)]
pub struct SessionRecord {
    homeserver: String,
    authentication: AuthenticationRecord,
}

#[derive(Clone, Deserialize, Serialize)]
#[serde(tag = "type", content = "session", rename_all = "snake_case")]
enum AuthenticationRecord {
    Matrix(MatrixSession),
    OAuth {
        client_id: ClientId,
        user: UserSession,
    },
}

impl SessionRecord {
    #[must_use]
    pub fn homeserver(&self) -> &str {
        &self.homeserver
    }
}

#[derive(Debug, Error)]
pub enum AccountError {
    #[error("{context} failed")]
    Operation {
        context: &'static str,
        #[source]
        source: Box<dyn std::error::Error + Send + Sync>,
    },
    #[error("the homeserver returned no authenticated session")]
    MissingSession,
    #[error("the SDK returned an unsupported authentication session")]
    UnsupportedAuthentication,
    #[error("the OAuth callback was already submitted")]
    OAuthCallbackAlreadySubmitted,
    #[error(transparent)]
    Engine(#[from] EngineError),
}

impl Session {
    pub async fn login_password(
        login: PasswordLogin,
        store: StoreConfig,
    ) -> Result<Self, AccountError> {
        let client = matrix_client_builder(&store)
            .server_name_or_homeserver_url(login.homeserver.trim())
            .build()
            .await
            .map_err(|error| account_error("building the Matrix client", error))?;
        let mut request = client
            .matrix_auth()
            .login_username(login.username.trim(), &login.password)
            .request_refresh_token();
        if let Some(name) = login.device_display_name.as_deref() {
            request = request.initial_device_display_name(name);
        }
        if let Some(device_id) = login.device_id.as_deref() {
            request = request.device_id(device_id.as_str());
        }
        request
            .send()
            .await
            .map_err(|error| account_error("password login", error))?;
        Self::from_client(client)
    }

    pub async fn begin_oauth(
        login: OAuthLogin,
        store: StoreConfig,
    ) -> Result<PendingOAuthLogin, AccountError> {
        let redirect_uri = Url::parse(login.redirect_uri.trim())
            .map_err(|error| account_error("parsing the OAuth redirect URI", error))?;
        let client_uri = Url::parse(login.client_uri.trim())
            .map_err(|error| account_error("parsing the OAuth client URI", error))?;
        let client = matrix_client_builder(&store)
            .server_name_or_homeserver_url(login.homeserver.trim())
            .build()
            .await
            .map_err(|error| account_error("building the Matrix client", error))?;
        let application_type = match login.application_type {
            OAuthApplicationType::Native => ApplicationType::Native,
            OAuthApplicationType::Web => ApplicationType::Web,
        };
        let mut metadata = ClientMetadata::new(
            application_type,
            vec![OAuthGrantType::AuthorizationCode {
                redirect_uris: vec![redirect_uri.clone()],
            }],
            Localized::new(client_uri, []),
        );
        metadata.client_name = login.client_name.map(|name| Localized::new(name, []));
        let metadata = Raw::new(&metadata)
            .map_err(|error| account_error("serializing OAuth client metadata", error))?;
        let mut registration = ClientRegistrationData::new(metadata);
        if let Some(client_id) = login.client_id {
            registration.static_registrations = Some(HashMap::from([(
                client.homeserver(),
                ClientId::new(client_id),
            )]));
        }
        let authorization = client
            .oauth()
            .login(redirect_uri, login.device_id, Some(registration), None)
            .build()
            .await
            .map_err(|error| account_error("starting OAuth login", error))?;
        Ok(PendingOAuthLogin {
            client,
            authorization_url: authorization.url.to_string(),
            state: Some(authorization.state),
        })
    }

    /// Restore a saved record with the same store location and passphrase.
    /// The caller must also handle refresh, rejection, and explicit logout.
    pub async fn restore(record: SessionRecord, store: StoreConfig) -> Result<Self, AccountError> {
        let client = matrix_client_builder(&store)
            .homeserver_url(&record.homeserver)
            .build()
            .await
            .map_err(|error| account_error("reopening the Matrix store", error))?;
        let authentication = match record.authentication.clone() {
            AuthenticationRecord::Matrix(session) => AuthSession::Matrix(session),
            AuthenticationRecord::OAuth { client_id, user } => {
                OAuthSession { client_id, user }.into()
            }
        };
        client
            .restore_session(authentication)
            .await
            .map_err(|error| account_error("restoring the authenticated session", error))?;
        Ok(Self { client, record })
    }

    /// Initial login credentials only. This snapshot does not track later
    /// token rotations; use `ClientHandle::session_record` after building.
    #[must_use]
    pub fn record(&self) -> SessionRecord {
        self.record.clone()
    }

    pub async fn build(
        self,
        config: EngineConfig,
    ) -> Result<(Engine, ClientHandle, EventStream), AccountError> {
        Engine::from_client(self.client, config)
            .await
            .map_err(AccountError::from)
    }

    fn from_client(client: Client) -> Result<Self, AccountError> {
        let record = session_record(&client)?;
        Ok(Self { client, record })
    }
}

pub(crate) fn session_record(client: &Client) -> Result<SessionRecord, AccountError> {
    let authentication = match client.session().ok_or(AccountError::MissingSession)? {
        AuthSession::Matrix(session) => AuthenticationRecord::Matrix(session),
        AuthSession::OAuth(session) => {
            let OAuthSession { client_id, user } = *session;
            AuthenticationRecord::OAuth { client_id, user }
        }
        _ => return Err(AccountError::UnsupportedAuthentication),
    };
    Ok(SessionRecord {
        homeserver: client.homeserver().to_string(),
        authentication,
    })
}

fn encryption_settings() -> EncryptionSettings {
    EncryptionSettings {
        auto_enable_cross_signing: true,
        backup_download_strategy: BackupDownloadStrategy::OneShot,
        auto_enable_backups: true,
    }
}

fn matrix_client_builder(store: &StoreConfig) -> ClientBuilder {
    Client::builder()
        .with_encryption_settings(encryption_settings())
        .with_threading_support(ThreadingSupport::Enabled {
            with_subscriptions: false,
        })
        .handle_refresh_tokens()
        .sqlite_store(&store.path, Some(&store.passphrase))
}

fn account_error(
    context: &'static str,
    source: impl std::error::Error + Send + Sync + 'static,
) -> AccountError {
    AccountError::Operation {
        context,
        source: Box::new(source),
    }
}

1	use std::{
2	collections::HashMap,
3	path::{Path, PathBuf},
4	};
5
6	use matrix_sdk::{
7	Client, ClientBuilder, ThreadingSupport,
8	authentication::{
9	AuthSession,
10	matrix::MatrixSession,
11	oauth::{
12	ClientId, ClientRegistrationData, CsrfToken, OAuthSession, UserSession,
13	registration::{ApplicationType, ClientMetadata, Localized, OAuthGrantType},
14	},
15	},
16	encryption::{BackupDownloadStrategy, EncryptionSettings},
17	ruma::{OwnedDeviceId, serde::Raw},
18	};
19	use serde::{Deserialize, Serialize};
20	use thiserror::Error;
21	use url::Url;
22
23	use crate::{ClientHandle, Engine, EngineConfig, EngineError, EventStream};
24
25	pub struct Session {
26	client: Client,
27	record: SessionRecord,
28	}
29
30	#[non_exhaustive]
31	pub struct PasswordLogin {
32	pub homeserver: String,
33	pub username: String,
34	pub password: String,
35	pub device_display_name: Option<String>,
36	/// set after a soft logout so the store's keys still belong to us
37	pub device_id: Option<OwnedDeviceId>,
38	}
39
40	impl PasswordLogin {
41	#[must_use]
42	pub fn new(
43	homeserver: impl Into<String>,
44	username: impl Into<String>,
45	password: impl Into<String>,
46	) -> Self {
47	Self {
48	homeserver: homeserver.into(),
49	username: username.into(),
50	password: password.into(),
51	device_display_name: None,
52	device_id: None,
53	}
54	}
55
56	#[must_use]
57	pub fn with_device_display_name(mut self, name: impl Into<String>) -> Self {
58	self.device_display_name = Some(name.into());
59	self
60	}
61
62	#[must_use]
63	pub fn with_device_id(mut self, device_id: OwnedDeviceId) -> Self {
64	self.device_id = Some(device_id);
65	self
66	}
67	}
68
69	#[non_exhaustive]
70	pub struct OAuthLogin {
71	pub homeserver: String,
72	pub redirect_uri: String,
73	pub client_uri: String,
74	pub client_name: Option<String>,
75	pub client_id: Option<String>,
76	pub device_id: Option<OwnedDeviceId>,
77	pub application_type: OAuthApplicationType,
78	}
79
80	impl OAuthLogin {
81	#[must_use]
82	pub fn new(
83	homeserver: impl Into<String>,
84	redirect_uri: impl Into<String>,
85	client_uri: impl Into<String>,
86	application_type: OAuthApplicationType,
87	) -> Self {
88	Self {
89	homeserver: homeserver.into(),
90	redirect_uri: redirect_uri.into(),
91	client_uri: client_uri.into(),
92	client_name: None,
93	client_id: None,
94	device_id: None,
95	application_type,
96	}
97	}
98
99	#[must_use]
100	pub fn with_client_name(mut self, name: impl Into<String>) -> Self {
101	self.client_name = Some(name.into());
102	self
103	}
104
105	/// a statically registered client id, skipping dynamic registration
106	#[must_use]
107	pub fn with_client_id(mut self, client_id: impl Into<String>) -> Self {
108	self.client_id = Some(client_id.into());
109	self
110	}
111
112	#[must_use]
113	pub fn with_device_id(mut self, device_id: OwnedDeviceId) -> Self {
114	self.device_id = Some(device_id);
115	self
116	}
117	}
118
119	#[derive(Clone, Copy, Debug, Eq, PartialEq)]
120	#[non_exhaustive]
121	pub enum OAuthApplicationType {
122	Native,
123	Web,
124	}
125
126	pub struct PendingOAuthLogin {
127	client: Client,
128	authorization_url: String,
129	state: Option<CsrfToken>,
130	}
131
132	impl PendingOAuthLogin {
133	#[must_use]
134	pub fn authorization_url(&self) -> &str {
135	&self.authorization_url
136	}
137
138	#[must_use]
139	pub fn registered_client_id(&self) -> String {
140	self.client
141	.oauth()
142	.client_id()
143	.expect("OAuth client was registered before authorization")
144	.as_str()
145	.to_owned()
146	}
147
148	#[must_use]
149	pub const fn is_pending(&self) -> bool {
150	self.state.is_some()
151	}
152
153	pub async fn finish(&mut self, callback_url: &str) -> Result<Session, AccountError> {
154	let callback_url = Url::parse(callback_url)
155	.map_err(\|error\| account_error("parsing the OAuth callback URL", error))?;
156	let state = self
157	.state
158	.take()
159	.ok_or(AccountError::OAuthCallbackAlreadySubmitted)?;
160	let oauth = self.client.oauth();
161	if let Err(error) = oauth.finish_login(callback_url.into()).await {
162	oauth.abort_login(&state).await;
163	return Err(account_error("OAuth login", error));
164	}
165	Session::from_client(self.client.clone())
166	}
167
168	pub async fn abort(mut self) {
169	if let Some(state) = self.state.take() {
170	self.client.oauth().abort_login(&state).await;
171	}
172	}
173	}
174
175	pub struct StoreConfig {
176	path: PathBuf,
177	passphrase: String,
178	}
179
180	impl StoreConfig {
181	#[must_use]
182	pub fn new(path: impl Into<PathBuf>, passphrase: impl Into<String>) -> Self {
183	Self {
184	path: path.into(),
185	passphrase: passphrase.into(),
186	}
187	}
188
189	#[must_use]
190	pub fn path(&self) -> &Path {
191	&self.path
192	}
193	}
194
195	/// Serializable authentication secrets. Store this separately from the
196	/// encrypted SQLite store in protected, preferably encrypted storage.
197	/// A saved record must be replaced when `Event::SessionTokensChanged` arrives.
198	#[derive(Clone, Deserialize, Serialize)]
199	pub struct SessionRecord {
200	homeserver: String,
201	authentication: AuthenticationRecord,
202	}
203
204	#[derive(Clone, Deserialize, Serialize)]
205	#[serde(tag = "type", content = "session", rename_all = "snake_case")]
206	enum AuthenticationRecord {
207	Matrix(MatrixSession),
208	OAuth {
209	client_id: ClientId,
210	user: UserSession,
211	},
212	}
213
214	impl SessionRecord {
215	#[must_use]
216	pub fn homeserver(&self) -> &str {
217	&self.homeserver
218	}
219	}
220
221	#[derive(Debug, Error)]
222	pub enum AccountError {
223	#[error("{context} failed")]
224	Operation {
225	context: &'static str,
226	#[source]
227	source: Box<dyn std::error::Error + Send + Sync>,
228	},
229	#[error("the homeserver returned no authenticated session")]
230	MissingSession,
231	#[error("the SDK returned an unsupported authentication session")]
232	UnsupportedAuthentication,
233	#[error("the OAuth callback was already submitted")]
234	OAuthCallbackAlreadySubmitted,
235	#[error(transparent)]
236	Engine(#[from] EngineError),
237	}
238
239	impl Session {
240	pub async fn login_password(
241	login: PasswordLogin,
242	store: StoreConfig,
243	) -> Result<Self, AccountError> {
244	let client = matrix_client_builder(&store)
245	.server_name_or_homeserver_url(login.homeserver.trim())
246	.build()
247	.await
248	.map_err(\|error\| account_error("building the Matrix client", error))?;
249	let mut request = client
250	.matrix_auth()
251	.login_username(login.username.trim(), &login.password)
252	.request_refresh_token();
253	if let Some(name) = login.device_display_name.as_deref() {
254	request = request.initial_device_display_name(name);
255	}
256	if let Some(device_id) = login.device_id.as_deref() {
257	request = request.device_id(device_id.as_str());
258	}
259	request
260	.send()
261	.await
262	.map_err(\|error\| account_error("password login", error))?;
263	Self::from_client(client)
264	}
265
266	pub async fn begin_oauth(
267	login: OAuthLogin,
268	store: StoreConfig,
269	) -> Result<PendingOAuthLogin, AccountError> {
270	let redirect_uri = Url::parse(login.redirect_uri.trim())
271	.map_err(\|error\| account_error("parsing the OAuth redirect URI", error))?;
272	let client_uri = Url::parse(login.client_uri.trim())
273	.map_err(\|error\| account_error("parsing the OAuth client URI", error))?;
274	let client = matrix_client_builder(&store)
275	.server_name_or_homeserver_url(login.homeserver.trim())
276	.build()
277	.await
278	.map_err(\|error\| account_error("building the Matrix client", error))?;
279	let application_type = match login.application_type {
280	OAuthApplicationType::Native => ApplicationType::Native,
281	OAuthApplicationType::Web => ApplicationType::Web,
282	};
283	let mut metadata = ClientMetadata::new(
284	application_type,
285	vec![OAuthGrantType::AuthorizationCode {
286	redirect_uris: vec![redirect_uri.clone()],
287	}],
288	Localized::new(client_uri, []),
289	);
290	metadata.client_name = login.client_name.map(\|name\| Localized::new(name, []));
291	let metadata = Raw::new(&metadata)
292	.map_err(\|error\| account_error("serializing OAuth client metadata", error))?;
293	let mut registration = ClientRegistrationData::new(metadata);
294	if let Some(client_id) = login.client_id {
295	registration.static_registrations = Some(HashMap::from([(
296	client.homeserver(),
297	ClientId::new(client_id),
298	)]));
299	}
300	let authorization = client
301	.oauth()
302	.login(redirect_uri, login.device_id, Some(registration), None)
303	.build()
304	.await
305	.map_err(\|error\| account_error("starting OAuth login", error))?;
306	Ok(PendingOAuthLogin {
307	client,
308	authorization_url: authorization.url.to_string(),
309	state: Some(authorization.state),
310	})
311	}
312
313	/// Restore a saved record with the same store location and passphrase.
314	/// The caller must also handle refresh, rejection, and explicit logout.
315	pub async fn restore(record: SessionRecord, store: StoreConfig) -> Result<Self, AccountError> {
316	let client = matrix_client_builder(&store)
317	.homeserver_url(&record.homeserver)
318	.build()
319	.await
320	.map_err(\|error\| account_error("reopening the Matrix store", error))?;
321	let authentication = match record.authentication.clone() {
322	AuthenticationRecord::Matrix(session) => AuthSession::Matrix(session),
323	AuthenticationRecord::OAuth { client_id, user } => {
324	OAuthSession { client_id, user }.into()
325	}
326	};
327	client
328	.restore_session(authentication)
329	.await
330	.map_err(\|error\| account_error("restoring the authenticated session", error))?;
331	Ok(Self { client, record })
332	}
333
334	/// Initial login credentials only. This snapshot does not track later
335	/// token rotations; use `ClientHandle::session_record` after building.
336	#[must_use]
337	pub fn record(&self) -> SessionRecord {
338	self.record.clone()
339	}
340
341	pub async fn build(
342	self,
343	config: EngineConfig,
344	) -> Result<(Engine, ClientHandle, EventStream), AccountError> {
345	Engine::from_client(self.client, config)
346	.await
347	.map_err(AccountError::from)
348	}
349
350	fn from_client(client: Client) -> Result<Self, AccountError> {
351	let record = session_record(&client)?;
352	Ok(Self { client, record })
353	}
354	}
355
356	pub(crate) fn session_record(client: &Client) -> Result<SessionRecord, AccountError> {
357	let authentication = match client.session().ok_or(AccountError::MissingSession)? {
358	AuthSession::Matrix(session) => AuthenticationRecord::Matrix(session),
359	AuthSession::OAuth(session) => {
360	let OAuthSession { client_id, user } = *session;
361	AuthenticationRecord::OAuth { client_id, user }
362	}
363	_ => return Err(AccountError::UnsupportedAuthentication),
364	};
365	Ok(SessionRecord {
366	homeserver: client.homeserver().to_string(),
367	authentication,
368	})
369	}
370
371	fn encryption_settings() -> EncryptionSettings {
372	EncryptionSettings {
373	auto_enable_cross_signing: true,
374	backup_download_strategy: BackupDownloadStrategy::OneShot,
375	auto_enable_backups: true,
376	}
377	}
378
379	fn matrix_client_builder(store: &StoreConfig) -> ClientBuilder {
380	Client::builder()
381	.with_encryption_settings(encryption_settings())
382	.with_threading_support(ThreadingSupport::Enabled {
383	with_subscriptions: false,
384	})
385	.handle_refresh_tokens()
386	.sqlite_store(&store.path, Some(&store.passphrase))
387	}
388
389	fn account_error(
390	context: &'static str,
391	source: impl std::error::Error + Send + Sync + 'static,
392	) -> AccountError {
393	AccountError::Operation {
394	context,
395	source: Box::new(source),
396	}
397	}