network_protocol/transport/

session_cache.rs

//! # TLS Session Cache
//!
//! This module provides in-memory TLS session ticket caching for session resumption.
//! Session resumption reduces the handshake overhead by ~50-70%, allowing clients to
//! reconnect without full TLS 1.3 handshakes.
//!
//! ## Features
//! - **Thread-safe**: Uses Arc<Mutex<>> for safe concurrent access
//! - **TTL-based expiration**: Sessions expire after configurable duration
//! - **Memory-bounded**: Configurable maximum entries to prevent unbounded growth
//! - **Non-blocking lookups**: Fast session retrieval without blocking other operations
//!
//! ## Performance
//! - Lookup: ~100-200ns (HashMap)
//! - Insertion: ~500-1000ns (with lock contention)
//! - Eviction: O(n) for expired entries, but typically O(1) for normal lookups
//!
//! ## Usage
//! ```ignore
//! use network_protocol::transport::session_cache::SessionCache;
//! use std::time::Duration;
//!
//! // Create a cache with 1000 max entries and 1-hour TTL
//! let cache = SessionCache::new(1000, Duration::from_secs(3600));
//!
//! // Store a session (typically handled internally by TLS layer)
//! cache.store(session_id.clone(), ticket.clone()).await;
//!
//! // Retrieve for resumption (returned as Arc<Vec<u8>>)
//! if let Some(ticket) = cache.get(&session_id).await {
//!     // Use ticket for session resumption
//! }
//! ```

use std::collections::HashMap;
use std::sync::Arc;
use std::time::{Duration, SystemTime};

use tokio::sync::Mutex;
use tracing::{debug, trace};

/// A cached TLS session ticket with metadata
#[derive(Clone, Debug)]
struct SessionEntry {
    /// The serialized session ticket from rustls
    ticket: Arc<Vec<u8>>,
    /// When this session was cached
    created_at: SystemTime,
    /// Time-to-live for this session
    ttl: Duration,
}

impl SessionEntry {
    /// Check if this entry has expired
    fn is_expired(&self) -> bool {
        match self.created_at.elapsed() {
            Ok(elapsed) => elapsed > self.ttl,
            Err(_) => true, // System time went backward, treat as expired
        }
    }
}

/// Thread-safe in-memory TLS session cache
///
/// Stores session tickets for TLS 1.3 resumption. This enables clients to reconnect
/// without performing full handshakes, reducing latency by 50-70%.
///
/// Pass the same cache (in an `Arc`) to every
/// [`TlsClient::connect_with_session`](crate::service::tls_client::TlsClient::connect_with_session)
/// call: the TLS layer stores the server's session tickets in it and offers them on
/// the next connection to the same server name. Those TLS tickets are kept in a
/// rustls session store bounded by `max_entries` and expire as the server's ticket
/// lifetime says; `default_ttl` applies to entries added with [`store`](Self::store).
///
/// rustls only resumes a session with the same certificate verifier that checked
/// the original connection, so the cache also keeps the TLS client configuration
/// built for each distinct `TlsClientConfig` and reuses it. Certificate, key and CA
/// files are therefore read on the first connection with a given configuration; call
/// [`clear`](Self::clear) after replacing those files.
#[derive(Clone)]
pub struct SessionCache {
    /// Maximum number of sessions to cache
    max_entries: usize,
    /// Default TTL for new sessions
    default_ttl: Duration,
    /// Inner cache protected by mutex
    inner: Arc<Mutex<SessionCacheInner>>,
    /// rustls session store and client configs used for TLS resumption. Reset by
    /// `clear()`.
    tls: Arc<std::sync::Mutex<TlsResumption>>,
}

/// TLS resumption state shared by every connection that uses one `SessionCache`.
struct TlsResumption {
    store: Arc<rustls::client::ClientSessionMemoryCache>,
    /// Built client configs, keyed by `TlsClientConfig::resumption_key()`
    configs: HashMap<String, Arc<rustls::ClientConfig>>,
}

impl TlsResumption {
    fn new(max_entries: usize) -> Self {
        Self {
            store: Arc::new(rustls::client::ClientSessionMemoryCache::new(
                max_entries.max(1),
            )),
            configs: HashMap::new(),
        }
    }
}

struct SessionCacheInner {
    /// Session ID -> cached ticket
    sessions: HashMap<String, SessionEntry>,
    /// Metadata for eviction policies
    total_inserts: u64,
}

impl SessionCache {
    /// Create a new session cache
    ///
    /// # Arguments
    /// * `max_entries` - Maximum number of sessions to cache (e.g., 1000)
    /// * `default_ttl` - Default time-to-live for each session (e.g., 1 hour)
    ///
    /// # Example
    /// ```ignore
    /// let cache = SessionCache::new(1000, Duration::from_secs(3600));
    /// ```
    pub fn new(max_entries: usize, default_ttl: Duration) -> Self {
        Self {
            max_entries,
            default_ttl,
            tls: Arc::new(std::sync::Mutex::new(TlsResumption::new(max_entries))),
            inner: Arc::new(Mutex::new(SessionCacheInner {
                sessions: HashMap::with_capacity(max_entries),
                total_inserts: 0,
            })),
        }
    }

    /// Store a session ticket in the cache
    ///
    /// This is typically called by the TLS layer after establishing a connection.
    /// Automatically manages eviction when cache is full.
    ///
    /// # Arguments
    /// * `session_id` - Unique session identifier
    /// * `ticket` - Serialized TLS session ticket
    pub async fn store<S: Into<String>>(&self, session_id: S, ticket: Vec<u8>) {
        let mut inner = self.inner.lock().await;

        let session_id = session_id.into();
        let entry = SessionEntry {
            ticket: Arc::new(ticket),
            created_at: SystemTime::now(),
            ttl: self.default_ttl,
        };

        // Clean expired entries before checking capacity
        self.evict_expired(&mut inner);

        // Store the session
        inner.sessions.insert(session_id.clone(), entry);
        inner.total_inserts += 1;

        // Evict oldest if we exceed capacity
        if inner.sessions.len() > self.max_entries {
            self.evict_oldest(&mut inner);
        }

        trace!(
            session_count = inner.sessions.len(),
            "Session stored in cache"
        );
    }

    /// Retrieve a session ticket from the cache
    ///
    /// Returns the ticket if found and not expired, None otherwise.
    ///
    /// # Example
    /// ```ignore
    /// if let Some(ticket) = cache.get("session-123").await {
    ///     // Use ticket for resumption
    /// }
    /// ```
    pub async fn get(&self, session_id: &str) -> Option<Arc<Vec<u8>>> {
        let mut inner = self.inner.lock().await;

        // Check if session exists and is not expired
        if let Some(entry) = inner.sessions.get(session_id) {
            if !entry.is_expired() {
                trace!("Session cache hit");
                return Some(entry.ticket.clone());
            }
        }

        // Remove expired session
        inner.sessions.remove(session_id);
        trace!("Session cache miss or expired");
        None
    }

    /// Clear all sessions from the cache
    pub async fn clear(&self) {
        *self
            .tls
            .lock()
            .unwrap_or_else(std::sync::PoisonError::into_inner) =
            TlsResumption::new(self.max_entries);
        let mut inner = self.inner.lock().await;
        let count = inner.sessions.len();
        inner.sessions.clear();
        debug!(cleared_count = count, "Session cache cleared");
    }

    /// The rustls client config for `config`, built on first use and then shared, with
    /// resumption backed by this cache's session store.
    pub(crate) fn tls_client_config(
        &self,
        config: &crate::transport::tls::TlsClientConfig,
    ) -> crate::error::Result<Arc<rustls::ClientConfig>> {
        let key = config.resumption_key();
        let mut tls = self
            .tls
            .lock()
            .unwrap_or_else(std::sync::PoisonError::into_inner);
        if let Some(built) = tls.configs.get(&key) {
            return Ok(built.clone());
        }
        let mut built = config.load_client_config()?;
        built.resumption = rustls::client::Resumption::store(tls.store.clone());
        let built = Arc::new(built);
        if tls.configs.len() >= self.max_entries.max(1) {
            tls.configs.clear();
        }
        tls.configs.insert(key, built.clone());
        Ok(built)
    }

    /// Get current cache statistics
    pub async fn stats(&self) -> SessionCacheStats {
        let inner = self.inner.lock().await;

        let expired_count = inner.sessions.values().filter(|e| e.is_expired()).count();

        SessionCacheStats {
            total_entries: inner.sessions.len(),
            max_entries: self.max_entries,
            expired_count,
            total_inserts: inner.total_inserts,
        }
    }

    /// Evict all expired entries from the cache
    #[allow(dead_code)]
    async fn evict_expired_async(&self) {
        let mut inner = self.inner.lock().await;
        self.evict_expired(&mut inner);
    }

    /// Internal: Evict expired entries (called with lock held)
    fn evict_expired(&self, inner: &mut SessionCacheInner) {
        let before = inner.sessions.len();
        inner.sessions.retain(|_, entry| !entry.is_expired());
        let after = inner.sessions.len();

        if before != after {
            debug!(
                removed_count = before - after,
                remaining_count = after,
                "Expired sessions evicted"
            );
        }
    }

    /// Internal: Evict oldest entry (called with lock held)
    fn evict_oldest(&self, inner: &mut SessionCacheInner) {
        if let Some(oldest_key) = inner
            .sessions
            .iter()
            .min_by_key(|(_, entry)| entry.created_at)
            .map(|(k, _)| k.clone())
        {
            inner.sessions.remove(&oldest_key);
            debug!("Oldest session evicted to make room");
        }
    }
}

/// Statistics about the session cache
#[derive(Debug, Clone, Copy)]
pub struct SessionCacheStats {
    /// Current number of valid sessions
    pub total_entries: usize,
    /// Maximum capacity
    pub max_entries: usize,
    /// Number of expired but not yet evicted entries
    pub expired_count: usize,
    /// Total sessions ever inserted
    pub total_inserts: u64,
}

#[cfg(test)]
mod tests {
    use super::*;

    #[tokio::test]
    #[allow(clippy::unwrap_used)]
    async fn test_store_and_retrieve() {
        let cache = SessionCache::new(10, Duration::from_secs(60));

        cache.store("session-1", vec![1, 2, 3, 4]).await;
        let ticket = cache.get("session-1").await;

        assert!(ticket.is_some());
        assert_eq!(*ticket.unwrap(), vec![1, 2, 3, 4]);
    }

    #[tokio::test]
    async fn test_missing_session() {
        let cache = SessionCache::new(10, Duration::from_secs(60));
        let ticket = cache.get("nonexistent").await;
        assert!(ticket.is_none());
    }

    #[tokio::test]
    async fn test_capacity_eviction() {
        let cache = SessionCache::new(3, Duration::from_secs(60));

        for i in 0..5 {
            cache.store(format!("session-{i}"), vec![i as u8]).await;
        }

        let stats = cache.stats().await;
        assert_eq!(stats.total_entries, 3);
        assert_eq!(stats.total_inserts, 5);
    }

    #[tokio::test]
    async fn test_clear() {
        let cache = SessionCache::new(10, Duration::from_secs(60));

        cache.store("session-1", vec![1, 2, 3]).await;
        cache.clear().await;

        let ticket = cache.get("session-1").await;
        assert!(ticket.is_none());
    }
}