Add extra_info field to GetAppKey for auth API pass-through

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")?;

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;
 }
 

kms/src/main_service.rs

@@ -127,14 +127,18 @@ impl RpcHandler {
 
     async fn ensure_kms_allowed(&self, vm_config: &str) -> Result<BootInfo> {
         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<BootConfig> {
+    async fn ensure_app_boot_allowed(
+        &self,
+        vm_config: &str,
+        extra_info: &str,
+    ) -> Result<BootConfig> {
         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<BootConfig> {
-        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

kms/src/main_service/upgrade_authority.rs

@@ -37,6 +37,8 @@ pub(crate) struct BootInfo {
     pub key_provider_info: Vec<u8>,
     pub tcb_status: String,
     pub advisory_ids: Vec<String>,
+    #[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(),
     })
 }