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

use toboggan_kv::{Toboggan, Tree};

use crate::{keys::private::AgentKey, AgentId, Error};

pub struct KeyManager<S: Toboggan> {
    _store: S,
    agent_keys: S::Tree,
    agent_id_config: S::Tree,
}

impl<S> KeyManager<S>
where
    S: Toboggan,
{
    pub fn new(store: S) -> Result<Self, Error> {
        Ok(Self {
            agent_keys: store.open_tree("agent_keys")?,
            agent_id_config: store.open_tree("agent_id_config")?,
            _store: store,
        })
    }
    pub fn list_agents(&self) -> Result<Vec<AgentId>, Error> {
        Ok(self
            .agent_keys
            .iter()
            .map(|kv| {
                let kv = kv.unwrap();
                let agentkey: AgentKey = bincode::deserialize(&kv.1).unwrap();
                AgentId {
                    pubkey: agentkey.pubkey(),
                }
            })
            .collect())
    }
    pub fn get_agent_key(&self, agent_id: &AgentId) -> Result<Option<AgentKey>, Error> {
        match self.agent_keys.get(&agent_id.pubkey)? {
            Some(ivec) => {
                let agentkey: AgentKey = bincode::deserialize(&ivec)?;
                Ok(Some(agentkey))
            }
            None => Ok(None),
        }
    }
    pub fn put_agent_key(&self, agentkey: AgentKey) -> Result<(), Error> {
        let key_ser: Vec<u8> = bincode::serialize(&agentkey).unwrap();

        let agent_id = agentkey.id();
        self.agent_keys.insert(&agent_id.pubkey[..], key_ser)?;
        self.agent_keys.flush()?;

        Ok(())
    }
    pub fn remove_all_agent_keys(&self) -> Result<(), Error> {
        self.agent_keys.clear()?;
        Ok(())
    }

    pub fn set_current_agent(&self, id: AgentId) -> Result<(), Error> {
        let id_ser: Vec<u8> = bincode::serialize(&id).unwrap();
        self.agent_id_config.insert("current", id_ser)?;
        Ok(())
    }
    pub fn current_agent_key(&self) -> Result<Option<AgentKey>, Error> {
        match self.agent_id_config.get("current")? {
            Some(value) => {
                let id: AgentId = bincode::deserialize(&value[..])?;

                match self.get_agent_key(&id)? {
                    Some(a) => Ok(Some(a)),
                    None => Err(Error::InvalidReferent),
                }
            }
            None => Ok(None),
        }
    }
}

impl<T> Default for KeyManager<T>
where
    T: Toboggan + Default,
{
    fn default() -> Self {
        KeyManager::new(Default::default())
            .expect("should not fail for store which implements default")
    }
}

#[cfg(test)]
mod test {
    use toboggan_kv::adapter::BTreeAdapter;

    use crate::{AgentKey, KeyManager};

    #[test]
    fn init() -> Result<(), std::io::Error> {
        let keymanager = KeyManager::new(BTreeAdapter::new())?;

        let agentkey = AgentKey::create(None);
        let id = agentkey.id();
        keymanager.put_agent_key(agentkey)?;
        keymanager.set_current_agent(id.clone())?;

        assert_eq!(
            keymanager
                .current_agent_key()
                .expect("is good")
                .expect("is some")
                .id(),
            id
        );

        Ok(())
    }
}