chanora/crates/chanora_server_prefetch/src/lib.rs

//! Server-address prefetch cache and policy for Chanora.
//!
//! This crate owns speculative server-resolution warming. It does not
//! decide whether a connection should use a prefetched address; callers
//! must still apply their own trust boundary before dialing.

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

use thiserror::Error;
use tokio::sync::Mutex;
use tracing::{debug, info, warn};

const SERVER_PREFETCH_TTL: Duration = Duration::from_secs(120);

#[derive(Debug, Error)]
pub enum ServerPrefetchError {
    #[error("resolver initialization failed: {0}")]
    ResolverInit(String),
    #[error("resolution failed: {0}")]
    Resolution(String),
    #[error("resolver returned invalid socket address '{resolved}': {reason}")]
    InvalidSocketAddress { resolved: String, reason: String },
}

#[derive(Debug, Clone)]
struct ServerPrefetchEntry {
    normalized_host: String,
    resolved_address: SocketAddr,
    completed_at: Instant,
}

#[derive(Debug, Default)]
struct ServerPrefetchCache {
    latest_generation: u64,
    entry: Option<ServerPrefetchEntry>,
}

impl ServerPrefetchCache {
    fn begin(&mut self, host: &str) -> u64 {
        if normalize_host(host).is_empty() {
            return self.latest_generation;
        }
        self.latest_generation = self.latest_generation.saturating_add(1);
        self.latest_generation
    }

    fn store_success(
        &mut self,
        generation: u64,
        host: &str,
        resolved_address: SocketAddr,
        completed_at: Instant,
    ) {
        if generation != self.latest_generation {
            return;
        }
        self.entry = Some(ServerPrefetchEntry {
            normalized_host: normalize_host(host),
            resolved_address,
            completed_at,
        });
    }

    fn fresh_match(&self, host: &str, now: Instant) -> Option<SocketAddr> {
        let normalized = normalize_host(host);
        let entry = self.entry.as_ref()?;
        if entry.normalized_host != normalized {
            return None;
        }
        if now.duration_since(entry.completed_at) > SERVER_PREFETCH_TTL {
            return None;
        }
        // Generation is not checked here: a fresh entry remains usable while
        // a newer prefetch is in flight. Stale async completions are still
        // rejected in store_success via the generation guard.
        Some(entry.resolved_address)
    }
}

#[derive(Debug, Clone, Default)]
pub struct ServerPrefetcher {
    cache: Arc<Mutex<ServerPrefetchCache>>,
    #[cfg(any(test, feature = "test-support"))]
    setup_error_for_test: Arc<Mutex<Option<String>>>,
}

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

    /// Schedules a fire-and-forget prefetch for `host`.
    ///
    /// The result only reports synchronous setup failures before scheduling,
    /// such as resolver initialization. DNS/resolution failures after the task
    /// is spawned are logged and do not complete this returned `Result`.
    pub async fn prefetch(&self, host: String) -> Result<(), ServerPrefetchError> {
        let normalized = normalize_host(&host);
        if normalized.is_empty() {
            return Ok(());
        }

        #[cfg(any(test, feature = "test-support"))]
        if let Some(err) = self.setup_error_for_test.lock().await.take() {
            return Err(ServerPrefetchError::ResolverInit(err));
        }

        let resolver = chanora_resolver::ChanoraResolver::new()
            .map_err(|err| ServerPrefetchError::ResolverInit(err.to_string()))?;

        let generation = {
            let mut cache = self.cache.lock().await;
            cache.begin(&normalized)
        };
        let cache = self.cache.clone();
        tokio::spawn(async move {
            info!(target: "chanora_server_prefetch", host = %normalized, "resolution prefetch started");
            let result = resolve_socket(resolver, &normalized).await;
            match result {
                Ok(addr) => {
                    info!(
                        target: "chanora_server_prefetch",
                        host = %normalized,
                        resolved = %addr,
                        "resolution prefetch result"
                    );
                    let mut guard = cache.lock().await;
                    guard.store_success(generation, &normalized, addr, Instant::now());
                }
                Err(err) => {
                    warn!(
                        target: "chanora_server_prefetch",
                        host = %normalized,
                        error = %err,
                        "resolution prefetch failed"
                    );
                }
            }
        });
        Ok(())
    }

    pub async fn fresh_match(&self, host: &str) -> Option<SocketAddr> {
        let resolved = {
            let cache = self.cache.lock().await;
            cache.fresh_match(host, Instant::now())
        };
        match resolved {
            Some(addr) => {
                info!(
                    target: "chanora_server_prefetch",
                    host = %host,
                    resolved = %addr,
                    "connect using prefetched resolution"
                );
                Some(addr)
            }
            None => {
                debug!(target: "chanora_server_prefetch", host = %host, "connect prefetch miss or stale");
                None
            }
        }
    }

    #[cfg(any(test, feature = "test-support"))]
    pub async fn begin_for_test(&self, host: &str) -> u64 {
        let mut cache = self.cache.lock().await;
        cache.begin(host)
    }

    #[cfg(any(test, feature = "test-support"))]
    pub async fn store_success_for_test(
        &self,
        generation: u64,
        host: &str,
        resolved_address: SocketAddr,
        completed_at: Instant,
    ) {
        let mut cache = self.cache.lock().await;
        cache.store_success(generation, host, resolved_address, completed_at);
    }

    #[cfg(any(test, feature = "test-support"))]
    pub async fn latest_generation_for_test(&self) -> u64 {
        let cache = self.cache.lock().await;
        cache.latest_generation
    }

    #[cfg(any(test, feature = "test-support"))]
    pub async fn fail_next_prefetch_setup_for_test(&self, error: impl Into<String>) {
        *self.setup_error_for_test.lock().await = Some(error.into());
    }
}

fn normalize_host(host: &str) -> String {
    host.trim().to_lowercase()
}

async fn resolve_socket(
    resolver: chanora_resolver::ChanoraResolver,
    host: &str,
) -> Result<SocketAddr, ServerPrefetchError> {
    let resolved = resolver
        .resolve_client_address(host)
        .await
        .map_err(|err| ServerPrefetchError::Resolution(err.to_string()))?;
    resolved
        .parse::<SocketAddr>()
        .map_err(|err| ServerPrefetchError::InvalidSocketAddress {
            resolved,
            reason: err.to_string(),
        })
}

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

    #[tokio::test]
    async fn fresh_exact_match_returns_socket_address() {
        let prefetcher = ServerPrefetcher::new();
        let generation = prefetcher.begin_for_test(" Example.COM ").await;
        let addr = "127.0.0.1:9987".parse().unwrap();
        prefetcher
            .store_success_for_test(generation, "example.com", addr, Instant::now())
            .await;

        assert_eq!(prefetcher.fresh_match("example.com").await, Some(addr));
        assert_eq!(prefetcher.fresh_match(" EXAMPLE.com ").await, Some(addr));
    }

    #[tokio::test]
    async fn stale_entries_are_ignored() {
        let prefetcher = ServerPrefetcher::new();
        let generation = prefetcher.begin_for_test("example.com").await;
        let addr = "127.0.0.1:9987".parse().unwrap();
        prefetcher
            .store_success_for_test(
                generation,
                "example.com",
                addr,
                Instant::now() - SERVER_PREFETCH_TTL - Duration::from_secs(1),
            )
            .await;

        assert_eq!(prefetcher.fresh_match("example.com").await, None);
    }

    #[tokio::test]
    async fn different_hosts_are_ignored() {
        let prefetcher = ServerPrefetcher::new();
        let generation = prefetcher.begin_for_test("example.com").await;
        let addr = "127.0.0.1:9987".parse().unwrap();
        prefetcher
            .store_success_for_test(generation, "example.com", addr, Instant::now())
            .await;

        assert_eq!(prefetcher.fresh_match("other.example.com").await, None);
    }

    #[tokio::test]
    async fn stale_generation_completions_are_ignored() {
        let prefetcher = ServerPrefetcher::new();
        let old_generation = prefetcher.begin_for_test("old.example.com").await;
        let _new_generation = prefetcher.begin_for_test("new.example.com").await;
        let old_addr = "127.0.0.1:9987".parse().unwrap();
        prefetcher
            .store_success_for_test(old_generation, "old.example.com", old_addr, Instant::now())
            .await;

        assert_eq!(prefetcher.fresh_match("old.example.com").await, None);
    }

    #[tokio::test]
    async fn fresh_entry_usable_during_warming() {
        let prefetcher = ServerPrefetcher::new();
        let first_generation = prefetcher.begin_for_test("example.com").await;
        let cached_addr = "127.0.0.1:9987".parse().unwrap();
        prefetcher
            .store_success_for_test(first_generation, "example.com", cached_addr, Instant::now())
            .await;

        let _pending_generation = prefetcher.begin_for_test("example.com").await;

        assert_eq!(
            prefetcher.fresh_match("example.com").await,
            Some(cached_addr)
        );

        let stale_addr = "127.0.0.2:9987".parse().unwrap();
        prefetcher
            .store_success_for_test(first_generation, "example.com", stale_addr, Instant::now())
            .await;

        assert_eq!(
            prefetcher.fresh_match("example.com").await,
            Some(cached_addr)
        );
    }

    #[tokio::test]
    async fn blank_host_skips_prefetch() {
        let prefetcher = ServerPrefetcher::new();
        let before = prefetcher.latest_generation_for_test().await;

        prefetcher.prefetch("   ".to_string()).await.unwrap();

        assert_eq!(prefetcher.latest_generation_for_test().await, before);
    }
}