AA: get_token API add runtime_data paramter

Signed-off-by: Jiale Zhang <[email protected]>
confidential-containers · Feb 2, 2024 · 25d8bae · 25d8bae
1 parent 8b0dbeb
commit 25d8bae
Show file tree

Hide file tree

Showing 15 changed files with 141 additions and 34 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/api-server-rest/README.md b/api-server-rest/README.md
@@ -13,6 +13,6 @@ $ curl http://127.0.0.1:8006/cdh/resource/default/key/1
 $ curl http://127.0.0.1:8006/aa/evidence\?runtime_data\=xxxx
 {"svn":"1","report_data":"eHh4eA=="}
 
-$ curl http://127.0.0.1:8006/aa/token\?token_type\=kbs
+$ curl http://127.0.0.1:8006/aa/token\?token_type\=kbs\?structured_runtime_data=xxx
 {"token":"eyJhbGciOiJFi...","tee_keypair":"-----BEGIN... "}
 ```
diff --git a/api-server-rest/build.rs b/api-server-rest/build.rs
@@ -12,7 +12,8 @@ use utoipa::OpenApi;
     get,
     path = "/aa/token",
     params(
-        ("token_type" = String, Query, description = "Token Type")
+        ("token_type" = String, Query, description = "Token Type"),
+        ("structured_runtime_data" = Option<String>, Query, description = "Structured data in JSON format, which will be hashed as runtime data")
     ),
     responses(
         (status = 200, description = "success response",

diff --git a/api-server-rest/openapi/api.json b/api-server-rest/openapi/api.json
@@ -80,6 +80,16 @@
             "schema": {
               "type": "string"
             }
+          },
+          {
+            "name": "structured_runtime_data",
+            "in": "query",
+            "description": "Structured data in JSON format, which will be hashed as runtime data",
+            "required": false,
+            "schema": {
+              "type": "string",
+              "nullable": true
+            }
           }
         ],
         "responses": {

diff --git a/api-server-rest/protos/attestation_agent.proto b/api-server-rest/protos/attestation_agent.proto
@@ -12,6 +12,7 @@ message GetEvidenceResponse {
 
 message GetTokenRequest {
     string TokenType = 1;
+    optional string StructuredRuntimeData = 2;
 }
 
 message GetTokenResponse {

diff --git a/api-server-rest/src/aa.rs b/api-server-rest/src/aa.rs
@@ -57,8 +57,12 @@ impl ApiHandler for AAClient {
         match url_path {
             AA_TOKEN_URL => match params.get("token_type") {
                 Some(token_type) => {
+                    let default_structed_runtime_data = String::from("{}");
+                    let structured_runtime_data = params
+                        .get("structured_runtime_data")
+                        .unwrap_or(&default_structed_runtime_data);
                     let results = self
-                        .get_token(token_type)
+                        .get_token(token_type, &structured_runtime_data)
                         .await
                         .unwrap_or_else(|e| e.to_string().into());
                     return self.octet_stream_response(results);
@@ -95,9 +99,14 @@ impl AAClient {
         })
     }
 
-    pub async fn get_token(&self, token_type: &str) -> Result<Vec<u8>> {
+    pub async fn get_token(
+        &self,
+        token_type: &str,
+        structured_runtime_data: &str,
+    ) -> Result<Vec<u8>> {
         let req = GetTokenRequest {
             TokenType: token_type.to_string(),
+            StructuredRuntimeData: Some(structured_runtime_data.to_string()),
             ..Default::default()
         };
         let res = self

diff --git a/api-server-rest/src/ttrpc_proto/attestation_agent.rs b/api-server-rest/src/ttrpc_proto/attestation_agent.rs
@@ -275,6 +275,8 @@ pub struct GetTokenRequest {
     // message fields
     // @@protoc_insertion_point(field:attestation_agent.GetTokenRequest.TokenType)
     pub TokenType: ::std::string::String,
+    // @@protoc_insertion_point(field:attestation_agent.GetTokenRequest.StructuredRuntimeData)
+    pub StructuredRuntimeData: ::std::option::Option<::std::string::String>,
     // special fields
     // @@protoc_insertion_point(special_field:attestation_agent.GetTokenRequest.special_fields)
     pub special_fields: ::protobuf::SpecialFields,
@@ -292,13 +294,18 @@ impl GetTokenRequest {
     }
 
     fn generated_message_descriptor_data() -> ::protobuf::reflect::GeneratedMessageDescriptorData {
-        let mut fields = ::std::vec::Vec::with_capacity(1);
+        let mut fields = ::std::vec::Vec::with_capacity(2);
         let mut oneofs = ::std::vec::Vec::with_capacity(0);
         fields.push(::protobuf::reflect::rt::v2::make_simpler_field_accessor::<_, _>(
             "TokenType",
             |m: &GetTokenRequest| { &m.TokenType },
             |m: &mut GetTokenRequest| { &mut m.TokenType },
         ));
+        fields.push(::protobuf::reflect::rt::v2::make_option_accessor::<_, _>(
+            "StructuredRuntimeData",
+            |m: &GetTokenRequest| { &m.StructuredRuntimeData },
+            |m: &mut GetTokenRequest| { &mut m.StructuredRuntimeData },
+        ));
         ::protobuf::reflect::GeneratedMessageDescriptorData::new_2::<GetTokenRequest>(
             "GetTokenRequest",
             fields,
@@ -320,6 +327,9 @@ impl ::protobuf::Message for GetTokenRequest {
                 10 => {
                     self.TokenType = is.read_string()?;
                 },
+                18 => {
+                    self.StructuredRuntimeData = ::std::option::Option::Some(is.read_string()?);
+                },
                 tag => {
                     ::protobuf::rt::read_unknown_or_skip_group(tag, is, self.special_fields.mut_unknown_fields())?;
                 },
@@ -335,6 +345,9 @@ impl ::protobuf::Message for GetTokenRequest {
         if !self.TokenType.is_empty() {
             my_size += ::protobuf::rt::string_size(1, &self.TokenType);
         }
+        if let Some(v) = self.StructuredRuntimeData.as_ref() {
+            my_size += ::protobuf::rt::string_size(2, &v);
+        }
         my_size += ::protobuf::rt::unknown_fields_size(self.special_fields.unknown_fields());
         self.special_fields.cached_size().set(my_size as u32);
         my_size
@@ -344,6 +357,9 @@ impl ::protobuf::Message for GetTokenRequest {
         if !self.TokenType.is_empty() {
             os.write_string(1, &self.TokenType)?;
         }
+        if let Some(v) = self.StructuredRuntimeData.as_ref() {
+            os.write_string(2, v)?;
+        }
         os.write_unknown_fields(self.special_fields.unknown_fields())?;
         ::std::result::Result::Ok(())
     }
@@ -362,12 +378,14 @@ impl ::protobuf::Message for GetTokenRequest {
 
     fn clear(&mut self) {
         self.TokenType.clear();
+        self.StructuredRuntimeData = ::std::option::Option::None;
         self.special_fields.clear();
     }
 
     fn default_instance() -> &'static GetTokenRequest {
         static instance: GetTokenRequest = GetTokenRequest {
             TokenType: ::std::string::String::new(),
+            StructuredRuntimeData: ::std::option::Option::None,
             special_fields: ::protobuf::SpecialFields::new(),
         };
         &instance
@@ -517,12 +535,14 @@ static file_descriptor_proto_data: &'static [u8] = b"\
     \n\x17attestation_agent.proto\x12\x11attestation_agent\"6\n\x12GetEviden\
     ceRequest\x12\x20\n\x0bRuntimeData\x18\x01\x20\x01(\x0cR\x0bRuntimeData\
     \"1\n\x13GetEvidenceResponse\x12\x1a\n\x08Evidence\x18\x01\x20\x01(\x0cR\
-    \x08Evidence\"/\n\x0fGetTokenRequest\x12\x1c\n\tTokenType\x18\x01\x20\
-    \x01(\tR\tTokenType\"(\n\x10GetTokenResponse\x12\x14\n\x05Token\x18\x01\
-    \x20\x01(\x0cR\x05Token2\xcc\x01\n\x17AttestationAgentService\x12\\\n\
-    \x0bGetEvidence\x12%.attestation_agent.GetEvidenceRequest\x1a&.attestati\
-    on_agent.GetEvidenceResponse\x12S\n\x08GetToken\x12\".attestation_agent.\
-    GetTokenRequest\x1a#.attestation_agent.GetTokenResponseb\x06proto3\
+    \x08Evidence\"\x84\x01\n\x0fGetTokenRequest\x12\x1c\n\tTokenType\x18\x01\
+    \x20\x01(\tR\tTokenType\x129\n\x15StructuredRuntimeData\x18\x02\x20\x01(\
+    \tH\0R\x15StructuredRuntimeData\x88\x01\x01B\x18\n\x16_StructuredRuntime\
+    Data\"(\n\x10GetTokenResponse\x12\x14\n\x05Token\x18\x01\x20\x01(\x0cR\
+    \x05Token2\xcc\x01\n\x17AttestationAgentService\x12\\\n\x0bGetEvidence\
+    \x12%.attestation_agent.GetEvidenceRequest\x1a&.attestation_agent.GetEvi\
+    denceResponse\x12S\n\x08GetToken\x12\".attestation_agent.GetTokenRequest\
+    \x1a#.attestation_agent.GetTokenResponseb\x06proto3\
 ";
 
 /// `FileDescriptorProto` object which was a source for this generated file

diff --git a/attestation-agent/app/src/rpc/attestation/mod.rs b/attestation-agent/app/src/rpc/attestation/mod.rs
@@ -45,7 +45,12 @@ pub mod grpc {
             debug!("Call AA to get token ...");
 
             let token = attestation_agent
-                .get_token(&request.token_type)
+                .get_token(
+                    &request.token_type,
+                    &request
+                        .structured_runtime_data
+                        .unwrap_or_else(|| String::from("{}")),
+                )
                 .await
                 .map_err(|e| {
                     error!("Call AA to get token failed: {}", e);
@@ -156,7 +161,11 @@ pub mod ttrpc {
             let mut attestation_agent = attestation_agent_mutex_clone.lock().await;
 
             let token = attestation_agent
-                .get_token(&req.TokenType)
+                .get_token(
+                    &req.TokenType,
+                    &req.StructuredRuntimeData
+                        .unwrap_or_else(|| String::from("{}")),
+                )
                 .await
                 .map_err(|e| {
                     error!("Call AA-KBC to get token failed: {}", e);

diff --git a/attestation-agent/kbs_protocol/src/token_provider/aa/attestation_agent.rs b/attestation-agent/kbs_protocol/src/token_provider/aa/attestation_agent.rs
@@ -275,6 +275,8 @@ pub struct GetTokenRequest {
     // message fields
     // @@protoc_insertion_point(field:attestation_agent.GetTokenRequest.TokenType)
     pub TokenType: ::std::string::String,
+    // @@protoc_insertion_point(field:attestation_agent.GetTokenRequest.StructuredRuntimeData)
+    pub StructuredRuntimeData: ::std::option::Option<::std::string::String>,
     // special fields
     // @@protoc_insertion_point(special_field:attestation_agent.GetTokenRequest.special_fields)
     pub special_fields: ::protobuf::SpecialFields,
@@ -292,13 +294,18 @@ impl GetTokenRequest {
     }
 
     fn generated_message_descriptor_data() -> ::protobuf::reflect::GeneratedMessageDescriptorData {
-        let mut fields = ::std::vec::Vec::with_capacity(1);
+        let mut fields = ::std::vec::Vec::with_capacity(2);
         let mut oneofs = ::std::vec::Vec::with_capacity(0);
         fields.push(::protobuf::reflect::rt::v2::make_simpler_field_accessor::<_, _>(
             "TokenType",
             |m: &GetTokenRequest| { &m.TokenType },
             |m: &mut GetTokenRequest| { &mut m.TokenType },
         ));
+        fields.push(::protobuf::reflect::rt::v2::make_option_accessor::<_, _>(
+            "StructuredRuntimeData",
+            |m: &GetTokenRequest| { &m.StructuredRuntimeData },
+            |m: &mut GetTokenRequest| { &mut m.StructuredRuntimeData },
+        ));
         ::protobuf::reflect::GeneratedMessageDescriptorData::new_2::<GetTokenRequest>(
             "GetTokenRequest",
             fields,
@@ -320,6 +327,9 @@ impl ::protobuf::Message for GetTokenRequest {
                 10 => {
                     self.TokenType = is.read_string()?;
                 },
+                18 => {
+                    self.StructuredRuntimeData = ::std::option::Option::Some(is.read_string()?);
+                },
                 tag => {
                     ::protobuf::rt::read_unknown_or_skip_group(tag, is, self.special_fields.mut_unknown_fields())?;
                 },
@@ -335,6 +345,9 @@ impl ::protobuf::Message for GetTokenRequest {
         if !self.TokenType.is_empty() {
             my_size += ::protobuf::rt::string_size(1, &self.TokenType);
         }
+        if let Some(v) = self.StructuredRuntimeData.as_ref() {
+            my_size += ::protobuf::rt::string_size(2, &v);
+        }
         my_size += ::protobuf::rt::unknown_fields_size(self.special_fields.unknown_fields());
         self.special_fields.cached_size().set(my_size as u32);
         my_size
@@ -344,6 +357,9 @@ impl ::protobuf::Message for GetTokenRequest {
         if !self.TokenType.is_empty() {
             os.write_string(1, &self.TokenType)?;
         }
+        if let Some(v) = self.StructuredRuntimeData.as_ref() {
+            os.write_string(2, v)?;
+        }
         os.write_unknown_fields(self.special_fields.unknown_fields())?;
         ::std::result::Result::Ok(())
     }
@@ -362,12 +378,14 @@ impl ::protobuf::Message for GetTokenRequest {
 
     fn clear(&mut self) {
         self.TokenType.clear();
+        self.StructuredRuntimeData = ::std::option::Option::None;
         self.special_fields.clear();
     }
 
     fn default_instance() -> &'static GetTokenRequest {
         static instance: GetTokenRequest = GetTokenRequest {
             TokenType: ::std::string::String::new(),
+            StructuredRuntimeData: ::std::option::Option::None,
             special_fields: ::protobuf::SpecialFields::new(),
         };
         &instance
@@ -760,18 +778,20 @@ static file_descriptor_proto_data: &'static [u8] = b"\
     \n\x17attestation-agent.proto\x12\x11attestation_agent\"6\n\x12GetEviden\
     ceRequest\x12\x20\n\x0bRuntimeData\x18\x01\x20\x01(\x0cR\x0bRuntimeData\
     \"1\n\x13GetEvidenceResponse\x12\x1a\n\x08Evidence\x18\x01\x20\x01(\x0cR\
-    \x08Evidence\"/\n\x0fGetTokenRequest\x12\x1c\n\tTokenType\x18\x01\x20\
-    \x01(\tR\tTokenType\"(\n\x10GetTokenResponse\x12\x14\n\x05Token\x18\x01\
-    \x20\x01(\x0cR\x05Token\"v\n\x1fExtendRuntimeMeasurementRequest\x12\x16\
-    \n\x06Events\x18\x01\x20\x03(\x0cR\x06Events\x12)\n\rRegisterIndex\x18\
-    \x02\x20\x01(\x04H\0R\rRegisterIndex\x88\x01\x01B\x10\n\x0e_RegisterInde\
-    x\"\"\n\x20ExtendRuntimeMeasurementResponse2\xd2\x02\n\x17AttestationAge\
-    ntService\x12\\\n\x0bGetEvidence\x12%.attestation_agent.GetEvidenceReque\
-    st\x1a&.attestation_agent.GetEvidenceResponse\x12S\n\x08GetToken\x12\".a\
-    ttestation_agent.GetTokenRequest\x1a#.attestation_agent.GetTokenResponse\
-    \x12\x83\x01\n\x18ExtendRuntimeMeasurement\x122.attestation_agent.Extend\
-    RuntimeMeasurementRequest\x1a3.attestation_agent.ExtendRuntimeMeasuremen\
-    tResponseb\x06proto3\
+    \x08Evidence\"\x84\x01\n\x0fGetTokenRequest\x12\x1c\n\tTokenType\x18\x01\
+    \x20\x01(\tR\tTokenType\x129\n\x15StructuredRuntimeData\x18\x02\x20\x01(\
+    \tH\0R\x15StructuredRuntimeData\x88\x01\x01B\x18\n\x16_StructuredRuntime\
+    Data\"(\n\x10GetTokenResponse\x12\x14\n\x05Token\x18\x01\x20\x01(\x0cR\
+    \x05Token\"v\n\x1fExtendRuntimeMeasurementRequest\x12\x16\n\x06Events\
+    \x18\x01\x20\x03(\x0cR\x06Events\x12)\n\rRegisterIndex\x18\x02\x20\x01(\
+    \x04H\0R\rRegisterIndex\x88\x01\x01B\x10\n\x0e_RegisterIndex\"\"\n\x20Ex\
+    tendRuntimeMeasurementResponse2\xd2\x02\n\x17AttestationAgentService\x12\
+    \\\n\x0bGetEvidence\x12%.attestation_agent.GetEvidenceRequest\x1a&.attes\
+    tation_agent.GetEvidenceResponse\x12S\n\x08GetToken\x12\".attestation_ag\
+    ent.GetTokenRequest\x1a#.attestation_agent.GetTokenResponse\x12\x83\x01\
+    \n\x18ExtendRuntimeMeasurement\x122.attestation_agent.ExtendRuntimeMeasu\
+    rementRequest\x1a3.attestation_agent.ExtendRuntimeMeasurementResponseb\
+    \x06proto3\
 ";
 
 /// `FileDescriptorProto` object which was a source for this generated file

diff --git a/attestation-agent/lib/Cargo.toml b/attestation-agent/lib/Cargo.toml
@@ -17,6 +17,7 @@ resource_uri.workspace = true
 reqwest = { workspace = true, features = ["json"], optional = true }
 serde.workspace = true
 serde_json.workspace = true
+sha2.workspace = true
 strum.workspace = true
 thiserror.workspace = true
 tokio = { workspace = true, features = ["fs"] }

diff --git a/attestation-agent/lib/src/lib.rs b/attestation-agent/lib/src/lib.rs
@@ -79,7 +79,11 @@ pub trait AttestationAPIs {
     ) -> Result<Vec<u8>>;
 
     /// Get attestation Token
-    async fn get_token(&mut self, token_type: &str) -> Result<Vec<u8>>;
+    async fn get_token(
+        &mut self,
+        token_type: &str,
+        structured_runtime_data: &str,
+    ) -> Result<Vec<u8>>;
 
     /// Get TEE hardware signed evidence that includes the runtime data.
     async fn get_evidence(&mut self, runtime_data: &[u8]) -> Result<Vec<u8>>;
@@ -188,7 +192,11 @@ impl AttestationAPIs for AttestationAgent {
     }
 
     #[allow(unreachable_code)]
-    async fn get_token(&mut self, _token_type: &str) -> Result<Vec<u8>> {
+    async fn get_token(
+        &mut self,
+        _token_type: &str,
+        _structured_runtime_data: &str,
+    ) -> Result<Vec<u8>> {
         let _uri = match self.config.as_ref() {
             Some(c) => c.as_uri.clone(),
             None => {
@@ -205,13 +213,13 @@ impl AttestationAPIs for AttestationAgent {
             #[cfg(feature = "kbs")]
             TokenType::Kbs => {
                 token::kbs::KbsTokenGetter::default()
-                    .get_token(_uri)
+                    .get_token(_uri, _structured_runtime_data)
                     .await?
             }
             #[cfg(feature = "coco_as")]
             TokenType::CoCoAS => {
                 token::coco_as::CoCoASTokenGetter::default()
-                    .get_token(_uri)
+                    .get_token(_uri, _structured_runtime_data)
                     .await?
             }
         };

diff --git a/attestation-agent/lib/src/token/coco_as.rs b/attestation-agent/lib/src/token/coco_as.rs
@@ -8,19 +8,35 @@ use anyhow::*;
 use async_trait::async_trait;
 use base64::engine::general_purpose::URL_SAFE_NO_PAD;
 use base64::Engine;
+use sha2::{Digest, Sha384};
 
 #[derive(Default)]
 pub struct CoCoASTokenGetter {}
 
 #[async_trait]
 impl GetToken for CoCoASTokenGetter {
-    async fn get_token(&self, as_uri: String) -> Result<Vec<u8>> {
+    async fn get_token(&self, as_uri: String, structured_runtime_data: &str) -> Result<Vec<u8>> {
+        let structured_value: serde_json::Value = serde_json::from_str(structured_runtime_data)
+            .context("Get Token Failed: Structured Runtime Data must be a JSON Map")?;
+
+        // TODO: Request AS to get Nonce and insert the Nonce into structured runtime data JSON Map.
+
+        let hash_materials =
+            serde_json::to_vec(&structured_value).context("parse JSON structured data")?;
+        let mut hasher = Sha384::new();
+        hasher.update(hash_materials);
+        let structured_data_digest = hasher.finalize().to_vec();
+
         let tee_type = attester::detect_tee_type();
         let attester = attester::BoxedAttester::try_from(tee_type)?;
-        let evidence = attester.get_evidence(vec![]).await?;
+        let evidence = attester.get_evidence(structured_data_digest.clone()).await?;
 
         let request_body = serde_json::json!({
             "tee": serde_json::to_string(&tee_type)?,
+            "runtime_data": {
+                "structured": structured_value
+            },
+            "runtime_data_hash_algorithm": "sha384",
             "evidence": URL_SAFE_NO_PAD.encode(evidence.as_bytes()),
         });