veridian_kernel/services/mesh/

proxy.rs

//! Sidecar Proxy
//!
//! Provides a connection-pooling sidecar proxy with health checking,
//! mTLS wrapping, and request routing for service mesh communication.

#![allow(dead_code)]

use alloc::{string::String, vec::Vec};
use core::sync::atomic::{AtomicU64, Ordering};

// ---------------------------------------------------------------------------
// Proxy Config
// ---------------------------------------------------------------------------

/// Sidecar proxy configuration.
#[derive(Debug, Clone)]
pub struct ProxyConfig {
    /// Port to listen on.
    pub listen_port: u16,
    /// Upstream service addresses (ip:port strings).
    pub upstream_addrs: Vec<String>,
    /// Health check interval in ticks.
    pub health_check_interval: u64,
    /// Connection timeout in ticks.
    pub connect_timeout: u64,
    /// Maximum retries on upstream failure.
    pub max_retries: u32,
}

impl Default for ProxyConfig {
    fn default() -> Self {
        ProxyConfig {
            listen_port: 15001,
            upstream_addrs: Vec::new(),
            health_check_interval: 10,
            connect_timeout: 5,
            max_retries: 3,
        }
    }
}

// ---------------------------------------------------------------------------
// Connection Pool
// ---------------------------------------------------------------------------

/// Connection state.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ConnectionState {
    /// Connection is idle and available.
    Idle,
    /// Connection is actively in use.
    Active,
    /// Connection is draining (finishing pending requests).
    Draining,
}

/// A pooled connection to an upstream.
#[derive(Debug, Clone)]
pub struct Connection {
    /// Unique connection identifier.
    pub id: u64,
    /// Index into upstream_addrs.
    pub upstream_idx: usize,
    /// Current state.
    pub state: ConnectionState,
    /// Number of requests handled.
    pub request_count: u64,
    /// Tick when created.
    pub created_tick: u64,
}

/// Next connection ID generator.
static NEXT_CONN_ID: AtomicU64 = AtomicU64::new(1);

/// Connection pool for upstream services.
#[derive(Debug)]
pub struct ConnectionPool {
    /// Pooled connections.
    connections: Vec<Connection>,
    /// Maximum connections per upstream.
    max_per_upstream: usize,
}

impl Default for ConnectionPool {
    fn default() -> Self {
        Self::new(10)
    }
}

impl ConnectionPool {
    /// Create a new connection pool.
    pub fn new(max_per_upstream: usize) -> Self {
        ConnectionPool {
            connections: Vec::new(),
            max_per_upstream,
        }
    }

    /// Get an idle connection to a specific upstream, or create one.
    pub fn get_connection(&mut self, upstream_idx: usize, current_tick: u64) -> &mut Connection {
        // Try to find an idle connection
        let idle_pos = self
            .connections
            .iter()
            .position(|c| c.upstream_idx == upstream_idx && c.state == ConnectionState::Idle);

        if let Some(pos) = idle_pos {
            self.connections[pos].state = ConnectionState::Active;
            return &mut self.connections[pos];
        }

        // Create a new connection
        let conn = Connection {
            id: NEXT_CONN_ID.fetch_add(1, Ordering::Relaxed),
            upstream_idx,
            state: ConnectionState::Active,
            request_count: 0,
            created_tick: current_tick,
        };
        self.connections.push(conn);
        let last = self.connections.len() - 1;
        &mut self.connections[last]
    }

    /// Release a connection back to the pool.
    pub fn release(&mut self, conn_id: u64) {
        if let Some(conn) = self.connections.iter_mut().find(|c| c.id == conn_id) {
            conn.state = ConnectionState::Idle;
        }
    }

    /// Drain connections to a specific upstream.
    pub fn drain_upstream(&mut self, upstream_idx: usize) {
        for conn in &mut self.connections {
            if conn.upstream_idx == upstream_idx {
                conn.state = ConnectionState::Draining;
            }
        }
    }

    /// Remove drained connections.
    pub fn cleanup_drained(&mut self) -> usize {
        let before = self.connections.len();
        self.connections
            .retain(|c| c.state != ConnectionState::Draining);
        before - self.connections.len()
    }

    /// Get the number of active connections.
    pub fn active_count(&self) -> usize {
        self.connections
            .iter()
            .filter(|c| c.state == ConnectionState::Active)
            .count()
    }

    /// Get total connection count.
    pub fn total_count(&self) -> usize {
        self.connections.len()
    }
}

// ---------------------------------------------------------------------------
// Proxy Stats
// ---------------------------------------------------------------------------

/// Proxy statistics.
#[derive(Debug, Default)]
pub struct ProxyStats {
    /// Total requests proxied.
    pub total_requests: AtomicU64,
    /// Total failed requests.
    pub failed_requests: AtomicU64,
    /// Total bytes sent upstream.
    pub bytes_sent: AtomicU64,
    /// Total bytes received from upstream.
    pub bytes_received: AtomicU64,
    /// Total health check passes.
    pub health_checks_passed: AtomicU64,
    /// Total health check failures.
    pub health_checks_failed: AtomicU64,
}

impl ProxyStats {
    /// Create new stats.
    pub fn new() -> Self {
        Self::default()
    }

    /// Record a successful request.
    pub fn record_success(&self, bytes_sent: u64, bytes_received: u64) {
        self.total_requests.fetch_add(1, Ordering::Relaxed);
        self.bytes_sent.fetch_add(bytes_sent, Ordering::Relaxed);
        self.bytes_received
            .fetch_add(bytes_received, Ordering::Relaxed);
    }

    /// Record a failed request.
    pub fn record_failure(&self) {
        self.total_requests.fetch_add(1, Ordering::Relaxed);
        self.failed_requests.fetch_add(1, Ordering::Relaxed);
    }
}

// ---------------------------------------------------------------------------
// Upstream Health
// ---------------------------------------------------------------------------

/// Upstream health status.
#[derive(Debug, Clone)]
pub struct UpstreamHealth {
    /// Upstream index.
    pub idx: usize,
    /// Whether the upstream is healthy.
    pub healthy: bool,
    /// Consecutive failure count.
    pub consecutive_failures: u32,
    /// Last check tick.
    pub last_check_tick: u64,
}

// ---------------------------------------------------------------------------
// Sidecar Proxy
// ---------------------------------------------------------------------------

/// Proxy error.
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum ProxyError {
    /// No healthy upstream available.
    NoHealthyUpstream,
    /// Connection failed.
    ConnectionFailed(String),
    /// Upstream timeout.
    UpstreamTimeout,
    /// Invalid request.
    InvalidRequest(String),
}

/// Sidecar proxy implementation.
#[derive(Debug)]
pub struct SidecarProxy {
    /// Proxy configuration.
    config: ProxyConfig,
    /// Connection pool.
    pool: ConnectionPool,
    /// Proxy statistics.
    stats: ProxyStats,
    /// Upstream health status.
    upstream_health: Vec<UpstreamHealth>,
    /// Round-robin counter for upstream selection.
    rr_counter: u64,
    /// Whether mTLS is enabled.
    mtls_enabled: bool,
}

impl SidecarProxy {
    /// Create a new sidecar proxy.
    pub fn new(config: ProxyConfig) -> Self {
        let health: Vec<UpstreamHealth> = (0..config.upstream_addrs.len())
            .map(|idx| UpstreamHealth {
                idx,
                healthy: true,
                consecutive_failures: 0,
                last_check_tick: 0,
            })
            .collect();

        SidecarProxy {
            config,
            pool: ConnectionPool::default(),
            stats: ProxyStats::new(),
            upstream_health: health,
            rr_counter: 0,
            mtls_enabled: false,
        }
    }

    /// Enable or disable mTLS.
    pub fn set_mtls(&mut self, enabled: bool) {
        self.mtls_enabled = enabled;
    }

    /// Select a healthy upstream using round-robin.
    fn select_upstream(&mut self) -> Result<usize, ProxyError> {
        let count = self.upstream_health.len();
        if count == 0 {
            return Err(ProxyError::NoHealthyUpstream);
        }

        for _ in 0..count {
            let idx = (self.rr_counter as usize) % count;
            self.rr_counter += 1;
            if self.upstream_health[idx].healthy {
                return Ok(idx);
            }
        }
        Err(ProxyError::NoHealthyUpstream)
    }

    /// Handle a TCP connection by routing to an upstream.
    pub fn handle_tcp(&mut self, payload: &[u8], current_tick: u64) -> Result<Vec<u8>, ProxyError> {
        let upstream_idx = self.select_upstream()?;
        let conn = self.pool.get_connection(upstream_idx, current_tick);
        conn.request_count += 1;
        let conn_id = conn.id;

        // Simulate forwarding
        self.stats.record_success(payload.len() as u64, 0);
        self.pool.release(conn_id);

        Ok(Vec::new()) // Response would come from actual upstream
    }

    /// Handle an HTTP request by routing to an upstream.
    pub fn handle_http(
        &mut self,
        _method: &str,
        _path: &str,
        payload: &[u8],
        current_tick: u64,
    ) -> Result<Vec<u8>, ProxyError> {
        let upstream_idx = self.select_upstream()?;
        let conn = self.pool.get_connection(upstream_idx, current_tick);
        conn.request_count += 1;
        let conn_id = conn.id;

        self.stats.record_success(payload.len() as u64, 0);
        self.pool.release(conn_id);

        Ok(Vec::new())
    }

    /// Run health checks on all upstreams.
    pub fn health_check(&mut self, current_tick: u64) {
        for health in &mut self.upstream_health {
            health.last_check_tick = current_tick;
            // Simulated: always passes unless manually marked unhealthy
            if health.healthy {
                health.consecutive_failures = 0;
                self.stats
                    .health_checks_passed
                    .fetch_add(1, Ordering::Relaxed);
            } else {
                health.consecutive_failures += 1;
                self.stats
                    .health_checks_failed
                    .fetch_add(1, Ordering::Relaxed);
            }
        }
    }

    /// Mark an upstream as unhealthy.
    pub fn mark_unhealthy(&mut self, idx: usize) {
        if idx < self.upstream_health.len() {
            self.upstream_health[idx].healthy = false;
        }
    }

    /// Mark an upstream as healthy.
    pub fn mark_healthy(&mut self, idx: usize) {
        if idx < self.upstream_health.len() {
            self.upstream_health[idx].healthy = true;
            self.upstream_health[idx].consecutive_failures = 0;
        }
    }

    /// Get the number of healthy upstreams.
    pub fn healthy_upstream_count(&self) -> usize {
        self.upstream_health.iter().filter(|h| h.healthy).count()
    }

    /// Get proxy stats.
    pub fn stats(&self) -> &ProxyStats {
        &self.stats
    }

    /// Get connection pool.
    pub fn pool(&self) -> &ConnectionPool {
        &self.pool
    }
}

// ---------------------------------------------------------------------------
// Tests
// ---------------------------------------------------------------------------

#[cfg(test)]
mod tests {
    #[allow(unused_imports)]
    use alloc::string::ToString;

    use super::*;

    fn make_proxy() -> SidecarProxy {
        let config = ProxyConfig {
            listen_port: 15001,
            upstream_addrs: alloc::vec![
                String::from("10.0.0.1:8080"),
                String::from("10.0.0.2:8080"),
                String::from("10.0.0.3:8080"),
            ],
            health_check_interval: 10,
            connect_timeout: 5,
            max_retries: 3,
        };
        SidecarProxy::new(config)
    }

    #[test]
    fn test_proxy_creation() {
        let proxy = make_proxy();
        assert_eq!(proxy.healthy_upstream_count(), 3);
    }

    #[test]
    fn test_handle_tcp() {
        let mut proxy = make_proxy();
        let result = proxy.handle_tcp(&[1, 2, 3], 100);
        assert!(result.is_ok());
        assert_eq!(proxy.stats.total_requests.load(Ordering::Relaxed), 1);
    }

    #[test]
    fn test_handle_http() {
        let mut proxy = make_proxy();
        let result = proxy.handle_http("GET", "/api/v1/pods", &[], 100);
        assert!(result.is_ok());
    }

    #[test]
    fn test_round_robin_selection() {
        let mut proxy = make_proxy();
        // Should cycle through upstreams
        proxy.handle_tcp(&[], 100).unwrap();
        proxy.handle_tcp(&[], 200).unwrap();
        proxy.handle_tcp(&[], 300).unwrap();
        assert_eq!(proxy.stats.total_requests.load(Ordering::Relaxed), 3);
    }

    #[test]
    fn test_no_healthy_upstream() {
        let mut proxy = make_proxy();
        proxy.mark_unhealthy(0);
        proxy.mark_unhealthy(1);
        proxy.mark_unhealthy(2);
        assert_eq!(
            proxy.handle_tcp(&[], 100),
            Err(ProxyError::NoHealthyUpstream)
        );
    }

    #[test]
    fn test_health_check() {
        let mut proxy = make_proxy();
        proxy.health_check(100);
        assert_eq!(proxy.stats.health_checks_passed.load(Ordering::Relaxed), 3);
    }

    #[test]
    fn test_mark_unhealthy_healthy() {
        let mut proxy = make_proxy();
        proxy.mark_unhealthy(1);
        assert_eq!(proxy.healthy_upstream_count(), 2);
        proxy.mark_healthy(1);
        assert_eq!(proxy.healthy_upstream_count(), 3);
    }

    #[test]
    fn test_connection_pool_lifecycle() {
        let mut pool = ConnectionPool::new(5);
        let conn = pool.get_connection(0, 100);
        let conn_id = conn.id;
        assert_eq!(pool.active_count(), 1);
        pool.release(conn_id);
        assert_eq!(pool.active_count(), 0);
    }

    #[test]
    fn test_connection_pool_drain() {
        let mut pool = ConnectionPool::new(5);
        pool.get_connection(0, 100);
        pool.get_connection(0, 200);
        pool.drain_upstream(0);
        let removed = pool.cleanup_drained();
        assert_eq!(removed, 2);
    }

    #[test]
    fn test_mtls_toggle() {
        let mut proxy = make_proxy();
        assert!(!proxy.mtls_enabled);
        proxy.set_mtls(true);
        assert!(proxy.mtls_enabled);
    }
}