Implement Dehydrated Devices MSC3814 (closes #200)

Signed-off-by: Jason Volk <jason@zemos.net>

Cargo.toml

@@ -339,6 +339,7 @@ features = [
     "unstable-msc2870",
     "unstable-msc3026",
     "unstable-msc3061",
+    "unstable-msc3814",
     "unstable-msc3245",
     "unstable-msc3381", # polls
     "unstable-msc3489", # beacon / live location

src/api/client/dehydrated_device.rs

@@ -0,0 +1,132 @@
+use axum::extract::State;
+use axum_client_ip::InsecureClientIp;
+use futures::StreamExt;
+use ruma::api::client::dehydrated_device::{
+	delete_dehydrated_device::unstable as delete_dehydrated_device,
+	get_dehydrated_device::unstable as get_dehydrated_device, get_events::unstable as get_events,
+	put_dehydrated_device::unstable as put_dehydrated_device,
+};
+use tuwunel_core::{Err, Result, at, utils::result::IsErrOr};
+
+use crate::Ruma;
+
+const MAX_BATCH_EVENTS: usize = 50;
+
+/// # `PUT /_matrix/client/../dehydrated_device`
+///
+/// Creates or overwrites the user's dehydrated device.
+#[tracing::instrument(skip_all, fields(%client))]
+pub(crate) async fn put_dehydrated_device_route(
+	State(services): State<crate::State>,
+	InsecureClientIp(client): InsecureClientIp,
+	body: Ruma<put_dehydrated_device::Request>,
+) -> Result<put_dehydrated_device::Response> {
+	let sender_user = body
+		.sender_user
+		.as_deref()
+		.expect("AccessToken authentication required");
+
+	let device_id = body.body.device_id.clone();
+
+	services
+		.users
+		.set_dehydrated_device(sender_user, body.body)
+		.await?;
+
+	Ok(put_dehydrated_device::Response { device_id })
+}
+
+/// # `DELETE /_matrix/client/../dehydrated_device`
+///
+/// Deletes the user's dehydrated device without replacement.
+#[tracing::instrument(skip_all, fields(%client))]
+pub(crate) async fn delete_dehydrated_device_route(
+	State(services): State<crate::State>,
+	InsecureClientIp(client): InsecureClientIp,
+	body: Ruma<delete_dehydrated_device::Request>,
+) -> Result<delete_dehydrated_device::Response> {
+	let sender_user = body.sender_user();
+
+	let device_id = services
+		.users
+		.get_dehydrated_device_id(sender_user)
+		.await?;
+
+	services
+		.users
+		.remove_device(sender_user, &device_id)
+		.await;
+
+	Ok(delete_dehydrated_device::Response { device_id })
+}
+
+/// # `GET /_matrix/client/../dehydrated_device`
+///
+/// Gets the user's dehydrated device
+#[tracing::instrument(skip_all, fields(%client))]
+pub(crate) async fn get_dehydrated_device_route(
+	State(services): State<crate::State>,
+	InsecureClientIp(client): InsecureClientIp,
+	body: Ruma<get_dehydrated_device::Request>,
+) -> Result<get_dehydrated_device::Response> {
+	let sender_user = body.sender_user();
+
+	let device = services
+		.users
+		.get_dehydrated_device(sender_user)
+		.await?;
+
+	Ok(get_dehydrated_device::Response {
+		device_id: device.device_id,
+		device_data: device.device_data,
+	})
+}
+
+/// # `GET /_matrix/client/../dehydrated_device/{device_id}/events`
+///
+/// Paginates the events of the dehydrated device.
+#[tracing::instrument(skip_all, fields(%client))]
+pub(crate) async fn get_dehydrated_events_route(
+	State(services): State<crate::State>,
+	InsecureClientIp(client): InsecureClientIp,
+	body: Ruma<get_events::Request>,
+) -> Result<get_events::Response> {
+	let sender_user = body.sender_user();
+
+	let device_id = &body.body.device_id;
+	let existing_id = services
+		.users
+		.get_dehydrated_device_id(sender_user)
+		.await;
+
+	if existing_id
+		.as_ref()
+		.is_err_or(|existing_id| existing_id != device_id)
+	{
+		return Err!(Request(Forbidden("Not the dehydrated device_id.")));
+	}
+
+	let since: Option<u64> = body
+		.body
+		.next_batch
+		.as_deref()
+		.map(str::parse)
+		.transpose()?;
+
+	let mut next_batch: Option<u64> = None;
+	let events = services
+		.users
+		.get_to_device_events(sender_user, device_id, since, None)
+		.take(MAX_BATCH_EVENTS)
+		.inspect(|&(count, _)| {
+			next_batch.replace(count);
+		})
+		.map(at!(1))
+		.collect()
+		.await;
+
+	Ok(get_events::Response {
+		events,
+		next_batch: next_batch.as_ref().map(ToString::to_string),
+	})
+}

src/api/client/mod.rs

@@ -5,6 +5,7 @@ pub(super) mod appservice;
 pub(super) mod backup;
 pub(super) mod capabilities;
 pub(super) mod context;
+pub(super) mod dehydrated_device;
 pub(super) mod device;
 pub(super) mod directory;
 pub(super) mod filter;
@@ -49,6 +50,7 @@ pub(super) use appservice::*;
 pub(super) use backup::*;
 pub(super) use capabilities::*;
 pub(super) use context::*;
+pub(super) use dehydrated_device::*;
 pub(super) use device::*;
 pub(super) use directory::*;
 pub(super) use filter::*;

src/api/client/unversioned.rs

@@ -48,6 +48,7 @@ pub(crate) async fn get_supported_versions_route(
 			("org.matrix.msc2946".to_owned(), true), /* spaces/hierarchy summaries (https://github.com/matrix-org/matrix-spec-proposals/pull/2946) */
 			("org.matrix.msc3026.busy_presence".to_owned(), true), /* busy presence status (https://github.com/matrix-org/matrix-spec-proposals/pull/3026) */
 			("org.matrix.msc3575".to_owned(), true), /* sliding sync (https://github.com/matrix-org/matrix-spec-proposals/pull/3575/files#r1588877046) */
+			("org.matrix.msc3814".to_owned(), true), /* dehydrated devices */
 			("org.matrix.msc3827".to_owned(), true), /* filtering of /publicRooms by room type (https://github.com/matrix-org/matrix-spec-proposals/pull/3827) */
 			("org.matrix.msc3916.stable".to_owned(), true), /* authenticated media (https://github.com/matrix-org/matrix-spec-proposals/pull/3916) */
 			("org.matrix.msc3952_intentional_mentions".to_owned(), true), /* intentional mentions (https://github.com/matrix-org/matrix-spec-proposals/pull/3952) */

src/api/router.rs

@@ -163,6 +163,10 @@ pub fn build(router: Router<State>, server: &Server) -> Router<State> {
 		.ruma_route(&client::update_device_route)
 		.ruma_route(&client::delete_device_route)
 		.ruma_route(&client::delete_devices_route)
+		.ruma_route(&client::put_dehydrated_device_route)
+		.ruma_route(&client::delete_dehydrated_device_route)
+		.ruma_route(&client::get_dehydrated_device_route)
+		.ruma_route(&client::get_dehydrated_events_route)
 		.ruma_route(&client::get_tags_route)
 		.ruma_route(&client::update_tag_route)
 		.ruma_route(&client::delete_tag_route)

src/database/maps.rs

@@ -374,6 +374,10 @@ pub(super) static MAPS: &[Descriptor] = &[
 		name: "userid_blurhash",
 		..descriptor::RANDOM_SMALL
 	},
+	Descriptor {
+		name: "userid_dehydrateddevice",
+		..descriptor::RANDOM_SMALL
+	},
 	Descriptor {
 		name: "userid_devicelistversion",
 		..descriptor::RANDOM_SMALL

src/service/users/dehydrated_device.rs

@@ -0,0 +1,155 @@
+use ruma::{
+	DeviceId, OwnedDeviceId, UserId,
+	api::client::dehydrated_device::{
+		DehydratedDeviceData, put_dehydrated_device::unstable::Request,
+	},
+	serde::Raw,
+};
+use serde::{Deserialize, Serialize};
+use tuwunel_core::{Err, Result, implement, trace};
+use tuwunel_database::{Deserialized, Json};
+
+#[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
+			.device_exists(user_id, &request.device_id)
+			.await
+	{
+		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,
+		(None, None),
+		None,
+		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);
+	self.add_one_time_keys(
+		user_id,
+		&request.device_id,
+		request
+			.one_time_keys
+			.iter()
+			.map(|(id, key)| (id.as_ref(), key)),
+	)
+	.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::<String>()
+		.and_then(|raw| serde_json::from_str(&raw).map_err(Into::into))
+}

src/service/users/device.rs

@@ -10,7 +10,7 @@ use ruma::{
 };
 use serde_json::json;
 use tuwunel_core::{
-	Err, Result, at, implement,
+	Err, Result, implement,
 	utils::{
 		self, ReadyExt,
 		stream::{IterStream, TryIgnore},
@@ -90,6 +90,11 @@ pub async fn remove_device(&self, user_id: &UserId, device_id: &DeviceId) {
 		})
 		.await;
 
+	// Removes the dehydrated device if the ID matches, otherwise no-op
+	self.remove_dehydrated_device(user_id, Some(device_id))
+		.await
+		.ok();
+
 	// TODO: Remove onetimekeys
 
 	increment(&self.db.userid_devicelistversion, user_id.as_bytes());

src/service/users/mod.rs

@@ -1,3 +1,4 @@
+mod dehydrated_device;
 pub mod device;
 mod keys;
 mod ldap;
@@ -45,6 +46,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>,
@@ -74,6 +76,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(),