devired_key: cleaning pass confidential-containers#1

Author: eldios

api-server-rest/src/aa.rs

@@ -67,6 +67,18 @@ impl ApiHandler for AAClient {
                 _ => {
                     return self.not_found();
                 }
+                None => return self.bad_request(),
+            },
+            AA_DERIVED_KEY_URL => match params.get() {
+                Some(key) => match self.get_derived_key().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(),
+            },
+
+            _ => {
+                return self.not_found();
             }
         }
 

attestation-agent/attestation-agent/src/bin/grpc-aa/server.rs

@@ -164,14 +164,10 @@ impl AttestationAgentService for AA {
 
         debug!("AA (grpc): get derived key ...");
 
-        let derived_key = self
-            .inner
-            .get_derived_key(&request.key_id)
-            .await
-            .map_err(|e| {
-                error!("AA (grpc): get derived key failed:\n{e:?}\nkey_id:\n{&request.key_id}");
-                Status::internal(format!("[ERROR:{AGENT_NAME}] AA get derived key failed"))
-            })?;
+        let derived_key = self.inner.get_derived_key().await.map_err(|e| {
+            error!("AA (grpc): get derived key failed:\n{e:?}");
+            Status::internal(format!("[ERROR:{AGENT_NAME}] AA get derived key failed"))
+        })?;
 
         debug!("AA (grpc): Get derived key successfully!");
 

attestation-agent/attestation-agent/src/bin/ttrpc-aa-client.rs

@@ -80,11 +80,7 @@ struct GetTokenArgs {
 
 #[derive(Args)]
 #[command(author, version, about, long_about = None)]
-struct GetDerivedKeyArgs {
-    /// base64 encodede runtime data
-    #[arg(short, long)]
-    key_id: String,
-}
+struct GetDerivedKeyArgs {}
 
 #[derive(Args)]
 #[command(author, version, about, long_about = None)]
@@ -152,15 +148,14 @@ pub async fn main() {
         }
         Operation::GetDerivedKey(get_derived_key_args) => {
             let req = GetDerivedKeyRequest {
-                KeyId: get_derived_key_args.key_id,
                 ..Default::default()
             };
             let res = client
                 .get_derived_key(context::with_timeout(TIMEOUT), &req)
                 .await
                 .expect("request to AA");
-            let key_id = String::from_utf8(res.KeyId).unwrap();
-            println!("{key_id}");
+            let key = String::from_utf8(res.Key).unwrap();
+            println!("{key}");
         }
         Operation::ExtendRuntimeMeasurement(extend_runtime_measurement_args) => {
             let req = ExtendRuntimeMeasurementRequest {

attestation-agent/attestation-agent/src/bin/ttrpc_dep/server.rs

@@ -84,21 +84,16 @@ impl AttestationAgentService for AA {
     ) -> ::ttrpc::Result<GetDerivedKeyResponse> {
         debug!("AA (ttrpc): get derived key ...");
 
+        let empty_context = Vec::new();
         let derived_key = self
             .inner
-            .get_derived_key(&req.KeyId, Vec::new())
+            .get_derived_key(empty_context)
             .await
             .map_err(|e| {
-                error!(
-                    "AA (ttrpc): get derived key failed:\n {e:?}\n key_id:\n {:#?}",
-                    &req.KeyId
-                );
+                error!("AA (ttrpc): get derived key failed:\n {e:?}");
                 let mut error_status = ::ttrpc::proto::Status::new();
                 error_status.set_code(Code::INTERNAL);
-                error_status.set_message(format!(
-                    "[ERROR:{AGENT_NAME}] AA-KBC get derived key failed. key_id: {:#?}",
-                    &req.KeyId
-                ));
+                error_status.set_message("[ERROR:{AGENT_NAME}] AA-KBC get derived key failed.");
                 ::ttrpc::Error::RpcStatus(error_status)
             })?;
 

attestation-agent/attestation-agent/src/lib.rs

@@ -60,8 +60,8 @@ pub trait AttestationAPIs {
     /// Get TEE hardware signed evidence that includes the runtime data.
     async fn get_evidence(&self, runtime_data: &[u8]) -> Result<Vec<u8>>;
 
-    /// Get a derived key using the provided key ID
-    async fn get_derived_key(&self, key_id: &[u8], context: Vec<u8>) -> Result<Vec<u8>>;
+    /// Get a derived key
+    async fn get_derived_key(&self, context: Vec<u8>) -> Result<Vec<u8>>;
 
     /// Extend runtime measurement register
     async fn extend_runtime_measurement(
@@ -180,8 +180,8 @@ impl AttestationAPIs for AttestationAgent {
         Ok(evidence.into_bytes())
     }
 
-    async fn get_derived_key(&self, key_id: &[u8], context: Vec<u8>) -> Result<Vec<u8>> {
-        self.attester.get_derived_key(key_id, context).await
+    async fn get_derived_key(&self, context: Vec<u8>) -> Result<Vec<u8>> {
+        self.attester.get_derived_key(context).await
     }
 
     /// Extend runtime measurement register. Parameters

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

@@ -71,10 +71,6 @@ impl Attester for SnpAttester {
         }
 
         context.resize(64, 0);
-        let _root_key: u8 = root_key_hinit
-            .first()
-            .copied()
-            .context("Invalid key or empty key specified")?;
 
         let mut firmware: Firmware = Firmware::open()?;
 

attestation-agent/kbs_protocol/src/evidence_provider/aa_ttrpc.rs

@@ -40,9 +40,8 @@ impl AAEvidenceProvider {
 #[async_trait]
 impl EvidenceProvider for AAEvidenceProvider {
     /// Get derived key using the provided key ID
-    async fn get_derived_key(&self, key_id: &[u8], _context: Vec<u8>) -> Result<Vec<u8>> {
+    async fn get_derived_key(&self, _context: Vec<u8>) -> Result<Vec<u8>> {
         let req = GetDerivedKeyRequest {
-            KeyId: key_id.to_vec(),
             ..Default::default()
         };
         let res = self

attestation-agent/kbs_protocol/src/evidence_provider/mock.rs

@@ -19,7 +19,7 @@ impl EvidenceProvider for MockedEvidenceProvider {
         Ok("test evidence".into())
     }
 
-    async fn get_derived_key(&self, _key_id: &[u8], _context: Vec<u8>) -> Result<Vec<u8>> {
+    async fn get_derived_key(&self, _context: Vec<u8>) -> Result<Vec<u8>> {
         Ok(vec![0u8; 32]) // Return a mock 32-byte key filled with zeros
     }
 

attestation-agent/kbs_protocol/src/evidence_provider/mod.rs

@@ -27,7 +27,6 @@ pub trait EvidenceProvider: Send + Sync {
     async fn get_tee_type(&self) -> Result<Tee>;
 
     /// Get a derived key using the hardware-specific key derivation function.
-    /// The parameter `root_key_hint` is the root key used for derivation,
-    /// and `context` is additional data used in the derivation process.
-    async fn get_derived_key(&self, key_id: &[u8], context: Vec<u8>) -> Result<Vec<u8>>;
+    /// The parameter `context` is data potentially used in the derivation process.
+    async fn get_derived_key(&self, context: Vec<u8>) -> Result<Vec<u8>>;
 }

attestation-agent/kbs_protocol/src/evidence_provider/native.rs

@@ -35,9 +35,9 @@ impl EvidenceProvider for NativeEvidenceProvider {
         Ok(detect_tee_type())
     }
 
-    async fn get_derived_key(&self, key_id: &[u8], context: Vec<u8>) -> Result<Vec<u8>> {
+    async fn get_derived_key(&self, context: Vec<u8>) -> Result<Vec<u8>> {
         self.0
-            .get_derived_key(key_id, context)
+            .get_derived_key(context)
             .await
             .map_err(|e| Error::GetDerivedKey(e.to_string()))
     }