From 8d549bec08dc8b02a64550a0c83da55ae74f4408 Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Tue, 14 Apr 2026 18:40:48 -0700 Subject: [PATCH 1/2] Add extra_info field to GetAppKey for auth API pass-through Add an extra_info string field to GetAppKeyRequest that gets threaded through to BootInfo and forwarded to the auth webhook, allowing custom information to be passed for authorization decisions. --- dstack-util/src/system_setup.rs | 1 + kms/rpc/proto/kms_rpc.proto | 2 ++ kms/src/main_service.rs | 18 ++++++++++++------ kms/src/main_service/upgrade_authority.rs | 3 +++ 4 files changed, 18 insertions(+), 6 deletions(-) diff --git a/dstack-util/src/system_setup.rs b/dstack-util/src/system_setup.rs index b7f59cfd..6e185080 100644 --- a/dstack-util/src/system_setup.rs +++ b/dstack-util/src/system_setup.rs @@ -851,6 +851,7 @@ impl<'a> Stage0<'a> { .get_app_key(rpc::GetAppKeyRequest { api_version: 1, vm_config: self.shared.sys_config.vm_config.clone(), + extra_info: String::new(), }) .await .context("Failed to get app key")?; diff --git a/kms/rpc/proto/kms_rpc.proto b/kms/rpc/proto/kms_rpc.proto index adaaa5f5..2e3c5145 100644 --- a/kms/rpc/proto/kms_rpc.proto +++ b/kms/rpc/proto/kms_rpc.proto @@ -11,6 +11,8 @@ package kms; message GetAppKeyRequest { uint32 api_version = 1; string vm_config = 2; + // Custom info to be passed through to the auth API for decision-making. + string extra_info = 3; } message AppId { diff --git a/kms/src/main_service.rs b/kms/src/main_service.rs index b738cd29..9352a8cf 100644 --- a/kms/src/main_service.rs +++ b/kms/src/main_service.rs @@ -127,14 +127,18 @@ impl RpcHandler { async fn ensure_kms_allowed(&self, vm_config: &str) -> Result { let att = self.ensure_attested()?; - self.ensure_app_attestation_allowed(att, true, false, vm_config) + self.ensure_app_attestation_allowed(att, true, false, vm_config, "") .await .map(|c| c.boot_info) } - async fn ensure_app_boot_allowed(&self, vm_config: &str) -> Result { + async fn ensure_app_boot_allowed( + &self, + vm_config: &str, + extra_info: &str, + ) -> Result { let att = self.ensure_attested()?; - self.ensure_app_attestation_allowed(att, false, false, vm_config) + self.ensure_app_attestation_allowed(att, false, false, vm_config, extra_info) .await } @@ -191,8 +195,10 @@ impl RpcHandler { is_kms: bool, use_boottime_mr: bool, vm_config_str: &str, + extra_info: &str, ) -> Result { - let boot_info = build_boot_info(att, use_boottime_mr, vm_config_str)?; + let mut boot_info = build_boot_info(att, use_boottime_mr, vm_config_str)?; + boot_info.extra_info = extra_info.to_string(); let response = self .state .config @@ -244,7 +250,7 @@ impl KmsRpc for RpcHandler { boot_info, gateway_app_id, } = self - .ensure_app_boot_allowed(&request.vm_config) + .ensure_app_boot_allowed(&request.vm_config, &request.extra_info) .await .context("App not allowed")?; let app_id = boot_info.app_id; @@ -402,7 +408,7 @@ impl KmsRpc for RpcHandler { .await .context("Quote verification failed")?; let app_info = self - .ensure_app_attestation_allowed(&attestation, false, true, &request.vm_config) + .ensure_app_attestation_allowed(&attestation, false, true, &request.vm_config, "") .await?; let app_ca = self.derive_app_ca(&app_info.boot_info.app_id)?; let cert = app_ca diff --git a/kms/src/main_service/upgrade_authority.rs b/kms/src/main_service/upgrade_authority.rs index e98b436f..d1debb24 100644 --- a/kms/src/main_service/upgrade_authority.rs +++ b/kms/src/main_service/upgrade_authority.rs @@ -37,6 +37,8 @@ pub(crate) struct BootInfo { pub key_provider_info: Vec, pub tcb_status: String, pub advisory_ids: Vec, + #[serde(default, skip_serializing_if = "String::is_empty")] + pub extra_info: String, } pub(crate) fn build_boot_info( @@ -69,6 +71,7 @@ pub(crate) fn build_boot_info( key_provider_info: app_info.key_provider_info, tcb_status, advisory_ids, + extra_info: String::new(), }) } From 9ca3957cbafa70a6e0b7c0cdf39282f98bef1d8e Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Tue, 14 Apr 2026 18:43:23 -0700 Subject: [PATCH 2/2] fix: remove trailing whitespace in kms_rpc.proto --- kms/rpc/proto/kms_rpc.proto | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/kms/rpc/proto/kms_rpc.proto b/kms/rpc/proto/kms_rpc.proto index 2e3c5145..0fa2f1bd 100644 --- a/kms/rpc/proto/kms_rpc.proto +++ b/kms/rpc/proto/kms_rpc.proto @@ -46,7 +46,7 @@ message AppKeyResponse { string tproxy_app_id = 6; // Reverse proxy app ID from DstackKms contract. string gateway_app_id = 7; - // OS Image hash + // OS Image hash bytes os_image_hash = 8; }