venice-e2ee-proxy 0.1.0

//! Venice E2EE encryption/decryption codec.
//!
//! Implements the Venice E2EE field codec:
//!
//! ```text
//! ephemeral_public_key[65 bytes] || nonce[12 bytes] || ciphertext_and_gcm_tag
//! ```
//!
//! The packed bytes are serialized as lowercase hex strings in request message
//! content fields and in encrypted Venice streaming response `delta.content`
//! fields.

use std::fmt;

use aes_gcm::{
    Aes256Gcm,
    aead::{Aead, KeyInit},
};
use hkdf::Hkdf;
use k256::{
    PublicKey, SecretKey,
    ecdh::{SharedSecret, diffie_hellman},
    elliptic_curve::sec1::ToEncodedPoint,
};
use rand_core::{OsRng, RngCore};
use serde::{Deserialize, Serialize};
use sha2::Sha256;
use thiserror::Error;
use zeroize::{ZeroizeOnDrop, Zeroizing};

use crate::config::E2eeConfig;

pub const EPHEMERAL_PUBLIC_KEY_LEN: usize = 65;
pub const NONCE_LEN: usize = 12;
pub const AES_256_KEY_LEN: usize = 32;
pub const AES_GCM_TAG_LEN: usize = 16;
pub const PACKED_PREFIX_LEN: usize = EPHEMERAL_PUBLIC_KEY_LEN + NONCE_LEN;
pub const MIN_PACKED_PAYLOAD_LEN: usize = PACKED_PREFIX_LEN + AES_GCM_TAG_LEN;

/// E2EE codec configured with the HKDF info string and fail-closed response policy.
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct E2eeCodec {
    hkdf_info: Vec<u8>,
    require_encrypted_response_content: bool,
}

impl E2eeCodec {
    /// Builds an E2EE codec from validated proxy configuration.
    pub fn from_config(config: &E2eeConfig) -> Result<Self, E2eeCodecError> {
        Self::new(
            config.hkdf_info.as_bytes(),
            config.require_encrypted_response_content,
        )
    }

    /// Builds an E2EE codec from HKDF context bytes and the missing-content response policy.
    pub fn new(
        hkdf_info: impl AsRef<[u8]>,
        require_encrypted_response_content: bool,
    ) -> Result<Self, E2eeCodecError> {
        let hkdf_info = hkdf_info.as_ref();
        if hkdf_info.is_empty() {
            return Err(E2eeCodecError::EmptyHkdfInfo);
        }

        Ok(Self {
            hkdf_info: hkdf_info.to_vec(),
            require_encrypted_response_content,
        })
    }

    /// Returns whether response chunks must contain encrypted content fields.
    pub fn require_encrypted_response_content(&self) -> bool {
        self.require_encrypted_response_content
    }

    /// Derives a Venice AES-256-GCM content key from a local secp256k1 private
    /// key and a peer uncompressed SEC1 public key hex string.
    pub fn derive_content_key(
        &self,
        local_private_key: &SecretKey,
        peer_public_key_hex: &str,
    ) -> Result<ContentEncryptionKey, E2eeCodecError> {
        let peer_public_key = decode_uncompressed_public_key_hex(peer_public_key_hex)?;
        Ok(self.derive_content_key_from_public_key(local_private_key, &peer_public_key))
    }

    /// Encrypts one normalized model-visible text field for a Venice E2EE request.
    pub fn encrypt_content(
        &self,
        plaintext: &str,
        peer_public_key_hex: &str,
    ) -> Result<EncryptedPayload, E2eeCodecError> {
        let peer_public_key = decode_uncompressed_public_key_hex(peer_public_key_hex)?;
        let ephemeral_private_key = SecretKey::random(&mut OsRng);
        let nonce = Nonce::generate();

        self.encrypt_content_with_parts(plaintext, &peer_public_key, ephemeral_private_key, nonce)
    }

    /// Decrypts a packed Venice E2EE hex payload into UTF-8 text.
    pub fn decrypt_content(
        &self,
        payload: &EncryptedPayload,
        recipient_private_key: &SecretKey,
    ) -> Result<String, E2eeCodecError> {
        let packed = PackedEncryptedPayload::unpack(payload)?;
        let key = self.derive_content_key_from_public_key(
            recipient_private_key,
            &packed.ephemeral_public_key,
        );
        let cipher = aes256_gcm_from_key(&key);
        let plaintext = Zeroizing::new(
            cipher
                .decrypt((&packed.nonce).into(), packed.ciphertext_and_tag.as_slice())
                .map_err(|_| E2eeCodecError::AuthenticationFailed)?,
        );

        String::from_utf8(plaintext.to_vec()).map_err(|_| E2eeCodecError::InvalidPlaintextUtf8)
    }

    /// Extracts and decrypts `choices[0].delta.content`-style encrypted response
    /// content. Missing content fails closed when configured to require encrypted
    /// response content, and otherwise returns `Ok(None)`.
    pub fn decrypt_response_content(
        &self,
        content: Option<&str>,
        recipient_private_key: &SecretKey,
    ) -> Result<Option<String>, E2eeCodecError> {
        let Some(content) = content else {
            return if self.require_encrypted_response_content {
                Err(E2eeCodecError::MissingEncryptedContent)
            } else {
                Ok(None)
            };
        };

        let payload = EncryptedPayload::from_hex(content)?;
        self.decrypt_content(&payload, recipient_private_key)
            .map(Some)
    }

    /// Encrypts plaintext using caller-supplied peer key, ephemeral key, and nonce values.
    fn encrypt_content_with_parts(
        &self,
        plaintext: &str,
        peer_public_key: &PublicKey,
        ephemeral_private_key: SecretKey,
        nonce: Nonce,
    ) -> Result<EncryptedPayload, E2eeCodecError> {
        let ephemeral_public_key = ephemeral_private_key.public_key();
        let ephemeral_public_key = ephemeral_public_key.to_encoded_point(false);
        let ephemeral_public_key_bytes = ephemeral_public_key.as_bytes();
        debug_assert_eq!(ephemeral_public_key_bytes.len(), EPHEMERAL_PUBLIC_KEY_LEN);

        let key = self.derive_content_key_from_public_key(&ephemeral_private_key, peer_public_key);
        let cipher = aes256_gcm_from_key(&key);
        let ciphertext_and_tag = cipher
            .encrypt(nonce.as_bytes().into(), plaintext.as_bytes())
            .map_err(|_| E2eeCodecError::EncryptionFailed)?;

        let mut packed = Vec::with_capacity(PACKED_PREFIX_LEN + ciphertext_and_tag.len());
        packed.extend_from_slice(ephemeral_public_key_bytes);
        packed.extend_from_slice(nonce.as_bytes());
        packed.extend_from_slice(&ciphertext_and_tag);

        Ok(EncryptedPayload::from_packed_bytes_unchecked(&packed))
    }

    /// Derives a content key from a local private key and parsed peer public key.
    fn derive_content_key_from_public_key(
        &self,
        local_private_key: &SecretKey,
        peer_public_key: &PublicKey,
    ) -> ContentEncryptionKey {
        let shared_secret = diffie_hellman(
            local_private_key.to_nonzero_scalar(),
            peer_public_key.as_affine(),
        );
        derive_aes_key(&shared_secret, &self.hkdf_info)
    }
}

impl Default for E2eeCodec {
    /// Returns the codec for the default E2EE configuration.
    fn default() -> Self {
        Self::from_config(&E2eeConfig::default()).expect("default E2EE config is valid")
    }
}

/// An AES-256-GCM content-encryption key derived from ECDH + HKDF-SHA256.
///
/// Debug output is redacted and the key bytes are zeroized on drop.
pub struct ContentEncryptionKey(Zeroizing<[u8; AES_256_KEY_LEN]>);

impl ContentEncryptionKey {
    /// Wraps derived key bytes in the content-encryption key type.
    fn new(bytes: Zeroizing<[u8; AES_256_KEY_LEN]>) -> Self {
        Self(bytes)
    }

    /// Returns the raw key bytes for cipher construction.
    fn as_slice(&self) -> &[u8] {
        &self.0[..]
    }
}

impl ZeroizeOnDrop for ContentEncryptionKey {}

impl fmt::Debug for ContentEncryptionKey {
    /// Formats the key without exposing secret key bytes.
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.write_str("ContentEncryptionKey([redacted])")
    }
}

impl PartialEq for ContentEncryptionKey {
    /// Compares two content keys by their raw key bytes.
    fn eq(&self, other: &Self) -> bool {
        self.as_slice() == other.as_slice()
    }
}

impl Eq for ContentEncryptionKey {}

/// Builds an AES-256-GCM cipher from a derived content-encryption key.
fn aes256_gcm_from_key(key: &ContentEncryptionKey) -> Aes256Gcm {
    // Zeroization note for the resolved RustCrypto stack used here:
    // - `aes-gcm` with `zeroize` wipes its temporary GHASH key during init.
    // - the direct `aes` dependency enables `aes/zeroize`, so AES-256 retained
    //   round keys implement safe zeroizing drop behavior.
    // - the direct `ghash` dependency enables zeroizing of GHASH temporary
    //   conversion buffers.
    // - the direct `polyval` dependency enables zeroizing drops for POLYVAL
    //   backends on targets/configurations where those drops are reachable.
    //
    // Residual concern: in the resolved `polyval` 0.6.x x86/x86_64 autodetect
    // path, the public wrapper uses `ManuallyDrop` and does not expose a safe
    // `Drop`/`ZeroizeOnDrop` implementation. That means `Aes256Gcm` still cannot
    // be proven to zeroize every retained GHASH/POLYVAL byte on this target. We
    // avoid brittle unsafe whole-object wiping and instead keep the cipher scoped
    // to one encrypt/decrypt operation while zeroizing the derived key material
    // and relying on verified safe drop behavior where the crates expose it.
    Aes256Gcm::new((&*key.0).into())
}

/// AES-GCM nonce used by the Venice field codec.
#[derive(Clone, Copy, PartialEq, Eq)]
pub struct Nonce([u8; NONCE_LEN]);

impl Nonce {
    /// Generates a fresh random AES-GCM nonce.
    pub fn generate() -> Self {
        let mut bytes = [0_u8; NONCE_LEN];
        OsRng.fill_bytes(&mut bytes);
        Self(bytes)
    }

    /// Builds a nonce from exactly 12 bytes.
    pub fn from_bytes(bytes: [u8; NONCE_LEN]) -> Self {
        Self(bytes)
    }

    /// Returns the nonce bytes for AES-GCM encryption or decryption.
    pub fn as_bytes(&self) -> &[u8; NONCE_LEN] {
        &self.0
    }
}

impl fmt::Debug for Nonce {
    /// Formats the nonce as hex for diagnostics.
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_tuple("Nonce")
            .field(&hex::encode(self.as_bytes()))
            .finish()
    }
}

/// Venice encrypted payload serialized as a lowercase hex string.
#[derive(Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(transparent)]
pub struct EncryptedPayload(String);

impl EncryptedPayload {
    /// Validates a packed encrypted payload hex string and stores it in lowercase form.
    pub fn from_hex(value: impl Into<String>) -> Result<Self, E2eeCodecError> {
        let value = value.into();
        validate_packed_payload_hex(&value)?;
        Ok(Self(value.to_ascii_lowercase()))
    }

    /// Returns the packed encrypted payload as a hex string slice.
    pub fn as_hex(&self) -> &str {
        &self.0
    }

    /// Consumes the payload and returns the owned packed hex string.
    pub fn into_hex(self) -> String {
        self.0
    }

    /// Encodes already-packed encrypted payload bytes as lowercase hex.
    fn from_packed_bytes_unchecked(bytes: &[u8]) -> Self {
        Self(hex::encode(bytes))
    }
}

impl fmt::Debug for EncryptedPayload {
    /// Formats payload metadata without printing the ciphertext bytes.
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("EncryptedPayload")
            .field("hex_len", &self.0.len())
            .finish()
    }
}

/// Parsed parts of a packed Venice E2EE payload.
struct PackedEncryptedPayload {
    ephemeral_public_key: PublicKey,
    nonce: [u8; NONCE_LEN],
    ciphertext_and_tag: Vec<u8>,
}

impl PackedEncryptedPayload {
    /// Splits a validated encrypted payload into public key, nonce, and ciphertext/tag parts.
    fn unpack(payload: &EncryptedPayload) -> Result<Self, E2eeCodecError> {
        let bytes = hex::decode(payload.as_hex()).map_err(|error| {
            E2eeCodecError::MalformedEncryptedPayload {
                message: error.to_string(),
            }
        })?;
        if bytes.len() < MIN_PACKED_PAYLOAD_LEN {
            return Err(E2eeCodecError::MalformedEncryptedPayload {
                message: format!(
                    "packed encrypted payload is too short: got {} bytes, need at least {MIN_PACKED_PAYLOAD_LEN}",
                    bytes.len()
                ),
            });
        }

        if bytes[0] != 0x04 {
            return Err(E2eeCodecError::MalformedEncryptedPayload {
                message: "ephemeral public key must be uncompressed SEC1 format".to_owned(),
            });
        }
        let ephemeral_public_key = PublicKey::from_sec1_bytes(&bytes[..EPHEMERAL_PUBLIC_KEY_LEN])
            .map_err(|_| E2eeCodecError::MalformedEncryptedPayload {
            message: "ephemeral public key is not a valid secp256k1 key".to_owned(),
        })?;

        let mut nonce = [0_u8; NONCE_LEN];
        nonce.copy_from_slice(&bytes[EPHEMERAL_PUBLIC_KEY_LEN..PACKED_PREFIX_LEN]);

        Ok(Self {
            ephemeral_public_key,
            nonce,
            ciphertext_and_tag: bytes[PACKED_PREFIX_LEN..].to_vec(),
        })
    }
}

/// Errors returned by E2EE request encryption and response decryption.
#[derive(Debug, Error, PartialEq, Eq)]
pub enum E2eeCodecError {
    #[error("configured E2EE HKDF info must not be empty")]
    EmptyHkdfInfo,
    #[error("encrypted response content is required but missing")]
    MissingEncryptedContent,
    #[error("encrypted payload is malformed: {message}")]
    MalformedEncryptedPayload { message: String },
    #[error("encrypted payload authentication failed")]
    AuthenticationFailed,
    #[error("invalid E2EE public key: {message}")]
    InvalidPublicKey { message: String },
    #[error("decrypted E2EE payload is not valid UTF-8")]
    InvalidPlaintextUtf8,
    #[error("E2EE encryption failed")]
    EncryptionFailed,
}

/// Derives an AES-256 content key from an ECDH shared secret and HKDF context bytes.
fn derive_aes_key(shared_secret: &SharedSecret, hkdf_info: &[u8]) -> ContentEncryptionKey {
    let hkdf = Hkdf::<Sha256>::new(None, shared_secret.raw_secret_bytes());
    let mut output_key = Zeroizing::new([0_u8; AES_256_KEY_LEN]);
    hkdf.expand(hkdf_info, output_key.as_mut_slice())
        .expect("32-byte HKDF-SHA256 output length is always valid");
    ContentEncryptionKey::new(output_key)
}

/// Parses an uncompressed SEC1 secp256k1 public key from a hex string.
fn decode_uncompressed_public_key_hex(value: &str) -> Result<PublicKey, E2eeCodecError> {
    let bytes = hex::decode(value).map_err(|error| E2eeCodecError::InvalidPublicKey {
        message: error.to_string(),
    })?;

    if bytes.len() != EPHEMERAL_PUBLIC_KEY_LEN {
        return Err(E2eeCodecError::InvalidPublicKey {
            message: format!(
                "expected {EPHEMERAL_PUBLIC_KEY_LEN} uncompressed SEC1 bytes, got {}",
                bytes.len()
            ),
        });
    }
    if bytes.first() != Some(&0x04) {
        return Err(E2eeCodecError::InvalidPublicKey {
            message: "public key must be uncompressed SEC1 format".to_owned(),
        });
    }

    PublicKey::from_sec1_bytes(&bytes).map_err(|_| E2eeCodecError::InvalidPublicKey {
        message: "public key is not a valid secp256k1 key".to_owned(),
    })
}

/// Validates that a hex string can contain the minimum packed Venice E2EE payload.
fn validate_packed_payload_hex(value: &str) -> Result<(), E2eeCodecError> {
    if value.is_empty() {
        return Err(E2eeCodecError::MalformedEncryptedPayload {
            message: "encrypted payload hex string is empty".to_owned(),
        });
    }
    if !value.len().is_multiple_of(2) {
        return Err(E2eeCodecError::MalformedEncryptedPayload {
            message: "encrypted payload hex string has odd length".to_owned(),
        });
    }
    if value.len() < MIN_PACKED_PAYLOAD_LEN * 2 {
        return Err(E2eeCodecError::MalformedEncryptedPayload {
            message: format!(
                "encrypted payload hex string is too short: got {} chars, need at least {}",
                value.len(),
                MIN_PACKED_PAYLOAD_LEN * 2
            ),
        });
    }
    if let Some((index, ch)) = value.char_indices().find(|(_, ch)| !ch.is_ascii_hexdigit()) {
        return Err(E2eeCodecError::MalformedEncryptedPayload {
            message: format!(
                "encrypted payload hex string contains non-hex character {ch:?} at index {index}"
            ),
        });
    }
    Ok(())
}

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

    const FIXED_NONCE: [u8; NONCE_LEN] = [
        0xa0, 0xa1, 0xa2, 0xa3, 0xa4, 0xa5, 0xa6, 0xa7, 0xa8, 0xa9, 0xaa, 0xab,
    ];
    const FIXED_RECIPIENT_PRIVATE_KEY_HEX: &str =
        "1111111111111111111111111111111111111111111111111111111111111111";
    const FIXED_EPHEMERAL_PRIVATE_KEY_HEX: &str =
        "2222222222222222222222222222222222222222222222222222222222222222";
    const DETERMINISTIC_PLAINTEXT: &str = "deterministic Venice E2EE fixture";
    const DETERMINISTIC_CIPHERTEXT_HEX: &str = "04466d7fcae563e5cb09a0d1870bb580344804617879a14949cf22285f1bae3f276728176c3c6431f8eeda4538dc37c865e2784f3a9e77d044f33e407797e1278aa0a1a2a3a4a5a6a7a8a9aaab3b364a6560dc6246955e1379bac6c7a0f453c5b2d9be6eabb00cad9955278b4c401f6793813d7f98ba8f163a5c51b87686";

    fn secret_key_from_hex(value: &str) -> SecretKey {
        let bytes = hex::decode(value).expect("test key hex should decode");
        SecretKey::from_slice(&bytes).expect("test key should be valid")
    }

    fn public_key_hex(secret_key: &SecretKey) -> String {
        hex::encode(secret_key.public_key().to_encoded_point(false).as_bytes())
    }

    #[test]
    fn encrypt_decrypt_round_trip() {
        let codec = E2eeCodec::default();
        let recipient_private_key = SecretKey::random(&mut OsRng);
        let recipient_public_key_hex = public_key_hex(&recipient_private_key);

        let encrypted = codec
            .encrypt_content("hello from local proxy", &recipient_public_key_hex)
            .expect("encryption should succeed");
        let decrypted = codec
            .decrypt_content(&encrypted, &recipient_private_key)
            .expect("decryption should succeed");

        assert_eq!(decrypted, "hello from local proxy");
        assert!(
            encrypted
                .as_hex()
                .chars()
                .all(|ch| !ch.is_ascii_uppercase())
        );
    }

    #[test]
    fn decryption_with_wrong_key_fails_authentication() {
        let codec = E2eeCodec::default();
        let recipient_private_key = SecretKey::random(&mut OsRng);
        let wrong_private_key = SecretKey::random(&mut OsRng);
        let recipient_public_key_hex = public_key_hex(&recipient_private_key);
        let encrypted = codec
            .encrypt_content("secret", &recipient_public_key_hex)
            .expect("encryption should succeed");

        let err = codec
            .decrypt_content(&encrypted, &wrong_private_key)
            .expect_err("wrong key must fail closed");

        assert_eq!(err, E2eeCodecError::AuthenticationFailed);
    }

    #[test]
    fn tampered_ciphertext_fails_authentication() {
        let codec = E2eeCodec::default();
        let recipient_private_key = SecretKey::random(&mut OsRng);
        let recipient_public_key_hex = public_key_hex(&recipient_private_key);
        let encrypted = codec
            .encrypt_content("secret", &recipient_public_key_hex)
            .expect("encryption should succeed");
        let mut packed = hex::decode(encrypted.as_hex()).expect("ciphertext should decode");
        let last = packed.last_mut().expect("ciphertext has tag byte");
        *last ^= 0x01;
        let tampered = EncryptedPayload::from_packed_bytes_unchecked(&packed);

        let err = codec
            .decrypt_content(&tampered, &recipient_private_key)
            .expect_err("tampered ciphertext must fail closed");

        assert_eq!(err, E2eeCodecError::AuthenticationFailed);
    }

    #[test]
    fn malformed_payload_fails_closed() {
        let codec = E2eeCodec::default();
        let recipient_private_key = SecretKey::random(&mut OsRng);

        let err = codec
            .decrypt_response_content(Some("not encrypted"), &recipient_private_key)
            .expect_err("non-hex payload should fail closed");
        assert!(matches!(
            err,
            E2eeCodecError::MalformedEncryptedPayload { .. }
        ));

        let too_short = "04".repeat(EPHEMERAL_PUBLIC_KEY_LEN + NONCE_LEN);
        let err =
            EncryptedPayload::from_hex(too_short).expect_err("short payload should be rejected");
        assert!(matches!(
            err,
            E2eeCodecError::MalformedEncryptedPayload { .. }
        ));
    }

    #[test]
    fn missing_encrypted_response_content_respects_config() {
        let recipient_private_key = SecretKey::random(&mut OsRng);

        let required = E2eeCodec::new("ecdsa_encryption", true).expect("config should be valid");
        let err = required
            .decrypt_response_content(None, &recipient_private_key)
            .expect_err("missing required encrypted content should fail");
        assert_eq!(err, E2eeCodecError::MissingEncryptedContent);

        let optional = E2eeCodec::new("ecdsa_encryption", false).expect("config should be valid");
        let decrypted = optional
            .decrypt_response_content(None, &recipient_private_key)
            .expect("missing optional content should be allowed");
        assert_eq!(decrypted, None);
    }

    #[test]
    fn deterministic_test_vector_with_fixed_nonce_and_ephemeral_key() {
        let codec = E2eeCodec::default();
        let recipient_private_key = secret_key_from_hex(FIXED_RECIPIENT_PRIVATE_KEY_HEX);
        let recipient_public_key = recipient_private_key.public_key();
        let ephemeral_private_key = secret_key_from_hex(FIXED_EPHEMERAL_PRIVATE_KEY_HEX);

        let encrypted = codec
            .encrypt_content_with_parts(
                DETERMINISTIC_PLAINTEXT,
                &recipient_public_key,
                ephemeral_private_key,
                Nonce::from_bytes(FIXED_NONCE),
            )
            .expect("deterministic encryption should succeed");

        assert_eq!(encrypted.as_hex(), DETERMINISTIC_CIPHERTEXT_HEX);
        let decrypted = codec
            .decrypt_content(&encrypted, &recipient_private_key)
            .expect("deterministic fixture should decrypt");
        assert_eq!(decrypted, DETERMINISTIC_PLAINTEXT);
    }

    #[test]
    fn derived_keys_match_from_both_sides_and_debug_is_redacted() {
        fn assert_zeroize_on_drop<T: ZeroizeOnDrop>() {}

        assert_zeroize_on_drop::<SecretKey>();
        assert_zeroize_on_drop::<SharedSecret>();
        assert_zeroize_on_drop::<aes::Aes256>();
        assert_zeroize_on_drop::<ContentEncryptionKey>();

        let codec = E2eeCodec::default();
        let local_private_key = SecretKey::random(&mut OsRng);
        let peer_private_key = SecretKey::random(&mut OsRng);
        let local_public_key_hex = public_key_hex(&local_private_key);
        let peer_public_key_hex = public_key_hex(&peer_private_key);

        let local_key = codec
            .derive_content_key(&local_private_key, &peer_public_key_hex)
            .expect("local derivation should succeed");
        let peer_key = codec
            .derive_content_key(&peer_private_key, &local_public_key_hex)
            .expect("peer derivation should succeed");

        assert_eq!(local_key, peer_key);
        assert_eq!(format!("{local_key:?}"), "ContentEncryptionKey([redacted])");
    }
}