1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158

use std::collections::HashMap;
use std::sync::Arc;

use tokio::sync::RwLock;

use crate::common::HeapSecretKey;

/// Represents the result of a request to the database.
#[derive(Copy, Clone, Debug, PartialEq, Eq)]
pub enum KeychainResult<T> {
    /// Id was not found in the database.
    InvalidId,

    /// Password match for an id failed.
    InvalidPassword,

    /// Successful match of id and password, removing from keychain and returning data `T`.
    Ok(T),
}

impl<T> KeychainResult<T> {
    pub fn is_invalid_id(&self) -> bool {
        matches!(self, Self::InvalidId)
    }

    pub fn is_invalid_password(&self) -> bool {
        matches!(self, Self::InvalidPassword)
    }

    pub fn is_invalid(&self) -> bool {
        matches!(self, Self::InvalidId | Self::InvalidPassword)
    }

    pub fn is_ok(&self) -> bool {
        matches!(self, Self::Ok(_))
    }

    pub fn into_ok(self) -> Option<T> {
        match self {
            Self::Ok(x) => Some(x),
            _ => None,
        }
    }
}

impl<T> From<KeychainResult<T>> for Option<T> {
    fn from(result: KeychainResult<T>) -> Self {
        result.into_ok()
    }
}

/// Manages keys with associated ids. Cloning will result in a copy pointing to the same underlying
/// storage, which enables support of managing the keys across multiple threads.
#[derive(Debug)]
pub struct Keychain<T = ()> {
    map: Arc<RwLock<HashMap<String, (HeapSecretKey, T)>>>,
}

impl<T> Clone for Keychain<T> {
    fn clone(&self) -> Self {
        Self {
            map: Arc::clone(&self.map),
        }
    }
}

impl<T> Keychain<T> {
    /// Creates a new keychain without any keys.
    pub fn new() -> Self {
        Self {
            map: Arc::new(RwLock::new(HashMap::new())),
        }
    }

    /// Stores a new `key` and `data` by a given `id`, returning the old data associated with the
    /// id if there was one already registered.
    pub async fn insert(&self, id: impl Into<String>, key: HeapSecretKey, data: T) -> Option<T> {
        self.map
            .write()
            .await
            .insert(id.into(), (key, data))
            .map(|(_, data)| data)
    }

    /// Checks if there is an `id` stored within the keychain.
    pub async fn has_id(&self, id: impl AsRef<str>) -> bool {
        self.map.read().await.contains_key(id.as_ref())
    }

    /// Checks if there is a key with the given `id` that matches the provided `key`.
    pub async fn has_key(&self, id: impl AsRef<str>, key: impl PartialEq<HeapSecretKey>) -> bool {
        self.map
            .read()
            .await
            .get(id.as_ref())
            .map(|(k, _)| key.eq(k))
            .unwrap_or(false)
    }

    /// Removes a key and its data by a given `id`, returning the data if the `id` exists.
    pub async fn remove(&self, id: impl AsRef<str>) -> Option<T> {
        self.map
            .write()
            .await
            .remove(id.as_ref())
            .map(|(_, data)| data)
    }

    /// Checks if there is a key with the given `id` that matches the provided `key`, returning the
    /// data if the `id` exists and the `key` matches.
    pub async fn remove_if_has_key(
        &self,
        id: impl AsRef<str>,
        key: impl PartialEq<HeapSecretKey>,
    ) -> KeychainResult<T> {
        let id = id.as_ref();
        let mut lock = self.map.write().await;

        match lock.get(id) {
            Some((k, _)) if key.eq(k) => KeychainResult::Ok(lock.remove(id).unwrap().1),
            Some(_) => KeychainResult::InvalidPassword,
            None => KeychainResult::InvalidId,
        }
    }
}

impl Keychain<()> {
    /// Stores a new `key by a given `id`.
    pub async fn put(&self, id: impl Into<String>, key: HeapSecretKey) {
        self.insert(id, key, ()).await;
    }
}

impl Default for Keychain {
    fn default() -> Self {
        Self::new()
    }
}

impl<T> From<HashMap<String, (HeapSecretKey, T)>> for Keychain<T> {
    /// Creates a new keychain populated with the provided `map`.
    fn from(map: HashMap<String, (HeapSecretKey, T)>) -> Self {
        Self {
            map: Arc::new(RwLock::new(map)),
        }
    }
}

impl From<HashMap<String, HeapSecretKey>> for Keychain<()> {
    /// Creates a new keychain populated with the provided `map`.
    fn from(map: HashMap<String, HeapSecretKey>) -> Self {
        Self::from(
            map.into_iter()
                .map(|(id, key)| (id, (key, ())))
                .collect::<HashMap<String, (HeapSecretKey, ())>>(),
        )
    }
}