all-smi 0.26.2

// Copyright 2025 Lablup Inc. and Jeongkyu Shin
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use axum::http::{HeaderName, HeaderValue, Method, header};
use axum::{Router, routing::get};
use std::time::Duration;
use tokio::net::TcpListener;
use tokio::sync::RwLock;
use tower_http::cors::{AllowOrigin, CorsLayer};
use tower_http::trace::TraceLayer;
use tracing_subscriber::{layer::SubscriberExt, util::SubscriberInitExt};

use crate::api::shutdown::{mark_serving, shutdown_signal};

#[cfg(unix)]
use std::path::PathBuf;
#[cfg(unix)]
use tokio::net::UnixListener;

use crate::api::FrameBus;
use crate::api::collection_loop::run_collection_loop;
use crate::api::handlers::events::events_handler;
use crate::api::handlers::snapshot::snapshot_handler;
use crate::api::handlers::{SharedState, metrics_handler, ready_handler};
use crate::api::server_state::ApiState;
use crate::app_state::AppState;
use crate::cli::ApiArgs;
use crate::common::config_file::Settings;

/// Get the default Unix domain socket path for the current platform.
/// - Linux: /var/run/all-smi.sock (fallback to /tmp/all-smi.sock if no permission)
/// - macOS: /tmp/all-smi.sock
#[cfg(unix)]
fn get_default_socket_path() -> PathBuf {
    #[cfg(target_os = "linux")]
    {
        let var_run_path = PathBuf::from("/var/run/all-smi.sock");
        // Check if we can write to /var/run
        if let Ok(metadata) = std::fs::metadata("/var/run")
            && metadata.is_dir()
        {
            // Try to create a test file to check write permission
            let test_path = PathBuf::from("/var/run/.all-smi-test");
            if std::fs::write(&test_path, b"").is_ok() {
                let _ = std::fs::remove_file(&test_path);
                return var_run_path;
            }
        }
        // Fallback to /tmp
        PathBuf::from("/tmp/all-smi.sock")
    }

    #[cfg(target_os = "macos")]
    {
        PathBuf::from("/tmp/all-smi.sock")
    }

    #[cfg(not(any(target_os = "linux", target_os = "macos")))]
    {
        PathBuf::from("/tmp/all-smi.sock")
    }
}

/// Remove stale socket file if it exists.
/// This is necessary because Unix sockets leave files on disk that prevent rebinding.
/// Uses atomic remove to avoid TOCTOU race conditions.
#[cfg(unix)]
fn remove_stale_socket(path: &std::path::Path) -> std::io::Result<()> {
    match std::fs::remove_file(path) {
        Ok(()) => {
            tracing::info!("Removed stale socket file: {}", path.display());
            Ok(())
        }
        Err(e) if e.kind() == std::io::ErrorKind::NotFound => {
            // File doesn't exist, that's fine
            Ok(())
        }
        Err(e) => Err(e),
    }
}

/// Set restrictive permissions (0o600) on the socket file.
/// This ensures only the owner can connect to the socket.
#[cfg(unix)]
fn set_socket_permissions(path: &std::path::Path) -> std::io::Result<()> {
    use std::os::unix::fs::PermissionsExt;
    let permissions = std::fs::Permissions::from_mode(0o600);
    std::fs::set_permissions(path, permissions)
}

/// Run the API server with TCP and optionally Unix Domain Socket listeners.
pub async fn run_api_mode(args: &ApiArgs, settings: &Settings) {
    // `try_init` rather than `init`: a host may have installed a
    // subscriber already, and `init` panics in that case. The Windows
    // service host does exactly that, because stdout is void under the
    // Service Control Manager and its logs have to reach a file under
    // %PROGRAMDATA% instead (issue #311). On every other entry point
    // nothing has registered a subscriber yet, so this still installs
    // the stdout layer as before.
    if tracing_subscriber::registry()
        .with(
            tracing_subscriber::EnvFilter::try_from_default_env()
                .unwrap_or_else(|_| "all_smi=debug,tower_http=debug".into()),
        )
        .with(tracing_subscriber::fmt::layer())
        .try_init()
        .is_err()
    {
        tracing::debug!("a tracing subscriber is already installed; keeping the host's");
    }

    println!("Starting API mode...");
    // Build the state with the merged energy config so the integrator's
    // `gap_interpolate_seconds` honours the TOML value. Without this
    // constructor, later assignments to `energy_config` would not
    // reconfigure the already-constructed `PowerIntegrator`.
    let mut initial_state = AppState::with_energy_config(&settings.energy);
    // Propagate `[display]` for symmetry with the TUI modes; the
    // HTTP/Prometheus surface currently ignores it but future
    // HTML/health-page renderers can read from a single location.
    initial_state.display_config = settings.display.clone();
    // Replay any persisted energy WAL so Prometheus counters stay
    // monotonic across restarts (issue #191). Failures are logged
    // but do not block startup — the integrator simply begins at
    // zero on this host.
    //
    // Path resolution flows through `resolve_wal_path` (issue #229) so
    // the same precedence rules — operator override → platform cache
    // dir → in-memory only — apply both here and at the flush task
    // below.
    if initial_state.energy_config.wal_enabled {
        match crate::metrics::energy_wal::resolve_wal_path(
            initial_state.energy_config.wal_path.as_deref(),
        ) {
            Some(wal_path) => {
                let path_display = wal_path.display().to_string();
                match crate::metrics::energy_wal::replay_from_path(
                    &wal_path,
                    initial_state.energy.integrator_mut(),
                ) {
                    Ok(index) => {
                        if !index.is_empty() {
                            tracing::info!(
                                "energy WAL: replayed {} records from {path_display}",
                                index.len()
                            );
                        }
                        initial_state.energy_wal_replay = index;
                    }
                    Err(e) => {
                        tracing::warn!("energy WAL: replay from {path_display} failed: {e}");
                    }
                }
            }
            None => {
                tracing::warn!(
                    "energy WAL: no cache directory available in environment; \
                     counters are in-memory only"
                );
            }
        }
    }
    let state = SharedState::new(RwLock::new(initial_state));
    let state_clone = state.clone();
    // `args.processes` is now `Option<bool>` so the CLI can force
    // `--processes=false` against a config file that sets it to true.
    // `main.rs` always resolves the option before reaching here; the
    // `unwrap_or(false)` is a defensive fallback so a direct
    // library-level caller cannot crash by passing `None` through.
    let processes = args.processes.unwrap_or(false);
    // args.interval was resolved against settings in main.rs; fall back
    // defensively to 3 (compiled default) when the caller somehow
    // passed `None`.
    let interval = args.interval.unwrap_or(3);

    // Spawn the WAL flush task if enabled. The returned handle owns a
    // oneshot sender used by the Ctrl+C / SIGTERM path so the task can
    // perform a final `flush_and_fsync` before the process exits
    // (issue #191).
    let wal_flush_handle = {
        let state = state.clone();
        let state_read = state.read().await;
        let cfg = state_read.energy_config.clone();
        drop(state_read);
        if cfg.wal_enabled {
            match crate::metrics::energy_wal::resolve_wal_path(cfg.wal_path.as_deref()) {
                Some(path) => Some(crate::metrics::energy_wal::spawn_wal_flush_task(
                    state,
                    path,
                    crate::metrics::energy_wal::DEFAULT_FLUSH_INTERVAL,
                )),
                None => {
                    tracing::warn!(
                        "energy WAL: no cache directory available in environment; \
                         skipping flush task (counters remain in-memory only)"
                    );
                    None
                }
            }
        } else {
            None
        }
    };

    // Publish one frame per collection cycle onto a shared broadcast
    // bus (issue #193). The SSE `/events` and one-shot `/snapshot`
    // handlers subscribe to this bus so they never need to run their
    // own reader loop — see `api::frame_bus` and `api::collection_loop`.
    let bus = FrameBus::new(Duration::from_secs(interval));

    // Spawn the background collection task. It owns the reader factories
    // and is the only place that calls them on the live server.
    tokio::spawn(run_collection_loop(
        state_clone.clone(),
        bus.clone(),
        interval,
        processes,
    ));

    // Compose the router state so each handler extracts only the
    // sub-state it needs (see `server_state::ApiState`).
    let api_state = ApiState::new(state, bus);

    // Create the router with shared state.
    //
    // `/-/ready` (issue #324) is the readiness gate: 503 until the
    // collection loop spawned above has written its first cycle into
    // `AppState`, 200 after. `/metrics` never joins that gate, it answers
    // 200 throughout and reports the same transition in-band through
    // `all_smi_up`, so existing scrapers are unaffected.
    let app = Router::new()
        .route("/metrics", get(metrics_handler))
        .route("/-/ready", get(ready_handler))
        .route("/events", get(events_handler))
        .route("/snapshot", get(snapshot_handler))
        .with_state(api_state)
        .layer(build_cors_layer())
        .layer(TraceLayer::new_for_http());

    // Determine which listeners to start
    #[cfg(unix)]
    {
        let socket_path = args.socket.as_ref().map(|s| {
            if s.is_empty() {
                get_default_socket_path()
            } else {
                PathBuf::from(s)
            }
        });

        let port = args.port.unwrap_or(9090);
        match (port, socket_path) {
            // Both TCP and UDS (port > 0 with socket)
            (1..=u16::MAX, Some(path)) => {
                run_dual_listeners(app, port, path).await;
            }
            // UDS only (port == 0 with socket)
            (0, Some(path)) => {
                run_unix_listener(app, path).await;
            }
            // TCP only (port > 0, no socket)
            (1..=u16::MAX, None) => {
                run_tcp_listener(app, port).await;
            }
            // No listeners - error (port == 0, no socket)
            (0, None) => {
                tracing::error!(
                    "No listeners configured. Use --port or --socket to specify a listener."
                );
                eprintln!(
                    "Error: No listeners configured. Use --port or --socket to specify a listener."
                );
            }
        }
    }

    #[cfg(not(unix))]
    {
        run_tcp_listener(app, args.port.unwrap_or(9090)).await;
    }

    // Signal the WAL flush task to perform a final flush and fsync
    // before we exit, so the last batch of pending Joule deltas is not
    // lost (issue #191). Has to run AFTER the listeners return because
    // that is the moment Ctrl+C / SIGTERM has propagated through axum.
    if let Some(handle) = wal_flush_handle {
        handle.shutdown().await;
    }
}

/// Run only the TCP listener
async fn run_tcp_listener(app: Router, port: u16) {
    let listener = match TcpListener::bind(&format!("0.0.0.0:{port}")).await {
        Ok(l) => l,
        Err(e) => {
            tracing::error!("Failed to bind TCP listener on port {port}: {e}");
            eprintln!("Error: Failed to bind TCP listener on port {port}: {e}");
            return;
        }
    };
    tracing::info!(
        "API server listening on {}",
        listener
            .local_addr()
            .unwrap_or_else(|_| "unknown".parse().unwrap())
    );
    mark_serving();
    if let Err(e) = axum::serve(listener, app)
        .with_graceful_shutdown(shutdown_signal())
        .await
    {
        tracing::error!("TCP server error: {e}");
    }
}

/// Run only the Unix Domain Socket listener
#[cfg(unix)]
async fn run_unix_listener(app: Router, path: PathBuf) {
    // Remove stale socket file if it exists
    if let Err(e) = remove_stale_socket(&path) {
        tracing::warn!("Failed to remove stale socket file: {e}");
    }

    // Create parent directory if it doesn't exist
    if let Some(parent) = path.parent()
        && !parent.exists()
        && let Err(e) = std::fs::create_dir_all(parent)
    {
        tracing::error!(
            "Failed to create socket directory {}: {e}",
            parent.display()
        );
        eprintln!(
            "Error: Failed to create socket directory {}: {e}",
            parent.display()
        );
        return;
    }

    let listener = match UnixListener::bind(&path) {
        Ok(l) => l,
        Err(e) => {
            tracing::error!("Failed to bind Unix socket at {}: {e}", path.display());
            eprintln!(
                "Error: Failed to bind Unix socket at {}: {e}",
                path.display()
            );
            return;
        }
    };

    // Set restrictive permissions (0o600) on the socket file
    if let Err(e) = set_socket_permissions(&path) {
        tracing::warn!("Failed to set socket permissions: {e}");
    }

    tracing::info!("API server listening on Unix socket: {}", path.display());
    mark_serving();

    // Serve the application with graceful shutdown so the caller can
    // run post-serve cleanup (WAL flush, socket cleanup) once we see a
    // SIGTERM / Ctrl+C.
    if let Err(e) = axum::serve(listener, app)
        .with_graceful_shutdown(shutdown_signal())
        .await
    {
        tracing::error!("Unix socket server error: {e}");
    }

    cleanup_socket(&path);
}

/// Run both TCP and Unix Domain Socket listeners simultaneously
#[cfg(unix)]
async fn run_dual_listeners(app: Router, port: u16, socket_path: PathBuf) {
    // Remove stale socket file if it exists
    if let Err(e) = remove_stale_socket(&socket_path) {
        tracing::warn!("Failed to remove stale socket file: {e}");
    }

    // Create parent directory if it doesn't exist
    if let Some(parent) = socket_path.parent()
        && !parent.exists()
        && let Err(e) = std::fs::create_dir_all(parent)
    {
        tracing::error!(
            "Failed to create socket directory {}: {e}",
            parent.display()
        );
        eprintln!(
            "Error: Failed to create socket directory {}: {e}",
            parent.display()
        );
        return;
    }

    // Create TCP listener
    let tcp_listener = match TcpListener::bind(&format!("0.0.0.0:{port}")).await {
        Ok(l) => l,
        Err(e) => {
            tracing::error!("Failed to bind TCP listener on port {port}: {e}");
            eprintln!("Error: Failed to bind TCP listener on port {port}: {e}");
            return;
        }
    };

    // Create Unix listener
    let unix_listener = match UnixListener::bind(&socket_path) {
        Ok(l) => l,
        Err(e) => {
            tracing::error!(
                "Failed to bind Unix socket at {}: {e}",
                socket_path.display()
            );
            eprintln!(
                "Error: Failed to bind Unix socket at {}: {e}",
                socket_path.display()
            );
            return;
        }
    };

    // Set restrictive permissions (0o600) on the socket file
    if let Err(e) = set_socket_permissions(&socket_path) {
        tracing::warn!("Failed to set socket permissions: {e}");
    }

    tracing::info!(
        "API server listening on TCP {} and Unix socket {}",
        tcp_listener
            .local_addr()
            .unwrap_or_else(|_| "unknown".parse().unwrap()),
        socket_path.display()
    );
    mark_serving();

    // Clone the app for the second server
    let app_clone = app.clone();

    // Run both servers concurrently; each installs its own graceful
    // shutdown listener so the select returns on SIGTERM / Ctrl+C and
    // the caller can run post-serve cleanup.
    tokio::select! {
        result = axum::serve(tcp_listener, app)
            .with_graceful_shutdown(shutdown_signal()) => {
            if let Err(e) = result {
                tracing::error!("TCP server error: {e}");
            }
        }
        result = axum::serve(unix_listener, app_clone)
            .with_graceful_shutdown(shutdown_signal()) => {
            if let Err(e) = result {
                tracing::error!("Unix socket server error: {e}");
            }
        }
    }

    cleanup_socket(&socket_path);
}

/// Clean up the Unix domain socket file.
/// Uses atomic remove to avoid TOCTOU race conditions.
#[cfg(unix)]
fn cleanup_socket(path: &std::path::Path) {
    match std::fs::remove_file(path) {
        Ok(()) => {
            tracing::info!("Cleaned up socket file: {}", path.display());
        }
        Err(e) if e.kind() == std::io::ErrorKind::NotFound => {
            // File already removed, that's fine
        }
        Err(e) => {
            tracing::warn!("Failed to remove socket file on shutdown: {e}");
        }
    }
}

/// Build the CORS layer for the API router.
///
/// Security posture:
///
/// * **No cross-origin access by default.** The previous wildcard
///   (`Allow-Origin: *`, `Allow-Methods: *`, `Allow-Headers: *`) let any
///   origin read `/metrics`, `/snapshot`, and the `/events` SSE stream
///   from a browsing context. That telemetry includes process command
///   lines, usernames, and power data — sensitive enough that a
///   malicious page loaded by any authenticated operator could
///   exfiltrate it via a cross-origin `fetch`. We therefore default
///   to the strict axum CORS posture: no extra headers, no allowed
///   origins, leaving same-origin requests unaffected.
/// * **Opt-in allowlist.** When the operator sets
///   `ALL_SMI_API_CORS_ALLOWED_ORIGINS` to a comma-separated list of
///   origins, those origins (and only those) are reflected in the
///   `Access-Control-Allow-Origin` header. `GET` + `Accept:
///   text/event-stream` remains sufficient to subscribe to `/events`
///   from a permitted origin.
/// * **Wildcard escape hatch.** Setting the env var to exactly `*`
///   restores the previous permissive behaviour for operators who
///   understand the exposure (public read-only dashboards on a
///   network-isolated host). A warning is logged so the risk is
///   discoverable in the operator logs.
fn build_cors_layer() -> CorsLayer {
    let raw = std::env::var("ALL_SMI_API_CORS_ALLOWED_ORIGINS").ok();
    let trimmed = raw.as_deref().map(str::trim).unwrap_or("");

    // Default: no CORS — browsers refuse cross-origin reads of our
    // telemetry, same-origin and non-browser clients (curl, Prometheus,
    // Tauri with file:// escape) are unaffected.
    if trimmed.is_empty() {
        return CorsLayer::new()
            .allow_methods([Method::GET, Method::OPTIONS])
            .allow_headers([
                header::ACCEPT,
                header::CONTENT_TYPE,
                HeaderName::from_static("last-event-id"),
            ]);
    }

    // Explicit opt-in to a wildcard. Noisy warning so operators running
    // this in a hostile environment see the audit trail in their logs.
    if trimmed == "*" {
        tracing::warn!(
            "ALL_SMI_API_CORS_ALLOWED_ORIGINS=* selected; every origin              may read /metrics, /snapshot, and /events. This exposes              GPU telemetry, process command lines, and usernames              cross-origin. Prefer an explicit origin list."
        );
        return CorsLayer::new()
            .allow_origin(AllowOrigin::any())
            .allow_methods([Method::GET, Method::OPTIONS])
            .allow_headers([
                header::ACCEPT,
                header::CONTENT_TYPE,
                HeaderName::from_static("last-event-id"),
            ]);
    }

    // Parse a comma-separated allowlist. Entries that cannot be parsed
    // as a `HeaderValue` are dropped with a warning rather than failing
    // startup — misconfiguration should not prevent the server from
    // booting.
    let origins: Vec<HeaderValue> = trimmed
        .split(',')
        .map(str::trim)
        .filter(|s| !s.is_empty())
        .filter_map(|o| match HeaderValue::from_str(o) {
            Ok(v) => Some(v),
            Err(e) => {
                tracing::warn!(origin = o, error = %e, "ignoring invalid CORS origin");
                None
            }
        })
        .collect();

    if origins.is_empty() {
        tracing::warn!(
            "ALL_SMI_API_CORS_ALLOWED_ORIGINS was set but contained no              valid origins; falling back to no-CORS default"
        );
        return CorsLayer::new()
            .allow_methods([Method::GET, Method::OPTIONS])
            .allow_headers([
                header::ACCEPT,
                header::CONTENT_TYPE,
                HeaderName::from_static("last-event-id"),
            ]);
    }

    tracing::info!(
        allowed_origins = origins.len(),
        "CORS: allowlist configured from ALL_SMI_API_CORS_ALLOWED_ORIGINS"
    );
    CorsLayer::new()
        .allow_origin(AllowOrigin::list(origins))
        .allow_methods([Method::GET, Method::OPTIONS])
        .allow_headers([
            header::ACCEPT,
            header::CONTENT_TYPE,
            HeaderName::from_static("last-event-id"),
        ])
}

#[cfg(test)]
mod cors_tests {
    //! CORS policy tests for `build_cors_layer`.
    //!
    //! These tests use `std::env::set_var` / `remove_var`. They are
    //! guarded by `#[serial_test]`... actually we lack that dependency,
    //! so each test reads and restores the env var manually via a
    //! scope guard.

    use super::*;

    /// RAII guard that saves the current `ALL_SMI_API_CORS_ALLOWED_ORIGINS`
    /// value on creation and restores it on drop, so parallel-running
    /// tests cannot leak env state to each other's assertions.
    struct EnvGuard {
        key: &'static str,
        original: Option<String>,
    }
    impl EnvGuard {
        fn set(key: &'static str, value: &str) -> Self {
            let original = std::env::var(key).ok();
            // SAFETY: the env var is set and restored inside a unit test
            // that runs serially within the test binary (the env guard
            // pattern preserves that invariant across the test body).
            unsafe { std::env::set_var(key, value) };
            Self { key, original }
        }
        fn unset(key: &'static str) -> Self {
            let original = std::env::var(key).ok();
            // SAFETY: same reasoning as `set`.
            unsafe { std::env::remove_var(key) };
            Self { key, original }
        }
    }
    impl Drop for EnvGuard {
        fn drop(&mut self) {
            match self.original.take() {
                // SAFETY: restoring env state the test captured.
                Some(v) => unsafe { std::env::set_var(self.key, v) },
                // SAFETY: restoring env state the test captured.
                None => unsafe { std::env::remove_var(self.key) },
            }
        }
    }

    #[test]
    fn default_builds_without_panic() {
        let _g = EnvGuard::unset("ALL_SMI_API_CORS_ALLOWED_ORIGINS");
        // We don't have a direct accessor into CorsLayer; the contract
        // this test pins is "constructing the default posture must
        // succeed without panic or env reads escaping the helper".
        let _layer = build_cors_layer();
    }

    #[test]
    fn wildcard_allowed_when_explicitly_requested() {
        let _g = EnvGuard::set("ALL_SMI_API_CORS_ALLOWED_ORIGINS", "*");
        let _layer = build_cors_layer();
    }

    #[test]
    fn invalid_origins_drop_to_default() {
        let _g = EnvGuard::set("ALL_SMI_API_CORS_ALLOWED_ORIGINS", "\n\tnot a url");
        let _layer = build_cors_layer();
    }
}