continuwuity/src/api/client/to_device.rs

use std::collections::BTreeMap;

use axum::extract::State;
use conduwuit::{Err, Result};
use conduwuit_service::sending::EduBuf;
use futures::StreamExt;
use ruma::{
	api::{
		client::to_device::send_event_to_device,
		federation::transactions::edu::{DirectDeviceContent, Edu},
	},
	assign,
	to_device::DeviceIdOrAllDevices,
};

use crate::Ruma;

/// # `PUT /_matrix/client/r0/sendToDevice/{eventType}/{txnId}`
///
/// Send a to-device event to a set of client devices.
pub(crate) async fn send_event_to_device_route(
	State(services): State<crate::State>,
	body: Ruma<send_event_to_device::v3::Request>,
) -> Result<send_event_to_device::v3::Response> {
	let sender_user = body.identity.expect_sender_user()?;
	let sender_device = body.identity.sender_device();

	// Check if this is a new transaction id
	if services
		.transactions
		.get_client_txn(sender_user, sender_device, &body.txn_id)
		.await
		.is_ok()
	{
		return Ok(send_event_to_device::v3::Response::new());
	}

	for (target_user_id, map) in &body.messages {
		for (target_device, event) in map {
			if !services.globals.user_is_local(target_user_id) {
				let mut map = BTreeMap::new();
				map.insert(target_device.clone(), event.clone());
				let mut messages = BTreeMap::new();
				messages.insert(target_user_id.clone(), map);
				let count = services.globals.next_count()?;

				let mut buf = EduBuf::new();
				serde_json::to_writer(
					&mut buf,
					&Edu::DirectToDevice(assign!(
						DirectDeviceContent::new(
							sender_user.to_owned(),
							body.event_type.clone(),
							count.to_string().into(),
						),
						{ messages }
					)),
				)
				.expect("DirectToDevice EDU can be serialized");

				services
					.sending
					.send_edu_server(target_user_id.server_name(), buf)?;

				continue;
			}

			let event_type = &body.event_type.to_string();

			let Ok(event) = event.deserialize_as() else {
				return Err!(Request(InvalidParam("Failed to deserialize event body.")));
			};

			match target_device {
				| DeviceIdOrAllDevices::DeviceId(target_device_id) => {
					services
						.users
						.add_to_device_event(
							sender_user,
							target_user_id,
							target_device_id,
							event_type,
							event,
						)
						.await;
				},
				| DeviceIdOrAllDevices::AllDevices => {
					let (event_type, event) = (&event_type, &event);
					services
						.users
						.all_device_ids(target_user_id)
						.for_each(async |target_device_id| {
							services
								.users
								.add_to_device_event(
									sender_user,
									target_user_id,
									&target_device_id,
									event_type,
									event.clone(),
								)
								.await;
						})
						.await;
				},
			}
		}
	}

	// Save transaction id with empty data
	services
		.transactions
		.add_client_txnid(sender_user, sender_device, &body.txn_id, &[]);

	Ok(send_event_to_device::v3::Response::new())
}