feat: Implement Dehydrated Devices MSC3814

Co-authored-by: Jade Ellis <jade@ellis.link>
Signed-off-by: Jason Volk <jason@zemos.net>

src/service/sending/sender.rs

@@ -10,7 +10,7 @@ use std::{
 
 use base64::{Engine as _, engine::general_purpose::URL_SAFE_NO_PAD};
 use conduwuit_core::{
-	Error, Event, Result, debug, err, error,
+	Error, Event, Result, at, debug, err, error,
 	result::LogErr,
 	trace,
 	utils::{
@@ -175,7 +175,7 @@ impl Service {
 		if !new_events.is_empty() {
 			self.db.mark_as_active(new_events.iter());
 
-			let new_events_vec = new_events.into_iter().map(|(_, event)| event).collect();
+			let new_events_vec = new_events.into_iter().map(at!(1)).collect();
 			futures.push(self.send_events(dest.clone(), new_events_vec));
 		} else {
 			statuses.remove(dest);

src/service/users/dehydrated_device.rs

@@ -0,0 +1,149 @@
+use conduwuit::{Err, Result, implement, trace};
+use conduwuit_database::{Deserialized, Json};
+use ruma::{
+	DeviceId, OwnedDeviceId, UserId,
+	api::client::dehydrated_device::{
+		DehydratedDeviceData, put_dehydrated_device::unstable::Request,
+	},
+	serde::Raw,
+};
+use serde::{Deserialize, Serialize};
+
+#[derive(Clone, Debug, Serialize, Deserialize)]
+pub struct DehydratedDevice {
+	/// Unique ID of the device.
+	pub device_id: OwnedDeviceId,
+
+	/// Contains serialized and encrypted private data.
+	pub device_data: Raw<DehydratedDeviceData>,
+}
+
+/// Creates or recreates the user's dehydrated device.
+#[implement(super::Service)]
+#[tracing::instrument(
+	level = "info",
+	skip_all,
+	fields(
+		%user_id,
+		device_id = %request.device_id,
+		display_name = ?request.initial_device_display_name,
+	)
+)]
+pub async fn set_dehydrated_device(&self, user_id: &UserId, request: Request) -> Result {
+	assert!(
+		self.exists(user_id).await,
+		"Tried to create dehydrated device for non-existent user"
+	);
+
+	let existing_id = self.get_dehydrated_device_id(user_id).await;
+
+	if existing_id.is_err()
+		&& self
+			.get_device_metadata(user_id, &request.device_id)
+			.await
+			.is_ok()
+	{
+		return Err!("A hydrated device already exists with that ID.");
+	}
+
+	if let Ok(existing_id) = existing_id {
+		self.remove_device(user_id, &existing_id).await;
+	}
+
+	self.create_device(
+		user_id,
+		&request.device_id,
+		"",
+		request.initial_device_display_name.clone(),
+		None,
+	)
+	.await?;
+
+	trace!(device_data = ?request.device_data);
+	self.db.userid_dehydrateddevice.raw_put(
+		user_id,
+		Json(&DehydratedDevice {
+			device_id: request.device_id.clone(),
+			device_data: request.device_data,
+		}),
+	);
+
+	trace!(device_keys = ?request.device_keys);
+	self.add_device_keys(user_id, &request.device_id, &request.device_keys)
+		.await;
+
+	trace!(one_time_keys = ?request.one_time_keys);
+	for (one_time_key_key, one_time_key_value) in &request.one_time_keys {
+		self.add_one_time_key(user_id, &request.device_id, one_time_key_key, one_time_key_value)
+			.await?;
+	}
+
+	Ok(())
+}
+
+/// Removes a user's dehydrated device.
+///
+/// Calling this directly will remove the dehydrated data but leak the frontage
+/// device. Thus this is called by the regular device interface such that the
+/// dehydrated data will not leak instead.
+///
+/// If device_id is given, the user's dehydrated device must match or this is a
+/// no-op, but an Err is still returned to indicate that. Otherwise returns the
+/// removed dehydrated device_id.
+#[implement(super::Service)]
+#[tracing::instrument(
+	level = "debug",
+	skip_all,
+	fields(
+		%user_id,
+		device_id = ?maybe_device_id,
+	)
+)]
+pub(super) async fn remove_dehydrated_device(
+	&self,
+	user_id: &UserId,
+	maybe_device_id: Option<&DeviceId>,
+) -> Result<OwnedDeviceId> {
+	let Ok(device_id) = self.get_dehydrated_device_id(user_id).await else {
+		return Err!(Request(NotFound("No dehydrated device for this user.")));
+	};
+
+	if let Some(maybe_device_id) = maybe_device_id {
+		if maybe_device_id != device_id {
+			return Err!(Request(NotFound("Not the user's dehydrated device.")));
+		}
+	}
+
+	self.db.userid_dehydrateddevice.remove(user_id);
+
+	Ok(device_id)
+}
+
+/// Get the device_id of the user's dehydrated device.
+#[implement(super::Service)]
+#[tracing::instrument(
+	level = "debug",
+	skip_all,
+	fields(%user_id)
+)]
+pub async fn get_dehydrated_device_id(&self, user_id: &UserId) -> Result<OwnedDeviceId> {
+	self.get_dehydrated_device(user_id)
+		.await
+		.map(|device| device.device_id)
+}
+
+/// Get the dehydrated device private data
+#[implement(super::Service)]
+#[tracing::instrument(
+	level = "debug",
+	skip_all,
+	fields(%user_id),
+	ret,
+)]
+pub async fn get_dehydrated_device(&self, user_id: &UserId) -> Result<DehydratedDevice> {
+	self.db
+		.userid_dehydrateddevice
+		.get(user_id)
+		.await
+		.deserialized()
+}

src/service/users/mod.rs

@@ -1,3 +1,5 @@
+pub(super) mod dehydrated_device;
+
 #[cfg(feature = "ldap")]
 use std::collections::HashMap;
 use std::{collections::BTreeMap, mem, net::IpAddr, sync::Arc};
@@ -5,7 +7,7 @@ use std::{collections::BTreeMap, mem, net::IpAddr, sync::Arc};
 #[cfg(feature = "ldap")]
 use conduwuit::result::LogErr;
 use conduwuit::{
-	Err, Error, Result, Server, at, debug_warn, err, is_equal_to, trace,
+	Err, Error, Result, Server, debug_warn, err, is_equal_to, trace,
 	utils::{self, ReadyExt, stream::TryIgnore, string::Unquoted},
 };
 #[cfg(feature = "ldap")]
@@ -70,6 +72,7 @@ struct Data {
 	userfilterid_filter: Arc<Map>,
 	userid_avatarurl: Arc<Map>,
 	userid_blurhash: Arc<Map>,
+	userid_dehydrateddevice: Arc<Map>,
 	userid_devicelistversion: Arc<Map>,
 	userid_displayname: Arc<Map>,
 	userid_lastonetimekeyupdate: Arc<Map>,
@@ -110,6 +113,7 @@ impl crate::Service for Service {
 				userfilterid_filter: args.db["userfilterid_filter"].clone(),
 				userid_avatarurl: args.db["userid_avatarurl"].clone(),
 				userid_blurhash: args.db["userid_blurhash"].clone(),
+				userid_dehydrateddevice: args.db["userid_dehydrateddevice"].clone(),
 				userid_devicelistversion: args.db["userid_devicelistversion"].clone(),
 				userid_displayname: args.db["userid_displayname"].clone(),
 				userid_lastonetimekeyupdate: args.db["userid_lastonetimekeyupdate"].clone(),
@@ -480,6 +484,11 @@ impl Service {
 
 	/// Removes a device from a user.
 	pub async fn remove_device(&self, user_id: &UserId, device_id: &DeviceId) {
+		// Remove dehydrated device if this is the dehydrated device
+		let _: Result<_> = self
+			.remove_dehydrated_device(user_id, Some(device_id))
+			.await;
+
 		let userdeviceid = (user_id, device_id);
 
 		// Remove tokens
@@ -1003,7 +1012,7 @@ impl Service {
 		device_id: &'a DeviceId,
 		since: Option<u64>,
 		to: Option<u64>,
-	) -> impl Stream<Item = Raw<AnyToDeviceEvent>> + Send + 'a {
+	) -> impl Stream<Item = (u64, Raw<AnyToDeviceEvent>)> + Send + 'a {
 		type Key<'a> = (&'a UserId, &'a DeviceId, u64);
 
 		let from = (user_id, device_id, since.map_or(0, |since| since.saturating_add(1)));
@@ -1017,7 +1026,7 @@ impl Service {
 					&& device_id == *device_id_
 					&& to.is_none_or(|to| *count <= to)
 			})
-			.map(at!(1))
+			.map(|((_, _, count), event)| (count, event))
 	}
 
 	pub async fn remove_to_device_events<Until>(