AA: Add API to extend measurement register at runtime

Signed-off-by: Jiale Zhang <[email protected]>

attestation-agent/app/build.rs

@@ -9,9 +9,17 @@ use ttrpc_codegen::{Codegen, Customize, ProtobufCustomize};
 fn main() -> std::io::Result<()> {
     #[cfg(feature = "grpc")]
     {
-        tonic_build::compile_protos("../protos/keyprovider.proto")?;
-        tonic_build::compile_protos("../protos/getresource.proto")?;
-        tonic_build::compile_protos("../protos/attestation-agent.proto")?;
+        tonic_build::configure()
+            .build_server(true)
+            .protoc_arg("--experimental_allow_proto3_optional")
+            .compile(
+                &[
+                    "../protos/keyprovider.proto",
+                    "../protos/getresource.proto",
+                    "../protos/attestation-agent.proto",
+                ],
+                &["../protos"],
+            )?;
     }
 
     #[cfg(feature = "ttrpc")]

attestation-agent/app/src/rpc/attestation/mod.rs

@@ -20,7 +20,10 @@ pub mod grpc {
     use attestation::attestation_agent_service_server::{
         AttestationAgentService, AttestationAgentServiceServer,
     };
-    use attestation::{GetEvidenceRequest, GetEvidenceResponse, GetTokenRequest, GetTokenResponse};
+    use attestation::{
+        ExtendRuntimeMeasurementRequest, ExtendRuntimeMeasurementResponse, GetEvidenceRequest,
+        GetEvidenceResponse, GetTokenRequest, GetTokenResponse,
+    };
     use std::net::SocketAddr;
     use tonic::{transport::Server, Request, Response, Status};
 
@@ -84,6 +87,35 @@ pub mod grpc {
 
             Result::Ok(Response::new(reply))
         }
+
+        async fn extend_runtime_measurement(
+            &self,
+            request: Request<ExtendRuntimeMeasurementRequest>,
+        ) -> Result<Response<ExtendRuntimeMeasurementResponse>, Status> {
+            let request = request.into_inner();
+
+            let attestation_agent_mutex_clone = Arc::clone(&ASYNC_ATTESTATION_AGENT);
+            let mut attestation_agent = attestation_agent_mutex_clone.lock().await;
+
+            debug!("Call AA to extend runtime measurement ...");
+
+            attestation_agent
+                .extend_runtime_measurement(request.events, request.register_index)
+                .await
+                .map_err(|e| {
+                    error!("Call AA to extend runtime measurement failed: {}", e);
+                    Status::internal(format!(
+                        "[ERROR:{}] AA extend runtime measurement failed: {}",
+                        AGENT_NAME, e
+                    ))
+                })?;
+
+            debug!("Extend runtime measurement successfully!");
+
+            let reply = ExtendRuntimeMeasurementResponse {};
+
+            Result::Ok(Response::new(reply))
+        }
     }
 
     pub async fn start_grpc_service(socket: SocketAddr) -> Result<()> {
@@ -176,6 +208,34 @@ pub mod ttrpc {
 
             ::ttrpc::Result::Ok(reply)
         }
+
+        async fn extend_runtime_measurement(
+            &self,
+            _ctx: &::ttrpc::r#async::TtrpcContext,
+            req: attestation_agent::ExtendRuntimeMeasurementRequest,
+        ) -> ::ttrpc::Result<attestation_agent::ExtendRuntimeMeasurementResponse> {
+            debug!("Call AA to extend runtime measurement ...");
+
+            let attestation_agent_mutex_clone = ASYNC_ATTESTATION_AGENT.clone();
+            let mut attestation_agent = attestation_agent_mutex_clone.lock().await;
+
+            attestation_agent
+                .extend_runtime_measurement(req.Events, req.RegisterIndex)
+                .await
+                .map_err(|e| {
+                    error!("Call AA-KBC to extend runtime measurement failed: {}", e);
+                    let mut error_status = ::ttrpc::proto::Status::new();
+                    error_status.set_code(Code::INTERNAL);
+                    error_status.set_message(format!(
+                        "[ERROR:{}] AA-KBC extend runtime measurement failed: {}",
+                        AGENT_NAME, e
+                    ));
+                    ::ttrpc::Error::RpcStatus(error_status)
+                })?;
+
+            let reply = attestation_agent::ExtendRuntimeMeasurementResponse::new();
+            ::ttrpc::Result::Ok(reply)
+        }
     }
 
     pub fn start_ttrpc_service() -> Result<HashMap<String, Service>> {

attestation-agent/attester/src/lib.rs

@@ -59,6 +59,16 @@ pub trait Attester {
     /// The parameter `report_data` will be used as the user input of the
     /// evidence to avoid reply attack.
     async fn get_evidence(&self, report_data: Vec<u8>) -> Result<String>;
+
+    /// Extend TEE specific dynamic measurement register
+    /// to enable dynamic measurement capabilities for input data at runtime.
+    async fn extend_runtime_measurement(
+        &self,
+        _events: Vec<Vec<u8>>,
+        _register_index: Option<u64>,
+    ) -> Result<()> {
+        bail!("Unimplement")
+    }
 }
 
 // Detect which TEE platform the KBC running environment is.

attestation-agent/attester/src/tdx/mod.rs

@@ -65,6 +65,26 @@ impl Attester for TdxAttester {
         serde_json::to_string(&evidence)
             .map_err(|e| anyhow!("Serialize TDX evidence failed: {:?}", e))
     }
+
+    async fn extend_runtime_measurement(
+        &self,
+        events: Vec<Vec<u8>>,
+        _register_index: Option<u64>,
+    ) -> Result<()> {
+        for event in events {
+            match tdx_attest_rs::tdx_att_extend(&event) {
+                tdx_attest_rs::tdx_attest_error_t::TDX_ATTEST_SUCCESS => todo!(),
+                error_code => {
+                    return Err(anyhow!(
+                        "TDX Attester: Failed to extend RTMR. Error code: {:?}",
+                        error_code
+                    ));
+                }
+            }
+        }
+
+        Ok(())
+    }
 }
 
 #[cfg(test)]