use super::{
fragment::MessageId,
sphinx::{SurbId, SurbPayloadEncryptionKeys, SURB_ID_SIZE},
};
use hashlink::{linked_hash_map, LinkedHashMap};
use log::warn;
use rand::{CryptoRng, Rng};
struct Value {
keys: SurbPayloadEncryptionKeys,
message_id: MessageId,
}
pub struct Entry<'a>(linked_hash_map::OccupiedEntry<'a, SurbId, Value>);
impl<'a> Entry<'a> {
pub fn keys(&self) -> &SurbPayloadEncryptionKeys {
&self.0.get().keys
}
pub fn message_id(&self) -> &MessageId {
&self.0.get().message_id
}
pub fn remove(self) {
self.0.remove();
}
}
pub struct SurbKeystore {
capacity: usize,
surbs: LinkedHashMap<SurbId, Value>,
}
impl SurbKeystore {
pub fn new(capacity: usize) -> Self {
debug_assert_ne!(capacity, 0);
Self { capacity, surbs: LinkedHashMap::with_capacity(capacity) }
}
pub fn insert(
&mut self,
rng: &mut (impl Rng + CryptoRng),
message_id: &MessageId,
log_target: &str,
) -> (SurbId, &mut SurbPayloadEncryptionKeys) {
debug_assert!(self.surbs.len() <= self.capacity);
if self.surbs.len() == self.capacity {
warn!(target: log_target, "Too many entries in SURB keystore; evicting oldest");
self.surbs.pop_front();
}
let mut id = [0; SURB_ID_SIZE];
rng.fill_bytes(&mut id);
match self.surbs.entry(id) {
linked_hash_map::Entry::Occupied(_) => panic!(
"Randomly generated SURB ID matches an existing SURB ID; something wrong with RNG?"
),
linked_hash_map::Entry::Vacant(entry) => {
let value = entry.insert(Value {
keys: SurbPayloadEncryptionKeys::new(),
message_id: *message_id,
});
(id, &mut value.keys)
},
}
}
pub fn entry(&mut self, id: &SurbId) -> Option<Entry> {
match self.surbs.entry(*id) {
linked_hash_map::Entry::Occupied(entry) => Some(Entry(entry)),
linked_hash_map::Entry::Vacant(_) => None,
}
}
}