iroh_relay/
key_cache.rs

1use std::{
2    num::NonZeroUsize,
3    sync::{Arc, Mutex},
4};
5
6use iroh_base::PublicKey;
7
8type SignatureError = <PublicKey as TryFrom<&'static [u8]>>::Error;
9type PublicKeyBytes = [u8; PublicKey::LENGTH];
10
11/// A cache for public keys.
12///
13/// This is used solely to make parsing public keys from byte slices more
14/// efficient for the very common case where a large number of identical keys
15/// are being parsed, like in the relay server.
16///
17/// The cache stores only successful parse results.
18#[derive(Debug, Clone)]
19pub enum KeyCache {
20    /// The key cache is disabled.
21    Disabled,
22    /// The key cache is enabled with a fixed capacity. It is shared between
23    /// multiple threads.
24    Shared(Arc<Mutex<lru::LruCache<PublicKey, ()>>>),
25}
26
27impl KeyCache {
28    /// Key cache to be used in tests.
29    #[cfg(all(test, feature = "server"))]
30    pub fn test() -> Self {
31        Self::Disabled
32    }
33
34    /// Create a new key cache with the given capacity.
35    ///
36    /// If the capacity is zero, the cache is disabled and has zero overhead.
37    pub fn new(capacity: usize) -> Self {
38        let Some(capacity) = NonZeroUsize::new(capacity) else {
39            return Self::Disabled;
40        };
41        let cache = lru::LruCache::new(capacity);
42        Self::Shared(Arc::new(Mutex::new(cache)))
43    }
44
45    /// Get a key from a slice of bytes.
46    pub fn key_from_slice(&self, slice: &[u8]) -> Result<PublicKey, SignatureError> {
47        let Self::Shared(cache) = self else {
48            return PublicKey::try_from(slice);
49        };
50        let Ok(bytes) = PublicKeyBytes::try_from(slice) else {
51            // if the size is wrong, use PublicKey::try_from to fail with a
52            // SignatureError.
53            return Err(PublicKey::try_from(slice).expect_err("invalid length"));
54        };
55        let mut cache = cache.lock().expect("not poisoned");
56        if let Some((key, _)) = cache.get_key_value(&bytes) {
57            return Ok(*key);
58        }
59        let key = PublicKey::from_bytes(&bytes)?;
60        cache.put(key, ());
61        Ok(key)
62    }
63}
iroh_relay/key_cache.rs

iroh_relay/
key_cache.rs
