product-os-proxy 0.0.19

//! MITM Proxy core implementation.
//!
//! This module provides the [`Proxy`] struct for running a man-in-the-middle
//! HTTP/HTTPS proxy server, along with the [`ProxyBuilder`] for configuration
//! and the [`ProxyState`] enum for lifecycle management.
//!
//! # Architecture
//!
//! The proxy operates by:
//! 1. Binding a TCP listener on the configured address
//! 2. Accepting incoming connections
//! 3. Spawning handler tasks for each connection via [`InternalProxy`]
//! 4. Supporting graceful shutdown through [`Proxy::stop`]
//!
//! # Usage
//!
//! ```ignore
//! let proxy = Proxy::https_builder()
//!     .use_address(addr)
//!     .use_certificate_authority(certs)
//!     .use_rustls_client(provider)
//!     .build()?;
//!
//! proxy.start().await?;
//! ```

mod internal;

pub mod builder;

pub use builder::ProxyBuilder;

use crate::mitm::{
    certificate_authority::CertificateAuthority, Error, HttpHandler, WebSocketHandler,
};

use builder::AddrOrListener;

use hyper_util::{
    client::legacy::{
        connect::{Connect, HttpConnector},
        Client,
    },
    rt::{TokioExecutor, TokioIo},
    server::conn::auto::Builder,
};

use internal::InternalProxy;
use std::sync::Arc;
use tokio::net::TcpListener;
use tokio::sync::oneshot;
use tokio::sync::watch;
use tokio_tungstenite::Connector;

use hyper::service::service_fn;
#[cfg(feature = "rustls_client")]
use hyper_rustls::HttpsConnector;
use product_os_http_body::BodyBytes;
use product_os_utilities::ProductOSError;

#[cfg(feature = "rcgen_ca")]
use crate::mitm::certificate_authority::RcgenAuthority;
use crate::ProxyMiddleware;

/// Represents the lifecycle state of the proxy server.
///
/// The proxy transitions through these states during its lifecycle:
///
/// ```text
/// Uninitialized -> Initialized -> Started -> StopRequested -> Stopped
/// ```
///
/// # Examples
///
/// ```ignore
/// use product_os_proxy::ProxyState;
///
/// let state = ProxyState::Started;
/// assert_eq!(format!("{:?}", state), "Started");
/// ```
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum ProxyState {
    /// Proxy has been created but not yet built or configured
    Uninitialized,
    /// Proxy has been built and configured but not started
    Initialized,
    /// Proxy is running and accepting connections
    Started,
    /// A stop has been requested but the proxy hasn't fully shut down yet
    StopRequested,
    /// Proxy has been stopped and is no longer accepting connections
    Stopped,
}

/// The core MITM proxy server.
///
/// This struct holds all the configuration and state needed to run a proxy server.
/// It is created through [`ProxyBuilder`] and supports starting, stopping, and
/// querying the state of the proxy.
///
/// # Type Parameters
///
/// * `C` - The HTTP connector type (e.g., `HttpsConnector<HttpConnector>`)
/// * `CA` - The certificate authority type (e.g., [`RcgenAuthority`])
/// * `H` - The HTTP handler type (e.g., [`ProxyMiddleware`])
/// * `W` - The WebSocket handler type (e.g., [`ProxyMiddleware`])
pub struct Proxy<C, CA, H, W> {
    pub(crate) address_or_listener: std::sync::Mutex<Option<AddrOrListener>>,
    pub(crate) ca: Arc<CA>,
    pub(crate) client: Client<C, BodyBytes>,
    pub(crate) websocket_connector: Option<Connector>,

    pub(crate) http_handler: H,
    pub(crate) websocket_handler: W,

    #[allow(dead_code)]
    pub(crate) certificates: product_os_security::certificates::Certificates,
    pub(crate) custom_requester: Option<product_os_request::ProductOSRequestClient>,
    pub(crate) compression: crate::config::NetworkProxyCompression,
    pub(crate) request_timeout_ms: Option<u64>,

    pub(crate) server: Option<Builder<TokioExecutor>>,

    /// Sender half of the shutdown signal channel
    shutdown_tx: watch::Sender<bool>,
    /// Receiver half of the shutdown signal channel
    shutdown_rx: watch::Receiver<bool>,
    /// Current state of the proxy
    state: std::sync::Arc<std::sync::atomic::AtomicU8>,
    /// Notified when the proxy transitions to Stopped state
    stopped_notify: Arc<tokio::sync::Notify>,

    #[cfg(feature = "tor")]
    pub(crate) tor_client: Option<arti_client::TorClient<tor_rtcompat::PreferredRuntime>>,
    #[cfg(feature = "vpn")]
    pub(crate) vpn_client: Option<product_os_vpn::ProductOSVPN>,
}

#[cfg(all(feature = "rustls_client", feature = "rcgen_ca"))]
impl Proxy<(), RcgenAuthority, ProxyMiddleware, ProxyMiddleware> {
    /// Creates an HTTPS proxy builder with rustls TLS support.
    ///
    /// This is the recommended way to create a proxy for intercepting HTTPS traffic.
    ///
    /// # Returns
    ///
    /// A [`ProxyBuilder`] configured for HTTPS connections
    pub fn https_builder(
    ) -> ProxyBuilder<HttpsConnector<HttpConnector>, RcgenAuthority, ProxyMiddleware, ProxyMiddleware>
    {
        ProxyBuilder::new()
    }

    /// Creates an HTTP-only proxy builder (no TLS).
    ///
    /// Use this when TLS interception is not needed.
    ///
    /// # Returns
    ///
    /// A [`ProxyBuilder`] configured for HTTP-only connections
    #[allow(dead_code)]
    pub fn http_builder(
    ) -> ProxyBuilder<HttpConnector, RcgenAuthority, ProxyMiddleware, ProxyMiddleware> {
        ProxyBuilder::new()
    }
}

/// Map from state enum to atomic u8
impl ProxyState {
    fn to_u8(self) -> u8 {
        match self {
            ProxyState::Uninitialized => 0,
            ProxyState::Initialized => 1,
            ProxyState::Started => 2,
            ProxyState::StopRequested => 3,
            ProxyState::Stopped => 4,
        }
    }

    fn from_u8(val: u8) -> Self {
        match val {
            1 => ProxyState::Initialized,
            2 => ProxyState::Started,
            3 => ProxyState::StopRequested,
            4 => ProxyState::Stopped,
            _ => ProxyState::Uninitialized,
        }
    }
}

struct ConnectionGuard {
    count: Arc<std::sync::atomic::AtomicUsize>,
    notify: Arc<tokio::sync::Notify>,
}

impl Drop for ConnectionGuard {
    fn drop(&mut self) {
        if self
            .count
            .fetch_sub(1, std::sync::atomic::Ordering::Relaxed)
            == 1
        {
            self.notify.notify_waiters();
        }
    }
}

impl<C, CA, H, W> Proxy<C, CA, H, W>
where
    C: Connect + Clone + Send + Sync + 'static,
    CA: CertificateAuthority,
    H: HttpHandler + Clone,
    W: WebSocketHandler + Clone,
{
    /// Returns the current state of the proxy.
    ///
    /// # Returns
    ///
    /// The current [`ProxyState`]
    #[must_use]
    pub fn get_state(&self) -> ProxyState {
        ProxyState::from_u8(self.state.load(std::sync::atomic::Ordering::Relaxed))
    }

    /// Requests the proxy to stop gracefully.
    ///
    /// This sends a shutdown signal to the running proxy. Active connections
    /// will be allowed to complete before the proxy fully shuts down.
    ///
    /// # Returns
    ///
    /// * `Ok(())` - Shutdown signal sent successfully
    /// * `Err(ProductOSError)` - If the proxy is not in a running state
    pub async fn stop(&self) -> Result<(), ProductOSError> {
        const GRACEFUL_SHUTDOWN_TIMEOUT_SECS: u64 = 30;

        let current = self.get_state();
        if current != ProxyState::Started {
            return Err(ProductOSError::GenericError(format!(
                "Cannot stop proxy in state {current:?}"
            )));
        }

        self.state.store(
            ProxyState::StopRequested.to_u8(),
            std::sync::atomic::Ordering::Relaxed,
        );

        let _ = self.shutdown_tx.send(true);

        // Wait for start() to finish draining and transition to Stopped, with a maximum wait
        let stopped = self.stopped_notify.notified();
        if self.get_state() != ProxyState::Stopped
            && tokio::time::timeout(
                std::time::Duration::from_secs(GRACEFUL_SHUTDOWN_TIMEOUT_SECS),
                stopped,
            )
            .await
            .is_err()
        {
            tracing::warn!(
                "Proxy did not stop within {}s",
                GRACEFUL_SHUTDOWN_TIMEOUT_SECS
            );
            return Err(ProductOSError::GenericError(format!(
                "Proxy did not stop within {GRACEFUL_SHUTDOWN_TIMEOUT_SECS}s"
            )));
        }

        Ok(())
    }

    pub(crate) async fn start_with_signal(
        self: Arc<Self>,
        started_tx: Option<oneshot::Sender<Result<(), String>>>,
    ) -> Result<(), Error> {
        let mut started_tx = started_tx;

        // Take the address_or_listener out of the mutex, allowing us to
        // consume a TcpListener by value even though `self` is behind Arc.
        let addr_or_listener = self
            .address_or_listener
            .lock()
            .map_err(|_| {
                Error::Io(std::io::Error::other(
                    "Mutex poisoned on address_or_listener",
                ))
            })?
            .take()
            .ok_or_else(|| {
                Error::Io(std::io::Error::other(
                    "Proxy has already been started (address_or_listener was already consumed)",
                ))
            })?;

        let listener = match addr_or_listener {
            AddrOrListener::Addr(addr) => match TcpListener::bind(addr).await {
                Ok(listener) => listener,
                Err(err) => {
                    if let Some(tx) = started_tx.take() {
                        let _ = tx.send(Err(err.to_string()));
                    }
                    return Err(err.into());
                }
            },
            AddrOrListener::Listener(listener) => listener,
        };

        let ca = Arc::clone(&self.ca);

        let client = self.client.clone();
        let websocket_connector = self.websocket_connector.clone();

        let http_handler = self.http_handler.clone();
        let websocket_handler = self.websocket_handler.clone();

        let mut shutdown_rx = self.shutdown_rx.clone();
        let request_timeout_ms = self.request_timeout_ms;

        #[cfg(feature = "tor")]
        let tor_client = self.tor_client.clone();

        #[cfg(feature = "vpn")]
        let vpn_client = self.vpn_client.clone();

        let server = self.server.clone().unwrap_or_else(|| {
            let mut builder = Builder::new(TokioExecutor::new());
            builder
                .http1()
                .title_case_headers(true)
                .preserve_header_case(true)
                .http2()
                .enable_connect_protocol();
            builder
        });

        // Mark as started once the listener is bound and the accept loop is ready.
        self.state.store(
            ProxyState::Started.to_u8(),
            std::sync::atomic::Ordering::Relaxed,
        );
        if let Some(tx) = started_tx.take() {
            let _ = tx.send(Ok(()));
        }

        // Track active connection tasks for graceful shutdown
        let active_connections = Arc::new(tokio::sync::Notify::new());
        let connection_count = Arc::new(std::sync::atomic::AtomicUsize::new(0));

        loop {
            let custom_requester = self.custom_requester.clone();
            let compression = self.compression;

            tokio::select! {
                res = listener.accept() => {
                    let (tcp, client_addr) = match res {
                        Ok((tcp, client_addr)) => (tcp, client_addr),
                        Err(e) => {
                            tracing::error!("Failed to accept incoming connection: {}", e);
                            continue;
                        }
                    };

                    let server = server.clone();
                    let client = client.clone();
                    let ca = Arc::clone(&ca);
                    let http_handler = http_handler.clone();
                    let websocket_handler = websocket_handler.clone();
                    let websocket_connector = websocket_connector.clone();
                    let mut conn_shutdown_rx = self.shutdown_rx.clone();
                    let active_connections = Arc::clone(&active_connections);
                    let connection_count = Arc::clone(&connection_count);

                    #[cfg(feature = "tor")]
                    let tor_client = tor_client.clone();
                    #[cfg(feature = "vpn")]
                    let vpn_client = vpn_client.clone();

                    connection_count.fetch_add(1, std::sync::atomic::Ordering::Relaxed);

                    tokio::spawn(async move {
                        let guard = ConnectionGuard {
                            count: Arc::clone(&connection_count),
                            notify: Arc::clone(&active_connections),
                        };

                        let conn = server.serve_connection_with_upgrades(
                            TokioIo::new(tcp),
                            service_fn(|req| {
                                InternalProxy {
                                    ca: Arc::clone(&ca),
                                    client: client.clone(),
                                    http_handler: http_handler.clone(),
                                    websocket_handler: websocket_handler.clone(),
                                    websocket_connector: websocket_connector.clone(),
                                    server: server.clone(),
                                    client_addr,
                                    request_timeout_ms,
                                    #[cfg(feature = "tor")]
                                    tor_client: tor_client.clone(),
                                    #[cfg(feature = "vpn")]
                                    vpn_client: vpn_client.clone(),
                                    custom_requester: custom_requester.clone(),
                                    compression
                                }
                                .proxy(req)
                            }),
                        );

                        let mut conn = std::pin::pin!(conn);

                        tokio::select! {
                            result = &mut conn => {
                                if let Err(err) = result {
                                    tracing::error!("Error serving connection: {}", err);
                                }
                            }
                            _ = conn_shutdown_rx.changed() => {
                                tracing::debug!("Connection received shutdown signal, draining...");
                                conn.as_mut().graceful_shutdown();
                                if let Err(err) = conn.await {
                                    tracing::error!("Error draining connection: {}", err);
                                }
                            }
                        }

                        drop(guard);
                    });
                }
                _ = shutdown_rx.changed() => {
                    tracing::info!("Shutdown signal received, stopping proxy accept loop");
                    break;
                }
            }
        }

        // Wait for all active connections to drain.
        // IMPORTANT: Register the waiter BEFORE loading the count to avoid a race
        // where the last connection finishes between the load and the notified() call.
        let notified = active_connections.notified();
        let remaining = connection_count.load(std::sync::atomic::Ordering::Relaxed);
        if remaining > 0 {
            tracing::info!("Waiting for {} active connection(s) to drain...", remaining);
            notified.await;
            tracing::info!("All connections drained");
        }

        // Mark as fully stopped now that all connections have drained
        self.state.store(
            ProxyState::Stopped.to_u8(),
            std::sync::atomic::Ordering::Relaxed,
        );
        self.stopped_notify.notify_waiters();

        Ok(())
    }

    /// Starts the proxy server and begins accepting connections.
    ///
    /// Binds to the configured address and enters the main accept loop.
    /// It will run until a shutdown signal is received via [`Proxy::stop`].
    ///
    /// # Returns
    ///
    /// * `Ok(())` - Proxy shut down gracefully
    /// * `Err(Error)` - If binding or serving failed
    ///
    /// # Panics
    ///
    /// Does not panic, but connection handler errors are logged
    /// and do not cause the proxy to stop.
    #[allow(dead_code)]
    #[allow(clippy::too_many_lines)]
    pub async fn start(self: Arc<Self>) -> Result<(), Error> {
        self.start_with_signal(None).await
    }
}

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

    #[test]
    fn test_proxy_state_debug() {
        assert_eq!(format!("{:?}", ProxyState::Uninitialized), "Uninitialized");
        assert_eq!(format!("{:?}", ProxyState::Initialized), "Initialized");
        assert_eq!(format!("{:?}", ProxyState::Started), "Started");
        assert_eq!(format!("{:?}", ProxyState::StopRequested), "StopRequested");
        assert_eq!(format!("{:?}", ProxyState::Stopped), "Stopped");
    }

    #[test]
    fn test_proxy_state_equality() {
        assert_eq!(ProxyState::Uninitialized, ProxyState::Uninitialized);
        assert_eq!(ProxyState::Initialized, ProxyState::Initialized);
        assert_eq!(ProxyState::Started, ProxyState::Started);
        assert_eq!(ProxyState::StopRequested, ProxyState::StopRequested);
        assert_eq!(ProxyState::Stopped, ProxyState::Stopped);

        assert_ne!(ProxyState::Uninitialized, ProxyState::Initialized);
        assert_ne!(ProxyState::Started, ProxyState::Stopped);
    }

    #[test]
    fn test_proxy_state_clone() {
        let state = ProxyState::Started;
        let cloned = state;
        assert_eq!(state, cloned);
    }

    #[test]
    fn test_proxy_state_roundtrip() {
        for state in [
            ProxyState::Uninitialized,
            ProxyState::Initialized,
            ProxyState::Started,
            ProxyState::StopRequested,
            ProxyState::Stopped,
        ] {
            assert_eq!(ProxyState::from_u8(state.to_u8()), state);
        }
    }
}