hyperhive/swarm-metrics-auth/src/main.rs

//! The metrics store gateway's `auth_request` responder for Grafana.
//!
//! nginx sends it every request to the store's query API that carries a
//! `Bearer` (`swarm-victoriametrics.nix`). It answers 200 when the bearer is
//! the value stored at `--store-path`, which Grafana's host creates and its
//! datasource presents, and 401 otherwise, store failures included:
//! `auth_request` turns any other status into a 500 for the caller.
//!
//! It reads the store under its own certificate, whose policy reaches that one
//! path. It serves a unix socket whose directory only the gateway's nginx may
//! enter.
//!
//! The stored value is kept for `--cache-ttl` seconds, and no request reaches
//! the store within that window: a replaced bearer is refused, and the one it
//! replaced still accepted, for up to that long. Lookups are serialised, so a
//! burst of requests against an expired copy costs one.

use std::future::Future;
use std::os::unix::fs::PermissionsExt as _;
use std::path::PathBuf;
use std::sync::Arc;
use std::time::{Duration, Instant};

use anyhow::Context as _;
use axum::extract::State;
use axum::http::{HeaderMap, StatusCode, header};
use clap::Parser;
use subtle::ConstantTimeEq as _;
use swarm_secret_client::client::{DEFAULT_CERT_MOUNT, Settings};

/// How long one store lookup may take before the request is denied. Well
/// under nginx's 60s `proxy_read_timeout`, so a slow store reads as a 401
/// rather than a gateway timeout.
const LOOKUP_TIMEOUT: Duration = Duration::from_secs(5);

#[derive(Parser)]
#[command(about = "auth_request responder: checks Grafana's metrics bearer against the store")]
struct Args {
    /// Unix socket to serve on.
    #[arg(long)]
    socket: PathBuf,
    /// Cert-auth role this responder logs in to the store as.
    #[arg(long, default_value = "swarm-metrics-auth")]
    store_cert_role: String,
    /// KV path, under the `secret` mount, holding the bearer in its `value`
    /// field.
    #[arg(long)]
    store_path: String,
    /// Seconds the stored bearer is used without reading the store again;
    /// 0 reads it on every request.
    #[arg(long, default_value_t = 60)]
    cache_ttl: u64,
}

/// Where the stored bearer comes from.
trait BearerSource {
    /// The stored bearer, `None` when nothing is stored, or an error when the
    /// store could not answer.
    fn lookup(&self) -> impl Future<Output = anyhow::Result<Option<String>>> + Send;
}

/// The object at the bearer's path. `bao kv put … value=-` on Grafana's host
/// writes this field.
#[derive(serde::Deserialize)]
struct Stored {
    value: String,
}

/// The swarm secret store, reached with the responder's own certificate.
struct Store {
    settings: Settings,
    cert_role: String,
    path: String,
}

impl BearerSource for Store {
    /// Logs in per lookup, so no token is left to expire inside a long-lived
    /// process.
    async fn lookup(&self) -> anyhow::Result<Option<String>> {
        let store = swarm_secret_client::SecretStore::connect(
            &self.settings,
            &self.cert_role,
            DEFAULT_CERT_MOUNT,
        )
        .await
        .context("logging in to the secret store")?;
        let stored: Option<Stored> = store
            .read_optional(&self.path)
            .await
            .with_context(|| format!("reading {}", self.path))?;
        Ok(stored.map(|s| s.value))
    }
}

struct Cached {
    value: String,
    read_at: Instant,
}

struct Responder<S> {
    source: S,
    ttl: Duration,
    cached: tokio::sync::Mutex<Option<Cached>>,
}

/// The answer for one subrequest's `Authorization` header, at `now`.
async fn decide(
    responder: &Responder<impl BearerSource>,
    authorization: Option<&str>,
    now: Instant,
) -> StatusCode {
    let Some(bearer) = authorization.and_then(|v| v.strip_prefix("Bearer ")) else {
        tracing::warn!("no bearer token; denying");
        return StatusCode::UNAUTHORIZED;
    };
    // Held across the lookup, so concurrent requests wait for one read
    // rather than each starting their own.
    let mut cached = responder.cached.lock().await;
    let fresh = cached
        .as_ref()
        .is_some_and(|c| now.saturating_duration_since(c.read_at) < responder.ttl);
    if !fresh {
        match tokio::time::timeout(LOOKUP_TIMEOUT, responder.source.lookup()).await {
            Ok(Ok(Some(value))) if !value.is_empty() => {
                *cached = Some(Cached {
                    value,
                    read_at: now,
                });
            }
            Ok(Ok(_)) => {
                tracing::warn!("no bearer is stored; denying");
                return StatusCode::UNAUTHORIZED;
            }
            Ok(Err(e)) => {
                tracing::warn!(error = format!("{e:#}"), "bearer lookup failed; denying");
                return StatusCode::UNAUTHORIZED;
            }
            Err(_) => {
                tracing::warn!("bearer lookup timed out; denying");
                return StatusCode::UNAUTHORIZED;
            }
        }
    }
    let Some(stored) = cached.as_ref() else {
        return StatusCode::UNAUTHORIZED;
    };
    // Constant time, so the time to refuse says nothing about how much of the
    // bearer was right.
    if bool::from(stored.value.as_bytes().ct_eq(bearer.as_bytes())) {
        StatusCode::OK
    } else {
        tracing::warn!("bearer does not match the stored one; denying");
        StatusCode::UNAUTHORIZED
    }
}

async fn handle(State(responder): State<Arc<Responder<Store>>>, headers: HeaderMap) -> StatusCode {
    let authorization = headers
        .get(header::AUTHORIZATION)
        .and_then(|v| v.to_str().ok());
    decide(&responder, authorization, Instant::now()).await
}

#[tokio::main]
async fn main() -> anyhow::Result<()> {
    swarm_queue_client::install_crypto_provider();
    tracing_subscriber::fmt()
        .with_env_filter(
            tracing_subscriber::EnvFilter::try_from_default_env()
                .unwrap_or_else(|_| tracing_subscriber::EnvFilter::new("info")),
        )
        // stdout is journald, which keeps ANSI escapes as raw bytes.
        .with_ansi(false)
        .init();
    let args = Args::parse();

    let settings = Settings::from_env()
        .context("reading the secret store's address and this responder's identity")?;
    let responder = Arc::new(Responder {
        source: Store {
            settings,
            cert_role: args.store_cert_role,
            path: args.store_path,
        },
        ttl: Duration::from_secs(args.cache_ttl),
        cached: tokio::sync::Mutex::new(None),
    });

    // `RuntimeDirectoryPreserve=yes` can leave the previous run's socket; the
    // directory is this unit's alone, so unlinking it is safe.
    if let Err(e) = std::fs::remove_file(&args.socket)
        && e.kind() != std::io::ErrorKind::NotFound
    {
        return Err(e)
            .with_context(|| format!("clearing stale socket at {}", args.socket.display()));
    }
    let listener = tokio::net::UnixListener::bind(&args.socket)
        .with_context(|| format!("binding {}", args.socket.display()))?;
    // nginx is another user and `connect(2)` needs write. The directory is
    // the access control, not this mode.
    std::fs::set_permissions(&args.socket, std::fs::Permissions::from_mode(0o666))
        .with_context(|| format!("chmod {}", args.socket.display()))?;
    tracing::info!(socket = %args.socket.display(), "swarm-metrics-auth listening");

    let app = axum::Router::new().fallback(handle).with_state(responder);
    axum::serve(listener, app).await.context("serving")
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::sync::Mutex;
    use std::sync::atomic::{AtomicUsize, Ordering};

    /// A store answering `stored` (or failing when it is `Err`), counting
    /// lookups.
    struct Fake {
        stored: Mutex<Result<Option<String>, ()>>,
        lookups: AtomicUsize,
    }

    impl Fake {
        fn holding(value: &str) -> Self {
            Self {
                stored: Mutex::new(Ok(Some(value.to_owned()))),
                lookups: AtomicUsize::new(0),
            }
        }

        fn set(&self, stored: Result<Option<&str>, ()>) {
            *self.stored.lock().unwrap() = stored.map(|v| v.map(str::to_owned));
        }

        fn lookups(&self) -> usize {
            self.lookups.load(Ordering::SeqCst)
        }
    }

    impl BearerSource for &Fake {
        async fn lookup(&self) -> anyhow::Result<Option<String>> {
            self.lookups.fetch_add(1, Ordering::SeqCst);
            self.stored
                .lock()
                .unwrap()
                .clone()
                .map_err(|()| anyhow::anyhow!("store unreachable"))
        }
    }

    const TTL: Duration = Duration::from_mins(1);
    const SECRET: &str = "0f1e2d3c4b5a";
    const GOOD: &str = "Bearer 0f1e2d3c4b5a";

    fn responder(store: &Fake) -> Responder<&Fake> {
        Responder {
            source: store,
            ttl: TTL,
            cached: tokio::sync::Mutex::new(None),
        }
    }

    #[tokio::test]
    async fn the_stored_bearer_is_200() {
        let store = Fake::holding(SECRET);
        assert_eq!(
            decide(&responder(&store), Some(GOOD), Instant::now()).await,
            StatusCode::OK
        );
    }

    #[tokio::test]
    async fn everything_else_is_401() {
        let store = Fake::holding(SECRET);
        let r = responder(&store);
        for authorization in [
            None,
            Some(SECRET),
            Some("Basic 0f1e2d3c4b5a"),
            Some("Bearer 0f1e2d3c4b5b"),
            Some("Bearer 0f1e2d3c4b5"),
            Some("Bearer "),
            Some("Bearer authelia_at_abc.def"),
        ] {
            assert_eq!(
                decide(&r, authorization, Instant::now()).await,
                StatusCode::UNAUTHORIZED,
                "{authorization:?}"
            );
        }
    }

    /// Nothing stored, an empty value, and a store that cannot answer are
    /// all denials, whatever the bearer.
    #[tokio::test]
    async fn no_usable_stored_value_is_401() {
        let store = Fake::holding(SECRET);
        let r = responder(&store);
        for stored in [Ok(None), Ok(Some("")), Err(())] {
            store.set(stored);
            assert_eq!(
                decide(&r, Some(GOOD), Instant::now()).await,
                StatusCode::UNAUTHORIZED,
                "{stored:?}"
            );
            assert_eq!(
                decide(&r, Some("Bearer "), Instant::now()).await,
                StatusCode::UNAUTHORIZED,
                "{stored:?}"
            );
        }
    }

    /// Within the TTL no request reaches the store, a wrong bearer included,
    /// and a replacement the store already holds goes unseen; after it the
    /// store is read again.
    #[tokio::test]
    async fn the_stored_value_is_reused_until_the_ttl() {
        let store = Fake::holding(SECRET);
        let r = responder(&store);
        let t0 = Instant::now();
        assert_eq!(decide(&r, Some(GOOD), t0).await, StatusCode::OK);
        store.set(Ok(Some("replaced")));
        let within = t0 + Duration::from_secs(59);
        assert_eq!(decide(&r, Some(GOOD), within).await, StatusCode::OK);
        assert_eq!(
            decide(&r, Some("Bearer replaced"), within).await,
            StatusCode::UNAUTHORIZED
        );
        assert_eq!(store.lookups(), 1);
        assert_eq!(
            decide(&r, Some(GOOD), t0 + TTL).await,
            StatusCode::UNAUTHORIZED
        );
        assert_eq!(
            decide(&r, Some("Bearer replaced"), t0 + TTL).await,
            StatusCode::OK
        );
        assert_eq!(store.lookups(), 2);
    }

    /// A failed read leaves nothing usable cached, so the next request reads
    /// the store again rather than waiting out a TTL.
    #[tokio::test]
    async fn a_failed_lookup_is_retried_on_the_next_request() {
        let store = Fake::holding(SECRET);
        store.set(Err(()));
        let r = responder(&store);
        let t0 = Instant::now();
        assert_eq!(decide(&r, Some(GOOD), t0).await, StatusCode::UNAUTHORIZED);
        store.set(Ok(Some(SECRET)));
        assert_eq!(decide(&r, Some(GOOD), t0).await, StatusCode::OK);
        assert_eq!(store.lookups(), 2);
    }
}