iroh-relay 1.0.2

use std::{
    num::NonZeroUsize,
    sync::{Arc, Mutex},
};

use iroh_base::PublicKey;

type SignatureError = <PublicKey as TryFrom<&'static [u8]>>::Error;
type PublicKeyBytes = [u8; PublicKey::LENGTH];

/// A cache for public keys.
///
/// This is used solely to make parsing public keys from byte slices more
/// efficient for the very common case where a large number of identical keys
/// are being parsed, like in the relay server.
///
/// The cache stores only successful parse results.
#[derive(Debug, Clone)]
pub struct KeyCache(Inner);

#[derive(Debug, Clone)]
enum Inner {
    /// The key cache is disabled.
    Disabled,
    /// The key cache is enabled with a fixed capacity. It is shared between
    /// multiple threads.
    Shared(Arc<Mutex<lru::LruCache<PublicKey, ()>>>),
}

impl KeyCache {
    /// Key cache to be used in tests.
    #[cfg(all(test, feature = "server"))]
    pub fn test() -> Self {
        Self(Inner::Disabled)
    }

    /// Create a new key cache with the given capacity.
    ///
    /// If the capacity is zero, the cache is disabled and has zero overhead.
    pub fn new(capacity: usize) -> Self {
        match NonZeroUsize::new(capacity) {
            None => Self(Inner::Disabled),
            Some(capacity) => {
                let cache = lru::LruCache::new(capacity);
                Self(Inner::Shared(Arc::new(Mutex::new(cache))))
            }
        }
    }

    /// Get a key from a slice of bytes.
    pub fn key_from_slice(&self, slice: &[u8]) -> Result<PublicKey, SignatureError> {
        let Inner::Shared(cache) = &self.0 else {
            return PublicKey::try_from(slice);
        };
        let Ok(bytes) = PublicKeyBytes::try_from(slice) else {
            // if the size is wrong, use PublicKey::try_from to fail with a
            // SignatureError.
            return Err(PublicKey::try_from(slice).expect_err("invalid length"));
        };
        let mut cache = cache.lock().expect("not poisoned");
        if let Some((key, _)) = cache.get_key_value(&bytes) {
            return Ok(*key);
        }
        let key = PublicKey::from_bytes(&bytes)?;
        cache.put(key, ());
        Ok(key)
    }
}