rialo-cdk 0.4.2

// Copyright (c) Subzero Labs, Inc.
// SPDX-License-Identifier: Apache-2.0

//! Secret encryption utilities for TEE-based rex systems.
//!
//! This module provides client-side encryption functions for encrypting secrets
//! (such as API keys) that will be decrypted inside TEEs during rex execution.
//! The encryption includes the creator's public key in the Additional Authenticated Data (AAD)
//! to bind the encrypted secret to its creator.

use rialo_tee_secret_sharing::{encrypt_user_payload, USER_SECRET_AAD};
use rialo_types::PublicKey;

use crate::{
    error::{Result, RialoError},
    rpc::types::Pubkey,
};

/// Maximum allowed secret length (64 KB)
pub const MAX_SECRET_LENGTH: usize = 64 * 1024;

/// Encrypts a secret using the TEE's public key and returns the raw encrypted bytes.
///
/// The encryption uses HPKE (Hybrid Public Key Encryption) with the creator's public key
/// included in the Additional Authenticated Data (AAD). This binds the encrypted secret
/// to the creator, ensuring that the TEE can verify the creator's identity during decryption.
///
/// # Arguments
///
/// * `secret` - The plaintext secret string to encrypt (max 64 KB)
/// * `creator_pubkey` - The public key of the secret creator (used in AAD for authentication)
/// * `secret_sharing_pubkey` - The secret sharing public key (used for encryption)
///
/// # Returns
///
/// * `Result<Vec<u8>>` - Encrypted payload in HPKE format as raw bytes
///
/// # Errors
///
/// Returns an error if:
/// - The secret is empty or exceeds MAX_SECRET_LENGTH
/// - Encryption fails
///
/// # Example
///
/// ```
/// use rialo_cdk::secret_encryption::encrypt_secret;
/// use rialo_types::PublicKey;
///
/// let secret = "Bearer sk-1234567890abcdef".to_string();
/// let creator_pub = "9h1HyLCW5dZnBVap8C5egQ9Z6pHyjsh5MNy83iPqqRuq".parse().unwrap();
///
/// // Decode hex string to bytes and convert to PublicKey
/// let sk_pub_bytes = hex::decode("a1b2c3d4e5f67890a1b2c3d4e5f67890a1b2c3d4e5f67890a1b2c3d4e5f67890")
///     .expect("valid hex string");
/// let sk_pub = PublicKey::from_bytes(
///     sk_pub_bytes.try_into().expect("32 bytes")
/// );
///
/// let encrypted_bytes = encrypt_secret(secret, &creator_pub, &sk_pub)?;
/// # Ok::<(), rialo_cdk::error::RialoError>(())
/// ```
pub fn encrypt_secret(
    secret: String,
    creator_pubkey: &Pubkey,
    secret_sharing_pubkey: &PublicKey,
) -> Result<Vec<u8>> {
    // Validate secret
    if secret.is_empty() {
        return Err(RialoError::InvalidInput(
            "Secret cannot be empty".to_string(),
        ));
    }

    let secret_bytes = secret.as_bytes();
    if secret_bytes.len() > MAX_SECRET_LENGTH {
        return Err(RialoError::InvalidInput(format!(
            "Secret exceeds maximum length of {} bytes (got {} bytes)",
            MAX_SECRET_LENGTH,
            secret_bytes.len()
        )));
    }

    // Decode the hex public key
    let aad = [USER_SECRET_AAD, &creator_pubkey.to_bytes()].concat();

    // Encrypt the secret using HPKE
    let ciphertext = encrypt_user_payload(secret_sharing_pubkey, secret_bytes, aad.as_ref())
        .map_err(|e| RialoError::Encryption(format!("Encryption failed: {}", e)))?;

    Ok(ciphertext)
}

#[cfg(test)]
mod tests {
    use rialo_tee_secret_sharing::{decrypt_user_message, initialize_secret_key, SecretSharingKey};
    use zeroize::Zeroizing;

    use super::*;

    fn create_test_key() -> SecretSharingKey {
        let key_bytes = Zeroizing::new([42u8; 32]);
        SecretSharingKey::from_private_key_bytes(key_bytes)
    }

    #[test]
    fn test_encrypt_secret_basic() {
        let sk = create_test_key();
        let creator_pubkey = Pubkey::from([1u8; 32]);
        let _ = initialize_secret_key(sk.clone());

        let secret = "Bearer test-token-12345".to_string();
        let result = encrypt_secret(secret.clone(), &creator_pubkey, sk.public_key());

        assert!(result.is_ok());
        let ciphertext = result.unwrap();

        // Verify decryption with creator public key in AAD
        let decrypted = decrypt_user_message(
            &ciphertext,
            &PublicKey::from_bytes(creator_pubkey.to_bytes()),
        )
        .unwrap();
        assert_eq!(decrypted.as_bytes(), secret.as_bytes());
    }

    #[test]
    fn test_encrypt_secret_empty() {
        let sk = create_test_key();
        let creator_pubkey = Pubkey::from([3u8; 32]);

        let result = encrypt_secret(String::new(), &creator_pubkey, sk.public_key());

        assert!(result.is_err());
        assert!(result.unwrap_err().to_string().contains("empty"));
    }

    #[test]
    fn test_encrypt_secret_oversized() {
        let sk = create_test_key();
        let creator_pubkey = Pubkey::from([4u8; 32]);

        let oversized_secret = "x".repeat(MAX_SECRET_LENGTH + 1);
        let result = encrypt_secret(oversized_secret, &creator_pubkey, sk.public_key());

        assert!(result.is_err());
        assert!(result.unwrap_err().to_string().contains("maximum length"));
    }

    #[test]
    fn test_encrypt_secret_with_different_creator() {
        let sk = create_test_key();
        let creator_pubkey1 = Pubkey::from([10u8; 32]);
        let creator_pubkey2 = Pubkey::from([20u8; 32]);
        let _ = initialize_secret_key(sk.clone());

        let secret = "test-secret-with-creator".to_string();

        // Encrypt with creator_pubkey1
        let ciphertext1 =
            encrypt_secret(secret.clone(), &creator_pubkey1, sk.public_key()).unwrap();

        // Encrypt with creator_pubkey2
        let ciphertext2 =
            encrypt_secret(secret.clone(), &creator_pubkey2, sk.public_key()).unwrap();

        // Different creator pubkeys should produce different ciphertexts (due to different AAD)
        assert_ne!(ciphertext1, ciphertext2);

        // Both should decrypt correctly with their respective AADs
        let decrypted1 = decrypt_user_message(
            &ciphertext1,
            &PublicKey::from_bytes(creator_pubkey1.to_bytes()),
        )
        .unwrap();
        assert_eq!(decrypted1.as_bytes(), secret.as_bytes());

        let decrypted2 = decrypt_user_message(
            &ciphertext2,
            &PublicKey::from_bytes(creator_pubkey2.to_bytes()),
        )
        .unwrap();
        assert_eq!(decrypted2.as_bytes(), secret.as_bytes());
    }
}