arminius_lib_user_keys/
one_time_keys.rs

1// SPDX-FileCopyrightText: 2025 BWI GmbH
2//
3// SPDX-License-Identifier: Apache-2.0
4
5//! Logic about users one-time keys.
6
7use std::collections::BTreeMap;
8
9use arminius_lib_component::error::{InternalServerError, ResultExt as _};
10use arminius_lib_db::{Database as _, DefaultDb};
11use arminius_lib_db_keys::{FallbackKeysQueries as _, OneTimeKeysQueries as _};
12use ruma::{
13	OneTimeKeyAlgorithm, OwnedDeviceId, OwnedOneTimeKeyId, OwnedUserId, encryption::OneTimeKey,
14	serde::Raw,
15};
16use serde::{Deserialize, Serialize};
17
18use crate::FederatedKeyFailure;
19
20/// Map of one-time keys fetched from the database.
21///
22/// This is the map from user IDs to maps from device IDs to maps from one-time
23/// keys IDs to raw one-time keys. We have to define manually here because both
24/// `claim_keys` from ruma (i.e. the
25/// [federation](ruma::api::federation::keys::claim_keys) and the
26/// [client](ruma::api::client::keys::claim_keys) ones)
27/// define different types for them.
28pub type RawOneTimeKeys =
29	BTreeMap<OwnedUserId, BTreeMap<OwnedDeviceId, BTreeMap<OwnedOneTimeKeyId, Raw<OneTimeKey>>>>;
30
31/// Federated one-time keys transiting over NATS.
32#[derive(Debug, Deserialize, Serialize)]
33pub struct FederatedOneTimeKeys {
34	/// The actual raw one-time keys received by foreign homeservers.
35	pub one_time_keys: RawOneTimeKeys,
36}
37
38/// Response received for a single foreign server one-time keys fetching.
39#[derive(Debug, Deserialize, Serialize)]
40pub enum FederatedOTKResponse {
41	/// The server was successfully queried and returned these keys.
42	Success(FederatedOneTimeKeys),
43	/// The server was unreachable.
44	Failure(FederatedKeyFailure),
45}
46
47/// Fetches every requested one-time keys from the database.
48///
49/// If found one-time keys will be used, and otherwise, fallback keys. These
50/// keys will then be converted to raw json values so that they can be used by
51/// their calling APIs.
52pub async fn claim_local_keys(
53	database: &mut DefaultDb,
54	requested_keys: BTreeMap<OwnedUserId, BTreeMap<OwnedDeviceId, OneTimeKeyAlgorithm>>,
55) -> Result<RawOneTimeKeys, InternalServerError> {
56	let mut transaction = database.begin_transaction().await?;
57
58	let mut onetime_keys = transaction.fetch_one_time_keys(&requested_keys).await?;
59
60	let mut result: RawOneTimeKeys = BTreeMap::new();
61	let mut otks_to_delete = Vec::new();
62	let mut missing_fallback_keys = BTreeMap::new();
63
64	for (req_user, req_devices) in requested_keys {
65		let Some(mut found_device_keys) = onetime_keys.remove(&req_user) else {
66			missing_fallback_keys.insert(req_user, req_devices);
67			continue;
68		};
69
70		for (req_device, req_algorithm) in req_devices {
71			let Some(mut found_device_otks) = found_device_keys.remove(&req_device) else {
72				missing_fallback_keys
73					.entry(req_user.clone())
74					.or_default()
75					.insert(req_device, req_algorithm);
76				continue;
77			};
78
79			if let Some((key_id, key)) = found_device_otks
80				.keys()
81				.find(|key_id| key_id.algorithm() == req_algorithm)
82				.cloned()
83				.and_then(|key_id| found_device_otks.remove_entry(&key_id))
84			{
85				let raw_key =
86					Raw::new(&key.key).unrecoverable_error("Converting one-time key to Raw")?;
87				otks_to_delete.push(key.nid);
88
89				result
90					.entry(req_user.clone())
91					.or_default()
92					.entry(req_device)
93					.or_default()
94					.insert(key_id, raw_key);
95			} else {
96				missing_fallback_keys
97					.entry(req_user.clone())
98					.or_default()
99					.entry(req_device)
100					.or_insert(req_algorithm);
101			}
102		}
103	}
104
105	// Remove the one-time keys that were used to construct the response
106	transaction.delete_one_time_keys(&otks_to_delete).await?;
107
108	// For every missing OTK, try and add a fallback key instead
109	let fallback_keys = transaction.fetch_fallback_keys(&missing_fallback_keys).await?;
110
111	let fallback_keys_nids: Vec<i32> = fallback_keys
112		.values()
113		.flat_map(|keys| {
114			keys.values()
115				.flat_map(|keys| keys.values().map(|key| key.nid).collect::<Vec<_>>())
116				.collect::<Vec<_>>()
117		})
118		.collect();
119
120	for (user_id, found_device_keys) in fallback_keys {
121		for (device_id, found_device_otks) in found_device_keys {
122			for (key_id, key) in found_device_otks {
123				let raw_key =
124					Raw::new(&key.key).unrecoverable_error("Converting one-time key to Raw")?;
125
126				result
127					.entry(user_id.clone())
128					.or_default()
129					.entry(device_id.clone())
130					.or_default()
131					.insert(key_id, raw_key);
132			}
133		}
134	}
135
136	// Mark the fallback keys as `used` to then notify the user.
137	transaction.set_fallback_keys_used(&fallback_keys_nids).await?;
138
139	transaction.commit().await?;
140
141	Ok(result)
142}
arminius_lib_user_keys/one_time_keys.rs

arminius_lib_user_keys/
one_time_keys.rs
