//! The log and metrics stores' gateway `auth_request` responder for agents.
//!
//! nginx sends it every `/select/logsql/` request (`swarm-victorialogs.nix`),
//! and every metrics read API request (`swarm-victoriametrics.nix`, as
//! `swarm-metrics-agent-auth`), whose `Authorization` is
//! `Bearer swarm-agent.<agent>.<secret>`. It answers 200 when the secret is
//! the one stored at `swarm/agents/<agent>/queue`, and 401 otherwise, store
//! failures included: `auth_request` turns any other status into a 500 for
//! the caller.
//!
//! It logs in to the store as the cert-auth role `--store-cert-role` names,
//! whose policy reaches every agent's `queue` secret. It serves a unix socket
//! whose directory only the gateway's nginx may enter.
//!
//! A grant is remembered for `--verdict-ttl` seconds ([`verdicts`]), so an
//! agent's queries within that window cost no store lookup, and a secret that
//! stops verifying is still accepted for up to that long.

mod verdicts;

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

use anyhow::Context as _;
use axum::extract::State;
use axum::http::{HeaderMap, StatusCode, header};
use clap::Parser;
use swarm_agent_verify::CredentialSource;
use swarm_agent_verify::agent_token::parse_agent_token;
use verdicts::Verdicts;

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

/// Most grants held at once. One live secret per agent, so this is a bound on
/// memory, not a size a swarm is expected to reach.
const VERDICT_CAP: usize = 4096;

#[derive(Parser)]
#[command(about = "auth_request responder: checks an agent's queue secret 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-logs-auth")]
    store_cert_role: String,
    /// Seconds a verified bearer is granted again without a store lookup;
    /// 0 looks up every request.
    #[arg(long, default_value_t = 60)]
    verdict_ttl: u64,
}

struct Responder<S> {
    source: S,
    verdicts: Mutex<Verdicts>,
}

/// The answer for one subrequest's `Authorization` header, at `now`.
async fn decide(
    responder: &Responder<impl CredentialSource>,
    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;
    };
    let token = match parse_agent_token(bearer) {
        Some(Ok(token)) => token,
        Some(Err(e)) => {
            tracing::warn!(error = %e, "malformed agent token; denying");
            return StatusCode::UNAUTHORIZED;
        }
        None => {
            tracing::warn!("not an agent token; denying");
            return StatusCode::UNAUTHORIZED;
        }
    };
    if lock(&responder.verdicts).granted(bearer, now) {
        tracing::debug!(agent = token.agent, "verified recently");
        return StatusCode::OK;
    }
    if swarm_agent_verify::verify(&responder.source, &token, LOOKUP_TIMEOUT).await {
        tracing::info!(agent = token.agent, "verified");
        lock(&responder.verdicts).grant(bearer, now);
        StatusCode::OK
    } else {
        StatusCode::UNAUTHORIZED
    }
}

/// The cache holds no invariant a panicking holder could break mid-way, so a
/// poisoned lock is used as is.
fn lock(verdicts: &Mutex<Verdicts>) -> std::sync::MutexGuard<'_, Verdicts> {
    verdicts
        .lock()
        .unwrap_or_else(std::sync::PoisonError::into_inner)
}

async fn handle(
    State(responder): State<Arc<Responder<swarm_agent_verify::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_secret_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 = swarm_secret_client::client::Settings::from_env()
        .context("reading the secret store's address and this responder's identity")?;
    let responder = Arc::new(Responder {
        source: swarm_agent_verify::Store::new(settings, args.store_cert_role),
        verdicts: Mutex::new(Verdicts::new(
            Duration::from_secs(args.verdict_ttl),
            VERDICT_CAP,
        )),
    });

    // `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-logs-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::atomic::{AtomicBool, AtomicUsize, Ordering};
    use swarm_secret_client::queue::AgentCredential;

    /// A store holding atlas's credential until `revoke`, counting lookups.
    #[derive(Default)]
    struct Atlas {
        lookups: AtomicUsize,
        revoked: AtomicBool,
    }

    impl Atlas {
        fn revoke(&self) {
            self.revoked.store(true, Ordering::SeqCst);
        }

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

    impl CredentialSource for &Atlas {
        async fn lookup(&self, agent: &str) -> anyhow::Result<Option<AgentCredential>> {
            self.lookups.fetch_add(1, Ordering::SeqCst);
            Ok(
                (agent == "atlas" && !self.revoked.load(Ordering::SeqCst)).then(|| {
                    AgentCredential {
                        value: "Ab9_-zSECRET".to_owned(),
                        agent: "atlas".to_owned(),
                        minted_at: None,
                    }
                }),
            )
        }
    }

    const TTL: Duration = Duration::from_mins(1);
    const GOOD: &str = "Bearer swarm-agent.atlas.Ab9_-zSECRET";
    const WRONG: &str = "Bearer swarm-agent.atlas.Ab9_-zSECREX";

    fn responder(store: &Atlas) -> Responder<&Atlas> {
        Responder {
            source: store,
            verdicts: Mutex::new(Verdicts::new(TTL, 16)),
        }
    }

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

    #[tokio::test]
    async fn everything_else_is_401() {
        let store = Atlas::default();
        let r = responder(&store);
        for authorization in [
            None,
            Some("swarm-agent.atlas.Ab9_-zSECRET"),
            Some("Basic swarm-agent.atlas.Ab9_-zSECRET"),
            Some(WRONG),
            Some("Bearer swarm-agent.argus.Ab9_-zSECRET"),
            Some("Bearer swarm-agent.atlas"),
            Some("Bearer authelia_at_abc.def"),
        ] {
            assert_eq!(
                decide(&r, authorization, Instant::now()).await,
                StatusCode::UNAUTHORIZED,
                "{authorization:?}"
            );
        }
    }

    /// A grant is reused within the TTL — a revocation the store already
    /// holds goes unseen — and checked against the store again after it.
    #[tokio::test]
    async fn a_grant_spares_the_store_until_the_ttl() {
        let store = Atlas::default();
        let r = responder(&store);
        let t0 = Instant::now();
        assert_eq!(decide(&r, Some(GOOD), t0).await, StatusCode::OK);
        store.revoke();
        let within = t0 + Duration::from_secs(59);
        assert_eq!(decide(&r, Some(GOOD), within).await, StatusCode::OK);
        assert_eq!(store.lookups(), 1);
        assert_eq!(
            decide(&r, Some(GOOD), t0 + TTL).await,
            StatusCode::UNAUTHORIZED
        );
        assert_eq!(store.lookups(), 2);
    }

    #[tokio::test]
    async fn a_denial_is_looked_up_every_time() {
        let store = Atlas::default();
        let r = responder(&store);
        let t0 = Instant::now();
        for _ in 0..3 {
            assert_eq!(decide(&r, Some(WRONG), t0).await, StatusCode::UNAUTHORIZED);
        }
        assert_eq!(store.lookups(), 3);
    }
}