stratumiops/crates/platform-nats/tests/test_nkey_auth.rs

/// Integration test: NKey authentication sign/verify round-trip.
///
/// Tests the full flow: key generation → signing → header injection → verification.
use async_nats::header::HeaderMap;
use bytes::Bytes;
use platform_nats::NKeyAuth;

/// Produce a `HeaderMap` with NKey authentication headers for the given payload.
fn signed_headers(seed: &str, payload: &[u8]) -> HeaderMap {
    let kp = nkeys::KeyPair::from_seed(seed).unwrap();
    let sig = kp.sign(payload).unwrap();
    let sig_b64 = b64url_encode(&sig);

    let pub_key = kp.public_key();
    let mut headers = HeaderMap::new();
    headers.insert("Nats-Nkey", pub_key.as_str());
    headers.insert("Nats-Signature", sig_b64.as_str());
    headers
}

/// Base64url encode (no padding, URL-safe alphabet: `-` and `_`).
///
/// Packs input into 3-byte chunks, encodes each as four 6-bit Base64 characters.
/// Partial trailing chunks emit 2 or 3 characters (no `=` padding).
fn b64url_encode(data: &[u8]) -> String {
    const CHARS: &[u8] = b"ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789-_";
    let mut out = String::with_capacity(data.len().div_ceil(3) * 4);
    for chunk in data.chunks(3) {
        let b0 = chunk[0] as u32;
        let b1 = chunk.get(1).copied().unwrap_or(0) as u32;
        let b2 = chunk.get(2).copied().unwrap_or(0) as u32;
        let n = (b0 << 16) | (b1 << 8) | b2;
        out.push(CHARS[((n >> 18) & 0x3F) as usize] as char);
        out.push(CHARS[((n >> 12) & 0x3F) as usize] as char);
        if chunk.len() > 1 {
            out.push(CHARS[((n >> 6) & 0x3F) as usize] as char);
        }
        if chunk.len() > 2 {
            out.push(CHARS[(n & 0x3F) as usize] as char);
        }
    }
    out
}

/// Generate a fresh user NKey seed for testing.
fn fresh_seed() -> String {
    let kp = nkeys::KeyPair::new_user();
    kp.seed().unwrap()
}

#[test]
fn test_trusted_key_signature_accepted() {
    let seed = fresh_seed();
    let kp = nkeys::KeyPair::from_seed(&seed).unwrap();
    let pub_key = kp.public_key();

    let auth = NKeyAuth::new(vec![pub_key]);
    let payload = Bytes::from_static(b"hello stratum");
    let headers = signed_headers(&seed, &payload);

    auth.verify_message(Some(&headers), &payload).unwrap();
}

#[test]
fn test_untrusted_key_rejected() {
    let trusted_seed = fresh_seed();
    let trusted_kp = nkeys::KeyPair::from_seed(&trusted_seed).unwrap();
    let auth = NKeyAuth::new(vec![trusted_kp.public_key()]);

    // Sign with a different (untrusted) key
    let attacker_seed = fresh_seed();
    let payload = Bytes::from_static(b"malicious payload");
    let headers = signed_headers(&attacker_seed, &payload);

    let err = auth.verify_message(Some(&headers), &payload).unwrap_err();
    assert!(
        err.to_string().contains("untrusted publisher"),
        "expected untrusted error, got: {err}"
    );
}

#[test]
fn test_tampered_payload_rejected() {
    let seed = fresh_seed();
    let kp = nkeys::KeyPair::from_seed(&seed).unwrap();
    let pub_key = kp.public_key();

    let auth = NKeyAuth::new(vec![pub_key]);
    let original_payload = Bytes::from_static(b"original payload");
    let headers = signed_headers(&seed, &original_payload);

    // Signature was for `original_payload`, but we verify against tampered one
    let tampered = Bytes::from_static(b"tampered payload!!!");
    let err = auth
        .verify_message(Some(&headers), &tampered)
        .unwrap_err();
    assert!(
        err.to_string().contains("signature verification failed")
            || err.to_string().contains("untrusted"),
        "expected signature error, got: {err}"
    );
}

#[test]
fn test_missing_headers_rejected() {
    let auth = NKeyAuth::new(vec!["some-key".to_string()]);
    let payload = Bytes::from_static(b"payload");

    let err = auth.verify_message(None, &payload).unwrap_err();
    assert!(
        err.to_string().contains("no headers"),
        "expected no-headers error, got: {err}"
    );
}

#[test]
fn test_missing_nkey_header_rejected() {
    let auth = NKeyAuth::new(vec!["some-key".to_string()]);
    let payload = Bytes::from_static(b"payload");

    let headers = HeaderMap::new(); // empty — no Nats-Nkey
    let err = auth
        .verify_message(Some(&headers), &payload)
        .unwrap_err();
    assert!(
        err.to_string().contains("missing Nats-Nkey"),
        "expected missing-header error, got: {err}"
    );
}

#[test]
fn test_auth_not_required_when_no_trusted_keys() {
    let auth = NKeyAuth::new(vec![]); // no trusted keys = auth not required
    assert!(
        !auth.is_auth_required(),
        "empty trusted set must mean auth is not required"
    );
}

#[test]
fn test_auth_required_when_trusted_keys_present() {
    let kp = nkeys::KeyPair::new_user();
    let auth = NKeyAuth::new(vec![kp.public_key()]);
    assert!(
        auth.is_auth_required(),
        "non-empty trusted set must mean auth is required"
    );
}