//! 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, and a
//! failed or timed-out one is remembered for `LOOKUP_FAILURE_WINDOW`, so a
//! burst of requests against an expired copy costs one lookup, success or
//! failure.

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

/// How long requests are denied without a lookup after one failed or timed
/// out. Each request queued behind a failing lookup would otherwise run its
/// own, and the k-th would wait k × `LOOKUP_TIMEOUT`, past nginx's 60s read
/// timeout. Short against `--cache-ttl`, so a recovered store is used again
/// within seconds.
const LOOKUP_FAILURE_WINDOW: Duration = Duration::from_secs(3);

#[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, except for a few seconds after a read
    /// failed.
    #[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,
}

#[derive(Default)]
struct Cache {
    stored: Option<Cached>,
    /// Until when requests are denied without a lookup, after one failed or
    /// timed out. Stamped when that lookup ends: the requests queued behind
    /// it carry the time they arrived, which is earlier.
    failed_until: Option<Instant>,
}

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

/// 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 cache = responder.cached.lock().await;
    let fresh = cache
        .stored
        .as_ref()
        .is_some_and(|c| now.saturating_duration_since(c.read_at) < responder.ttl);
    if !fresh {
        if cache.failed_until.is_some_and(|until| now < until) {
            tracing::warn!("bearer lookup failed a moment ago; denying without another");
            return StatusCode::UNAUTHORIZED;
        }
        match tokio::time::timeout(LOOKUP_TIMEOUT, responder.source.lookup()).await {
            Ok(Ok(Some(value))) if !value.is_empty() => {
                cache.stored = Some(Cached {
                    value,
                    read_at: now,
                });
                cache.failed_until = None;
            }
            Ok(Ok(_)) => {
                tracing::warn!("no bearer is stored; denying");
                return StatusCode::UNAUTHORIZED;
            }
            Ok(Err(e)) => {
                tracing::warn!(error = format!("{e:#}"), "bearer lookup failed; denying");
                cache.failed_until = Some(Instant::now() + LOOKUP_FAILURE_WINDOW);
                return StatusCode::UNAUTHORIZED;
            }
            Err(_) => {
                tracing::warn!("bearer lookup timed out; denying");
                cache.failed_until = Some(Instant::now() + LOOKUP_FAILURE_WINDOW);
                return StatusCode::UNAUTHORIZED;
            }
        }
    }
    let Some(stored) = cache.stored.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(Cache::default()),
    });

    // `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`) after `delay`,
    /// counting lookups.
    struct Fake {
        stored: Mutex<Result<Option<String>, ()>>,
        delay: Duration,
        lookups: AtomicUsize,
    }

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

        fn hanging() -> Self {
            Self {
                delay: Duration::from_hours(1),
                ..Self::holding(SECRET)
            }
        }

        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);
            tokio::time::sleep(self.delay).await;
            self.stored
                .lock()
                .unwrap()
                .clone()
                .map_err(|()| anyhow::anyhow!("store unreachable"))
        }
    }

    impl BearerSource for Arc<Fake> {
        async fn lookup(&self) -> anyhow::Result<Option<String>> {
            (&**self).lookup().await
        }
    }

    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(Cache::default()),
        }
    }

    #[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 is remembered for `LOOKUP_FAILURE_WINDOW`, not for a TTL:
    /// requests inside it are denied without a lookup, the first one after it
    /// reads the store again.
    #[tokio::test]
    async fn a_failed_lookup_is_retried_after_the_failure_window() {
        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::UNAUTHORIZED);
        assert_eq!(store.lookups(), 1);
        let after = t0 + LOOKUP_FAILURE_WINDOW + Duration::from_secs(1);
        assert_eq!(decide(&r, Some(GOOD), after).await, StatusCode::OK);
        assert_eq!(store.lookups(), 2);
    }

    /// A burst against a store that errors, or that sleeps past
    /// `LOOKUP_TIMEOUT`, costs one lookup, and every request is denied
    /// within that one timeout rather than queueing behind one timeout each.
    #[tokio::test(start_paused = true)]
    async fn a_burst_against_a_failing_store_costs_one_lookup() {
        const BURST: usize = 20;
        let erroring = Fake::holding(SECRET);
        erroring.set(Err(()));
        for (name, store) in [("erroring", erroring), ("hanging", Fake::hanging())] {
            let store = Arc::new(store);
            let r = Arc::new(Responder {
                source: Arc::clone(&store),
                ttl: TTL,
                cached: tokio::sync::Mutex::new(Cache::default()),
            });
            let started = tokio::time::Instant::now();
            let now = Instant::now();
            let mut burst = tokio::task::JoinSet::new();
            for _ in 0..BURST {
                let r = Arc::clone(&r);
                burst.spawn(async move { decide(&r, Some(GOOD), now).await });
            }
            while let Some(status) = burst.join_next().await {
                assert_eq!(status.unwrap(), StatusCode::UNAUTHORIZED, "{name}");
            }
            assert_eq!(store.lookups(), 1, "{name}");
            assert!(started.elapsed() <= LOOKUP_TIMEOUT, "{name}");
        }
    }
}