arminius_lib_user_keys/
identity_keys.rs1use 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#[derive(Debug, Default)]
26pub struct UsersKeys {
27 pub master_keys: BTreeMap<OwnedUserId, CrossSigningKeyRow>,
29 pub self_signing_keys: BTreeMap<OwnedUserId, CrossSigningKeyRow>,
31 pub user_signing_keys: BTreeMap<OwnedUserId, CrossSigningKeyRow>,
33 pub device_keys: BTreeMap<OwnedUserId, BTreeMap<OwnedDeviceId, DeviceKeysRow>>,
35}
36
37#[derive(Debug, Default, Deserialize, Serialize)]
43pub struct RawUserKeys {
44 #[serde(default)]
46 pub device_keys: BTreeMap<OwnedUserId, BTreeMap<OwnedDeviceId, Raw<DeviceKeys>>>,
47 #[serde(default)]
49 pub master_keys: BTreeMap<OwnedUserId, Raw<CrossSigningKey>>,
50 #[serde(default)]
52 pub self_signing_keys: BTreeMap<OwnedUserId, Raw<CrossSigningKey>>,
53 #[serde(default)]
58 pub user_signing_keys: BTreeMap<OwnedUserId, Raw<CrossSigningKey>>,
59}
60
61impl 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#[derive(Debug, Deserialize, Serialize)]
102pub enum FederatedIdentityKeysResponse {
103 Success(RawUserKeys),
105 Failure(FederatedKeyFailure),
107}
108
109pub 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 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 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 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}