zinit 0.3.6

//! Xinet proxy - handles multiple frontend sockets and a backend service
//!
//! The proxy:
//! 1. Listens on one or more frontend sockets (TCP and/or Unix)
//! 2. When a client connects, ensures the backend service is running
//! 3. Connects to the backend socket
//! 4. Forwards data bidirectionally
//! 5. Tracks idle time and can stop the service after timeout

use std::sync::Arc;
use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering};
use std::time::{Duration, Instant};
use tokio::net::{TcpListener, UnixListener};
use tokio::sync::{RwLock, Semaphore};
use tracing::{debug, error, info, warn};

use crate::sdk::xinet::{SocketAddr, XinetConfig};
use crate::server::xinet::connection::{Stream, forward_bidirectional, wait_for_socket};

/// Callback type for starting a zinit service
pub type StartServiceFn = Arc<dyn Fn(&str) -> bool + Send + Sync>;

/// Callback type for stopping a zinit service
pub type StopServiceFn = Arc<dyn Fn(&str) -> bool + Send + Sync>;

/// Callback type for checking if a service is running
pub type IsRunningFn = Arc<dyn Fn(&str) -> bool + Send + Sync>;

/// Statistics for a proxy
#[derive(Debug, Default)]
pub struct ProxyStats {
    /// Total connections handled
    pub total_connections: AtomicU64,
    /// Currently active connections
    pub active_connections: AtomicUsize,
    /// Total bytes forwarded client → backend
    pub bytes_to_backend: AtomicU64,
    /// Total bytes forwarded backend → client
    pub bytes_from_backend: AtomicU64,
    /// Last connection timestamp (unix epoch seconds)
    pub last_connection: AtomicU64,
}

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

    pub fn record_connection_start(&self) {
        self.total_connections.fetch_add(1, Ordering::Relaxed);
        self.active_connections.fetch_add(1, Ordering::Relaxed);
        self.last_connection.store(
            std::time::SystemTime::now()
                .duration_since(std::time::UNIX_EPOCH)
                .unwrap_or_default()
                .as_secs(),
            Ordering::Relaxed,
        );
    }

    pub fn record_connection_end(&self, to_backend: u64, from_backend: u64) {
        self.active_connections.fetch_sub(1, Ordering::Relaxed);
        self.bytes_to_backend
            .fetch_add(to_backend, Ordering::Relaxed);
        self.bytes_from_backend
            .fetch_add(from_backend, Ordering::Relaxed);
    }
}

/// A single xinet proxy instance
pub struct XinetProxy {
    /// Configuration
    config: XinetConfig,

    /// Statistics
    stats: Arc<ProxyStats>,

    /// Callback to start the backend service
    start_service: StartServiceFn,

    /// Callback to stop the backend service
    stop_service: StopServiceFn,

    /// Callback to check if service is running
    is_running: IsRunningFn,

    /// Semaphore for single-connection mode
    connection_semaphore: Option<Arc<Semaphore>>,

    /// Last activity timestamp for idle timeout
    last_activity: Arc<RwLock<Instant>>,

    /// Shutdown signal
    shutdown: Arc<RwLock<bool>>,
}

impl XinetProxy {
    /// Create a new proxy
    pub fn new(
        config: XinetConfig,
        start_service: StartServiceFn,
        stop_service: StopServiceFn,
        is_running: IsRunningFn,
    ) -> Self {
        let connection_semaphore = if config.single_connection {
            Some(Arc::new(Semaphore::new(1)))
        } else {
            None
        };

        Self {
            config,
            stats: Arc::new(ProxyStats::new()),
            start_service,
            stop_service,
            is_running,
            connection_semaphore,
            last_activity: Arc::new(RwLock::new(Instant::now())),
            shutdown: Arc::new(RwLock::new(false)),
        }
    }

    /// Get the proxy configuration
    pub fn config(&self) -> &XinetConfig {
        &self.config
    }

    /// Get statistics
    pub fn stats(&self) -> &Arc<ProxyStats> {
        &self.stats
    }

    /// Signal the proxy to shutdown
    pub async fn shutdown(&self) {
        *self.shutdown.write().await = true;
    }

    /// Ensure the backend service is running
    async fn ensure_service_running(&self) -> Result<(), String> {
        // Check if already running
        if (self.is_running)(&self.config.service) {
            return Ok(());
        }

        // Start the service
        info!(
            "[xinet:{}] Starting backend service '{}'",
            self.config.name, self.config.service
        );

        if !(self.start_service)(&self.config.service) {
            return Err(format!("Failed to start service '{}'", self.config.service));
        }

        // Wait for the backend socket to become available
        info!(
            "[xinet:{}] Waiting for backend socket {} (timeout: {}s)",
            self.config.name, self.config.backend, self.config.connect_timeout
        );

        if !wait_for_socket(&self.config.backend, self.config.connect_timeout).await {
            return Err(format!(
                "Backend socket {} not available after {}s",
                self.config.backend, self.config.connect_timeout
            ));
        }

        info!(
            "[xinet:{}] Backend socket {} is ready",
            self.config.name, self.config.backend
        );

        Ok(())
    }

    /// Update last activity timestamp
    async fn touch_activity(&self) {
        *self.last_activity.write().await = Instant::now();
    }

    /// Handle a single client connection
    async fn handle_connection(&self, client_stream: Stream) {
        self.stats.record_connection_start();
        self.touch_activity().await;

        let result = self.do_handle_connection(client_stream).await;

        match result {
            Ok((to_backend, from_backend)) => {
                debug!(
                    "[xinet:{}] Connection completed: {} bytes to backend, {} bytes from backend",
                    self.config.name, to_backend, from_backend
                );
                self.stats.record_connection_end(to_backend, from_backend);
            }
            Err(e) => {
                warn!("[xinet:{}] Connection error: {}", self.config.name, e);
                self.stats.record_connection_end(0, 0);
            }
        }

        self.touch_activity().await;
    }

    /// Inner connection handler
    async fn do_handle_connection(&self, client_stream: Stream) -> Result<(u64, u64), String> {
        // Ensure service is running
        self.ensure_service_running().await?;

        // Connect to backend
        let backend_stream = Stream::connect(&self.config.backend)
            .await
            .map_err(|e| format!("Failed to connect to backend: {}", e))?;

        // Split streams
        let (client_read, client_write) = client_stream.into_split();
        let (backend_read, backend_write) = backend_stream.into_split();

        // Forward data bidirectionally
        let (to_backend, from_backend) =
            forward_bidirectional(client_read, client_write, backend_read, backend_write)
                .await
                .map_err(|e| format!("Forward error: {}", e))?;

        Ok((to_backend, from_backend))
    }

    /// Run the idle timeout monitor
    async fn run_idle_monitor(self: Arc<Self>) {
        if self.config.idle_timeout == 0 {
            return; // Disabled
        }

        let idle_duration = Duration::from_secs(self.config.idle_timeout);

        loop {
            tokio::time::sleep(Duration::from_secs(10)).await;

            if *self.shutdown.read().await {
                break;
            }

            // Check if we have active connections
            if self.stats.active_connections.load(Ordering::Relaxed) > 0 {
                continue;
            }

            // Check idle time
            let last = *self.last_activity.read().await;
            if last.elapsed() >= idle_duration {
                // Check if service is running
                if (self.is_running)(&self.config.service) {
                    info!(
                        "[xinet:{}] Idle timeout reached ({}s), stopping service '{}'",
                        self.config.name, self.config.idle_timeout, self.config.service
                    );
                    (self.stop_service)(&self.config.service);
                }
            }
        }
    }

    /// Run a single listener for a specific address
    async fn run_listener(self: Arc<Self>, addr: SocketAddr) -> Result<(), String> {
        match addr {
            SocketAddr::Unix(ref path) => {
                // Clean up old socket file
                if path.exists() {
                    let _ = std::fs::remove_file(path);
                }
                // Ensure parent directory exists
                if let Some(parent) = path.parent() {
                    let _ = std::fs::create_dir_all(parent);
                }

                let listener = UnixListener::bind(path)
                    .map_err(|e| format!("Failed to bind Unix socket {}: {}", path.display(), e))?;

                info!("[xinet:{}] Listening on {}", self.config.name, addr);

                self.accept_loop_unix(listener).await
            }
            SocketAddr::Tcp(ref tcp_addr) => {
                let listener = TcpListener::bind(tcp_addr)
                    .await
                    .map_err(|e| format!("Failed to bind TCP socket {}: {}", tcp_addr, e))?;

                info!("[xinet:{}] Listening on {}", self.config.name, addr);

                self.accept_loop_tcp(listener).await
            }
        }
    }

    /// Run the proxy (blocking) - spawns listeners for all configured addresses
    pub async fn run(self: Arc<Self>) -> Result<(), String> {
        // Start idle monitor task
        let idle_self = Arc::clone(&self);
        tokio::spawn(async move {
            idle_self.run_idle_monitor().await;
        });

        // Start listeners for all addresses
        let listen_addrs = self.config.listen.clone();

        if listen_addrs.is_empty() {
            return Err("No listen addresses configured".to_string());
        }

        // If only one listener, run it directly
        if listen_addrs.len() == 1 {
            return self.run_listener(listen_addrs[0].clone()).await;
        }

        // Multiple listeners - spawn tasks for each
        let mut handles = Vec::new();

        for addr in listen_addrs {
            let proxy = Arc::clone(&self);
            let handle = tokio::spawn(async move {
                if let Err(e) = proxy.run_listener(addr.clone()).await {
                    error!("[xinet] Listener error for {}: {}", addr, e);
                }
            });
            handles.push(handle);
        }

        // Wait for shutdown signal
        loop {
            if *self.shutdown.read().await {
                break;
            }
            tokio::time::sleep(Duration::from_millis(100)).await;
        }

        // Abort all listener tasks
        for handle in handles {
            handle.abort();
        }

        Ok(())
    }

    /// Accept loop for Unix sockets
    async fn accept_loop_unix(self: Arc<Self>, listener: UnixListener) -> Result<(), String> {
        loop {
            if *self.shutdown.read().await {
                break;
            }

            let accept_result = tokio::select! {
                result = listener.accept() => result,
                _ = tokio::time::sleep(Duration::from_millis(100)) => continue,
            };

            match accept_result {
                Ok((stream, _addr)) => {
                    info!("[xinet:{}] Accepted Unix connection", self.config.name);

                    let proxy = Arc::clone(&self);
                    let permit = if let Some(ref sem) = self.connection_semaphore {
                        match sem.clone().try_acquire_owned() {
                            Ok(p) => Some(p),
                            Err(_) => {
                                warn!(
                                    "[xinet:{}] Connection rejected: single connection mode active",
                                    self.config.name
                                );
                                continue;
                            }
                        }
                    } else {
                        None
                    };

                    tokio::spawn(async move {
                        proxy.handle_connection(Stream::Unix(stream)).await;
                        drop(permit);
                    });
                }
                Err(e) => {
                    warn!("[xinet:{}] Accept error: {}", self.config.name, e);
                }
            }
        }

        Ok(())
    }

    /// Accept loop for TCP sockets
    async fn accept_loop_tcp(self: Arc<Self>, listener: TcpListener) -> Result<(), String> {
        loop {
            if *self.shutdown.read().await {
                break;
            }

            let accept_result = tokio::select! {
                result = listener.accept() => result,
                _ = tokio::time::sleep(Duration::from_millis(100)) => continue,
            };

            match accept_result {
                Ok((stream, addr)) => {
                    debug!(
                        "[xinet:{}] Accepted connection from {}",
                        self.config.name, addr
                    );

                    let proxy = Arc::clone(&self);
                    let permit = if let Some(ref sem) = self.connection_semaphore {
                        match sem.clone().try_acquire_owned() {
                            Ok(p) => Some(p),
                            Err(_) => {
                                warn!(
                                    "[xinet:{}] Connection rejected: single connection mode active",
                                    self.config.name
                                );
                                continue;
                            }
                        }
                    } else {
                        None
                    };

                    tokio::spawn(async move {
                        proxy.handle_connection(Stream::Tcp(stream)).await;
                        drop(permit);
                    });
                }
                Err(e) => {
                    warn!("[xinet:{}] Accept error: {}", self.config.name, e);
                }
            }
        }

        Ok(())
    }
}