underhill_attestation: refactoring

openhcl/underhill_attestation/src/hardware_key_sealing.rs

@@ -24,18 +24,18 @@ pub(crate) enum HardwareDerivedKeysError {
 pub(crate) enum HardwareKeySealingError {
     #[error("failed to encrypt the egress key")]
     EncryptEgressKey(#[source] crypto::aes_256_cbc::Aes256CbcError),
-    #[error("invalid egress key encryption size {0}, expected {1}")]
-    InvalidEgressKeyEncryptionSize(usize, usize),
+    #[error("invalid egress key encryption size {size}, expected {expected_size}")]
+    InvalidEgressKeyEncryptionSize { size: usize, expected_size: usize },
     #[error("HMAC-SHA-256 after encryption failed")]
     HmacAfterEncrypt(#[source] crypto::hmac_sha_256::HmacSha256Error),
-    #[error("HMAC-SHA-256 before ecryption failed")]
+    #[error("HMAC-SHA-256 before decryption failed")]
     HmacBeforeDecrypt(#[source] crypto::hmac_sha_256::HmacSha256Error),
-    #[error("Hardware key protector HMAC verification failed")]
+    #[error("hardware key protector HMAC verification failed")]
     HardwareKeyProtectorHmacVerificationFailed,
     #[error("failed to decrypt the ingress key")]
     DecryptIngressKey(#[source] crypto::aes_256_cbc::Aes256CbcError),
-    #[error("invalid ingress key decryption size {0}, expected {1}")]
-    InvalidIngressKeyDecryptionSize(usize, usize),
+    #[error("invalid ingress key decryption size {size}, expected {expected_size}")]
+    InvalidIngressKeyDecryptionSize { size: usize, expected_size: usize },
 }
 
 /// Hold the hardware-derived keys.
@@ -89,7 +89,7 @@ pub trait HardwareKeyProtectorExt: Sized {
         egress_key: &[u8],
     ) -> Result<Self, HardwareKeySealingError>;
 
-    /// Unseal the `inress_key` with verify-mac-then-decrypt.
+    /// Unseal the `ingress_key` with verify-mac-then-decrypt.
     fn unseal_key(
         &self,
         hardware_derived_keys: &HardwareDerivedKeys,
@@ -115,10 +115,10 @@ impl HardwareKeyProtectorExt for HardwareKeyProtector {
             .and_then(|aes| aes.encrypt()?.cipher(&iv, egress_key))
             .map_err(HardwareKeySealingError::EncryptEgressKey)?;
         if output.len() != vmgs::AES_GCM_KEY_LENGTH {
-            Err(HardwareKeySealingError::InvalidEgressKeyEncryptionSize(
-                output.len(),
-                vmgs::AES_GCM_KEY_LENGTH,
-            ))?
+            return Err(HardwareKeySealingError::InvalidEgressKeyEncryptionSize {
+                size: output.len(),
+                expected_size: vmgs::AES_GCM_KEY_LENGTH,
+            });
         }
         encrypted_egress_key.copy_from_slice(&output[..vmgs::AES_GCM_KEY_LENGTH]);
 
@@ -152,18 +152,18 @@ impl HardwareKeyProtectorExt for HardwareKeyProtector {
         .map_err(HardwareKeySealingError::HmacBeforeDecrypt)?;
 
         if !constant_time_eq::constant_time_eq_32(&hmac, &self.hmac) {
-            Err(HardwareKeySealingError::HardwareKeyProtectorHmacVerificationFailed)?
+            return Err(HardwareKeySealingError::HardwareKeyProtectorHmacVerificationFailed);
         }
 
         let mut decrypted_ingress_key = [0u8; vmgs::AES_GCM_KEY_LENGTH];
         let output = crypto::aes_256_cbc::Aes256Cbc::new(&hardware_derived_keys.aes_key)
             .and_then(|aes| aes.decrypt()?.cipher(&self.iv, &self.ciphertext))
             .map_err(HardwareKeySealingError::DecryptIngressKey)?;
         if output.len() != vmgs::AES_GCM_KEY_LENGTH {
-            Err(HardwareKeySealingError::InvalidIngressKeyDecryptionSize(
-                output.len(),
-                vmgs::AES_GCM_KEY_LENGTH,
-            ))?
+            return Err(HardwareKeySealingError::InvalidIngressKeyDecryptionSize {
+                size: output.len(),
+                expected_size: vmgs::AES_GCM_KEY_LENGTH,
+            });
         }
         decrypted_ingress_key.copy_from_slice(&output[..vmgs::AES_GCM_KEY_LENGTH]);
 
@@ -179,7 +179,7 @@ impl HardwareKeyProtectorExt for HardwareKeyProtector {
 #[cfg(test)]
 mod tests {
     use super::*;
-    use crate::test_utils::MockTeeCall;
+    use crate::test_helpers::MockTeeCall;
     use zerocopy::FromBytes;
 
     #[test]

openhcl/underhill_attestation/src/igvm_attest/ak_cert.rs

@@ -8,23 +8,49 @@ use crate::igvm_attest::parse_response_header;
 
 use thiserror::Error;
 
-/// AkCertError is returned by parse_ak_cert_response() in emuplat/tpm.rs
+/// Errors returned when parsing an `AK_CERT_REQUEST` response.
+///
+/// Returned by [`parse_response`], which is re-exported from the crate root
+/// as `parse_ak_cert_response` and used by `emuplat/tpm.rs`.
 #[derive(Debug, Error)]
 pub enum AkCertError {
+    /// The response buffer is shorter than the minimum expected size to fit
+    /// the response header.
     #[error(
         "AK cert response is too small to parse. Found {size} bytes but expected at least {minimum_size}"
     )]
-    SizeTooSmall { size: usize, minimum_size: usize },
+    SizeTooSmall {
+        /// Actual length of the response buffer in bytes.
+        size: usize,
+        /// Minimum size in bytes required to parse the response header.
+        minimum_size: usize,
+    },
+    /// The size declared in the response header is larger than the actual
+    /// length of the response buffer received from the host.
     #[error(
         "AK cert response size {specified_size} specified in the header is larger then the actual size {size}"
     )]
-    SizeMismatch { size: usize, specified_size: usize },
+    SizeMismatch {
+        /// Actual length of the response buffer in bytes.
+        size: usize,
+        /// Length declared inside the response header.
+        specified_size: usize,
+    },
+    /// The response header version does not match the version expected by
+    /// this build.
     #[error(
         "AK cert response header version {version} does match the expected version {expected_version}"
     )]
-    HeaderVersionMismatch { version: u32, expected_version: u32 },
+    HeaderVersionMismatch {
+        /// Header version reported in the response.
+        version: u32,
+        /// Header version expected by this build.
+        expected_version: u32,
+    },
+    /// Parsing the common response header failed.
     #[error("error in parsing response header")]
     ParseHeader(#[source] CommonError),
+    /// The response header version is not a value recognized by this build.
     #[error("invalid response header version: {0}")]
     InvalidResponseVersion(u32),
 }

openhcl/underhill_attestation/src/igvm_attest/key_release.rs

@@ -60,7 +60,7 @@ pub fn parse_response(
     let minimum_payload_size = CIPHER_TEXT_KEY.len() + wrapped_key_base64_url_size - 1;
 
     if payload.len() < minimum_payload_size {
-        Err(KeyReleaseError::PayloadSizeTooSmall)?
+        return Err(KeyReleaseError::PayloadSizeTooSmall);
     }
     let data_utf8 = String::from_utf8_lossy(payload);
     let wrapped_key = match serde_json::from_str::<akv::AkvKeyReleaseKeyBlob>(&data_utf8) {
@@ -79,7 +79,7 @@ pub fn parse_response(
                     .verify_signature()
                     .map_err(KeyReleaseError::VerifyAkvJwtSignature)?
                 {
-                    Err(KeyReleaseError::VerifyAkvJwtSignatureFailed)?
+                    return Err(KeyReleaseError::VerifyAkvJwtSignatureFailed);
                 }
             }
             get_wrapped_key_blob(result)?

openhcl/underhill_attestation/src/igvm_attest/mod.rs

@@ -28,38 +28,69 @@ pub mod wrapped_key;
 
 base64_serde_type!(Base64Url, base64::engine::general_purpose::URL_SAFE_NO_PAD);
 
-#[expect(missing_docs)] // self-explanatory fields
+/// Errors returned by IGVM attest request preparation and response parsing.
 #[derive(Debug, Error)]
 pub enum Error {
+    /// The attestation report supplied to the request helper does not match
+    /// the expected size for the active TEE.
     #[error(
         "the size of the attestation report {report_size} is invalid, expected {expected_size}"
     )]
     InvalidAttestationReportSize {
+        /// Actual size of the report blob in bytes.
         report_size: usize,
+        /// Expected size of the report blob in bytes.
         expected_size: usize,
     },
+    /// The attestation response is shorter than the minimum header size.
     #[error("the size of the attestation response {response_size} is too small to parse")]
-    ResponseSizeTooSmall { response_size: usize },
+    ResponseSizeTooSmall {
+        /// Length of the response that failed to parse.
+        response_size: usize,
+    },
+    /// The attestation response header could not be deserialized.
     #[error(
         "the header of the attestation response (size {response_size}) is not in correct format"
     )]
-    ResponseHeaderInvalidFormat { response_size: usize },
+    ResponseHeaderInvalidFormat {
+        /// Length of the response whose header is malformed.
+        response_size: usize,
+    },
+    /// The size declared in the response header does not match the actual
+    /// length of the buffer received from the host.
     #[error(
         "response size {specified_size} specified in the header not match the actual size {size}"
     )]
-    ResponseSizeMismatch { size: usize, specified_size: usize },
+    ResponseSizeMismatch {
+        /// Actual length of the response buffer.
+        size: usize,
+        /// Length declared inside the header.
+        specified_size: usize,
+    },
+    /// The response header advertises a version newer than this build supports.
     #[error("response header version {version:?} larger than current version {latest_version:?}")]
     InvalidResponseHeaderVersion {
+        /// Version reported in the response header.
         version: IgvmAttestResponseVersion,
+        /// Latest version known to this build.
         latest_version: IgvmAttestResponseVersion,
     },
+    /// The host-side IGVM agent reported an attestation failure. The
+    /// `retry_signal` and `skip_hw_unsealing_signal` fields convey hints
+    /// from the agent about how the caller should proceed.
     #[error(
         "attest failed ({igvm_error_code}-{http_status_code}), retry recommendation ({retry_signal}), skip hw unsealing recommendation ({skip_hw_unsealing_signal})"
     )]
     Attestation {
+        /// IGVM-specific error code returned by the agent.
         igvm_error_code: u32,
+        /// HTTP status code from the underlying call to the attestation
+        /// service, when applicable.
         http_status_code: u32,
+        /// Hint from the agent that the operation may succeed on retry.
         retry_signal: bool,
+        /// Hint from the agent that hardware key unsealing should be skipped
+        /// because the protected secret is no longer recoverable.
         skip_hw_unsealing_signal: bool,
     },
 }
@@ -89,9 +120,16 @@ impl ReportType {
 }
 
 /// Helper struct to create `IgvmAttestRequest` in raw bytes.
+///
+/// The helper captures the immutable runtime-claims context (report type,
+/// serialized runtime claims, claims hash, and hash type). The specific
+/// `IgvmAttestRequestType` for each call is passed to [`create_request`] so
+/// that the same helper can be reused across multiple related requests
+/// (e.g. the wrapped-key flow issues both a `WRAPPED_KEY_REQUEST` and a
+/// `KEY_RELEASE_REQUEST` from the same helper).
+///
+/// [`create_request`]: IgvmAttestRequestHelper::create_request
 pub struct IgvmAttestRequestHelper {
-    /// The request type.
-    request_type: IgvmAttestRequestType,
     /// The report type.
     report_type: ReportType,
     /// Raw bytes of `RuntimeClaims`.
@@ -130,7 +168,6 @@ impl IgvmAttestRequestHelper {
         runtime_claims_hash[0..hash.len()].copy_from_slice(&hash);
 
         Self {
-            request_type: IgvmAttestRequestType::KEY_RELEASE_REQUEST,
             report_type,
             runtime_claims,
             runtime_claims_hash,
@@ -173,7 +210,6 @@ impl IgvmAttestRequestHelper {
         runtime_claims_hash[0..hash.len()].copy_from_slice(&hash);
 
         Self {
-            request_type: IgvmAttestRequestType::AK_CERT_REQUEST,
             report_type,
             runtime_claims,
             runtime_claims_hash,
@@ -186,20 +222,20 @@ impl IgvmAttestRequestHelper {
         &self.runtime_claims_hash
     }
 
-    /// Set the `request_type`.
-    pub fn set_request_type(&mut self, request_type: IgvmAttestRequestType) {
-        self.request_type = request_type
-    }
-
-    /// Create the request in raw bytes.
+    /// Create the request in raw bytes for the given `request_type`.
+    ///
+    /// The runtime claims captured by this helper are shared across all
+    /// requests built from it; only `request_type` and `attestation_report`
+    /// vary per call.
     pub fn create_request(
         &self,
         version: IgvmAttestRequestVersion,
+        request_type: IgvmAttestRequestType,
         attestation_report: &[u8],
     ) -> Result<Vec<u8>, Error> {
         create_request(
             version,
-            self.request_type,
+            request_type,
             &self.runtime_claims,
             attestation_report,
             &self.report_type,
@@ -220,16 +256,16 @@ pub fn parse_response_header(response: &[u8]) -> Result<IgvmAttestCommonResponse
 
     // Check header data_size and version
     if header.data_size as usize > response.len() {
-        Err(Error::ResponseSizeMismatch {
+        return Err(Error::ResponseSizeMismatch {
             size: response.len(),
             specified_size: header.data_size as usize,
-        })?
+        });
     }
     if header.version > IGVM_ATTEST_RESPONSE_CURRENT_VERSION {
-        Err(Error::InvalidResponseHeaderVersion {
+        return Err(Error::InvalidResponseHeaderVersion {
             version: header.version,
             latest_version: IGVM_ATTEST_RESPONSE_CURRENT_VERSION,
-        })?
+        });
     }
 
     // IgvmErrorInfo is added in response header since version 2
@@ -244,12 +280,12 @@ pub fn parse_response_header(response: &[u8]) -> Result<IgvmAttestCommonResponse
         .0; // TODO: zerocopy: err (https://github.com/microsoft/openvmm/issues/759)
 
         if 0 != igvm_error_info.error_code {
-            Err(Error::Attestation {
+            return Err(Error::Attestation {
                 igvm_error_code: igvm_error_info.error_code,
                 http_status_code: igvm_error_info.http_status_code,
                 retry_signal: igvm_error_info.igvm_signal.retry(),
                 skip_hw_unsealing_signal: igvm_error_info.igvm_signal.skip_hw_unsealing(),
-            })?
+            });
         }
     }
     Ok(IgvmAttestCommonResponseHeader {
@@ -276,10 +312,10 @@ fn create_request(
 
     let expected_report_size = get_report_size(report_type);
     if attestation_report.len() != expected_report_size {
-        Err(Error::InvalidAttestationReportSize {
+        return Err(Error::InvalidAttestationReportSize {
             report_size: attestation_report.len(),
             expected_size: expected_report_size,
-        })?
+        });
     }
 
     let runtime_claims_len = runtime_claims.len();
@@ -347,8 +383,7 @@ fn attestation_vm_config_with_time(
 fn runtime_claims_to_bytes(
     runtime_claims: &openhcl_attestation_protocol::igvm_attest::get::runtime_claims::RuntimeClaims,
 ) -> Vec<u8> {
-    let runtime_claims = serde_json::to_string(runtime_claims).expect("JSON serialization failed");
-    runtime_claims.as_bytes().to_vec()
+    serde_json::to_vec(runtime_claims).expect("JSON serialization failed")
 }
 
 #[cfg(test)]

openhcl/underhill_attestation/src/igvm_attest/wrapped_key.rs

@@ -23,6 +23,11 @@ pub(crate) enum WrappedKeyError {
     ParseHeader(#[source] CommonError),
     #[error("invalid response header version: {0}")]
     InvalidResponseVersion(u32),
+    #[error("response data_size {data_size} is smaller than header size {header_size}")]
+    InvalidDataSize {
+        data_size: usize,
+        header_size: usize,
+    },
 }
 
 /// Return value of the [`parse_response`].
@@ -56,10 +61,17 @@ pub fn parse_response(response: &[u8]) -> Result<IgvmWrappedKeyParsedResponse, W
         IgvmAttestResponseVersion::VERSION_2 => size_of::<IgvmAttestWrappedKeyResponseHeader>(),
         invalid_version => return Err(WrappedKeyError::InvalidResponseVersion(invalid_version.0)),
     };
-    let payload = &response[header_size..header.data_size as usize];
+    let data_size = header.data_size as usize;
+    if data_size < header_size {
+        return Err(WrappedKeyError::InvalidDataSize {
+            data_size,
+            header_size,
+        });
+    }
+    let payload = &response[header_size..data_size];
 
     if payload.len() < MINIMUM_PAYLOAD_SIZE {
-        Err(WrappedKeyError::PayloadSizeTooSmall)?
+        return Err(WrappedKeyError::PayloadSizeTooSmall);
     }
     let payload = String::from_utf8_lossy(payload);
     let payload: cps::VmmdBlob = serde_json::from_str(&payload).map_err(|json_err| {