kms: enforce self authorization on trusted RPCs

Author: kvinwang

kms/src/main_service.rs

@@ -22,8 +22,9 @@ use ra_tls::{
 };
 use scale::Decode;
 use sha2::Digest;
+use tokio::sync::OnceCell;
 use tracing::info;
-use upgrade_authority::{build_boot_info, BootInfo};
+use upgrade_authority::{build_boot_info, local_kms_boot_info, BootInfo};
 
 use crate::{
     config::KmsConfig,
@@ -52,6 +53,7 @@ pub struct KmsStateInner {
     temp_ca_cert: String,
     temp_ca_key: String,
     verifier: CvmVerifier,
+    self_boot_info: OnceCell<BootInfo>,
 }
 
 impl KmsState {
@@ -79,6 +81,7 @@ impl KmsState {
                 temp_ca_cert,
                 temp_ca_key,
                 verifier,
+                self_boot_info: OnceCell::new(),
             }),
         })
     }
@@ -95,6 +98,31 @@ struct BootConfig {
 }
 
 impl RpcHandler {
+    async fn ensure_self_allowed(&self) -> Result<()> {
+        if !self.state.config.onboard.quote_enabled {
+            return Ok(());
+        }
+        let boot_info = self
+            .state
+            .self_boot_info
+            .get_or_try_init(|| async {
+                local_kms_boot_info(self.state.config.pccs_url.as_deref()).await
+            })
+            .await
+            .context("Failed to load cached self boot info")?;
+        let response = self
+            .state
+            .config
+            .auth_api
+            .is_app_allowed(boot_info, true)
+            .await
+            .context("Failed to call self KMS auth check")?;
+        if !response.is_allowed {
+            bail!("KMS is not allowed: {}", response.reason);
+        }
+        Ok(())
+    }
+
     fn ensure_attested(&self) -> Result<&VerifiedAttestation> {
         let Some(attestation) = &self.attestation else {
             bail!("No attestation provided");
@@ -214,6 +242,9 @@ impl KmsRpc for RpcHandler {
         if request.api_version > 1 {
             bail!("Unsupported API version: {}", request.api_version);
         }
+        self.ensure_self_allowed()
+            .await
+            .context("KMS self authorization failed")?;
         let BootConfig {
             boot_info,
             gateway_app_id,
@@ -254,6 +285,9 @@ impl KmsRpc for RpcHandler {
     }
 
     async fn get_app_env_encrypt_pub_key(self, request: AppId) -> Result<PublicKeyResponse> {
+        self.ensure_self_allowed()
+            .await
+            .context("KMS self authorization failed")?;
         let secret = kdf::derive_dh_secret(
             &self.state.root_ca.key,
             &[&request.app_id[..], "env-encrypt-key".as_bytes()],
@@ -320,6 +354,9 @@ impl KmsRpc for RpcHandler {
     }
 
     async fn get_kms_key(self, request: GetKmsKeyRequest) -> Result<KmsKeyResponse> {
+        self.ensure_self_allowed()
+            .await
+            .context("KMS self authorization failed")?;
         if self.state.config.onboard.quote_enabled {
             let _info = self.ensure_kms_allowed(&request.vm_config).await?;
         }
@@ -333,6 +370,9 @@ impl KmsRpc for RpcHandler {
     }
 
     async fn get_temp_ca_cert(self) -> Result<GetTempCaCertResponse> {
+        self.ensure_self_allowed()
+            .await
+            .context("KMS self authorization failed")?;
         Ok(GetTempCaCertResponse {
             temp_ca_cert: self.state.inner.temp_ca_cert.clone(),
             temp_ca_key: self.state.inner.temp_ca_key.clone(),
@@ -341,6 +381,9 @@ impl KmsRpc for RpcHandler {
     }
 
     async fn sign_cert(self, request: SignCertRequest) -> Result<SignCertResponse> {
+        self.ensure_self_allowed()
+            .await
+            .context("KMS self authorization failed")?;
         let csr = match request.api_version {
             1 => {
                 let csr = CertSigningRequestV1::decode(&mut &request.csr[..])

kms/src/main_service/upgrade_authority.rs

@@ -4,8 +4,13 @@
 
 use crate::config::AuthApi;
 use anyhow::{bail, Context, Result};
+use dstack_guest_agent_rpc::{
+    dstack_guest_client::DstackGuestClient, AttestResponse, RawQuoteArgs,
+};
+use http_client::prpc::PrpcClient;
 use ra_tls::attestation::AttestationMode;
 use ra_tls::attestation::VerifiedAttestation;
+use ra_tls::attestation::VersionedAttestation;
 use serde::de::DeserializeOwned;
 use serde::{Deserialize, Serialize};
 use serde_human_bytes as hex_bytes;
@@ -67,6 +72,20 @@ pub(crate) fn build_boot_info(
     })
 }
 
+pub(crate) async fn local_kms_boot_info(pccs_url: Option<&str>) -> Result<BootInfo> {
+    let response = app_attest(pad64([0u8; 32]))
+        .await
+        .context("Failed to get local KMS attestation")?;
+    let attestation = VersionedAttestation::from_scale(&response.attestation)
+        .context("Failed to decode local KMS attestation")?
+        .into_inner();
+    let verified = attestation
+        .verify(pccs_url)
+        .await
+        .context("Failed to verify local KMS attestation")?;
+    build_boot_info(&verified, false, "")
+}
+
 #[derive(Debug, Serialize, Deserialize)]
 #[serde(rename_all = "camelCase")]
 pub(crate) struct BootResponse {
@@ -168,3 +187,20 @@ fn url_join(url: &str, path: &str) -> String {
     url.push_str(path);
     url
 }
+
+fn dstack_client() -> DstackGuestClient<PrpcClient> {
+    let address = dstack_types::dstack_agent_address();
+    let http_client = PrpcClient::new(address);
+    DstackGuestClient::new(http_client)
+}
+
+async fn app_attest(report_data: Vec<u8>) -> Result<AttestResponse> {
+    dstack_client().attest(RawQuoteArgs { report_data }).await
+}
+
+fn pad64(hash: [u8; 32]) -> Vec<u8> {
+    let mut padded = Vec::with_capacity(64);
+    padded.extend_from_slice(&hash);
+    padded.resize(64, 0);
+    padded
+}

kms/src/onboard_service.rs

@@ -28,7 +28,10 @@ use ra_tls::{
 };
 use safe_write::safe_write;
 
-use crate::{config::KmsConfig, main_service::upgrade_authority::build_boot_info};
+use crate::{
+    config::KmsConfig,
+    main_service::upgrade_authority::{build_boot_info, local_kms_boot_info},
+};
 
 #[derive(Clone)]
 pub struct OnboardState {
@@ -400,17 +403,18 @@ async fn app_attest(report_data: Vec<u8>) -> Result<AttestResponse> {
 }
 
 async fn ensure_self_kms_allowed(cfg: &KmsConfig) -> Result<()> {
-    let response = app_attest(pad64([0u8; 32]))
+    let boot_info = local_kms_boot_info(cfg.pccs_url.as_deref())
         .await
-        .context("Failed to get local KMS attestation")?;
-    let attestation = VersionedAttestation::from_scale(&response.attestation)
-        .context("Failed to decode local KMS attestation")?
-        .into_inner();
-    let verified = attestation
-        .verify(cfg.pccs_url.as_deref())
+        .context("Failed to build local KMS boot info")?;
+    let response = cfg
+        .auth_api
+        .is_app_allowed(&boot_info, true)
         .await
-        .context("Failed to verify local KMS attestation")?;
-    ensure_kms_allowed(cfg, &verified).await
+        .context("Failed to call KMS auth check")?;
+    if !response.is_allowed {
+        bail!("Boot denied: {}", response.reason);
+    }
+    Ok(())
 }
 
 async fn ensure_remote_kms_allowed(