nixos-exporter/src/main.rs

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

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

#[derive(Clone, PartialEq)]
enum OperationMode {
    None,
    Exporter,
    Validator,
}

#[derive(Clone)]
struct AppState {
    listen: String,
    operationmode: OperationMode,
    prometheus_url: String,
    prometheus_query_tag_template: String,
    hydra_url: String,
    hydra_job_template: String,
}

impl AppState{
    pub fn new() -> Self {
        Self {
            listen: String::from("[::]:9152"),
            operationmode: OperationMode::None,
            prometheus_url: String::from(""),
            prometheus_query_tag_template: String::from("instance=\"{}\""),
            hydra_url: String::from(""),
            hydra_job_template: String::from("nixfiles/nixfiles/nixosConfigurations.{}"),
        }
    }

    pub fn is_valid(self) -> bool {
        let mut valid = true;
        if self.operationmode == OperationMode::None {
            println!("operationmode is not set");
            valid = false;
        }
        if self.prometheus_url == String::from("") && self.operationmode == OperationMode::Validator {
            println!("Prometheus url is not specified");
            valid = false;
        }
        if self.hydra_url == String::from("") && self.operationmode == OperationMode::Validator {
            println!("Hydra url is not specified");
            valid = false;
        }

        return valid;
    }
}

#[derive(Clone)]
struct NixStorePath {
    hash: String,
    name: String,
}

impl NixStorePath {
    pub fn from_str_symlink(path: &str) -> Result<Self, String> {
        Ok(Self::from_path_buf_symlink(PathBuf::from(path))?)
    }

    pub fn from_path_buf_symlink(path: PathBuf) -> Result<Self, String> {
        Ok(Self::from_path_buf(path.read_link().map_err(|err| err.to_string())?)?)
    }

    pub fn from_path_buf(path: PathBuf) -> Result<Self, String> {
        let store_path_name = path.iter().nth(3)
            .ok_or_else(|| String::from("Can't read store path name"))?
            .to_str()
            .ok_or_else(|| String::from("Failed converting store path name to string"))?
            .to_string();
        Ok(Self::from_store_path_name(store_path_name)?)
    }

    pub fn from_store_path_name(store_path_name: String) -> Result<Self, String> {
        let (hash, name) = store_path_name
            .split_once("-")
            .ok_or_else(|| String::from("Failed splitting store path name for hash and name"))?;
        Ok(Self {
            hash: hash.to_string(),
            name: name.to_string(),
        })
    }

    pub fn to_prometheus_metric(self, infix: String) -> Result<String, String> {
        return Ok(format!("nixos_{}_hash{{hash=\"{}\"}} 1\nnixos_{}_name{{name=\"{}\"}} 1\n", infix, self.hash, infix, self.name));
    }
}

#[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();
            }
            "--prometheus-url" => {
                app_state.prometheus_url = args.next().unwrap();
            }
            "--prometheus-query-tag-template" => {
                app_state.prometheus_query_tag_template = args.next().unwrap();
            }
            "--hydra-url" => {
                app_state.hydra_url = args.next().unwrap();
            }
            "--hydra-job-template" => {
                app_state.hydra_job_template = args.next().unwrap();
            }
            "exporter" => {
                app_state.operationmode = OperationMode::Exporter;
            }
            "validator" => {
                app_state.operationmode = OperationMode::Validator;
            }
            unknown => {
                println!("unknown option: {}", unknown);
                std::process::exit(1)
            }
        }
    }

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

    let mut app = Router::new();
    if app_state.operationmode == OperationMode::Exporter {
        println!("Running NixOS Exporter in Exporter mode");
        app = app.route("/metrics", get(metrics));
    } else if app_state.operationmode == OperationMode::Validator {
        println!("Running NixOS Exporter in Validator mode");
        app = app.route("/metrics", get(check));
    } else {
        println!("Run mode not specified, do {} --help", name);
        std::process::exit(1);
    };

    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,
    ));
}

async fn check(State(app_state): State<AppState>, Query(params): Query<HashMap<String, String>>) -> Result<(StatusCode, impl IntoResponse), (StatusCode, impl IntoResponse)> {
    let target = params.get("target")
        .ok_or_else(|| (StatusCode::NOT_FOUND, "specify target"))?;

    if target.contains("\"") {
        return Err((StatusCode::INTERNAL_SERVER_ERROR, "Invalid target name"));
    }

    let client = reqwest::Client::new();

    let prometheus_req = client.get(format!("{}/api/v1/query?query=nixos_current_system_hash{{{}}}", app_state.prometheus_url, app_state.prometheus_query_tag_template.clone().replace("{}", target)))
        .header("Accept", "application/json")
        .send().await
        .map_err(|_err| (StatusCode::INTERNAL_SERVER_ERROR, "Promehteus can't get reached"))?;

    if prometheus_req.status() != reqwest::StatusCode::OK {
       return Err((StatusCode::NOT_FOUND, "Target does not exist in Hydra"));
    }

    let prometheus_body = prometheus_req.json::<serde_json::Value>().await
        .map_err(|_err| (StatusCode::INTERNAL_SERVER_ERROR, "Invalid response from Hydra"))?;

    let current_system_hash = prometheus_body["data"]["result"][0]["metric"]["hash"].as_str()
        .ok_or_else(|| (StatusCode::INTERNAL_SERVER_ERROR, "No buildoutput found in Hydra"))?;

    let hydra_req = client.get(format!("{}/job/{}/latest", app_state.hydra_url, app_state.hydra_job_template.clone().replace("{}", target)))
        .header("Accept", "application/json")
        .send().await
        .map_err(|_err| (StatusCode::INTERNAL_SERVER_ERROR, "Hydra can't get reached"))?;

    if hydra_req.status() != reqwest::StatusCode::OK {
        return Err((StatusCode::NOT_FOUND, "Target does not exist in Hydra"));
    }

    let hydra_body = hydra_req.json::<serde_json::Value>().await
        .map_err(|_err| (StatusCode::INTERNAL_SERVER_ERROR, "Invalid response from Hydra"))?;

    let nix_store_path = hydra_body["buildoutputs"]["out"]["path"].as_str()
        .ok_or_else(|| (StatusCode::INTERNAL_SERVER_ERROR, "No buildoutput found in Hydra"))?;

    let hydra_system_hash = NixStorePath::from_path_buf(std::path::PathBuf::from(nix_store_path))
        .map_err(|_err| (StatusCode::INTERNAL_SERVER_ERROR, "Invalid store path returned by Hydra"))?
        .hash;

    let mut status = "0";

    if current_system_hash == hydra_system_hash {
        status = "1";
    }

    return Ok((
        StatusCode::OK,
        format!("nixos_current_system_is_sync{{}} {}\n", status)
    ));
}