Merge pull request #2 from cohere-ai/nvidia-tdx-attestation

feat(api-server-rest): add additional evidence endpoint

Author: yousef-cohere

Cargo.lock

@@ -11,6 +11,7 @@ async-trait.workspace = true
 base64.workspace = true
 clap = { workspace = true, features = ["derive"] }
 form_urlencoded = "1.2.2"
+hex.workspace = true
 hyper = { version = "0.14.27", features = ["server", "http1", "runtime"] }
 serde.workspace = true
 serde_json.workspace = true

api-server-rest/Cargo.toml

@@ -31,7 +31,8 @@ fn _token() {}
     get,
     path = "/aa/evidence",
     params(
-        ("runtime_data" = String, Query, description = "Runtime Data")
+        ("runtime_data" = String, Query, description = "Runtime Data"),
+        ("encoding" = Option<String>, Query, description = "Encoding of runtime_data: 'hex', 'base64', or omit for raw UTF-8 string")
     ),
     responses(
         (status = 200, description = "success response",
@@ -46,6 +47,26 @@ fn _token() {}
 )]
 fn _evidence() {}
 
+#[utoipa::path(
+    get,
+    path = "/aa/additional_evidence",
+    params(
+        ("runtime_data" = String, Query, description = "Runtime Data"),
+        ("encoding" = Option<String>, Query, description = "Encoding of runtime_data: 'hex', 'base64', or omit for raw UTF-8 string")
+    ),
+    responses(
+        (status = 200, description = "success response",
+                content_type = "application/octet-stream",
+                body = String,
+                example = json!({"svn":"1","report_data":"eHh4eA=="})),
+        (status = 400, description = "bad request for invalid query param"),
+        (status = 403, description = "forbid external access"),
+        (status = 404, description = "resource not found"),
+        (status = 405, description = "only Get method allowed")
+    )
+)]
+fn _additional_evidence() {}
+
 #[derive(ToSchema)]
 pub struct AaelEvent {
     /// Attestation Agent Event Log Domain
@@ -127,7 +148,7 @@ fn generate_openapi_document() -> std::io::Result<()> {
         (url = "http://127.0.0.1:8006", description = "CoCo RESTful API")
      ),
 
-    paths(_token, _evidence, _aael, _resource, _version)
+    paths(_token, _evidence, _additional_evidence, _aael, _resource, _version)
  )]
     struct ApiDoc;
     let mut file = File::create("openapi/api.json")?;

api-server-rest/build.rs

@@ -49,6 +49,60 @@
         }
       }
     },
+    "/aa/additional_evidence": {
+      "get": {
+        "tags": [],
+        "operationId": "_additional_evidence",
+        "parameters": [
+          {
+            "name": "runtime_data",
+            "in": "query",
+            "description": "Runtime Data",
+            "required": true,
+            "schema": {
+              "type": "string"
+            }
+          },
+          {
+            "name": "encoding",
+            "in": "query",
+            "description": "Encoding of runtime_data: 'hex', 'base64', or omit for raw UTF-8 string",
+            "required": false,
+            "schema": {
+              "type": "string"
+            }
+          }
+        ],
+        "responses": {
+          "200": {
+            "description": "success response",
+            "content": {
+              "application/octet-stream": {
+                "schema": {
+                  "type": "string"
+                },
+                "example": {
+                  "report_data": "eHh4eA==",
+                  "svn": "1"
+                }
+              }
+            }
+          },
+          "400": {
+            "description": "bad request for invalid query param"
+          },
+          "403": {
+            "description": "forbid external access"
+          },
+          "404": {
+            "description": "resource not found"
+          },
+          "405": {
+            "description": "only Get method allowed"
+          }
+        }
+      }
+    },
     "/aa/evidence": {
       "get": {
         "tags": [],
@@ -62,6 +116,15 @@
             "schema": {
               "type": "string"
             }
+          },
+          {
+            "name": "encoding",
+            "in": "query",
+            "description": "Encoding of runtime_data: 'hex', 'base64', or omit for raw UTF-8 string",
+            "required": false,
+            "schema": {
+              "type": "string"
+            }
           }
         ],
         "responses": {

api-server-rest/openapi/api.json

@@ -5,8 +5,8 @@
 
 use anyhow::*;
 use protos::ttrpc::aa::attestation_agent::{
-    ExtendRuntimeMeasurementRequest, GetAdditionalTeesRequest, GetEvidenceRequest,
-    GetTeeTypeRequest, GetTokenRequest,
+    ExtendRuntimeMeasurementRequest, GetAdditionalEvidenceRequest, GetAdditionalTeesRequest,
+    GetEvidenceRequest, GetTeeTypeRequest, GetTokenRequest,
 };
 use protos::ttrpc::aa::attestation_agent_ttrpc::AttestationAgentServiceClient;
 use serde::Deserialize;
@@ -19,6 +19,7 @@ pub const AA_ROOT: &str = "/aa";
 /// URL for querying CDH get resource API
 pub const AA_TOKEN_URL: &str = "/token";
 pub const AA_EVIDENCE_URL: &str = "/evidence";
+pub const AA_ADDITIONAL_EVIDENCE_URL: &str = "/additional_evidence";
 pub const AA_AAEL_URL: &str = "/aael";
 
 pub struct AAClient {
@@ -66,6 +67,18 @@ impl AAClient {
         Ok(res.Evidence)
     }
 
+    pub async fn get_additional_evidence(&self, runtime_data: &[u8]) -> Result<Vec<u8>> {
+        let req = GetAdditionalEvidenceRequest {
+            RuntimeData: runtime_data.to_vec(),
+            ..Default::default()
+        };
+        let res = self
+            .client
+            .get_additional_evidence(ttrpc::context::with_timeout(TTRPC_TIMEOUT), &req)
+            .await?;
+        Ok(res.Evidence)
+    }
+
     pub async fn extend_aael_entry(
         &self,
         domain: &str,

api-server-rest/src/client/aa.rs

@@ -4,15 +4,32 @@
 //
 
 use anyhow::*;
+use base64::Engine;
 use hyper::body::HttpBody;
 use hyper::{header, Body, Method, Request, Response, StatusCode};
 use serde::Serialize;
 use std::collections::HashMap;
 use std::net::SocketAddr;
 use tracing::{debug, info};
 
+fn decode_runtime_data(raw: &str, encoding: Option<&str>) -> Result<Vec<u8>> {
+    match encoding {
+        Some("hex") => {
+            hex::decode(raw).map_err(|e| anyhow!("invalid hex in runtime_data: {e}"))
+        }
+        Some("base64") => base64::engine::general_purpose::STANDARD
+            .decode(raw)
+            .map_err(|e| anyhow!("invalid base64 in runtime_data: {e}")),
+        Some(other) => bail!("unsupported encoding: {other} (expected hex, base64, or omit)"),
+        None => Ok(raw.as_bytes().to_vec()),
+    }
+}
+
 use crate::client::{
-    aa::{AAClient, AaelEvent, AA_AAEL_URL, AA_EVIDENCE_URL, AA_ROOT, AA_TOKEN_URL},
+    aa::{
+        AAClient, AaelEvent, AA_AAEL_URL, AA_ADDITIONAL_EVIDENCE_URL, AA_EVIDENCE_URL, AA_ROOT,
+        AA_TOKEN_URL,
+    },
     cdh::{CDHClient, CDH_RESOURCE_URL, CDH_ROOT},
 };
 use crate::utils::split_nth_slash;
@@ -168,8 +185,36 @@ impl Router {
                             info!("Get evidence");
                             match params.get("runtime_data") {
                                 Some(runtime_data) => {
+                                    let data = match decode_runtime_data(
+                                        runtime_data,
+                                        params.get("encoding").map(|s| s.as_str()),
+                                    ) {
+                                        std::result::Result::Ok(d) => d,
+                                        Err(e) => return self.internal_error(e.to_string()),
+                                    };
+                                    match client.get_evidence(&data).await {
+                                        std::result::Result::Ok(results) => {
+                                            return self.octet_stream_response(results)
+                                        }
+                                        Err(e) => return self.internal_error(e.to_string()),
+                                    }
+                                }
+                                None => return self.bad_request(),
+                            }
+                        }
+                        (AA_ADDITIONAL_EVIDENCE_URL, &Method::GET) => {
+                            info!("Get additional evidence");
+                            match params.get("runtime_data") {
+                                Some(runtime_data) => {
+                                    let data = match decode_runtime_data(
+                                        runtime_data,
+                                        params.get("encoding").map(|s| s.as_str()),
+                                    ) {
+                                        std::result::Result::Ok(d) => d,
+                                        Err(e) => return self.internal_error(e.to_string()),
+                                    };
                                     match client
-                                        .get_evidence(&runtime_data.clone().into_bytes())
+                                        .get_additional_evidence(&data)
                                         .await
                                     {
                                         std::result::Result::Ok(results) => {