arminius_lib_db_todevice/

lib.rs

// SPDX-FileCopyrightText: 2025 BWI GmbH
//
// SPDX-License-Identifier: Apache-2.0

//! Queries related to the `to_device_messages` table.

use arminius_lib_db::{Database, DatabaseError, ResultExt as _};
use arminius_lib_model::types::StreamId;
use itertools::multiunzip;
use ruma::{
	DeviceId, OwnedUserId, UserId, events::ToDeviceEventType, to_device::DeviceIdOrAllDevices,
};
use sqlx::postgres::PgQueryResult;

/// A message that an user wants to send to one or more devices
#[derive(Debug)]
pub struct ToDeviceMessage {
	/// The user to send the message to
	pub user: OwnedUserId,
	/// Destination of the to-device message.
	pub device: DeviceIdOrAllDevices,
	/// The content of the event
	pub message: serde_json::Value,
}

/// A to_device event received as part of a sync response
#[derive(Debug)]
pub struct SyncToDeviceEvent {
	/// Sender of the event
	pub sender: OwnedUserId,
	/// Type of the to_device event
	pub event_type: ToDeviceEventType,
	/// Content of the event
	pub content: serde_json::Value,
	/// The stream_id of the event (See <https://bundesmessenger.gitlab.io/backend/neo/arminius/arminius/architecture/sync.html>)
	pub stream_id: StreamId,
}

/// Queries related to the room table.
// Queries are usually used as extension traits so suggest `import as _`.
#[rust_analyzer::prefer_underscore_import]
pub trait ToDeviceQueries: Database {
	/// Inserts a list of send-to-device messages into the database
	#[tracing::instrument(skip(self))]
	async fn insert_to_device_messages(
		&mut self,
		event_type: ToDeviceEventType,
		sender: &UserId,
		messages: Vec<ToDeviceMessage>,
	) -> Result<PgQueryResult, DatabaseError> {
		// sqlx does not support `Vec<Option<_>>` neatly, so we pass along a
		// `Vec<bool>` next to our actual values
		let (users, all_devices, devices, messages): (
			Vec<String>,
			Vec<bool>,
			Vec<String>,
			Vec<serde_json::Value>,
		) = multiunzip(messages.into_iter().map(|msg| {
			(
				msg.user.to_string(),
				matches!(msg.device, DeviceIdOrAllDevices::AllDevices),
				msg.device.to_string(),
				msg.message,
			)
		}));
		Ok(sqlx::query!(
			"
			WITH raw_rows AS (
			    SELECT
			        *
			    FROM
			        UNNEST($1::text[], $2::bool[], $3::text[], $4::json[]) AS t (user_id,
			            all_devices,
			            device_id,
			            message)
			),
			inserting_rows AS (
			    SELECT
			        d.nid,
			        message,
			        $5 AS event_type,
			        $6 AS sender
			    FROM
			        raw_rows r
			        LEFT JOIN users u ON u.user_id = r.user_id
			        LEFT JOIN devices d ON d.user_nid = u.nid
			            AND r.all_devices
			            OR (NOT r.all_devices
			                AND d.device_id = r.device_id))
			    INSERT INTO to_device_messages (device_nid, event_content, event_type, sender) (
			        SELECT
			            *
			        FROM
			            inserting_rows);
		",
			&users,
			&all_devices,
			&devices,
			&messages,
			&event_type.to_string(),
			sender.as_str(),
		)
		.execute(self.connection().await?)
		.await?)
	}

	/// Insert a to-device transaction, used in federation
	#[tracing::instrument(skip(self))]
	async fn insert_to_device_transaction(
		&mut self,
		origin: &ruma::ServerName,
		txnid: &ruma::TransactionId,
	) -> Result<(), DatabaseError> {
		sqlx::query!(
			r#"
			INSERT INTO federated_to_device_transactions (origin, txnid)
				VALUES ($1, $2)
			"#,
			origin.as_str(),
			txnid.as_str()
		)
		.execute(self.connection().await?)
		.await?;
		Ok(())
	}

	/// Returns true if a homeserver has already used the transaction ID for to-device messages
	#[tracing::instrument(skip(self))]
	async fn to_device_transaction_exists(
		&mut self,
		origin: &ruma::ServerName,
		txnid: &ruma::TransactionId,
	) -> Result<bool, DatabaseError> {
		let result = sqlx::query!(
			r#"
			SELECT
				*
			FROM
				federated_to_device_transactions
			WHERE
				origin = $1
				AND txnid = $2
			"#,
			origin.as_str(),
			txnid.as_str()
		)
		.fetch_optional(self.connection().await?)
		.await?;
		Ok(result.is_some())
	}

	/// Gets the todevice messages sent to a specific device since last sync.
	#[tracing::instrument(skip(self))]
	async fn fetch_sync_todevice_messages(
		&mut self,
		user_id: &UserId,
		device_id: &DeviceId,
		from: StreamId,
	) -> Result<Vec<SyncToDeviceEvent>, DatabaseError> {
		let rows = sqlx::query!(
			"
			SELECT
			    sender,
			    m.event_type,
			    m.event_content,
			    m.stream_id
			FROM
			    to_device_messages m
			    JOIN devices d ON d.nid = m.device_nid
			    JOIN users u ON u.nid = d.user_nid
			WHERE
			    user_id = $1
			    AND d.device_id = $2
			    AND m.stream_id > $3
			    AND m.stream_id <= max_safe_seq ('device_message_streamids')
			",
			user_id.as_str(),
			device_id.as_str(),
			from,
		)
		.fetch_all(self.connection().await?)
		.await?;

		rows.into_iter()
			.map(|row| {
				Ok(SyncToDeviceEvent {
					sender: OwnedUserId::try_from(row.sender)
						.invalid_data("Invalid user ID in database")?,
					event_type: ToDeviceEventType::from(row.event_type),
					content: row.event_content,
					stream_id: row.stream_id,
				})
			})
			.collect()
	}

	/// Imports a single to-device message from a homeserver export. Returns
	/// `true` if a row was inserted.
	///
	/// arminius's `stream_id` is a per-row UNIQUE sync cursor, but synapse
	/// shares one `device_inbox` stream_id across an entire fan-out (many
	/// recipient devices), so it cannot be preserved. Each recipient row gets a
	/// fresh per-row `stream_id` from the `device_message_streamids` sequence
	/// via the column default, exactly like a live send, so no sequence
	/// realignment is needed after the import.
	///
	/// If the recipient of the message does not exist, `Ok(false)` is returned.
	#[tracing::instrument(skip(self, content))]
	async fn import_to_device_message(
		&mut self,
		user_id: &UserId,
		device_id: &DeviceId,
		sender: &UserId,
		event_type: &ToDeviceEventType,
		content: serde_json::Value,
	) -> Result<bool, DatabaseError> {
		let result = sqlx::query!(
			r#"
			INSERT INTO to_device_messages (device_nid, event_content, event_type, sender)
			SELECT
				d.nid,
				$4::json,
				$5,
				$3
			FROM
				devices d
				JOIN users u ON d.user_nid = u.nid
			WHERE
				u.user_id = $1
				AND d.device_id = $2
				AND NOT EXISTS (
					SELECT
						1
					FROM
						to_device_messages existing
					WHERE
						existing.device_nid = d.nid
						AND existing.sender = $3
						AND existing.event_type = $5
						AND existing.event_content::text = $4::json::text)
			"#,
			user_id.as_str(),       // $1
			device_id.as_str(),     // $2
			sender.as_str(),        // $3
			content,                // $4
			event_type.to_string(), // $5
		)
		.execute(self.connection().await?)
		.await?;
		Ok(result.rows_affected() > 0)
	}

	/// When we are syncing past old send-to-device messages, these old messages
	/// should be deleted from the server, per the spec.
	///
	/// See <https://spec.matrix.org/v1.15/client-server-api/#server-behaviour-6>.
	#[tracing::instrument(skip(self))]
	async fn delete_synced_todevice_messages(
		&mut self,
		user_id: &UserId,
		device_id: &DeviceId,
		delete_to_inclusive: StreamId,
	) -> Result<PgQueryResult, DatabaseError> {
		Ok(sqlx::query!(
			r#"
			DELETE FROM to_device_messages m USING devices d, users u
			WHERE u.nid = d.user_nid
				AND d.nid = m.device_nid
				AND u.user_id = $1
				AND d.device_id = $2
				AND m.stream_id <= $3
				AND m.stream_id <= max_safe_seq ('device_message_streamids')
			"#,
			user_id.as_str(),
			device_id.as_str(),
			delete_to_inclusive,
		)
		.execute(self.connection().await?)
		.await?)
	}
}

impl<T> ToDeviceQueries for T where T: Database {}