kca/ocsp-responder/src/cache.rs

//! One-second, single-flight cache of per-CA status snapshots.

use std::sync::Arc;
use std::time::{Duration, Instant};

use tokio::sync::{Mutex, RwLock};

use crate::snapshot::Snapshot;
use crate::store::ObjectStore;

#[derive(Clone)]
pub struct CachedSnapshot {
    snapshot: Arc<Snapshot>,
    fetched_at: Instant,
}

#[derive(Default)]
pub struct SnapshotCache {
    cached: RwLock<Option<CachedSnapshot>>,
    refresh: Mutex<()>,
}

impl SnapshotCache {
    pub fn new() -> Self {
        Self::default()
    }

    /// Return a snapshot, refreshing from S3 at most once per `ttl`.
    ///
    /// On refresh failure, data up to `ttl + stale_grace` old is still used.
    pub async fn get<S: ObjectStore + ?Sized>(
        &self,
        store: &S,
        key: &str,
        expected_sha256: &str,
        ttl: Duration,
        stale_grace: Duration,
    ) -> Result<Arc<Snapshot>, CacheError> {
        if let Some(value) = self.cached.read().await.as_ref() {
            if value.fetched_at.elapsed() < ttl {
                return Ok(value.snapshot.clone());
            }
        }

        // Only one request refreshes a given CA. Everybody else waits here
        // and then observes the refreshed value below.
        let _guard = self.refresh.lock().await;
        let stale = {
            let cache = self.cached.read().await;
            if let Some(value) = cache.as_ref() {
                if value.fetched_at.elapsed() < ttl {
                    return Ok(value.snapshot.clone());
                }
                Some(value.clone())
            } else {
                None
            }
        };

        match store.get(key).await {
            Ok(bytes) => match Snapshot::parse(&bytes, expected_sha256) {
                Ok(snapshot) => {
                    let snapshot = Arc::new(snapshot);
                    *self.cached.write().await = Some(CachedSnapshot {
                        snapshot: snapshot.clone(),
                        fetched_at: Instant::now(),
                    });
                    Ok(snapshot)
                }
                Err(error) => self.fallback(stale, error.to_string(), stale_grace),
            },
            Err(error) => self.fallback(stale, error.to_string(), stale_grace),
        }
    }

    fn fallback(
        &self,
        stale: Option<CachedSnapshot>,
        error: String,
        stale_grace: Duration,
    ) -> Result<Arc<Snapshot>, CacheError> {
        if let Some(value) = stale {
            if value.fetched_at.elapsed() <= stale_grace {
                tracing::warn!(error = %error, "serving stale OCSP snapshot");
                return Ok(value.snapshot);
            }
        }
        Err(CacheError::Unavailable(error))
    }
}

#[derive(Debug, thiserror::Error)]
pub enum CacheError {
    #[error("OCSP snapshot unavailable: {0}")]
    Unavailable(String),
}

#[cfg(test)]
mod tests {
    use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};

    use async_trait::async_trait;
    use bytes::Bytes;

    use super::*;
    use crate::snapshot::sha256_hex;
    use crate::store::{ObjectStore, StoreError};

    struct FakeStore {
        body: String,
        calls: AtomicUsize,
        fail: AtomicBool,
    }

    impl FakeStore {
        fn new(body: String) -> Self {
            Self {
                body,
                calls: AtomicUsize::new(0),
                fail: AtomicBool::new(false),
            }
        }
    }

    #[async_trait]
    impl ObjectStore for FakeStore {
        async fn get(&self, _key: &str) -> Result<Bytes, StoreError> {
            self.calls.fetch_add(1, Ordering::SeqCst);
            if self.fail.load(Ordering::SeqCst) {
                return Err(StoreError::S3("offline".to_string()));
            }
            Ok(Bytes::from(self.body.clone()))
        }
    }

    fn body(sha: &str) -> String {
        format!(
            r#"{{"schema":"kca-ocsp-status/v1","issuer":{{"sha256":"{sha}"}},"entries":[{{"serial_hex":"01","status":"valid"}}]}}"#
        )
    }

    #[tokio::test]
    async fn single_flight_refreshes_once() {
        let sha = sha256_hex(b"ca");
        let store = FakeStore::new(body(&sha));
        let cache = SnapshotCache::new();
        let (a, b) = tokio::join!(
            cache.get(
                &store,
                "key",
                &sha,
                Duration::from_millis(100),
                Duration::from_secs(1)
            ),
            cache.get(
                &store,
                "key",
                &sha,
                Duration::from_millis(100),
                Duration::from_secs(1)
            )
        );
        assert!(a.is_ok());
        assert!(b.is_ok());
        assert_eq!(store.calls.load(Ordering::SeqCst), 1);
    }

    #[tokio::test]
    async fn refreshes_after_ttl_and_serves_stale_on_error() {
        let sha = sha256_hex(b"ca");
        let store = FakeStore::new(body(&sha));
        let cache = SnapshotCache::new();
        cache
            .get(
                &store,
                "key",
                &sha,
                Duration::from_millis(20),
                Duration::from_secs(1),
            )
            .await
            .unwrap();
        tokio::time::sleep(Duration::from_millis(30)).await;
        store.fail.store(true, Ordering::SeqCst);
        let stale = cache
            .get(
                &store,
                "key",
                &sha,
                Duration::from_millis(20),
                Duration::from_secs(1),
            )
            .await
            .unwrap();
        assert_eq!(stale.entries.len(), 1);
        assert_eq!(store.calls.load(Ordering::SeqCst), 2);
    }
}