distant_net/common/
keychain.rs

1use std::collections::HashMap;
2use std::sync::Arc;
3
4use tokio::sync::RwLock;
5
6use crate::common::HeapSecretKey;
7
8/// Represents the result of a request to the database.
9#[derive(Copy, Clone, Debug, PartialEq, Eq)]
10pub enum KeychainResult<T> {
11    /// Id was not found in the database.
12    InvalidId,
13
14    /// Password match for an id failed.
15    InvalidPassword,
16
17    /// Successful match of id and password, removing from keychain and returning data `T`.
18    Ok(T),
19}
20
21impl<T> KeychainResult<T> {
22    pub fn is_invalid_id(&self) -> bool {
23        matches!(self, Self::InvalidId)
24    }
25
26    pub fn is_invalid_password(&self) -> bool {
27        matches!(self, Self::InvalidPassword)
28    }
29
30    pub fn is_invalid(&self) -> bool {
31        matches!(self, Self::InvalidId | Self::InvalidPassword)
32    }
33
34    pub fn is_ok(&self) -> bool {
35        matches!(self, Self::Ok(_))
36    }
37
38    pub fn into_ok(self) -> Option<T> {
39        match self {
40            Self::Ok(x) => Some(x),
41            _ => None,
42        }
43    }
44}
45
46impl<T> From<KeychainResult<T>> for Option<T> {
47    fn from(result: KeychainResult<T>) -> Self {
48        result.into_ok()
49    }
50}
51
52/// Manages keys with associated ids. Cloning will result in a copy pointing to the same underlying
53/// storage, which enables support of managing the keys across multiple threads.
54#[derive(Debug)]
55pub struct Keychain<T = ()> {
56    map: Arc<RwLock<HashMap<String, (HeapSecretKey, T)>>>,
57}
58
59impl<T> Clone for Keychain<T> {
60    fn clone(&self) -> Self {
61        Self {
62            map: Arc::clone(&self.map),
63        }
64    }
65}
66
67impl<T> Keychain<T> {
68    /// Creates a new keychain without any keys.
69    pub fn new() -> Self {
70        Self {
71            map: Arc::new(RwLock::new(HashMap::new())),
72        }
73    }
74
75    /// Stores a new `key` and `data` by a given `id`, returning the old data associated with the
76    /// id if there was one already registered.
77    pub async fn insert(&self, id: impl Into<String>, key: HeapSecretKey, data: T) -> Option<T> {
78        self.map
79            .write()
80            .await
81            .insert(id.into(), (key, data))
82            .map(|(_, data)| data)
83    }
84
85    /// Checks if there is an `id` stored within the keychain.
86    pub async fn has_id(&self, id: impl AsRef<str>) -> bool {
87        self.map.read().await.contains_key(id.as_ref())
88    }
89
90    /// Checks if there is a key with the given `id` that matches the provided `key`.
91    pub async fn has_key(&self, id: impl AsRef<str>, key: impl PartialEq<HeapSecretKey>) -> bool {
92        self.map
93            .read()
94            .await
95            .get(id.as_ref())
96            .map(|(k, _)| key.eq(k))
97            .unwrap_or(false)
98    }
99
100    /// Removes a key and its data by a given `id`, returning the data if the `id` exists.
101    pub async fn remove(&self, id: impl AsRef<str>) -> Option<T> {
102        self.map
103            .write()
104            .await
105            .remove(id.as_ref())
106            .map(|(_, data)| data)
107    }
108
109    /// Checks if there is a key with the given `id` that matches the provided `key`, returning the
110    /// data if the `id` exists and the `key` matches.
111    pub async fn remove_if_has_key(
112        &self,
113        id: impl AsRef<str>,
114        key: impl PartialEq<HeapSecretKey>,
115    ) -> KeychainResult<T> {
116        let id = id.as_ref();
117        let mut lock = self.map.write().await;
118
119        match lock.get(id) {
120            Some((k, _)) if key.eq(k) => KeychainResult::Ok(lock.remove(id).unwrap().1),
121            Some(_) => KeychainResult::InvalidPassword,
122            None => KeychainResult::InvalidId,
123        }
124    }
125}
126
127impl Keychain<()> {
128    /// Stores a new `key by a given `id`.
129    pub async fn put(&self, id: impl Into<String>, key: HeapSecretKey) {
130        self.insert(id, key, ()).await;
131    }
132}
133
134impl Default for Keychain {
135    fn default() -> Self {
136        Self::new()
137    }
138}
139
140impl<T> From<HashMap<String, (HeapSecretKey, T)>> for Keychain<T> {
141    /// Creates a new keychain populated with the provided `map`.
142    fn from(map: HashMap<String, (HeapSecretKey, T)>) -> Self {
143        Self {
144            map: Arc::new(RwLock::new(map)),
145        }
146    }
147}
148
149impl From<HashMap<String, HeapSecretKey>> for Keychain<()> {
150    /// Creates a new keychain populated with the provided `map`.
151    fn from(map: HashMap<String, HeapSecretKey>) -> Self {
152        Self::from(
153            map.into_iter()
154                .map(|(id, key)| (id, (key, ())))
155                .collect::<HashMap<String, (HeapSecretKey, ())>>(),
156        )
157    }
158}