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..0fa2f1bd 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 { @@ -44,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; } 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(), }) }