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 struct KeyCache(Inner);
20
21#[derive(Debug, Clone)]
22enum Inner {
23    /// The key cache is disabled.
24    Disabled,
25    /// The key cache is enabled with a fixed capacity. It is shared between
26    /// multiple threads.
27    Shared(Arc<Mutex<lru::LruCache<PublicKey, ()>>>),
28}
29
30impl KeyCache {
31    /// Key cache to be used in tests.
32    #[cfg(all(test, feature = "server"))]
33    pub fn test() -> Self {
34        Self(Inner::Disabled)
35    }
36
37    /// Create a new key cache with the given capacity.
38    ///
39    /// If the capacity is zero, the cache is disabled and has zero overhead.
40    pub fn new(capacity: usize) -> Self {
41        match NonZeroUsize::new(capacity) {
42            None => Self(Inner::Disabled),
43            Some(capacity) => {
44                let cache = lru::LruCache::new(capacity);
45                Self(Inner::Shared(Arc::new(Mutex::new(cache))))
46            }
47        }
48    }
49
50    /// Get a key from a slice of bytes.
51    pub fn key_from_slice(&self, slice: &[u8]) -> Result<PublicKey, SignatureError> {
52        let Inner::Shared(cache) = &self.0 else {
53            return PublicKey::try_from(slice);
54        };
55        let Ok(bytes) = PublicKeyBytes::try_from(slice) else {
56            // if the size is wrong, use PublicKey::try_from to fail with a
57            // SignatureError.
58            return Err(PublicKey::try_from(slice).expect_err("invalid length"));
59        };
60        let mut cache = cache.lock().expect("not poisoned");
61        if let Some((key, _)) = cache.get_key_value(&bytes) {
62            return Ok(*key);
63        }
64        let key = PublicKey::from_bytes(&bytes)?;
65        cache.put(key, ());
66        Ok(key)
67    }
68}
iroh_relay/key_cache.rs

iroh_relay/
key_cache.rs
