Merge pull request #503 from Dstack-TEE/feat/onboard-attestation-info

feat(kms): add GetAttestationInfo RPC to onboard service

Author: kvinwang

kms/rpc/proto/kms_rpc.proto

@@ -133,12 +133,26 @@ message OnboardRequest {
 message OnboardResponse {
 }
 
+// Attestation info needed for on-chain KMS authorization.
+message AttestationInfoResponse {
+  // Device ID (SHA256 of platform device identifier)
+  bytes device_id = 1;
+  // Aggregated measurement of the VM execution environment
+  bytes mr_aggregated = 2;
+  // OS image hash
+  bytes os_image_hash = 3;
+  // Attestation mode (e.g. "dstack-tdx", "dstack-gcp-tdx")
+  string attestation_mode = 4;
+}
+
 // The Onboard RPC service.
 service Onboard {
   // Bootstrap a new KMS
   rpc Bootstrap(BootstrapRequest) returns (BootstrapResponse);
   // Onboard from existing KMS
   rpc Onboard(OnboardRequest) returns (OnboardResponse);
+  // Get attestation info for on-chain KMS authorization
+  rpc GetAttestationInfo(google.protobuf.Empty) returns (AttestationInfoResponse);
   // Finish onboarding
   rpc Finish(google.protobuf.Empty) returns (google.protobuf.Empty);
 }

kms/src/onboard_service.rs

@@ -9,7 +9,8 @@ use dstack_guest_agent_rpc::{
 use dstack_kms_rpc::{
     kms_client::KmsClient,
     onboard_server::{OnboardRpc, OnboardServer},
-    BootstrapRequest, BootstrapResponse, GetKmsKeyRequest, OnboardRequest, OnboardResponse,
+    AttestationInfoResponse, BootstrapRequest, BootstrapResponse, GetKmsKeyRequest, OnboardRequest,
+    OnboardResponse,
 };
 use fs_err as fs;
 use http_client::prpc::PrpcClient;
@@ -90,6 +91,47 @@ impl OnboardRpc for OnboardHandler {
         Ok(OnboardResponse {})
     }
 
+    async fn get_attestation_info(self) -> Result<AttestationInfoResponse> {
+        let pccs_url = self.state.config.pccs_url.clone();
+
+        // Get attestation from guest agent
+        let report_data = pad64([0u8; 32]);
+        let response = app_attest(report_data)
+            .await
+            .context("Failed to get attestation")?;
+
+        // Decode and verify the attestation to get real device ID
+        let attestation = VersionedAttestation::from_scale(&response.attestation)
+            .context("Failed to decode attestation")?
+            .into_inner();
+        let attestation_mode = serde_json::to_value(attestation.quote.mode())
+            .ok()
+            .and_then(|v| v.as_str().map(String::from))
+            .unwrap_or_else(|| format!("{:?}", attestation.quote.mode()));
+        let verified = attestation
+            .verify(pccs_url.as_deref())
+            .await
+            .context("Failed to verify attestation")?;
+
+        // Get vm_config from guest agent
+        let info = dstack_client()
+            .info()
+            .await
+            .context("Failed to get VM info")?;
+
+        // Decode app info to get device_id, mr_aggregated, os_image_hash, mr_system
+        let app_info = verified
+            .decode_app_info_ex(false, &info.vm_config)
+            .context("Failed to decode app info")?;
+
+        Ok(AttestationInfoResponse {
+            device_id: app_info.device_id,
+            mr_aggregated: app_info.mr_aggregated.to_vec(),
+            os_image_hash: app_info.os_image_hash,
+            attestation_mode,
+        })
+    }
+
     async fn finish(self) -> anyhow::Result<()> {
         std::process::exit(0);
     }

kms/src/www/onboard.html

@@ -113,13 +113,71 @@
         .fade-leave-to {
             opacity: 0;
         }
+
+        .attestation-info {
+            background-color: #f0f4f8;
+            border: 1px solid #ccd;
+            border-radius: 4px;
+            padding: 15px;
+            margin-bottom: 20px;
+        }
+
+        .attestation-info h3 {
+            margin-top: 0;
+            color: #444;
+        }
+
+        .info-row {
+            display: flex;
+            margin-bottom: 8px;
+            font-size: 0.9em;
+        }
+
+        .info-label {
+            font-weight: bold;
+            min-width: 160px;
+            color: #555;
+        }
+
+        .info-value {
+            font-family: monospace;
+            word-break: break-all;
+            color: #333;
+        }
+
+        .loading {
+            color: #888;
+            font-style: italic;
+        }
     </style>
 </head>
 
 <body>
     <div id="app" class="container">
         <h1>dstack KMS Setup</h1>
 
+        <div v-if="attestationLoading" class="loading">Loading attestation info...</div>
+        <div v-else-if="attestationError" class="error">Attestation info: {{ attestationError }}</div>
+        <div v-else-if="attestationInfo" class="attestation-info">
+            <h3>Attestation Info (for on-chain registration)</h3>
+            <div class="info-row">
+                <span class="info-label">Attestation Mode:</span>
+                <span class="info-value">{{ attestationInfo.attestation_mode }}</span>
+            </div>
+            <div class="info-row">
+                <span class="info-label">Device ID:</span>
+                <span class="info-value">0x{{ attestationInfo.device_id }}</span>
+            </div>
+            <div class="info-row">
+                <span class="info-label">MR Aggregated:</span>
+                <span class="info-value">0x{{ attestationInfo.mr_aggregated }}</span>
+            </div>
+            <div class="info-row">
+                <span class="info-label">OS Image Hash:</span>
+                <span class="info-value">0x{{ attestationInfo.os_image_hash }}</span>
+            </div>
+        </div>
+
         <div v-if="!setupFinished">
             <div v-if="!selectedOption" class="initial-buttons">
                 <button @click="selectedOption = 'bootstrap'">Bootstrap</button>
@@ -200,7 +258,24 @@ <h2>Onboard from an Existing KMS Instance</h2>
                     error: '',
                     success: '',
                     result: '',
-                    setupFinished: false
+                    setupFinished: false,
+                    attestationInfo: null,
+                    attestationLoading: true,
+                    attestationError: ''
+                }
+            },
+            async mounted() {
+                try {
+                    const data = await rpcCall('GetAttestationInfo', {});
+                    if (data.error) {
+                        this.attestationError = data.error;
+                    } else {
+                        this.attestationInfo = data;
+                    }
+                } catch (err) {
+                    this.attestationError = err.message;
+                } finally {
+                    this.attestationLoading = false;
                 }
             },
             methods: {