Claude/add enclave verify endpoint ac vqi

enclave/Cargo.toml

@@ -34,6 +34,7 @@ vsock = { version = "=0.5.2", default-features = false }
 zeroize = { version = "=1.8.2", default-features = false, features = ["zeroize_derive"] }
 
 [target.'cfg(target_env = "musl")'.dependencies]
+aws-nitro-enclaves-nsm-api = { version = "=0.4.0", default-features = false }
 mimalloc = { version = "=0.1.48", default-features = false, features = ["secure"] }
 
 [dev-dependencies]

enclave/src/constants.rs

@@ -26,3 +26,19 @@ pub const P521: &[u8; 10] = &[72, 80, 75, 69, 0, 18, 0, 3, 0, 2];
 
 pub const ENCODING_HEX: &str = "1";
 pub const ENCODING_BINARY: &str = "2";
+
+// NSM (Nitro Secure Module) constants for attestation
+
+/// Minimum nonce length in bytes (128 bits) per Trail of Bits recommendations.
+///
+/// Reference: <https://blog.trailofbits.com/2024/09/24/notes-on-aws-nitro-enclaves-attack-surface/>
+pub const MIN_NONCE_LENGTH: usize = 16;
+
+/// Maximum nonce length in bytes (NSM limit)
+pub const MAX_NONCE_LENGTH: usize = 512;
+
+/// Maximum user_data length in bytes (NSM limit)
+pub const MAX_USER_DATA_LENGTH: usize = 512;
+
+/// Maximum public_key length in bytes (NSM limit)
+pub const MAX_PUBLIC_KEY_LENGTH: usize = 1024;

enclave/src/lib.rs

@@ -8,5 +8,6 @@ pub mod functions;
 pub mod hpke;
 pub mod kms;
 pub mod models;
+pub mod nsm;
 pub mod protocol;
 pub mod utils;

enclave/src/main.rs

@@ -10,8 +10,13 @@ use anyhow::{Error, Result, anyhow};
 use enclave_vault::{
     constants::{ENCLAVE_PORT, MAX_CONCURRENT_CONNECTIONS},
     expressions::execute_expressions,
-    models::{EnclaveRequest, EnclaveResponse},
+    models::{
+        AttestationRequest, AttestationResponse, EnclaveRequest, EnclaveRequestType,
+        EnclaveResponse,
+    },
+    nsm,
     protocol::{recv_message, send_message},
+    utils::base64_decode,
 };
 use vsock::VsockListener;
 
@@ -20,11 +25,39 @@ use vsock::VsockListener;
 #[global_allocator]
 static GLOBAL: mimalloc::MiMalloc = mimalloc::MiMalloc;
 
+/// Parse the incoming payload, supporting both new tagged format and legacy format.
+///
+/// The new format uses a "type" tag to discriminate between request types:
+/// - `{"type": "decrypt", ...}` for decrypt requests
+/// - `{"type": "attestation", ...}` for attestation requests
+///
+/// For backward compatibility, payloads without a "type" field are treated
+/// as legacy decrypt requests.
 #[inline]
-fn parse_payload(payload_buffer: &[u8]) -> Result<EnclaveRequest> {
-    let payload: EnclaveRequest = serde_json::from_slice(payload_buffer)
+fn parse_payload(payload_buffer: &[u8]) -> Result<EnclaveRequestType> {
+    // First try to parse as the new tagged format
+    if let Ok(request_type) = serde_json::from_slice::<EnclaveRequestType>(payload_buffer) {
+        return Ok(request_type);
+    }
+
+    // Fall back to legacy format (EnclaveRequest without type tag)
+    let legacy_request: EnclaveRequest = serde_json::from_slice(payload_buffer)
         .map_err(|err| anyhow!("failed to deserialize payload: {err:?}"))?;
-    Ok(payload)
+
+    Ok(EnclaveRequestType::Decrypt(Box::new(legacy_request)))
+}
+
+/// Sanitizes error messages to prevent sensitive data leakage in logs.
+/// Removes potential field values, keys, or other sensitive content.
+#[inline]
+fn sanitize_error_message(err: &Error) -> String {
+    let msg = err.to_string();
+    // Truncate very long error messages that might contain data
+    if msg.len() > 200 {
+        format!("{}... (truncated)", &msg[..200])
+    } else {
+        msg
+    }
 }
 
 #[inline]
@@ -46,39 +79,17 @@ fn send_error<W: Write>(mut stream: W, err: Error) -> Result<()> {
     Ok(())
 }
 
-/// Sanitizes error messages to prevent sensitive data leakage in logs.
-/// Removes potential field values, keys, or other sensitive content.
-#[inline]
-fn sanitize_error_message(err: &Error) -> String {
-    let msg = err.to_string();
-    // Truncate very long error messages that might contain data
-    if msg.len() > 200 {
-        format!("{}... (truncated)", &msg[..200])
-    } else {
-        msg
-    }
-}
-
-fn handle_client<S: Read + Write>(mut stream: S) -> Result<()> {
-    println!("[enclave] handling client");
-
-    let payload: EnclaveRequest = match recv_message(&mut stream)
-        .map_err(|err| anyhow!("failed to receive message: {err:?}"))
-    {
-        Ok(payload_buffer) => match parse_payload(&payload_buffer) {
-            Ok(payload) => payload,
-            Err(err) => return send_error(stream, err),
-        },
-        Err(err) => return send_error(stream, err),
-    };
+/// Handle a decrypt request (existing functionality).
+fn handle_decrypt<S: Read + Write>(mut stream: S, request: EnclaveRequest) -> Result<()> {
+    println!("[enclave] handling decrypt request");
 
     // Decrypt the individual field values (uses rayon for parallelization internally)
-    let (decrypted_fields, errors) = match payload.decrypt_fields() {
+    let (decrypted_fields, errors) = match request.decrypt_fields() {
         Ok(result) => result,
         Err(err) => return send_error(stream, err),
     };
 
-    let final_fields = match payload.request.expressions {
+    let final_fields = match request.request.expressions {
         Some(expressions) => match execute_expressions(&decrypted_fields, &expressions) {
             Ok(fields) => fields,
             Err(err) => {
@@ -98,19 +109,106 @@ fn handle_client<S: Read + Write>(mut stream: S) -> Result<()> {
     let payload: String = serde_json::to_string(&response)
         .map_err(|err| anyhow!("failed to serialize response: {err:?}"))?;
 
-    println!("[enclave] sending response to parent");
+    println!("[enclave] sending decrypt response to parent");
 
     if let Err(err) = send_message(&mut stream, &payload)
         .map_err(|err| anyhow!("Failed to send message: {err:?}"))
     {
         return send_error(stream, err);
     }
 
-    println!("[enclave] finished client");
+    println!("[enclave] finished decrypt request");
 
     Ok(())
 }
 
+/// Handle an attestation request.
+fn handle_attestation<S: Read + Write>(mut stream: S, request: AttestationRequest) -> Result<()> {
+    println!("[enclave] handling attestation request");
+
+    // Decode nonce from base64
+    let nonce = match base64_decode(&request.nonce) {
+        Ok(n) => n,
+        Err(err) => {
+            let response = AttestationResponse::error(format!("invalid nonce base64: {err}"));
+            let payload = serde_json::to_string(&response)
+                .map_err(|err| anyhow!("failed to serialize response: {err:?}"))?;
+            send_message(&mut stream, &payload)?;
+            return Ok(());
+        }
+    };
+
+    // Decode optional user_data from base64
+    let user_data = match &request.user_data {
+        Some(ud) => match base64_decode(ud) {
+            Ok(d) => Some(d),
+            Err(err) => {
+                let response =
+                    AttestationResponse::error(format!("invalid user_data base64: {err}"));
+                let payload = serde_json::to_string(&response)
+                    .map_err(|err| anyhow!("failed to serialize response: {err:?}"))?;
+                send_message(&mut stream, &payload)?;
+                return Ok(());
+            }
+        },
+        None => None,
+    };
+
+    // Generate attestation document
+    let response = match nsm::get_attestation_document(
+        user_data.as_deref(),
+        Some(&nonce),
+        None, // public_key not used for this endpoint
+    ) {
+        Ok(document) => {
+            // Encode document as base64
+            let document_b64 = data_encoding::BASE64.encode(&document);
+            AttestationResponse::success(document_b64)
+        }
+        Err(err) => {
+            println!("[enclave error] attestation failed");
+            #[cfg(debug_assertions)]
+            println!("[enclave debug] attestation error: {:?}", err);
+            AttestationResponse::error(err.to_string())
+        }
+    };
+
+    let payload = serde_json::to_string(&response)
+        .map_err(|err| anyhow!("failed to serialize attestation response: {err:?}"))?;
+
+    println!("[enclave] sending attestation response to parent");
+
+    if let Err(err) = send_message(&mut stream, &payload)
+        .map_err(|err| anyhow!("Failed to send message: {err:?}"))
+    {
+        return send_error(stream, err);
+    }
+
+    println!("[enclave] finished attestation request");
+
+    Ok(())
+}
+
+fn handle_client<S: Read + Write>(mut stream: S) -> Result<()> {
+    println!("[enclave] handling client");
+
+    let request_type: EnclaveRequestType = match recv_message(&mut stream)
+        .map_err(|err| anyhow!("failed to receive message: {err:?}"))
+    {
+        Ok(payload_buffer) => match parse_payload(&payload_buffer) {
+            Ok(request) => request,
+            Err(err) => return send_error(stream, err),
+        },
+        Err(err) => return send_error(stream, err),
+    };
+
+    // Dispatch based on request type
+    match request_type {
+        EnclaveRequestType::Decrypt(request) => handle_decrypt(stream, *request),
+        EnclaveRequestType::Attestation(request) => handle_attestation(stream, request),
+    }
+}
+
 fn main() -> Result<()> {
     println!("[enclave] init");
 

enclave/src/models.rs

@@ -446,6 +446,84 @@ impl TryFrom<String> for Suite {
     }
 }
 
+// =============================================================================
+// Attestation Models
+// =============================================================================
+
+/// Request for attestation document generation.
+///
+/// This request is sent from the parent to the enclave to request an
+/// attestation document from the Nitro Secure Module.
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct AttestationRequest {
+    /// Nonce for freshness guarantee (base64 encoded, min 16 bytes decoded).
+    ///
+    /// Per Trail of Bits recommendations, a minimum nonce length is enforced
+    /// to prevent replay attacks.
+    pub nonce: String,
+
+    /// Optional application-specific data to include (base64 encoded, max 512 bytes decoded).
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub user_data: Option<String>,
+}
+
+/// Response containing an attestation document.
+///
+/// The document is a COSE Sign1 structure that can be verified by clients
+/// to prove the enclave's identity and configuration.
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct AttestationResponse {
+    /// Base64-encoded COSE Sign1 attestation document.
+    ///
+    /// This document contains:
+    /// - PCR values (enclave image hash, kernel, application)
+    /// - Module ID
+    /// - Timestamp (from hypervisor)
+    /// - Certificate chain to AWS Nitro root
+    /// - Echoed nonce, user_data, public_key (if provided)
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub document: Option<String>,
+
+    /// Error message if attestation generation failed.
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub error: Option<String>,
+}
+
+impl AttestationResponse {
+    /// Create a successful attestation response.
+    pub fn success(document: String) -> Self {
+        Self {
+            document: Some(document),
+            error: None,
+        }
+    }
+
+    /// Create an error attestation response.
+    pub fn error(message: impl Into<String>) -> Self {
+        Self {
+            document: None,
+            error: Some(message.into()),
+        }
+    }
+}
+
+/// Request envelope that discriminates between different request types.
+///
+/// This allows the enclave to handle multiple types of requests over
+/// the same vsock connection.
+#[derive(Debug, Clone, Serialize, Deserialize)]
+#[serde(tag = "type")]
+pub enum EnclaveRequestType {
+    /// Decrypt request (existing functionality)
+    /// Boxed to reduce enum size difference with AttestationRequest
+    #[serde(rename = "decrypt")]
+    Decrypt(Box<EnclaveRequest>),
+
+    /// Attestation document request
+    #[serde(rename = "attestation")]
+    Attestation(AttestationRequest),
+}
+
 #[cfg(test)]
 #[allow(clippy::unwrap_used, clippy::expect_used, clippy::indexing_slicing)]
 mod tests {