refactor: extract AttestationV1::with_report_data to deduplicate patching logic

The same report_data patching logic was duplicated in three places:
- guest-agent-simulator/src/simulator.rs
- guest-agent/src/rpc_service.rs (test)
- dstack-attest/src/attestation.rs (test)

Extract it into AttestationV1::with_report_data() method and replace
all three copies with calls to it.

Author: kvinwang

dstack-attest/src/attestation.rs

@@ -1221,50 +1221,7 @@ mod tests {
     use super::*;
 
     fn patch_v1_report_data(attestation: AttestationV1, report_data: [u8; 64]) -> AttestationV1 {
-        let AttestationV1 {
-            version,
-            platform,
-            stack,
-        } = attestation;
-        let platform = match platform {
-            PlatformEvidence::Tdx {
-                mut quote,
-                event_log,
-            } => {
-                if quote.len() >= TDX_QUOTE_REPORT_DATA_RANGE.end {
-                    quote[TDX_QUOTE_REPORT_DATA_RANGE].copy_from_slice(&report_data);
-                }
-                PlatformEvidence::Tdx { quote, event_log }
-            }
-            other => other,
-        };
-        let stack = match stack {
-            StackEvidence::Dstack {
-                runtime_events,
-                config,
-                ..
-            } => StackEvidence::Dstack {
-                report_data: report_data.to_vec(),
-                runtime_events,
-                config,
-            },
-            StackEvidence::DstackPod {
-                runtime_events,
-                config,
-                report_data_payload,
-                ..
-            } => StackEvidence::DstackPod {
-                report_data: report_data.to_vec(),
-                runtime_events,
-                config,
-                report_data_payload,
-            },
-        };
-        AttestationV1 {
-            version,
-            platform,
-            stack,
-        }
+        attestation.with_report_data(report_data)
     }
 
     fn dummy_tdx_attestation(report_data: [u8; 64]) -> Attestation {

dstack-attest/src/v1.rs

@@ -192,6 +192,51 @@ impl Attestation {
             stack: self.stack.into_dstack_pod(report_data_payload),
         }
     }
+
+    /// Return a new attestation with the report_data patched in both platform quote and stack.
+    pub fn with_report_data(self, report_data: [u8; 64]) -> Self {
+        use crate::attestation::TDX_QUOTE_REPORT_DATA_RANGE;
+
+        let platform = match self.platform {
+            PlatformEvidence::Tdx {
+                mut quote,
+                event_log,
+            } => {
+                if quote.len() >= TDX_QUOTE_REPORT_DATA_RANGE.end {
+                    quote[TDX_QUOTE_REPORT_DATA_RANGE].copy_from_slice(&report_data);
+                }
+                PlatformEvidence::Tdx { quote, event_log }
+            }
+            other => other,
+        };
+        let stack = match self.stack {
+            StackEvidence::Dstack {
+                runtime_events,
+                config,
+                ..
+            } => StackEvidence::Dstack {
+                report_data: report_data.to_vec(),
+                runtime_events,
+                config,
+            },
+            StackEvidence::DstackPod {
+                runtime_events,
+                config,
+                report_data_payload,
+                ..
+            } => StackEvidence::DstackPod {
+                report_data: report_data.to_vec(),
+                runtime_events,
+                config,
+                report_data_payload,
+            },
+        };
+        Self {
+            version: self.version,
+            platform,
+            stack,
+        }
+    }
 }
 
 #[cfg(test)]

guest-agent-simulator/src/simulator.rs

@@ -7,8 +7,7 @@ use std::path::Path;
 use anyhow::{anyhow, Context, Result};
 use dstack_guest_agent_rpc::{AttestResponse, GetQuoteResponse};
 use ra_tls::attestation::{
-    AttestationV1, PlatformEvidence, QuoteContentType, StackEvidence, TdxAttestationExt,
-    VersionedAttestation, TDX_QUOTE_REPORT_DATA_RANGE,
+    AttestationV1, QuoteContentType, TdxAttestationExt, VersionedAttestation,
 };
 use std::fs;
 use tracing::warn;
@@ -89,52 +88,5 @@ fn maybe_patch_report_data(
         );
         return attestation.clone().into_v1();
     }
-    patch_v1_report_data(attestation.clone().into_v1(), report_data)
-}
-
-fn patch_v1_report_data(attestation: AttestationV1, report_data: [u8; 64]) -> AttestationV1 {
-    let AttestationV1 {
-        version,
-        platform,
-        stack,
-    } = attestation;
-    let platform = match platform {
-        PlatformEvidence::Tdx {
-            mut quote,
-            event_log,
-        } => {
-            if quote.len() >= TDX_QUOTE_REPORT_DATA_RANGE.end {
-                quote[TDX_QUOTE_REPORT_DATA_RANGE].copy_from_slice(&report_data);
-            }
-            PlatformEvidence::Tdx { quote, event_log }
-        }
-        other => other,
-    };
-    let stack = match stack {
-        StackEvidence::Dstack {
-            runtime_events,
-            config,
-            ..
-        } => StackEvidence::Dstack {
-            report_data: report_data.to_vec(),
-            runtime_events,
-            config,
-        },
-        StackEvidence::DstackPod {
-            runtime_events,
-            config,
-            report_data_payload,
-            ..
-        } => StackEvidence::DstackPod {
-            report_data: report_data.to_vec(),
-            runtime_events,
-            config,
-            report_data_payload,
-        },
-    };
-    AttestationV1 {
-        version,
-        platform,
-        stack,
-    }
+    attestation.clone().into_v1().with_report_data(report_data)
 }

guest-agent/src/rpc_service.rs

@@ -670,10 +670,7 @@ mod tests {
         Signature as Ed25519Signature, Verifier, VerifyingKey as Ed25519VerifyingKey,
     };
     use k256::ecdsa::{Signature as K256Signature, VerifyingKey};
-    use ra_tls::attestation::{
-        AttestationV1, PlatformEvidence, StackEvidence, VersionedAttestation,
-        TDX_QUOTE_REPORT_DATA_RANGE,
-    };
+    use ra_tls::attestation::{AttestationV1, VersionedAttestation};
     use sha2::Sha256;
     use std::collections::HashSet;
     use std::convert::TryFrom;
@@ -817,51 +814,7 @@ pNs85uhOZE8z2jr8Pg==
             attestation: &VersionedAttestation,
             report_data: [u8; 64],
         ) -> AttestationV1 {
-            let attestation = attestation.clone().into_v1();
-            let AttestationV1 {
-                version,
-                platform,
-                stack,
-            } = attestation;
-            let platform = match platform {
-                PlatformEvidence::Tdx {
-                    mut quote,
-                    event_log,
-                } => {
-                    if quote.len() >= TDX_QUOTE_REPORT_DATA_RANGE.end {
-                        quote[TDX_QUOTE_REPORT_DATA_RANGE].copy_from_slice(&report_data);
-                    }
-                    PlatformEvidence::Tdx { quote, event_log }
-                }
-                other => other,
-            };
-            let stack = match stack {
-                StackEvidence::Dstack {
-                    runtime_events,
-                    config,
-                    ..
-                } => StackEvidence::Dstack {
-                    report_data: report_data.to_vec(),
-                    runtime_events,
-                    config,
-                },
-                StackEvidence::DstackPod {
-                    runtime_events,
-                    config,
-                    report_data_payload,
-                    ..
-                } => StackEvidence::DstackPod {
-                    report_data: report_data.to_vec(),
-                    runtime_events,
-                    config,
-                    report_data_payload,
-                },
-            };
-            AttestationV1 {
-                version,
-                platform,
-                stack,
-            }
+            attestation.clone().into_v1().with_report_data(report_data)
         }
 
         impl PlatformBackend for TestSimulatorPlatform {