crank/crates/crank-observability/src/prometheus.rs

use std::{env, net::SocketAddr};

use axum::{
    Router,
    extract::{Request, State},
    http::{
        HeaderMap, StatusCode,
        header::{self, HeaderValue},
    },
    middleware::{self, Next},
    response::{IntoResponse, Response},
    routing::get,
};
use metrics_exporter_prometheus::{PrometheusBuilder, PrometheusHandle, PrometheusRecorder};
use sha2::{Digest, Sha256};
use subtle::ConstantTimeEq;
use thiserror::Error;
use tokio::net::TcpListener;

use crate::{DURATION_BUCKETS_SECONDS, ServiceIdentity};

const METRICS_ENABLED_ENV: &str = "CRANK_METRICS_ENABLED";
const METRICS_TOKEN_ENV: &str = "CRANK_METRICS_BEARER_TOKEN";
const PROMETHEUS_CONTENT_TYPE: &str = "text/plain; version=0.0.4; charset=utf-8";

#[derive(Clone)]
pub struct MetricsConfig {
    enabled: bool,
    bind_addr: SocketAddr,
    token_digest: Option<[u8; 32]>,
}

impl std::fmt::Debug for MetricsConfig {
    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        formatter
            .debug_struct("MetricsConfig")
            .field("enabled", &self.enabled)
            .field("bind_addr", &self.bind_addr)
            .field("authentication_configured", &self.token_digest.is_some())
            .finish()
    }
}

impl MetricsConfig {
    pub fn new(
        enabled: bool,
        bind_addr: SocketAddr,
        bearer_token: Option<String>,
    ) -> Result<Self, MetricsConfigError> {
        let token_digest = bearer_token
            .filter(|token| !token.is_empty())
            .map(|token| token_digest(token.as_bytes()));

        if enabled && !bind_addr.ip().is_loopback() && token_digest.is_none() {
            return Err(MetricsConfigError::MissingTokenForExternalBind);
        }

        Ok(Self {
            enabled,
            bind_addr,
            token_digest,
        })
    }

    pub fn from_env(
        bind_env: &'static str,
        default_bind: SocketAddr,
    ) -> Result<Self, MetricsConfigError> {
        let enabled = parse_enabled(env::var(METRICS_ENABLED_ENV))?;
        let bind_addr = match env::var(bind_env) {
            Ok(raw) => raw
                .parse()
                .map_err(|_| MetricsConfigError::InvalidBindAddress { field: bind_env })?,
            Err(env::VarError::NotPresent) => default_bind,
            Err(env::VarError::NotUnicode(_)) => {
                return Err(MetricsConfigError::InvalidEnvironmentEncoding { field: bind_env });
            }
        };
        let bearer_token = match env::var(METRICS_TOKEN_ENV) {
            Ok(token) => Some(token),
            Err(env::VarError::NotPresent) => None,
            Err(env::VarError::NotUnicode(_)) => {
                return Err(MetricsConfigError::InvalidEnvironmentEncoding {
                    field: METRICS_TOKEN_ENV,
                });
            }
        };

        Self::new(enabled, bind_addr, bearer_token)
    }

    pub fn enabled(&self) -> bool {
        self.enabled
    }

    pub fn bind_addr(&self) -> SocketAddr {
        self.bind_addr
    }

    pub fn requires_authentication(&self) -> bool {
        !self.bind_addr.ip().is_loopback()
    }
}

#[derive(Debug, Error)]
pub enum MetricsConfigError {
    #[error("metrics environment variable is not valid UTF-8: {field}")]
    InvalidEnvironmentEncoding { field: &'static str },
    #[error("metrics bind address is invalid: {field}")]
    InvalidBindAddress { field: &'static str },
    #[error("metrics enabled flag must be one of true, false, 1, 0")]
    InvalidEnabledFlag,
    #[error("external metrics bind requires a bearer token")]
    MissingTokenForExternalBind,
}

#[derive(Clone)]
struct MetricsState {
    handle: PrometheusHandle,
    token_digest: Option<[u8; 32]>,
    requires_authentication: bool,
}

pub struct MetricsSurface {
    config: MetricsConfig,
    state: MetricsState,
    _recorder: Option<PrometheusRecorder>,
}

impl MetricsSurface {
    pub(crate) fn new(config: MetricsConfig, handle: PrometheusHandle) -> Self {
        Self {
            state: MetricsState {
                handle,
                token_digest: config.token_digest,
                requires_authentication: config.requires_authentication(),
            },
            config,
            _recorder: None,
        }
    }

    pub fn for_test(
        config: MetricsConfig,
        identity: ServiceIdentity,
    ) -> Result<Self, MetricsSurfaceError> {
        let recorder = prometheus_builder(&identity)?.build_recorder();
        let handle = recorder.handle();
        let mut surface = Self::new(config, handle);
        surface._recorder = Some(recorder);
        Ok(surface)
    }

    pub fn router(&self) -> Router {
        Router::new()
            .route("/metrics", get(render_metrics))
            .route("/health", get(metrics_health))
            .route_layer(middleware::from_fn_with_state(
                self.state.clone(),
                authorize_metrics,
            ))
            .with_state(self.state.clone())
    }

    pub async fn bind(self) -> Result<MetricsServer, MetricsServeError> {
        let listener = TcpListener::bind(self.config.bind_addr)
            .await
            .map_err(|_| MetricsServeError::Bind)?;
        Ok(MetricsServer {
            listener,
            router: self.router(),
        })
    }
}

pub struct MetricsServer {
    listener: TcpListener,
    router: Router,
}

impl MetricsServer {
    pub fn local_addr(&self) -> Result<SocketAddr, MetricsServeError> {
        self.listener
            .local_addr()
            .map_err(|_| MetricsServeError::LocalAddress)
    }

    pub async fn serve(self) -> Result<(), MetricsServeError> {
        axum::serve(self.listener, self.router)
            .await
            .map_err(|_| MetricsServeError::Serve)
    }
}

#[derive(Debug, Error)]
pub enum MetricsSurfaceError {
    #[error("failed to configure Prometheus recorder")]
    RecorderConfiguration,
}

#[derive(Debug, Error)]
pub enum MetricsServeError {
    #[error("failed to bind metrics listener")]
    Bind,
    #[error("metrics listener stopped unexpectedly")]
    Serve,
    #[error("failed to read metrics listener address")]
    LocalAddress,
}

pub(crate) fn install_prometheus_recorder(
    identity: &ServiceIdentity,
) -> Result<PrometheusHandle, MetricsSurfaceError> {
    prometheus_builder(identity)?
        .install_recorder()
        .map_err(|_| MetricsSurfaceError::RecorderConfiguration)
}

fn prometheus_builder(
    identity: &ServiceIdentity,
) -> Result<PrometheusBuilder, MetricsSurfaceError> {
    PrometheusBuilder::new()
        .set_buckets(DURATION_BUCKETS_SECONDS)
        .map(|builder| {
            builder
                .add_global_label("service", identity.service())
                .add_global_label("version", identity.version())
                .add_global_label("environment", identity.environment())
        })
        .map_err(|_| MetricsSurfaceError::RecorderConfiguration)
}

async fn render_metrics(State(state): State<MetricsState>) -> Response {
    let mut response = state.handle.render().into_response();
    response.headers_mut().insert(
        header::CONTENT_TYPE,
        HeaderValue::from_static(PROMETHEUS_CONTENT_TYPE),
    );
    response
}

async fn metrics_health() -> impl IntoResponse {
    (StatusCode::OK, "ok\n")
}

async fn authorize_metrics(
    State(state): State<MetricsState>,
    request: Request,
    next: Next,
) -> Response {
    if !state.requires_authentication {
        return next.run(request).await;
    }

    let authorized = bearer_token(request.headers())
        .map(token_digest)
        .zip(state.token_digest)
        .is_some_and(|(actual, expected)| bool::from(actual.ct_eq(&expected)));

    if authorized {
        next.run(request).await
    } else {
        StatusCode::UNAUTHORIZED.into_response()
    }
}

fn bearer_token(headers: &HeaderMap) -> Option<&[u8]> {
    let value = headers.get(header::AUTHORIZATION)?.as_bytes();
    let separator = value.iter().position(|byte| *byte == b' ')?;
    let (scheme, token_with_spaces) = value.split_at(separator);
    let token_start = token_with_spaces.iter().position(|byte| *byte != b' ')?;
    let token = &token_with_spaces[token_start..];

    scheme
        .eq_ignore_ascii_case(b"bearer")
        .then_some(token)
        .filter(|token| !token.is_empty())
}

fn token_digest(token: &[u8]) -> [u8; 32] {
    Sha256::digest(token).into()
}

fn parse_enabled(value: Result<String, env::VarError>) -> Result<bool, MetricsConfigError> {
    match value {
        Ok(raw) => match raw.to_ascii_lowercase().as_str() {
            "true" | "1" => Ok(true),
            "false" | "0" => Ok(false),
            _ => Err(MetricsConfigError::InvalidEnabledFlag),
        },
        Err(env::VarError::NotPresent) => Ok(true),
        Err(env::VarError::NotUnicode(_)) => Err(MetricsConfigError::InvalidEnvironmentEncoding {
            field: METRICS_ENABLED_ENV,
        }),
    }
}