newton-core 0.4.16

//! SecureEnvelope for encrypted data transport.
//!
//! Provides a serializable container that bundles HPKE ciphertext with the
//! metadata needed for decryption: encapsulated key, policy client address,
//! chain ID, and recipient public key. The additional authenticated data (AAD)
//! is derived from `keccak256(abi.encodePacked(policy_client, chain_id))` to
//! bind the ciphertext to a specific policy context.

use alloy::primitives::keccak256;
use serde::{Deserialize, Serialize};
use zeroize::Zeroizing;

use super::{
    error::CryptoError,
    hpke::{self, HpkePrivateKey, HpkePublicKey},
};

/// Encrypted data envelope for privacy-preserving transport between clients
/// and operators.
///
/// Contains HPKE ciphertext along with the context needed for decryption and
/// policy binding. The AAD is deterministically computed from `policy_client`
/// and `chain_id`, so tampering with either field causes decryption to fail.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SecureEnvelope {
    /// HPKE encapsulated key (hex-encoded, no 0x prefix).
    pub enc: String,
    /// HPKE ciphertext including the Poly1305 authentication tag (hex-encoded, no 0x prefix).
    pub ciphertext: String,
    /// Policy client address (0x-prefixed, 20-byte Ethereum address).
    pub policy_client: String,
    /// Chain ID for context binding.
    pub chain_id: u64,
    /// Recipient operator Ed25519 public key (hex-encoded, no 0x prefix).
    pub recipient_pubkey: String,
}

impl SecureEnvelope {
    /// Encrypt `plaintext` and package it into a `SecureEnvelope`.
    ///
    /// The AAD is computed as `keccak256(abi.encodePacked(policy_client_bytes, chain_id_be))`,
    /// binding the ciphertext to the specific policy client and chain.
    pub fn seal(
        plaintext: &[u8],
        policy_client: &str,
        chain_id: u64,
        recipient_hpke_pk: &HpkePublicKey,
        recipient_ed25519_pubkey: &[u8],
    ) -> Result<Self, CryptoError> {
        let aad = compute_aad(policy_client, chain_id)?;

        let (enc_bytes, ct_bytes) = hpke::encrypt(recipient_hpke_pk, plaintext, aad.as_slice())?;

        Ok(Self {
            enc: hex::encode(&enc_bytes),
            ciphertext: hex::encode(&ct_bytes),
            policy_client: policy_client.to_string(),
            chain_id,
            recipient_pubkey: hex::encode(recipient_ed25519_pubkey),
        })
    }

    /// Decrypt the envelope contents using the recipient's HPKE private key.
    ///
    /// Recomputes AAD from `policy_client` and `chain_id` stored in the
    /// envelope, so any tampering with those fields causes decryption to fail.
    pub fn open(&self, recipient_hpke_sk: &HpkePrivateKey) -> Result<Zeroizing<Vec<u8>>, CryptoError> {
        let aad = compute_aad(&self.policy_client, self.chain_id)?;

        let enc_bytes =
            hex::decode(&self.enc).map_err(|e| CryptoError::Deserialization(format!("invalid enc hex: {e}")))?;
        let ct_bytes = hex::decode(&self.ciphertext)
            .map_err(|e| CryptoError::Deserialization(format!("invalid ciphertext hex: {e}")))?;

        hpke::decrypt(recipient_hpke_sk, &enc_bytes, &ct_bytes, aad.as_slice())
    }
}

/// Compute the AAD as `keccak256(abi.encodePacked(policy_client_bytes, chain_id_be_bytes))`.
///
/// Public so the gateway can recompute AAD for threshold decryption without
/// access to a full `SecureEnvelope`.
pub fn compute_aad(policy_client: &str, chain_id: u64) -> Result<Vec<u8>, CryptoError> {
    let client_hex = policy_client.strip_prefix("0x").unwrap_or(policy_client);
    let policy_client_bytes =
        hex::decode(client_hex).map_err(|e| CryptoError::Serialization(format!("invalid policy_client hex: {e}")))?;

    let mut packed = Vec::with_capacity(policy_client_bytes.len() + 8);
    packed.extend_from_slice(&policy_client_bytes);
    packed.extend_from_slice(&chain_id.to_be_bytes());

    let hash = keccak256(&packed);
    Ok(hash.to_vec())
}

/// Derive an HPKE keypair from an ECDSA private key.
///
/// Follows the full key derivation chain:
/// ECDSA (secp256k1) → Ed25519 (HKDF-SHA256) → X25519 → HPKE keypair.
///
/// Both gateway and operator use this to obtain HPKE keys from their existing
/// ECDSA key material, avoiding the need for separate HPKE key management.
pub fn derive_hpke_keypair_from_ecdsa(ecdsa_key: &[u8; 32]) -> Result<(HpkePrivateKey, HpkePublicKey), CryptoError> {
    let ed25519_sk = super::ed25519::derive_ed25519_from_ecdsa(ecdsa_key)?;

    let x25519_sk_bytes = super::ed25519::ed25519_to_x25519_private(&ed25519_sk)?;
    let x25519_pk_bytes = super::ed25519::ed25519_to_x25519_public(&ed25519_sk)?;

    let hpke_sk = HpkePrivateKey::from_bytes(&x25519_sk_bytes)?;
    let hpke_pk = HpkePublicKey::from_bytes(&x25519_pk_bytes)?;

    Ok((hpke_sk, hpke_pk))
}

/// A decrypted envelope: `(ref_id, zeroize-on-drop plaintext bytes)`.
#[cfg(feature = "database")]
pub type DecryptedEnvelope = (uuid::Uuid, Zeroizing<Vec<u8>>);

/// Decrypt encrypted data references using an ECDSA-derived HPKE key.
///
/// Derives the full ECDSA → Ed25519 → X25519 → HPKE key chain, then decrypts
/// each [`SecureEnvelope`] contained in the provided database records.
///
/// This is the shared decryption path used by the operator (HPKE
/// decryption before broadcast) and the operator (fallback / direct mode).
///
/// Returns a vector of `(ref_id, decrypted_plaintext)` pairs.
#[cfg(feature = "database")]
pub fn decrypt_envelopes(
    ecdsa_key: &[u8; 32],
    refs: &[crate::database::EncryptedDataRefRecord],
) -> Result<Vec<DecryptedEnvelope>, CryptoError> {
    let (hpke_sk, _) = derive_hpke_keypair_from_ecdsa(ecdsa_key)?;

    let mut results = Vec::with_capacity(refs.len());
    for record in refs {
        let envelope: SecureEnvelope = serde_json::from_slice(&record.envelope).map_err(|e| {
            CryptoError::Deserialization(format!("failed to deserialize envelope for ref {}: {e}", record.id))
        })?;

        let plaintext = envelope.open(&hpke_sk)?;
        results.push((record.id, plaintext));
    }

    Ok(results)
}

/// Derive an HPKE keypair from a raw 32-byte key (not ECDSA-derived).
///
/// Uses the same derivation chain as `derive_hpke_keypair_from_ecdsa` but
/// accepts a standalone raw key instead of an ECDSA key. This is for the
/// `ENCRYPTION_PRIVATE_KEY` path, decoupled from the `task_generator_signer`.
pub fn derive_hpke_keypair_from_raw(raw_key: &[u8; 32]) -> Result<(HpkePrivateKey, HpkePublicKey), CryptoError> {
    // Reuse the existing ECDSA derivation pipeline:
    // raw_key → HKDF-SHA256 → Ed25519 seed → X25519 private key → X25519 public key
    derive_hpke_keypair_from_ecdsa(raw_key)
}

/// Decrypt SecureEnvelopes using a pre-derived HPKE private key.
///
/// This is the operator-side decryption path when a standalone
/// `ENCRYPTION_PRIVATE_KEY` is configured: the caller derives the HPKE keypair
/// once at startup and passes the private key directly, avoiding repeated
/// ECDSA key derivation on every task.
///
/// Returns a merged JSON object of all decrypted envelopes. Returns an empty
/// object if `refs` is empty. Returns an error if any envelope fails to
/// deserialize or decrypt — callers should handle the error rather than
/// silently dropping data.
#[cfg(feature = "database")]
pub fn decrypt_envelopes_with_key(
    hpke_sk: &HpkePrivateKey,
    refs: &[crate::database::EncryptedDataRefRecord],
) -> Result<serde_json::Value, CryptoError> {
    if refs.is_empty() {
        return Ok(serde_json::Value::Object(serde_json::Map::new()));
    }

    let mut merged = serde_json::Map::new();

    for (i, r) in refs.iter().enumerate() {
        let envelope: SecureEnvelope = serde_json::from_slice(&r.envelope)
            .map_err(|e| CryptoError::HpkeDecrypt(format!("ref {i}: failed to deserialize envelope: {e}")))?;

        let plaintext = envelope
            .open(hpke_sk)
            .map_err(|e| CryptoError::HpkeDecrypt(format!("ref {i}: HPKE decryption failed: {e}")))?;

        let key = format!("ref_{}", i);
        if let Ok(json_val) = serde_json::from_slice::<serde_json::Value>(&plaintext) {
            merged.insert(key, json_val);
        } else {
            let text = String::from_utf8_lossy(&plaintext).to_string();
            merged.insert(key, serde_json::Value::String(text));
        }
    }

    Ok(serde_json::Value::Object(merged))
}

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

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

    const TEST_POLICY_CLIENT: &str = "0xaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa";
    const TEST_CHAIN_ID: u64 = 31337;

    fn test_keypair() -> (HpkePrivateKey, HpkePublicKey) {
        hpke::generate_keypair()
    }

    #[test]
    fn seal_open_roundtrip() {
        let (sk, pk) = test_keypair();
        let ed25519_pubkey = [0xABu8; 32];
        let plaintext = b"confidential policy input data";

        let envelope = SecureEnvelope::seal(plaintext, TEST_POLICY_CLIENT, TEST_CHAIN_ID, &pk, &ed25519_pubkey)
            .expect("seal failed");

        let recovered = envelope.open(&sk).expect("open failed");
        assert_eq!(&*recovered, plaintext);
    }

    #[test]
    fn open_with_wrong_key_fails() {
        let (_sk, pk) = test_keypair();
        let (wrong_sk, _wrong_pk) = test_keypair();
        let ed25519_pubkey = [0xBBu8; 32];
        let plaintext = b"secret";

        let envelope = SecureEnvelope::seal(plaintext, TEST_POLICY_CLIENT, TEST_CHAIN_ID, &pk, &ed25519_pubkey)
            .expect("seal failed");

        let result = envelope.open(&wrong_sk);
        assert!(result.is_err(), "open with wrong key should fail");
    }

    #[test]
    fn tampered_ciphertext_fails() {
        let (sk, pk) = test_keypair();
        let ed25519_pubkey = [0xCCu8; 32];
        let plaintext = b"integrity check";

        let mut envelope = SecureEnvelope::seal(plaintext, TEST_POLICY_CLIENT, TEST_CHAIN_ID, &pk, &ed25519_pubkey)
            .expect("seal failed");

        // Flip a byte in the ciphertext hex
        let mut ct_bytes = hex::decode(&envelope.ciphertext).expect("hex decode");
        if let Some(b) = ct_bytes.first_mut() {
            *b ^= 0xFF;
        }
        envelope.ciphertext = hex::encode(&ct_bytes);

        let result = envelope.open(&sk);
        assert!(result.is_err(), "tampered ciphertext should fail");
    }

    #[test]
    fn tampered_policy_client_fails() {
        let (sk, pk) = test_keypair();
        let ed25519_pubkey = [0xDDu8; 32];
        let plaintext = b"aad binding test";

        let mut envelope = SecureEnvelope::seal(plaintext, TEST_POLICY_CLIENT, TEST_CHAIN_ID, &pk, &ed25519_pubkey)
            .expect("seal failed");

        // Change the policy client (AAD will differ on open)
        envelope.policy_client = "0xbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb".to_string();

        let result = envelope.open(&sk);
        assert!(result.is_err(), "open with tampered policy_client should fail");
    }

    #[test]
    fn tampered_chain_id_fails() {
        let (sk, pk) = test_keypair();
        let ed25519_pubkey = [0xEEu8; 32];
        let plaintext = b"chain binding test";

        let mut envelope = SecureEnvelope::seal(plaintext, TEST_POLICY_CLIENT, TEST_CHAIN_ID, &pk, &ed25519_pubkey)
            .expect("seal failed");

        envelope.chain_id = 99999;

        let result = envelope.open(&sk);
        assert!(result.is_err(), "open with tampered chain_id should fail");
    }

    #[test]
    fn serde_json_roundtrip() {
        let (_sk, pk) = test_keypair();
        let ed25519_pubkey = [0x11u8; 32];
        let plaintext = b"json serialization test";

        let envelope = SecureEnvelope::seal(plaintext, TEST_POLICY_CLIENT, TEST_CHAIN_ID, &pk, &ed25519_pubkey)
            .expect("seal failed");

        let json = serde_json::to_string(&envelope).expect("serialize failed");
        let deserialized: SecureEnvelope = serde_json::from_str(&json).expect("deserialize failed");

        assert_eq!(envelope.enc, deserialized.enc);
        assert_eq!(envelope.ciphertext, deserialized.ciphertext);
        assert_eq!(envelope.policy_client, deserialized.policy_client);
        assert_eq!(envelope.chain_id, deserialized.chain_id);
        assert_eq!(envelope.recipient_pubkey, deserialized.recipient_pubkey);
    }
}