newton-core 0.4.16

//! Threshold HPKE decryption — RFC 9180 post-DH operations.
//!
//! In standard HPKE, `single_shot_open` bundles DH + KDF + AEAD. For threshold
//! decryption, operators compute partial DH outputs (D_i = s_i * enc) which the
//! gateway combines via Lagrange interpolation. The combined DH output is then
//! fed into [`decrypt_with_precomputed_dh`] to complete HPKE decryption.
//!
//! This module implements only the post-DH pipeline (ExtractAndExpand → KeySchedule
//! → AEAD.Open) for the X25519HkdfSha256 + HkdfSha256 + ChaCha20Poly1305
//! ciphersuite (KEM 0x0020, KDF 0x0001, AEAD 0x0003).
//!
//! The threshold DH operations (partial decryption, DLEQ proofs, Lagrange
//! interpolation) live in the [`dkg`](crate::dkg) module.
//!
//! Reference: RFC 9180 Sections 4, 5.1, and 5.2.

use chacha20poly1305::{aead::Aead, ChaCha20Poly1305, KeyInit as AeadKeyInit};
use hkdf::Hkdf;
use hmac::Mac;
use sha2::Sha256;

type HmacSha256 = hmac::Hmac<Sha256>;

use zeroize::Zeroizing;

use super::error::CryptoError;

// ---------------------------------------------------------------------------
// RFC 9180 constants for DHKEM(X25519, HKDF-SHA256) + HKDF-SHA256 +
// ChaCha20Poly1305
// ---------------------------------------------------------------------------

/// KEM suite ID: "KEM" || I2OSP(0x0020, 2)
const KEM_SUITE_ID: &[u8] = b"KEM\x00\x20";

/// HPKE suite ID: "HPKE" || I2OSP(kem_id=0x0020, 2) || I2OSP(kdf_id=0x0001, 2) || I2OSP(aead_id=0x0003, 2)
const HPKE_SUITE_ID: &[u8] = b"HPKE\x00\x20\x00\x01\x00\x03";

/// Version label for all labeled HPKE operations (RFC 9180 Section 4).
const VERSION_LABEL: &[u8] = b"HPKE-v1";

/// HKDF-SHA256 hash output length.
const NH: usize = 32;

/// X25519 shared secret size.
const N_SECRET: usize = 32;

/// ChaCha20Poly1305 key size.
const NK: usize = 32;

/// ChaCha20Poly1305 nonce size.
const NN: usize = 12;

/// Expected size of X25519 encapsulated key and public key.
const N_PK: usize = 32;

// ---------------------------------------------------------------------------
// HPKE Base mode constant
// ---------------------------------------------------------------------------

/// Base mode identifier (no PSK, no auth).
const MODE_BASE: u8 = 0x00;

// ---------------------------------------------------------------------------
// Core labeled HKDF operations (RFC 9180 Section 4)
// ---------------------------------------------------------------------------

/// Build `labeled_ikm = "HPKE-v1" || suite_id || label || ikm`.
fn build_labeled_ikm(suite_id: &[u8], label: &[u8], ikm: &[u8]) -> Vec<u8> {
    let mut buf = Vec::with_capacity(VERSION_LABEL.len() + suite_id.len() + label.len() + ikm.len());
    buf.extend_from_slice(VERSION_LABEL);
    buf.extend_from_slice(suite_id);
    buf.extend_from_slice(label);
    buf.extend_from_slice(ikm);
    buf
}

/// Build `labeled_info = I2OSP(L, 2) || "HPKE-v1" || suite_id || label || info`.
fn build_labeled_info(suite_id: &[u8], label: &[u8], info: &[u8], len: u16) -> Vec<u8> {
    let mut buf = Vec::with_capacity(2 + VERSION_LABEL.len() + suite_id.len() + label.len() + info.len());
    buf.extend_from_slice(&len.to_be_bytes());
    buf.extend_from_slice(VERSION_LABEL);
    buf.extend_from_slice(suite_id);
    buf.extend_from_slice(label);
    buf.extend_from_slice(info);
    buf
}

/// `LabeledExtract(salt, label, ikm)` → PRK bytes (RFC 9180 Section 4).
///
/// HKDF-Extract(salt, "HPKE-v1" || suite_id || label || ikm)
/// which is HMAC-SHA256(salt, labeled_ikm).
fn labeled_extract(suite_id: &[u8], salt: &[u8], label: &[u8], ikm: &[u8]) -> [u8; NH] {
    let labeled_ikm = build_labeled_ikm(suite_id, label, ikm);

    // HKDF-Extract(salt, IKM) = HMAC-Hash(salt, IKM)
    // Empty salt is replaced with a zero-filled string of length NH per RFC 5869.
    let effective_salt: &[u8] = if salt.is_empty() { &[0u8; NH] } else { salt };

    let mut mac = <HmacSha256 as Mac>::new_from_slice(effective_salt).expect("HMAC-SHA256 accepts any key length");
    Mac::update(&mut mac, &labeled_ikm);
    let result = mac.finalize().into_bytes();

    let mut prk = [0u8; NH];
    prk.copy_from_slice(&result);
    prk
}

/// `LabeledExpand(prk, label, info, L)` → output key material (RFC 9180 Section 4).
///
/// HKDF-Expand(prk, "I2OSP(L,2) || HPKE-v1 || suite_id || label || info", L)
fn labeled_expand(
    suite_id: &[u8],
    prk: &[u8; NH],
    label: &[u8],
    info: &[u8],
    len: usize,
) -> Result<Vec<u8>, CryptoError> {
    let labeled_info = build_labeled_info(suite_id, label, info, len as u16);

    let hkdf = Hkdf::<Sha256>::from_prk(prk)
        .map_err(|_| CryptoError::ThresholdDecrypt("invalid PRK for HKDF-Expand".into()))?;

    let mut okm = vec![0u8; len];
    hkdf.expand(&labeled_info, &mut okm)
        .map_err(|_| CryptoError::ThresholdDecrypt("HKDF-Expand failed".into()))?;
    Ok(okm)
}

// ---------------------------------------------------------------------------
// KEM-level: ExtractAndExpand (RFC 9180 Section 4.1)
// ---------------------------------------------------------------------------

/// `ExtractAndExpand(dh, kem_context)` → shared_secret.
///
/// Derives the HPKE shared secret from a raw DH output and KEM context.
/// For DHKEM(X25519, HKDF-SHA256), `kem_context = enc || pkR`.
fn extract_and_expand(dh: &[u8], kem_context: &[u8]) -> Result<[u8; N_SECRET], CryptoError> {
    // prk = LabeledExtract("", "eae_prk", dh) with KEM suite ID
    let prk = labeled_extract(KEM_SUITE_ID, &[], b"eae_prk", dh);

    // shared_secret = LabeledExpand(prk, "shared_secret", kem_context, Nsecret)
    let shared_secret_vec = labeled_expand(KEM_SUITE_ID, &prk, b"shared_secret", kem_context, N_SECRET)?;

    let mut shared_secret = [0u8; N_SECRET];
    shared_secret.copy_from_slice(&shared_secret_vec);
    Ok(shared_secret)
}

// ---------------------------------------------------------------------------
// HPKE-level: KeyScheduleR (RFC 9180 Section 5.1)
// ---------------------------------------------------------------------------

/// Derives the AEAD key and base nonce from a shared secret (Base mode, no PSK).
///
/// Returns `(key, base_nonce)` for ChaCha20Poly1305.
fn key_schedule_base(shared_secret: &[u8; N_SECRET], info: &[u8]) -> Result<([u8; NK], [u8; NN]), CryptoError> {
    // psk_id_hash = LabeledExtract("", "psk_id_hash", "") — Base mode: psk_id is empty
    let psk_id_hash = labeled_extract(HPKE_SUITE_ID, &[], b"psk_id_hash", &[]);

    // info_hash = LabeledExtract("", "info_hash", info)
    let info_hash = labeled_extract(HPKE_SUITE_ID, &[], b"info_hash", info);

    // ks_context = concat(mode, psk_id_hash, info_hash)
    let mut ks_context = Vec::with_capacity(1 + NH + NH);
    ks_context.push(MODE_BASE);
    ks_context.extend_from_slice(&psk_id_hash);
    ks_context.extend_from_slice(&info_hash);

    // secret = LabeledExtract(shared_secret, "secret", psk) — Base mode: psk is empty
    let secret = labeled_extract(HPKE_SUITE_ID, shared_secret, b"secret", &[]);

    // key = LabeledExpand(secret, "key", ks_context, Nk)
    let key_vec = labeled_expand(HPKE_SUITE_ID, &secret, b"key", &ks_context, NK)?;
    let mut key = [0u8; NK];
    key.copy_from_slice(&key_vec);

    // base_nonce = LabeledExpand(secret, "base_nonce", ks_context, Nn)
    let nonce_vec = labeled_expand(HPKE_SUITE_ID, &secret, b"base_nonce", &ks_context, NN)?;
    let mut base_nonce = [0u8; NN];
    base_nonce.copy_from_slice(&nonce_vec);

    Ok((key, base_nonce))
}

// ---------------------------------------------------------------------------
// Public API
// ---------------------------------------------------------------------------

/// Decrypt HPKE ciphertext using a precomputed DH shared secret.
///
/// This is the core building block for threshold HPKE decryption. In the
/// threshold scheme:
///
/// 1. Each operator computes a partial DH: `D_i = s_i * enc_point`
/// 2. The gateway combines partials via Lagrange interpolation to recover
///    `dh = sk * enc_point` (the full X25519 DH output)
/// 3. This function takes that `dh` output and completes HPKE decryption:
///    `ExtractAndExpand(dh) → KeySchedule → AEAD.Open`
///
/// For the equivalence test: `decrypt_with_precomputed_dh(X25519(sk, enc), enc, pk_R, ct, aad)`
/// produces the same plaintext as `hpke::decrypt(sk, enc, ct, aad)`.
///
/// # Arguments
///
/// * `dh` - Raw 32-byte X25519 Diffie-Hellman output
/// * `enc` - 32-byte encapsulated key (ephemeral sender public key)
/// * `pk_r` - 32-byte recipient public key
/// * `ciphertext` - AEAD ciphertext with appended Poly1305 tag (16 bytes)
/// * `aad` - Additional authenticated data (must match what was used during encryption)
pub fn decrypt_with_precomputed_dh(
    dh: &[u8],
    enc: &[u8],
    pk_r: &[u8],
    ciphertext: &[u8],
    aad: &[u8],
) -> Result<Zeroizing<Vec<u8>>, CryptoError> {
    if enc.len() != N_PK {
        return Err(CryptoError::ThresholdDecrypt(format!(
            "enc must be {} bytes, got {}",
            N_PK,
            enc.len()
        )));
    }
    if pk_r.len() != N_PK {
        return Err(CryptoError::ThresholdDecrypt(format!(
            "pk_r must be {} bytes, got {}",
            N_PK,
            pk_r.len()
        )));
    }

    // kem_context = enc || pkR (RFC 9180 Section 4.1, Base mode)
    let mut kem_context = Vec::with_capacity(N_PK + N_PK);
    kem_context.extend_from_slice(enc);
    kem_context.extend_from_slice(pk_r);

    // Step 1: ExtractAndExpand(dh, kem_context) → shared_secret
    let shared_secret = extract_and_expand(dh, &kem_context)?;

    // Step 2: KeyScheduleR(shared_secret, info="") → (key, base_nonce)
    // info is empty in our usage (matches single_shot_open(&OpModeR::Base, ..., &[], ...))
    let (key, base_nonce) = key_schedule_base(&shared_secret, &[])?;

    // Step 3: AEAD.Open(key, nonce=base_nonce XOR 0, aad, ciphertext)
    // For single-shot (seq=0), nonce = base_nonce XOR encode_big_endian(0) = base_nonce
    let cipher = <ChaCha20Poly1305 as AeadKeyInit>::new(chacha20poly1305::Key::from_slice(&key));
    let nonce = chacha20poly1305::Nonce::from_slice(&base_nonce);

    let plaintext = cipher
        .decrypt(nonce, chacha20poly1305::aead::Payload { msg: ciphertext, aad })
        .map_err(|e| CryptoError::ThresholdDecrypt(format!("AEAD decryption failed: {e}")))?;

    Ok(Zeroizing::new(plaintext))
}

// ---------------------------------------------------------------------------
// Tests
// ---------------------------------------------------------------------------

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

    /// Compute raw X25519 DH: sk * pk → 32-byte shared point.
    fn x25519_dh(sk_bytes: &[u8], pk_bytes: &[u8]) -> [u8; 32] {
        let sk = x25519_dalek::StaticSecret::from(<[u8; 32]>::try_from(sk_bytes).unwrap());
        let pk = x25519_dalek::PublicKey::from(<[u8; 32]>::try_from(pk_bytes).unwrap());
        *sk.diffie_hellman(&pk).as_bytes()
    }

    #[test]
    fn equivalence_with_standard_hpke_decrypt() {
        // Generate a keypair and encrypt some data using standard HPKE
        let (sk, pk) = hpke::generate_keypair();
        let plaintext = b"threshold decryption equivalence test";
        let aad = b"newton-privacy-aad";

        let (enc, ct) = hpke::encrypt(&pk, plaintext, aad).expect("encrypt failed");

        // Compute DH(sk, enc) manually using x25519-dalek
        let dh = x25519_dh(sk.to_bytes(), &enc);

        // Decrypt using our precomputed DH function
        let recovered =
            decrypt_with_precomputed_dh(&dh, &enc, pk.to_bytes(), &ct, aad).expect("threshold decrypt failed");

        assert_eq!(&*recovered, plaintext, "threshold decrypt must match standard decrypt");

        // Also verify standard decrypt produces the same result
        let standard_recovered = hpke::decrypt(&sk, &enc, &ct, aad).expect("standard decrypt failed");
        assert_eq!(&*recovered, &*standard_recovered);
    }

    #[test]
    fn equivalence_with_empty_plaintext() {
        let (sk, pk) = hpke::generate_keypair();
        let plaintext = b"";
        let aad = b"empty-test";

        let (enc, ct) = hpke::encrypt(&pk, plaintext, aad).expect("encrypt failed");
        let dh = x25519_dh(sk.to_bytes(), &enc);

        let recovered =
            decrypt_with_precomputed_dh(&dh, &enc, pk.to_bytes(), &ct, aad).expect("threshold decrypt failed");

        assert_eq!(&*recovered, plaintext);
    }

    #[test]
    fn equivalence_with_empty_aad() {
        let (sk, pk) = hpke::generate_keypair();
        let plaintext = b"test data with no aad";
        let aad = b"";

        let (enc, ct) = hpke::encrypt(&pk, plaintext, aad).expect("encrypt failed");
        let dh = x25519_dh(sk.to_bytes(), &enc);

        let recovered =
            decrypt_with_precomputed_dh(&dh, &enc, pk.to_bytes(), &ct, aad).expect("threshold decrypt failed");

        assert_eq!(&*recovered, plaintext);
    }

    #[test]
    fn equivalence_with_large_plaintext() {
        let (sk, pk) = hpke::generate_keypair();
        let plaintext = vec![0xABu8; 8192]; // 8KB payload
        let aad = b"large-payload-test";

        let (enc, ct) = hpke::encrypt(&pk, &plaintext, aad).expect("encrypt failed");
        let dh = x25519_dh(sk.to_bytes(), &enc);

        let recovered =
            decrypt_with_precomputed_dh(&dh, &enc, pk.to_bytes(), &ct, aad).expect("threshold decrypt failed");

        assert_eq!(recovered[..], plaintext[..]);
    }

    #[test]
    fn wrong_dh_output_fails() {
        let (sk, pk) = hpke::generate_keypair();
        let plaintext = b"should fail with wrong DH";
        let aad = b"ctx";

        let (enc, ct) = hpke::encrypt(&pk, plaintext, aad).expect("encrypt failed");

        // Use a different keypair's DH — should fail AEAD auth check
        let (wrong_sk, _) = hpke::generate_keypair();
        let wrong_dh = x25519_dh(wrong_sk.to_bytes(), &enc);

        let result = decrypt_with_precomputed_dh(&wrong_dh, &enc, pk.to_bytes(), &ct, aad);
        assert!(result.is_err(), "wrong DH output must cause decryption failure");
    }

    #[test]
    fn wrong_aad_fails() {
        let (sk, pk) = hpke::generate_keypair();
        let plaintext = b"should fail with wrong AAD";
        let aad = b"correct";

        let (enc, ct) = hpke::encrypt(&pk, plaintext, aad).expect("encrypt failed");
        let dh = x25519_dh(sk.to_bytes(), &enc);

        let result = decrypt_with_precomputed_dh(&dh, &enc, pk.to_bytes(), &ct, b"wrong");
        assert!(result.is_err(), "wrong AAD must cause decryption failure");
    }

    #[test]
    fn wrong_pk_r_fails() {
        let (sk, pk) = hpke::generate_keypair();
        let plaintext = b"should fail with wrong pkR";
        let aad = b"ctx";

        let (enc, ct) = hpke::encrypt(&pk, plaintext, aad).expect("encrypt failed");
        let dh = x25519_dh(sk.to_bytes(), &enc);

        // Use a different public key for kem_context — changes the shared_secret derivation
        let (_, wrong_pk) = hpke::generate_keypair();
        let result = decrypt_with_precomputed_dh(&dh, &enc, wrong_pk.to_bytes(), &ct, aad);
        assert!(result.is_err(), "wrong pkR must cause decryption failure");
    }

    #[test]
    fn invalid_enc_length_rejected() {
        let result = decrypt_with_precomputed_dh(&[0u8; 32], &[0u8; 16], &[0u8; 32], &[0u8; 32], &[]);
        assert!(result.is_err());
        let err_msg = result.unwrap_err().to_string();
        assert!(err_msg.contains("enc must be 32 bytes"));
    }

    #[test]
    fn invalid_pk_r_length_rejected() {
        let result = decrypt_with_precomputed_dh(&[0u8; 32], &[0u8; 32], &[0u8; 16], &[0u8; 32], &[]);
        assert!(result.is_err());
        let err_msg = result.unwrap_err().to_string();
        assert!(err_msg.contains("pk_r must be 32 bytes"));
    }

    #[test]
    fn multiple_encryptions_same_key_all_decrypt() {
        let (sk, pk) = hpke::generate_keypair();
        let aad = b"multi-msg";

        for i in 0..5 {
            let plaintext = format!("message number {i}");
            let (enc, ct) = hpke::encrypt(&pk, plaintext.as_bytes(), aad).expect("encrypt failed");
            let dh = x25519_dh(sk.to_bytes(), &enc);

            let recovered =
                decrypt_with_precomputed_dh(&dh, &enc, pk.to_bytes(), &ct, aad).expect("threshold decrypt failed");
            assert_eq!(&*recovered, plaintext.as_bytes());
        }
    }
}