Merge pull request #651 from Dstack-TEE/feat/kms-optional-self-auth

feat(kms): make self-authorization enforcement configurable

Author: kvinwang

kms/kms.toml

@@ -26,6 +26,11 @@ mandatory = false
 cert_dir = "/etc/kms/certs"
 subject_postfix = ".dstack"
 admin_token_hash = ""
+# Whether trusted RPCs require the KMS to first attest itself to its own
+# auth API. Defaults to true (strict). Set to false ONLY when running KMS
+# outside a TEE (e.g. local dev/testing) where the local guest agent socket
+# is unavailable.
+enforce_self_authorization = true
 
 [core.image]
 verify = true

kms/src/config.rs

@@ -40,6 +40,16 @@ pub(crate) struct KmsConfig {
     pub image: ImageConfig,
     #[serde(with = "serde_human_bytes")]
     pub admin_token_hash: Vec<u8>,
+    /// Whether trusted RPCs require the KMS to first attest itself to its
+    /// own auth API. Defaults to `true` (strict). Set `false` only for local
+    /// dev/testing where the KMS runs outside a TEE and cannot reach a guest
+    /// agent socket.
+    #[serde(default = "default_true")]
+    pub enforce_self_authorization: bool,
+}
+
+fn default_true() -> bool {
+    true
 }
 
 impl KmsConfig {

kms/src/main_service.rs

@@ -26,7 +26,7 @@ use ra_tls::{
 use scale::Decode;
 use sha2::Digest;
 use tokio::sync::OnceCell;
-use tracing::info;
+use tracing::{info, warn};
 use upgrade_authority::{build_boot_info, local_kms_boot_info, BootInfo};
 
 use crate::{
@@ -76,6 +76,9 @@ impl KmsState {
             config.image.download_timeout,
             config.pccs_url.clone(),
         );
+        if !config.enforce_self_authorization {
+            warn!("self-authorization is disabled; trusted RPCs will not be gated by KMS self-attestation - do not use in production TEE deployments");
+        }
         Ok(Self {
             inner: Arc::new(KmsStateInner {
                 config,
@@ -102,6 +105,9 @@ struct BootConfig {
 
 impl RpcHandler {
     async fn ensure_self_allowed(&self) -> Result<()> {
+        if !self.state.config.enforce_self_authorization {
+            return Ok(());
+        }
         let boot_info = self
             .state
             .self_boot_info

kms/src/main_service/upgrade_authority.rs

@@ -206,6 +206,9 @@ pub(crate) fn pad64(hash: [u8; 32]) -> Vec<u8> {
 }
 
 pub(crate) async fn ensure_self_kms_allowed(cfg: &KmsConfig) -> Result<()> {
+    if !cfg.enforce_self_authorization {
+        return Ok(());
+    }
     let boot_info = local_kms_boot_info(cfg.pccs_url.as_deref())
         .await
         .context("failed to build local KMS boot info")?;