arminius_lib_db_keys/

room_keys.rs

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

//! Queries related to room keys backups.
//!
//! There are two tables for these backups:
//! - `room_keys_versions`: Contains the backup declarations.
//! - `room_keys`: Contains the actual keys, linked to the `room_keys_versions` entry they belong
//!   to.

use std::collections::BTreeMap;

use arminius_lib_db::{Database, DatabaseError, ResultExt as _, TransactedDb};
use ruma::{OwnedRoomId, RoomId, UserId, api::client::backup::KeyBackupData};

/// The information about a stored backup.
#[derive(Debug)]
pub struct DbRoomKeysBackup {
	/// Algorithm-dependent data about the backup.
	pub auth_data: serde_json::Value,
	/// Number of keys stored in the backup.
	pub keys_count: u64,
	/// Integer representing the backup content's version.
	pub etag: i64,
	/// Version of the backup.
	pub version: i64,
	/// The actual keys stored in the backup.
	pub rooms: BTreeMap<OwnedRoomId, BTreeMap<String, KeyBackupData>>,
}

/// Queries related to room keys tables.
// Queries are usually used as extension traits so suggest `import as _`.
#[rust_analyzer::prefer_underscore_import]
pub trait RoomKeysQueries: Database {
	/// Creates a new room keys backup in the database and return its version as
	/// a string.
	///
	/// Note that the version will be set to `1` by default.
	#[tracing::instrument(skip(self))]
	async fn insert_backup(
		&mut self,
		user_id: &UserId,
		auth_data: &serde_json::Value,
	) -> Result<String, DatabaseError> {
		Ok(sqlx::query!(
			"
			WITH new_backup_version AS (
				UPDATE users SET keys_backup_seq = keys_backup_seq + 1
				WHERE user_id = $1
				RETURNING nid, keys_backup_seq
			)
			INSERT INTO room_keys_versions (user_nid, version, auth_data)
			SELECT nid, keys_backup_seq, $2 FROM new_backup_version
			RETURNING version
			",
			user_id.to_string(),
			auth_data
		)
		.fetch_one(self.connection().await?)
		.await?
		.version
		.to_string())
	}

	/// Updates a specific user's backup version.
	///
	/// Errors if the specified backup does not exist.
	#[tracing::instrument(skip(self))]
	async fn update_user_backup(
		&mut self,
		user_id: &UserId,
		version: i64,
		new_auth_data: &serde_json::Value,
	) -> Result<(), DatabaseError> {
		if sqlx::query!(
			r#"
			UPDATE
				room_keys_versions
			SET
				auth_data = $3,
				etag = etag + 1
			FROM
				users
			WHERE
				room_keys_versions.user_nid = users.nid
				AND users.user_id = $1
				AND room_keys_versions.version = $2
			"#,
			user_id.to_string(),
			version,
			new_auth_data
		)
		.execute(self.connection().await?)
		.await?
		.rows_affected()
			!= 1
		{
			return Err(DatabaseError::SqlxError(sqlx::Error::RowNotFound));
		}

		Ok(())
	}

	/// Fetches the `keys_backup_seq` of an user. This value is the highest room
	/// keys backup version ever created for that user.
	#[tracing::instrument(skip(self))]
	async fn fetch_max_user_backup_version(
		&mut self,
		user_id: &UserId,
	) -> Result<i64, DatabaseError> {
		Ok(sqlx::query_scalar!(
			"SELECT keys_backup_seq FROM users WHERE user_id = $1",
			user_id.to_string()
		)
		.fetch_one(self.connection().await?)
		.await?)
	}

	/// Deletes a specific user's backup version.
	#[tracing::instrument(skip(self))]
	async fn delete_user_backup(
		&mut self,
		user_id: &UserId,
		version: i64,
	) -> Result<(), DatabaseError> {
		sqlx::query!(
			"
			DELETE FROM room_keys_versions
			USING users
			WHERE
				room_keys_versions.user_nid = users.nid
				AND users.user_id = $1
				AND room_keys_versions.version = $2
			",
			user_id.to_string(),
			version
		)
		.execute(self.connection().await?)
		.await?;

		Ok(())
	}

	/// Deletes `room_keys` entries belonging to the given user and backup
	/// version.
	///
	/// If `room_id` and `session_id` are `Some`, the query will use them as
	/// filters.
	///
	/// This returns the new `etag` for the backup, as well as the count of keys
	/// left in it.
	#[tracing::instrument(skip(self))]
	async fn delete_backup_keys(
		&mut self,
		user_id: &UserId,
		version: i64,
		room_id: Option<OwnedRoomId>,
		session_id: Option<String>,
	) -> Result<Option<(i64, u64)>, DatabaseError> {
		Ok(sqlx::query_file!(
			"src/delete_backup_keys.sql",
			user_id.to_string(),
			version,
			room_id.map(|room_id| room_id.to_string()),
			session_id
		)
		.fetch_optional(self.connection().await?)
		.await?
		.map(|row| (row.etag, row.keys_count as u64)))
	}
}

impl<T: Database> RoomKeysQueries for T {}

/// SQL queries related to room keys that must occur in a database transaction.
pub trait TransactedRoomKeysQueries: RoomKeysQueries {
	/// Fetches a specific user's backup.
	///
	/// If `version` is `Some()`, it will filter using this version, otherwise
	/// it will take the highest version one (i.e. the latest backup).
	///
	/// If `room_id` and/or `session_id` are `Some()`, fetched keys will be
	/// filtered by these values.
	#[tracing::instrument(skip(self))]
	async fn fetch_user_backup(
		&mut self,
		user_id: &UserId,
		version: Option<i64>,
		room_id: Option<OwnedRoomId>,
		session_id: Option<String>,
	) -> Result<Option<DbRoomKeysBackup>, DatabaseError> {
		let mut backup_and_nid = sqlx::query!(
			r#"
			SELECT
				room_keys_versions.nid,
				auth_data,
				etag,
				version
			FROM
				room_keys_versions
				JOIN users ON room_keys_versions.user_nid = users.nid
			WHERE
				users.user_id = $1
				AND ($2::bigint IS NULL
					OR version = $2)
			ORDER BY
				version DESC
			LIMIT 1
			"#,
			user_id.to_string(),
			version
		)
		.fetch_optional(self.connection().await?)
		.await?
		.map(|row| {
			(
				row.nid,
				DbRoomKeysBackup {
					auth_data: row.auth_data,
					keys_count: 0,
					etag: row.etag,
					version: row.version,
					rooms: BTreeMap::new(),
				},
			)
		});

		// If we found the backup, fetch all the keys that belong to it
		if let Some((backup_nid, backup)) = backup_and_nid.as_mut() {
			let rows = sqlx::query!(
				r#"
				SELECT
					rooms.room_id,
					session_id,
					backup_data
				FROM
					room_keys
					JOIN room_keys_versions ON room_keys.room_keys_version_nid = room_keys_versions.nid
					JOIN users ON room_keys_versions.user_nid = users.nid
					JOIN rooms ON room_keys.room_nid = rooms.nid
				WHERE
					users.user_id = $1
					AND room_keys_versions.nid = $2
					AND ($3::text IS NULL
						OR rooms.room_id = $3)
					AND ($4::text IS NULL
						OR room_keys.session_id = $4)
				"#,
				user_id.to_string(),
				*backup_nid,
				room_id.map(|room_id| room_id.to_string()),
				session_id
			)
			.fetch_all(self.connection().await?)
			.await?;

			backup.keys_count = rows.len() as u64;

			for row in rows {
				let room_id =
					RoomId::parse(row.room_id).invalid_data("Invalid room ID in database")?;

				let backup_data: KeyBackupData = serde_json::from_value(row.backup_data)
					.invalid_data("Invalid room keys backup data in database")?;

				let room_keys = backup.rooms.entry(room_id).or_default();
				room_keys.insert(row.session_id, backup_data);
			}
		}

		Ok(backup_and_nid.map(|(_nid, backup)| backup))
	}

	/// Upserts keys in an existing backup.
	#[tracing::instrument(skip(self))]
	async fn upsert_backup_data(
		&mut self,
		user_id: &UserId,
		version: i64,
		new_backup_data: BTreeMap<OwnedRoomId, BTreeMap<String, serde_json::Value>>,
	) -> Result<(i64, u64), DatabaseError> {
		for (room_id, room_backups) in new_backup_data {
			for (session_id, backup_data) in room_backups {
				sqlx::query!(
					r#"
					WITH the_backup AS (
						SELECT
							room_keys_versions.nid,
							etag
						FROM
							room_keys_versions
							JOIN users ON room_keys_versions.user_nid = users.nid
						WHERE
							users.user_id = $1
							AND room_keys_versions.version = $2)
					INSERT INTO room_keys (room_keys_version_nid, room_nid, session_id, backup_data)
					VALUES
						((
								SELECT
									nid
								FROM
									the_backup),
								(
									SELECT
										nid
									FROM
										rooms
									WHERE
										room_id = $3), $4, $5)
						ON CONFLICT (room_keys_version_nid,
							room_nid,
							session_id)
						DO UPDATE SET
							backup_data = EXCLUDED.backup_data
					"#,
					user_id.to_string(),
					version,
					room_id.to_string(),
					session_id,
					backup_data
				)
				.execute(self.connection().await?)
				.await?;
			}
		}

		// Increment and get the `etag` of the backup, and the count of keys
		// associated to it
		let etag_and_keys_count = sqlx::query!(
			r#"
			WITH updated_backup AS (
				UPDATE
					room_keys_versions
				SET
					etag = etag + 1
				FROM
					users
				WHERE
					room_keys_versions.user_nid = users.nid
					AND users.user_id = $1
					AND room_keys_versions.version = $2
				RETURNING
					room_keys_versions.nid,
					etag
			)
			SELECT
				etag,
				COUNT(room_keys) AS "keys_count!"
			FROM
				updated_backup
				LEFT JOIN room_keys ON updated_backup.nid = room_keys.room_keys_version_nid
			GROUP BY
				etag
			"#,
			user_id.to_string(),
			version
		)
		.fetch_one(self.connection().await?)
		.await?;

		Ok((etag_and_keys_count.etag, etag_and_keys_count.keys_count as u64))
	}
}

impl<'c> TransactedRoomKeysQueries for TransactedDb<'c> {}