//! The callout xkey: the curve keypair `nats-server` seals every auth request
//! to, and that this responder seals every reply back with.
//!
//! Both halves live in the secret store, and only the `swarm-nats-auth` role
//! reads either. The seed is at [`SEED_PATH`]; the public half is at
//! [`PUBLIC_PATH`], and the queue's host renders it into `auth_callout.xkey`
//! from there.
//!
//! [`ensure`] creates the pair when the seed is absent and otherwise only
//! reads it. [`load`] only ever reads, and is what the responder uses: a
//! responder that minted a new pair on its own start would hold a key the
//! running server does not know, and every request would fail to open.

use std::future::Future;

use anyhow::{Context, Result, anyhow};
use nkeys::XKey;
use swarm_secret_client::SecretStore;

/// The seed's path under the store's KV mount. The nix half is
/// `xkeySeedLeaf` in `nix/host-modules/swarm-bao.nix`.
pub const SEED_PATH: &str = "swarm/services/swarm-nats/callout/xkey-seed";

/// The public half's path under the store's KV mount. The nix half is
/// `xkeyPublicLeaf` in `nix/host-modules/swarm-bao.nix`, and the fetch in
/// `nix/host-modules/swarm-nats.nix` reads it.
pub const PUBLIC_PATH: &str = "swarm/public/swarm-nats/callout-xkey";

/// The header `nats-server` sets on a request it sealed. Its value is the
/// server's own public xkey, which the reply is sealed to.
pub const SERVER_XKEY_HEADER: &str = "Nats-Server-Xkey";

/// The two operations [`ensure`] and [`load`] need from the store, so a test
/// can stand in for it.
pub trait KeyStore {
    /// The value at `path`, or `None` when nothing is stored there.
    fn read(&self, path: &str) -> impl Future<Output = Result<Option<String>>> + Send;
    /// Store `value` at `path`, failing when anything is already there.
    fn create(&self, path: &str, value: &str) -> impl Future<Output = Result<()>> + Send;
}

/// One field, `value`, as every single-value secret in the store is kept.
/// No `Debug`: one of the values held here is the seed.
#[derive(serde::Serialize, serde::Deserialize)]
struct Stored {
    value: String,
}

impl KeyStore for SecretStore {
    async fn read(&self, path: &str) -> Result<Option<String>> {
        let stored: Option<Stored> = self
            .read_optional(path)
            .await
            .with_context(|| format!("reading {path}"))?;
        Ok(stored.map(|s| s.value))
    }

    async fn create(&self, path: &str, value: &str) -> Result<()> {
        SecretStore::create(
            self,
            path,
            &Stored {
                value: value.to_owned(),
            },
        )
        .await
        .with_context(|| format!("creating {path}"))
    }
}

/// The stored pair, created first when no seed is stored.
///
/// Safe to run from two places at once. Each writes with a create that fails
/// once the path holds anything, so one seed lands, and the run whose write
/// failed reads that one back. The public half is written the same way, and
/// is written even when the seed already exists, so a run that stopped
/// between the two writes is completed by the next.
///
/// # Errors
/// When the store cannot be read or written, when the stored seed is not a
/// curve seed, and when the stored public half is not the seed's.
pub async fn ensure(store: &impl KeyStore) -> Result<XKey> {
    let key = match store.read(SEED_PATH).await? {
        Some(seed) => from_seed(&seed)?,
        None => create_seed(store).await?,
    };
    publish(store, &key.public_key()).await?;
    Ok(key)
}

/// The stored pair. Never writes.
///
/// # Errors
/// When the store cannot be read, holds no seed, or holds one that is not a
/// curve seed.
pub async fn load(store: &impl KeyStore) -> Result<XKey> {
    let seed = store.read(SEED_PATH).await?.with_context(|| {
        format!("nothing is stored at {SEED_PATH}; swarm-bao-nats-callout-xkey creates it")
    })?;
    from_seed(&seed)
}

async fn create_seed(store: &impl KeyStore) -> Result<XKey> {
    let fresh = XKey::new();
    let seed = fresh
        .seed()
        .map_err(|_| anyhow!("a freshly generated xkey has no seed"))?;
    match store.create(SEED_PATH, &seed).await {
        Ok(()) => {
            tracing::info!(public_key = %fresh.public_key(), "created the callout xkey");
            Ok(fresh)
        }
        Err(e) => match store.read(SEED_PATH).await? {
            Some(theirs) => from_seed(&theirs),
            None => Err(e),
        },
    }
}

async fn publish(store: &impl KeyStore, public: &str) -> Result<()> {
    let stored = match store.read(PUBLIC_PATH).await? {
        Some(stored) => stored,
        None => match store.create(PUBLIC_PATH, public).await {
            Ok(()) => return Ok(()),
            Err(e) => store.read(PUBLIC_PATH).await?.ok_or(e)?,
        },
    };
    anyhow::ensure!(
        stored == public,
        "{PUBLIC_PATH} holds a public key that is not the one for the seed at {SEED_PATH}"
    );
    Ok(())
}

/// Never carries the parse error: it is about the bytes of the seed.
fn from_seed(seed: &str) -> Result<XKey> {
    XKey::from_seed(seed.trim()).map_err(|_| anyhow!("{SEED_PATH} does not hold a curve seed"))
}

/// `payload` opened with `own`, and the server key `server_key` names, which
/// the reply is sealed to.
///
/// # Errors
/// When there is no [`SERVER_XKEY_HEADER`], when it is not a public xkey, and
/// when the payload does not open with `own`.
pub fn open(own: &XKey, server_key: Option<&str>, payload: &[u8]) -> Result<(XKey, Vec<u8>)> {
    let server_key = server_key
        .with_context(|| format!("the request is not sealed: no {SERVER_XKEY_HEADER}"))?;
    let server = XKey::from_public_key(server_key)
        .map_err(|_| anyhow!("{SERVER_XKEY_HEADER} is not a public xkey"))?;
    let opened = own
        .open(payload, &server)
        .map_err(|_| anyhow!("the request does not open with this responder's xkey"))?;
    Ok((server, opened))
}

/// `reply` sealed by `own` to `server`.
///
/// # Errors
/// When sealing fails, which it does only for a key with no seed.
pub fn seal(own: &XKey, server: &XKey, reply: &[u8]) -> Result<Vec<u8>> {
    own.seal(reply, server)
        .map_err(|_| anyhow!("sealing the reply"))
}

#[cfg(test)]
pub(crate) mod tests {
    use std::collections::HashMap;
    use std::sync::Mutex;
    use std::sync::atomic::{AtomicUsize, Ordering};

    use super::*;

    /// A store that keeps values in memory and counts writes. `rival` is a
    /// value another start lands at a path between this one's read and its
    /// create, which the create then loses to.
    #[derive(Default)]
    struct Fake {
        values: Mutex<HashMap<String, String>>,
        rival: Mutex<HashMap<String, String>>,
        writes: AtomicUsize,
    }

    impl Fake {
        fn holding(entries: &[(&str, &str)]) -> Self {
            let fake = Self::default();
            for (path, value) in entries {
                fake.put(path, value);
            }
            fake
        }

        fn put(&self, path: &str, value: &str) {
            self.values
                .lock()
                .expect("unpoisoned")
                .insert(path.to_owned(), value.to_owned());
        }

        fn get(&self, path: &str) -> Option<String> {
            self.values.lock().expect("unpoisoned").get(path).cloned()
        }
    }

    impl KeyStore for Fake {
        async fn read(&self, path: &str) -> Result<Option<String>> {
            Ok(self.get(path))
        }

        async fn create(&self, path: &str, value: &str) -> Result<()> {
            if let Some(theirs) = self.rival.lock().expect("unpoisoned").remove(path) {
                self.put(path, &theirs);
            }
            if self.get(path).is_some() {
                anyhow::bail!("403 permission denied");
            }
            self.writes.fetch_add(1, Ordering::SeqCst);
            self.put(path, value);
            Ok(())
        }
    }

    fn seed_of(key: &XKey) -> String {
        key.seed().expect("a generated key has a seed")
    }

    #[tokio::test]
    async fn an_empty_store_gets_a_seed_and_its_public_half() {
        let store = Fake::default();
        let key = ensure(&store).await.expect("ensured");
        assert_eq!(store.get(SEED_PATH), Some(seed_of(&key)));
        assert_eq!(store.get(PUBLIC_PATH), Some(key.public_key()));
        assert!(key.public_key().starts_with('X'), "a curve public key");
        assert_eq!(store.writes.load(Ordering::SeqCst), 2);
    }

    #[tokio::test]
    async fn an_existing_pair_is_only_read() {
        let existing = XKey::new();
        let store = Fake::holding(&[
            (SEED_PATH, &seed_of(&existing)),
            (PUBLIC_PATH, &existing.public_key()),
        ]);
        let key = ensure(&store).await.expect("ensured");
        assert_eq!(key.public_key(), existing.public_key());
        assert_eq!(store.writes.load(Ordering::SeqCst), 0);
    }

    /// A run that stopped between its two writes left a seed and no public
    /// half; the next one writes the public half for that seed.
    #[tokio::test]
    async fn an_existing_seed_without_its_public_half_gets_one() {
        let existing = XKey::new();
        let store = Fake::holding(&[(SEED_PATH, &seed_of(&existing))]);
        let key = ensure(&store).await.expect("ensured");
        assert_eq!(key.public_key(), existing.public_key());
        assert_eq!(store.get(PUBLIC_PATH), Some(existing.public_key()));
        assert_eq!(store.writes.load(Ordering::SeqCst), 1);
    }

    /// Two starts read an empty store at once and both generate; the one
    /// whose create fails takes the seed that landed.
    #[tokio::test]
    async fn a_start_that_loses_the_create_race_takes_the_winners_seed() {
        let winner = XKey::new();
        let store = Fake::default();
        store
            .rival
            .lock()
            .expect("unpoisoned")
            .insert(SEED_PATH.to_owned(), seed_of(&winner));
        store
            .rival
            .lock()
            .expect("unpoisoned")
            .insert(PUBLIC_PATH.to_owned(), winner.public_key());
        let key = ensure(&store).await.expect("ensured");
        assert_eq!(key.public_key(), winner.public_key());
        assert_eq!(store.get(SEED_PATH), Some(seed_of(&winner)));
        assert_eq!(store.get(PUBLIC_PATH), Some(winner.public_key()));
        assert_eq!(store.writes.load(Ordering::SeqCst), 0);
    }

    #[tokio::test]
    async fn a_public_half_that_is_not_the_seeds_is_refused() {
        let existing = XKey::new();
        let store = Fake::holding(&[
            (SEED_PATH, &seed_of(&existing)),
            (PUBLIC_PATH, &XKey::new().public_key()),
        ]);
        assert!(ensure(&store).await.is_err());
    }

    #[tokio::test]
    async fn a_seed_that_is_not_a_curve_seed_is_refused_without_quoting_it() {
        let account = nkeys::KeyPair::new_account();
        let not_curve = account.seed().expect("a seed");
        let store = Fake::holding(&[(SEED_PATH, &not_curve)]);
        let e = ensure(&store).await.expect_err("an account seed");
        assert!(!format!("{e:#}").contains(&not_curve), "{e:#}");
        let e = load(&store).await.expect_err("an account seed");
        assert!(!format!("{e:#}").contains(&not_curve), "{e:#}");
    }

    #[tokio::test]
    async fn loading_reads_the_stored_seed_and_never_creates_one() {
        let store = Fake::default();
        assert!(load(&store).await.is_err());
        assert_eq!(store.writes.load(Ordering::SeqCst), 0);
        assert_eq!(store.get(SEED_PATH), None);

        let existing = XKey::new();
        store.put(SEED_PATH, &seed_of(&existing));
        let key = load(&store).await.expect("loaded");
        assert_eq!(key.public_key(), existing.public_key());
    }

    pub(crate) fn request_jwt() -> Vec<u8> {
        let b = |s: &str| data_encoding::BASE64URL_NOPAD.encode(s.as_bytes());
        format!(
            "{}.{}.{}",
            b("{}"),
            b(r#"{"nats":{"connect_opts":{"auth_token":"t0k"},"user_nkey":"UCLIENT","server_id":{"id":"NSERVER"}}}"#),
            b("sig"),
        )
        .into_bytes()
    }

    /// The exchange as the server runs it: it seals the request to our public
    /// key under a key of its own, which it names in the header, and opens
    /// our reply with that same key.
    #[test]
    fn a_sealed_request_opens_and_its_grant_is_sealed_back_to_the_server() {
        let own = XKey::new();
        let server = XKey::new();
        let own_public = XKey::from_public_key(&own.public_key()).expect("public");
        let sealed = server.seal(&request_jwt(), &own_public).expect("sealed");

        let (server_key, opened) = open(&own, Some(&server.public_key()), &sealed).expect("opens");
        assert_eq!(server_key.public_key(), server.public_key());
        let req = crate::request::decode(&opened).expect("a request");
        assert_eq!(req.user_nkey, "UCLIENT");

        let issuer = nkeys::KeyPair::new_account();
        let permissions = crate::policy::Permissions {
            publish: vec!["$KV.hive-status.alpha".to_owned()],
            subscribe: None,
        };
        let token = crate::respond::grant(
            &issuer,
            "APP",
            &req.server_id.id,
            &req.user_nkey,
            &permissions,
        );
        let reply = seal(&own, &server_key, token.as_bytes()).expect("sealed");
        assert_ne!(reply, token.as_bytes(), "the reply leaves sealed");
        let read_by_server = server
            .open(&reply, &own_public)
            .expect("the server opens it");
        assert_eq!(read_by_server, token.as_bytes());
    }

    #[test]
    fn a_request_without_the_header_is_refused() {
        let own = XKey::new();
        let e = open(&own, None, &request_jwt()).expect_err("no header");
        assert!(format!("{e:#}").contains(SERVER_XKEY_HEADER), "{e:#}");
    }

    #[test]
    fn a_request_sealed_to_another_key_does_not_open() {
        let own = XKey::new();
        let server = XKey::new();
        let someone_else = XKey::from_public_key(&XKey::new().public_key()).expect("public");
        let sealed = server.seal(&request_jwt(), &someone_else).expect("sealed");
        assert!(open(&own, Some(&server.public_key()), &sealed).is_err());
        // The header present over a payload that was never sealed.
        assert!(open(&own, Some(&server.public_key()), &request_jwt()).is_err());
        assert!(open(&own, Some("not-a-key"), &sealed).is_err());
    }
}