tuwunel_service/appservice/
keys.rs

1use std::collections::BTreeMap;
2
3use futures::StreamExt;
4use ruma::{
5	OneTimeKeyAlgorithm, OwnedDeviceId, OwnedOneTimeKeyId, UserId,
6	api::appservice::{
7		Registration,
8		keys::{
9			claim_keys::unstable::Request as ClaimRequest,
10			query_keys::unstable::Request as QueryRequest,
11		},
12	},
13	encryption::{DeviceKeys, OneTimeKey},
14	serde::Raw,
15};
16use tuwunel_core::{
17	implement,
18	smallvec::SmallVec,
19	utils::{IterStream, stream::ReadyExt},
20};
21
22type ClaimedKeys = BTreeMap<OwnedDeviceId, BTreeMap<OwnedOneTimeKeyId, Raw<OneTimeKey>>>;
23type QueriedKeys = BTreeMap<OwnedDeviceId, Raw<DeviceKeys>>;
24type Registrations = SmallVec<[Registration; 1]>;
25
26#[implement(super::Service)]
27#[tracing::instrument(level = "debug", skip(self, one_time_keys))]
28pub async fn claim_keys(
29	&self,
30	user_id: &UserId,
31	one_time_keys: &BTreeMap<OwnedDeviceId, OneTimeKeyAlgorithm>,
32) -> ClaimedKeys {
33	self.registrations_for_user(user_id)
34		.await
35		.into_iter()
36		.stream()
37		.filter_map(async |registration| {
38			let devices = one_time_keys
39				.iter()
40				.map(|(device_id, algorithm)| (device_id.clone(), vec![algorithm.clone()]))
41				.collect();
42
43			let request = ClaimRequest {
44				one_time_keys: [(user_id.to_owned(), devices)].into(),
45			};
46
47			self.send_request(registration, request)
48				.await
49				.ok()
50				.flatten()
51				.and_then(|mut response| response.one_time_keys.remove(user_id))
52		})
53		.ready_fold(ClaimedKeys::new(), |claimed, response| {
54			response
55				.into_iter()
56				.fold(claimed, |mut claimed, (device_id, keys)| {
57					let Some(algorithm) = one_time_keys.get(&device_id) else {
58						return claimed;
59					};
60
61					let keys = keys
62						.into_iter()
63						.filter(|(key_id, _)| key_id.algorithm() == *algorithm)
64						.take(1)
65						.collect::<BTreeMap<_, _>>();
66
67					if !keys.is_empty() {
68						claimed.insert(device_id, keys);
69					}
70
71					claimed
72				})
73		})
74		.await
75}
76
77#[implement(super::Service)]
78#[tracing::instrument(level = "debug", skip(self, devices))]
79pub async fn query_keys(&self, user_id: &UserId, devices: &[OwnedDeviceId]) -> QueriedKeys {
80	self.registrations_for_user(user_id)
81		.await
82		.into_iter()
83		.stream()
84		.filter_map(async |registration| {
85			let request = QueryRequest {
86				device_keys: [(user_id.to_owned(), devices.to_vec())].into(),
87			};
88
89			self.send_request(registration, request)
90				.await
91				.ok()
92				.flatten()
93				.and_then(|mut response| response.device_keys.remove(user_id))
94		})
95		.ready_fold(QueriedKeys::new(), |mut queried, response| {
96			queried.extend(response);
97			queried
98		})
99		.await
100}
101
102#[implement(super::Service)]
103async fn registrations_for_user(&self, user_id: &UserId) -> Registrations {
104	self.read()
105		.await
106		.values()
107		.filter(|info| info.is_user_match(user_id))
108		.map(|info| info.registration.clone())
109		.collect()
110}
tuwunel_service/appservice/keys.rs

tuwunel_service/appservice/
keys.rs
