Use Variant proto message in Endorsements

Fixes: 384532107
Change-Id: I1c9051336a8040879a938eea11b90b0539bae83a

oak_attestation_verification/BUILD

@@ -98,7 +98,6 @@ rust_test_suite(
         "@oak_crates_index//:hex",
         "@oak_crates_index//:lazy_static",
         "@oak_crates_index//:prost",
-        "@oak_crates_index//:prost-types",
         "@oak_crates_index//:x509-cert",
         "@oak_crates_index//:zerocopy",
     ],
@@ -181,7 +180,6 @@ rust_test_suite(
         "@oak_crates_index//:hex",
         "@oak_crates_index//:lazy_static",
         "@oak_crates_index//:prost",
-        "@oak_crates_index//:prost-types",
         "@oak_crates_index//:x509-cert",
         "@oak_crates_index//:zerocopy",
     ],

oak_attestation_verification/src/policy/application.rs

@@ -15,12 +15,14 @@
 //
 
 use anyhow::Context;
-use oak_attestation_verification_types::policy::Policy;
-use oak_proto_rust::oak::attestation::v1::{
-    ApplicationLayerData, ApplicationLayerEndorsements, ApplicationLayerReferenceValues,
-    EventAttestationResults,
+use oak_attestation_verification_types::{policy::Policy, APPLICATION_ENDORSEMENT_ID};
+use oak_proto_rust::oak::{
+    attestation::v1::{
+        ApplicationLayerData, ApplicationLayerEndorsements, ApplicationLayerReferenceValues,
+        EventAttestationResults,
+    },
+    Variant,
 };
-use prost_types::Any;
 
 use crate::{
     compare::compare_application_layer_measurement_digests,
@@ -38,22 +40,23 @@ impl ApplicationPolicy {
     }
 }
 
-// We have to use [`Policy<[u8], Any>`] instead of [`EventPolicy`], because
+// We have to use [`Policy<[u8], Variant>`] instead of [`EventPolicy`], because
 // Rust doesn't yet support implementing trait aliases.
 // <https://github.com/rust-lang/rfcs/blob/master/text/1733-trait-alias.md>
-impl Policy<[u8], Any> for ApplicationPolicy {
+impl Policy<[u8], Variant> for ApplicationPolicy {
     fn verify(
         &self,
         encoded_event: &[u8],
-        encoded_event_endorsement: &Any,
+        encoded_event_endorsement: &Variant,
         milliseconds_since_epoch: i64,
     ) -> anyhow::Result<EventAttestationResults> {
         let event = decode_event_proto::<ApplicationLayerData>(
             "type.googleapis.com/oak.attestation.v1.ApplicationLayerData",
             encoded_event,
         )?;
+        // TODO: b/375137648 - Decode into new endorsement protos.
         let event_endorsement = decode_endorsement_proto::<ApplicationLayerEndorsements>(
-            "type.googleapis.com/oak.attestation.v1.ApplicationLayerEndorsements",
+            APPLICATION_ENDORSEMENT_ID,
             encoded_event_endorsement,
         )?;
 

oak_attestation_verification/src/policy/binary.rs

@@ -16,10 +16,10 @@
 
 use anyhow::Context;
 use oak_attestation_verification_types::policy::Policy;
-use oak_proto_rust::oak::attestation::v1::{
-    EventAttestationResults, EventData, EventReferenceValues,
+use oak_proto_rust::oak::{
+    attestation::v1::{EventAttestationResults, EventData, EventReferenceValues},
+    Variant,
 };
-use prost_types::Any;
 
 use crate::{
     compare::compare_event_measurement_digests, expect::get_event_expected_values,
@@ -36,11 +36,11 @@ impl BinaryPolicy {
     }
 }
 
-impl Policy<[u8], Any> for BinaryPolicy {
+impl Policy<[u8], Variant> for BinaryPolicy {
     fn verify(
         &self,
         encoded_event: &[u8],
-        _encoded_event_endorsement: &Any,
+        _encoded_event_endorsement: &Variant,
         milliseconds_since_epoch: i64,
     ) -> anyhow::Result<EventAttestationResults> {
         let event = decode_event_proto::<EventData>(

oak_attestation_verification/src/policy/container.rs

@@ -15,12 +15,14 @@
 //
 
 use anyhow::Context;
-use oak_attestation_verification_types::policy::Policy;
-use oak_proto_rust::oak::attestation::v1::{
-    ContainerLayerData, ContainerLayerEndorsements, ContainerLayerReferenceValues,
-    EventAttestationResults,
+use oak_attestation_verification_types::{policy::Policy, CONTAINER_ENDORSEMENT_ID};
+use oak_proto_rust::oak::{
+    attestation::v1::{
+        ContainerLayerData, ContainerLayerEndorsements, ContainerLayerReferenceValues,
+        EventAttestationResults,
+    },
+    Variant,
 };
-use prost_types::Any;
 
 use crate::{
     compare::compare_container_layer_measurement_digests,
@@ -38,22 +40,23 @@ impl ContainerPolicy {
     }
 }
 
-// We have to use [`Policy<[u8], Any>`] instead of [`EventPolicy`], because
+// We have to use [`Policy<[u8], Variant>`] instead of [`EventPolicy`], because
 // Rust doesn't yet support implementing trait aliases.
 // <https://github.com/rust-lang/rfcs/blob/master/text/1733-trait-alias.md>
-impl Policy<[u8], Any> for ContainerPolicy {
+impl Policy<[u8], Variant> for ContainerPolicy {
     fn verify(
         &self,
         encoded_event: &[u8],
-        encoded_event_endorsement: &Any,
+        encoded_event_endorsement: &Variant,
         milliseconds_since_epoch: i64,
     ) -> anyhow::Result<EventAttestationResults> {
         let event = decode_event_proto::<ContainerLayerData>(
             "type.googleapis.com/oak.attestation.v1.ContainerLayerData",
             encoded_event,
         )?;
+        // TODO: b/375137648 - Decode into new endorsement protos.
         let event_endorsement = decode_endorsement_proto::<ContainerLayerEndorsements>(
-            "type.googleapis.com/oak.attestation.v1.ContainerLayerEndorsements",
+            CONTAINER_ENDORSEMENT_ID,
             encoded_event_endorsement,
         )?;
 

oak_attestation_verification/src/policy/firmware.rs

@@ -15,14 +15,15 @@
 //
 
 use anyhow::Context;
-use oak_attestation_verification_types::policy::Policy;
-use oak_proto_rust::oak::attestation::v1::{
-    BinaryReferenceValue, EventAttestationResults, FirmwareEndorsement,
+use oak_attestation_verification_types::{policy::Policy, FIRMWARE_ENDORSEMENT_ID};
+use oak_proto_rust::oak::{
+    attestation::v1::{BinaryReferenceValue, EventAttestationResults, FirmwareEndorsement},
+    Variant,
 };
 
 use crate::{
     compare::compare_measurement_digest, expect::get_stage0_expected_values,
-    platform::convert_amd_sev_snp_initial_measurement,
+    platform::convert_amd_sev_snp_initial_measurement, util::decode_endorsement_proto,
 };
 
 pub struct FirmwarePolicy {
@@ -35,14 +36,19 @@ impl FirmwarePolicy {
     }
 }
 
-impl Policy<[u8], FirmwareEndorsement> for FirmwarePolicy {
+impl Policy<[u8], Variant> for FirmwarePolicy {
     fn verify(
         &self,
         firmware_measurement: &[u8],
-        _firmware_endorsement: &FirmwareEndorsement,
+        encoded_firmware_endorsement: &Variant,
         milliseconds_since_epoch: i64,
     ) -> anyhow::Result<EventAttestationResults> {
         let initial_measurement = convert_amd_sev_snp_initial_measurement(firmware_measurement);
+        let _firmware_endorsement = decode_endorsement_proto::<FirmwareEndorsement>(
+            FIRMWARE_ENDORSEMENT_ID,
+            encoded_firmware_endorsement,
+        )?;
+
         let initial_measurement_expected_values = get_stage0_expected_values(
             milliseconds_since_epoch,
             // TODO: b/375137648 - Use firmware endorsement, once we switch to new endorsment

oak_attestation_verification/src/policy/kernel.rs

@@ -15,12 +15,14 @@
 //
 
 use anyhow::Context;
-use oak_attestation_verification_types::policy::Policy;
-use oak_proto_rust::oak::attestation::v1::{
-    EventAttestationResults, KernelLayerEndorsements, KernelLayerReferenceValues,
-    Stage0Measurements,
+use oak_attestation_verification_types::{policy::Policy, KERNEL_ENDORSEMENT_ID};
+use oak_proto_rust::oak::{
+    attestation::v1::{
+        EventAttestationResults, KernelLayerEndorsements, KernelLayerReferenceValues,
+        Stage0Measurements,
+    },
+    Variant,
 };
-use prost_types::Any;
 
 use crate::{
     compare::compare_kernel_layer_measurement_digests,
@@ -39,20 +41,21 @@ impl KernelPolicy {
     }
 }
 
-impl Policy<[u8], Any> for KernelPolicy {
+impl Policy<[u8], Variant> for KernelPolicy {
     fn verify(
         &self,
         encoded_event: &[u8],
-        encoded_event_endorsement: &Any,
+        encoded_event_endorsement: &Variant,
         milliseconds_since_epoch: i64,
     ) -> anyhow::Result<EventAttestationResults> {
         let event =
             stage0_measurements_to_kernel_layer_data(decode_event_proto::<Stage0Measurements>(
                 "type.googleapis.com/oak.attestation.v1.Stage0Measurements",
                 encoded_event,
             )?);
+        // TODO: b/375137648 - Decode into new endorsement protos.
         let event_endorsements = decode_endorsement_proto::<KernelLayerEndorsements>(
-            "type.googleapis.com/oak.attestation.v1.KernelLayerEndorsements",
+            KERNEL_ENDORSEMENT_ID,
             encoded_event_endorsement,
         )?;
 

oak_attestation_verification/src/policy/platform.rs

@@ -15,9 +15,10 @@
 //
 
 use anyhow::Context;
-use oak_attestation_verification_types::policy::Policy;
-use oak_proto_rust::oak::attestation::v1::{
-    AmdSevReferenceValues, AmdSevSnpEndorsement, EventAttestationResults,
+use oak_attestation_verification_types::{policy::Policy, AMD_SEV_SNP_PLATFORM_ENDORSEMENT_ID};
+use oak_proto_rust::oak::{
+    attestation::v1::{AmdSevReferenceValues, AmdSevSnpEndorsement, EventAttestationResults},
+    Variant,
 };
 use oak_sev_snp_attestation_report::AttestationReport;
 
@@ -27,6 +28,7 @@ use crate::{
         convert_amd_sev_snp_attestation_report, verify_amd_sev_attestation_report_values,
         verify_amd_sev_snp_attestation_report_validity,
     },
+    util::decode_endorsement_proto,
 };
 
 pub struct AmdSevSnpPolicy {
@@ -39,13 +41,18 @@ impl AmdSevSnpPolicy {
     }
 }
 
-impl Policy<AttestationReport, AmdSevSnpEndorsement> for AmdSevSnpPolicy {
+impl Policy<AttestationReport, Variant> for AmdSevSnpPolicy {
     fn verify(
         &self,
         attestation_report: &AttestationReport,
-        platform_endorsement: &AmdSevSnpEndorsement,
+        encoded_platform_endorsement: &Variant,
         milliseconds_since_epoch: i64,
     ) -> anyhow::Result<EventAttestationResults> {
+        let platform_endorsement = decode_endorsement_proto::<AmdSevSnpEndorsement>(
+            AMD_SEV_SNP_PLATFORM_ENDORSEMENT_ID,
+            encoded_platform_endorsement,
+        )?;
+
         // Ensure the Attestation report is properly signed by the platform and the
         // corresponding certificate is signed by AMD.
         verify_amd_sev_snp_attestation_report_validity(

oak_attestation_verification/src/policy/system.rs

@@ -15,11 +15,14 @@
 //
 
 use anyhow::Context;
-use oak_attestation_verification_types::policy::Policy;
-use oak_proto_rust::oak::attestation::v1::{
-    EventAttestationResults, SystemLayerData, SystemLayerEndorsements, SystemLayerReferenceValues,
+use oak_attestation_verification_types::{policy::Policy, SYSTEM_ENDORSEMENT_ID};
+use oak_proto_rust::oak::{
+    attestation::v1::{
+        EventAttestationResults, SystemLayerData, SystemLayerEndorsements,
+        SystemLayerReferenceValues,
+    },
+    Variant,
 };
-use prost_types::Any;
 
 use crate::{
     compare::compare_system_layer_measurement_digests,
@@ -37,19 +40,20 @@ impl SystemPolicy {
     }
 }
 
-impl Policy<[u8], Any> for SystemPolicy {
+impl Policy<[u8], Variant> for SystemPolicy {
     fn verify(
         &self,
         encoded_event: &[u8],
-        encoded_event_endorsement: &Any,
+        encoded_event_endorsement: &Variant,
         milliseconds_since_epoch: i64,
     ) -> anyhow::Result<EventAttestationResults> {
         let event = decode_event_proto::<SystemLayerData>(
             "type.googleapis.com/oak.attestation.v1.SystemLayerData",
             encoded_event,
         )?;
+        // TODO: b/375137648 - Decode into new endorsement protos.
         let event_endorsements = decode_endorsement_proto::<SystemLayerEndorsements>(
-            "type.googleapis.com/oak.attestation.v1.SystemLayerEndorsements",
+            SYSTEM_ENDORSEMENT_ID,
             encoded_event_endorsement,
         )?;
 

oak_attestation_verification/src/util.rs

@@ -30,7 +30,7 @@ use oak_proto_rust::oak::{
         RootLayerData, RootLayerReferenceValues, Signature, SkipVerification, StringLiterals,
         SystemLayerReferenceValues, TextReferenceValue, Validity, VerifyingKeySet,
     },
-    HexDigest, RawDigest,
+    HexDigest, RawDigest, Variant,
 };
 use p256::pkcs8::{der::Decode, DecodePublicKey};
 use prost::Message;
@@ -428,10 +428,16 @@ pub fn decode_event_proto<M: Message + Default>(
 
 /// Decodes serialized endorsement into a specified [`Message`].
 pub fn decode_endorsement_proto<M: Message + Default>(
-    expected_type_url: &str,
-    endorsement_proto: &Any,
+    id: &[u8],
+    message: &Variant,
 ) -> anyhow::Result<M> {
-    decode_protobuf_any::<M>(expected_type_url, endorsement_proto)
+    if message.id == id {
+        let decoded_message = M::decode(message.value.as_ref())
+            .map_err(|error| anyhow::anyhow!("couldn't decode endorsement: {:?}", error))?;
+        Ok(decoded_message)
+    } else {
+        anyhow::bail!("unexpected endorsement ID, expected {:?}, found {:?}", id, message.id);
+    }
 }
 
 /// Decodes [`Any`] message into a specified [`Message`].