rialo_tee_secret_sharing/

lib.rs

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

//! Secret sharing key management and wrapping for TEEs
//!
//! This crate provides secure key management and secret sharing capabilities
//! for TEEs. It includes key generation,
//! cryptographic operations, and wrapping mechanisms for secure communication.

pub mod key;
pub mod traits;
pub mod types;
pub mod wrapping;

pub mod constants;
mod error;
pub mod hw_rng;
#[cfg(test)]
pub mod tests;

pub use rialo_types::PublicKey;

pub use self::{
    constants::{SECRET_SHARING_HPKE_INFO, USER_SECRET_AAD},
    error::SecretSharingError,
    hw_rng::{HwRng, HwRngError},
    key::SecretSharingKey,
    traits::SecretSharingEngine,
    types::{Aad, KeyId, WrappedKeyEnvelope},
    wrapping::{aad_for_dkg, unwrap_scalar, wrap_scalar_for_recipient, HpkeX25519ChaChaEngine},
};

/// Client-side HPKE encryption function for user payloads.
/// This is a convenience wrapper around the engine's encrypt function that:
/// - Provides a stable public API independent of engine implementation
/// - Allows future engine replacement without breaking client code  
/// - Maintains consistent error handling and parameter validation
///
/// # Arguments
/// * `sk_pub` - The SecretSharingKey public key bytes (from oracle-registry)
/// * `payload` - The plaintext payload to encrypt
/// * `aad` - Associated authenticated data for binding (e.g. USER_SECRET_AAD || creator_pubkey(32 bytes))
///
/// # Returns
/// * `Ok(Vec<u8>)` - Encrypted payload in HPKE format: [enc (32 bytes)] || [ciphertext + tag]
/// * `Err(SecretSharingError)` - If encryption fails
///
/// # Example
/// ```ignore
/// use rialo_tee_secret_sharing::{encrypt_user_payload, USER_SECRET_AAD};
///
/// let sk_pub = get_secret_sharing_pubkey(); // From oracle registry
/// let payload = b"secret data";
/// let aad = USER_SECRET_AAD || creator_pubkey_bytes;
/// let ciphertext = encrypt_user_payload(&sk_pub, payload, &aad)?;
/// ```
pub fn encrypt_user_payload(
    sk_pub: &PublicKey,
    payload: &[u8],
    aad: &[u8],
) -> Result<Vec<u8>, SecretSharingError> {
    HpkeX25519ChaChaEngine::hpke_encrypt_user_payload(sk_pub, payload, aad)
}

/// Default engine = HPKE X25519 + HKDF-SHA256 + ChaCha20-Poly1305
pub type DefaultEngine = HpkeX25519ChaChaEngine;

use std::sync::OnceLock;

/// Global storage for the secret sharing key in this TEE instance.
/// This key is used to decrypt user-submitted secret payloads.
/// The key never leaves this module boundary for security.
static SECRET_KEY: OnceLock<SecretSharingKey> = OnceLock::new();

/// Initialize the secret sharing key for this TEE instance.
/// This should be called once during TEE startup after key generation or distribution.
/// The key remains isolated within this module for security.
///
/// Returns an error if the key is already initialized.
pub fn initialize_secret_key(sk: SecretSharingKey) -> Result<(), SecretSharingError> {
    SECRET_KEY
        .set(sk)
        .map_err(|_| SecretSharingError::UnwrapFailed("Secret key already initialized".into()))
}

/// Returns `true` if a secret key has been initialized in this TEE instance.
pub fn has_secret_key() -> bool {
    SECRET_KEY.get().is_some()
}

/// Common helper function to decrypt user ciphertext using HPKE.
/// This function uses the stored secret sharing key to perform HPKE decryption with the
/// user provided AAD for authenticated encryption.
///
/// # Arguments
/// * `ciphertext` - The encrypted data in HPKE format
/// * `aad` - The associated authenticated data used during encryption
///
/// # Returns
/// * `Ok(Zeroizing<Vec<u8>>)` - The decrypted plaintext bytes
/// * `Err(SecretSharingError)` - If decryption fails or key not initialized
pub fn decrypt_user_ciphertext(
    ciphertext: &[u8],
    aad: &[u8],
) -> Result<zeroize::Zeroizing<Vec<u8>>, SecretSharingError> {
    let sk = SECRET_KEY
        .get()
        .ok_or_else(|| SecretSharingError::UnwrapFailed("Secret key not initialized".into()))?;

    // For user payloads, we expect the ciphertext to be in a format similar to HPKE:
    // [enc (32 bytes)] || [ciphertext + tag (variable length)]
    if ciphertext.len() < 32 {
        return Err(SecretSharingError::EnvelopeInvalid(
            "Ciphertext too short for HPKE format".into(),
        ));
    }

    let (enc_bytes, ct_with_tag) = ciphertext.split_at(32);
    let mut enc_array = [0u8; 32];
    enc_array.copy_from_slice(enc_bytes);

    // Use the existing HPKE implementation from wrapping module with the caller-bound AAD
    let plaintext = wrapping::HpkeX25519ChaChaEngine::hpke_decrypt_user_payload(
        &sk.private_key().to_bytes(),
        &enc_array,
        ct_with_tag,
        aad,
    )?;

    Ok(plaintext)
}

/// Decrypt a user-submitted encrypted message using the stored SecretSharingKey.
///
/// This function takes encrypted user data (containing a plain text message) and decrypts it
/// inside the TEE using the module's stored SK_priv. The decryption uses HPKE in base mode
/// with the caller-bound AAD for authenticated encryption.
/// The AAD is derived from the USER_SECRET_AAD and the creator's public key.
/// # Arguments
/// * `ciphertext` - The encrypted message from the user in HPKE format
/// * `creator_pubkey` - The 32-byte public key of the oracle creator who created the message
///
/// # Returns
/// * `Ok(String)` - The decrypted plain text message
/// * `Err(SecretSharingError)` - If decryption fails, key not initialized, or invalid UTF-8
///
/// # Security Notes
/// - The encrypted message is decrypted inside the TEE using the stored secret key
/// - The SecretSharingKey never leaves this module boundary for security
/// - The decrypted content is validated as UTF-8 before returning
/// - Only the creator's matching public key can decrypt the message
pub fn decrypt_user_message(
    ciphertext: &[u8],
    creator_pubkey: &PublicKey,
) -> Result<String, SecretSharingError> {
    let aad = [USER_SECRET_AAD, &creator_pubkey.to_bytes()].concat();
    let plaintext = decrypt_user_ciphertext(ciphertext, &aad)?;

    // Convert bytes to string using std::str::from_utf8
    let s = std::str::from_utf8(plaintext.as_ref()).map_err(|err| {
        SecretSharingError::UnwrapFailed(format!(
            "Failed to parse decrypted message as utf8: {}",
            err
        ))
    })?;
    Ok(s.to_string())
}

#[cfg(test)]
mod secret_sharing_tests {
    use assert_matches::assert_matches;

    use super::*;

    fn create_test_secret_key() -> SecretSharingKey {
        SecretSharingKey::from_test_seed(b"test-secret-key")
    }

    #[test]
    fn test_initialize_secret_key_double_initialization() {
        let sk1 = create_test_secret_key();
        let sk2 = create_test_secret_key();

        // First initialization
        let _ = initialize_secret_key(sk1);

        // Second initialization should fail
        let result = initialize_secret_key(sk2);
        assert_matches!(
            result,
            Err(SecretSharingError::UnwrapFailed(msg)) if msg.contains("already initialized")
        );
    }

    #[test]
    fn test_decrypt_user_message_ciphertext_too_short() {
        // Ensure key is initialized
        let sk = create_test_secret_key();
        let _ = initialize_secret_key(sk);

        // Try with ciphertext shorter than 32 bytes
        let short_ciphertext = vec![0u8; 16];

        let creator_pubkey = PublicKey::from_bytes([2u8; 32]);
        let result = decrypt_user_message(&short_ciphertext, &creator_pubkey);

        assert_matches!(
            result,
            Err(SecretSharingError::EnvelopeInvalid(msg)) if msg.contains("too short")
        );
    }

    #[test]
    fn test_decrypt_user_message_invalid_utf8() {
        // Initialize key if not already done (safe to call multiple times)
        let sk = create_test_secret_key();
        let _ = initialize_secret_key(sk);

        // Create ciphertext that will decrypt to invalid UTF-8
        let mut invalid_ciphertext = vec![0u8; 32]; // 32-byte "enc"
        invalid_ciphertext.extend_from_slice(b"invalid data that will cause decryption to fail");

        let creator_pubkey = PublicKey::from_bytes([3u8; 32]);
        let result = decrypt_user_message(&invalid_ciphertext, &creator_pubkey);

        // This will likely fail during HPKE decryption since we're using dummy data
        // The exact error depends on the internal HPKE implementation
        assert!(result.is_err());
    }
}