socktop/socktop_agent/src/main.rs

//! socktop agent entrypoint: sets up sysinfo handles and serves a WebSocket endpoint at /ws.

mod gpu;
mod metrics;
mod proto;
// sampler module removed (metrics now purely request-driven)
mod state;
mod types;
mod ws;

use axum::{Router, http::StatusCode, routing::get};
use std::net::SocketAddr;
use std::str::FromStr;

mod tls;

use state::AppState;

fn arg_flag(name: &str) -> bool {
    std::env::args().any(|a| a == name)
}
fn arg_value(name: &str) -> Option<String> {
    let mut it = std::env::args();
    while let Some(a) = it.next() {
        if a == name {
            return it.next();
        }
    }
    None
}

#[tokio::main]
async fn main() -> anyhow::Result<()> {
    tracing_subscriber::fmt::init();

    // Version flag (print and exit). Keep before heavy initialization.
    if arg_flag("--version") || arg_flag("-V") {
        println!("socktop_agent {}", env!("CARGO_PKG_VERSION"));
        return Ok(());
    }

    let state = AppState::new();

    // No background samplers: metrics collected on-demand per websocket request.

    // Web app: route /ws to the websocket handler
    async fn healthz() -> StatusCode {
        println!("/healthz request");
        StatusCode::OK
    }
    let app = Router::new()
        .route("/ws", get(ws::ws_handler))
        .route("/healthz", get(healthz))
        .with_state(state.clone());

    let enable_ssl =
        arg_flag("--enableSSL") || std::env::var("SOCKTOP_ENABLE_SSL").ok().as_deref() == Some("1");
    if enable_ssl {
        // Port can be overridden by --port or SOCKTOP_PORT; default to 8443 when SSL
        let port = arg_value("--port")
            .or_else(|| arg_value("-p"))
            .or_else(|| std::env::var("SOCKTOP_PORT").ok())
            .and_then(|s| s.parse::<u16>().ok())
            .unwrap_or(8443);

        let (cert_path, key_path) = tls::ensure_self_signed_cert()?;
        let cfg = axum_server::tls_rustls::RustlsConfig::from_pem_file(cert_path, key_path).await?;

        let addr = SocketAddr::from_str(&format!("0.0.0.0:{port}"))?;
        println!("socktop_agent: TLS enabled. Listening on wss://{addr}/ws");
        axum_server::bind_rustls(addr, cfg)
            .serve(app.into_make_service())
            .await?;
        return Ok(());
    }

    // Non-TLS HTTP/WS path
    let port = arg_value("--port")
        .or_else(|| arg_value("-p"))
        .or_else(|| std::env::var("SOCKTOP_PORT").ok())
        .and_then(|s| s.parse::<u16>().ok())
        .unwrap_or(3000);
    let addr = SocketAddr::from(([0, 0, 0, 0], port));
    println!("socktop_agent: Listening on ws://{addr}/ws");
    axum_server::bind(addr)
        .serve(app.into_make_service())
        .await?;
    Ok(())
}

// Unit tests for CLI parsing moved to `tests/port_parse.rs`.