lb_rs/service/
keychain.rs

1use std::{
2    collections::HashMap,
3    sync::{Arc, RwLock},
4};
5
6use crate::{
7    model::{
8        account::Account,
9        crypto::AESKey,
10        errors::{LbErrKind, LbResult},
11    },
12    Lb,
13};
14use libsecp256k1::PublicKey;
15use tokio::sync::OnceCell;
16use uuid::Uuid;
17
18pub type KeyCache = Arc<RwLock<HashMap<Uuid, AESKey>>>;
19
20#[derive(Default, Clone)]
21pub struct Keychain {
22    key_cache: KeyCache,
23    account: Arc<OnceCell<Account>>,
24    public_key: Arc<OnceCell<PublicKey>>,
25}
26
27impl From<Option<&Account>> for Keychain {
28    fn from(value: Option<&Account>) -> Self {
29        match value {
30            Some(account) => {
31                let account = account.clone();
32                let pk = account.public_key();
33                let key_cache = Default::default();
34
35                Self {
36                    account: Arc::new(OnceCell::from(account)),
37                    public_key: Arc::new(OnceCell::from(pk)),
38                    key_cache,
39                }
40            }
41            None => Self::default(),
42        }
43    }
44}
45
46impl Lb {
47    pub fn get_account(&self) -> LbResult<&Account> {
48        self.keychain.get_account()
49    }
50}
51
52impl Keychain {
53    pub fn get_account(&self) -> LbResult<&Account> {
54        self.account
55            .get()
56            .ok_or_else(|| LbErrKind::AccountNonexistent.into())
57    }
58
59    pub fn get_pk(&self) -> LbResult<PublicKey> {
60        self.public_key
61            .get()
62            .copied()
63            .ok_or_else(|| LbErrKind::AccountNonexistent.into())
64    }
65
66    #[doc(hidden)]
67    pub async fn cache_account(&self, account: Account) -> LbResult<()> {
68        let pk = account.public_key();
69        self.account
70            .set(account)
71            .map_err(|_| LbErrKind::AccountExists)?;
72        self.public_key
73            .set(pk)
74            .map_err(|_| LbErrKind::AccountExists)?;
75
76        Ok(())
77    }
78
79    pub fn contains_aes_key(&self, id: &Uuid) -> LbResult<bool> {
80        Ok(self.key_cache.read()?.contains_key(id))
81    }
82
83    pub fn insert_aes_key(&self, id: Uuid, key: AESKey) -> LbResult<()> {
84        self.key_cache.write()?.insert(id, key);
85        Ok(())
86    }
87
88    pub fn get_aes_key(&self, id: &Uuid) -> LbResult<Option<AESKey>> {
89        Ok(self.key_cache.read()?.get(id).copied())
90    }
91}