attestation: use direct CBOR V1 format

Author: kvinwang

Cargo.lock

@@ -41,7 +41,6 @@ members = [
     "guest-api",
     "load_config",
     "key-provider-client",
-    "dstack-attestation-types",
     "dstack-types",
     "cert-client",
     "lspci",
@@ -73,7 +72,6 @@ supervisor = { path = "supervisor" }
 supervisor-client = { path = "supervisor/client" }
 tdx-attest = { path = "tdx-attest" }
 dstack-attest = { path = "dstack-attest" }
-dstack-attestation-types = { path = "dstack-attestation-types" }
 certbot = { path = "certbot" }
 rocket-vsock-listener = { path = "rocket-vsock-listener" }
 host-api = { path = "host-api", default-features = false }

Cargo.toml

@@ -14,7 +14,7 @@ anyhow.workspace = true
 cc-eventlog.workspace = true
 ciborium.workspace = true
 dcap-qvl.workspace = true
-dstack-attestation-types.workspace = true
+dstack-attestation-types = { path = "../../../dstack-attestation-types" }
 dstack-types.workspace = true
 ez-hash.workspace = true
 fs-err.workspace = true

dstack-attest/Cargo.toml

@@ -27,12 +27,17 @@ use dstack_types::SysConfig;
 use dstack_types::{Platform, VmConfig};
 use ez_hash::{sha256, Hasher, Sha384};
 use or_panic::ResultOrPanic;
-use scale::{Decode, Encode};
+use scale::{Decode, Encode, Error as ScaleError, Input, Output};
 use serde::de::DeserializeOwned;
 use serde::{Deserialize, Serialize};
 use serde_human_bytes as hex_bytes;
 use sha2::Digest as _;
 
+const NATIVE_PLATFORM_TDX: &str = "tdx";
+const NATIVE_PLATFORM_GCP_TDX: &str = "gcp-tdx";
+const NATIVE_PLATFORM_NITRO_ENCLAVE: &str = "nitro-enclave";
+const DSTACK_VARIANT_DSTACK: &str = "dstack";
+const DSTACK_VARIANT_DSTACK_K8S: &str = "dstack-k8s";
 const DSTACK_TDX: &str = "dstack-tdx";
 const DSTACK_GCP_TDX: &str = "dstack-gcp-tdx";
 const DSTACK_NITRO_ENCLAVE: &str = "dstack-nitro-enclave";
@@ -132,7 +137,15 @@ fn report_data_payload_v1(attestation: &AttestationV1) -> Option<&str> {
     }
 }
 
-fn legacy_into_v1(attestation: Attestation, extensions: Vec<V1Extension>) -> AttestationV1 {
+fn is_cbor_map_prefix(byte: u8) -> bool {
+    matches!(byte, 0xa0..=0xbb | 0xbf)
+}
+
+fn legacy_into_v1(
+    attestation: Attestation,
+    variant: &str,
+    extensions: Vec<V1Extension>,
+) -> AttestationV1 {
     let Attestation {
         quote,
         runtime_events,
@@ -141,7 +154,12 @@ fn legacy_into_v1(attestation: Attestation, extensions: Vec<V1Extension>) -> Att
         report: _,
     } = attestation;
 
-    let mut platform = PlatformEvidence::new(quote.mode().as_str());
+    let native_platform = match &quote {
+        AttestationQuote::DstackTdx(_) => NATIVE_PLATFORM_TDX,
+        AttestationQuote::DstackGcpTdx => NATIVE_PLATFORM_GCP_TDX,
+        AttestationQuote::DstackNitroEnclave => NATIVE_PLATFORM_NITRO_ENCLAVE,
+    };
+    let mut platform = PlatformEvidence::new(native_platform, variant);
     match quote {
         AttestationQuote::DstackTdx(TdxQuote { quote, event_log }) => {
             platform
@@ -182,8 +200,8 @@ fn v1_into_legacy(attestation: AttestationV1) -> Result<Attestation> {
             .unwrap_or_default();
     let config = optional_text_attribute(&attestation.optional_attributes, ATTR_CONFIG)?;
 
-    let quote = match attestation.platform.platform_id.as_str() {
-        DSTACK_TDX => {
+    let quote = match attestation.platform.platform.as_str() {
+        NATIVE_PLATFORM_TDX => {
             let quote =
                 required_bytes_attribute(&attestation.platform.basic_attributes, ATTR_TDX_QUOTE)?;
             let event_log: Vec<TdxEvent> = optional_decoded_attribute(
@@ -193,9 +211,9 @@ fn v1_into_legacy(attestation: AttestationV1) -> Result<Attestation> {
             .unwrap_or_default();
             AttestationQuote::DstackTdx(TdxQuote { quote, event_log })
         }
-        DSTACK_GCP_TDX => AttestationQuote::DstackGcpTdx,
-        DSTACK_NITRO_ENCLAVE => AttestationQuote::DstackNitroEnclave,
-        other => bail!("Unsupported platform_id in AttestationV1: {other}"),
+        NATIVE_PLATFORM_GCP_TDX => AttestationQuote::DstackGcpTdx,
+        NATIVE_PLATFORM_NITRO_ENCLAVE => AttestationQuote::DstackNitroEnclave,
+        other => bail!("Unsupported platform in AttestationV1: {other}"),
     };
 
     Ok(Attestation {
@@ -426,8 +444,13 @@ pub struct NsmQuote {
     pub document: Vec<u8>,
 }
 
-/// Represents a versioned attestation.
 #[derive(Clone, Encode, Decode)]
+enum LegacyVersionedAttestation {
+    V0 { attestation: Attestation },
+}
+
+/// Represents a versioned attestation.
+#[derive(Clone)]
 pub enum VersionedAttestation {
     /// Legacy SCALE-encoded attestation.
     V0 {
@@ -441,15 +464,69 @@ pub enum VersionedAttestation {
     },
 }
 
+impl Encode for VersionedAttestation {
+    fn size_hint(&self) -> usize {
+        self.to_scale().len()
+    }
+
+    fn encode_to<T: Output + ?Sized>(&self, dest: &mut T) {
+        dest.write(&self.to_scale());
+    }
+}
+
+impl Decode for VersionedAttestation {
+    fn decode<I: Input>(input: &mut I) -> Result<Self, ScaleError> {
+        let Some(remaining_len) = input.remaining_len()? else {
+            return Err(ScaleError::from(
+                "VersionedAttestation requires a bounded input to decode",
+            ));
+        };
+        let mut bytes = vec![0u8; remaining_len];
+        input.read(&mut bytes)?;
+        Self::from_scale(&bytes).map_err(|err| {
+            ScaleError::from(std::io::Error::new(
+                std::io::ErrorKind::InvalidData,
+                err.to_string(),
+            ))
+        })
+    }
+}
+
 impl VersionedAttestation {
     /// Decode versioned attestation bytes.
     pub fn from_scale(scale: &[u8]) -> Result<Self> {
-        Self::decode(&mut &scale[..]).context("Failed to decode VersionedAttestation")
+        let Some(first) = scale.first().copied() else {
+            bail!("Empty attestation bytes");
+        };
+        if first == 0x00 {
+            let legacy = LegacyVersionedAttestation::decode(&mut &scale[..])
+                .context("Failed to decode legacy VersionedAttestation")?;
+            return match legacy {
+                LegacyVersionedAttestation::V0 { attestation } => Ok(Self::V0 { attestation }),
+            };
+        }
+        if is_cbor_map_prefix(first) {
+            let attestation = AttestationV1::from_cbor(scale)?;
+            return Ok(Self::V1 { attestation });
+        }
+        if first == 0x01 && scale.get(1).is_some_and(|byte| is_cbor_map_prefix(*byte)) {
+            let attestation = AttestationV1::from_cbor(&scale[1..])?;
+            return Ok(Self::V1 { attestation });
+        }
+        bail!("Unknown attestation wire format");
     }
 
     /// Encode versioned attestation bytes.
     pub fn to_scale(&self) -> Vec<u8> {
-        self.encode()
+        match self {
+            Self::V0 { attestation } => LegacyVersionedAttestation::V0 {
+                attestation: attestation.clone(),
+            }
+            .encode(),
+            Self::V1 { attestation } => attestation
+                .to_cbor()
+                .or_panic("AttestationV1 should encode as CBOR"),
+        }
     }
 
     /// Try to project any version into the legacy attestation structure.
@@ -884,11 +961,19 @@ impl Attestation {
 
 impl Attestation {
     pub fn into_v1(self) -> AttestationV1 {
-        legacy_into_v1(self, Vec::new())
+        legacy_into_v1(self, DSTACK_VARIANT_DSTACK, Vec::new())
     }
 
     pub fn into_v1_with_extensions(self, extensions: Vec<V1Extension>) -> AttestationV1 {
-        legacy_into_v1(self, extensions)
+        legacy_into_v1(self, DSTACK_VARIANT_DSTACK, extensions)
+    }
+
+    pub fn into_v1_with_variant_and_extensions(
+        self,
+        variant: &str,
+        extensions: Vec<V1Extension>,
+    ) -> AttestationV1 {
+        legacy_into_v1(self, variant, extensions)
     }
 
     /// Verify the quote with optional custom time (testing hook)
@@ -928,11 +1013,14 @@ impl Attestation {
         report_data_payload: String,
     ) -> VersionedAttestation {
         VersionedAttestation::V1 {
-            attestation: self.into_v1_with_extensions(vec![V1Extension::new(
-                EXT_REPORT_DATA_PAYLOAD,
-                true,
-                text_value(report_data_payload),
-            )]),
+            attestation: self.into_v1_with_variant_and_extensions(
+                DSTACK_VARIANT_DSTACK_K8S,
+                vec![V1Extension::new(
+                    EXT_REPORT_DATA_PAYLOAD,
+                    true,
+                    text_value(report_data_payload),
+                )],
+            ),
         }
     }
 
@@ -1119,9 +1207,15 @@ mod tests {
         let encoded = dummy_tdx_attestation(report_data)
             .into_versioned_with_report_data_payload(payload.clone())
             .to_scale();
+        assert!(matches!(encoded.first(), Some(0xa0..=0xbf)));
         let decoded = VersionedAttestation::from_scale(&encoded).expect("decode attestation");
         assert_eq!(decoded.report_data_payload(), Some(payload.as_str()));
         assert_eq!(decoded.clone().into_inner().report_data, report_data);
+        let VersionedAttestation::V1 { attestation } = decoded else {
+            panic!("expected V1 attestation");
+        };
+        assert_eq!(attestation.platform.platform, NATIVE_PLATFORM_TDX);
+        assert_eq!(attestation.platform.variant, DSTACK_VARIANT_DSTACK_K8S);
     }
 
     #[test]
@@ -1132,4 +1226,11 @@ mod tests {
         assert_eq!(patched.report_data_payload(), Some("payload"));
         assert_eq!(patched.into_inner().report_data, [9u8; 64]);
     }
+
+    #[test]
+    fn legacy_v0_upgrade_uses_dstack_variant() {
+        let upgraded = dummy_tdx_attestation([3u8; 64]).into_v1();
+        assert_eq!(upgraded.platform.platform, NATIVE_PLATFORM_TDX);
+        assert_eq!(upgraded.platform.variant, DSTACK_VARIANT_DSTACK);
+    }
 }