From 40b07b58b5ec52418e7b7eaf9e96901f123ed214 Mon Sep 17 00:00:00 2001
From: Kevin Wang <wy721@qq.com>
Date: Mon, 3 Aug 2026 19:47:35 -0700
Subject: [PATCH] fix(kms): send the Finish response before shutting down

`Onboard.Finish` called `std::process::exit(0)` from inside the RPC handler, so
the process died before Rocket could write the response. The onboarding client
saw an EOF and could not tell whether finalization had succeeded.

Use Rocket's graceful shutdown instead, which is what the web UI's `GET /finish`
route has always done: the server finishes in-flight responses and then stops
accepting connections. `run_onboard_service` then returns normally and `main`
continues into the regular KMS startup path, exactly as it does after the web-UI
flow.

The shutdown handle only exists after ignition, so the launch is split into
`ignite()` and `launch()` with the handle stashed in `OnboardState` in between.
It is written once, before any request can be served, so a `OnceLock` is enough.
---
 dstack/kms/src/main.rs            |  9 +++++++--
 dstack/kms/src/onboard_service.rs | 22 ++++++++++++++++++++--
 2 files changed, 27 insertions(+), 4 deletions(-)

diff --git a/dstack/kms/src/main.rs b/dstack/kms/src/main.rs
index 478c5c52a..74921a698 100644
--- a/dstack/kms/src/main.rs
+++ b/dstack/kms/src/main.rs
@@ -63,13 +63,18 @@ async fn run_onboard_service(kms_config: KmsConfig, figment: Figment) -> Result<
 
     // Remove section tls
 
-    let _ = rocket::custom(figment)
+    let rocket = rocket::custom(figment)
         .mount("/", rocket::routes![index, finish])
         .mount(
             "/prpc",
             ra_rpc::prpc_routes!(OnboardState, OnboardHandler, trim: "Onboard."),
         )
-        .manage(state)
+        .manage(state.clone())
+        .ignite()
+        .await
+        .map_err(|err| anyhow!(err.to_string()))?;
+    state.set_shutdown(rocket.shutdown())?;
+    let _ = rocket
         .launch()
         .await
         .map_err(|err| anyhow!(err.to_string()))?;
diff --git a/dstack/kms/src/onboard_service.rs b/dstack/kms/src/onboard_service.rs
index 334a6c6db..58651833a 100644
--- a/dstack/kms/src/onboard_service.rs
+++ b/dstack/kms/src/onboard_service.rs
@@ -2,7 +2,7 @@
 //
 // SPDX-License-Identifier: Apache-2.0
 
-use std::sync::{Arc, Mutex};
+use std::sync::{Arc, Mutex, OnceLock};
 
 use anyhow::{bail, Context, Result};
 use dstack_kms_rpc::{
@@ -45,6 +45,7 @@ pub struct OnboardState {
     config: KmsConfig,
     attestation_verifier: Arc<AttestationVerifier>,
     bootstrap_lock: Arc<AsyncMutex<()>>,
+    shutdown: Arc<OnceLock<rocket::Shutdown>>,
 }
 
 impl OnboardState {
@@ -57,8 +58,17 @@ impl OnboardState {
             config,
             attestation_verifier,
             bootstrap_lock: Arc::new(AsyncMutex::new(())),
+            shutdown: Arc::new(OnceLock::new()),
         })
     }
+
+    /// Hand the Rocket shutdown handle to the service so `finish` can stop the
+    /// server after its response has been sent.
+    pub fn set_shutdown(&self, shutdown: rocket::Shutdown) -> Result<()> {
+        self.shutdown
+            .set(shutdown)
+            .map_err(|_| anyhow::anyhow!("onboard shutdown handle is already set"))
+    }
 }
 
 pub struct OnboardHandler {
@@ -199,7 +209,15 @@ impl OnboardRpc for OnboardHandler {
     }
 
     async fn finish(self) -> anyhow::Result<()> {
-        std::process::exit(0);
+        let shutdown = self
+            .state
+            .shutdown
+            .get()
+            .context("onboard shutdown handle is unavailable")?;
+        // Graceful shutdown lets Rocket finish sending this response before the
+        // server stops, so the client learns that onboarding succeeded.
+        shutdown.clone().notify();
+        Ok(())
     }
 }