arminius_lib_db_keys/

one_time_keys.rs

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

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

mod tests;

use std::collections::{BTreeMap, HashMap};

use arminius_lib_db::{Database, DatabaseError, ResultExt as _};
use arminius_lib_model::user_keys::{OneTimeKeyMap, OneTimeKeyRow};
use ruma::{
	DeviceId, KeyId, OneTimeKeyAlgorithm, OneTimeKeyName, OwnedDeviceId, OwnedOneTimeKeyId,
	OwnedOneTimeKeyName, OwnedUserId, UInt, UserId, encryption::OneTimeKey,
};

/// Queries related to the `one_time_keys` table.
// Queries are usually used as extension traits so suggest `import as _`.
#[rust_analyzer::prefer_underscore_import]
pub trait OneTimeKeysQueries: Database {
	/// Tries inserting one-time keys for a user into the database.
	///
	/// NOTE: If the given key already exists, does nothing.
	#[tracing::instrument(skip(self))]
	async fn try_insert_one_time_key(
		&mut self,
		key_algorithm: &OneTimeKeyAlgorithm,
		key_name: &OneTimeKeyName,
		user_id: &UserId,
		device_id: &DeviceId,
		key: &OneTimeKey,
	) -> Result<(), DatabaseError> {
		let key_object =
			serde_json::to_value(key).invalid_data("Failed to serialize one-time key")?;
		sqlx::query!(
			"INSERT INTO one_time_keys (key_algorithm, key_name, device_nid, key_object)
				VALUES (
					$1,
					$2,
					(
						SELECT devices.nid
						FROM devices
						JOIN users ON devices.user_nid = users.nid
						WHERE user_id = $3 AND device_id = $4
					),
					$5)
				ON CONFLICT DO NOTHING",
			key_algorithm.as_str(),
			key_name.as_str(),
			user_id.as_str(),
			device_id.as_str(),
			key_object,
		)
		.execute(self.connection().await?)
		.await?;

		Ok(())
	}

	/// Idempotently inserts a one-time key for import purposes.
	///
	/// Returns `true` if a new row was inserted, `false` if it already existed
	/// (i.e. the `ON CONFLICT DO NOTHING` branch fired). This makes imports
	/// safe to resume and replay.
	#[tracing::instrument(skip(self))]
	async fn import_one_time_key(
		&mut self,
		key_algorithm: &OneTimeKeyAlgorithm,
		key_name: &OneTimeKeyName,
		user_id: &UserId,
		device_id: &DeviceId,
		key: &OneTimeKey,
	) -> Result<bool, DatabaseError> {
		let key_object =
			serde_json::to_value(key).invalid_data("Failed to serialize one-time key")?;
		let result = sqlx::query!(
			r#"
			INSERT INTO one_time_keys (key_algorithm, key_name, device_nid, key_object)
			SELECT
				$1,
				$2,
				devices.nid,
				$5
			FROM
				devices
				JOIN users ON devices.user_nid = users.nid
			WHERE
				user_id = $3
				AND device_id = $4
			ON CONFLICT
				DO NOTHING
			"#,
			key_algorithm.as_str(),
			key_name.as_str(),
			user_id.as_str(),
			device_id.as_str(),
			key_object,
		)
		.execute(self.connection().await?)
		.await?;
		Ok(result.rows_affected() > 0)
	}

	/// Given a [`UserId`], returns the counts of OTKs per algorithm they
	/// own.
	#[tracing::instrument(skip(self))]
	async fn fetch_one_time_keys_counts(
		&mut self,
		user_id: &UserId,
		device_id: &DeviceId,
	) -> Result<BTreeMap<OneTimeKeyAlgorithm, UInt>, DatabaseError> {
		Ok(sqlx::query!(
			"SELECT
				key_algorithm,
				COUNT(*) AS count
			 FROM
			 	one_time_keys
				JOIN devices ON one_time_keys.device_nid = devices.nid
				JOIN users ON devices.user_nid = users.nid
			 WHERE
			 	user_id = $1
			 	AND device_id = $2
			 GROUP BY
				key_algorithm
		",
			user_id.as_str(),
			device_id.as_str(),
		)
		.fetch_all(self.connection().await?)
		.await?
		.into_iter()
		.map(|row| {
			(
				OneTimeKeyAlgorithm::from(row.key_algorithm),
				// SAFETY: `UInt::try_from(0)` should never fail.
				UInt::try_from(row.count.unwrap_or(0)).unwrap_or(UInt::MAX),
			)
		})
		.collect())
	}

	/// Fetch **one** one-time key for each requested user's device and
	/// algorithm.
	#[tracing::instrument(skip(self))]
	async fn fetch_one_time_keys(
		&mut self,
		one_time_keys: &BTreeMap<OwnedUserId, BTreeMap<OwnedDeviceId, OneTimeKeyAlgorithm>>,
	) -> Result<OneTimeKeyMap, DatabaseError> {
		let mut users_ids = Vec::new();
		let mut devices_ids = Vec::new();
		let mut otks_algorithms = Vec::new();

		for (user_id, devices) in one_time_keys {
			for (device_id, otk_algorithm) in devices {
				users_ids.push(user_id.to_string());
				devices_ids.push(device_id.to_string());
				otks_algorithms.push(otk_algorithm.to_string());
			}
		}

		sqlx::query!(
			"SELECT
				otks.nid,
				users.user_id,
				devices.device_id,
				otks.key_algorithm,
				otks.key_name,
				otks.key_object
			FROM
				users
				JOIN devices ON devices.user_nid = users.nid
				CROSS JOIN LATERAL (
					SELECT
						nid, key_algorithm, key_name, key_object
					FROM
						one_time_keys
					WHERE
						one_time_keys.device_nid = devices.nid
						AND key_algorithm = ANY($3::text[])
					ORDER BY
						upload_revision ASC
					LIMIT 1
				) AS otks
			WHERE
				(user_id, device_id) IN
				(SELECT * FROM UNNEST($1::text[], $2::text[]))
				",
			&users_ids,
			&devices_ids,
			&otks_algorithms
		)
		.fetch_all(self.connection().await?)
		.await?
		.into_iter()
		.try_fold(
			HashMap::new(),
			|mut result: HashMap<
				OwnedUserId,
				HashMap<OwnedDeviceId, HashMap<OwnedOneTimeKeyId, OneTimeKeyRow>>,
			>,
			 row| {
				let user_id = row.user_id.parse().invalid_data("Invalid user ID in database")?;
				let key_id = KeyId::from_parts(
					OneTimeKeyAlgorithm::from(row.key_algorithm),
					OwnedOneTimeKeyName::from(row.key_name).as_ref(),
				);
				let key = serde_json::from_value(row.key_object)
					.invalid_data("Invalid one-time key in database")?;

				result
					.entry(user_id)
					.or_default()
					.entry(row.device_id.into())
					.or_default()
					.entry(key_id)
					.or_insert(OneTimeKeyRow { nid: row.nid, key, used: None });

				Ok(result)
			},
		)
	}

	/// Deletes a one-time key from the database by its ID.
	///
	/// Returns the number of affected (e.g deleted) rows.
	#[tracing::instrument(skip(self))]
	async fn delete_one_time_keys(&mut self, nids: &[i32]) -> Result<(), DatabaseError> {
		if sqlx::query!("DELETE FROM one_time_keys WHERE nid = ANY($1::int[])", nids)
			.execute(self.connection().await?)
			.await?
			.rows_affected() as usize
			!= nids.len()
		{
			return Err(DatabaseError::SqlxError(sqlx::Error::RowNotFound));
		}

		Ok(())
	}
}

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