newton-enclave 0.4.16

use alloy::primitives::{Address, Bytes, FixedBytes};
use newton_aggregator::{rpc_server::PublicShare, PartialDecryptionData};
use newton_core::{
    common::ResolvedPolicyInputs, crypto::SecureEnvelope, newton_prover_task_manager::NewtonMessage, TaskId,
};
use serde::{Deserialize, Serialize};

/// wire protocol version. bump on any change to request/response struct field order or types —
/// bincode is positional, not self-describing, so silent deserialization corruption is otherwise
/// possible. the version check in `handle_request` surfaces mismatches as a clean error.
pub const ENCLAVE_PROTOCOL_VERSION: u16 = 3;
/// fixed vsock port for compute traffic.
pub const VSOCK_PORT_COMPUTE: u32 = 5005;
/// fixed vsock port for egress HTTP proxy (enclave → host → external API).
pub const VSOCK_PORT_EGRESS: u32 = 5006;
/// vsock CID of the parent instance (host) from inside an AWS Nitro enclave.
pub const VSOCK_CID_HOST: u32 = 3;
/// largest accepted framed message.
pub const MAX_FRAME_LEN: usize = 16 * 1024 * 1024;
/// largest accepted egress response body (1 MB).
pub const MAX_EGRESS_RESPONSE_BYTES: usize = 1024 * 1024;
/// maximum HTTP requests per WASM execution via egress proxy (per-connection).
/// intentionally lower than `DataProviderConfig.max_http_calls` (50) because this
/// is a security boundary — the proxy is the last enforcement point before external
/// network access. real data providers make 1-3 calls; 10 is generous headroom.
/// the enclave's per-plugin `HttpProvider.max_calls` (configurable per PrepareEvalRequest)
/// is the primary control; this is a secondary safety net on the proxy side.
pub const MAX_EGRESS_REQUESTS_PER_EXECUTION: u32 = 10;

/// encrypted envelope plus optional domain metadata.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct EnclaveEnvelope {
    /// off-chain reference id, if this envelope came from a ref.
    pub ref_id: Option<String>,
    /// identity or confidential domain.
    pub domain: Option<FixedBytes<32>>,
    /// encrypted payload.
    pub envelope: SecureEnvelope,
}

/// policy evaluation request for secret-bearing tasks.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct EnclaveEvalRequest {
    /// task id.
    pub task_id: TaskId,
    /// policy client address.
    pub policy_client: Address,
    /// user intent.
    pub intent: NewtonMessage::Intent,
    /// user intent signature.
    pub intent_signature: Bytes,
    /// consensus policy task data.
    pub policy_task_data: NewtonMessage::PolicyTaskData,
    /// public policy inputs resolved by the host before enclave evaluation.
    pub resolved_policy: ResolvedPolicyInputs,
    /// task initialization timestamp.
    pub initialization_timestamp: u64,
    /// verified tls proof data.
    pub proof_data: Option<serde_json::Value>,
    /// inline encrypted privacy envelopes.
    pub ephemeral_envelopes: Vec<SecureEnvelope>,
    /// identity encrypted envelopes.
    pub identity_envelopes: Vec<EnclaveEnvelope>,
    /// confidential encrypted envelopes.
    pub confidential_envelopes: Vec<EnclaveEnvelope>,
    /// threshold inputs, when threshold envelopes are used.
    pub threshold: Option<ThresholdEvalInput>,
}

/// threshold decryption inputs for enclave evaluation.
///
/// encrypted peer partials are the sole input path. the enclave decrypts
/// them using its own ephemeral private key, splits by envelope type
/// (ephemeral/identity/confidential), then feeds into the threshold
/// decrypt pipeline. plaintext partials never exist outside the enclave.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ThresholdEvalInput {
    /// task id for AAD reconstruction when decrypting peer partials.
    pub task_id: TaskId,
    /// encrypted peer partial DH blobs from all operators for all envelope types.
    /// each blob is HPKE-encrypted to this enclave's ephemeral pubkey.
    pub encrypted_peer_partials: Vec<EncryptedPartialDH>,
    /// number of ephemeral envelopes (partials[0..eph_count] are ephemeral).
    pub ephemeral_count: usize,
    /// number of identity envelopes (partials[eph_count..eph_count+id_count] are identity).
    pub identity_count: usize,
    /// number of confidential envelopes (remainder are confidential).
    pub confidential_count: usize,
    /// operator public shares.
    pub public_shares: Vec<PublicShare>,
    /// threshold config.
    pub config: EnclaveThresholdConfig,
}

/// threshold config with named fields for wire stability.
#[derive(Debug, Clone, Copy, Serialize, Deserialize)]
pub struct EnclaveThresholdConfig {
    /// minimum partials needed.
    pub threshold: u32,
    /// total share count.
    pub total: u32,
}

/// policy evaluation response.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct EnclaveEvalResponse {
    /// boolean policy result.
    pub verified: bool,
}

pub use newton_core::crypto::encrypted_partial::{EnclaveOperatorId, EncryptedPartialDH};

/// threshold partial dh request.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct EnclavePartialDhRequest {
    /// task id for AAD binding (prevents cross-task replay of encrypted partials).
    pub task_id: TaskId,
    /// hpke enc points from envelopes.
    pub enc_points: Vec<Vec<u8>>,
    /// peer enclave ephemeral X25519 pubkeys (32 bytes each) for per-peer encryption.
    /// the enclave encrypts its partial DH outputs to each peer so the gateway cannot read them.
    pub peer_enclave_pubkeys: Vec<(EnclaveOperatorId, [u8; 32])>,
}

/// threshold partial dh response.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct EnclavePartialDhResponse {
    /// per-peer encrypted partial DH blobs. one entry per peer enclave.
    pub encrypted_partials: Vec<EncryptedPartialDH>,
}

/// request for an attestation document from the enclave's Nitro Secure Module.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct GetAttestationRequest {
    /// caller-provided nonce bound into the attestation document's `user_data` field.
    /// currently: random 32 bytes at init time for freshness + PCR0 extraction.
    /// Phase 3: per-task `keccak256(task_id || response_digest)` for hardware-bound
    /// task output binding (requires NSM hardware and on-chain COSE_Sign1 verifier).
    pub nonce: Vec<u8>,
}

/// attestation document returned by the enclave.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct GetAttestationResponse {
    /// raw attestation document bytes.
    /// real Nitro: COSE_Sign1-wrapped CBOR (parseable via `attestation::nitro`).
    /// loopback: `NEWT_LBK` magic + 48-byte zero PCR0 + nonce (parseable via `attestation::sentinel`).
    pub attestation_doc: Vec<u8>,
}

/// framed request body sent over vsock.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct EnclaveWireRequest {
    /// protocol version.
    pub version: u16,
    /// request id chosen by caller.
    pub request_id: u64,
    /// request body.
    pub body: EnclaveWireRequestBody,
}

/// framed response body sent over vsock.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct EnclaveWireResponse {
    /// protocol version.
    pub version: u16,
    /// matching request id.
    pub request_id: u64,
    /// response body.
    pub body: EnclaveWireResponseBody,
}

/// supported enclave requests.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum EnclaveWireRequestBody {
    /// check server liveness.
    Health,
    /// initialize enclave-held keys.
    Init(EnclaveInitRequest),
    /// return enclave-held HPKE public key.
    GetPublicKey,
    /// evaluate policy inside enclave.
    Evaluate(Box<EnclaveEvalRequest>),
    /// compute threshold partial dh values.
    PartialDh(EnclavePartialDhRequest),
    /// request a fresh attestation document from the Nitro Secure Module.
    GetAttestation(GetAttestationRequest),
    /// WASM execution + secrets decryption + partial DH in one call (Phase 2).
    PrepareEval(Box<PrepareEvalRequest>),
}

/// supported enclave responses.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum EnclaveWireResponseBody {
    /// health ok.
    Health,
    /// init ok.
    Init,
    /// HPKE public key bytes.
    PublicKey(Vec<u8>),
    /// policy evaluation response.
    Evaluate(EnclaveEvalResponse),
    /// threshold partial dh response.
    PartialDh(EnclavePartialDhResponse),
    /// attestation document from the Nitro Secure Module.
    Attestation(GetAttestationResponse),
    /// WASM + secrets + partial DH combined response (Phase 2).
    PrepareEval(PrepareEvalResponse),
    /// typed failure.
    Error(EnclaveWireError),
}

/// key initialization request.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct EnclaveInitRequest {
    /// raw centralized HPKE private key bytes (legacy path, Phase 1).
    pub hpke_private_key: Option<Vec<u8>>,
    /// encrypted FROST threshold share keystore JSON.
    pub threshold_keystore: Option<Vec<u8>>,
    /// password for the threshold share keystore.
    pub threshold_keystore_password: Option<Vec<u8>>,
    /// KMS-encrypted seed for HPKE key derivation (Phase 2 path).
    /// when present, the enclave decrypts via KMS attestation and derives
    /// the HPKE keypair via HKDF. takes precedence over `hpke_private_key`.
    pub kms_seed_ciphertext: Option<Vec<u8>>,
}

// ---------------------------------------------------------------------------
// PrepareEval: WASM execution + secrets + partial DH (Phase 2)
// ---------------------------------------------------------------------------

/// Combined Prepare phase request: WASM execution with secrets + partial DH.
///
/// The enclave decrypts secrets, executes WASM plugins in parallel via the
/// Egress proxy, then computes encrypted per-peer partial DH outputs.
///
/// Invariants:
/// - `task_id` is used as AAD for partial DH peer encryption (prevents cross-task replay).
/// - Response `plugin_results` preserves request order.
/// - Secrets are zeroized on drop inside the enclave via `Zeroizing<Vec<u8>>`.
/// - Plugin count is capped at `MAX_PLUGINS_PER_BATCH` by the executor.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct PrepareEvalRequest {
    /// task id for partial DH AAD binding.
    pub task_id: TaskId,
    /// WASM plugins to execute (each with encrypted secrets).
    pub wasm_plugins: Vec<WasmPluginInput>,
    /// enc_points for partial DH computation (all envelope types combined).
    pub enc_points: Vec<Vec<u8>>,
    /// peer enclave X25519 pubkeys (32 bytes each) for per-peer encrypted partials.
    pub peer_enclave_pubkeys: Vec<(EnclaveOperatorId, [u8; 32])>,
}

/// Input for one WASM plugin execution inside the enclave.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct WasmPluginInput {
    /// raw WASM component binary bytes.
    pub wasm_bytes: Vec<u8>,
    /// expected `keccak256(wasm_bytes)` for integrity verification.
    pub wasm_code_hash: FixedBytes<32>,
    /// JSON input string for the WASM `run()` function.
    pub wasm_args: String,
    /// HPKE-encrypted secrets envelope. Decrypted inside enclave.
    /// `None` for plugins that do not use secrets (skips HPKE decrypt).
    pub encrypted_secrets: Option<SecureEnvelope>,
    /// max HTTP calls this plugin may make via egress.
    pub max_http_calls: u32,
}

/// Combined Prepare phase response.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct PrepareEvalResponse {
    /// per-plugin WASM outputs, same order as request.
    pub plugin_results: Vec<WasmPluginOutput>,
    /// per-peer encrypted partial DH blobs.
    pub encrypted_partials: Vec<EncryptedPartialDH>,
}

/// Output from one WASM plugin execution.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct WasmPluginOutput {
    /// WASM run() result — Ok(json_string) or classified error.
    pub result: Result<String, WasmPluginError>,
}

/// Classified WASM plugin execution error.
#[derive(Debug, Clone, thiserror::Error, Serialize, Deserialize)]
pub enum WasmPluginError {
    /// HPKE secrets decryption failed (AAD mismatch, corrupted envelope).
    #[error("secret decrypt failed: {0}")]
    SecretDecryptFailed(String),
    /// WASM compilation or instantiation failed (bad binary, missing imports).
    #[error("compilation failed: {0}")]
    CompilationFailed(String),
    /// WASM `run()` returned Err (plugin-level logic error).
    #[error("execution failed: {0}")]
    ExecutionFailed(String),
    /// Plugin exceeded its time budget.
    #[error("timeout: {0}")]
    Timeout(String),
    /// Plugin exceeded WASM binary size limit.
    #[error("oversized binary: {0}")]
    OversizedBinary(String),
    /// Batch-level resource exhaustion (semaphore closed, egress budget).
    #[error("resource exhausted: {0}")]
    ResourceExhausted(String),
    /// Plugin task panicked or was cancelled.
    #[error("cancelled: {0}")]
    Cancelled(String),
}

// ---------------------------------------------------------------------------
// Egress HTTP proxy wire types (VSOCK_PORT_EGRESS)
// ---------------------------------------------------------------------------

/// HTTP request from enclave to host egress proxy.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct EgressRequest {
    /// full URL including scheme (e.g. `https://api.example.com/v1/data`).
    pub url: String,
    /// HTTP method (GET, POST, PUT, DELETE).
    pub method: String,
    /// request headers as key-value pairs.
    pub headers: Vec<(String, String)>,
    /// optional request body.
    pub body: Option<Vec<u8>>,
}

/// HTTP response from host egress proxy back to enclave.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct EgressResponse {
    /// HTTP status code.
    pub status: u16,
    /// response headers as key-value pairs.
    pub headers: Vec<(String, String)>,
    /// response body (truncated to `MAX_EGRESS_RESPONSE_BYTES`).
    pub body: Vec<u8>,
}

/// egress proxy error sent back to enclave when the request cannot be fulfilled.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct EgressError {
    /// error message describing why the request was rejected.
    pub message: String,
}

/// wire envelope for egress proxy communication.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum EgressWireMessage {
    /// enclave → host: HTTP request.
    Request(EgressRequest),
    /// host → enclave: successful HTTP response.
    Response(EgressResponse),
    /// host → enclave: request rejected (security, rate limit, etc).
    Error(EgressError),
}

/// wire-safe enclave error.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct EnclaveWireError {
    /// machine-readable error code.
    pub code: EnclaveErrorCode,
    /// sanitized error text.
    pub message: String,
    /// whether caller may retry.
    pub retryable: bool,
}

/// enclave error class.
#[derive(Debug, Clone, Copy, Serialize, Deserialize)]
pub enum EnclaveErrorCode {
    /// invalid request.
    InvalidRequest,
    /// protocol version is unsupported.
    UnsupportedVersion,
    /// required input missing.
    MissingInput,
    /// decrypt failed.
    DecryptFailed,
    /// threshold operation failed.
    ThresholdFailed,
    /// policy evaluation failed.
    PolicyEvalFailed,
    /// enclave is not initialized.
    Uninitialized,
    /// internal failure.
    Internal,
    /// KMS or HKDF key derivation failed.
    KeyDerivation,
    /// NSM attestation request failed.
    AttestationFailed,
}