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

//! Traits for secret sharing operations.
//!
//! This module defines the engine abstraction that powers our secret
//! distribution flow between TEEs.
//!
//! Key concepts:
//! - SecretSharingKey (SK): A long-lived X25519 keypair used by TEEs to
//!   decrypt user secrets (encrypted to `SK_pub` by clients).
//! - SecretSharingEngine: Provides two capabilities:
//!   1. Generate an SK inside a TEE.
//!   2. Wrap (encrypt) `SK_priv` for a recipient TEE and later unwrap
//!      it inside that recipient TEE.
//!
//! Implementations MUST ensure confidentiality and integrity of `SK_priv`,
//! bind any caller-provided AAD (Associated (Authenticated) Data) into the encryption/auth path,
//! and avoid leaking secret bytes via logs, panics, or formatting.
//!
//! The default engine in this crate uses HPKE
//! (Diffie-Hellman KEM(X25519,HKDF-SHA256) + ChaCha20-Poly1305) and represents the wrapped
//! key as a [`WrappedKeyEnvelope`].
use rialo_types::PublicKey;

use crate::{
    key::SecretSharingKey,
    types::{Aad, WrappedKeyEnvelope},
    SecretSharingError,
};

/// Engine interface for creating the long-lived SK and wrapping/unwrapping it
/// for transfer to other TEEs.
///
/// # Security expectations
///
/// - RNG: `generate` must draw entropy from a cryptographically secure RNG
///   available inside the enclave.
/// - Zeroization: Implementations should zeroize intermediate buffers that
///   temporarily hold `SK_priv`.
/// - AAD binding: `wrap_for_recipient` must bind the supplied [`Aad`] into
///   the AEAD/HPKE context so `unwrap_and_install` will fail if AAD differs.
/// - Recipient check: `unwrap_and_install` should verify that the envelope
///   is addressed to *this* TEE (e.g., compare recipient public key bytes)
///   before attempting decryption.
/// - Replay/misdelivery: Callers are encouraged to put monotonic or unique
///   values inside [`Aad`] (e.g., `request_id`), and engines should fully
///   authenticate that AAD to prevent replay across contexts.
/// - Error handling: Implementations must not leak secret bytes in error
///   messages or via panics.
pub trait SecretSharingEngine {
    /// Generate a fresh long-lived [`SecretSharingKey`] (SK).
    fn generate_secret_key(&self) -> Result<SecretSharingKey, SecretSharingError>;

    /// Return the raw `SK_pub` bytes for the given key.
    ///
    /// This will be published on-chain (with attestation) so users can
    /// encrypt their secrets to `SK_pub`.
    fn secret_key_public_key(&self, sk: &SecretSharingKey) -> PublicKey;

    /// Return this engine’s wrapping public key bytes (e.g., X25519).
    ///
    /// Every TEE will publish its own wrapping public key so peers
    /// can encrypt SK_priv to it.
    fn wrapping_public_key(&self) -> PublicKey;

    /// Encrypt (wrap) `SK_priv` for a recipient TEE, producing a [`WrappedKeyEnvelope`].
    ///
    /// - `recipient_wrapping_pub_x25519`: the recipient TEE’s wrapping/public key
    ///   that the engine will use to encrypt the payload (e.g., the recipient’s
    ///   HPKE KEM public key in the default implementation).
    /// - `aad`: caller-provided associated data (e.g., request id, authority
    ///   index, key id). MUST be bound to the encryption so tampering or
    ///   mismatch causes decryption failure.
    ///
    /// The resulting envelope includes enough information for the recipient to
    /// decrypt (e.g., HPKE encapped key) and metadata such as sender/recipient
    /// public keys for provenance.
    fn wrap_for_recipient(
        &self,
        sk: &SecretSharingKey,
        recipient_wrapping_pub_x25519: PublicKey,
        aad: Aad,
    ) -> Result<WrappedKeyEnvelope, SecretSharingError>;

    /// Decrypt (unwrap) a WrappedKeyEnvelope addressed to *this* TEE and get the [`SecretSharingKey`].
    ///
    /// Implementations should:
    /// - First verify the envelope’s recipient matches this engine’s identity.
    /// - Reconstruct the crypto context and authenticate the same `aad`.
    /// - Return an error on any failure (do NOT partially return secrets).
    ///
    /// The returned `SecretSharingKey` owns its secret and will zeroize on drop.
    fn unwrap_and_install(
        &self,
        env: &WrappedKeyEnvelope,
    ) -> Result<SecretSharingKey, SecretSharingError>;
}
