//! 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>`. `--surface` says which store this
//! instance guards. It answers 200 when the secret is the one stored at
//! `swarm/agents/<agent>/queue` and the agent holds the permission for that
//! surface (`logs` or `metrics`), 401 when the secret does not verify, and 403
//! when it does but the permission is not held. A grant it cannot read is not
//! held. Any other status becomes a 500 for the caller in `auth_request`.
//!
//! The permission is read from the leaf swarm-controller mirrors
//! (`swarm_secret_client::observability`), because this responder has no queue
//! access.
//!
//! It logs in to the store as the cert-auth role `--store-cert-role` names,
//! whose policy reaches every agent's `queue` secret and grant leaf. 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, or a permission that is revoked, is still honoured for up
//! to that long. Only an allow is remembered; a denial is looked up again.

mod verdicts;

use std::future::Future;
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, ValueEnum};
use swarm_agent_verify::CredentialSource;
use swarm_agent_verify::agent_token::parse_agent_token;
use swarm_secret_client::observability::{Grant, grant_path};
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,
    /// Which store this instance guards, and so which permission it checks.
    #[arg(long, value_enum)]
    surface: Surface,
    /// 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,
}

/// The store a gateway fronts.
#[derive(Clone, Copy, Debug, PartialEq, Eq, ValueEnum)]
enum Surface {
    /// The log store, behind the `logs` permission.
    Logs,
    /// The metrics store, behind the `metrics` permission.
    Metrics,
}

impl Surface {
    fn allowed_by(self, grant: Grant) -> bool {
        match self {
            Self::Logs => grant.logs,
            Self::Metrics => grant.metrics,
        }
    }
}

/// Where an agent's mirrored [`Grant`] comes from.
trait GrantSource {
    /// The grant mirrored for `agent`, `None` when nothing is stored, or an
    /// error when the store could not answer.
    fn grant(&self, agent: &str) -> impl Future<Output = anyhow::Result<Option<Grant>>> + Send;
}

impl GrantSource for swarm_agent_verify::Store {
    async fn grant(&self, agent: &str) -> anyhow::Result<Option<Grant>> {
        self.read_optional(&grant_path(agent)?).await
    }
}

struct Responder<S> {
    source: S,
    surface: Surface,
    /// How long one store lookup may take before the request is denied.
    lookup_timeout: Duration,
    verdicts: Mutex<Verdicts>,
}

/// The answer for one subrequest's `Authorization` header, at `now`.
async fn decide(
    responder: &Responder<impl CredentialSource + GrantSource>,
    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, responder.lookup_timeout).await {
        return StatusCode::UNAUTHORIZED;
    }
    if !holds_permission(responder, token.agent).await {
        return StatusCode::FORBIDDEN;
    }
    tracing::info!(agent = token.agent, surface = ?responder.surface, "verified");
    lock(&responder.verdicts).grant(bearer, now);
    StatusCode::OK
}

/// Whether `agent` holds the permission for the surface. A grant that is
/// absent, unreadable or slow is not held.
async fn holds_permission(responder: &Responder<impl GrantSource>, agent: &str) -> bool {
    let grant =
        match tokio::time::timeout(responder.lookup_timeout, responder.source.grant(agent)).await {
            Ok(Ok(Some(grant))) => grant,
            Ok(Ok(None)) => {
                tracing::warn!(
                    agent,
                    "no permission grant is mirrored for this agent; denying"
                );
                return false;
            }
            Ok(Err(e)) => {
                tracing::warn!(
                    agent,
                    error = format!("{e:#}"),
                    "grant lookup failed; denying"
                );
                return false;
            }
            Err(_) => {
                tracing::warn!(agent, "grant lookup timed out; denying");
                return false;
            }
        };
    let held = responder.surface.allowed_by(grant);
    if !held {
        tracing::info!(agent, surface = ?responder.surface, "agent lacks the permission; denying");
    }
    held
}

/// 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),
        surface: args.surface,
        lookup_timeout: LOOKUP_TIMEOUT,
        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;

    /// What the mirrored grant leaf answers.
    #[derive(Clone, Copy)]
    enum Leaf {
        Holds(Grant),
        Absent,
        Unreadable,
        /// Accepts the read and never answers.
        Hangs,
    }

    /// A store holding atlas's credential until `revoke` and her grant leaf,
    /// counting lookups of each.
    struct Atlas {
        lookups: AtomicUsize,
        revoked: AtomicBool,
        leaf: Mutex<Leaf>,
        grant_lookups: AtomicUsize,
    }

    impl Atlas {
        fn with(leaf: Leaf) -> Self {
            Self {
                lookups: AtomicUsize::new(0),
                revoked: AtomicBool::new(false),
                leaf: Mutex::new(leaf),
                grant_lookups: AtomicUsize::new(0),
            }
        }

        /// Atlas holds `logs` and `metrics`.
        fn both() -> Self {
            Self::with(Leaf::Holds(Grant {
                logs: true,
                metrics: true,
            }))
        }

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

        fn set_leaf(&self, leaf: Leaf) {
            *self.leaf.lock().unwrap() = leaf;
        }

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

    impl GrantSource for &Atlas {
        async fn grant(&self, agent: &str) -> anyhow::Result<Option<Grant>> {
            self.grant_lookups.fetch_add(1, Ordering::SeqCst);
            let leaf = *self.leaf.lock().unwrap();
            match (agent, leaf) {
                ("atlas", Leaf::Holds(grant)) => Ok(Some(grant)),
                ("atlas", Leaf::Unreadable) => anyhow::bail!("store unreachable"),
                ("atlas", Leaf::Hangs) => std::future::pending().await,
                _ => Ok(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, surface: Surface) -> Responder<&Atlas> {
        Responder {
            source: store,
            surface,
            lookup_timeout: Duration::from_millis(100),
            verdicts: Mutex::new(Verdicts::new(TTL, 16)),
        }
    }

    #[tokio::test]
    async fn the_stored_secret_with_the_grant_is_200() {
        let store = Atlas::both();
        for surface in [Surface::Logs, Surface::Metrics] {
            assert_eq!(
                decide(&responder(&store, surface), Some(GOOD), Instant::now()).await,
                StatusCode::OK,
                "{surface:?}"
            );
        }
    }

    /// Each surface checks its own permission: a `logs` grant opens the log
    /// store and not the metrics store, and the reverse.
    #[tokio::test]
    async fn each_surface_needs_its_own_permission() {
        async fn at(store: &Atlas, surface: Surface) -> StatusCode {
            decide(&responder(store, surface), Some(GOOD), Instant::now()).await
        }
        let logs_only = Atlas::with(Leaf::Holds(Grant {
            logs: true,
            metrics: false,
        }));
        let metrics_only = Atlas::with(Leaf::Holds(Grant {
            logs: false,
            metrics: true,
        }));
        assert_eq!(at(&logs_only, Surface::Logs).await, StatusCode::OK);
        assert_eq!(
            at(&logs_only, Surface::Metrics).await,
            StatusCode::FORBIDDEN
        );
        assert_eq!(at(&metrics_only, Surface::Metrics).await, StatusCode::OK);
        assert_eq!(
            at(&metrics_only, Surface::Logs).await,
            StatusCode::FORBIDDEN
        );
    }

    /// No grant, an empty grant and a grant that cannot be read all deny.
    #[tokio::test]
    async fn without_a_readable_grant_the_answer_is_403() {
        for leaf in [
            Leaf::Absent,
            Leaf::Unreadable,
            Leaf::Hangs,
            Leaf::Holds(Grant::default()),
        ] {
            let store = Atlas::with(leaf);
            assert_eq!(
                decide(
                    &responder(&store, Surface::Logs),
                    Some(GOOD),
                    Instant::now()
                )
                .await,
                StatusCode::FORBIDDEN
            );
        }
    }

    #[tokio::test]
    async fn everything_else_is_401() {
        let store = Atlas::both();
        let r = responder(&store, Surface::Logs);
        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 wrong secret is refused before the grant is read.
        assert_eq!(store.grant_lookups.load(Ordering::SeqCst), 0);
    }

    /// An allow is reused within the TTL — a revocation the store already
    /// holds, of the secret or of the permission, goes unseen — and checked
    /// against the store again after it.
    #[tokio::test]
    async fn an_allow_spares_the_store_until_the_ttl() {
        let store = Atlas::both();
        let r = responder(&store, Surface::Logs);
        let t0 = Instant::now();
        assert_eq!(decide(&r, Some(GOOD), t0).await, StatusCode::OK);
        store.revoke();
        store.set_leaf(Leaf::Holds(Grant::default()));
        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);
    }

    /// A revoked permission is denied once the allow expires; a denial is not
    /// remembered, so a permission granted later is honoured at once.
    #[tokio::test]
    async fn a_revoked_or_added_permission_takes_effect_by_the_ttl() {
        let store = Atlas::both();
        let r = responder(&store, Surface::Logs);
        let t0 = Instant::now();
        assert_eq!(decide(&r, Some(GOOD), t0).await, StatusCode::OK);
        store.set_leaf(Leaf::Holds(Grant::default()));
        assert_eq!(
            decide(&r, Some(GOOD), t0 + TTL).await,
            StatusCode::FORBIDDEN
        );
        assert_eq!(
            decide(&r, Some(GOOD), t0 + TTL).await,
            StatusCode::FORBIDDEN
        );
        store.set_leaf(Leaf::Holds(Grant {
            logs: true,
            metrics: false,
        }));
        assert_eq!(decide(&r, Some(GOOD), t0 + TTL).await, StatusCode::OK);
    }

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