kms: register root key epochs on chain and report them from auth backends

.github/workflows/sdk.yaml

@@ -84,6 +84,12 @@ jobs:
           # the entries that are actually shared.
           save-if: ${{ github.event_name == 'push' }}
 
+      - name: Set up Python
+        if: steps.changes.outputs.relevant == 'true'
+        uses: actions/setup-python@v5
+        with:
+          python-version: '3.11'
+
       - name: SDK tests
         if: steps.changes.outputs.relevant == 'true'
         run: cd sdk && ./run-tests.sh

dstack/crates/dstack-auth/Cargo.toml

@@ -14,6 +14,8 @@ path = "src/main.rs"
 
 [dependencies]
 anyhow.workspace = true
+dstack-types.workspace = true
+hex.workspace = true
 clap.workspace = true
 rocket = { workspace = true, features = ["json"] }
 serde.workspace = true

dstack/crates/dstack-auth/src/main.rs

@@ -20,8 +20,9 @@
 //! (capture `tcbStatus`, deny unless `UpToDate`) if this grows into a
 //! multi-tenant / hosted deployment.
 
-use anyhow::Result;
+use anyhow::{Context, Result};
 use clap::Parser;
+use rocket::http::Status;
 use rocket::serde::json::Json;
 use rocket::{get, post, routes, State};
 use serde::{Deserialize, Serialize};
@@ -58,6 +59,10 @@ struct BootInfo {
     app_id: String,
     compose_hash: String,
     device_id: String,
+    /// the root epoch the asking KMS serves and its chain hash; absent from a
+    /// KMS that predates root epochs, and from KMS checks made while onboarding.
+    root_epoch: Option<usize>,
+    root_epoch_hash: Option<String>,
 }
 
 #[derive(Serialize)]
@@ -68,13 +73,61 @@ struct BootResponse {
     reason: String,
 }
 
+/// the current KMS root epoch and its chain hash, which commits to the root
+/// public keys of epochs 0..=number.
+struct EpochInfo {
+    number: usize,
+    hash: String,
+}
+
 #[derive(Deserialize, Default)]
 #[serde(rename_all = "camelCase", default)]
 struct Allowlist {
     os_images: Vec<String>,
     gateway_app_id: String,
     kms: KmsRules,
     apps: HashMap<String, AppRules>,
+    /// root public keys of each KMS root epoch, indexed by epoch; append only.
+    root_epochs: Vec<RootEpoch>,
+    /// the root epoch KMS instances serve; raise it once every node holds it.
+    current_root_epoch: usize,
+}
+
+#[derive(Deserialize, Default, Clone)]
+#[serde(rename_all = "camelCase", default)]
+struct RootEpoch {
+    ca_pubkey: String,
+    k256_pubkey: String,
+}
+
+impl Allowlist {
+    /// `None` only when no epoch is configured; invalid configuration is an error.
+    fn current_root_epoch(&self) -> Result<Option<EpochInfo>> {
+        if self.root_epochs.is_empty() && self.current_root_epoch == 0 {
+            return Ok(None);
+        }
+        let epochs = self
+            .root_epochs
+            .get(..=self.current_root_epoch)
+            .context("currentRootEpoch is not a configured root epoch")?;
+        let keys = epochs
+            .iter()
+            .map(|e| {
+                Ok((
+                    hex::decode(norm(&e.ca_pubkey)).context("invalid root epoch CA public key")?,
+                    hex::decode(norm(&e.k256_pubkey))
+                        .context("invalid root epoch k256 public key")?,
+                ))
+            })
+            .collect::<Result<Vec<_>>>()?;
+        let hashes = dstack_types::root_epoch::chain_hashes(
+            keys.iter().map(|(ca, k256)| (&ca[..], &k256[..])),
+        );
+        Ok(Some(EpochInfo {
+            number: self.current_root_epoch,
+            hash: format!("0x{}", hex::encode(hashes[self.current_root_epoch])),
+        }))
+    }
 }
 
 #[derive(Deserialize, Default)]
@@ -125,11 +178,39 @@ fn deny(al: &Allowlist, reason: &str) -> BootResponse {
     }
 }
 
-fn allow(al: &Allowlist) -> BootResponse {
-    BootResponse {
-        is_allowed: true,
-        gateway_app_id: al.gateway_app_id.clone(),
-        reason: "ok".to_string(),
+/// allow, provided the asking KMS serves the current root epoch with the
+/// registered keys up to it. A KMS that predates root epochs reports none and
+/// holds only epoch 0; KMS checks made while onboarding carry none.
+fn allow(info: &BootInfo, is_kms: bool, al: &Allowlist) -> BootResponse {
+    let current = match al.current_root_epoch() {
+        Ok(current) => current,
+        Err(error) => return deny(al, &format!("invalid root epoch configuration: {error}")),
+    };
+    let current_epoch = current.as_ref().map_or(0, |c| c.number);
+    let denial = match (info.root_epoch, &current) {
+        (None, _) if is_kms || current_epoch == 0 => None,
+        (None, _) => Some(format!(
+            "root epoch {current_epoch} is current, but the KMS reports none"
+        )),
+        (Some(0), None) => None,
+        (Some(epoch), None) => Some(format!("root epoch {epoch} is not registered")),
+        (Some(epoch), Some(_)) if epoch != current_epoch => Some(format!(
+            "the KMS serves root epoch {epoch}, but {current_epoch} is current"
+        )),
+        (Some(_), Some(c)) if info.root_epoch_hash.as_deref().map(norm) != Some(norm(&c.hash)) => {
+            Some(format!(
+                "the KMS root keys up to epoch {current_epoch} are not the registered ones"
+            ))
+        }
+        _ => None,
+    };
+    match denial {
+        Some(reason) => deny(al, &reason),
+        None => BootResponse {
+            is_allowed: true,
+            gateway_app_id: al.gateway_app_id.clone(),
+            reason: "ok".to_string(),
+        },
     }
 }
 
@@ -152,7 +233,7 @@ fn check_app(info: &BootInfo, al: &Allowlist) -> BootResponse {
     if !device_ok(app.allow_any_device, &app.devices, &info.device_id) {
         return deny(al, "device not allowed");
     }
-    allow(al)
+    allow(info, false, al)
 }
 
 fn check_kms(info: &BootInfo, al: &Allowlist) -> BootResponse {
@@ -162,7 +243,7 @@ fn check_kms(info: &BootInfo, al: &Allowlist) -> BootResponse {
     if !device_ok(al.kms.allow_any_device, &al.kms.devices, &info.device_id) {
         return deny(al, "device not allowed");
     }
-    allow(al)
+    allow(info, true, al)
 }
 
 /// load the allowlist, failing closed (deny-all) if it's missing or invalid.
@@ -206,15 +287,27 @@ fn boot_kms(info: Json<BootInfo>, cli: &State<Cli>) -> Json<BootResponse> {
 
 /// info endpoint the KMS GETs to populate its metadata. Single-node: no chain.
 #[get("/")]
-fn info() -> Json<serde_json::Value> {
-    Json(json!({
+fn info(cli: &State<Cli>) -> Result<Json<serde_json::Value>, Status> {
+    info_json(&load(&cli.config)).map(Json).map_err(|error| {
+        rocket::warn!("invalid root epoch configuration: {error}");
+        Status::ServiceUnavailable
+    })
+}
+
+fn info_json(al: &Allowlist) -> Result<serde_json::Value> {
+    let mut info = json!({
         "status": "ok",
         "kmsContractAddr": "",
         "ethRpcUrl": "",
         "gatewayAppId": "",
         "chainId": 0,
-        "appImplementation": ""
-    }))
+        "appImplementation": "",
+    });
+    if let Some(epoch) = al.current_root_epoch()? {
+        info["rootEpoch"] = epoch.number.into();
+        info["rootEpochHash"] = epoch.hash.into();
+    }
+    Ok(info)
 }
 
 #[rocket::main]
@@ -283,6 +376,89 @@ mod tests {
         assert!(!check_kms(&info, &Allowlist::default()).is_allowed);
     }
 
+    #[test]
+    fn allows_only_a_kms_serving_the_current_root_epoch() {
+        // Shared with the DstackKms contract tests.
+        let h0 = "0x10ab6b6c8e7601901b0f1767985bf40d35bf3129e0866fe12dcc7c8f04a298f2";
+        let h1 = "0x50152a82023f2219a31723b476c13bc5f608ab4147946dabdec15d4f3d648ed4";
+        let mut al: Allowlist = serde_json::from_str(
+            r#"{"kms":{"mrAggregated":["0xMR"]},
+                "rootEpochs":[{"caPubkey":"0x3059301306072a8648ce3d020106082a8648ce3d03010703420004a0a0a0a0a0a0a0a0a0a0a0a0a0a0a0a0a0a0a0a0a0a0a0a0a0a0a0a0a0a0a0a0a0a0a0a0a0a0a0a0a0a0a0a0a0a0a0a0a0a0a0a0a0a0a0a0a0a0a0a0a0a0a0a0","k256Pubkey":"0x02b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0"},
+                              {"caPubkey":"0x3059301306072a8648ce3d020106082a8648ce3d03010703420004a1a1a1a1a1a1a1a1a1a1a1a1a1a1a1a1a1a1a1a1a1a1a1a1a1a1a1a1a1a1a1a1a1a1a1a1a1a1a1a1a1a1a1a1a1a1a1a1a1a1a1a1a1a1a1a1a1a1a1a1a1a1a1a1","k256Pubkey":"0x02b1b1b1b1b1b1b1b1b1b1b1b1b1b1b1b1b1b1b1b1b1b1b1b1b1b1b1b1b1b1b1b1"}]}"#,
+        )
+        .unwrap();
+        let info = |epoch: Option<usize>, hash: &str| BootInfo {
+            mr_aggregated: "0xMR".into(),
+            root_epoch: epoch,
+            root_epoch_hash: Some(hash.into()),
+            ..Default::default()
+        };
+        assert_eq!(info_json(&al).unwrap()["rootEpochHash"], h0);
+        assert!(check_kms(&info(Some(0), h0), &al).is_allowed);
+
+        al.current_root_epoch = 1;
+        assert_eq!(info_json(&al).unwrap()["rootEpoch"], 1);
+        assert_eq!(info_json(&al).unwrap()["rootEpochHash"], h1);
+        assert!(check_kms(&info(Some(1), &h1.to_uppercase()), &al).is_allowed);
+        assert!(!check_kms(&info(Some(0), h0), &al).is_allowed);
+        assert!(!check_kms(&info(Some(1), h0), &al).is_allowed);
+        // onboarding checks carry no epoch
+        assert!(check_kms(&info(None, ""), &al).is_allowed);
+
+        al.current_root_epoch = 2;
+        assert!(al.current_root_epoch().is_err());
+        assert!(info_json(&al).is_err());
+        assert!(!check_kms(&info(Some(1), h1), &al).is_allowed);
+    }
+
+    #[test]
+    fn a_kms_predating_root_epochs_is_cut_off_once_one_is_current() {
+        let mut al = allowlist();
+        let info = boot("0xapp1", "0xhash", "0ximg");
+        assert!(check_app(&info, &al).is_allowed);
+        al.root_epochs = vec![
+            RootEpoch {
+                ca_pubkey: "0xa0".into(),
+                k256_pubkey: "0xb0".into(),
+            },
+            RootEpoch {
+                ca_pubkey: "0xa1".into(),
+                k256_pubkey: "0xb1".into(),
+            },
+        ];
+        assert!(check_app(&info, &al).is_allowed);
+        al.current_root_epoch = 1;
+        assert!(!check_app(&info, &al).is_allowed);
+    }
+
+    #[test]
+    fn invalid_root_epochs_deny_legacy_and_epoch_aware_kms() {
+        let invalid_configs = [
+            json!({"currentRootEpoch": 1}),
+            json!({"currentRootEpoch": 1, "rootEpochs": [{"caPubkey": "a0", "k256Pubkey": "b0"}]}),
+            json!({"rootEpochs": [{"caPubkey": "invalid", "k256Pubkey": "b0"}]}),
+            json!({"rootEpochs": [{"caPubkey": "a0", "k256Pubkey": "invalid"}]}),
+        ];
+        for config in invalid_configs {
+            let epochs: Allowlist = serde_json::from_value(config).unwrap();
+            let al = Allowlist {
+                root_epochs: epochs.root_epochs,
+                current_root_epoch: epochs.current_root_epoch,
+                ..allowlist()
+            };
+            assert!(info_json(&al).is_err());
+            for epoch in [None, Some(0), Some(1)] {
+                let info = BootInfo {
+                    mr_aggregated: "0xMR".into(),
+                    root_epoch: epoch,
+                    ..boot("0xapp1", "0xhash", "0ximg")
+                };
+                assert!(!check_app(&info, &al).is_allowed);
+                assert!(!check_kms(&info, &al).is_allowed);
+            }
+        }
+    }
+
     // wire-contract snapshot: BootInfo as the KMS serializes it (camelCase).
     // Keep these field names in sync with the kms BootInfo. `#[serde(default)]`
     // means extra fields are ignored AND a renamed field deserializes to "" —

dstack/dstack-types/src/lib.rs

@@ -2541,6 +2541,7 @@ pub struct ImageInfo {
 }
 
 pub mod mr_config;
+pub mod root_epoch;
 pub mod sha256sum;
 pub mod shared_filenames;
 pub mod version;

dstack/dstack-types/src/root_epoch.rs

@@ -0,0 +1,57 @@
+// SPDX-FileCopyrightText: © 2026 Phala Network <dstack@phala.network>
+//
+// SPDX-License-Identifier: Apache-2.0
+
+//! KMS root epoch chain hashes, computed as the DstackKms contract does.
+
+use sha3::{Digest, Keccak256};
+
+/// Chain hashes of root epochs 0, 1, ... given each epoch's
+/// `(ca_pubkey, k256_pubkey)`. Epoch `n`'s hash is
+/// `keccak256(previous || keccak256(ca_pubkey) || keccak256(k256_pubkey))`
+/// with a zero previous hash for epoch 0, so it commits to every key up to `n`;
+/// the position in the chain fixes the epoch number.
+pub fn chain_hashes<'a>(epochs: impl IntoIterator<Item = (&'a [u8], &'a [u8])>) -> Vec<[u8; 32]> {
+    let mut previous = [0u8; 32];
+    epochs
+        .into_iter()
+        .map(|(ca_pubkey, k256_pubkey)| {
+            previous = Keccak256::new()
+                .chain_update(previous)
+                .chain_update(Keccak256::digest(ca_pubkey))
+                .chain_update(Keccak256::digest(k256_pubkey))
+                .finalize()
+                .into();
+            previous
+        })
+        .collect()
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    /// Well-formed root public keys whose point bytes are all `fill`.
+    fn test_keys(ca_fill: u8, k256_fill: u8) -> (Vec<u8>, Vec<u8>) {
+        let mut ca = hex::decode("3059301306072a8648ce3d020106082a8648ce3d03010703420004").unwrap();
+        ca.resize(91, ca_fill);
+        let mut k256 = vec![0x02];
+        k256.resize(33, k256_fill);
+        (ca, k256)
+    }
+
+    /// Shared with the DstackKms contract and auth-simple tests.
+    #[test]
+    fn matches_the_contract_vectors() {
+        let epochs = [test_keys(0xa0, 0xb0), test_keys(0xa1, 0xb1)];
+        let hashes = chain_hashes(epochs.iter().map(|(ca, k256)| (&ca[..], &k256[..])));
+        assert_eq!(
+            hex::encode(hashes[0]),
+            "10ab6b6c8e7601901b0f1767985bf40d35bf3129e0866fe12dcc7c8f04a298f2"
+        );
+        assert_eq!(
+            hex::encode(hashes[1]),
+            "50152a82023f2219a31723b476c13bc5f608ab4147946dabdec15d4f3d648ed4"
+        );
+    }
+}

dstack/kms/auth-eth-bun/README.md

@@ -17,6 +17,14 @@ a single-file implementation of the dstack KMS ethereum backend using bun + hono
 bun install
 ```
 
+## Upgrade order
+
+Upgrade the `DstackKms` contract before deploying the new auth backend or KMS.
+An empty root epoch registry is supported on the upgraded contract; there is no
+fallback for legacy contracts. Old KMS instances that omit the root epoch can
+receive app authorization only while no epoch is registered or epoch 0 is current,
+subject to the existing authorization policy.
+
 ## usage
 
 ### development