fix(kms): unstick onboard UI after a repeated click

dstack/kms/rpc/proto/kms_rpc.proto

@@ -185,6 +185,18 @@ message AttestationInfoResponse {
   bytes ppid = 8;
 }
 
+// Public material already written by Bootstrap or Onboard, so the setup page
+// can resume after a refresh without repeating those RPCs.
+message SetupStatusResponse {
+  bool provisioned = 1;
+  bytes k256_pubkey = 2;
+  bytes ca_pubkey = 3;
+  bytes attestation = 4;
+  // All key and certificate files are present, so Finish can start the KMS.
+  // If ready is true, provisioned is also true.
+  bool ready = 5;
+}
+
 // The Onboard RPC service.
 service Onboard {
   // Bootstrap a new KMS
@@ -193,6 +205,8 @@ service Onboard {
   rpc Onboard(OnboardRequest) returns (OnboardResponse);
   // Get attestation info for on-chain KMS authorization
   rpc GetAttestationInfo(google.protobuf.Empty) returns (AttestationInfoResponse);
+  // Public keys already stored for this KMS, if any
+  rpc GetSetupStatus(google.protobuf.Empty) returns (SetupStatusResponse);
   // Finish onboarding
   rpc Finish(google.protobuf.Empty) returns (google.protobuf.Empty);
 }

dstack/kms/src/onboard_service.rs

@@ -13,7 +13,7 @@ use dstack_kms_rpc::{
     kms_client::KmsClient,
     onboard_server::{OnboardRpc, OnboardServer},
     AttestationInfoResponse, BootstrapRequest, BootstrapResponse, GetKmsKeyRequest, KmsKeyResponse,
-    OnboardRequest, OnboardResponse,
+    OnboardRequest, OnboardResponse, SetupStatusResponse,
 };
 use fs_err as fs;
 use k256::ecdsa::SigningKey;
@@ -107,6 +107,31 @@ fn validate_onboarding_domain(domain: &str) -> Result<()> {
     Ok(())
 }
 
+fn stored_k256_pubkey(cfg: &KmsConfig) -> Result<Vec<u8>> {
+    let bytes = fs::read(cfg.k256_key()).context("failed to read k256 key")?;
+    let key = SigningKey::from_slice(&bytes).context("failed to parse k256 key")?;
+    Ok(key.verifying_key().to_sec1_bytes().to_vec())
+}
+
+fn setup_status(cfg: &KmsConfig) -> Result<SetupStatusResponse> {
+    if !cfg.root_keys_exist() {
+        return Ok(SetupStatusResponse::default());
+    }
+    let k256_pubkey = stored_k256_pubkey(cfg)?;
+    let info = fs::read(cfg.bootstrap_info())
+        .ok()
+        .and_then(|bytes| serde_json::from_slice::<BootstrapResponse>(&bytes).ok())
+        .filter(|info| info.k256_pubkey == k256_pubkey)
+        .unwrap_or_default();
+    Ok(SetupStatusResponse {
+        provisioned: true,
+        ready: cfg.keys_exists(),
+        k256_pubkey,
+        ca_pubkey: info.ca_pubkey,
+        attestation: info.attestation,
+    })
+}
+
 impl OnboardRpc for OnboardHandler {
     async fn bootstrap(self, request: BootstrapRequest) -> Result<BootstrapResponse> {
         validate_onboarding_domain(&request.domain)?;
@@ -216,6 +241,11 @@ impl OnboardRpc for OnboardHandler {
         )
     }
 
+    async fn get_setup_status(self) -> Result<SetupStatusResponse> {
+        let _bootstrap_guard = self.state.bootstrap_lock.lock().await;
+        setup_status(&self.state.config)
+    }
+
     async fn finish(self) -> anyhow::Result<()> {
         let shutdown = self
             .state
@@ -494,6 +524,78 @@ mod tests {
         assert_ne!(fs::read(cfg.root_ca_key()).unwrap(), b"stale");
     }
 
+    #[test]
+    fn setup_status_is_empty_before_root_keys_exist() {
+        let cert_dir = tempfile::tempdir().unwrap();
+        let cfg = auto_bootstrap_config(cert_dir.path());
+        let status = setup_status(&cfg).unwrap();
+        assert!(!status.provisioned);
+        assert!(!status.ready);
+        assert!(status.k256_pubkey.is_empty());
+        assert!(status.ca_pubkey.is_empty());
+        assert!(status.attestation.is_empty());
+    }
+
+    #[rocket::async_test]
+    async fn setup_status_includes_bootstrap_info_when_present() {
+        let cert_dir = tempfile::tempdir().unwrap();
+        let cfg = auto_bootstrap_config(cert_dir.path());
+        let keys = Keys::generate("kms.example.com", false).await.unwrap();
+        keys.store(&cfg).unwrap();
+        let expected_k256 = keys.k256_key.verifying_key().to_sec1_bytes().to_vec();
+        let stored = BootstrapResponse {
+            ca_pubkey: vec![1, 2, 3],
+            k256_pubkey: expected_k256.clone(),
+            attestation: vec![7, 8, 9],
+        };
+        fs::write(cfg.bootstrap_info(), serde_json::to_vec(&stored).unwrap()).unwrap();
+        let status = setup_status(&cfg).unwrap();
+        assert!(status.provisioned);
+        assert!(status.ready);
+        assert_eq!(status.ca_pubkey, stored.ca_pubkey);
+        assert_eq!(status.attestation, stored.attestation);
+        assert_eq!(status.k256_pubkey, expected_k256);
+    }
+
+    #[rocket::async_test]
+    async fn setup_status_ignores_absent_or_invalid_bootstrap_info() {
+        let cert_dir = tempfile::tempdir().unwrap();
+        let cfg = auto_bootstrap_config(cert_dir.path());
+        let keys = Keys::generate("kms.example.com", false).await.unwrap();
+        let expected_k256 = keys.k256_key.verifying_key().to_sec1_bytes().to_vec();
+        keys.store(&cfg).unwrap();
+        let stored = BootstrapResponse {
+            ca_pubkey: vec![1, 2, 3],
+            k256_pubkey: vec![4, 5, 6],
+            attestation: vec![7, 8, 9],
+        };
+        for ready in [true, false] {
+            if !ready {
+                fs::remove_file(cfg.rpc_cert()).unwrap();
+                fs::remove_file(cfg.bootstrap_info()).unwrap();
+            }
+            for info in [
+                None,
+                Some(b"invalid json".to_vec()),
+                Some(serde_json::to_vec(&stored).unwrap()),
+            ] {
+                if let Some(info) = info {
+                    fs::write(cfg.bootstrap_info(), info).unwrap();
+                }
+                let status = setup_status(&cfg).unwrap();
+                assert!(status.provisioned);
+                assert_eq!(status.ready, ready);
+                assert_eq!(status.k256_pubkey, expected_k256);
+                assert!(status.ca_pubkey.is_empty());
+                assert!(status.attestation.is_empty());
+            }
+        }
+
+        fs::write(cfg.k256_key(), b"invalid key").unwrap();
+        let err = setup_status(&cfg).unwrap_err();
+        assert!(format!("{err:#}").contains("failed to parse k256 key"));
+    }
+
     #[test]
     fn ca_certificate_is_renewed_only_within_the_renewal_window() {
         let now = UNIX_EPOCH + Duration::from_secs(2_000_000_000);