//! Auth-callout responder for the swarm's NATS queue.
//!
//! `nix/host-modules/swarm-nats.nix` configures `nats-server` with an
//! `auth_callout` block and no responder, which is the fail-closed state: the
//! server answers `"auth_required":true` and admits nobody. This binary is what
//! makes it able to say *yes*.
//!
//! It connects as the one callout-exempt user (by nkey, never by name — the
//! server refuses to start if that entry carries a username), subscribes to
//! `$SYS.REQ.USER.AUTH`, validates the presented bearer token against
//! authelia's introspection endpoint (or, for an agent's own token, against the
//! secret store: see `agent_token`), and answers with a NATS user JWT signed
//! by the account key. A rejection is answered explicitly: silence is
//! indistinguishable from the responder being down, and the queue is the
//! swarm's control path.
//!
//! Every request is opened with the callout xkey, and its reply is sealed
//! back to the server; see `xkey`. A request that does not open gets no
//! reply: the server this answers always seals. The `ensure-callout-keys`
//! verb creates that key and the two nkeys (see `nkey`) in the secret store,
//! and runs before the server does.
//!
//! # Secrets
//!
//! Every credential is taken as a **path**, never a value: a value in nix
//! config is rendered into the world-readable store, and a value in `argv` is
//! readable by anyone via `/proc/<pid>/cmdline`, which is `0444`. Paths are
//! not secrets, so passing them as flags is fine. The callout's three seeds
//! and the responder's own OIDC client secret are not files: they are read
//! from the secret store, at the paths `xkey`, `nkey` and `oidc_secret` name.

use std::path::PathBuf;

use anyhow::Context;
use clap::{Parser, Subcommand};
use futures_util::StreamExt;

mod agent_token;
mod introspect;
mod keystore;
mod nkey;
mod oidc_secret;
mod policy;
mod request;
mod respond;
mod xkey;

/// Subject the NATS server publishes authorization requests on.
const AUTH_SUBJECT: &str = "$SYS.REQ.USER.AUTH";

/// With no verb this is the responder, and its flags are required.
#[derive(Debug, Parser)]
#[command(
    name = "swarm-nats-auth",
    about = "Auth-callout responder for the swarm NATS queue",
    args_conflicts_with_subcommands = true,
    subcommand_negates_reqs = true
)]
struct Cli {
    #[command(subcommand)]
    command: Option<Command>,

    #[command(flatten)]
    responder: Option<Args>,
}

#[derive(Debug, Subcommand)]
enum Command {
    /// Create the callout xkey, issuer and user keys in the secret store when
    /// it holds none, and check that each public half is stored beside its
    /// seed. Needs no NATS server.
    EnsureCalloutKeys {
        /// Role on the store's `cert` auth mount to log in as.
        #[arg(long, default_value = "swarm-nats-auth")]
        store_cert_role: String,

        /// Issuer seed to store when the store holds none. Removed once the
        /// store holds an issuer seed. Absent or empty, a new key is made.
        #[arg(long)]
        import_issuer_seed_file: Option<PathBuf>,

        /// User seed to store when the store holds none, handled as
        /// `--import-issuer-seed-file` is.
        #[arg(long)]
        import_user_seed_file: Option<PathBuf>,
    },
}

#[derive(Debug, clap::Args)]
struct Args {
    /// NATS server to connect to.
    #[arg(long, default_value = "nats://127.0.0.1:4222")]
    nats_url: String,

    /// Authelia's OIDC introspection endpoint.
    #[arg(long)]
    introspection_url: String,

    /// `OAuth2` client id this responder introspects as, and whose secret it
    /// reads from the store. Must match
    /// `services.hyperhive.swarm.nats.clientId`, whose default this mirrors.
    #[arg(long, default_value = "swarm-nats")]
    client_id: String,

    /// Account an admitted client is placed in. Must name an entry in the
    /// server's own `accounts` block — in server-config mode the account is
    /// resolved by *name*, so a value the server does not know is a grant it
    /// refuses. The module passes its `clientAccount`; the default mirrors it.
    #[arg(long, default_value = "APP")]
    account: String,

    /// Client-id prefix that marks a hive. `swarm-authelia.nix` mints one
    /// machine client per roster entry as `hive-<name>`, while the KV key is
    /// the bare `<name>` — this is the contract between the two, declared
    /// rather than inferred from the shape of an id.
    #[arg(long, default_value = "hive-")]
    hive_client_prefix: String,

    /// Suffix that marks an agent container, appended to the hive's own client
    /// id: `swarm-authelia.nix` mints one machine client per roster entry as
    /// `hive-<name>-agent` — per **hive**, not per agent, because agents are
    /// created at runtime and a per-agent client would make creating one a
    /// config change plus a reload.
    ///
    /// So this identity says which hive an agent belongs to and never which
    /// agent: two agents on one hive are indistinguishable under it. The
    /// responder refuses it; the suffix is what tells it apart from a hive.
    ///
    /// A suffix on the hive's id rather than a prefix of its own, because
    /// `agent-<name>` reads as *the agent called `<name>`* — the one thing
    /// this identity does not carry.
    ///
    /// A suffix starts with `-`, so it has to be accepted as a value rather
    /// than read as a short-flag cluster — without this the default below
    /// cannot be passed back in on the command line that produced it.
    #[arg(long, default_value = "-agent", allow_hyphen_values = true)]
    agent_client_suffix: String,

    /// Client ids allowed to read every hive's status. Repeatable. The
    /// default is the swarm controller, which is the only reader that exists.
    #[arg(long = "reader-client", default_values_t = [String::from("swarm-controller")])]
    reader_clients: Vec<String>,

    /// Additional subjects a hive may publish to, with `{hive}` standing for
    /// its own name. Repeatable, empty by default.
    ///
    /// The extension point for a second stream published by the same
    /// `hive-<name>` identity — lifecycle notices, say. Without it, adding one
    /// means changing this responder; with it, a deployment says so and the
    /// subject still lands inside that hive's own namespace.
    #[arg(long = "hive-publish-subject")]
    hive_publish_subjects: Vec<String>,

    /// Subjects an agent that presented its own credential may publish to,
    /// with `{agent}` standing for its name. Repeatable, empty by default,
    /// which denies every agent token.
    #[arg(long = "agent-token-publish-subject")]
    agent_token_publish_subjects: Vec<String>,

    /// Client ids that answer the review classifier's try-out requests.
    /// Repeatable, empty by default, which denies such a client.
    #[arg(long = "review-classifier-client")]
    review_classifier_clients: Vec<String>,

    /// Role on the secret store's `cert` auth mount this responder logs in
    /// as to read the callout keys, its own OIDC client secret and agent
    /// credentials. The store is found
    /// through the `BAO_*` environment, which the responder does not start
    /// without.
    #[arg(long, default_value = "swarm-nats-auth")]
    store_cert_role: String,
}

#[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")),
        )
        // This is a systemd-managed daemon — stdout always goes to journald,
        // never a human terminal, and journald doesn't strip ANSI escapes:
        // they land in victorialogs as raw byte-array spam otherwise.
        .with_ansi(false)
        .init();
    let cli = Cli::parse();
    match (cli.command, cli.responder) {
        (
            Some(Command::EnsureCalloutKeys {
                store_cert_role,
                import_issuer_seed_file,
                import_user_seed_file,
            }),
            _,
        ) => {
            ensure_callout_keys(
                &store_cert_role,
                import_issuer_seed_file.as_deref(),
                import_user_seed_file.as_deref(),
            )
            .await
        }
        (None, Some(args)) => serve(&args).await,
        (None, None) => unreachable!("clap requires the responder's flags when no verb is given"),
    }
}

/// The responder: answers every authorization request until its
/// subscription ends.
async fn serve(args: &Args) -> anyhow::Result<()> {
    // The bucket is NOT a flag. `swarm_queue_client::status`'s own docs say
    // why: reader and writer must name the same bucket, and an option is a
    // way for two deployments to disagree about which one that is. This
    // responder is the third end that names it, so it takes the same
    // constant rather than a copy of the literal.
    // Fails the process rather than warning: a policy that cannot express a
    // per-hive namespace is not a policy this responder should run with, and
    // the queue's fail-closed state (no responder) is a legible outage where a
    // silently over-broad grant is not.
    let policy = policy::Policy::new(
        args.hive_client_prefix.clone(),
        args.agent_client_suffix.clone(),
        swarm_queue_client::status::BUCKET.to_owned(),
        args.reader_clients.clone(),
        args.hive_publish_subjects.clone(),
    )?
    .with_agent_token_subjects(args.agent_token_publish_subjects.clone())?
    .with_review_classifiers(args.review_classifier_clients.clone());
    let http = reqwest::Client::new();
    let settings = swarm_secret_client::client::Settings::from_env()
        .context("reading the secret store's address and this responder's identity")?;
    // Before the connect: a responder that cannot open a sealed request, sign
    // a grant, authenticate to the queue or introspect must not start
    // answering.
    let keys = swarm_secret_client::SecretStore::connect(
        &settings,
        &args.store_cert_role,
        swarm_secret_client::client::DEFAULT_CERT_MOUNT,
    )
    .await
    .context("logging in to the secret store for the callout keys and the OIDC client secret")?;
    let (callout_xkey, issuer, user_seed) = load_callout_keys(&keys).await?;
    let client_secret = load_client_secret(&keys, args).await?;
    let store = agent_token::Store::new(settings, args.store_cert_role.clone());
    let client = async_nats::ConnectOptions::with_nkey(user_seed)
        .name("swarm-nats-auth")
        .connect(&args.nats_url)
        .await
        .with_context(|| format!("connect to {}", args.nats_url))?;
    let mut requests = client
        .subscribe(AUTH_SUBJECT)
        .await
        .with_context(|| format!("subscribe to {AUTH_SUBJECT}"))?;
    tracing::info!(
        nats_url = %args.nats_url,
        subject = AUTH_SUBJECT,
        "swarm-nats-auth: connected, awaiting authorization requests"
    );

    while let Some(msg) = requests.next().await {
        let Some((server_key, req)) = open_request(&callout_xkey, &msg) else {
            continue;
        };
        // No token is a denial, not an error: an anonymous connect is a
        // normal thing for a client to attempt and an abnormal thing to
        // grant. Neither the store nor introspection is reached until
        // something was presented.
        let (caller, permissions) = match req
            .connect_opts
            .auth_token
            .as_deref()
            .map(agent_token::classify)
        {
            // The caller is the name the token claims, logged whether or not
            // the secret proved it; `granted` says which.
            Some(agent_token::Presented::Agent(token)) => (
                Some(format!("agent:{}", token.agent)),
                agent_token::authorize(&policy, &store, &token).await,
            ),
            Some(agent_token::Presented::Malformed(e)) => {
                tracing::warn!(error = %e, "malformed agent token; denying");
                (None, None)
            }
            Some(agent_token::Presented::Bearer(token)) => {
                introspected(&policy, &http, args, &client_secret, token).await
            }
            None => (None, None),
        };
        // The caller is an identifier, not a credential: a client id, or
        // `agent:<name>` for an agent token. One line per auth request, a connect
        // or a reconnect, so `hive-<h>-agent` lines count requests made on a
        // hive's shared credential, not agents.
        tracing::info!(
            user_nkey = %req.user_nkey,
            server_id = %req.server_id.id,
            granted = permissions.is_some(),
            caller = caller.as_deref().unwrap_or("-"),
            "auth request"
        );

        // Always reply, including on a denial. A server that hears nothing
        // cannot tell a refusing responder from a dead one, so silence turns
        // every rejection into a 2s timeout and hides an outage inside what
        // looks like ordinary denials.
        let Some(reply_to) = msg.reply.clone() else {
            tracing::warn!("auth request had no reply subject; dropping");
            continue;
        };
        let token = match &permissions {
            Some(permissions) => respond::grant(
                &issuer,
                &args.account,
                &req.server_id.id,
                &req.user_nkey,
                permissions,
            ),
            None => respond::deny(&issuer, &req.server_id.id, &req.user_nkey),
        };
        let reply = match xkey::seal(&callout_xkey, &server_key, token.as_bytes()) {
            Ok(reply) => reply,
            Err(e) => {
                tracing::warn!(error = %e, "could not seal the auth response; dropping");
                continue;
            }
        };
        if let Err(e) = client.publish(reply_to, reply.into()).await {
            tracing::warn!(error = ?e, "failed to publish auth response");
        }
    }
    anyhow::bail!("subscription to {AUTH_SUBJECT} ended")
}

/// The request in `msg`, opened, and the server key its reply is sealed to.
/// `None`, logged, for anything else, a request with no
/// [`xkey::SERVER_XKEY_HEADER`] included.
///
/// Logged and dropped, never propagated: the loop calling this is the
/// swarm's login path, and exiting on one malformed payload would let any
/// client take authentication down for everyone.
fn open_request(
    callout_xkey: &nkeys::XKey,
    msg: &async_nats::Message,
) -> Option<(nkeys::XKey, request::AuthRequest)> {
    let header = msg
        .headers
        .as_ref()
        .and_then(|h| h.get(xkey::SERVER_XKEY_HEADER))
        .map(async_nats::HeaderValue::as_str);
    let (server_key, payload) = match xkey::open(callout_xkey, header, &msg.payload) {
        Ok(opened) => opened,
        Err(e) => {
            tracing::warn!(error = %e, "unopenable auth request, ignoring");
            return None;
        }
    };
    match request::decode(&payload) {
        Ok(req) => Some((server_key, req)),
        Err(e) => {
            tracing::warn!(error = ?e, "undecodable auth request, ignoring");
            None
        }
    }
}

/// The callout xkey, the issuer key, and the user key's seed, read from the
/// store and never created here; see `xkey` and `nkey`.
async fn load_callout_keys(
    keys: &impl keystore::KeyStore,
) -> anyhow::Result<(nkeys::XKey, nkeys::KeyPair, String)> {
    let callout_xkey = xkey::load(keys).await?;
    tracing::info!(public_key = %callout_xkey.public_key(), "loaded the callout xkey");
    let issuer = nkey::load(keys, &nkey::ISSUER).await?;
    tracing::info!(public_key = %issuer.public_key(), "loaded the callout issuer key");
    let user = nkey::load(keys, &nkey::USER).await?;
    tracing::info!(public_key = %user.public_key(), "loaded the callout user key");
    let user_seed = user
        .seed()
        .map_err(|_| anyhow::anyhow!("the callout user key has no seed"))?;
    Ok((callout_xkey, issuer, user_seed))
}

/// This responder's own OIDC client secret, read from `keys`, and re-read
/// later as `args.store_cert_role`; see `oidc_secret`.
async fn load_client_secret(
    keys: &impl keystore::KeyStore,
    args: &Args,
) -> anyhow::Result<oidc_secret::OidcSecret<oidc_secret::Store>> {
    let path = oidc_secret::path(&args.client_id)?;
    let secret = oidc_secret::read(keys, &path).await?;
    tracing::info!(%path, "loaded the OIDC client secret");
    let source = oidc_secret::Store::new(
        swarm_secret_client::client::Settings::from_env()?,
        args.store_cert_role.clone(),
        path,
    );
    Ok(oidc_secret::OidcSecret::new(source, secret))
}

/// The `ensure-callout-keys` verb: see [`xkey::ensure`] and [`nkey::ensure`].
async fn ensure_callout_keys(
    store_cert_role: &str,
    import_issuer_seed_file: Option<&std::path::Path>,
    import_user_seed_file: Option<&std::path::Path>,
) -> anyhow::Result<()> {
    let store = swarm_secret_client::SecretStore::from_env(store_cert_role)
        .await
        .context("logging in to the secret store")?;
    let key = xkey::ensure(&store).await?;
    tracing::info!(
        public_key = %key.public_key(),
        path = xkey::PUBLIC_PATH,
        "the callout xkey is in the secret store"
    );
    for (pair, import) in [
        (&nkey::ISSUER, import_issuer_seed_file),
        (&nkey::USER, import_user_seed_file),
    ] {
        let key = nkey::ensure(&store, pair, import).await?;
        tracing::info!(
            public_key = %key.public_key(),
            path = pair.public_path,
            "the {} key is in the secret store", pair.name
        );
    }
    Ok(())
}

/// The OIDC path: who the `IdP` says presented `token`, and what the policy
/// grants that client. `None` permissions is a denial.
///
/// The caller is an identity or nothing — see `introspect`'s module docs.
/// There is no "admitted, identity unknown" branch to write here because
/// there is no such value to receive.
async fn introspected(
    policy: &policy::Policy,
    http: &reqwest::Client,
    args: &Args,
    client_secret: &oidc_secret::OidcSecret<impl oidc_secret::Source>,
    token: &str,
) -> (Option<String>, Option<policy::Permissions>) {
    let deadline = tokio::time::Instant::now() + introspect::INTROSPECTION_TIMEOUT;
    let caller = client_secret
        .introspect(deadline, |secret| async move {
            introspect::identify_caller(
                http,
                &args.introspection_url,
                &args.client_id,
                &secret,
                token,
                deadline,
            )
            .await
        })
        .await
        // An introspection that could not be *made* is a denial too. The
        // failure modes of an HTTP call are exactly the conditions under
        // which an attacker would most like this to fall open.
        .unwrap_or_else(|e| {
            tracing::warn!(error = ?e, "introspection failed; denying");
            None
        });
    // Admission said who; the policy says what. A caller the `IdP`
    // vouches for but no rule matches is denied — see `policy`'s module
    // docs for why that is deny and not "connect with nothing".
    let permissions = caller.as_deref().and_then(|id| policy.permissions(id));
    if let (Some(id), None) = (caller.as_deref(), permissions.as_ref()) {
        // Loud, and the one case an operator has to be able to find: a
        // valid credential refused by our own policy. The alternative is
        // a client that authenticates fine and mysteriously cannot work.
        tracing::warn!(
            caller = %id,
            "authenticated client matches no policy rule; denying"
        );
    }
    (caller, permissions)
}

#[cfg(test)]
mod tests {
    use super::*;

    /// The responder's flags, parsed the way `main` parses them.
    fn responder<const N: usize>(argv: [&str; N]) -> Result<Args, clap::Error> {
        Cli::try_parse_from(argv).map(|cli| {
            assert!(cli.command.is_none(), "no verb was given");
            cli.responder.expect("with no verb, the responder's flags")
        })
    }

    /// The argument vector `swarm-nats.nix` actually builds, in the form it
    /// builds it: every flag space-separated, and the suffix's value opening
    /// with `-`. Read as a short-flag cluster it aborts before `main` runs,
    /// and a responder that never starts denies every client — which reaches
    /// the operator as `authorization violation` on an unrelated process.
    #[test]
    fn the_units_own_argument_vector_parses() {
        let args = responder([
            "swarm-nats-auth",
            "--introspection-url",
            "https://auth.example/api/oidc/introspection",
            "--hive-client-prefix",
            "hive-",
            "--agent-client-suffix",
            "-agent",
            "--agent-token-publish-subject",
            "$SWARM.term.{agent}",
            "--agent-token-publish-subject",
            "$SWARM.agent-state.{agent}",
            "--agent-token-publish-subject",
            "$KV.agent-icons.{agent}",
            "--agent-token-publish-subject",
            "$SWARM.term.{agent}.sub.>",
            "--agent-token-publish-subject",
            "$JS.API.DIRECT.GET.KV_agent-schedules.$KV.agent-schedules.{agent}",
            "--agent-token-publish-subject",
            "$JS.API.DIRECT.GET.KV_agent-inference.$KV.agent-inference.{agent}",
            "--store-cert-role",
            "swarm-nats-auth",
        ])
        .expect("the unit's own argument vector must parse");
        assert_eq!(args.agent_client_suffix, "-agent");
        assert_eq!(args.agent_token_publish_subjects.len(), 6);
    }

    /// A request with no `Nats-Server-Xkey` is dropped before anything reads
    /// it, so it gets no reply. The same request sealed, with the header,
    /// opens.
    #[test]
    fn a_request_without_the_server_xkey_header_is_dropped() {
        let own = nkeys::XKey::new();
        let message = |headers, payload: Vec<u8>| async_nats::Message {
            subject: AUTH_SUBJECT.into(),
            reply: Some("_INBOX.reply".into()),
            payload: payload.into(),
            headers,
            status: None,
            description: None,
            length: 0,
        };
        let request = xkey::tests::request_jwt();
        assert!(open_request(&own, &message(None, request.clone())).is_none());

        let server = nkeys::XKey::new();
        let own_public = nkeys::XKey::from_public_key(&own.public_key()).expect("public");
        let sealed = server.seal(&request, &own_public).expect("sealed");
        let mut headers = async_nats::HeaderMap::new();
        headers.insert(xkey::SERVER_XKEY_HEADER, server.public_key());
        assert!(open_request(&own, &message(Some(headers), sealed)).is_some());
    }

    /// The control for the case above: an ordinary value parses through the
    /// same flag, so a failure there is about the leading `-` and not about
    /// the argument being unknown.
    #[test]
    fn an_agent_suffix_without_a_leading_dash_parses_too() {
        let args = responder([
            "swarm-nats-auth",
            "--introspection-url",
            "https://auth.example/api/oidc/introspection",
            "--agent-client-suffix",
            "agent",
        ])
        .expect("an ordinary suffix must parse");
        assert_eq!(args.agent_client_suffix, "agent");
    }

    /// The default is the value the deployment passes, so the two cannot
    /// drift into disagreeing about what an agent client is called.
    #[test]
    fn the_default_suffix_is_what_the_module_passes() {
        let args = responder([
            "swarm-nats-auth",
            "--introspection-url",
            "https://auth.example/api/oidc/introspection",
        ])
        .expect("the required flags alone must parse");
        assert_eq!(args.agent_client_suffix, "-agent");
    }

    /// The verb as `swarm-bao-nats-callout-xkey` runs it, with none of the
    /// responder's required flags.
    #[test]
    fn the_ensure_verb_parses_as_the_unit_runs_it() {
        let cli = Cli::try_parse_from([
            "swarm-nats-auth",
            "ensure-callout-keys",
            "--store-cert-role",
            "swarm-nats-auth",
            "--import-issuer-seed-file",
            "/var/lib/swarm-nats-callout/issuer.seed",
            "--import-user-seed-file",
            "/var/lib/swarm-nats-callout/callout-user.seed",
        ])
        .expect("the verb needs no responder flags");
        assert!(cli.responder.is_none());
        let Some(Command::EnsureCalloutKeys {
            store_cert_role,
            import_issuer_seed_file,
            import_user_seed_file,
        }) = cli.command
        else {
            panic!("the verb parses as itself");
        };
        assert_eq!(store_cert_role, "swarm-nats-auth");
        assert_eq!(
            import_issuer_seed_file.as_deref(),
            Some(std::path::Path::new(
                "/var/lib/swarm-nats-callout/issuer.seed"
            ))
        );
        assert_eq!(
            import_user_seed_file.as_deref(),
            Some(std::path::Path::new(
                "/var/lib/swarm-nats-callout/callout-user.seed"
            ))
        );
    }

    /// A bare invocation is neither: the responder's flags stay required
    /// when no verb is given.
    #[test]
    fn no_verb_and_no_responder_flags_is_refused() {
        Cli::try_parse_from(["swarm-nats-auth"]).expect_err("the responder's flags are required");
    }

    #[test]
    fn the_verb_does_not_take_the_responders_flags() {
        Cli::try_parse_from([
            "swarm-nats-auth",
            "--introspection-url",
            "https://auth.example/api/oidc/introspection",
            "ensure-callout-keys",
        ])
        .expect_err("a verb and the responder's flags conflict");
    }
}