fix: Properly handle appservice device creation

changelog.d/+39fcec3a.bugfix.md

@@ -0,0 +1 @@
+Appservices are now properly able to create devices for E2EE.

src/api/client/account/mod.rs

@@ -26,7 +26,7 @@
 };
 use service::{mailer::messages, uiaa::UiaaInitiator, users::HashedPassword};
 
-use super::{DEVICE_ID_LENGTH, TOKEN_LENGTH};
+use super::DEVICE_ID_LENGTH;
 use crate::{Ruma, router::ClientIdentity};
 
 pub(crate) mod register;

src/api/client/account/register.rs

@@ -20,9 +20,12 @@
 	assign,
 };
 use serde_json::value::RawValue;
-use service::{mailer::messages, users::HashedPassword};
+use service::{
+	mailer::messages,
+	users::{DeviceToken, HashedPassword},
+};
 
-use super::{DEVICE_ID_LENGTH, TOKEN_LENGTH};
+use super::DEVICE_ID_LENGTH;
 use crate::Ruma;
 
 /// # `POST /_matrix/client/v3/register`
@@ -119,7 +122,7 @@ pub(crate) async fn register_route(
 			.unwrap_or_else(|| utils::random_string(DEVICE_ID_LENGTH).into());
 
 		// Generate new token for the device
-		let new_token = utils::random_string(TOKEN_LENGTH);
+		let new_token = DeviceToken::new_random();
 
 		// Create device for this account
 		services
@@ -127,8 +130,7 @@ pub(crate) async fn register_route(
 			.create_device(
 				&user_id,
 				&device_id,
-				&new_token,
-				None,
+				Some(new_token.clone()),
 				body.initial_device_display_name.clone(),
 				Some(client.to_string()),
 			)
@@ -142,7 +144,7 @@ pub(crate) async fn register_route(
 	debug_info!(%user_id, ?device, "New account created via legacy registration");
 
 	Ok(assign!(register::v3::Response::new(user_id), {
-		access_token: token,
+		access_token: token.map(DeviceToken::into_token),
 		device_id: device,
 		refresh_token: None,
 		expires_in: None,

src/api/client/device.rs

@@ -89,14 +89,7 @@ pub(crate) async fn update_device_route(
 
 			services
 				.users
-				.create_device(
-					sender_user,
-					&device_id,
-					&appservice.registration.as_token,
-					None,
-					None,
-					Some(client.to_string()),
-				)
+				.create_device(sender_user, &device_id, None, None, Some(client.to_string()))
 				.await?;
 
 			return Ok(update_device::v3::Response::new());

src/api/client/session.rs

@@ -29,8 +29,9 @@
 	},
 	assign,
 };
+use service::users::DeviceToken;
 
-use super::{DEVICE_ID_LENGTH, TOKEN_LENGTH};
+use super::DEVICE_ID_LENGTH;
 use crate::Ruma;
 
 /// # `GET /_matrix/client/v3/login`
@@ -196,8 +197,8 @@ pub(crate) async fn login_route(
 		.clone()
 		.unwrap_or_else(|| utils::random_string(DEVICE_ID_LENGTH).into());
 
-	// Generate a new token for the device (ensuring no collisions)
-	let token = services.users.generate_unique_token().await;
+	// Generate a new token for the device
+	let token = DeviceToken::new_random();
 
 	// Determine if device_id was provided and exists in the db for this user
 	let device_exists = if body.device_id.is_some() {
@@ -213,7 +214,7 @@ pub(crate) async fn login_route(
 	if device_exists {
 		services
 			.users
-			.set_token(&user_id, &device_id, &token, None)
+			.set_token(&user_id, &device_id, token.clone())
 			.await?;
 	} else {
 		services
@@ -221,8 +222,7 @@ pub(crate) async fn login_route(
 			.create_device(
 				&user_id,
 				&device_id,
-				&token,
-				None,
+				Some(token.clone()),
 				body.initial_device_display_name.clone(),
 				Some(client.to_string()),
 			)
@@ -241,7 +241,7 @@ pub(crate) async fn login_route(
 	info!("{user_id} logged in");
 
 	#[allow(deprecated)]
-	Ok(assign!(login::v3::Response::new(user_id, token, device_id), {
+	Ok(assign!(login::v3::Response::new(user_id, token.into_token(), device_id), {
 		well_known: client_discovery_info,
 		expires_in: None,
 		home_server: Some(services.config.server_name.clone()),
@@ -273,7 +273,7 @@ pub(crate) async fn login_token_route(
 		.authenticate_password(&body.auth, sender_user, body.identity.sender_device(), None)
 		.await?;
 
-	let login_token = utils::random_string(TOKEN_LENGTH);
+	let login_token = DeviceToken::new_random().into_token();
 	let expires_in = services.users.create_login_token(sender_user, &login_token);
 
 	Ok(get_login_token::v1::Response::new(

src/service/oauth/mod.rs

@@ -25,7 +25,7 @@
 			OAuthError, ResponseMode, Scope, TokenRequest, TokenResponse, TokenType,
 		},
 	},
-	users,
+	users::{self, DeviceToken},
 };
 
 pub mod client_metadata;
@@ -343,7 +343,7 @@ async fn create_session(
 		client_name: Option<String>,
 		client_id: String,
 	) -> Result<TokenResponse, OAuthError> {
-		let access_token = Self::generate_token();
+		let access_token = DeviceToken::new_random().with_max_age(Self::ACCESS_TOKEN_MAX_AGE);
 		let refresh_token = Self::generate_token();
 
 		let device_id = requested_scopes
@@ -376,8 +376,7 @@ async fn create_session(
 			.create_device(
 				&authorizing_user,
 				device_id,
-				&access_token,
-				Some(Self::ACCESS_TOKEN_MAX_AGE),
+				Some(access_token.clone()),
 				client_name,
 				None,
 			)
@@ -413,7 +412,7 @@ async fn create_session(
 		);
 
 		Ok(TokenResponse {
-			access_token,
+			access_token: access_token.into_token(),
 			token_type: TokenType::Bearer,
 			expires_in: Self::ACCESS_TOKEN_MAX_AGE.as_secs(),
 			scope: requested_scopes.iter().join(" "),
@@ -449,7 +448,7 @@ async fn refresh_session(
 
 		assert_eq!(&client_id, &session_info.client_id, "session info client id mismatch");
 
-		let new_access_token = Self::generate_token();
+		let new_access_token = DeviceToken::new_random().with_max_age(Self::ACCESS_TOKEN_MAX_AGE);
 		let new_refresh_token = Self::generate_token();
 		let scope = session_info.scopes.iter().join(" ");
 		session_info
@@ -461,8 +460,7 @@ async fn refresh_session(
 			.set_token(
 				&refresh_token_info.user_id,
 				&refresh_token_info.device_id,
-				&new_access_token,
-				Some(Self::ACCESS_TOKEN_MAX_AGE),
+				new_access_token.clone(),
 			)
 			.await
 			.expect("should be able to set token");
@@ -479,7 +477,7 @@ async fn refresh_session(
 			.raw_put(&new_refresh_token, Json(refresh_token_info));
 
 		Ok(TokenResponse {
-			access_token: new_access_token,
+			access_token: new_access_token.into_token(),
 			token_type: TokenType::Bearer,
 			expires_in: Self::ACCESS_TOKEN_MAX_AGE.as_secs(),
 			scope,

src/service/users/dehydrated_device.rs

@@ -41,7 +41,6 @@ pub async fn set_dehydrated_device(&self, user_id: &UserId, request: Request) ->
 		self.create_device(
 			user_id,
 			&request.device_id,
-			"",
 			None,
 			request.initial_device_display_name.clone(),
 			None,

src/service/users/device.rs

@@ -11,20 +11,44 @@
 use futures::{Stream, StreamExt};
 use ruma::{
 	DeviceId, MilliSecondsSinceUnixEpoch, OwnedDeviceId, OwnedUserId, UserId,
-	api::client::device::Device, events::AnyToDeviceEvent, serde::Raw, uint,
+	api::client::device::Device, assign, events::AnyToDeviceEvent, serde::Raw, uint,
 };
 use serde_json::json;
 
 use crate::users::increment;
 
+#[must_use]
+#[derive(Clone)]
+pub struct DeviceToken {
+	token: String,
+	max_age: Option<Duration>,
+}
+
+impl DeviceToken {
+	const DEVICE_ID_LENGTH: usize = 10;
+
+	pub fn new(token: String) -> Self { Self { token, max_age: None } }
+
+	pub fn new_random() -> Self { Self::new(utils::random_string(Self::DEVICE_ID_LENGTH)) }
+
+	pub fn with_max_age(self, max_age: Duration) -> Self {
+		assign!(self, { max_age: Some(max_age) })
+	}
+
+	#[must_use]
+	pub fn into_token(self) -> String { self.token }
+}
+
 impl super::Service {
 	/// Adds a new device to a user.
+	///
+	/// If no `token` is provided, the device will not be able to be logged
+	/// into.
 	pub async fn create_device(
 		&self,
 		user_id: &UserId,
 		device_id: &DeviceId,
-		token: &str,
-		token_max_age: Option<Duration>,
+		token: Option<DeviceToken>,
 		initial_device_display_name: Option<String>,
 		client_ip: Option<String>,
 	) -> Result<()> {
@@ -38,8 +62,12 @@ pub async fn create_device(
 
 		increment(&self.db.userid_devicelistversion, user_id.as_bytes());
 		self.db.userdeviceid_metadata.put(key, Json(device));
-		self.set_token(user_id, device_id, token, token_max_age)
-			.await
+
+		if let Some(token) = token {
+			self.set_token(user_id, device_id, token).await?;
+		}
+
+		Ok(())
 	}
 
 	/// Removes a device from a user.
@@ -97,31 +125,12 @@ pub async fn get_token(&self, user_id: &UserId, device_id: &DeviceId) -> Result<
 		self.db.userdeviceid_token.qry(&key).await.deserialized()
 	}
 
-	/// Generate a unique access token that doesn't collide with existing tokens
-	pub async fn generate_unique_token(&self) -> String {
-		loop {
-			let token = utils::random_string(32);
-
-			// Check for collision with existing appservice and user tokens
-			let (appservice, usr) = tokio::join!(
-				self.services.appservice.find_from_token(&token),
-				self.db.token_userdeviceid.get(&token)
-			);
-			if appservice.is_ok() || usr.is_ok() {
-				continue;
-			}
-
-			return token;
-		}
-	}
-
 	/// Replaces the access token of one device.
 	pub async fn set_token(
 		&self,
 		user_id: &UserId,
 		device_id: &DeviceId,
-		token: &str,
-		token_max_age: Option<Duration>,
+		DeviceToken { token, max_age }: DeviceToken,
 	) -> Result<()> {
 		let key = (user_id, device_id);
 		if self.db.userdeviceid_metadata.qry(&key).await.is_err() {
@@ -136,7 +145,7 @@ pub async fn set_token(
 		if self
 			.services
 			.appservice
-			.find_from_token(token)
+			.find_from_token(&token)
 			.await
 			.is_ok()
 		{
@@ -153,10 +162,10 @@ pub async fn set_token(
 		}
 
 		// Assign token to user device combination
-		self.db.userdeviceid_token.put_raw(key, token);
-		self.db.token_userdeviceid.raw_put(token, key);
+		self.db.userdeviceid_token.put_raw(key, &token);
+		self.db.token_userdeviceid.raw_put(&token, key);
 
-		if let Some(max_age) = token_max_age {
+		if let Some(max_age) = max_age {
 			let expires = SystemTime::now()
 				.duration_since(SystemTime::UNIX_EPOCH)
 				.expect("system time should not be before the epoch")

src/service/users/mod.rs

@@ -14,6 +14,7 @@
 	utils::{self},
 };
 use database::Map;
+pub use device::DeviceToken;
 pub use profile::ProfileFieldChange;
 use ruma::{UserId, api::error::ErrorKind, encryption::CrossSigningKey, serde::Raw};
 use serde::{Deserialize, Serialize};