newton-enclave 0.4.13

use newton_core::crypto::HpkePrivateKey;
use zeroize::Zeroizing;

use crate::error::EnclaveError;

const HPKE_SECRETS_INFO: &[u8] = b"newton-hpke-secrets-v1";

/// Derives an HPKE private key deterministically from a 32-byte seed via HKDF-SHA256.
///
/// All enclaves with access to the same KMS seed produce the same keypair.
/// The seed is zeroized immediately after derivation.
pub fn derive_hpke_from_seed(seed: &Zeroizing<[u8; 32]>) -> Result<HpkePrivateKey, EnclaveError> {
    let hkdf = hkdf::Hkdf::<sha2::Sha256>::new(None, seed.as_ref());
    let mut okm = Zeroizing::new([0u8; 32]);
    hkdf.expand(HPKE_SECRETS_INFO, okm.as_mut())
        .map_err(|_| EnclaveError::KeyDerivation("HKDF expand failed".to_string()))?;
    HpkePrivateKey::from_bytes(okm.as_ref())
        .map_err(|e| EnclaveError::KeyDerivation(format!("HPKE key from HKDF output: {e}")))
}

/// Provider for KMS-attested HPKE key derivation.
///
/// Production: `NitroKmsKeyProvider` (in the enclave binary crate, Linux-only)
/// calls NSM for attestation, KMS Decrypt to get the seed, then `derive_hpke_from_seed`.
///
/// Testing: `MockKmsKeyProvider` uses a hardcoded seed directly.
pub trait KmsKeyProvider: Send + Sync {
    /// derive the HPKE keypair from the provider's key material.
    fn derive_hpke_keypair(&self) -> Result<HpkePrivateKey, EnclaveError>;
}

/// Mock provider for LoopbackEnclave and tests.
/// Derives from a fixed seed without KMS or attestation.
#[derive(Debug)]
pub struct MockKmsKeyProvider {
    seed: Zeroizing<[u8; 32]>,
}

impl MockKmsKeyProvider {
    /// create a mock provider from a raw 32-byte seed.
    pub fn new(seed: [u8; 32]) -> Self {
        Self {
            seed: Zeroizing::new(seed),
        }
    }
}

impl KmsKeyProvider for MockKmsKeyProvider {
    fn derive_hpke_keypair(&self) -> Result<HpkePrivateKey, EnclaveError> {
        derive_hpke_from_seed(&self.seed)
    }
}

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

    #[test]
    fn deterministic_derivation_from_same_seed() {
        let seed = [0xAB; 32];
        let provider1 = MockKmsKeyProvider::new(seed);
        let provider2 = MockKmsKeyProvider::new(seed);

        let key1 = provider1.derive_hpke_keypair().unwrap();
        let key2 = provider2.derive_hpke_keypair().unwrap();
        assert_eq!(key1.to_bytes(), key2.to_bytes());
    }

    #[test]
    fn different_seeds_produce_different_keys() {
        let provider1 = MockKmsKeyProvider::new([0xAA; 32]);
        let provider2 = MockKmsKeyProvider::new([0xBB; 32]);

        let key1 = provider1.derive_hpke_keypair().unwrap();
        let key2 = provider2.derive_hpke_keypair().unwrap();
        assert_ne!(key1.to_bytes(), key2.to_bytes());
    }

    #[test]
    fn public_key_derivable_from_result() {
        let provider = MockKmsKeyProvider::new([0xCC; 32]);
        let sk = provider.derive_hpke_keypair().unwrap();
        let pk = sk.to_public_key();
        assert_eq!(pk.to_bytes().len(), 32);
    }
}