feat: add agent attestation (#188)

Co-authored-by: Robert Yan <46699230+think-in-universe@users.noreply.github.com>

Author: nickpismenkov

crates/api/src/models.rs

@@ -144,12 +144,48 @@ pub struct ModelAttestation {
     pub info: Option<serde_json::Value>,
 }
 
+/// Agent attestation from agent instance
+#[derive(Debug, Clone, Serialize, Deserialize, ToSchema)]
+pub struct AgentAttestation {
+    /// Agent instance name
+    pub name: String,
+
+    /// Container image digest
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub image_digest: Option<String>,
+
+    /// TDX event log
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub event_log: Option<String>,
+
+    /// Additional TDX/tappd info (structured JSON)
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub info: Option<serde_json::Value>,
+
+    /// Intel TDX quote in hex format
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub intel_quote: Option<String>,
+
+    /// Request nonce
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub request_nonce: Option<String>,
+
+    /// TLS certificate
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub tls_certificate: Option<String>,
+
+    /// TLS certificate fingerprint
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub tls_certificate_fingerprint: Option<String>,
+}
+
 /// Complete attestation report combining all layers
 ///
 /// This report proves the entire trust chain:
 /// 1. This chat-api service runs in a TEE (your_gateway_attestation)
-/// 2. The cloud-api dependency runs in a TEE (cloud_api_gateway_attestation)  
+/// 2. The cloud-api dependency runs in a TEE (cloud_api_gateway_attestation)
 /// 3. The model inference providers run on trusted hardware (model_attestations)
+/// 4. Optional agent instance attestations when agent parameter is provided
 #[derive(Debug, Clone, Serialize, Deserialize, ToSchema)]
 pub struct CombinedAttestationReport {
     /// This chat-api's own CPU attestation (proves this service runs in a TEE)
@@ -161,6 +197,10 @@ pub struct CombinedAttestationReport {
     /// Model provider attestations (can be multiple when routing to different models)
     #[serde(skip_serializing_if = "Option::is_none")]
     pub model_attestations: Option<Vec<ModelAttestation>>,
+
+    /// Agent instance attestations (included when agent query parameter is provided)
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub agent_attestations: Option<Vec<AgentAttestation>>,
 }
 
 /// Attestation report structure from proxy_service

crates/api/src/routes/attestation.rs

@@ -4,13 +4,19 @@ use crate::{
     state::AppState,
     ApiError,
 };
-use axum::{extract::Query, extract::State, response::Json, routing::get, Router};
+use axum::{
+    extract::{Query, State},
+    response::Json,
+    routing::get,
+    Router,
+};
 use futures::TryStreamExt;
 use http::Method;
+use reqwest::Client;
 use serde::{Deserialize, Serialize};
 use services::vpc::load_vpc_info;
 
-#[derive(Debug, Deserialize, Serialize)]
+#[derive(Debug, Clone, Deserialize, Serialize)]
 pub struct AttestationQuery {
     /// Optional model name to get specific attestations
     #[serde(skip_serializing_if = "Option::is_none")]
@@ -27,6 +33,10 @@ pub struct AttestationQuery {
     /// Signing address
     #[serde(skip_serializing_if = "Option::is_none")]
     pub signing_address: Option<String>,
+
+    /// Optional agent instance ID to get agent attestations
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub agent: Option<String>,
 }
 
 /// GET /v1/attestation/report
@@ -46,7 +56,8 @@ pub struct AttestationQuery {
         ("model" = Option<String>, Query, description = "Optional model name to filter model attestations"),
         ("signing_algo" = Option<String>, Query, description = "Signing algorithm: 'ecdsa' or 'ed25519'"),
         ("nonce" = Option<String>, Query, description = "64 length (32 bytes) hex string"),
-        ("signing_address" = Option<String>, Query, description = "Query the attestation of the specific model that owns this signing address")
+        ("signing_address" = Option<String>, Query, description = "Query the attestation of the specific model that owns this signing address"),
+        ("agent" = Option<String>, Query, description = "Optional agent instance ID to include agent attestations in response")
     ),
     responses(
         (status = 200, description = "Combined attestation report", body = CombinedAttestationReport),
@@ -58,7 +69,11 @@ pub async fn get_attestation_report(
     State(app_state): State<AppState>,
     Query(params): Query<AttestationQuery>,
 ) -> Result<Json<CombinedAttestationReport>, ApiError> {
-    let query = serde_urlencoded::to_string(&params).expect("Failed to serialize query string");
+    // Exclude agent parameter from cloud-api query since it's not relevant there
+    let mut cloud_api_params = params.clone();
+    cloud_api_params.agent = None;
+    let query =
+        serde_urlencoded::to_string(&cloud_api_params).expect("Failed to serialize query string");
 
     // Build the path for proxy_service attestation endpoint
     let path = format!("attestation/report?{}", query);
@@ -131,7 +146,7 @@ pub async fn get_attestation_report(
             signing_algo: None,
             intel_quote: "0x1234567890abcdef".to_string(),
             event_log: None,
-            request_nonce,
+            request_nonce: request_nonce.clone(),
             info: None,
             vpc: vpc_info,
         }
@@ -163,7 +178,7 @@ pub async fn get_attestation_report(
             info: Some(serde_json::to_value(info).map_err(|_| {
                 ApiError::internal_server_error("Failed to serialize attestation info")
             })?),
-            request_nonce,
+            request_nonce: request_nonce.clone(),
             vpc: vpc_info,
         }
     };
@@ -172,15 +187,252 @@ pub async fn get_attestation_report(
 
     let model_attestations = proxy_report.model_attestations;
 
+    // Fetch agent attestations if agent parameter is provided (no user auth required)
+    let agent_attestations = if let Some(agent_id) = &params.agent {
+        match fetch_agent_attestations(&app_state, agent_id, &request_nonce).await {
+            Ok(attestations) => Some(attestations),
+            Err(e) => {
+                tracing::warn!("Failed to fetch agent attestations: {:?}", e);
+                // Don't fail the entire request if agent attestation fetch fails
+                None
+            }
+        }
+    } else {
+        None
+    };
+
     let report = CombinedAttestationReport {
         chat_api_gateway_attestation,
         cloud_api_gateway_attestation,
         model_attestations,
+        agent_attestations,
     };
 
     Ok(Json(report))
 }
 
+/// Fetch agent attestations from compose-api
+#[derive(Debug, Deserialize)]
+struct AgentAttestationResponse {
+    event_log: Option<String>,
+    quote: Option<String>,
+    #[serde(default)]
+    info: Option<serde_json::Value>,
+    tls_certificate: Option<String>,
+    tls_certificate_fingerprint: Option<String>,
+}
+
+#[derive(Debug, Deserialize)]
+struct AgentInstanceAttestationResponse {
+    image_digest: Option<String>,
+    name: String,
+}
+
+/// Validate nonce is properly formatted and reasonable length (replay protection)
+fn validate_nonce(nonce: &str) -> Result<(), ApiError> {
+    // Nonce should be a valid hex string of reasonable length (64 chars = 32 bytes)
+    const EXPECTED_NONCE_LEN: usize = 64;
+    const MAX_NONCE_LEN: usize = 256;
+
+    if nonce.len() > MAX_NONCE_LEN {
+        tracing::warn!("Nonce exceeds maximum length: {}", nonce.len());
+        return Err(ApiError::bad_request("Nonce is too long"));
+    }
+
+    if !nonce.chars().all(|c| c.is_ascii_hexdigit()) {
+        tracing::warn!("Nonce contains non-hex characters");
+        return Err(ApiError::bad_request("Nonce must be a valid hex string"));
+    }
+
+    if nonce.len() != EXPECTED_NONCE_LEN {
+        tracing::warn!(
+            "Nonce has unexpected length: {} (expected {})",
+            nonce.len(),
+            EXPECTED_NONCE_LEN
+        );
+        return Err(ApiError::bad_request(format!(
+            "Nonce must be exactly {} characters",
+            EXPECTED_NONCE_LEN
+        )));
+    }
+
+    Ok(())
+}
+
+/// Validate instance name doesn't contain path traversal sequences
+fn validate_instance_name(name: &str) -> Result<(), ApiError> {
+    // Reject names containing path traversal sequences
+    if name.contains("..") || name.contains("/") || name.contains("\\") {
+        tracing::warn!("Instance name contains invalid characters: {}", name);
+        return Err(ApiError::bad_request(
+            "Instance name contains invalid characters",
+        ));
+    }
+
+    if name.is_empty() {
+        return Err(ApiError::bad_request("Instance name cannot be empty"));
+    }
+
+    Ok(())
+}
+
+/// Build full URL for agent manager request (handles base URL with/without trailing slash)
+fn build_manager_url(base_url: &str, path: &str) -> Result<String, ApiError> {
+    let base = url::Url::parse(base_url).map_err(|e| {
+        tracing::error!("Invalid agent manager URL {}: {}", base_url, e);
+        ApiError::internal_server_error("Invalid agent manager URL")
+    })?;
+    let full = base.join(path).map_err(|e| {
+        tracing::error!("Failed to build manager URL: {}", e);
+        ApiError::internal_server_error("Failed to build manager URL")
+    })?;
+    Ok(full.to_string())
+}
+
+/// Helper to handle HTTP response from agent manager (status check + body)
+async fn handle_manager_response(
+    response: reqwest::Response,
+    context: &str,
+) -> Result<bytes::Bytes, ApiError> {
+    let status = response.status();
+    if !status.is_success() {
+        tracing::error!(
+            "Agent manager returned error status {} for {}",
+            status,
+            context
+        );
+        return Err(ApiError::service_unavailable(format!(
+            "{} service returned error: {}",
+            context, status
+        )));
+    }
+    response.bytes().await.map_err(|e| {
+        tracing::error!("Failed to read {} response: {}", context, e);
+        ApiError::internal_server_error(format!("Failed to read {} response", context))
+    })
+}
+
+async fn fetch_agent_attestations(
+    app_state: &AppState,
+    agent_id: &str,
+    request_nonce: &str,
+) -> Result<Vec<crate::models::AgentAttestation>, ApiError> {
+    use uuid::Uuid;
+
+    // Security: Validate nonce to prevent panic/DoS from malformed input
+    validate_nonce(request_nonce)?;
+
+    // Parse the agent_id as UUID
+    let agent_uuid = Uuid::parse_str(agent_id).map_err(|e| {
+        tracing::error!("Invalid agent ID format: {}", e);
+        ApiError::bad_request(format!("Invalid agent ID format: {}", e))
+    })?;
+
+    // Fetch the agent instance from database (no user_id check - attestation is public)
+    let agent_instance = app_state
+        .agent_repository
+        .get_instance(agent_uuid)
+        .await
+        .map_err(|e| {
+            tracing::error!("Failed to fetch agent instance from database: {}", e);
+            ApiError::internal_server_error("Failed to fetch agent instance")
+        })?
+        .ok_or_else(|| {
+            tracing::warn!("Agent instance not found: {}", agent_id);
+            ApiError::not_found("Agent instance not found")
+        })?;
+
+    // Security: Validate instance name to prevent path traversal attacks
+    validate_instance_name(&agent_instance.name)?;
+
+    // Get the agent manager URL - each instance is hosted on a specific manager
+    let manager_base_url = agent_instance
+        .agent_api_base_url
+        .as_deref()
+        .ok_or_else(|| {
+            tracing::warn!("Agent instance has no agent_api_base_url: {}", agent_id);
+            ApiError::bad_gateway("Agent instance has no manager URL; cannot fetch attestation")
+        })?;
+
+    let instance_name = &agent_instance.name;
+    // URL-encode instance name for safe URL construction
+    let encoded_instance_name = urlencoding::encode(instance_name);
+
+    // Build URLs for the manager that hosts this instance
+    // NOTE: Nonce is critical for replay protection - bind the quote to the client's nonce
+    let attestation_url = build_manager_url(
+        manager_base_url,
+        &format!("attestation/report?nonce={}", request_nonce),
+    )?;
+    let instance_attestation_url = build_manager_url(
+        manager_base_url,
+        &format!("instances/{}/attestation", encoded_instance_name),
+    )?;
+
+    // Fetch both attestations concurrently from the corresponding agent manager
+    let http_client = Client::builder()
+        .timeout(std::time::Duration::from_secs(30))
+        .build()
+        .map_err(|e| {
+            tracing::error!("Failed to create HTTP client: {}", e);
+            ApiError::internal_server_error("Failed to create HTTP client")
+        })?;
+
+    let (attestation_response, instance_response) = tokio::join!(
+        http_client.get(&attestation_url).send(),
+        http_client.get(&instance_attestation_url).send(),
+    );
+
+    let attestation_response = attestation_response.map_err(|e| {
+        tracing::error!(
+            "Failed to fetch agent attestation from manager {}: {}",
+            manager_base_url,
+            e
+        );
+        ApiError::bad_gateway(format!("Failed to fetch agent attestation: {}", e))
+    })?;
+
+    let attestation_bytes =
+        handle_manager_response(attestation_response, "Agent attestation").await?;
+
+    let attestation_data: AgentAttestationResponse = serde_json::from_slice(&attestation_bytes)
+        .map_err(|e| {
+            tracing::error!("Failed to parse agent attestation response: {}", e);
+            ApiError::internal_server_error("Failed to parse agent attestation")
+        })?;
+
+    let instance_response = instance_response.map_err(|e| {
+        tracing::error!(
+            "Failed to fetch instance attestation from manager {}: {}",
+            manager_base_url,
+            e
+        );
+        ApiError::bad_gateway(format!("Failed to fetch instance attestation: {}", e))
+    })?;
+
+    let instance_bytes = handle_manager_response(instance_response, "Instance attestation").await?;
+
+    let instance_data: AgentInstanceAttestationResponse = serde_json::from_slice(&instance_bytes)
+        .map_err(|e| {
+        tracing::error!("Failed to parse instance attestation response: {}", e);
+        ApiError::internal_server_error("Failed to parse instance attestation")
+    })?;
+
+    // Combine the data
+    let agent_attestation = crate::models::AgentAttestation {
+        name: instance_data.name,
+        image_digest: instance_data.image_digest,
+        event_log: attestation_data.event_log,
+        info: attestation_data.info,
+        intel_quote: attestation_data.quote,
+        request_nonce: Some(request_nonce.to_string()),
+        tls_certificate: attestation_data.tls_certificate,
+        tls_certificate_fingerprint: attestation_data.tls_certificate_fingerprint,
+    };
+
+    Ok(vec![agent_attestation])
+}
+
 /// Create the attestation router
 pub fn create_attestation_router() -> Router<AppState> {
     Router::new().route("/v1/attestation/report", get(get_attestation_report))