vl-convert-server 2.0.0-rc2

//! Library crate for serving Vega-Lite and Vega conversion endpoints over HTTP.
//!
//! The crate exposes construction and serving primitives used by the
//! `vl-convert serve` CLI, and can also be embedded in another Rust binary.
//! Callers provide process lifecycle policy such as signal handling, readiness
//! output, and forced shutdown deadlines.

mod accept;
mod admin;
mod budget;
mod bundling;
mod config;
mod health;
mod json_fmt;
mod listen;
mod listener;
mod middleware;
mod reconfig;
mod router;
mod svg;
mod themes;
mod types;
mod util;
mod vega;
mod vegalite;

#[cfg(test)]
mod test_support;

pub use admin::admin_openapi;
pub use config::{init_tracing, BuiltApp, LogFormat, ServeConfig};
pub(crate) use config::{validate_serve_config, AdminConfig, ApiKey, AppState, RuntimeSnapshot};
pub use listen::ListenAddr;
#[cfg(unix)]
pub use listener::UdsCleanup;
pub use listener::{bind_listener, BoundListener, EndpointInfo};
pub use router::public_openapi;
pub(crate) use router::{build_middleware_stack, build_router};

use std::future::Future;
use std::net::SocketAddr;
use std::sync::Arc;
use std::time::Duration;
use tokio::task::JoinSet;
use tokio_util::sync::CancellationToken;
use tower_http::catch_panic::CatchPanicLayer;
use tower_http::request_id::{MakeRequestUuid, PropagateRequestIdLayer, SetRequestIdLayer};
use tower_http::trace::TraceLayer;
use vl_convert_rs::anyhow;
use vl_convert_rs::converter::VlcConfig;

/// Serve a [`BuiltApp`] on a bound main listener.
///
/// Runs the main listener, optional admin listener, and budget refill loop
/// on the current runtime. When `shutdown` resolves, all listeners receive
/// the shared shutdown token and this function returns after every spawned
/// task exits. Signal handling and drain-timeout escalation are the caller's
/// responsibility.
pub async fn serve(
    listener: BoundListener,
    built: BuiltApp,
    shutdown: impl Future<Output = ()> + Send + 'static,
) -> Result<(), anyhow::Error> {
    // The shutdown token is constructed in `build_app` and shared with
    // the reconfig coordinator so mid-reconfig cancellations abort cleanly.
    // `serve` fires the same token when its shutdown future resolves.
    let token = built.shutdown_token.clone();
    let mut tasks: JoinSet<()> = JoinSet::new();

    if let Some(tracker) = built.tracker {
        let cancel = token.clone();
        tasks.spawn(async move {
            let mut interval = tokio::time::interval(Duration::from_secs(1));
            loop {
                tokio::select! {
                    _ = interval.tick() => tracker.refill(),
                    _ = cancel.cancelled() => break,
                }
            }
        });
    }

    if let Some(admin) = built.admin {
        let cancel = token.clone();
        tasks.spawn(async move {
            log::info!("Admin API listening on {}", admin.addr);
            let admin_shutdown = async move { cancel.cancelled().await };
            match admin.listener {
                BoundListener::Tcp(l) => {
                    let _ = axum::serve(l, admin.router)
                        .with_graceful_shutdown(admin_shutdown)
                        .await;
                }
                #[cfg(unix)]
                BoundListener::Uds(l, _cleanup) => {
                    let _ = listener::serve_uds(l, admin.router, admin_shutdown).await;
                    // `_cleanup` drops here, unlinking the admin socket.
                }
            }
        });
    }

    let shutdown_token = token;
    let shutdown = async move {
        shutdown.await;
        shutdown_token.cancel();
    };

    match listener {
        BoundListener::Tcp(l) => {
            axum::serve(
                l,
                built
                    .router
                    .into_make_service_with_connect_info::<SocketAddr>(),
            )
            .with_graceful_shutdown(shutdown)
            .await?;
        }
        #[cfg(unix)]
        BoundListener::Uds(l, _cleanup) => {
            listener::serve_uds(l, built.router, shutdown).await?;
            // `_cleanup` drops here, unlinking the main socket.
        }
    }

    while tasks.join_next().await.is_some() {}
    Ok(())
}

/// Build a [`BuiltApp`] from converter and server configuration.
///
/// This validates the server config, initializes and warms the converter
/// worker pool, builds the main router, and binds the optional admin listener.
/// The main listener is bound separately with [`bind_listener`].
pub async fn build_app(
    config: VlcConfig,
    serve_config: &ServeConfig,
) -> Result<BuiltApp, anyhow::Error> {
    use arc_swap::ArcSwap;

    validate_serve_config(serve_config)?;

    // Build the shared shutdown token first. The reconfig coordinator needs
    // it so a mid-drain cancellation unwinds through the scope guard. The
    // same token is clone-stashed on `BuiltApp` for `serve()` to drive.
    let shutdown_token = CancellationToken::new();

    let reconfig_drain_timeout = Duration::from_secs(serve_config.reconfig_drain_timeout_secs);
    let coordinator =
        reconfig::ReconfigCoordinator::new(shutdown_token.clone(), reconfig_drain_timeout);

    let num_workers = config.num_workers;
    log::info!("Initializing converter with {num_workers} worker(s)...");
    let converter = vl_convert_rs::converter::VlConverter::with_config(config)?;
    converter.warm_up()?;
    let local_tz = match converter.get_local_tz().await {
        Ok(tz) => tz,
        Err(err) => {
            log::warn!("Failed to determine local timezone for /infoz: {err}");
            None
        }
    };
    log::info!("Workers initialized");

    // Seed the admin baseline and initial runtime snapshot from the normalized
    // config returned by `VlConverter::with_config` (resolved plugins, locale
    // aliases, and derived import domains). `DELETE /admin/config` resets to
    // this same startup baseline.
    let normalized = converter.config();
    let baseline = Arc::new(normalized.clone());

    let runtime = Arc::new(ArcSwap::from_pointee(RuntimeSnapshot {
        converter,
        config: Arc::new(normalized),
        generation: 0,
    }));

    let api_key = serve_config
        .api_key
        .as_ref()
        .map(|k| ApiKey::new(k.clone()));
    let readiness = Arc::new(health::ReadinessState::default());
    let state = Arc::new(AppState {
        runtime: runtime.clone(),
        api_key,
        opaque_errors: serve_config.opaque_errors,
        require_user_agent: serve_config.require_user_agent,
        readiness: readiness.clone(),
        local_tz,
        coordinator: coordinator.clone(),
    });

    let tracker = if serve_config.per_ip_budget_ms.is_some()
        || serve_config.global_budget_ms.is_some()
        || serve_config.admin.is_some()
    {
        Some(budget::BudgetTracker::new(
            serve_config.per_ip_budget_ms.unwrap_or(0),
            serve_config.global_budget_ms.unwrap_or(0),
            serve_config.budget_hold_ms,
        ))
    } else {
        None
    };

    let admin = if let (Some(admin_addr), Some(t)) = (&serve_config.admin, &tracker) {
        // `socket_mode` applies only to UDS listeners. Admin and main
        // share the same --socket-mode; a future --admin-socket-mode
        // flag could be added if asymmetric permissions are ever needed.
        let bound: BoundListener = bind_listener(admin_addr, serve_config.socket_mode).await?;
        let addr = bound.endpoint_label();
        // AdminState shares the runtime, coordinator, and readiness handles
        // with AppState so admin commits are visible to the main listener.
        let admin_state = Arc::new(admin::AdminState {
            runtime: runtime.clone(),
            baseline: baseline.clone(),
            coordinator: coordinator.clone(),
            readiness: readiness.clone(),
            admin_api_key: serve_config
                .admin_api_key
                .as_ref()
                .map(|k| ApiKey::new(k.clone())),
            tracker: t.clone(),
            opaque_errors: serve_config.opaque_errors,
        });
        let admin_router = admin::admin_router(admin_state)
            .layer(PropagateRequestIdLayer::x_request_id())
            .layer(TraceLayer::new_for_http())
            .layer(SetRequestIdLayer::x_request_id(MakeRequestUuid))
            .layer(CatchPanicLayer::new());
        Some(AdminConfig {
            listener: bound,
            addr,
            router: admin_router,
        })
    } else {
        None
    };

    let router = build_router(
        state,
        tracker.clone(),
        serve_config.opaque_errors,
        serve_config.trust_proxy,
        serve_config.google_font_cache_miss_penalty_ms,
    );
    let app = build_middleware_stack(router, serve_config);

    Ok(BuiltApp {
        router: app,
        runtime,
        shutdown_token,
        tracker,
        admin,
    })
}

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

    #[tokio::test]
    async fn test_build_app_rejects_non_positive_budget_hold_ms() {
        let config = VlcConfig::default();
        let mut serve_config = default_serve_config();
        serve_config.budget_hold_ms = 0;

        let err = build_app(config, &serve_config).await.err().unwrap();
        assert!(
            err.to_string().contains("budget_hold_ms must be positive"),
            "unexpected error: {err}"
        );
    }

    #[tokio::test]
    async fn test_build_app_rejects_empty_main_api_key() {
        for key in ["", "   ", "\t"] {
            let mut serve_config = default_serve_config();
            serve_config.api_key = Some(key.to_string());

            let err = build_app(VlcConfig::default(), &serve_config)
                .await
                .err()
                .unwrap();
            assert!(
                err.to_string().contains("api_key"),
                "key {key:?} must trigger validator: {err}"
            );
        }
    }

    #[tokio::test]
    async fn test_build_app_rejects_empty_admin_api_key_on_safe_admin_bind() {
        for key in ["", "   ", "\t"] {
            let mut serve_config = default_serve_config();
            serve_config.admin = Some(crate::ListenAddr::loopback_tcp(0));
            serve_config.admin_api_key = Some(key.to_string());

            let err = build_app(VlcConfig::default(), &serve_config)
                .await
                .err()
                .unwrap();
            assert!(
                err.to_string().contains("admin_api_key"),
                "loopback admin key {key:?} must trigger validator: {err}"
            );
        }

        #[cfg(unix)]
        for key in ["", "   ", "\t"] {
            let mut serve_config = default_serve_config();
            serve_config.admin = Some(crate::ListenAddr::Uds {
                path: "/tmp/vl-convert-admin-test.sock".into(),
            });
            serve_config.admin_api_key = Some(key.to_string());

            let err = build_app(VlcConfig::default(), &serve_config)
                .await
                .err()
                .unwrap();
            assert!(
                err.to_string().contains("admin_api_key"),
                "UDS admin key {key:?} must trigger validator: {err}"
            );
        }
    }

    #[tokio::test]
    async fn test_build_app_rejects_non_loopback_admin_without_usable_key() {
        // None triggers the non-loopback admin guard; whitespace-only
        // keys are rejected earlier as invalid configured secrets.
        for key in [None, Some(""), Some("   "), Some("\t")] {
            let mut serve_config = default_serve_config();
            serve_config.admin = Some(crate::ListenAddr::Tcp {
                host: "0.0.0.0".to_string(),
                port: 0,
            });
            serve_config.admin_api_key = key.map(String::from);

            let err = build_app(VlcConfig::default(), &serve_config)
                .await
                .err()
                .unwrap();
            assert!(
                err.to_string().contains("admin_api_key"),
                "key {key:?} must trigger validator: {err}"
            );
        }
    }

    #[tokio::test]
    async fn test_build_app_accepts_non_loopback_admin_with_key() {
        // The validation pre-bind should pass; an actual bind may still
        // fail at runtime depending on the test environment, but the
        // error must come from `bind_listener`, not from the validator.
        let mut serve_config = default_serve_config();
        serve_config.admin = Some(crate::ListenAddr::Tcp {
            host: "0.0.0.0".to_string(),
            port: 0,
        });
        serve_config.admin_api_key = Some("supersecret".to_string());
        // Need budget tracker so the admin listener is wired up at all.
        serve_config.global_budget_ms = Some(1_000);

        if let Err(err) = build_app(VlcConfig::default(), &serve_config).await {
            assert!(
                !err.to_string().contains("admin_api_key"),
                "validator rejected admin with key: {err}"
            );
        }
    }
}