nixos-exporter/src/main.rs

use axum::{
    http::StatusCode,
    response::IntoResponse,
    routing::get,
    Router,
};

use std::collections::HashMap;
use std::net::SocketAddr;
use std::str::FromStr;

use nixos_exporter::nixos::NixStorePath;

#[derive(Clone)]
struct AppState {
    listen: String,
}

impl AppState{
    pub fn new() -> Self {
        Self {
            listen: String::from("[::]:9152"),
        }
    }

    pub fn is_valid(self) -> bool {
        return true;
    }
}

#[tokio::main]
async fn main() {
    let mut app_state = AppState::new();

    let mut args = std::env::args();
    let _name = args.next().unwrap();
    loop {
        let arg = if let Some(arg) = args.next() {
            arg
        } else {
            break;
        };

        match arg.as_str() {
            "--help" | "-h" => {
                println!("Prometheus exporter for NixOS systems");
                println!("Use --listen <addr:port> bind the web service.");
                println!("Output will be on /metrics endpoint. HTTP 500 if something broke while scraping.");
                std::process::exit(0);
            }
            "--listen" => {
                app_state.listen = args.next().unwrap();
            }
            unknown => {
                println!("unknown option: {}", unknown);
                std::process::exit(1)
            }
        }
    }

    if !app_state.clone().is_valid() {
        std::process::exit(1);
    }

    let app = Router::new();
    let app = app.route("/metrics", get(metrics));
    let app = app.with_state(app_state.clone());

    let addr = SocketAddr::from_str(&app_state.listen.clone()).unwrap();
    println!("listening on http://{}", addr);
    axum::Server::bind(&addr)
        .serve(app.into_make_service())
        .await
        .unwrap();
}

async fn metrics() -> Result<(StatusCode, impl IntoResponse), (StatusCode, impl IntoResponse)> {
    let nix_store_paths = HashMap::from([
        ("current_system", NixStorePath::from_str_symlink("/run/current-system")
            .map_err(|err| (StatusCode::INTERNAL_SERVER_ERROR, err))?),
        ("current_system_kernel", NixStorePath::from_str_symlink("/run/current-system/kernel")
            .map_err(|err| (StatusCode::INTERNAL_SERVER_ERROR, err))?),
        ("booted_system", NixStorePath::from_str_symlink("/run/booted-system")
            .map_err(|err| (StatusCode::INTERNAL_SERVER_ERROR, err))?),
        ("booted_system_kernel", NixStorePath::from_str_symlink("/run/booted-system/kernel")
            .map_err(|err| (StatusCode::INTERNAL_SERVER_ERROR, err))?),
    ]);

    let mut out = String::new();
    for (infix, nix_store_path) in nix_store_paths.iter() {
        out.push_str(nix_store_path.clone().to_prometheus_metric(infix.to_string())
            .map_err(|err| (StatusCode::INTERNAL_SERVER_ERROR, err))?.as_str());
    }

    out.push_str(format!("nixos_current_system_kernel_is_booted_system_kernel{{}} {}", (
            nix_store_paths.get("current_system_kernel").ok_or_else(|| (StatusCode::INTERNAL_SERVER_ERROR, String::from("")))?.hash
            == nix_store_paths.get("booted_system_kernel").ok_or_else(|| (StatusCode::INTERNAL_SERVER_ERROR, String::from("")))?.hash
        ) as i32).as_str()
    );

    return Ok((
        StatusCode::OK,
        out,
    ));
}