use std::collections::HashMap;
use std::sync::Arc;
use std::time::Duration;
use adk_core::AdkError;
use async_trait::async_trait;
use tokio::sync::RwLock;
use tokio::time::Instant;
use zeroize::Zeroize;
use super::provider::SecretProvider;
pub const DEFAULT_MAX_ENTRIES: usize = 128;
struct CachedEntry {
value: String,
expires_at: Instant,
last_access: u64,
}
impl CachedEntry {
fn is_expired(&self, now: Instant) -> bool {
self.expires_at <= now
}
}
impl Drop for CachedEntry {
fn drop(&mut self) {
self.value.zeroize();
}
}
pub struct CachedSecretProvider<P: SecretProvider> {
inner: P,
cache: Arc<RwLock<HashMap<String, CachedEntry>>>,
ttl: Duration,
max_entries: usize,
access_counter: std::sync::atomic::AtomicU64,
}
impl<P: SecretProvider> CachedSecretProvider<P> {
pub fn new(inner: P, ttl: Duration) -> Self {
Self {
inner,
cache: Arc::new(RwLock::new(HashMap::new())),
ttl,
max_entries: DEFAULT_MAX_ENTRIES,
access_counter: std::sync::atomic::AtomicU64::new(0),
}
}
#[must_use]
pub fn with_max_entries(mut self, max_entries: usize) -> Self {
self.max_entries = max_entries;
self
}
pub async fn invalidate(&self, name: &str) {
self.cache.write().await.remove(name);
}
pub async fn invalidate_all(&self) {
self.cache.write().await.clear();
}
pub async fn purge_expired(&self) -> usize {
let now = Instant::now();
let mut cache = self.cache.write().await;
let before = cache.len();
cache.retain(|_, entry| !entry.is_expired(now));
before - cache.len()
}
pub async fn len(&self) -> usize {
self.cache.read().await.len()
}
pub async fn is_empty(&self) -> bool {
self.cache.read().await.is_empty()
}
fn next_access(&self) -> u64 {
self.access_counter.fetch_add(1, std::sync::atomic::Ordering::Relaxed)
}
async fn store(&self, name: &str, value: &str) {
if self.max_entries == 0 {
return;
}
let now = Instant::now();
let mut cache = self.cache.write().await;
cache.retain(|_, entry| !entry.is_expired(now));
while cache.len() >= self.max_entries {
let victim = cache
.iter()
.min_by_key(|(_, entry)| entry.last_access)
.map(|(name, _)| name.clone());
match victim {
Some(victim) => {
cache.remove(&victim);
}
None => break,
}
}
cache.insert(
name.to_string(),
CachedEntry {
value: value.to_string(),
expires_at: now + self.ttl,
last_access: self.next_access(),
},
);
}
}
impl<P: SecretProvider> std::fmt::Debug for CachedSecretProvider<P> {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("CachedSecretProvider")
.field("ttl", &self.ttl)
.field("max_entries", &self.max_entries)
.field("cache", &"<redacted>")
.finish()
}
}
#[async_trait]
impl<P: SecretProvider> SecretProvider for CachedSecretProvider<P> {
async fn get_secret(&self, name: &str) -> Result<String, AdkError> {
{
let mut cache = self.cache.write().await;
let now = Instant::now();
if let Some(entry) = cache.get_mut(name) {
if entry.is_expired(now) {
cache.remove(name);
} else {
entry.last_access = self.next_access();
return Ok(entry.value.clone());
}
}
}
let value = self.inner.get_secret(name).await?;
self.store(name, &value).await;
Ok(value)
}
}