arminius_lib_db_keys/
fallback_keys.rs

1// SPDX-FileCopyrightText: 2025 BWI GmbH
2//
3// SPDX-License-Identifier: Apache-2.0
4
5//! Queries related to the `fallback_keys` table.
6
7mod tests;
8
9use std::collections::{BTreeMap, HashMap};
10
11use arminius_lib_db::{Database, DatabaseError, ResultExt as _};
12use arminius_lib_model::user_keys::{OneTimeKeyMap, OneTimeKeyRow};
13use ruma::{
14	DeviceId, KeyId, OneTimeKeyAlgorithm, OneTimeKeyName, OwnedDeviceId, OwnedOneTimeKeyId,
15	OwnedOneTimeKeyName, OwnedUserId, UserId, encryption::OneTimeKey,
16};
17
18/// Queries related to the `fallback_keys` table.
19// Queries are usually used as extension traits so suggest `import as _`.
20#[rust_analyzer::prefer_underscore_import]
21pub trait FallbackKeysQueries: Database {
22	/// Inserts fallback keys for a user into the database.
23	///
24	/// # Conflict resolution
25	///
26	/// If a fallback key for the same device and algorithm is already present
27	/// in the database, it will get overwritten by the given new one.
28	#[tracing::instrument(skip(self))]
29	async fn upsert_fallback_key(
30		&mut self,
31		key_algorithm: &OneTimeKeyAlgorithm,
32		key_name: &OneTimeKeyName,
33		user_id: &UserId,
34		device_id: &DeviceId,
35		key: &OneTimeKey,
36	) -> Result<(), DatabaseError> {
37		let key_object =
38			serde_json::to_value(key).invalid_data("Failed to serialize fallback key")?;
39		sqlx::query!(
40			"INSERT INTO fallback_keys (key_algorithm, key_name, device_nid, key_object)
41				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)
42				ON CONFLICT (device_nid, key_algorithm)
43				DO UPDATE SET
44					key_name = EXCLUDED.key_name,
45					key_object = EXCLUDED.key_object
46				WHERE
47					fallback_keys.key_name IS DISTINCT FROM EXCLUDED.key_name OR
48					fallback_keys.key_object IS DISTINCT FROM EXCLUDED.key_object
49				",
50			key_algorithm.as_str(),
51			key_name.as_str(),
52			user_id.as_str(),
53			device_id.as_str(),
54			key_object
55		)
56		.execute(self.connection().await?)
57		.await?;
58
59		Ok(())
60	}
61
62	/// Idempotently inserts a fallback key for import purposes, preserving the
63	/// `used` flag from the export.
64	///
65	/// Returns `true` if a new row was inserted, `false` if a fallback key
66	/// already exists with the same device ID and algorithm. This makes imports
67	/// safe to resume and replay an import.
68	#[tracing::instrument(skip(self))]
69	async fn import_fallback_key(
70		&mut self,
71		key_algorithm: &OneTimeKeyAlgorithm,
72		key_name: &OneTimeKeyName,
73		user_id: &UserId,
74		device_id: &DeviceId,
75		key: &OneTimeKey,
76		used: bool,
77	) -> Result<bool, DatabaseError> {
78		let key_object =
79			serde_json::to_value(key).invalid_data("Failed to serialize fallback key")?;
80		let result = sqlx::query!(
81			r#"
82			INSERT INTO fallback_keys (key_algorithm, key_name, device_nid, key_object, used)
83			SELECT
84				$1,
85				$2,
86				devices.nid,
87				$5,
88				$6
89			FROM
90				devices
91				JOIN users ON devices.user_nid = users.nid
92			WHERE
93				user_id = $3
94				AND device_id = $4
95			ON CONFLICT (device_nid,
96				key_algorithm)
97				DO NOTHING
98			"#,
99			key_algorithm.as_str(),
100			key_name.as_str(),
101			user_id.as_str(),
102			device_id.as_str(),
103			key_object,
104			used,
105		)
106		.execute(self.connection().await?)
107		.await?;
108		Ok(result.rows_affected() > 0)
109	}
110
111	/// Fetches all the requested fallback keys
112	#[tracing::instrument(skip(self))]
113	async fn fetch_fallback_keys(
114		&mut self,
115		fallback_keys: &BTreeMap<OwnedUserId, BTreeMap<OwnedDeviceId, OneTimeKeyAlgorithm>>,
116	) -> Result<OneTimeKeyMap, DatabaseError> {
117		let mut users_ids = Vec::new();
118		let mut devices_ids = Vec::new();
119		let mut keys_algorithms = Vec::new();
120
121		for (user_id, devices) in fallback_keys {
122			for (device_id, key_algorithm) in devices {
123				users_ids.push(user_id.to_string());
124				devices_ids.push(device_id.to_string());
125				keys_algorithms.push(key_algorithm.to_string());
126			}
127		}
128
129		sqlx::query!(
130			r#"
131			SELECT
132				fallback_keys.nid,
133				users.user_id,
134				devices.device_id,
135				fallback_keys.key_algorithm,
136				fallback_keys.key_name,
137				fallback_keys.key_object,
138				fallback_keys.used
139			FROM
140				fallback_keys
141				JOIN devices ON fallback_keys.device_nid = devices.nid
142				JOIN users ON devices.user_nid = users.nid
143			WHERE (user_id, device_id, key_algorithm) IN (
144				SELECT
145					*
146				FROM
147					UNNEST($1::text[], $2::text[], $3::text[]))
148			"#,
149			&users_ids,
150			&devices_ids,
151			&keys_algorithms
152		)
153		.fetch_all(self.connection().await?)
154		.await?
155		.into_iter()
156		.try_fold(
157			HashMap::new(),
158			|mut result: HashMap<
159				OwnedUserId,
160				HashMap<OwnedDeviceId, HashMap<OwnedOneTimeKeyId, OneTimeKeyRow>>,
161			>,
162			 row| {
163				let user_id = row.user_id.parse().invalid_data("Invalid user ID in database")?;
164				let key_id = KeyId::from_parts(
165					OneTimeKeyAlgorithm::from(row.key_algorithm),
166					OwnedOneTimeKeyName::from(row.key_name).as_ref(),
167				);
168				let key = serde_json::from_value(row.key_object)
169					.invalid_data("Invalid one-time key in database")?;
170
171				result
172					.entry(user_id)
173					.or_default()
174					.entry(row.device_id.into())
175					.or_default()
176					.entry(key_id)
177					.or_insert(OneTimeKeyRow { nid: row.nid, key, used: Some(row.used) });
178
179				Ok(result)
180			},
181		)
182	}
183
184	/// Sets all the given fallback keys' `used` field to true.
185	///
186	/// Returns the number of affected (i.e. modified) rows.
187	#[tracing::instrument(skip(self))]
188	async fn set_fallback_keys_used(&mut self, nids: &[i32]) -> Result<(), DatabaseError> {
189		if sqlx::query!("UPDATE fallback_keys SET used = true WHERE nid = ANY($1::int[])", nids)
190			.execute(self.connection().await?)
191			.await?
192			.rows_affected() as usize
193			!= nids.len()
194		{
195			return Err(DatabaseError::SqlxError(sqlx::Error::RowNotFound));
196		}
197
198		Ok(())
199	}
200
201	/// Fetches the list of unused fallback keys algorithms stored for the given
202	/// user's device.
203	#[tracing::instrument(skip(self))]
204	async fn fetch_unused_fallback_keys_algorithms(
205		&mut self,
206		user_id: &UserId,
207		device_id: &DeviceId,
208	) -> Result<Vec<OneTimeKeyAlgorithm>, DatabaseError> {
209		Ok(sqlx::query!(
210			"
211			SELECT DISTINCT
212				key_algorithm
213			FROM
214				fallback_keys
215				JOIN devices ON fallback_keys.device_nid = devices.nid
216				JOIN users ON devices.user_nid = users.nid
217			WHERE
218				users.user_id = $1
219				AND devices.device_id = $2
220				-- Only consider the unused keys
221				AND fallback_keys.used = FALSE
222			",
223			user_id.to_string(),
224			device_id.to_string()
225		)
226		.fetch_all(self.connection().await?)
227		.await?
228		.into_iter()
229		.map(|row| row.key_algorithm.into())
230		.collect())
231	}
232}
233
234impl<T> FallbackKeysQueries for T where T: Database {}
arminius_lib_db_keys/fallback_keys.rs

arminius_lib_db_keys/
fallback_keys.rs
