Skip to main content

arminius_lib_user_keys/
identity_keys.rs

1// SPDX-FileCopyrightText: 2025 BWI GmbH
2//
3// SPDX-License-Identifier: Apache-2.0
4
5//! Logic about users identity keys (i.e. device keys and cross-signing 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::{CrossSigningKeyRow, DeviceKeysRow, UserKeysQueries as _};
12use arminius_lib_db_signatures::LocalSignatureQueries as _;
13use arminius_lib_model::user_keys::DbKeyType;
14use ruma::{
15	OwnedDeviceId, OwnedUserId, UserId,
16	encryption::{CrossSigningKey, DeviceKeys},
17	serde::{Base64, Raw, base64::Standard},
18};
19use serde::{Deserialize, Serialize};
20
21use crate::FederatedKeyFailure;
22
23/// Helper struct containing identity keys (i.e. device keys and cross-signing
24/// keys) of users in maps, in their database formats.
25#[derive(Debug, Default)]
26pub struct UsersKeys {
27	/// Map of users to master keys
28	pub master_keys: BTreeMap<OwnedUserId, CrossSigningKeyRow>,
29	/// Map of users to self-signing keys
30	pub self_signing_keys: BTreeMap<OwnedUserId, CrossSigningKeyRow>,
31	/// Map of users to user-signing keys
32	pub user_signing_keys: BTreeMap<OwnedUserId, CrossSigningKeyRow>,
33	/// Map of users to devices to device keys
34	pub device_keys: BTreeMap<OwnedUserId, BTreeMap<OwnedDeviceId, DeviceKeysRow>>,
35}
36
37/// Map of identity keys (e.g. device keys and cross-signing keys), in their
38/// API-ready formats.
39///
40/// This is used by both the federation-input and the client APIs when being
41/// requested keys for local users.
42#[derive(Debug, Default, Deserialize, Serialize)]
43pub struct RawUserKeys {
44	/// Users device keys.
45	#[serde(default)]
46	pub device_keys: BTreeMap<OwnedUserId, BTreeMap<OwnedDeviceId, Raw<DeviceKeys>>>,
47	/// Users Master keys.
48	#[serde(default)]
49	pub master_keys: BTreeMap<OwnedUserId, Raw<CrossSigningKey>>,
50	/// Users self-signing keys.
51	#[serde(default)]
52	pub self_signing_keys: BTreeMap<OwnedUserId, Raw<CrossSigningKey>>,
53	/// Users user-signing keys.
54	///
55	/// NOTE: These will only be fetched for the client API, but not for the
56	/// federation input (which is why we default it to an empty map).
57	#[serde(default)]
58	pub user_signing_keys: BTreeMap<OwnedUserId, Raw<CrossSigningKey>>,
59}
60
61/// Converts [UsersKeys] (database format) into [RawUserKeys] (API format) by
62/// simply changing all the keys to [Raw] ones.
63impl TryFrom<UsersKeys> for RawUserKeys {
64	type Error = InternalServerError;
65
66	fn try_from(users_keys: UsersKeys) -> Result<Self, Self::Error> {
67		let mut result = RawUserKeys::default();
68
69		for (user_id, msk_row) in users_keys.master_keys {
70			result
71				.master_keys
72				.insert(user_id, Raw::new(&msk_row.key).unrecoverable_error("Converting to raw")?);
73		}
74
75		for (user_id, ssk_row) in users_keys.self_signing_keys {
76			result
77				.self_signing_keys
78				.insert(user_id, Raw::new(&ssk_row.key).unrecoverable_error("Converting to raw")?);
79		}
80
81		for (user_id, usk_row) in users_keys.user_signing_keys {
82			result
83				.user_signing_keys
84				.insert(user_id, Raw::new(&usk_row.key).unrecoverable_error("Converting to raw")?);
85		}
86
87		for (user_id, user_device_keys) in users_keys.device_keys {
88			for (device_id, device_key_row) in user_device_keys {
89				result.device_keys.entry(user_id.clone()).or_default().insert(
90					device_id,
91					Raw::new(&device_key_row.key).unrecoverable_error("Converting to raw")?,
92				);
93			}
94		}
95
96		Ok(result)
97	}
98}
99
100/// Response received for a single foreign server identity keys fetching.
101#[derive(Debug, Deserialize, Serialize)]
102pub enum FederatedIdentityKeysResponse {
103	/// The server was successfully queried and returned these keys.
104	Success(RawUserKeys),
105	/// The server was unreachable.
106	Failure(FederatedKeyFailure),
107}
108
109/// Gets identity keys for local users, along their signatures.
110///
111/// This function will filter-in keys and signatures the given `requester` is
112/// allowed to see.
113///
114/// NOTE: Use `RawUserKeys::try_from()` to convert these database keys to raw
115/// keys (ready for API responses).
116pub async fn get_local_keys(
117	database: &mut DefaultDb,
118	requester: Option<&UserId>,
119	requested_keys: BTreeMap<OwnedUserId, Vec<OwnedDeviceId>>,
120) -> Result<UsersKeys, InternalServerError> {
121	let mut transaction = database.begin_transaction().await?;
122
123	let mut device_keys = transaction.fetch_device_keys_for_users(&requested_keys).await?;
124
125	let requested_users = requested_keys.into_keys().collect::<Vec<_>>();
126
127	let mut cross_signing_keys =
128		transaction.fetch_allowed_cross_signing_keys(requester, &requested_users).await?;
129
130	let all_keys_nids: Vec<i64> = device_keys
131		.values()
132		.flat_map(|keys| keys.values().map(|key| key.nid))
133		.chain(cross_signing_keys.values().flat_map(|keys| keys.values().map(|key| key.nid)))
134		.collect();
135
136	let signatures = transaction.fetch_allowed_local_signatures(requester, &all_keys_nids).await?;
137
138	transaction.commit().await?;
139
140	// Add the signatures to their corresponding keys
141	for (signed_key_nid, signatures) in signatures {
142		if let Some(device_keys) = device_keys
143			.values_mut()
144			.find_map(|keys| keys.values_mut().find(|keys| keys.nid == signed_key_nid))
145		{
146			for signature in signatures {
147				let entries = device_keys
148					.key
149					.signatures
150					.entry(signature.signing_user_id.clone())
151					.or_default();
152				entries.insert(
153					signature.signing_key_id.clone(),
154					Base64::<Standard>::new(signature.signature).encode(),
155				);
156			}
157		} else {
158			let cross_signing_key = cross_signing_keys.values_mut().find_map(|keys| {
159				keys.iter_mut()
160					.find_map(|(_key_type, key)| (key.nid == signed_key_nid).then_some(key))
161			});
162
163			if let Some(cross_signing_key) = cross_signing_key {
164				for signature in signatures {
165					let entries = cross_signing_key
166						.key
167						.signatures
168						.entry(signature.signing_user_id.clone())
169						.or_default();
170					entries.insert(
171						signature.signing_key_id.clone(),
172						Base64::<Standard>::new(signature.signature).encode(),
173					);
174				}
175			}
176		}
177	}
178
179	// Construct the final mappings
180	let mut result = UsersKeys { device_keys, ..Default::default() };
181
182	for (user_id, user_keys) in cross_signing_keys {
183		for (key_type, cross_signing_key) in user_keys {
184			match key_type {
185				DbKeyType::Master => {
186					result.master_keys.insert(user_id.clone(), cross_signing_key);
187				}
188				DbKeyType::SelfSigning => {
189					result.self_signing_keys.insert(user_id.clone(), cross_signing_key);
190				}
191				DbKeyType::UserSigning => {
192					result.user_signing_keys.insert(user_id.clone(), cross_signing_key);
193				}
194
195				// This should never happen.
196				DbKeyType::Device => {
197					return None.unrecoverable_error(
198						"Cross-signing key stored as device keys in database",
199					)?;
200				}
201			}
202		}
203	}
204
205	Ok(result)
206}