From 73bab94e2630df9470b44ec112a7d3037eae9b1b Mon Sep 17 00:00:00 2001 From: Emrul Islam Date: Sun, 12 Apr 2026 22:45:59 +0100 Subject: [PATCH] fix(store): bound MemPublicKeyStore cache to prevent memory exhaustion MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The in-memory public key cache was unbounded — every distinct 32-byte public key passed through `public_key()` added an entry that persisted for the lifetime of the store. The existing source already carried a `TODO: Make max number of keys stored configurable.` comment acknowledging this. This matters because the cache is fed from sync-peer input via the `insert_remote_entry → validate_entry → SignedEntry::verify → public_key` call chain. A peer that sends entries signed by many distinct authors grows the victim's map one entry per distinct author forever. Fix: cap the cache at `MAX_CACHED_KEYS = 10_000` using an `lru::LruCache`. When the cache is full, the least-recently-used entry is evicted. The `lru` crate is already transitively present in the dependency tree via `iroh-relay`, so this adds no new compile. --- Cargo.lock | 1 + Cargo.toml | 1 + src/store/pubkeys.rs | 112 ++++++++++++++++++++++++++++++++++++++----- 3 files changed, 103 insertions(+), 11 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index bd1e0d6c..d40aefc9 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2101,6 +2101,7 @@ dependencies = [ "iroh-metrics", "iroh-tickets", "irpc", + "lru", "n0-error", "n0-future", "nested_enum_utils", diff --git a/Cargo.toml b/Cargo.toml index f5657192..815fa8ce 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -37,6 +37,7 @@ iroh-blobs = { version = "0.100", default-features = false } iroh-gossip = { version = "0.98", features = ["net"], default-features = false } iroh-metrics = { version = "0.38", default-features = false } irpc = { version = "0.14", default-features = false } +lru = "0.16" n0-error = "0.1.0" n0-future = { version = "0.3.1", features = ["serde"] } num_enum = "0.7" diff --git a/src/store/pubkeys.rs b/src/store/pubkeys.rs index 005e8675..2dd473a5 100644 --- a/src/store/pubkeys.rs +++ b/src/store/pubkeys.rs @@ -1,9 +1,10 @@ use std::{ - collections::HashMap, - sync::{Arc, RwLock}, + num::NonZeroUsize, + sync::{Arc, Mutex}, }; use ed25519_dalek::{SignatureError, VerifyingKey}; +use lru::LruCache; use crate::{AuthorId, AuthorPublicKey, NamespaceId, NamespacePublicKey}; @@ -51,20 +52,109 @@ impl PublicKeyStore for () { } } -/// In-memory key storage -// TODO: Make max number of keys stored configurable. -#[derive(Debug, Clone, Default)] +/// Maximum number of public keys cached in [`MemPublicKeyStore`]. +/// +/// The cache is populated by [`PublicKeyStore::public_key`], which is called +/// during signature verification of entries received from sync peers (see +/// `validate_entry` in `sync.rs`). Without a bound, a peer that sends entries +/// signed by many distinct authors causes unbounded memory growth — one cache +/// entry per distinct 32-byte public key, forever. +/// +/// When the cache is full, the least-recently-used entry is evicted. +const MAX_CACHED_KEYS: usize = 10_000; + +/// In-memory key storage with a bounded LRU cache. +#[derive(Debug, Clone)] pub struct MemPublicKeyStore { - keys: Arc>>, + keys: Arc>>, +} + +impl Default for MemPublicKeyStore { + fn default() -> Self { + let cap = NonZeroUsize::new(MAX_CACHED_KEYS).expect("MAX_CACHED_KEYS is non-zero"); + Self { + keys: Arc::new(Mutex::new(LruCache::new(cap))), + } + } } impl PublicKeyStore for MemPublicKeyStore { fn public_key(&self, bytes: &[u8; 32]) -> Result { - if let Some(id) = self.keys.read().unwrap().get(bytes) { - return Ok(*id); + let mut guard = self.keys.lock().expect("MemPublicKeyStore mutex poisoned"); + if let Some(vk) = guard.get(bytes) { + return Ok(*vk); + } + let vk = VerifyingKey::from_bytes(bytes)?; + guard.put(*bytes, vk); + Ok(vk) + } +} + +#[cfg(test)] +mod tests { + use ed25519_dalek::SigningKey; + use rand::{rngs::StdRng, SeedableRng}; + + use super::*; + + /// Regression test for unbounded cache growth. + /// + /// Before the [`MAX_CACHED_KEYS`] bound was added, a peer that sent entries + /// signed by many distinct authors could grow [`MemPublicKeyStore`]'s + /// internal map without limit via the + /// `insert_remote_entry → validate_entry → entry.verify → public_key` call + /// chain. This test inserts more distinct keys than the bound and asserts + /// that the cache size stays bounded. + #[test] + fn cache_is_bounded_under_unique_key_flood() { + let store = MemPublicKeyStore::default(); + let mut rng = StdRng::seed_from_u64(0xDEAD_BEEF); + + let n = MAX_CACHED_KEYS + MAX_CACHED_KEYS / 2; + for _ in 0..n { + let sk = SigningKey::generate(&mut rng); + let bytes = sk.verifying_key().to_bytes(); + store.public_key(&bytes).expect("valid key"); + } + + let cache_size = store.keys.lock().unwrap().len(); + assert!( + cache_size <= MAX_CACHED_KEYS, + "cache should be bounded by MAX_CACHED_KEYS ({MAX_CACHED_KEYS}) but contains {cache_size} entries after {n} distinct keys", + ); + } + + /// Recently-used keys survive eviction when the cache overflows. + #[test] + fn recently_used_keys_survive_eviction() { + let store = MemPublicKeyStore::default(); + let mut rng = StdRng::seed_from_u64(0xBEEF_CAFE); + + // Fill the cache exactly to capacity. + let keys: Vec<[u8; 32]> = (0..MAX_CACHED_KEYS) + .map(|_| SigningKey::generate(&mut rng).verifying_key().to_bytes()) + .collect(); + for k in &keys { + store.public_key(k).expect("valid key"); + } + + // Touch the first inserted key to promote it to most-recently-used. + store.public_key(&keys[0]).expect("valid key"); + + // Insert 100 more distinct keys, forcing 100 LRU evictions. + for _ in 0..100 { + let bytes = SigningKey::generate(&mut rng).verifying_key().to_bytes(); + store.public_key(&bytes).expect("valid key"); } - let id = VerifyingKey::from_bytes(bytes)?; - self.keys.write().unwrap().insert(*bytes, id); - Ok(id) + + let guard = store.keys.lock().unwrap(); + assert!( + guard.contains(&keys[0]), + "touched key should survive LRU eviction" + ); + assert!( + !guard.contains(&keys[1]), + "oldest untouched key should have been evicted" + ); } }