newton-core 0.4.16

//! HPKE encryption and decryption (RFC 9180).
//!
//! Uses X25519 KEM + HKDF-SHA256 + ChaCha20-Poly1305 for hybrid public key
//! encryption. Provides wrapper types for keys and convenience functions for
//! single-shot encrypt/decrypt operations.

use ::hpke::{
    aead::ChaCha20Poly1305 as AeadImpl, kdf::HkdfSha256 as KdfImpl, kem::X25519HkdfSha256 as KemImpl, Deserializable,
    Kem as KemTrait, OpModeR, OpModeS, Serializable,
};

use rand_core::OsRng;
use zeroize::Zeroizing;

use super::error::CryptoError;

// ---------------------------------------------------------------------------
// Type aliases for the chosen ciphersuite
// ---------------------------------------------------------------------------

type Kem = KemImpl;
type Kdf = KdfImpl;
type Aead = AeadImpl;

// ---------------------------------------------------------------------------
// Wrapper types
// ---------------------------------------------------------------------------

/// An HPKE public key (32 bytes for X25519).
#[derive(Debug, Clone)]
pub struct HpkePublicKey(Vec<u8>);

impl HpkePublicKey {
    /// Construct an `HpkePublicKey` from raw bytes.
    ///
    /// The slice must be exactly 32 bytes for X25519.
    pub fn from_bytes(bytes: &[u8]) -> Result<Self, CryptoError> {
        // Validate by attempting to deserialize through the HPKE crate
        <Kem as KemTrait>::PublicKey::from_bytes(bytes).map_err(|e| CryptoError::InvalidPublicKey {
            reason: format!("failed to deserialize X25519 public key: {e}"),
        })?;
        Ok(Self(bytes.to_vec()))
    }

    /// Return the raw key bytes.
    pub fn to_bytes(&self) -> &[u8] {
        &self.0
    }
}

/// An HPKE private key (32 bytes for X25519).
///
/// Inner bytes are zeroized on drop to prevent key material from lingering in memory.
#[derive(Debug, Clone)]
pub struct HpkePrivateKey(Zeroizing<Vec<u8>>);

impl HpkePrivateKey {
    /// Construct an `HpkePrivateKey` from raw bytes.
    ///
    /// The slice must be exactly 32 bytes for X25519.
    pub fn from_bytes(bytes: &[u8]) -> Result<Self, CryptoError> {
        <Kem as KemTrait>::PrivateKey::from_bytes(bytes).map_err(|e| CryptoError::InvalidPrivateKey {
            reason: format!("failed to deserialize X25519 private key: {e}"),
        })?;
        Ok(Self(Zeroizing::new(bytes.to_vec())))
    }

    /// Return the raw key bytes.
    pub fn to_bytes(&self) -> &[u8] {
        &self.0
    }

    /// Derive the corresponding X25519 public key.
    ///
    /// Safe to unwrap because `from_bytes` validates the key is exactly 32
    /// bytes at construction time.
    pub fn to_public_key(&self) -> HpkePublicKey {
        let sk_bytes: [u8; 32] = self.0[..].try_into().expect("X25519 key is always 32 bytes");
        let sk = x25519_dalek::StaticSecret::from(sk_bytes);
        let pk = x25519_dalek::PublicKey::from(&sk);
        HpkePublicKey(pk.as_bytes().to_vec())
    }
}

// ---------------------------------------------------------------------------
// Key generation
// ---------------------------------------------------------------------------

/// Generate a fresh X25519 HPKE keypair.
pub fn generate_keypair() -> (HpkePrivateKey, HpkePublicKey) {
    let mut rng = OsRng;
    let (sk, pk) = Kem::gen_keypair(&mut rng);
    let sk_bytes = Zeroizing::new(sk.to_bytes().to_vec());
    let pk_bytes = pk.to_bytes().to_vec();
    (HpkePrivateKey(sk_bytes), HpkePublicKey(pk_bytes))
}

// ---------------------------------------------------------------------------
// Encrypt / Decrypt (single-shot, allocating)
// ---------------------------------------------------------------------------

/// Encrypt `plaintext` to `recipient_pk` with associated data `aad`.
///
/// Returns `(encapped_key_bytes, ciphertext)` where `ciphertext` includes
/// the ChaCha20-Poly1305 authentication tag appended by the HPKE crate.
pub fn encrypt(recipient_pk: &HpkePublicKey, plaintext: &[u8], aad: &[u8]) -> Result<(Vec<u8>, Vec<u8>), CryptoError> {
    let pk = <Kem as KemTrait>::PublicKey::from_bytes(recipient_pk.to_bytes())
        .map_err(|e| CryptoError::InvalidPublicKey { reason: format!("{e}") })?;

    let mut rng = OsRng;
    let (encapped_key, ciphertext) =
        ::hpke::single_shot_seal::<Aead, Kdf, Kem, _>(&OpModeS::Base, &pk, &[], plaintext, aad, &mut rng)
            .map_err(|e| CryptoError::HpkeEncrypt(format!("{e}")))?;

    Ok((encapped_key.to_bytes().to_vec(), ciphertext))
}

/// Decrypt `ciphertext` (which includes the Poly1305 tag) using the
/// recipient's private key and the encapsulated key produced during
/// encryption.
pub fn decrypt(
    recipient_sk: &HpkePrivateKey,
    encapped_key: &[u8],
    ciphertext: &[u8],
    aad: &[u8],
) -> Result<Zeroizing<Vec<u8>>, CryptoError> {
    let sk = <Kem as KemTrait>::PrivateKey::from_bytes(recipient_sk.to_bytes())
        .map_err(|e| CryptoError::InvalidPrivateKey { reason: format!("{e}") })?;

    let enc = <Kem as KemTrait>::EncappedKey::from_bytes(encapped_key)
        .map_err(|e| CryptoError::HpkeDecrypt(format!("invalid encapped key: {e}")))?;

    let plaintext = ::hpke::single_shot_open::<Aead, Kdf, Kem>(&OpModeR::Base, &sk, &enc, &[], ciphertext, aad)
        .map_err(|e| CryptoError::HpkeDecrypt(format!("{e}")))?;

    Ok(Zeroizing::new(plaintext))
}

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

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

    #[test]
    fn encrypt_decrypt_roundtrip() {
        let (sk, pk) = generate_keypair();
        let plaintext = b"hello, newton privacy layer";
        let aad = b"task-context";

        let (enc, ct) = encrypt(&pk, plaintext, aad).expect("encrypt failed");
        let recovered = decrypt(&sk, &enc, &ct, aad).expect("decrypt failed");
        assert_eq!(&*recovered, plaintext);
    }

    #[test]
    fn decrypt_with_wrong_key_fails() {
        let (_sk, pk) = generate_keypair();
        let (wrong_sk, _wrong_pk) = generate_keypair();
        let plaintext = b"secret data";
        let aad = b"ctx";

        let (enc, ct) = encrypt(&pk, plaintext, aad).expect("encrypt failed");
        let result = decrypt(&wrong_sk, &enc, &ct, aad);
        assert!(result.is_err(), "decryption with wrong key should fail");
    }

    #[test]
    fn decrypt_with_wrong_aad_fails() {
        let (sk, pk) = generate_keypair();
        let plaintext = b"secret data";
        let aad = b"correct-aad";

        let (enc, ct) = encrypt(&pk, plaintext, aad).expect("encrypt failed");
        let result = decrypt(&sk, &enc, &ct, b"wrong-aad");
        assert!(result.is_err(), "decryption with wrong AAD should fail");
    }

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

        let (enc, ct) = encrypt(&pk, plaintext, aad).expect("encrypt failed");
        let recovered = decrypt(&sk, &enc, &ct, aad).expect("decrypt failed");
        assert_eq!(&*recovered, plaintext);
    }

    #[test]
    fn key_from_bytes_roundtrip() {
        let (sk, pk) = generate_keypair();
        let sk2 = HpkePrivateKey::from_bytes(sk.to_bytes()).expect("sk from_bytes failed");
        let pk2 = HpkePublicKey::from_bytes(pk.to_bytes()).expect("pk from_bytes failed");
        assert_eq!(sk.to_bytes(), sk2.to_bytes());
        assert_eq!(pk.to_bytes(), pk2.to_bytes());
    }

    #[test]
    fn invalid_key_bytes_rejected() {
        let result = HpkePublicKey::from_bytes(&[0u8; 5]);
        assert!(result.is_err(), "short bytes should be rejected for public key");

        let result = HpkePrivateKey::from_bytes(&[0u8; 5]);
        assert!(result.is_err(), "short bytes should be rejected for private key");
    }
}