From c5fd176eab8633c0bad79a0c0af952797f026b4c Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Sat, 7 Nov 2026 01:30:58 -0800 Subject: [PATCH 01/11] feat(kms): rotate application keys independently of root epochs Signed-off-by: Kevin Wang --- docs/app-key-epochs.md | 70 +++++++ docs/sealed-env.md | 13 +- dstack/cert-client/src/lib.rs | 16 +- dstack/dstack-types/src/lib.rs | 12 +- dstack/dstack-util/src/main.rs | 2 + dstack/dstack-util/src/system_setup.rs | 172 +++++++++++++----- dstack/dstack-util/src/system_setup/luks.rs | 41 ++++- .../src/system_setup/sealed_env.rs | 27 ++- .../guest-agent/rpc/proto/agent_rpc_v1.proto | 2 + dstack/guest-agent/src/rpc_service.rs | 76 ++++++-- dstack/guest-agent/src/rpc_service_v1.rs | 31 +++- dstack/kms/README.md | 12 +- dstack/kms/auth-eth-bun/README.md | 7 +- dstack/kms/auth-eth-bun/index.test.ts | 2 +- dstack/kms/auth-eth-bun/index.ts | 22 ++- dstack/kms/auth-eth-bun/openapi.json | 41 ++++- dstack/kms/auth-eth/README.md | 7 +- dstack/kms/auth-eth/contracts/DstackApp.sol | 20 +- dstack/kms/auth-eth/contracts/DstackKms.sol | 17 +- dstack/kms/auth-eth/contracts/IAppEpoch.sol | 11 ++ dstack/kms/auth-eth/script/Manage.s.sol | 11 ++ dstack/kms/auth-eth/script/Query.s.sol | 8 + dstack/kms/auth-eth/src/ethereum.ts | 5 + dstack/kms/auth-eth/src/server.ts | 6 +- dstack/kms/auth-eth/test/DstackKms.t.sol | 58 +++++- dstack/kms/auth-mock/README.md | 7 +- dstack/kms/auth-mock/index.ts | 25 ++- dstack/kms/auth-mock/openapi.json | 41 ++++- dstack/kms/auth-simple/README.md | 7 +- dstack/kms/auth-simple/index.test.ts | 14 +- dstack/kms/auth-simple/index.ts | 26 ++- dstack/kms/rpc/proto/kms_rpc.proto | 2 + dstack/kms/src/crypto.rs | 35 +++- dstack/kms/src/main_service.rs | 156 ++++++++++++---- .../kms/src/main_service/upgrade_authority.rs | 39 +++- sdk/go/dstack/client_v1.go | 20 +- sdk/js/src/client-v1.ts | 14 +- sdk/python/src/dstack_sdk/dstack_client_v1.py | 16 +- sdk/rust/src/dstack_client_v1.rs | 29 ++- sdk/rust/types/src/dstack_v1.rs | 5 +- 40 files changed, 932 insertions(+), 193 deletions(-) create mode 100644 docs/app-key-epochs.md create mode 100644 dstack/kms/auth-eth/contracts/IAppEpoch.sol diff --git a/docs/app-key-epochs.md b/docs/app-key-epochs.md new file mode 100644 index 000000000..4e13f0652 --- /dev/null +++ b/docs/app-key-epochs.md @@ -0,0 +1,70 @@ +# Application key epochs + +Application owners can rotate their app's derived keys without rotating KMS root +keys or other apps. This extends SealedEnv (#1528); the PR dependency order is +#1519 → #1517 → #1520 → #1521 → #1528 → application epochs. + +## Authorization and rotation + +`checkAppAuthorization` returns `currentAppEpoch` together with the current root +epoch and its chain hash. `requestedRootEpoch` and `requestedAppEpoch` are +independent `OptionalEpoch` values: omitted selects current, explicit zero selects +historical epoch zero. Future epochs are denied through `rejectionReason`. +`checkKmsAuthorization` remains separate and uses app epoch zero. + +`DstackApp.currentAppEpoch` starts at zero. Its owner calls +`advanceAppEpoch(expectedCurrentEpoch)` to increment it. The expected value +prevents a stale or retried transaction from advancing twice. Epochs cannot be +rolled back and are limited to uint32. `AdvanceAppEpoch` in `Manage.s.sol` exposes +this operation; `GetAppEpoch` in `Query.s.sol` reads it. Upgrade the app policy +implementation before advancing, and upgrade the KMS/auth backend and guests +before using nonzero epochs. Legacy policies without the ERC-165 `IAppEpoch` +capability remain at zero; failures of an advertised getter propagate. + +The Ethereum auth backends query the policy through the KMS contract. The simple +backend reads `apps[appId].currentAppEpoch` (default zero); operators must preserve +its monotonicity. Public encryption-key lookup uses `GET /app/:appId` to read this +value. Legacy boot authorization and legacy environment public-key issuance are +rejected once the app epoch is nonzero, preventing old clients from silently +continuing with epoch-zero keys. + +## Derivation and lifetime + +Every key is identified by `(root_epoch, app_epoch)`. App epoch zero keeps existing derived keys and app-key signatures. Nonzero app epochs append +`"app-epoch-v1:" || app_epoch[4, big-endian]` to the existing derivation context +for disk, environment, SealedEnv, app CA and secp256k1 keys. + +The root signature for a nonzero app secp256k1 key signs Keccak256 of +`"dstack-kms-issued-v2:" || app_id[20] || app_epoch[4, big-endian] || public_key[33]`. +Relying parties must verify this format against the trusted root key for the +returned root epoch. Epoch zero keeps the original `dstack-kms-issued` format. + +Restart CVMs after advancing the policy. A running CVM continues to use its boot +key pair; omitted guest-agent GetKey selectors mean that CVM's boot epochs, not a +fresh chain lookup. Explicit historical requests select each dimension +independently and cache by the complete pair. Historical access still requires +current app authorization. Supplied `trusted_ca` always pins the KMS connection, +regardless of epoch selection; historical responses cannot replace that pin. + +The guest migrates disk headers and environment ciphertext if either dimension +changes and rejects future epochs in either dimension. LUKS tokens record both +epochs (missing app epoch means zero). Rekeying replaces the keyslot wrapping key; +it does not rotate the volume encryption key or erase host snapshots. App +rotation isolates future app keys from an exposed old derived key, but cannot +recover from compromise of the KMS root or revoke already disclosed secrets. + +## SDK selection + +Both selectors are optional; specify zero explicitly to recover an old key. +For example, Python `client.get_key("storage", "secp256k1", root_epoch=3, +app_epoch=0)` selects exactly that pair. JavaScript uses +`client.getKey("storage", "secp256k1", 3, 0)`. Rust provides +`get_key_at_epochs(..., Some(3), Some(0))`; Go provides `GetKeyAtEpochs` with +pointer selectors. Responses report both epochs. Existing root-only helpers +continue to use the CVM's current app epoch. + +Sealing helpers also accept app epochs: `seal_env_at_epochs` (Rust), +`SealEnvAtEpochs` (Go), and the optional final argument of `sealEnv` / `seal_env` +(JavaScript / Python). See [SealedEnv](sealed-env.md) for the wire and signature +formats. The VMM UI and CLI use the epochs returned with the encryption key. SealedEnv +uses a single version-1 format that always includes both epochs. diff --git a/docs/sealed-env.md b/docs/sealed-env.md index de45c86d7..a343901ed 100644 --- a/docs/sealed-env.md +++ b/docs/sealed-env.md @@ -44,7 +44,8 @@ that epoch's P-256 CA private scalar using HKDF-SHA256: - info: the raw 20-byte app ID. `GetAppKey` returns it as `sealed_env_key` alongside the legacy key. The legacy -key derivation remains unchanged. This PR reserves app epochs and accepts only zero. +key derivation remains unchanged at app epoch zero. `GetAppEnvEncryptPubKey` +rejects nonzero app epochs; use `GetSealedEnvPubKey` after app rotation. The public-key response signs Keccak256 of this unambiguous byte sequence: @@ -77,10 +78,12 @@ with nonempty legacy `encrypted_env` while switching to an older image. Empty en fields preserve the existing format. Guest `UpdateSealedEnv` uploads also require a capable running image. -SealedEnv has one format, version 1, with both epoch fields always present. -This PR accepts app epoch zero only; the field is reserved for application key -rotation without changing this format. Signature and HPKE contexts always include -both epochs, including zero. Changing either epoch invalidates authentication. +SealedEnv has one format, version 1, with both epoch fields always present. The +receiver-key KDF appends `"app-epoch-v1:" || app_epoch[4, big-endian]` to its info +for nonzero app epochs; at zero its info is just the app ID. Signature and HPKE +contexts always include both epochs, including zero. Changing either epoch +invalidates the signature or ciphertext. See [application epochs](app-key-epochs.md) +for owner operations. The browser UI receives the key through the same-origin VMM API and does not independently verify the response signer or freshness. Its trust boundary includes diff --git a/dstack/cert-client/src/lib.rs b/dstack/cert-client/src/lib.rs index 9e222bb04..c95aa3caf 100644 --- a/dstack/cert-client/src/lib.rs +++ b/dstack/cert-client/src/lib.rs @@ -86,10 +86,14 @@ pub async fn request_app_keys( .build() .into_client()?; let requested_root_epoch = request.root_epoch; + let requested_app_epoch = request.app_epoch; let response = KmsClient::new(client).get_app_key(request).await?; if requested_root_epoch.is_some_and(|epoch| epoch != response.root_epoch) { bail!("KMS returned a different root epoch than requested"); } + if requested_app_epoch.is_some_and(|epoch| epoch != response.app_epoch) { + bail!("KMS returned a different app epoch than requested"); + } let cert = server_cert .lock() .map_err(|_| anyhow::anyhow!("KMS certificate lock poisoned"))?; @@ -170,21 +174,23 @@ impl CertRequestClient { /// The app keys of an older KMS root epoch, from the KMS that served the /// current ones. - pub async fn get_app_key(&self, root_epoch: u32) -> Result { + pub async fn get_app_key(&self, root_epoch: u32, app_epoch: u32) -> Result { let CertRequestClient::Kms { client, vm_config } = self else { bail!("only KMS-provided keys have older root epochs"); }; let response = client .get_app_key(GetAppKeyRequest { + app_epoch: Some(app_epoch), api_version: 1, vm_config: vm_config.clone(), root_epoch: Some(root_epoch), }) .await?; - if response.root_epoch != root_epoch { + if response.root_epoch != root_epoch || response.app_epoch != app_epoch { bail!( - "KMS returned keys of root epoch {} instead of {root_epoch}", - response.root_epoch + "KMS returned key epochs ({}, {}) instead of ({root_epoch}, {app_epoch})", + response.root_epoch, + response.app_epoch ); } Ok(response) @@ -594,6 +600,7 @@ mod tests { } }; let request = || GetAppKeyRequest { + app_epoch: None, api_version: 1, vm_config: String::new(), root_epoch: None, @@ -613,6 +620,7 @@ mod tests { let other_ca = kms_root().pem_cert; for requested_root_epoch in [None, Some(0), Some(1)] { let request = || GetAppKeyRequest { + app_epoch: None, api_version: 1, vm_config: String::new(), root_epoch: requested_root_epoch, diff --git a/dstack/dstack-types/src/lib.rs b/dstack/dstack-types/src/lib.rs index c7419103a..f1496c502 100644 --- a/dstack/dstack-types/src/lib.rs +++ b/dstack/dstack-types/src/lib.rs @@ -2412,6 +2412,13 @@ pub struct AppKeys { } impl AppKeys { + pub fn app_epoch(&self) -> u32 { + match self.key_provider { + KeyProvider::Kms { app_epoch, .. } => app_epoch, + _ => 0, + } + } + /// KMS root epoch the keys were derived from; 0 for other key providers, /// which have a single root. pub fn root_epoch(&self) -> u32 { @@ -2447,6 +2454,8 @@ pub enum KeyProvider { /// KMS root key epoch the keys were derived from. #[serde(default)] root_epoch: u32, + #[serde(default)] + app_epoch: u32, }, } @@ -2509,7 +2518,7 @@ mod key_provider_tests { assert_eq!( serde_json::to_value(provider).unwrap(), serde_json::json!({ - "Kms": { "url": "https://kms.example/prpc", "pubkey": "abcd", "root_epoch": 0 } + "Kms": { "url": "https://kms.example/prpc", "pubkey": "abcd", "root_epoch": 0, "app_epoch": 0 } }) ); } @@ -2526,6 +2535,7 @@ mod key_provider_tests { assert_eq!(tpm.kind(), KeyProviderKind::Tpm); let kms = KeyProvider::Kms { + app_epoch: 0, url: "https://kms.example".into(), pubkey: vec![0xab; 32], root_epoch: 0, diff --git a/dstack/dstack-util/src/main.rs b/dstack/dstack-util/src/main.rs index c8bc9e0df..e6c49f271 100644 --- a/dstack/dstack-util/src/main.rs +++ b/dstack/dstack-util/src/main.rs @@ -655,6 +655,7 @@ async fn cmd_get_keys(args: GetKeysArgs) -> Result<()> { let response = cert_client::request_app_keys( kms_url.clone(), dstack_kms_rpc::GetAppKeyRequest { + app_epoch: None, api_version: 1, vm_config: String::new(), root_epoch: None, @@ -682,6 +683,7 @@ async fn cmd_get_keys(args: GetKeysArgs) -> Result<()> { k256_signature: response.k256_signature, gateway_app_id: response.gateway_app_id, key_provider: KeyProvider::Kms { + app_epoch: response.app_epoch, url: kms_url, pubkey: root_pubkey, root_epoch: response.root_epoch, diff --git a/dstack/dstack-util/src/system_setup.rs b/dstack/dstack-util/src/system_setup.rs index 172e2b70b..8a961c9eb 100644 --- a/dstack/dstack-util/src/system_setup.rs +++ b/dstack/dstack-util/src/system_setup.rs @@ -28,8 +28,8 @@ use dstack_types::{ }; use fs_err as fs; use luks::{ - rekey_luks_header, validate_luks2_headers, write_luks_root_epoch, LUKS_INITIALIZING_LABEL, - PAYLOAD_OFFSET, + rekey_luks_header_at_epochs, validate_luks2_headers, write_luks_epochs, + LUKS_INITIALIZING_LABEL, PAYLOAD_OFFSET, }; use ra_rpc::{ client::{CertInfo, RaClient, RaClientConfig}, @@ -98,7 +98,9 @@ mod config_id_verifier; mod luks; mod sealed_env; -use sealed_env::{env_crypt_key, seal_env, sealed_env_info, SealedEnv}; +#[cfg(test)] +use sealed_env::seal_env; +use sealed_env::{env_crypt_key, seal_env_at_epochs, sealed_env_info_at_epochs, SealedEnv}; #[derive(clap::Parser)] /// Prepare full disk encryption @@ -2404,6 +2406,7 @@ impl HostShared { &self, kms_url: String, root_epoch: Option, + app_epoch: Option, trusted_ca: Option, ) -> Result<(AppKeys, Vec)> { info!("Requesting app keys of root epoch {root_epoch:?} from KMS: {kms_url}"); @@ -2412,6 +2415,7 @@ impl HostShared { let response = request_app_keys( kms_url.clone(), rpc::GetAppKeyRequest { + app_epoch, api_version: 1, vm_config: self.sys_config.vm_config.clone(), root_epoch, @@ -2439,6 +2443,7 @@ impl HostShared { k256_signature: response.k256_signature, gateway_app_id: response.gateway_app_id, key_provider: KeyProvider::Kms { + app_epoch: response.app_epoch, url: kms_url, pubkey: root_pubkey, root_epoch: served, @@ -2449,12 +2454,22 @@ impl HostShared { /// Request the app keys of an older root epoch from the KMS that served /// `keys`. - async fn request_historical_keys(&self, keys: &AppKeys, epoch: u32) -> Result { + async fn request_historical_keys( + &self, + keys: &AppKeys, + epoch: u32, + app_epoch: u32, + ) -> Result { let KeyProvider::Kms { url, .. } = &keys.key_provider else { bail!("only KMS-provided keys have older root epochs"); }; Ok(self - .request_app_keys_from_kms(url.clone(), Some(epoch), Some(keys.ca_cert.clone())) + .request_app_keys_from_kms( + url.clone(), + Some(epoch), + Some(app_epoch), + Some(keys.ca_cert.clone()), + ) .await? .0) } @@ -2531,7 +2546,7 @@ impl<'a> Stage0<'a> { let kms_url = kms_rpc_url(kms_url); match self .shared - .request_app_keys_from_kms(kms_url.clone(), None, None) + .request_app_keys_from_kms(kms_url.clone(), None, None, None) .await { Ok((keys, os_image_hash)) => { @@ -2985,8 +3000,9 @@ impl<'a> Stage0<'a> { { bail!("Failed to setup luks volume"); } - if let Some(epoch) = kms_root_epoch(keys).filter(|epoch| *epoch > 0) { - write_luks_root_epoch(root_hd, epoch)?; + if let Some(epoch) = kms_root_epoch(keys).filter(|epoch| *epoch > 0 || keys.app_epoch() > 0) + { + write_luks_epochs(root_hd, epoch, keys.app_epoch())?; } self.open_encrypted_volume(keys, name, discard).await?; Ok(()) @@ -3003,16 +3019,20 @@ impl<'a> Stage0<'a> { &self, keys: &AppKeys, header: &Path, - from: u32, - to: u32, + from: (u32, u32), + to: (u32, u32), ) -> Result<()> { - info!("re-keying the data disk from root epoch {from} to {to}"); - let old = self.shared.request_historical_keys(keys, from).await?; - rekey_luks_header( + info!("re-keying the data disk from epochs {from:?} to {to:?}"); + let old = self + .shared + .request_historical_keys(keys, from.0, from.1) + .await?; + rekey_luks_header_at_epochs( header, &hex::encode(&old.disk_crypt_key), &hex::encode(&keys.disk_crypt_key), - to, + to.0, + to.1, )?; let root_hd = &self.args.device; cmd!(cryptsetup luksHeaderRestore --batch-mode --header-backup-file=$header $root_hd) @@ -3056,16 +3076,25 @@ impl<'a> Stage0<'a> { let mut header = validate()?; if let Some(current) = kms_root_epoch(keys) { let stored = header.root_epoch.unwrap_or(0); - if stored > current { - bail!("the data disk is at root epoch {stored}, newer than the current {current}"); + let current_app = keys.app_epoch(); + if stored > current || header.app_epoch > current_app { + bail!( + "the data disk epochs ({stored}, {}) exceed current ({current}, {current_app})", + header.app_epoch + ); } - if stored < current { - self.migrate_disk_key(keys, Path::new(&in_mem_hdr), stored, current) - .await - .context("failed to migrate the data disk key")?; + if stored != current || header.app_epoch != current_app { + self.migrate_disk_key( + keys, + Path::new(&in_mem_hdr), + (stored, header.app_epoch), + (current, current_app), + ) + .await + .context("failed to migrate the data disk key")?; header = validate()?; - if header.root_epoch != Some(current) { - bail!("the re-keyed LUKS header does not record root epoch {current}"); + if header.root_epoch != Some(current) || header.app_epoch != current_app { + bail!("the re-keyed LUKS header does not record epochs ({current}, {current_app})"); } } } @@ -3329,7 +3358,7 @@ impl Stage1<'_> { if let Some(sealed_env) = &self.shared.sealed_env { let envelope: SealedEnv = serde_json::from_slice(sealed_env).context("failed to parse sealed env")?; - if envelope.version != 1 || envelope.app_epoch != 0 { + if envelope.version != 1 { bail!("unsupported sealed env version: {}", envelope.version); } if envelope.enc.len() != 32 { @@ -3337,31 +3366,40 @@ impl Stage1<'_> { } let payload = [envelope.enc, envelope.ciphertext].concat(); let current = self.keys.root_epoch(); - if envelope.root_epoch > current { + if envelope.root_epoch > current || envelope.app_epoch > self.keys.app_epoch() { bail!( "sealed env epoch {} exceeds current epoch {current}", envelope.root_epoch ); } let historical; - let secret = if envelope.root_epoch == current { + let secret = if envelope.root_epoch == current + && envelope.app_epoch == self.keys.app_epoch() + { current_secret } else { historical = self .shared - .request_historical_keys(&self.keys, envelope.root_epoch) + .request_historical_keys(&self.keys, envelope.root_epoch, envelope.app_epoch) .await?; historical.sealed_env_key.as_slice() }; let plaintext = crate::crypto::sealed_env_decrypt( secret, - &sealed_env_info(&envelope.root_epoch), + &sealed_env_info_at_epochs(envelope.root_epoch, envelope.app_epoch), &payload, )?; - if envelope.root_epoch == self.keys.root_epoch() { + if envelope.root_epoch == self.keys.root_epoch() + && envelope.app_epoch == self.keys.app_epoch() + { return Ok((plaintext, Vec::new())); } - let sealed = seal_env(current_secret, self.keys.root_epoch(), &plaintext)?; + let sealed = seal_env_at_epochs( + current_secret, + self.keys.root_epoch(), + self.keys.app_epoch(), + &plaintext, + )?; return Ok((plaintext, sealed)); } let ciphertext = &self.shared.encrypted_env; @@ -3369,21 +3407,27 @@ impl Stage1<'_> { let plaintext = match dh_decrypt(current_key, ciphertext) { Ok(plaintext) => return Ok((plaintext, Vec::new())), Err(_) => { - decrypt_historical_env( - ciphertext, - (0..self.keys.root_epoch()).rev(), - async |epoch| { - Ok(self - .shared - .request_historical_keys(&self.keys, epoch) - .await? - .env_crypt_key) - }, - ) + let current = self.keys.root_epoch(); + // Legacy encryption is only issued at app epoch zero. + let epochs = (0..=current) + .rev() + .filter(|root| *root != current || self.keys.app_epoch() != 0); + decrypt_historical_env(ciphertext, epochs, async |root| { + Ok(self + .shared + .request_historical_keys(&self.keys, root, 0) + .await? + .env_crypt_key) + }) .await? } }; - let sealed = seal_env(current_secret, self.keys.root_epoch(), &plaintext)?; + let sealed = seal_env_at_epochs( + current_secret, + self.keys.root_epoch(), + self.keys.app_epoch(), + &plaintext, + )?; Ok((plaintext, sealed)) } @@ -4378,6 +4422,7 @@ mod sealed_env_file_tests { gateway_app_id: String::new(), ca_cert: String::new(), key_provider: KeyProvider::Kms { + app_epoch: 0, url: "http://127.0.0.1:1".into(), pubkey: vec![], root_epoch: 3, @@ -4463,6 +4508,51 @@ mod sealed_env_file_tests { ); } + #[tokio::test] + async fn sealed_env_binds_both_epochs_and_rejects_mislabeled_versions() { + let dir = tempfile::tempdir().unwrap(); + let args = SetupArgs { + work_dir: dir.path().into(), + device: "unused".into(), + mount_point: "unused".into(), + }; + let plaintext = br#"{"env":[{"key":"VALUE","value":"epoch-test"}]}"#; + let sealed = seal_env_at_epochs(&[7; 32], 3, 2, plaintext).unwrap(); + let mut current = stage(&args, Some(sealed.clone()), vec![]); + if let KeyProvider::Kms { app_epoch, .. } = &mut current.keys.key_provider { + *app_epoch = 2; + } + let (opened, update) = current.decrypt_env().await.unwrap(); + assert_eq!(opened, plaintext); + assert!(update.is_empty()); + for (field, value) in [ + ("version", 2), + ("root_epoch", 2), + ("app_epoch", 1), + ("app_epoch", 3), + ] { + let mut tampered: Value = serde_json::from_slice(&sealed).unwrap(); + tampered[field] = value.into(); + current.shared.sealed_env = Some(serde_json::to_vec(&tampered).unwrap()); + assert!(current.decrypt_env().await.is_err()); + } + let envelope: SealedEnv = serde_json::from_slice(&sealed).unwrap(); + assert_eq!(envelope.version, 1); + let payload = [envelope.enc.clone(), envelope.ciphertext.clone()].concat(); + assert!(crate::crypto::sealed_env_decrypt( + &[7; 32], + &sealed_env_info_at_epochs(3, 1), + &payload + ) + .is_err()); + assert!(crate::crypto::sealed_env_decrypt( + &[7; 32], + &sealed_env_info_at_epochs(2, 2), + &payload + ) + .is_err()); + } + #[test] fn sealed_file_selection_ignores_legacy_but_rejects_links() { let dir = tempfile::tempdir().unwrap(); diff --git a/dstack/dstack-util/src/system_setup/luks.rs b/dstack/dstack-util/src/system_setup/luks.rs index 5b5e2fef1..3cfc1027b 100644 --- a/dstack/dstack-util/src/system_setup/luks.rs +++ b/dstack/dstack-util/src/system_setup/luks.rs @@ -73,6 +73,8 @@ struct LuksToken { pub keyslots: Vec, #[serde(default)] pub root_epoch: Option, + #[serde(default)] + pub app_epoch: u32, } #[derive(Deserialize)] @@ -205,11 +207,12 @@ fn run_with_input(mut command: Command, input: &[u8]) -> Result> { .code()) } -pub(super) fn write_luks_root_epoch(device: &Path, epoch: u32) -> Result<()> { +pub(super) fn write_luks_epochs(device: &Path, epoch: u32, app_epoch: u32) -> Result<()> { let token = serde_json::json!({ "type": ROOT_EPOCH_TOKEN_TYPE, "keyslots": ["0"], "root_epoch": epoch, + "app_epoch": app_epoch, }); let args = ["token", "import", "--token-id", "0", "--token-replace"]; let code = run_with_input(cryptsetup(&args, device), token.to_string().as_bytes())?; @@ -222,11 +225,12 @@ pub(super) fn write_luks_root_epoch(device: &Path, epoch: u32) -> Result<()> { /// Replace the keyslot 0 passphrase of a detached LUKS `header` and record /// `epoch` in it. The new key goes over stdin; the old one is written next to /// the header, which must be in a private in-memory directory. -pub(super) fn rekey_luks_header( +pub(super) fn rekey_luks_header_at_epochs( header: &Path, old_key: &str, new_key: &str, epoch: u32, + app_epoch: u32, ) -> Result<()> { use std::os::unix::fs::PermissionsExt; @@ -252,7 +256,17 @@ pub(super) fn rekey_luks_header( if run_with_input(command, new_key.as_bytes())? != Some(0) { bail!("cryptsetup luksChangeKey failed"); } - write_luks_root_epoch(header, epoch) + write_luks_epochs(header, epoch, app_epoch) +} + +#[cfg(test)] +fn write_luks_root_epoch(device: &Path, epoch: u32) -> Result<()> { + write_luks_epochs(device, epoch, 0) +} + +#[cfg(test)] +fn rekey_luks_header(header: &Path, old_key: &str, new_key: &str, epoch: u32) -> Result<()> { + rekey_luks_header_at_epochs(header, old_key, new_key, epoch, 0) } /// Set by luksFormat and cleared once the filesystem exists, so a disk still @@ -265,12 +279,13 @@ pub(super) struct LuksHeaderState { pub(super) initializing: bool, /// Root epoch recorded in the dstack token; `None` means epoch 0. pub(super) root_epoch: Option, + pub(super) app_epoch: u32, } pub(super) fn validate_luks2_headers(mut reader: impl std::io::Read) -> Result { let primary = validate_single_luks2_header(&mut reader, 0)?; let secondary = validate_single_luks2_header(&mut reader, 1)?; - if primary.root_epoch != secondary.root_epoch { + if primary.root_epoch != secondary.root_epoch || primary.app_epoch != secondary.app_epoch { bail!("LUKS header copies disagree on the key epoch"); } Ok(primary) @@ -479,6 +494,7 @@ fn validate_single_luks2_header( Ok(LuksHeaderState { initializing, root_epoch, + app_epoch: tokens.get(&0).map(|token| token.app_epoch).unwrap_or(0), }) } @@ -491,7 +507,8 @@ fn test_validate_luks2_header() { state, LuksHeaderState { initializing: false, - root_epoch: None + root_epoch: None, + app_epoch: 0 } ); header_data[24..24 + LUKS_INITIALIZING_LABEL.len()] @@ -625,6 +642,20 @@ fn test_disk_key_migration_with_cryptsetup() { assert!(backup.success()); let state = validate_luks2_headers(fs::File::open(&header).unwrap()).unwrap(); assert_eq!(state.root_epoch, Some(3)); + rekey_luks_header_at_epochs(&header, "new", "app-epoch-one", 3, 1).unwrap(); + let rotated = validate_luks2_headers(fs::File::open(&header).unwrap()).unwrap(); + assert_eq!(rotated.root_epoch, Some(3)); + assert_eq!(rotated.app_epoch, 1); + assert!(cryptsetup( + &["luksHeaderRestore", "--batch-mode", "--header-backup-file"], + &header + ) + .arg(&image) + .status() + .unwrap() + .success()); + assert!(!unlocks("new")); + assert!(unlocks("app-epoch-one")); } #[test] diff --git a/dstack/dstack-util/src/system_setup/sealed_env.rs b/dstack/dstack-util/src/system_setup/sealed_env.rs index e3e0e1adc..ef7a717f6 100644 --- a/dstack/dstack-util/src/system_setup/sealed_env.rs +++ b/dstack/dstack-util/src/system_setup/sealed_env.rs @@ -10,28 +10,36 @@ use serde::{Deserialize, Serialize}; #[derive(Serialize, Deserialize)] pub(super) struct SealedEnv { pub(super) version: u32, - pub(super) app_epoch: u32, pub(super) root_epoch: u32, + pub(super) app_epoch: u32, #[serde(with = "serde_human_bytes::base64")] pub(super) enc: Vec, #[serde(with = "serde_human_bytes::base64")] pub(super) ciphertext: Vec, } -pub(super) fn sealed_env_info(epoch: &u32) -> Vec { - format!("dstack-sealed-env-v1:{epoch}:0").into_bytes() +pub(super) fn sealed_env_info_at_epochs(root: u32, app: u32) -> Vec { + format!("dstack-sealed-env-v1:{root}:{app}").into_bytes() } -pub(super) fn seal_env(secret: &[u8], epoch: u32, plaintext: &[u8]) -> Result> { +pub(super) fn seal_env_at_epochs( + secret: &[u8], + epoch: u32, + app_epoch: u32, + plaintext: &[u8], +) -> Result> { let secret: [u8; 32] = secret.try_into().context("invalid sealed env key")?; let public = x25519_dalek::PublicKey::from(&x25519_dalek::StaticSecret::from(secret)); - let payload = - crate::crypto::sealed_env_encrypt(&public.to_bytes(), &sealed_env_info(&epoch), plaintext)?; + let payload = crate::crypto::sealed_env_encrypt( + &public.to_bytes(), + &sealed_env_info_at_epochs(epoch, app_epoch), + plaintext, + )?; let (enc, ciphertext) = payload.split_at(32); Ok(serde_json::to_vec(&SealedEnv { version: 1, - app_epoch: 0, root_epoch: epoch, + app_epoch, enc: enc.to_vec(), ciphertext: ciphertext.to_vec(), })?) @@ -42,3 +50,8 @@ pub(super) fn env_crypt_key(key: &[u8]) -> Result<[u8; 32]> { .ok() .context("encrypted env is present but the key provider gave no valid env crypt key") } + +#[cfg(test)] +pub(super) fn seal_env(secret: &[u8], epoch: u32, plaintext: &[u8]) -> Result> { + seal_env_at_epochs(secret, epoch, 0, plaintext) +} diff --git a/dstack/guest-agent/rpc/proto/agent_rpc_v1.proto b/dstack/guest-agent/rpc/proto/agent_rpc_v1.proto index bebc85f99..1a36a6692 100644 --- a/dstack/guest-agent/rpc/proto/agent_rpc_v1.proto +++ b/dstack/guest-agent/rpc/proto/agent_rpc_v1.proto @@ -194,6 +194,7 @@ message GetKeyRequest { // epoch only to reach a key from before a KMS root key rotation, for example // to migrate data or assets off it. optional uint32 root_epoch = 3; + optional uint32 app_epoch = 4; } message GetKeyResponse { @@ -217,6 +218,7 @@ message GetKeyResponse { // KMS root key epoch the key was derived from. Link [1] is signed by that // epoch's KMS root key. uint32 root_epoch = 4; + uint32 app_epoch = 5; } // --------------------------------------------------------------------------- diff --git a/dstack/guest-agent/src/rpc_service.rs b/dstack/guest-agent/src/rpc_service.rs index 58ff5a865..db0983c83 100644 --- a/dstack/guest-agent/src/rpc_service.rs +++ b/dstack/guest-agent/src/rpc_service.rs @@ -97,17 +97,19 @@ pub struct AppState { #[derive(Clone)] pub(crate) struct AppRoot { pub epoch: u32, + pub app_epoch: u32, pub k256_key: Vec, pub signing_key: SigningKey, pub kms_signature: Vec, } impl AppRoot { - fn new(epoch: u32, k256_key: Vec, kms_signature: Vec) -> Result { + fn new(epoch: u32, app_epoch: u32, k256_key: Vec, kms_signature: Vec) -> Result { let signing_key = SigningKey::from_slice(&k256_key) .context("the app root k256 key is not a valid secp256k1 scalar")?; Ok(Self { epoch, + app_epoch, k256_key, signing_key, kms_signature, @@ -121,7 +123,7 @@ struct AppStateInner { vm_config: String, cert_client: CertRequestClient, /// Held across a fetch so concurrent requests for one epoch fetch it once. - historical_roots: tokio::sync::Mutex>, + historical_roots: tokio::sync::Mutex>, demo_cert: tokio::sync::OnceCell, /// When the last demo-cert request failed; held while one is running. demo_cert_last_failure: tokio::sync::Mutex>, @@ -411,12 +413,22 @@ impl AppState { /// The app root key and its KMS signature of `root_epoch`; the current /// ones when unset. An older epoch is requested from the KMS once and kept /// in memory. - pub(crate) async fn app_root(&self, root_epoch: Option) -> Result { + pub(crate) async fn app_root( + &self, + root_epoch: Option, + app_epoch: Option, + ) -> Result { let current = self.root_epoch(); let epoch = root_epoch.unwrap_or(current); - if epoch == current { + let current_app = self.inner.keys.app_epoch(); + let app_epoch = app_epoch.unwrap_or(current_app); + if app_epoch > current_app { + anyhow::bail!("app epoch {app_epoch} is newer than current {current_app}"); + } + if epoch == current && app_epoch == current_app { return Ok(AppRoot { epoch, + app_epoch, k256_key: self.app_root_k256_key().to_vec(), signing_key: self.app_root_signing_key()?.clone(), kms_signature: self.kms_k256_signature().to_vec(), @@ -425,16 +437,16 @@ impl AppState { if epoch > current { anyhow::bail!("root epoch {epoch} is newer than the current one, {current}"); } - self.historical_root(epoch).await + self.historical_root(epoch, app_epoch).await } - pub(crate) async fn historical_root(&self, epoch: u32) -> Result { + pub(crate) async fn historical_root(&self, epoch: u32, app_epoch: u32) -> Result { if let Some(root) = self .inner .historical_roots .lock() .await - .get(&epoch) + .get(&(epoch, app_epoch)) .cloned() { return Ok(root); @@ -442,12 +454,15 @@ impl AppState { let keys = self .inner .cert_client - .get_app_key(epoch) + .get_app_key(epoch, app_epoch) .await .with_context(|| format!("failed to fetch the app keys of root epoch {epoch}"))?; - let root = AppRoot::new(epoch, keys.k256_key, keys.k256_signature)?; + let root = AppRoot::new(epoch, app_epoch, keys.k256_key, keys.k256_signature)?; let mut roots = self.inner.historical_roots.lock().await; - Ok(roots.entry(epoch).or_insert_with(|| root.clone()).clone()) + Ok(roots + .entry((epoch, app_epoch)) + .or_insert_with(|| root.clone()) + .clone()) } #[cfg(test)] @@ -457,8 +472,12 @@ impl AppState { k256_key: Vec, kms_signature: Vec, ) { - let root = AppRoot::new(epoch, k256_key, kms_signature).unwrap(); - self.inner.historical_roots.lock().await.insert(epoch, root); + let root = AppRoot::new(epoch, 0, k256_key, kms_signature).unwrap(); + self.inner + .historical_roots + .lock() + .await + .insert((epoch, 0), root); } /// The KMS root key's signature over the app root public key: the second @@ -1242,6 +1261,39 @@ pub(crate) mod tests { setup_test_state_with_platform(None).await } + #[tokio::test] + async fn app_root_selectors_and_cache_keep_epochs_independent() { + let (mut state, _guard) = setup_test_state().await; + Arc::get_mut(&mut state.inner).unwrap().keys.key_provider = KeyProvider::Kms { + url: "https://unused.invalid".into(), + pubkey: vec![], + root_epoch: 2, + app_epoch: 2, + }; + for (root, app, byte) in [(2, 0, 0x11), (0, 2, 0x22), (0, 0, 0x33)] { + state.inner.historical_roots.lock().await.insert( + (root, app), + AppRoot::new(root, app, vec![byte; 32], vec![]).unwrap(), + ); + } + let current = state.app_root(None, None).await.unwrap(); + assert_eq!((current.epoch, current.app_epoch), (2, 2)); + for (root, app, byte) in [ + (None, Some(0), 0x11), + (Some(0), None, 0x22), + (Some(0), Some(0), 0x33), + ] { + let key = state.app_root(root, app).await.unwrap(); + assert_eq!( + (key.epoch, key.app_epoch), + (root.unwrap_or(2), app.unwrap_or(2)) + ); + assert_eq!(key.k256_key, vec![byte; 32]); + } + assert!(state.app_root(Some(3), None).await.is_err()); + assert!(state.app_root(None, Some(3)).await.is_err()); + } + /// How many times the fixture platform was asked to attest for `Info`, and /// whether it should refuse. Counting the calls is the only way to see the /// identity throttle work: what it changes is how often the platform is diff --git a/dstack/guest-agent/src/rpc_service_v1.rs b/dstack/guest-agent/src/rpc_service_v1.rs index 8f178ccbf..7eadc40ea 100644 --- a/dstack/guest-agent/src/rpc_service_v1.rs +++ b/dstack/guest-agent/src/rpc_service_v1.rs @@ -166,12 +166,13 @@ impl V1RpcHandler { domain: &str, algorithm: &str, root_epoch: Option, - ) -> Result<(AppKey, Vec>, u32)> { + app_epoch: Option, + ) -> Result<(AppKey, Vec>, u32, u32)> { let algorithm = Algorithm::parse(algorithm)?; - let root = self.state.app_root(root_epoch).await?; + let root = self.state.app_root(root_epoch, app_epoch).await?; let key = AppKey::derive(&root.k256_key, domain, algorithm)?; let chain = vec![key.claim_signature(&root.signing_key)?, root.kms_signature]; - Ok((key, chain, root.epoch)) + Ok((key, chain, root.epoch, root.app_epoch)) } } @@ -199,14 +200,20 @@ impl DstackGuestRpc for V1RpcHandler { } async fn get_key(self, request: GetKeyRequest) -> Result { - let (key, signature_chain, root_epoch) = self - .derive(&request.domain, &request.algorithm, request.root_epoch) + let (key, signature_chain, root_epoch, app_epoch) = self + .derive( + &request.domain, + &request.algorithm, + request.root_epoch, + request.app_epoch, + ) .await?; Ok(GetKeyResponse { key: key.secret(), public_key: key.public_key(), signature_chain, root_epoch, + app_epoch, }) } @@ -349,6 +356,7 @@ mod tests { domain: domain.to_string(), algorithm: algorithm.to_string(), root_epoch: None, + app_epoch: None, } } @@ -363,6 +371,14 @@ mod tests { let current = handler().get_key(request(None)).await.unwrap(); assert_eq!(current.root_epoch, 0); + assert_eq!(current.app_epoch, 0); + assert!(handler() + .get_key(GetKeyRequest { + app_epoch: Some(1), + ..request(None) + }) + .await + .is_err()); assert_eq!( handler().get_key(request(Some(0))).await.unwrap().key, current.key @@ -379,11 +395,12 @@ mod tests { state .insert_historical_root(1, key.clone(), vec![0xab; 65]) .await; - let root = state.historical_root(1).await.unwrap(); + let root = state.historical_root(1, 0).await.unwrap(); assert_eq!(root.k256_key, key); assert_eq!(root.kms_signature, [0xab; 65]); // An uncached epoch needs the KMS, which a local key provider lacks. - assert!(state.historical_root(2).await.is_err()); + assert!(state.historical_root(2, 0).await.is_err()); + assert!(state.historical_root(1, 1).await.is_err()); } /// The fixture's app root key is what `keys::tests` committed its vectors diff --git a/dstack/kms/README.md b/dstack/kms/README.md index 7e5d4b710..368704d45 100644 --- a/dstack/kms/README.md +++ b/dstack/kms/README.md @@ -257,12 +257,12 @@ sent. The KMS method checks the KMS policy; the app method applies registered ap policy. KMS callers do not fetch contract metadata to decide which policy to use. -`requestedAppEpoch` reserves the application's own key epoch, independently of KMS root -epochs. HTTP requests default it to 0; KMS currently sends 0. Solidity uses the same -`OptionalEpoch` tuple as the root selector, reserving omission for the current -app epoch. Responses include `currentAppEpoch`, currently 0. Nonzero values are -rejected until app-epoch support is implemented, and rejection responses still -carry the current root epoch and chain hash. +`requestedAppEpoch` selects the application's own key epoch, independently of KMS +root epochs. Omission uses `currentAppEpoch`; explicit zero selects app epoch zero. +Solidity uses `OptionalEpoch { isSet, value }` for both selectors. Future epochs +are denied; historical epochs remain subject to current authorization policy. +Rejection responses still carry the current epochs and root chain hash. See +[application key epochs](../../docs/app-key-epochs.md) for rotation and migration. `requestedRootEpoch` is optional in the JSON request. Omit it to use the epoch in this authorization result; an explicit `0` selects epoch 0. Solidity represents diff --git a/dstack/kms/auth-eth-bun/README.md b/dstack/kms/auth-eth-bun/README.md index bc9c72e19..066ccf73b 100644 --- a/dstack/kms/auth-eth-bun/README.md +++ b/dstack/kms/auth-eth-bun/README.md @@ -204,10 +204,11 @@ the test suite (`index.test.ts`) validates: ### Authorization endpoints `POST /authorizeKms` and `POST /authorizeApp` return `AuthorizationResult`, including -`rejectionReason`, `gatewayAppId`, `currentRootEpoch`, and `currentRootEpochHash`. +`rejectionReason`, `gatewayAppId`, `currentRootEpoch`, `currentRootEpochHash`, and `currentAppEpoch`. The selected endpoint determines the policy; the backend checks that the verified -`appId` matches that scope. `requestedAppEpoch` defaults to 0; nonzero values are currently -rejected. The optional +`appId` matches that scope. Optional `requestedAppEpoch` selects the app epoch; +omission uses the current app epoch, explicit zero selects epoch zero, and future +values are denied. The optional `requestedRootEpoch` selects historical keys; omission uses the returned current epoch, and an explicit `0` selects epoch 0. See the [KMS authorization protocol](../README.md#authorization-protocol) for registration, diff --git a/dstack/kms/auth-eth-bun/index.test.ts b/dstack/kms/auth-eth-bun/index.test.ts index dda821bba..224040730 100644 --- a/dstack/kms/auth-eth-bun/index.test.ts +++ b/dstack/kms/auth-eth-bun/index.test.ts @@ -858,7 +858,7 @@ describe('POST /authorizeApp', () => { expect(mockReadContract).toHaveBeenCalledTimes(1); expect(mockReadContract).toHaveBeenCalledWith(expect.objectContaining({ functionName: 'authorizeApp', blockNumber: 98n, - args: [expect.objectContaining({ appInfo: expect.objectContaining({ appId: appInfo.appId }), requestedAppEpoch: { isSet: true, value: 0n }, requestedRootEpoch: { isSet: false, value: 0n } })], + args: [expect.objectContaining({ appInfo: expect.objectContaining({ appId: appInfo.appId }), requestedAppEpoch: { isSet: false, value: 0n }, requestedRootEpoch: { isSet: false, value: 0n } })], })); }); diff --git a/dstack/kms/auth-eth-bun/index.ts b/dstack/kms/auth-eth-bun/index.ts index f4e12ab3a..376c79a0b 100644 --- a/dstack/kms/auth-eth-bun/index.ts +++ b/dstack/kms/auth-eth-bun/index.ts @@ -38,7 +38,7 @@ const BootInfoSchema = z.object({ const AuthorizationRequestSchema = BootInfoSchema.omit({ rootEpoch: true, rootEpochHash: true }).extend({ operation: z.enum(['kms_bootstrap', 'kms_serve', 'kms_export_keys', 'kms_import_keys', 'app_get_key', 'app_get_cert']), - requestedAppEpoch: z.number().int().min(0).max(0xffffffff).default(0), + requestedAppEpoch: z.number().int().min(0).max(0xffffffff).optional(), requestedRootEpoch: z.number().int().min(0).max(0xffffffff).optional(), }); type AuthorizationRequest = z.infer; @@ -119,6 +119,7 @@ const AUTHORIZATION_ABI = { // DstackKms contract ABI (minimal required functions) const DSTACK_KMS_ABI = [ + { name: 'currentAppEpoch', type: 'function', stateMutability: 'view', inputs: [{ name: 'appId', type: 'address' }], outputs: [{ name: '', type: 'uint256' }] }, { ...AUTHORIZATION_ABI, name: 'authorizeKms' }, { ...AUTHORIZATION_ABI, name: 'authorizeApp' }, { @@ -324,6 +325,15 @@ class EthereumBackend { return { rootEpoch: Number(epoch), rootEpochHash: chainHash }; } + async getAppEpoch(appId: string): Promise { + const epoch = Number(await this.client.readContract({ + address: this.kmsContractAddr, abi: DSTACK_KMS_ABI, functionName: 'currentAppEpoch', + args: [this.decodeHex(appId, 20)], blockNumber: await this.finalizedBlockNumber(), + })); + if (!Number.isSafeInteger(epoch)) throw new Error('app epoch exceeds the supported range'); + return epoch; + } + async getGatewayAppId(): Promise { const result = await this.client.readContract({ address: this.kmsContractAddr, @@ -425,6 +435,16 @@ app.get('/', async (c) => { } }); +app.get('/app/:appId', async (c) => { + const appId = c.req.param('appId'); + if (!/^0x[0-9a-fA-F]{40}$/.test(appId)) return c.json({ error: 'invalid app ID' }, 400); + try { + return c.json({ currentAppEpoch: await ethereum.getAppEpoch(appId) }); + } catch { + return c.json({ error: 'authorization backend unavailable' }, 503); + } +}); + for (const method of ['authorizeKms', 'authorizeApp'] as const) { app.post(`/${method}`, bodyLimit({ maxSize: MAX_BOOT_INFO_BYTES }), diff --git a/dstack/kms/auth-eth-bun/openapi.json b/dstack/kms/auth-eth-bun/openapi.json index e431225a9..3b12cefa1 100644 --- a/dstack/kms/auth-eth-bun/openapi.json +++ b/dstack/kms/auth-eth-bun/openapi.json @@ -160,6 +160,44 @@ } } } + }, + "/app/{appId}": { + "get": { + "summary": "Read the current application key epoch", + "parameters": [ + { + "name": "appId", + "in": "path", + "required": true, + "schema": { + "type": "string", + "pattern": "^0x[0-9a-fA-F]{40}$" + } + } + ], + "responses": { + "200": { + "description": "Current application epoch", + "content": { + "application/json": { + "schema": { + "type": "object", + "required": [ + "currentAppEpoch" + ], + "properties": { + "currentAppEpoch": { + "type": "integer", + "minimum": 0, + "maximum": 4294967295 + } + } + } + } + } + } + } + } } }, "components": { @@ -426,8 +464,7 @@ "type": "integer", "minimum": 0, "maximum": 4294967295, - "default": 0, - "description": "Application epoch, independent of KMS root epochs. Only zero is currently supported." + "description": "Omit to use the current application epoch. Explicit zero selects epoch 0; future epochs are denied." }, "operation": { "type": "string", diff --git a/dstack/kms/auth-eth/README.md b/dstack/kms/auth-eth/README.md index c7a8d6e7b..85eb6a52e 100644 --- a/dstack/kms/auth-eth/README.md +++ b/dstack/kms/auth-eth/README.md @@ -175,10 +175,11 @@ Documentation: https://book.getfoundry.sh/ ### Authorization endpoints `POST /authorizeKms` and `POST /authorizeApp` return `AuthorizationResult`, including -`rejectionReason`, `gatewayAppId`, `currentRootEpoch`, and `currentRootEpochHash`. +`rejectionReason`, `gatewayAppId`, `currentRootEpoch`, `currentRootEpochHash`, and `currentAppEpoch`. The selected endpoint determines the policy; the backend checks that the verified -`appId` matches that scope. `requestedAppEpoch` defaults to 0; nonzero values are currently -rejected. The optional +`appId` matches that scope. Optional `requestedAppEpoch` selects the app epoch; +omission uses the current app epoch, explicit zero selects epoch zero, and future +values are denied. The optional `requestedRootEpoch` selects historical keys; omission uses the returned current epoch, and an explicit `0` selects epoch 0. See the [KMS authorization protocol](../README.md#authorization-protocol) for registration, diff --git a/dstack/kms/auth-eth/contracts/DstackApp.sol b/dstack/kms/auth-eth/contracts/DstackApp.sol index d55a6c883..951fe7868 100644 --- a/dstack/kms/auth-eth/contracts/DstackApp.sol +++ b/dstack/kms/auth-eth/contracts/DstackApp.sol @@ -7,6 +7,7 @@ pragma solidity ^0.8.24; import "./IAppAuth.sol"; +import "./IAppEpoch.sol"; import "./IAppAuthBasicManagement.sol"; import "@openzeppelin/contracts-upgradeable/proxy/utils/Initializable.sol"; import "@openzeppelin/contracts-upgradeable/access/Ownable2StepUpgradeable.sol"; @@ -19,7 +20,8 @@ contract DstackApp is UUPSUpgradeable, ERC165Upgradeable, IAppAuth, - IAppAuthBasicManagement + IAppAuthBasicManagement, + IAppEpoch { // Mapping of allowed compose hashes for this app mapping(bytes32 => bool) public allowedComposeHashes; @@ -36,6 +38,9 @@ contract DstackApp is // Whether to require TCB status to be UpToDate bool public requireTcbUpToDate; + uint256 public override currentAppEpoch; + event AppEpochAdvanced(uint256 indexed epoch); + // Additional events specific to DstackApp event UpgradesDisabled(); event AllowAnyDeviceSet(bool allowAny); @@ -128,7 +133,7 @@ contract DstackApp is override(ERC165Upgradeable, IERC165) returns (bool) { - return interfaceId == 0x1e079198 // IAppAuth + return interfaceId == type(IAppEpoch).interfaceId || interfaceId == 0x1e079198 // IAppAuth || interfaceId == 0x8fd37527 // IAppAuthBasicManagement || super.supportsInterface(interfaceId); } @@ -215,6 +220,15 @@ contract DstackApp is return (true, ""); } + /// Advance only from the epoch the owner intended, rejecting stale transactions. + function advanceAppEpoch(uint256 expectedCurrentEpoch) external onlyOwner { + require(currentAppEpoch == expectedCurrentEpoch, "App epoch changed"); + require(currentAppEpoch < type(uint32).max, "App epoch exhausted"); + currentAppEpoch++; + emit AppEpochAdvanced(currentAppEpoch); + _emitPolicy("app-epoch", bytes32(currentAppEpoch), true); + } + // Function to permanently disable upgrades function disableUpgrades() external onlyOwner { _upgradesDisabled = true; @@ -223,5 +237,5 @@ contract DstackApp is } // Add storage gap for upgradeable contracts - uint256[49] private __gap; + uint256[48] private __gap; } diff --git a/dstack/kms/auth-eth/contracts/DstackKms.sol b/dstack/kms/auth-eth/contracts/DstackKms.sol index a31483e60..77b3e2fbc 100644 --- a/dstack/kms/auth-eth/contracts/DstackKms.sol +++ b/dstack/kms/auth-eth/contracts/DstackKms.sol @@ -8,6 +8,8 @@ pragma solidity ^0.8.24; import "./IAppAuth.sol"; import "./AuthorizationOperation.sol"; +import "./IAppEpoch.sol"; +import "@openzeppelin/contracts/utils/introspection/ERC165Checker.sol"; import "@openzeppelin/contracts-upgradeable/proxy/utils/Initializable.sol"; import "@openzeppelin/contracts-upgradeable/access/Ownable2StepUpgradeable.sol"; import "@openzeppelin/contracts-upgradeable/proxy/utils/UUPSUpgradeable.sol"; @@ -393,6 +395,13 @@ contract DstackKms is Initializable, Ownable2StepUpgradeable, UUPSUpgradeable, E _emitPolicy("os-image", osImageHash, false); } + /// Legacy policies have epoch zero; an advertised epoch getter must succeed. + function currentAppEpoch(address appId) public view returns (uint256) { + if (!registeredApps[appId] || appId.code.length == 0) return 0; + if (!ERC165Checker.supportsInterface(appId, type(IAppEpoch).interfaceId)) return 0; + return IAppEpoch(appId).currentAppEpoch(); + } + /// Authorize an attested identity and return the root epoch from the same state. /// The caller selects KMS or app policy; KMS authorization requires the proxy app ID. function authorizeKms(AuthorizationRequest calldata request) external view returns (AuthorizationResult memory) { @@ -412,6 +421,7 @@ contract DstackKms is Initializable, Ownable2StepUpgradeable, UUPSUpgradeable, E returns (AuthorizationResult memory result) { result.gatewayAppId = gatewayAppId; + result.currentAppEpoch = isKms ? 0 : currentAppEpoch(request.appInfo.appId); if (!(isKms ? AuthorizationOperation.isKms(request.operation) : AuthorizationOperation.isApp(request.operation))) { @@ -431,8 +441,8 @@ contract DstackKms is Initializable, Ownable2StepUpgradeable, UUPSUpgradeable, E result.rejectionReason = "Requested root epoch is not active"; return result; } - if (request.requestedAppEpoch.isSet && request.requestedAppEpoch.value != 0) { - result.rejectionReason = "Requested app epoch must be zero"; + if (request.requestedAppEpoch.isSet && request.requestedAppEpoch.value > result.currentAppEpoch) { + result.rejectionReason = "Requested app epoch is not active"; return result; } AppInfo calldata appInfo = request.appInfo; @@ -490,6 +500,9 @@ contract DstackKms is Initializable, Ownable2StepUpgradeable, UUPSUpgradeable, E override returns (bool isAllowed, string memory reason) { + if (currentAppEpoch(bootInfo.appId) != 0) { + return (false, "App epoch requires epoch-aware authorization"); + } return _isAppAllowed(bootInfo); } diff --git a/dstack/kms/auth-eth/contracts/IAppEpoch.sol b/dstack/kms/auth-eth/contracts/IAppEpoch.sol new file mode 100644 index 000000000..546c0e26a --- /dev/null +++ b/dstack/kms/auth-eth/contracts/IAppEpoch.sol @@ -0,0 +1,11 @@ +/* + * SPDX-FileCopyrightText: © 2026 Phala Network + * + * SPDX-License-Identifier: Apache-2.0 + */ + +pragma solidity ^0.8.24; + +interface IAppEpoch { + function currentAppEpoch() external view returns (uint256); +} diff --git a/dstack/kms/auth-eth/script/Manage.s.sol b/dstack/kms/auth-eth/script/Manage.s.sol index 7f57f1d62..be924ecf2 100644 --- a/dstack/kms/auth-eth/script/Manage.s.sol +++ b/dstack/kms/auth-eth/script/Manage.s.sol @@ -363,3 +363,14 @@ contract DisableAppUpgrades is BaseScript { // a live proxy to the current `contracts/DstackKms.sol` / `contracts/DstackApp.sol` // source. The `UpgradeKmsToV2` / `UpgradeAppToV2` variants target the V2 test // mocks under `contracts/test-utils/` and are for upgrade-flow tests only. + +contract AdvanceAppEpoch is BaseScript { + function run() external { + require(address(app) != address(0), "APP_CONTRACT_ADDR not set"); + uint256 expected = vm.envUint("EXPECTED_APP_EPOCH"); + vm.startBroadcast(vm.envUint("PRIVATE_KEY")); + app.advanceAppEpoch(expected); + vm.stopBroadcast(); + console.log("Current app epoch:", app.currentAppEpoch()); + } +} diff --git a/dstack/kms/auth-eth/script/Query.s.sol b/dstack/kms/auth-eth/script/Query.s.sol index 933d13578..534b2de46 100644 --- a/dstack/kms/auth-eth/script/Query.s.sol +++ b/dstack/kms/auth-eth/script/Query.s.sol @@ -250,3 +250,11 @@ contract GetAppSettings is BaseQueryScript { console.log("\nProxy Implementation:", implementation); } } + +contract GetAppEpoch is BaseQueryScript { + function run() external view { + require(address(kms) != address(0), "KMS_CONTRACT_ADDR not set"); + require(address(app) != address(0), "APP_CONTRACT_ADDR not set"); + console.log("Current app epoch:", kms.currentAppEpoch(address(app))); + } +} diff --git a/dstack/kms/auth-eth/src/ethereum.ts b/dstack/kms/auth-eth/src/ethereum.ts index 2aedc8b54..24d0fa8ca 100644 --- a/dstack/kms/auth-eth/src/ethereum.ts +++ b/dstack/kms/auth-eth/src/ethereum.ts @@ -7,6 +7,7 @@ import { BootInfo, BootResponse, RootEpochInfo, AuthorizationResult, Authorizati // Minimal ABI for DstackKms contract const DSTACK_KMS_ABI = [ + "function currentAppEpoch(address appId) view returns (uint256)", "function authorizeKms(((address appId,bytes32 composeHash,address instanceId,bytes32 deviceId,bytes32 mrAggregated,bytes32 mrSystem,bytes32 osImageHash,string tcbStatus,string[] advisoryIds) appInfo,bytes32 operation,(bool isSet,uint256 value) requestedAppEpoch,(bool isSet,uint256 value) requestedRootEpoch) request) view returns ((bool isAllowed,string rejectionReason,string gatewayAppId,uint256 currentRootEpoch,bytes32 currentRootEpochHash,uint256 currentAppEpoch))", "function authorizeApp(((address appId,bytes32 composeHash,address instanceId,bytes32 deviceId,bytes32 mrAggregated,bytes32 mrSystem,bytes32 osImageHash,string tcbStatus,string[] advisoryIds) appInfo,bytes32 operation,(bool isSet,uint256 value) requestedAppEpoch,(bool isSet,uint256 value) requestedRootEpoch) request) view returns ((bool isAllowed,string rejectionReason,string gatewayAppId,uint256 currentRootEpoch,bytes32 currentRootEpochHash,uint256 currentAppEpoch))", "function isAppAllowed((address appId,bytes32 composeHash,address instanceId,bytes32 deviceId,bytes32 mrAggregated,bytes32 mrSystem,bytes32 osImageHash,string tcbStatus,string[] advisoryIds) bootInfo) view returns (bool, string)", @@ -126,6 +127,10 @@ export class EthereumBackend { return { rootEpoch: Number(epoch), rootEpochHash: chainHash }; } + async getAppEpoch(appId: string): Promise { + return ethers.toNumber(await this.kmsContract.currentAppEpoch(this.decodeHex(appId, 20))); + } + async getGatewayAppId(): Promise { return await this.kmsContract.gatewayAppId(); } diff --git a/dstack/kms/auth-eth/src/server.ts b/dstack/kms/auth-eth/src/server.ts index e3d81424a..8aab70c55 100644 --- a/dstack/kms/auth-eth/src/server.ts +++ b/dstack/kms/auth-eth/src/server.ts @@ -57,7 +57,7 @@ export async function build(): Promise { properties: { ...appInfoSchema.properties, operation: { type: 'string', enum: ['kms_bootstrap', 'kms_serve', 'kms_export_keys', 'kms_import_keys', 'app_get_key', 'app_get_cert'] }, - requestedAppEpoch: { type: 'integer', minimum: 0, maximum: 4294967295, default: 0 }, + requestedAppEpoch: { type: 'integer', minimum: 0, maximum: 4294967295 }, requestedRootEpoch: { type: 'integer', minimum: 0, maximum: 4294967295 }, }, }); @@ -122,6 +122,10 @@ export async function build(): Promise { } }); + server.get<{ Params: { appId: string } }>('/app/:appId', { + schema: { params: { type: 'object', required: ['appId'], properties: { appId: { type: 'string', pattern: '^0x[0-9a-fA-F]{40}$' } } } }, + }, async (request) => ({ currentAppEpoch: await server.ethereum.getAppEpoch(request.params.appId) })); + for (const method of ['authorizeKms', 'authorizeApp'] as const) { server.post<{ Body: AuthorizationRequest; Reply: AuthorizationResult }>(`/${method}`, { schema: { diff --git a/dstack/kms/auth-eth/test/DstackKms.t.sol b/dstack/kms/auth-eth/test/DstackKms.t.sol index 138759983..1be41706a 100644 --- a/dstack/kms/auth-eth/test/DstackKms.t.sol +++ b/dstack/kms/auth-eth/test/DstackKms.t.sol @@ -573,7 +573,7 @@ contract DstackKmsTest is Test { ); result = kms.authorizeApp(appRequest); assertFalse(result.isAllowed); - assertEq(result.rejectionReason, "Requested app epoch must be zero"); + assertEq(result.rejectionReason, "Requested app epoch is not active"); assertEq(result.currentRootEpochHash, chainHash(bytes32(0), ca(0xa0), k256(0xb0))); info.appId = address(kms); result = authorize(info, false); @@ -584,6 +584,62 @@ contract DstackKmsTest is Test { assertEq(result.rejectionReason, "Aggregated MR not allowed"); } + function test_AppEpochCapabilityIsOptionalButGetterFailuresPropagate() public { + address policy = address(new DstackApp()); + kms.registerApp(policy); + vm.mockCall(policy, abi.encodeCall(IERC165.supportsInterface, (type(IAppEpoch).interfaceId)), abi.encode(false)); + assertEq(kms.currentAppEpoch(policy), 0); + vm.clearMockedCalls(); + vm.mockCallRevert(policy, abi.encodeCall(IAppEpoch.currentAppEpoch, ()), "epoch unavailable"); + vm.expectRevert("epoch unavailable"); + kms.currentAppEpoch(policy); + } + + function test_AppEpochRotationIsOwnerScopedAndDoesNotChangeRootEpoch() public { + DstackKms.AppInfo memory info = authorizationInfo(address(0)); + vm.startPrank(owner); + kms.registerRootEpoch(ca(0xa0), k256(0xb0), "https://example.com/epoch-0"); + kms.addOsImageHash(info.osImageHash); + address appId = kms.deployAndRegisterApp(user, false, false, info.deviceId, info.composeHash); + address otherApp = kms.deployAndRegisterApp(user, false, false, info.deviceId, info.composeHash); + vm.stopPrank(); + info.appId = appId; + assertEq(kms.currentAppEpoch(appId), 0); + vm.prank(owner); + vm.expectRevert(); + DstackApp(appId).advanceAppEpoch(0); + vm.prank(user); + DstackApp(appId).advanceAppEpoch(0); + vm.prank(user); + vm.expectRevert("App epoch changed"); + DstackApp(appId).advanceAppEpoch(0); + assertEq(kms.currentAppEpoch(appId), 1); + assertEq(kms.currentAppEpoch(otherApp), 0); + (uint256 rootEpoch,) = kms.currentRootEpoch(); + assertEq(rootEpoch, 0); + DstackKms.AuthorizationRequest memory request = DstackKms.AuthorizationRequest( + info, DstackKms.OptionalEpoch(false, 0), DstackKms.OptionalEpoch(false, 0) + ); + DstackKms.AuthorizationResult memory result = kms.checkAppAuthorization(request); + assertTrue(result.isAllowed); + assertEq(result.currentAppEpoch, 1); + request.requestedAppEpoch = DstackKms.OptionalEpoch(true, 0); + assertTrue(kms.checkAppAuthorization(request).isAllowed); + request.requestedAppEpoch.value = 2; + result = kms.checkAppAuthorization(request); + assertFalse(result.isAllowed); + assertEq(result.currentAppEpoch, 1); + assertEq(result.rejectionReason, "Requested app epoch is not active"); + request.requestedAppEpoch = DstackKms.OptionalEpoch(false, 0); + request.appInfo.composeHash = bytes32(0); + result = kms.checkAppAuthorization(request); + assertFalse(result.isAllowed); + assertEq(result.currentAppEpoch, 1); + vm.prank(user); + DstackApp(appId).advanceAppEpoch(1); + assertEq(kms.currentAppEpoch(appId), 2); + } + function test_CheckAuthorization_RejectionIncludesActivatedEpoch() public { vm.startPrank(owner); kms.registerRootEpoch(ca(0xa0), k256(0xb0), "https://example.com/epoch-0"); diff --git a/dstack/kms/auth-mock/README.md b/dstack/kms/auth-mock/README.md index 4ddb11646..bec80cd4a 100644 --- a/dstack/kms/auth-mock/README.md +++ b/dstack/kms/auth-mock/README.md @@ -172,10 +172,11 @@ the API remains identical, so no client code changes are needed. ### Authorization endpoints `POST /authorizeKms` and `POST /authorizeApp` return `AuthorizationResult`, including -`rejectionReason`, `gatewayAppId`, `currentRootEpoch`, and `currentRootEpochHash`. +`rejectionReason`, `gatewayAppId`, `currentRootEpoch`, `currentRootEpochHash`, and `currentAppEpoch`. The selected endpoint determines the policy; the backend checks that the verified -`appId` matches that scope. `requestedAppEpoch` defaults to 0; nonzero values are currently -rejected. The optional +`appId` matches that scope. Optional `requestedAppEpoch` selects the app epoch; +omission uses the current app epoch, explicit zero selects epoch zero, and future +values are denied. The optional `requestedRootEpoch` selects historical keys; omission uses the returned current epoch, and an explicit `0` selects epoch 0. See the [KMS authorization protocol](../README.md#authorization-protocol) for registration, diff --git a/dstack/kms/auth-mock/index.ts b/dstack/kms/auth-mock/index.ts index 518b3e334..d29a84a8d 100644 --- a/dstack/kms/auth-mock/index.ts +++ b/dstack/kms/auth-mock/index.ts @@ -31,7 +31,7 @@ const BootInfoSchema = z.object({ const AuthorizationRequestSchema = BootInfoSchema.omit({ rootEpoch: true, rootEpochHash: true }).extend({ operation: z.enum(['kms_bootstrap', 'kms_serve', 'kms_export_keys', 'kms_import_keys', 'app_get_key', 'app_get_cert']), - requestedAppEpoch: z.number().int().min(0).max(0xffffffff).default(0), + requestedAppEpoch: z.number().int().min(0).max(0xffffffff).optional(), requestedRootEpoch: z.number().int().min(0).max(0xffffffff).optional(), }); type AuthorizationRequest = z.infer; @@ -113,10 +113,14 @@ class MockBackend { this.mockAppImplementation = process.env.MOCK_APP_IMPLEMENTATION || '0xmockapp9876543210987654321098765432109'; } + async getAppEpoch(_appId: string): Promise { + return z.coerce.number().int().min(0).max(0xffffffff).parse(process.env.MOCK_APP_EPOCH ?? 0); + } + async checkAuthorization(appInfo: AuthorizationRequest, isKms: boolean) { const current = getRootEpoch(); const epoch = { - currentAppEpoch: 0, + currentAppEpoch: isKms ? 0 : await this.getAppEpoch(appInfo.appId), currentRootEpoch: current.rootEpoch ?? 0, currentRootEpochHash: current.rootEpochHash ?? '0x' + '00'.repeat(32), }; @@ -130,8 +134,8 @@ class MockBackend { if (appInfo.requestedRootEpoch !== undefined && appInfo.requestedRootEpoch > (current.rootEpoch ?? 0)) { return { isAllowed: false, rejectionReason: 'Requested root epoch is not active', gatewayAppId: this.mockGatewayAppId, ...epoch }; } - if (appInfo.requestedAppEpoch !== 0) { - return { isAllowed: false, rejectionReason: 'Requested app epoch must be zero', gatewayAppId: this.mockGatewayAppId, ...epoch }; + if (appInfo.requestedAppEpoch !== undefined && appInfo.requestedAppEpoch > epoch.currentAppEpoch) { + return { isAllowed: false, rejectionReason: 'Requested app epoch is not active', gatewayAppId: this.mockGatewayAppId, ...epoch }; } const result = this.evaluateBoot(appInfo, isKms); return { isAllowed: result.isAllowed, rejectionReason: result.isAllowed ? '' : result.reason, @@ -139,6 +143,9 @@ class MockBackend { } async checkBoot(bootInfo: BootInfo, isKms: boolean): Promise { + if (!isKms && await this.getAppEpoch(bootInfo.appId) !== 0) { + return { isAllowed: false, reason: "App epoch requires epoch-aware authorization", gatewayAppId: "" }; + } const response = this.evaluateBoot(bootInfo, isKms); const epochDenial = response.isAllowed ? rootEpochDenial(bootInfo, isKms) : undefined; return epochDenial ? { isAllowed: false, reason: epochDenial, gatewayAppId: '' } : response; @@ -234,6 +241,16 @@ app.get('/', async (c) => { } }); +app.get('/app/:appId', async (c) => { + const appId = c.req.param('appId'); + if (!/^0x[0-9a-fA-F]{40}$/.test(appId)) return c.json({ error: 'invalid app ID' }, 400); + try { + return c.json({ currentAppEpoch: await mockBackend.getAppEpoch(appId) }); + } catch { + return c.json({ error: 'authorization backend unavailable' }, 503); + } +}); + for (const method of ['authorizeKms', 'authorizeApp'] as const) { app.post(`/${method}`, zValidator('json', AuthorizationRequestSchema), async (c) => { try { diff --git a/dstack/kms/auth-mock/openapi.json b/dstack/kms/auth-mock/openapi.json index 73b1b16c5..ff3505a8e 100644 --- a/dstack/kms/auth-mock/openapi.json +++ b/dstack/kms/auth-mock/openapi.json @@ -160,6 +160,44 @@ } } } + }, + "/app/{appId}": { + "get": { + "summary": "Read the current application key epoch", + "parameters": [ + { + "name": "appId", + "in": "path", + "required": true, + "schema": { + "type": "string", + "pattern": "^0x[0-9a-fA-F]{40}$" + } + } + ], + "responses": { + "200": { + "description": "Current application epoch", + "content": { + "application/json": { + "schema": { + "type": "object", + "required": [ + "currentAppEpoch" + ], + "properties": { + "currentAppEpoch": { + "type": "integer", + "minimum": 0, + "maximum": 4294967295 + } + } + } + } + } + } + } + } } }, "components": { @@ -435,8 +473,7 @@ "type": "integer", "minimum": 0, "maximum": 4294967295, - "default": 0, - "description": "Application epoch, independent of KMS root epochs. Only zero is currently supported." + "description": "Omit to use the current application epoch. Explicit zero selects epoch 0; future epochs are denied." }, "operation": { "type": "string", diff --git a/dstack/kms/auth-simple/README.md b/dstack/kms/auth-simple/README.md index bfa0951b1..3a9b4c8b2 100644 --- a/dstack/kms/auth-simple/README.md +++ b/dstack/kms/auth-simple/README.md @@ -257,10 +257,11 @@ bun run test:run ### Authorization endpoints `POST /authorizeKms` and `POST /authorizeApp` return `AuthorizationResult`, including -`rejectionReason`, `gatewayAppId`, `currentRootEpoch`, and `currentRootEpochHash`. +`rejectionReason`, `gatewayAppId`, `currentRootEpoch`, `currentRootEpochHash`, and `currentAppEpoch`. The selected endpoint determines the policy; the backend checks that the verified -`appId` matches that scope. `requestedAppEpoch` defaults to 0; nonzero values are currently -rejected. The optional +`appId` matches that scope. Optional `requestedAppEpoch` selects the app epoch; +omission uses the current app epoch, explicit zero selects epoch zero, and future +values are denied. The optional `requestedRootEpoch` selects historical keys; omission uses the returned current epoch, and an explicit `0` selects epoch 0. See the [KMS authorization protocol](../README.md#authorization-protocol) for registration, diff --git a/dstack/kms/auth-simple/index.test.ts b/dstack/kms/auth-simple/index.test.ts index 2a08a5877..5772c85a2 100644 --- a/dstack/kms/auth-simple/index.test.ts +++ b/dstack/kms/auth-simple/index.test.ts @@ -486,7 +486,7 @@ describe('auth-simple', () => { }); expect(await authorize({ requestedRootEpoch: 0 })).toMatchObject({ isAllowed: true, currentRootEpoch: 1 }); expect(await authorize({ requestedRootEpoch: 1 })).toMatchObject({ isAllowed: true }); - expect(await authorize({ requestedAppEpoch: 1 })).toMatchObject({ isAllowed: false, rejectionReason: 'Requested app epoch must be zero', currentRootEpoch: 1 }); + expect(await authorize({ requestedAppEpoch: 1 })).toMatchObject({ isAllowed: false, rejectionReason: 'Requested app epoch is not active', currentRootEpoch: 1 }); expect(await authorize({ requestedAppEpoch: 0 })).toMatchObject({ isAllowed: true }); expect(await authorize({ requestedRootEpoch: 2 })).toMatchObject({ isAllowed: false, rejectionReason: 'Requested root epoch is not active', currentRootEpoch: 1, currentRootEpochHash: h1, @@ -498,6 +498,18 @@ describe('auth-simple', () => { expect(await authorize({ appId: '0x' + '22'.repeat(20) })).toMatchObject({ isAllowed: false }); }); + it('rotates only the selected app and allows historical app epochs', async () => { + writeTestConfig({ ...config, apps: { + ...config.apps, + '0xapp123': { ...config.apps['0xapp123'], currentAppEpoch: 2 }, + } }); + expect(await authorize()).toMatchObject({ isAllowed: true, currentAppEpoch: 2, currentRootEpoch: 0 }); + expect(await authorize({ requestedAppEpoch: 0 })).toMatchObject({ isAllowed: true, currentAppEpoch: 2 }); + expect(await authorize({ requestedAppEpoch: 2 })).toMatchObject({ isAllowed: true }); + expect(await authorize({ requestedAppEpoch: 3 })).toMatchObject({ isAllowed: false, rejectionReason: 'Requested app epoch is not active', currentAppEpoch: 2 }); + expect(await boot('app', {})).toMatchObject({ isAllowed: false, reason: 'App epoch requires epoch-aware authorization' }); + }); + it('denies authorization without registered roots, while bootstrap remains available', async () => { writeTestConfig({ ...config, rootEpochs: [] }); expect(await authorize()).toMatchObject({ diff --git a/dstack/kms/auth-simple/index.ts b/dstack/kms/auth-simple/index.ts index 7444a38ab..8c0c108f1 100644 --- a/dstack/kms/auth-simple/index.ts +++ b/dstack/kms/auth-simple/index.ts @@ -29,7 +29,7 @@ const BootInfoSchema = z.object({ const AuthorizationRequestSchema = BootInfoSchema.omit({ rootEpoch: true, rootEpochHash: true }).extend({ operation: z.enum(['kms_bootstrap', 'kms_serve', 'kms_export_keys', 'kms_import_keys', 'app_get_key', 'app_get_cert']), - requestedAppEpoch: z.number().int().min(0).max(0xffffffff).default(0), + requestedAppEpoch: z.number().int().min(0).max(0xffffffff).optional(), requestedRootEpoch: z.number().int().min(0).max(0xffffffff).optional(), }); type AuthorizationRequest = z.infer; @@ -47,6 +47,7 @@ const RootEpochSchema = z.object({ // config schema const AppConfigSchema = z.object({ + currentAppEpoch: z.number().int().min(0).max(0xffffffff).default(0), composeHashes: z.array(z.string()).default([]), devices: z.array(z.string()).default([]), allowAnyDevice: z.boolean().default(false) @@ -159,11 +160,15 @@ class ConfigBackend { } } + async getAppEpoch(appId: string): Promise { + return this.loadConfig().apps[normalizeHex(appId)]?.currentAppEpoch ?? 0; + } + async checkAuthorization(appInfo: AuthorizationRequest, isKms: boolean) { const config = this.loadConfig(); const current = currentRootEpoch(config); const epoch = { - currentAppEpoch: 0, + currentAppEpoch: isKms ? 0 : (config.apps[normalizeHex(appInfo.appId)]?.currentAppEpoch ?? 0), currentRootEpoch: current.rootEpoch ?? 0, currentRootEpochHash: current.rootEpochHash ?? '0x' + '00'.repeat(32), }; @@ -177,8 +182,8 @@ class ConfigBackend { if (appInfo.requestedRootEpoch !== undefined && appInfo.requestedRootEpoch > (current.rootEpoch ?? 0)) { return { isAllowed: false, rejectionReason: 'Requested root epoch is not active', gatewayAppId: config.gatewayAppId, ...epoch }; } - if (appInfo.requestedAppEpoch !== 0) { - return { isAllowed: false, rejectionReason: 'Requested app epoch must be zero', gatewayAppId: config.gatewayAppId, ...epoch }; + if (appInfo.requestedAppEpoch !== undefined && appInfo.requestedAppEpoch > epoch.currentAppEpoch) { + return { isAllowed: false, rejectionReason: 'Requested app epoch is not active', gatewayAppId: config.gatewayAppId, ...epoch }; } const result = this.evaluateBoot(appInfo, isKms, config); return { isAllowed: result.isAllowed, rejectionReason: result.reason, gatewayAppId: result.gatewayAppId, ...epoch }; @@ -186,6 +191,9 @@ class ConfigBackend { async checkBoot(bootInfo: BootInfo, isKms: boolean): Promise { const config = this.loadConfig(); + if (!isKms && (config.apps[normalizeHex(bootInfo.appId)]?.currentAppEpoch ?? 0) !== 0) { + return { isAllowed: false, reason: 'App epoch requires epoch-aware authorization', gatewayAppId: config.gatewayAppId }; + } const response = this.evaluateBoot(bootInfo, isKms, config); const epochDenial = response.isAllowed ? rootEpochDenial(bootInfo, isKms, currentRootEpoch(config)) @@ -364,6 +372,16 @@ app.get('/', async (c) => { } }); +app.get('/app/:appId', async (c) => { + const appId = c.req.param('appId'); + if (!/^0x[0-9a-fA-F]{40}$/.test(appId)) return c.json({ error: 'invalid app ID' }, 400); + try { + return c.json({ currentAppEpoch: await backend.getAppEpoch(appId) }); + } catch { + return c.json({ error: 'authorization backend unavailable' }, 503); + } +}); + for (const method of ['authorizeKms', 'authorizeApp'] as const) { app.post(`/${method}`, zValidator('json', AuthorizationRequestSchema), async (c) => { try { diff --git a/dstack/kms/rpc/proto/kms_rpc.proto b/dstack/kms/rpc/proto/kms_rpc.proto index 79b85419a..8085b2892 100644 --- a/dstack/kms/rpc/proto/kms_rpc.proto +++ b/dstack/kms/rpc/proto/kms_rpc.proto @@ -14,6 +14,7 @@ message GetAppKeyRequest { // Root key epoch to derive from; unset means the current epoch. Epoch 0 is // the original root keys. Older epochs remain available so CVMs can migrate. optional uint32 root_epoch = 3; + optional uint32 app_epoch = 4; } message AppId { @@ -61,6 +62,7 @@ message AppKeyResponse { uint32 root_epoch = 9; // X25519 private key for the SealedEnv HPKE receiver, derived for this epoch. bytes sealed_env_key = 10; + uint32 app_epoch = 11; } message GetMetaResponse { diff --git a/dstack/kms/src/crypto.rs b/dstack/kms/src/crypto.rs index faa8e01f8..93a5c5054 100644 --- a/dstack/kms/src/crypto.rs +++ b/dstack/kms/src/crypto.rs @@ -8,11 +8,21 @@ use k256::ecdsa::SigningKey; use ra_tls::api_v1::sign_recoverable_keccak256; use ra_tls::kdf; +pub(crate) fn app_epoch_context(epoch: u32) -> Vec { + if epoch == 0 { + Vec::new() + } else { + [b"app-epoch-v1:".as_slice(), &epoch.to_be_bytes()].concat() + } +} + pub(crate) fn derive_k256_key( parent_key: &SigningKey, app_id: &[u8], + app_epoch: u32, ) -> Result<(SigningKey, Vec)> { - let context_data = [app_id, b"app-key"]; + let epoch_context = app_epoch_context(app_epoch); + let context_data = [app_id, b"app-key", epoch_context.as_slice()]; let derived_key_bytes: [u8; 32] = kdf::derive_key(&parent_key.to_bytes(), &context_data, 32)? .try_into() .ok() @@ -20,12 +30,21 @@ pub(crate) fn derive_k256_key( let derived_signing_key = SigningKey::from_bytes(&derived_key_bytes.into())?; let pubkey = derived_signing_key.verifying_key(); - let signature = sign_message( - parent_key, - b"dstack-kms-issued", - app_id, - &pubkey.to_sec1_bytes(), - )?; + let signature = if app_epoch == 0 { + sign_message( + parent_key, + b"dstack-kms-issued", + app_id, + &pubkey.to_sec1_bytes(), + )? + } else { + sign_message( + parent_key, + b"dstack-kms-issued-v2", + app_id, + &[app_epoch.to_be_bytes().as_slice(), &pubkey.to_sec1_bytes()].concat(), + )? + }; Ok((derived_signing_key, signature)) } @@ -141,7 +160,7 @@ mod tests { #[test] fn issued_app_key_signatures_match_their_golden_vectors() { let root = SigningKey::from_slice(&[7_u8; 32]).unwrap(); - let (derived, signature) = derive_k256_key(&root, b"app-a").unwrap(); + let (derived, signature) = derive_k256_key(&root, b"app-a", 0).unwrap(); assert_eq!( hex::encode(derived.to_bytes()), diff --git a/dstack/kms/src/main_service.rs b/dstack/kms/src/main_service.rs index 2ef95ebee..1c5a851a4 100644 --- a/dstack/kms/src/main_service.rs +++ b/dstack/kms/src/main_service.rs @@ -38,7 +38,7 @@ use upgrade_authority::{ use crate::{ config::{KmsConfig, RootKeyPaths}, - crypto::{derive_k256_key, sign_message, sign_message_with_timestamp}, + crypto::{app_epoch_context, derive_k256_key, sign_message, sign_message_with_timestamp}, }; pub(crate) mod amd_attest; @@ -305,9 +305,6 @@ impl KmsState { result: &AuthorizationResult, requested: Option, ) -> Result { - if result.current_app_epoch != 0 { - bail!("this KMS does not support nonzero app epochs"); - } let current = if self.config.auth_api.is_dev() { self.keyring.latest() } else { @@ -370,13 +367,13 @@ struct AppKeys { sealed_env_key: Vec, } -fn derive_sealed_env_key(root: &RootKeys, app_id: &[u8]) -> Result<[u8; 32]> { +fn derive_sealed_env_key(root: &RootKeys, app_id: &[u8], app_epoch: u32) -> Result<[u8; 32]> { ensure_app_id_len(app_id)?; let ca_secret = p256::SecretKey::from_pkcs8_der(root.ca.key.serialized_der())?; kdf::derive_key_with_salt( b"dstack-sealed-env-v1", &ca_secret.to_bytes(), - &[app_id], + &[app_id, &app_epoch_context(app_epoch)], 32, ) .map_err(|_| anyhow!("failed to derive sealed env key"))? @@ -384,16 +381,31 @@ fn derive_sealed_env_key(root: &RootKeys, app_id: &[u8]) -> Result<[u8; 32]> { .map_err(|_| anyhow!("invalid sealed env key length")) } -fn derive_app_keys(root: &RootKeys, app_id: &[u8], instance_id: &[u8]) -> Result { - let disk_crypt_key = - kdf::derive_dh_secret(&root.ca.key, &[app_id, instance_id, b"app-disk-crypt-key"]) - .context("Failed to derive app disk key")?; - let env_crypt_key = kdf::derive_dh_secret(&root.ca.key, &[app_id, b"env-encrypt-key"]) - .context("Failed to derive env encrypt key")?; +fn derive_app_keys( + root: &RootKeys, + app_id: &[u8], + instance_id: &[u8], + app_epoch: u32, +) -> Result { + let disk_crypt_key = kdf::derive_dh_secret( + &root.ca.key, + &[ + app_id, + instance_id, + b"app-disk-crypt-key", + &app_epoch_context(app_epoch), + ], + ) + .context("Failed to derive app disk key")?; + let env_crypt_key = kdf::derive_dh_secret( + &root.ca.key, + &[app_id, b"env-encrypt-key", &app_epoch_context(app_epoch)], + ) + .context("Failed to derive env encrypt key")?; let env_crypt_key = x25519_dalek::StaticSecret::from(env_crypt_key).to_bytes(); let (k256_key, k256_signature) = - derive_k256_key(&root.k256, app_id).context("Failed to derive app ecdsa key")?; - let sealed_env_key = derive_sealed_env_key(root, app_id)?.to_vec(); + derive_k256_key(&root.k256, app_id, app_epoch).context("Failed to derive app ecdsa key")?; + let sealed_env_key = derive_sealed_env_key(root, app_id, app_epoch)?.to_vec(); Ok(AppKeys { disk_crypt_key: disk_crypt_key.to_vec(), env_crypt_key: env_crypt_key.to_vec(), @@ -446,10 +458,17 @@ impl RpcHandler { } } +#[derive(Default)] +struct RequestedAppEpochs { + root: Option, + app: Option, +} + struct BootConfig { boot_info: BootInfo, gateway_app_id: String, root_epoch: u32, + app_epoch: u32, } pub(crate) fn build_boot_info_for_attestation( @@ -565,7 +584,7 @@ impl RpcHandler { async fn ensure_app_boot_allowed( &self, vm_config: &str, - requested_root_epoch: Option, + epochs: RequestedAppEpochs, ) -> Result { let att = self.ensure_attested()?; let boot_info = build_boot_info_for_attestation(att, false, vm_config)?; @@ -574,7 +593,7 @@ impl RpcHandler { boot_info, vm_config, Operation::AppGetKey, - requested_root_epoch, + epochs, ) .await } @@ -585,23 +604,26 @@ impl RpcHandler { boot_info: BootInfo, vm_config: &str, operation: Operation, - requested_root_epoch: Option, + epochs: RequestedAppEpochs, ) -> Result { let response = self .state .config .auth_api - .authorize_app(&boot_info, operation, requested_root_epoch) + .authorize_app(&boot_info, operation, epochs.root, epochs.app) .await?; - let root_epoch = self - .state - .authorized_epoch(&response, requested_root_epoch)?; + let root_epoch = self.state.authorized_epoch(&response, epochs.root)?; + let app_epoch = epochs.app.unwrap_or(response.current_app_epoch); + if app_epoch > response.current_app_epoch { + bail!("requested app epoch is not active"); + } self.ensure_attestation_policy(att, vm_config, &boot_info, &response) .await?; Ok(BootConfig { boot_info, gateway_app_id: response.gateway_app_id, root_epoch, + app_epoch, }) } @@ -631,8 +653,9 @@ impl RpcHandler { Ok(()) } - fn derive_app_ca(&self, app_id: &[u8], root: &RootKeys) -> Result { - let context_data = vec![app_id, b"app-ca"]; + fn derive_app_ca(&self, app_id: &[u8], root: &RootKeys, app_epoch: u32) -> Result { + let epoch_context = app_epoch_context(app_epoch); + let context_data = vec![app_id, b"app-ca", epoch_context.as_slice()]; let app_key = kdf::derive_p256_key_pair(&root.ca.key, &context_data) .context("Failed to derive app disk key")?; let req = CertRequest::builder() @@ -660,13 +683,20 @@ impl KmsRpc for RpcHandler { boot_info, gateway_app_id, root_epoch, + app_epoch, } = self - .ensure_app_boot_allowed(&request.vm_config, request.root_epoch) + .ensure_app_boot_allowed( + &request.vm_config, + RequestedAppEpochs { + root: request.root_epoch, + app: request.app_epoch, + }, + ) .await .context("App not allowed")?; ensure_key_release_allowed(&boot_info, (&self.state.config).into())?; let root = self.state.epoch_root(root_epoch)?; - let keys = derive_app_keys(root, &boot_info.app_id, &boot_info.instance_id)?; + let keys = derive_app_keys(root, &boot_info.app_id, &boot_info.instance_id, app_epoch)?; Ok(AppKeyResponse { ca_cert: root.ca.pem_cert.clone(), @@ -679,6 +709,7 @@ impl KmsRpc for RpcHandler { os_image_hash: boot_info.os_image_hash, root_epoch, sealed_env_key: keys.sealed_env_key, + app_epoch, }) } @@ -689,11 +720,18 @@ impl KmsRpc for RpcHandler { .context("KMS self authorization failed")?; let root = self.state.epoch_root(epoch)?; ensure_app_id_len(&request.app_id)?; - let secret = kdf::derive_dh_secret( - &root.ca.key, - &[&request.app_id[..], "env-encrypt-key".as_bytes()], - ) - .context("Failed to derive env encrypt key")?; + let app_epoch = self + .state + .config + .auth_api + .app_epoch(&request.app_id) + .await?; + if app_epoch != 0 { + bail!("epoch-aware applications must use GetSealedEnvPubKey"); + } + let secret = + kdf::derive_dh_secret(&root.ca.key, &[&request.app_id[..], b"env-encrypt-key"]) + .context("Failed to derive env encrypt key")?; let secret = x25519_dalek::StaticSecret::from(secret); let pubkey = x25519_dalek::PublicKey::from(&secret); @@ -734,7 +772,13 @@ impl KmsRpc for RpcHandler { let epoch = self.ensure_self_allowed(Operation::KmsServe).await?; let root = self.state.epoch_root(epoch)?; ensure_app_id_len(&request.app_id)?; - let secret = derive_sealed_env_key(root, &request.app_id)?; + let app_epoch = self + .state + .config + .auth_api + .app_epoch(&request.app_id) + .await?; + let secret = derive_sealed_env_key(root, &request.app_id, app_epoch)?; let public_key = x25519_dalek::PublicKey::from(&x25519_dalek::StaticSecret::from(secret)).to_bytes(); let timestamp = std::time::SystemTime::now() @@ -743,7 +787,7 @@ impl KmsRpc for RpcHandler { .as_secs(); let mut message = Vec::with_capacity(8 + 32 + 8); message.extend_from_slice(&epoch.to_be_bytes()); - message.extend_from_slice(&0u32.to_be_bytes()); + message.extend_from_slice(&app_epoch.to_be_bytes()); message.extend_from_slice(&public_key); message.extend_from_slice(×tamp.to_be_bytes()); let signature = sign_message( @@ -753,9 +797,9 @@ impl KmsRpc for RpcHandler { &message, )?; Ok(SealedEnvPubKeyResponse { - app_epoch: 0, public_key: public_key.to_vec(), root_epoch: epoch, + app_epoch, timestamp, signature, }) @@ -847,12 +891,12 @@ impl KmsRpc for RpcHandler { boot_info, &request.vm_config, Operation::AppGetCert, - None, + RequestedAppEpochs::default(), ) .await?; ensure_key_release_allowed(&app_info.boot_info, (&self.state.config).into())?; let root = self.state.epoch_root(app_info.root_epoch)?; - let app_ca = self.derive_app_ca(&app_info.boot_info.app_id, root)?; + let app_ca = self.derive_app_ca(&app_info.boot_info.app_id, root, app_info.app_epoch)?; // Stamp the identity verified above, not the requester's own claims. // Runtime measurements (`false`) are what certificates have always carried. let cert_app_info = attestation @@ -1588,13 +1632,47 @@ mod tests { let state = kms_state(dir.path(), "http://127.0.0.1:1"); let root = state.epoch_root(0).unwrap(); let app = [1; 20]; - let sealed = derive_sealed_env_key(root, &app).unwrap(); - let app_keys = derive_app_keys(root, &app, &[2; 20]).unwrap(); + let sealed = derive_sealed_env_key(root, &app, 0).unwrap(); + let app_keys = derive_app_keys(root, &app, &[2; 20], 0).unwrap(); assert_eq!(sealed.as_slice(), app_keys.sealed_env_key); assert_ne!(sealed.as_slice(), app_keys.env_crypt_key); assert_ne!(sealed.as_slice(), app_keys.disk_crypt_key); - assert_ne!(sealed, derive_sealed_env_key(root, &[3; 20]).unwrap()); - assert!(derive_sealed_env_key(root, &[1; 19]).is_err()); + assert_ne!(sealed, derive_sealed_env_key(root, &[3; 20], 0).unwrap()); + assert!(derive_sealed_env_key(root, &[1; 19], 0).is_err()); + } + + #[test] + fn app_epoch_derivation_is_separate_and_epoch_zero_is_unchanged() { + let dir = tempfile::tempdir().unwrap(); + let state = kms_state(dir.path(), "http://127.0.0.1:1"); + let root = state.epoch_root(0).unwrap(); + let app = [1; 20]; + let instance = [2; 20]; + let zero = derive_app_keys(root, &app, &instance, 0).unwrap(); + assert_eq!( + zero.disk_crypt_key, + kdf::derive_dh_secret(&root.ca.key, &[&app, &instance, b"app-disk-crypt-key"]).unwrap() + ); + assert_eq!( + zero.env_crypt_key, + x25519_dalek::StaticSecret::from( + kdf::derive_dh_secret(&root.ca.key, &[&app, b"env-encrypt-key"]).unwrap() + ) + .to_bytes() + ); + for epoch in [1, 2, u32::MAX] { + let rotated = derive_app_keys(root, &app, &instance, epoch).unwrap(); + let again = derive_app_keys(root, &app, &instance, epoch).unwrap(); + assert_eq!(rotated.disk_crypt_key, again.disk_crypt_key); + assert_ne!(zero.disk_crypt_key, rotated.disk_crypt_key); + assert_ne!(zero.env_crypt_key, rotated.env_crypt_key); + assert_ne!(zero.sealed_env_key, rotated.sealed_env_key); + assert_ne!(zero.k256_key, rotated.k256_key); + } + assert_ne!( + derive_app_keys(root, &app, &instance, 1).unwrap().k256_key, + derive_app_keys(root, &app, &instance, 2).unwrap().k256_key + ); } #[rocket::async_test] diff --git a/dstack/kms/src/main_service/upgrade_authority.rs b/dstack/kms/src/main_service/upgrade_authority.rs index c265a555b..f464a4e85 100644 --- a/dstack/kms/src/main_service/upgrade_authority.rs +++ b/dstack/kms/src/main_service/upgrade_authority.rs @@ -266,7 +266,7 @@ impl AuthApi { if !operation.is_kms() { bail!("invalid KMS authorization operation"); } - self.check_authorization(app_info, "authorizeKms", operation, requested_root_epoch) + self.check_authorization(app_info, "authorizeKms", operation, requested_root_epoch, None) .await } @@ -275,11 +275,12 @@ impl AuthApi { app_info: &BootInfo, operation: Operation, requested_root_epoch: Option, + requested_app_epoch: Option, ) -> Result { if operation.is_kms() { bail!("invalid app authorization operation"); } - self.check_authorization(app_info, "authorizeApp", operation, requested_root_epoch) + self.check_authorization(app_info, "authorizeApp", operation, requested_root_epoch, requested_app_epoch) .await } @@ -289,6 +290,7 @@ impl AuthApi { method: &str, operation: Operation, requested_root_epoch: Option, + requested_app_epoch: Option, ) -> Result { #[derive(Serialize)] #[serde(rename_all = "camelCase")] @@ -296,7 +298,8 @@ impl AuthApi { #[serde(flatten)] app_info: &'a BootInfo, operation: Operation, - requested_app_epoch: u32, + #[serde(skip_serializing_if = "Option::is_none")] + requested_app_epoch: Option, #[serde(skip_serializing_if = "Option::is_none")] requested_root_epoch: Option, } @@ -317,7 +320,7 @@ impl AuthApi { &AuthorizationRequest { app_info, operation, - requested_app_epoch: 0, + requested_app_epoch, requested_root_epoch, }, webhook.timeout, @@ -327,6 +330,23 @@ impl AuthApi { } } + pub async fn app_epoch(&self, app_id: &[u8]) -> Result { + ensure_app_id_len(app_id)?; + #[derive(Deserialize)] + #[serde(rename_all = "camelCase")] + struct AppEpoch { + current_app_epoch: u32, + } + match self { + AuthApi::Dev { .. } => Ok(0), + AuthApi::Webhook { webhook } => { + let url = url_join(&webhook.url, &format!("app/0x{}", hex::encode(app_id))); + let response: AppEpoch = http_get(&url, webhook.timeout).await?; + Ok(response.current_app_epoch) + } + } + } + #[cfg(test)] pub async fn is_app_allowed( &self, @@ -734,7 +754,7 @@ pub(crate) mod tests { .await .unwrap(); } else { - auth.authorize_app(&boot_info(1), operation, None) + auth.authorize_app(&boot_info(1), operation, None, None) .await .unwrap(); } @@ -773,18 +793,18 @@ pub(crate) mod tests { let auth = webhook(url); let info = boot_info(50); let first = auth - .authorize_app(&info, Operation::AppGetKey, None) + .authorize_app(&info, Operation::AppGetKey, None, None) .await .unwrap(); let second = auth - .authorize_app(&info, Operation::AppGetKey, None) + .authorize_app(&info, Operation::AppGetKey, None, None) .await .unwrap(); assert!(first.is_allowed); assert!(!second.is_allowed); assert_eq!(first.current_root().unwrap().unwrap().epoch, 0); assert_eq!(second.current_root().unwrap().unwrap().epoch, 1); - auth.authorize_app(&info, Operation::AppGetKey, Some(0)) + auth.authorize_app(&info, Operation::AppGetKey, Some(0), Some(0)) .await .unwrap(); auth.authorize_kms(&info, Operation::KmsServe, None) @@ -797,7 +817,6 @@ pub(crate) mod tests { .all(|(path, _)| path == "/authorizeApp")); assert_eq!(requests[3].0, "/authorizeKms"); let mut expected = serde_json::to_value(info).unwrap(); - expected["requestedAppEpoch"] = 0.into(); expected["operation"] = "app_get_key".into(); assert_eq!(requests[0].1, expected); assert_eq!(requests[2].1["requestedAppEpoch"], 0); @@ -811,7 +830,7 @@ pub(crate) mod tests { r#"{"isAllowed":true,"reason":"","gatewayAppId":"gw"}"#, ]); assert!(webhook(url) - .authorize_app(&boot_info(50), Operation::AppGetKey, None) + .authorize_app(&boot_info(50), Operation::AppGetKey, None, None) .await .is_err()); assert_eq!(server.join().unwrap().len(), 1); diff --git a/sdk/go/dstack/client_v1.go b/sdk/go/dstack/client_v1.go index 5b27c29da..f35d64565 100644 --- a/sdk/go/dstack/client_v1.go +++ b/sdk/go/dstack/client_v1.go @@ -44,6 +44,7 @@ type GetKeyV1Response struct { SignatureChain [][]byte // KMS root key epoch the key was derived from. RootEpoch uint32 + AppEpoch uint32 } // Response from a v1 attestation request. @@ -444,7 +445,7 @@ func (c *DstackClientV1) IssueCert(ctx context.Context, options ...IssueCertV1Op // Derivation is flat -- `a/b` is not a child of `a` -- and binds both arguments, // so these keys differ from the v0 keys of the same name. func (c *DstackClientV1) GetKey(ctx context.Context, domain string, algorithm string) (*GetKeyV1Response, error) { - return c.getKey(ctx, domain, algorithm, nil) + return c.getKey(ctx, domain, algorithm, nil, nil) } // GetKeyAtRootEpoch is GetKey under the KMS root key epoch rootEpoch instead of @@ -452,10 +453,15 @@ func (c *DstackClientV1) GetKey(ctx context.Context, domain string, algorithm st // to reach a key from before a KMS root key rotation, for example to migrate // off it. func (c *DstackClientV1) GetKeyAtRootEpoch(ctx context.Context, domain string, algorithm string, rootEpoch uint32) (*GetKeyV1Response, error) { - return c.getKey(ctx, domain, algorithm, &rootEpoch) + return c.getKey(ctx, domain, algorithm, &rootEpoch, nil) } -func (c *DstackClientV1) getKey(ctx context.Context, domain string, algorithm string, rootEpoch *uint32) (*GetKeyV1Response, error) { +// GetKeyAtEpochs selects the two independent key epochs. Nil means current. +func (c *DstackClientV1) GetKeyAtEpochs(ctx context.Context, domain, algorithm string, rootEpoch, appEpoch *uint32) (*GetKeyV1Response, error) { + return c.getKey(ctx, domain, algorithm, rootEpoch, appEpoch) +} + +func (c *DstackClientV1) getKey(ctx context.Context, domain string, algorithm string, rootEpoch *uint32, appEpoch *uint32) (*GetKeyV1Response, error) { if algorithm == "" { return nil, fmt.Errorf("algorithm is required, use `secp256k1` or `ed25519`") } @@ -468,6 +474,9 @@ func (c *DstackClientV1) getKey(ctx context.Context, domain string, algorithm st payload["root_epoch"] = *rootEpoch } + if appEpoch != nil { + payload["app_epoch"] = *appEpoch + } data, err := c.sendRPCRequest(ctx, "/v1/GetKey", payload) if err != nil { return nil, err @@ -482,6 +491,7 @@ func (c *DstackClientV1) getKey(ctx context.Context, domain string, algorithm st PublicKey *string `json:"public_key"` SignatureChain *[]*string `json:"signature_chain"` RootEpoch uint32 `json:"root_epoch"` + AppEpoch uint32 `json:"app_epoch"` } if err := json.Unmarshal(data, &response); err != nil { return nil, err @@ -507,11 +517,15 @@ func (c *DstackClientV1) getKey(ctx context.Context, domain string, algorithm st } } + if (rootEpoch != nil && response.RootEpoch != *rootEpoch) || (appEpoch != nil && response.AppEpoch != *appEpoch) { + return nil, fmt.Errorf("GetKey returned different epochs than requested") + } return &GetKeyV1Response{ Key: key, PublicKey: publicKey, SignatureChain: chain, RootEpoch: response.RootEpoch, + AppEpoch: response.AppEpoch, }, nil } diff --git a/sdk/js/src/client-v1.ts b/sdk/js/src/client-v1.ts index 9d0f94b15..4c39f4902 100644 --- a/sdk/js/src/client-v1.ts +++ b/sdk/js/src/client-v1.ts @@ -70,6 +70,7 @@ export interface GetKeyResponseV1 { signature_chain: Uint8Array[] /** KMS root key epoch the key was derived from. */ root_epoch: number + app_epoch: number } export interface AttestResponseV1 { @@ -335,7 +336,7 @@ export class DstackClientV1 { * @param domain Caller-chosen domain-separation string. May be empty. * @param algorithm Exactly `secp256k1` or `ed25519`. No default, no `k256` alias. */ - async getKey(domain: string, algorithm: string, rootEpoch?: number): Promise { + async getKey(domain: string, algorithm: string, rootEpoch?: number, appEpoch?: number): Promise { // v0 defaulted an empty algorithm to secp256k1, which let a typo hand back a // key of the wrong type under a name the caller thought meant something else. if (!algorithm) { @@ -345,17 +346,22 @@ export class DstackClientV1 { if (rootEpoch !== undefined && (!Number.isInteger(rootEpoch) || rootEpoch < 0 || rootEpoch > 0xffffffff)) { throw new Error('rootEpoch must be a uint32') } - const payload = JSON.stringify( - rootEpoch === undefined ? { domain, algorithm } : { domain, algorithm, root_epoch: rootEpoch }) - const result = await send_rpc_request<{ key: string, public_key: string, signature_chain: string[], root_epoch?: number }>( + if (appEpoch !== undefined && (!Number.isInteger(appEpoch) || appEpoch < 0 || appEpoch > 0xffffffff)) throw new Error('appEpoch must be a uint32') + const payload = JSON.stringify({ domain, algorithm, root_epoch: rootEpoch, app_epoch: appEpoch }) + const result = await send_rpc_request<{ key: string, public_key: string, signature_chain: string[], root_epoch?: number, app_epoch?: number }>( this.endpoint, '/v1/GetKey', payload) throwOnRpcError(result) + if ((rootEpoch !== undefined && (result.root_epoch ?? 0) !== rootEpoch) || + (appEpoch !== undefined && (result.app_epoch ?? 0) !== appEpoch)) { + throw new Error('GetKey returned different epochs than requested') + } return Object.freeze({ key: from_hex(result.key, 'key'), public_key: from_hex(result.public_key, 'public_key'), signature_chain: to_list(result.signature_chain, 'signature_chain', 'error') .map((link, i) => from_hex(link, `signature_chain[${i}]`)), root_epoch: result.root_epoch ?? 0, + app_epoch: result.app_epoch ?? 0, __name__: 'GetKeyResponseV1' as const, }) } diff --git a/sdk/python/src/dstack_sdk/dstack_client_v1.py b/sdk/python/src/dstack_sdk/dstack_client_v1.py index 750b0b942..24cc07f36 100644 --- a/sdk/python/src/dstack_sdk/dstack_client_v1.py +++ b/sdk/python/src/dstack_sdk/dstack_client_v1.py @@ -109,6 +109,7 @@ class GetKeyResponseV1(BaseModel): signature_chain: List[HexBytes] # KMS root key epoch the key was derived from. root_epoch: int = 0 + app_epoch: int = 0 class GpuEvidenceBundleV1(BaseModel): @@ -240,7 +241,7 @@ async def issue_cert( return IssueCertResponseV1(**result) async def get_key( - self, domain: str, algorithm: str, root_epoch: Optional[int] = None + self, domain: str, algorithm: str, root_epoch: Optional[int] = None, app_epoch: Optional[int] = None ) -> GetKeyResponseV1: """Derive an application key from ``(domain, algorithm)``. @@ -262,8 +263,17 @@ async def get_key( if type(root_epoch) is not int or not 0 <= root_epoch <= 0xFFFFFFFF: raise ValueError("root_epoch must be a uint32") data["root_epoch"] = root_epoch + if app_epoch is not None: + if type(app_epoch) is not int or not 0 <= app_epoch <= 0xFFFFFFFF: + raise ValueError("app_epoch must be a uint32") + data["app_epoch"] = app_epoch result = await self._send_rpc_request("GetKey", data) - return GetKeyResponseV1(**result) + response = GetKeyResponseV1(**result) + if (root_epoch is not None and response.root_epoch != root_epoch) or ( + app_epoch is not None and response.app_epoch != app_epoch + ): + raise ValueError("GetKey returned different epochs than requested") + return response async def attest( self, @@ -364,7 +374,7 @@ def issue_cert( @call_async def get_key( - self, domain: str, algorithm: str, root_epoch: Optional[int] = None + self, domain: str, algorithm: str, root_epoch: Optional[int] = None, app_epoch: Optional[int] = None ) -> GetKeyResponseV1: """Derive an application key from ``(domain, algorithm)``.""" raise NotImplementedError diff --git a/sdk/rust/src/dstack_client_v1.rs b/sdk/rust/src/dstack_client_v1.rs index 71f9c5f33..e9d2d39f1 100644 --- a/sdk/rust/src/dstack_client_v1.rs +++ b/sdk/rust/src/dstack_client_v1.rs @@ -95,7 +95,7 @@ impl DstackClientV1 { /// default and no `k256` alias -- a typo is an error rather than a key of /// the wrong type under a name the caller misread. pub async fn get_key(&self, domain: &str, algorithm: &str) -> Result { - self.request_key(domain, algorithm, None).await + self.request_key(domain, algorithm, None, None).await } /// [`Self::get_key`] under the KMS root key epoch `root_epoch` instead of @@ -108,7 +108,20 @@ impl DstackClientV1 { algorithm: &str, root_epoch: u32, ) -> Result { - self.request_key(domain, algorithm, Some(root_epoch)).await + self.request_key(domain, algorithm, Some(root_epoch), None) + .await + } + + /// Select root and application epochs independently; omitted values use this CVM's current keys. + pub async fn get_key_at_epochs( + &self, + domain: &str, + algorithm: &str, + root_epoch: Option, + app_epoch: Option, + ) -> Result { + self.request_key(domain, algorithm, root_epoch, app_epoch) + .await } async fn request_key( @@ -116,6 +129,7 @@ impl DstackClientV1 { domain: &str, algorithm: &str, root_epoch: Option, + app_epoch: Option, ) -> Result { if algorithm.is_empty() { anyhow::bail!("algorithm is required, use `secp256k1` or `ed25519`") @@ -124,8 +138,17 @@ impl DstackClientV1 { if let Some(root_epoch) = root_epoch { data["root_epoch"] = root_epoch.into(); } + if let Some(app_epoch) = app_epoch { + data["app_epoch"] = app_epoch.into(); + } let response = self.send_rpc_request("GetKey", &data).await?; - Ok(serde_json::from_value::(response)?) + let response = serde_json::from_value::(response)?; + if root_epoch.is_some_and(|epoch| epoch != response.root_epoch) + || app_epoch.is_some_and(|epoch| epoch != response.app_epoch) + { + anyhow::bail!("GetKey returned different epochs than requested"); + } + Ok(response) } /// Produce a versioned attestation over `report_data`. diff --git a/sdk/rust/types/src/dstack_v1.rs b/sdk/rust/types/src/dstack_v1.rs index a31a647ec..5a16a5393 100644 --- a/sdk/rust/types/src/dstack_v1.rs +++ b/sdk/rust/types/src/dstack_v1.rs @@ -128,6 +128,8 @@ pub struct GetKeyResponse { /// KMS root key epoch the key was derived from. #[serde(default)] pub root_epoch: u32, + #[serde(default)] + pub app_epoch: u32, } /// The value an attestation binds to: bytes, and only bytes. @@ -348,7 +350,8 @@ mod tests { #[test] fn signature_chain_round_trips_through_hex() { let response = GetKeyResponse { - root_epoch: 0, + root_epoch: 3, + app_epoch: 2, key: vec![0x01; 32], public_key: vec![0x02; 33], signature_chain: vec![vec![0xaa, 0xbb], vec![0xcc]], From 8fc8a7435dd7548ba4dab882ae4b29c5abdb2881 Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Sat, 7 Nov 2026 17:39:20 -0800 Subject: [PATCH 02/11] fix(auth-simple): normalize app IDs when selecting epochs Signed-off-by: Kevin Wang --- dstack/kms/auth-simple/index.test.ts | 16 ++++++++++++++++ dstack/kms/auth-simple/index.ts | 16 +++++++++------- 2 files changed, 25 insertions(+), 7 deletions(-) diff --git a/dstack/kms/auth-simple/index.test.ts b/dstack/kms/auth-simple/index.test.ts index 5772c85a2..914ae90d3 100644 --- a/dstack/kms/auth-simple/index.test.ts +++ b/dstack/kms/auth-simple/index.test.ts @@ -561,6 +561,22 @@ describe('auth-simple', () => { } }); + it.each(['0x' + 'ab'.repeat(20), '0X' + 'AB'.repeat(20), 'Ab'.repeat(20)])( + 'uses the same normalized app policy for epoch lookup and authorization: %s', async (id) => { + const appId = '0x' + 'ab'.repeat(20); + writeTestConfig({ ...config, apps: { [id]: { + composeHashes: [baseBootInfo.composeHash], allowAnyDevice: true, currentAppEpoch: 1, + } } }); + const lookup = await app.fetch(new Request(`http://localhost/app/${appId}`)); + expect((await lookup.json()).currentAppEpoch).toBe(1); + expect(await authorize({ appId })).toMatchObject({ isAllowed: true, currentAppEpoch: 1 }); + expect((await authorize({ appId, requestedAppEpoch: 1 })).isAllowed).toBe(true); + expect((await authorize({ appId, requestedAppEpoch: 0 })).isAllowed).toBe(true); + expect((await authorize({ appId, requestedAppEpoch: 2 })).isAllowed).toBe(false); + expect((await boot('app', { appId })).isAllowed).toBe(false); + } + ); + it('serves only epoch 0 when none is configured', async () => { writeTestConfig({ ...config, rootEpochs: [] }); expect((await info()).rootEpoch).toBeUndefined(); diff --git a/dstack/kms/auth-simple/index.ts b/dstack/kms/auth-simple/index.ts index 8c0c108f1..0206d07e6 100644 --- a/dstack/kms/auth-simple/index.ts +++ b/dstack/kms/auth-simple/index.ts @@ -137,6 +137,11 @@ function normalizeHex(hex: string): string { return hex; } +function findApp(config: AuthConfig, appId: string) { + const normalized = normalizeHex(appId); + return Object.entries(config.apps).find(([id]) => normalizeHex(id) === normalized)?.[1]; +} + // config-based backend class ConfigBackend { private getConfigPath(): string { @@ -161,14 +166,14 @@ class ConfigBackend { } async getAppEpoch(appId: string): Promise { - return this.loadConfig().apps[normalizeHex(appId)]?.currentAppEpoch ?? 0; + return findApp(this.loadConfig(), appId)?.currentAppEpoch ?? 0; } async checkAuthorization(appInfo: AuthorizationRequest, isKms: boolean) { const config = this.loadConfig(); const current = currentRootEpoch(config); const epoch = { - currentAppEpoch: isKms ? 0 : (config.apps[normalizeHex(appInfo.appId)]?.currentAppEpoch ?? 0), + currentAppEpoch: isKms ? 0 : (findApp(config, appInfo.appId)?.currentAppEpoch ?? 0), currentRootEpoch: current.rootEpoch ?? 0, currentRootEpochHash: current.rootEpochHash ?? '0x' + '00'.repeat(32), }; @@ -191,7 +196,7 @@ class ConfigBackend { async checkBoot(bootInfo: BootInfo, isKms: boolean): Promise { const config = this.loadConfig(); - if (!isKms && (config.apps[normalizeHex(bootInfo.appId)]?.currentAppEpoch ?? 0) !== 0) { + if (!isKms && (findApp(config, bootInfo.appId)?.currentAppEpoch ?? 0) !== 0) { return { isAllowed: false, reason: 'App epoch requires epoch-aware authorization', gatewayAppId: config.gatewayAppId }; } const response = this.evaluateBoot(bootInfo, isKms, config); @@ -275,13 +280,10 @@ class ConfigBackend { } private checkAppBoot(bootInfo: BootInfo, config: AuthConfig, deviceId: string): BootResponse { - const appId = normalizeHex(bootInfo.appId); const composeHash = normalizeHex(bootInfo.composeHash); // check app exists - const appConfig = Object.entries(config.apps).find( - ([id]) => normalizeHex(id) === appId - )?.[1]; + const appConfig = findApp(config, bootInfo.appId); if (!appConfig) { return { From 05e004ef3edaaceec4f94b1af580f9a55387a1a2 Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Sat, 7 Nov 2026 17:39:20 -0800 Subject: [PATCH 03/11] fix(auth): support application epochs in the local backend Signed-off-by: Kevin Wang --- docs/onboarding.md | 6 ++ dstack/crates/dstack-auth/src/main.rs | 118 +++++++++++++++++++++++++- 2 files changed, 120 insertions(+), 4 deletions(-) diff --git a/docs/onboarding.md b/docs/onboarding.md index ff948d2e3..2aa0dc588 100644 --- a/docs/onboarding.md +++ b/docs/onboarding.md @@ -329,3 +329,9 @@ registered-root requirement as the other backends. Missing registration denies k release, and invalid epoch configuration returns HTTP 503. Legacy `/bootAuth/*` routes remain available for bootstrap. Local compose/device allowlists apply to all permitted historical roots; they are not the on-chain per-epoch matrices. + +Application entries also accept `currentAppEpoch` (default zero). Keep this value +monotonic, restart the application after advancing it, and retain permissions for +historical keys needed during migration. `/app/:appId` exposes the current app +epoch; legacy app boot authorization is denied once it is nonzero. The Rust +backend uses the same normalized app ID lookup for policy and epoch selection. diff --git a/dstack/crates/dstack-auth/src/main.rs b/dstack/crates/dstack-auth/src/main.rs index a55f21b13..3d2b4ac27 100644 --- a/dstack/crates/dstack-auth/src/main.rs +++ b/dstack/crates/dstack-auth/src/main.rs @@ -173,6 +173,7 @@ struct KmsRules { #[derive(Deserialize, Default)] #[serde(rename_all = "camelCase", default)] struct AppRules { + current_app_epoch: u32, compose_hashes: Vec, devices: Vec, allow_any_device: bool, @@ -281,6 +282,12 @@ fn kms_policy(info: &BootInfo, al: &Allowlist) -> BootResponse { } fn check_app(info: &BootInfo, al: &Allowlist) -> BootResponse { + if al + .app(&info.app_id) + .is_some_and(|app| app.current_app_epoch != 0) + { + return deny(al, "app epoch requires epoch-aware authorization"); + } let result = app_policy(info, al); if result.is_allowed { allow(info, false, al) @@ -313,7 +320,12 @@ fn authorize( || format!("0x{}", "00".repeat(32)), |epoch| epoch.hash.clone(), ), - current_app_epoch: 0, + current_app_epoch: if is_kms { + 0 + } else { + al.app(&request.boot_info.app_id) + .map_or(0, |app| app.current_app_epoch) + }, }; let info = &request.boot_info; result.rejection_reason = if request.operation.is_kms() != is_kms { @@ -325,8 +337,11 @@ fn authorize( .is_some_and(|epoch| epoch as usize > result.current_root_epoch) { "requested root epoch is not active".into() - } else if request.requested_app_epoch.unwrap_or(0) != 0 { - "requested app epoch must be zero".into() + } else if request + .requested_app_epoch + .is_some_and(|epoch| epoch > result.current_app_epoch) + { + "requested app epoch is not active".into() } else if request.operation != Operation::KmsBootstrap && norm(&al.kms_contract_addr).is_empty() { "kmsContractAddr is not configured".into() @@ -371,6 +386,18 @@ fn authorize_kms( .map_err(unavailable) } +#[get("/app/")] +fn app_epoch(app_id: &str, cli: &State) -> Result, Status> { + let id = norm(app_id); + if hex::decode(&id).map_or(true, |bytes| bytes.len() != 20) { + return Err(Status::BadRequest); + } + let al = load(&cli.config); + Ok(Json( + json!({ "currentAppEpoch": al.app(&id).map_or(0, |app| app.current_app_epoch) }), + )) +} + /// load the allowlist, failing closed (deny-all) if it's missing or invalid. fn load(path: &PathBuf) -> Allowlist { match std::fs::read_to_string(path) { @@ -441,7 +468,14 @@ fn server(cli: Cli) -> rocket::Rocket { .merge(("port", cli.port)); rocket::custom(figment).manage(cli).mount( "/", - routes![info, boot_app, boot_kms, authorize_app, authorize_kms], + routes![ + info, + boot_app, + boot_kms, + authorize_app, + authorize_kms, + app_epoch + ], ) } @@ -561,6 +595,82 @@ mod tests { ); } + #[rocket::async_test] + async fn app_epochs_share_normalized_policy_and_reject_legacy_boot() { + use rocket::{http::ContentType, local::asynchronous::Client}; + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("allowlist.json"); + let client = Client::tracked(server(Cli { + config: path.clone(), + address: "127.0.0.1".into(), + port: 0, + })) + .await + .unwrap(); + let id = "ab".repeat(20); + for configured_id in [ + id.clone(), + format!("0X{}", id.to_uppercase()), + format!("0x{id}"), + ] { + let policy = json!({ + "kmsContractAddr": "11".repeat(20), + "rootEpochs": [{"caPubkey": "aa", "k256Pubkey": "bb"}, {"caPubkey": "cc", "k256Pubkey": "dd"}], + "currentRootEpoch": 1, + "apps": {configured_id: {"currentAppEpoch": 2, "composeHashes": ["ee"]}} + }); + std::fs::write(&path, policy.to_string()).unwrap(); + let response: serde_json::Value = client + .get(format!("/app/0x{id}")) + .dispatch() + .await + .into_json() + .await + .unwrap(); + assert_eq!(response["currentAppEpoch"], 2); + for selectors in [ + json!({}), + json!({"requestedRootEpoch": 0, "requestedAppEpoch": 0}), + json!({"requestedRootEpoch": 1, "requestedAppEpoch": 2}), + json!({"requestedAppEpoch": 3}), + json!({"requestedRootEpoch": 2}), + ] { + let allowed = + selectors["requestedAppEpoch"] != 3 && selectors["requestedRootEpoch"] != 2; + let mut request = selectors; + request["appId"] = id.clone().into(); + request["composeHash"] = "ee".into(); + let response: serde_json::Value = client + .post("/authorizeApp") + .header(ContentType::JSON) + .body(request.to_string()) + .dispatch() + .await + .into_json() + .await + .unwrap(); + assert_eq!(response["isAllowed"], allowed); + assert_eq!(response["currentAppEpoch"], 2); + assert_eq!(response["currentRootEpoch"], 1); + } + let body = json!({"appId": id, "composeHash": "ee"}); + let response: serde_json::Value = client + .post("/bootAuth/app") + .header(ContentType::JSON) + .body(body.to_string()) + .dispatch() + .await + .into_json() + .await + .unwrap(); + assert_eq!(response["isAllowed"], false); + } + assert_eq!( + client.get("/app/not-hex").dispatch().await.status(), + Status::BadRequest + ); + } + fn allowlist() -> Allowlist { serde_json::from_str( r#"{ From 04b58ca620b10504ce0c32e0fba950a6f3ce30a2 Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Sat, 7 Nov 2026 18:15:46 -0800 Subject: [PATCH 04/11] refactor(auth-simple): reuse the app policy across authorization checks Signed-off-by: Kevin Wang --- dstack/kms/auth-simple/index.ts | 23 ++++++++++++++--------- 1 file changed, 14 insertions(+), 9 deletions(-) diff --git a/dstack/kms/auth-simple/index.ts b/dstack/kms/auth-simple/index.ts index 0206d07e6..0ad01af9a 100644 --- a/dstack/kms/auth-simple/index.ts +++ b/dstack/kms/auth-simple/index.ts @@ -81,6 +81,7 @@ const AuthConfigSchema = z.object({ type BootInfo = z.infer; type BootResponse = z.infer; type AuthConfig = z.infer; +type AppConfig = z.infer; const bytesOf = (hex: string) => hexToBytes(hex.replace(/^0x/i, '')); @@ -172,8 +173,9 @@ class ConfigBackend { async checkAuthorization(appInfo: AuthorizationRequest, isKms: boolean) { const config = this.loadConfig(); const current = currentRootEpoch(config); + const appConfig = isKms ? undefined : findApp(config, appInfo.appId); const epoch = { - currentAppEpoch: isKms ? 0 : (findApp(config, appInfo.appId)?.currentAppEpoch ?? 0), + currentAppEpoch: appConfig?.currentAppEpoch ?? 0, currentRootEpoch: current.rootEpoch ?? 0, currentRootEpochHash: current.rootEpochHash ?? '0x' + '00'.repeat(32), }; @@ -190,23 +192,26 @@ class ConfigBackend { if (appInfo.requestedAppEpoch !== undefined && appInfo.requestedAppEpoch > epoch.currentAppEpoch) { return { isAllowed: false, rejectionReason: 'Requested app epoch is not active', gatewayAppId: config.gatewayAppId, ...epoch }; } - const result = this.evaluateBoot(appInfo, isKms, config); + const result = this.evaluateBoot(appInfo, isKms, config, appConfig); return { isAllowed: result.isAllowed, rejectionReason: result.reason, gatewayAppId: result.gatewayAppId, ...epoch }; } async checkBoot(bootInfo: BootInfo, isKms: boolean): Promise { const config = this.loadConfig(); - if (!isKms && (findApp(config, bootInfo.appId)?.currentAppEpoch ?? 0) !== 0) { + const appConfig = isKms ? undefined : findApp(config, bootInfo.appId); + if ((appConfig?.currentAppEpoch ?? 0) !== 0) { return { isAllowed: false, reason: 'App epoch requires epoch-aware authorization', gatewayAppId: config.gatewayAppId }; } - const response = this.evaluateBoot(bootInfo, isKms, config); + const response = this.evaluateBoot(bootInfo, isKms, config, appConfig); const epochDenial = response.isAllowed ? rootEpochDenial(bootInfo, isKms, currentRootEpoch(config)) : undefined; return epochDenial ? { ...response, isAllowed: false, reason: epochDenial } : response; } - private evaluateBoot(bootInfo: BootInfo, isKms: boolean, config: AuthConfig): BootResponse { + private evaluateBoot( + bootInfo: BootInfo, isKms: boolean, config: AuthConfig, appConfig: AppConfig | undefined + ): BootResponse { const osImageHash = normalizeHex(bootInfo.osImageHash); const deviceId = normalizeHex(bootInfo.deviceId); @@ -243,7 +248,7 @@ class ConfigBackend { if (isKms) { return this.checkKmsBoot(bootInfo, config, deviceId); } else { - return this.checkAppBoot(bootInfo, config, deviceId); + return this.checkAppBoot(bootInfo, config, deviceId, appConfig); } } @@ -279,12 +284,12 @@ class ConfigBackend { }; } - private checkAppBoot(bootInfo: BootInfo, config: AuthConfig, deviceId: string): BootResponse { + private checkAppBoot( + bootInfo: BootInfo, config: AuthConfig, deviceId: string, appConfig: AppConfig | undefined + ): BootResponse { const composeHash = normalizeHex(bootInfo.composeHash); // check app exists - const appConfig = findApp(config, bootInfo.appId); - if (!appConfig) { return { isAllowed: false, From 5543e7efc6008fcc27c756eea62bb729db2ca6ea Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Sat, 7 Nov 2026 18:15:46 -0800 Subject: [PATCH 05/11] docs(kms): include the application epoch in authorization responses Signed-off-by: Kevin Wang --- dstack/kms/README.md | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/dstack/kms/README.md b/dstack/kms/README.md index 368704d45..e38cc0ea2 100644 --- a/dstack/kms/README.md +++ b/dstack/kms/README.md @@ -259,14 +259,14 @@ KMS callers do not fetch contract metadata to decide which policy to use. `requestedAppEpoch` selects the application's own key epoch, independently of KMS root epochs. Omission uses `currentAppEpoch`; explicit zero selects app epoch zero. -Solidity uses `OptionalEpoch { isSet, value }` for both selectors. Future epochs +Solidity uses `OptionalUint256 { isSet, value }` for both selectors. Future epochs are denied; historical epochs remain subject to current authorization policy. Rejection responses still carry the current epochs and root chain hash. See [application key epochs](../../docs/app-key-epochs.md) for rotation and migration. `requestedRootEpoch` is optional in the JSON request. Omit it to use the epoch in this authorization result; an explicit `0` selects epoch 0. Solidity represents -it as `OptionalEpoch { isSet, value }` alongside `AppInfo` in the request. The +it as `OptionalUint256 { isSet, value }` alongside `AppInfo` in the request. The contract rejects a requested epoch newer than the current one; historical epochs remain available for migration. No serving epoch or serving hash is sent. KMS-to-KMS handover transfers the complete keyring. The receiver explicitly @@ -301,6 +301,7 @@ The response is an `AuthorizationResult`: "rejectionReason": "", "gatewayAppId": "0x...", "currentRootEpoch": 0, + "currentAppEpoch": 0, "currentRootEpochHash": "0x..." } ``` From ae044751f234399b934435621ca126352a92f4a8 Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Sat, 7 Nov 2026 19:04:23 -0800 Subject: [PATCH 06/11] style(sdk): format epoch-aware Python key requests Signed-off-by: Kevin Wang --- sdk/python/src/dstack_sdk/dstack_client_v1.py | 12 ++++++++++-- 1 file changed, 10 insertions(+), 2 deletions(-) diff --git a/sdk/python/src/dstack_sdk/dstack_client_v1.py b/sdk/python/src/dstack_sdk/dstack_client_v1.py index 24cc07f36..29f17fc94 100644 --- a/sdk/python/src/dstack_sdk/dstack_client_v1.py +++ b/sdk/python/src/dstack_sdk/dstack_client_v1.py @@ -241,7 +241,11 @@ async def issue_cert( return IssueCertResponseV1(**result) async def get_key( - self, domain: str, algorithm: str, root_epoch: Optional[int] = None, app_epoch: Optional[int] = None + self, + domain: str, + algorithm: str, + root_epoch: Optional[int] = None, + app_epoch: Optional[int] = None, ) -> GetKeyResponseV1: """Derive an application key from ``(domain, algorithm)``. @@ -374,7 +378,11 @@ def issue_cert( @call_async def get_key( - self, domain: str, algorithm: str, root_epoch: Optional[int] = None, app_epoch: Optional[int] = None + self, + domain: str, + algorithm: str, + root_epoch: Optional[int] = None, + app_epoch: Optional[int] = None, ) -> GetKeyResponseV1: """Derive an application key from ``(domain, algorithm)``.""" raise NotImplementedError From 0dd4fadd57e3b5ea85e952d170e00102e1930016 Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Sun, 8 Nov 2026 01:07:45 -0800 Subject: [PATCH 07/11] fix(cert-client): preserve optional app key epoch selectors Signed-off-by: Kevin Wang --- dstack/Cargo.lock | 1 + dstack/cert-client/Cargo.toml | 1 + dstack/cert-client/src/lib.rs | 162 ++++++++++++++++++++------ dstack/guest-agent/src/rpc_service.rs | 2 +- 4 files changed, 131 insertions(+), 35 deletions(-) diff --git a/dstack/Cargo.lock b/dstack/Cargo.lock index 419c7c44f..39936e3ae 100644 --- a/dstack/Cargo.lock +++ b/dstack/Cargo.lock @@ -868,6 +868,7 @@ dependencies = [ "hex", "http-client", "mock-attestation", + "prpc", "ra-rpc", "ra-tls", "rocket", diff --git a/dstack/cert-client/Cargo.toml b/dstack/cert-client/Cargo.toml index 8a715e379..558389531 100644 --- a/dstack/cert-client/Cargo.toml +++ b/dstack/cert-client/Cargo.toml @@ -23,6 +23,7 @@ tokio.workspace = true x509-parser = { workspace = true, features = ["verify"] } [dev-dependencies] +prpc.workspace = true ra-rpc = { workspace = true, features = ["rocket", "client"] } mock-attestation = { path = "../crates/mock-attestation" } hex.workspace = true diff --git a/dstack/cert-client/src/lib.rs b/dstack/cert-client/src/lib.rs index c95aa3caf..beccbaf55 100644 --- a/dstack/cert-client/src/lib.rs +++ b/dstack/cert-client/src/lib.rs @@ -172,26 +172,28 @@ impl CertRequestClient { } } - /// The app keys of an older KMS root epoch, from the KMS that served the - /// current ones. - pub async fn get_app_key(&self, root_epoch: u32, app_epoch: u32) -> Result { + /// Fetch app keys from the startup KMS, using its current epochs for omitted selectors. + pub async fn get_app_key( + &self, + root_epoch: Option, + app_epoch: Option, + ) -> Result { let CertRequestClient::Kms { client, vm_config } = self else { bail!("only KMS-provided keys have older root epochs"); }; let response = client .get_app_key(GetAppKeyRequest { - app_epoch: Some(app_epoch), + app_epoch, api_version: 1, vm_config: vm_config.clone(), - root_epoch: Some(root_epoch), + root_epoch, }) .await?; - if response.root_epoch != root_epoch || response.app_epoch != app_epoch { - bail!( - "KMS returned key epochs ({}, {}) instead of ({root_epoch}, {app_epoch})", - response.root_epoch, - response.app_epoch - ); + if root_epoch.is_some_and(|epoch| epoch != response.root_epoch) { + bail!("KMS returned a different root epoch than requested"); + } + if app_epoch.is_some_and(|epoch| epoch != response.app_epoch) { + bail!("KMS returned a different app epoch than requested"); } Ok(response) } @@ -469,13 +471,54 @@ mod tests { ) } - #[rocket::post("/GetAppKey", data = "<_request>")] + #[test] + fn generated_epoch_selectors_preserve_omission_and_explicit_zero() { + use prpc::Message; + + let default = GetAppKeyRequest::default(); + assert_eq!((default.root_epoch, default.app_epoch), (None, None)); + for json in ["{}", r#"{"api_version":1,"vm_config":""}"#] { + let request: GetAppKeyRequest = serde_json::from_str(json).unwrap(); + assert_eq!((request.root_epoch, request.app_epoch), (None, None)); + } + for root_epoch in [None, Some(0), Some(7)] { + for app_epoch in [None, Some(0), Some(9)] { + let request = GetAppKeyRequest { + root_epoch, + app_epoch, + ..Default::default() + }; + let json = serde_json::to_vec(&request).unwrap(); + let from_json: GetAppKeyRequest = serde_json::from_slice(&json).unwrap(); + let from_proto = + GetAppKeyRequest::decode(request.encode_to_vec().as_slice()).unwrap(); + for decoded in [from_json, from_proto] { + assert_eq!( + (decoded.root_epoch, decoded.app_epoch), + (root_epoch, app_epoch) + ); + } + } + } + let explicit: GetAppKeyRequest = + serde_json::from_str(r#"{"root_epoch":0,"app_epoch":0}"#).unwrap(); + assert_eq!( + (explicit.root_epoch, explicit.app_epoch), + (Some(0), Some(0)) + ); + } + + #[rocket::post("/GetAppKey", data = "")] fn direct_app_keys( - _request: String, + request: String, _cert: rocket::mtls::Certificate<'_>, - state: &rocket::State<(Arc>, std::sync::atomic::AtomicUsize)>, + state: &rocket::State<(Arc>, Arc>>)>, ) -> rocket::response::content::RawJson { - state.1.fetch_add(1, std::sync::atomic::Ordering::SeqCst); + state + .1 + .lock() + .unwrap() + .push(serde_json::from_str(&request).unwrap()); rocket::response::content::RawJson(state.0.lock().unwrap().clone()) } @@ -549,12 +592,10 @@ mod tests { let (tx, rx) = tokio::sync::oneshot::channel(); let tx = Mutex::new(Some(tx)); let response_body = Arc::new(Mutex::new(serde_json::to_string(&response).unwrap())); + let requests = Arc::new(Mutex::new(Vec::::new())); let server = rocket::custom(figment) .attach(ra_rpc::ratls_client_verifier::RaTlsClientAuth::fairing()) - .manage(( - response_body.clone(), - std::sync::atomic::AtomicUsize::new(0), - )) + .manage((response_body.clone(), requests.clone())) .mount("/", rocket::routes![direct_app_keys]) .attach(rocket::fairing::AdHoc::on_liftoff("port", move |rocket| { if let Some(tx) = tx.lock().unwrap().take() { @@ -616,6 +657,8 @@ mod tests { .unwrap(); assert_eq!(result.ca_cert, response.ca_cert); assert_eq!(result.root_epoch, 1); + let received = requests.lock().unwrap().last().unwrap().clone(); + assert_eq!((received.root_epoch, received.app_epoch), (None, None)); let other_ca = kms_root().pem_cert; for requested_root_epoch in [None, Some(0), Some(1)] { @@ -668,15 +711,73 @@ mod tests { .unwrap(); assert_eq!(result.ca_cert, other_ca); // First-contact validation still requires the response CA to sign the RPC cert. - assert!(request_app_keys( + assert!( + request_app_keys( + url.clone(), + request(), + client_cert(), + verifier.clone(), + None + ) + .await + .is_err() + ); + } + assert_eq!(requests.lock().unwrap().len(), 10); + + let cert = client_cert(); + let client = CertRequestClient::Kms { + client: kms_client( url.clone(), - request(), - client_cert(), + root.pem_cert.clone(), + cert.cert_pem, + cert.key_pem, verifier.clone(), - None ) - .await - .is_err()); + .unwrap(), + vm_config: "selector-test".into(), + }; + for root_epoch in [None, Some(0), Some(1)] { + for app_epoch in [None, Some(0), Some(2)] { + let selected = AppKeyResponse { + root_epoch: root_epoch.unwrap_or(3), + app_epoch: app_epoch.unwrap_or(4), + ..response.clone() + }; + *response_body.lock().unwrap() = serde_json::to_string(&selected).unwrap(); + let result = client.get_app_key(root_epoch, app_epoch).await.unwrap(); + assert_eq!( + (result.root_epoch, result.app_epoch), + (selected.root_epoch, selected.app_epoch) + ); + let received = requests.lock().unwrap().last().unwrap().clone(); + assert_eq!( + (received.root_epoch, received.app_epoch), + (root_epoch, app_epoch) + ); + assert_eq!(received.api_version, 1); + assert_eq!(received.vm_config, "selector-test"); + + for (wrong_root, wrong_app, message) in [ + (selected.root_epoch + 1, selected.app_epoch, "root epoch"), + (selected.root_epoch, selected.app_epoch + 1, "app epoch"), + ] { + let mismatched = AppKeyResponse { + root_epoch: wrong_root, + app_epoch: wrong_app, + ..selected.clone() + }; + *response_body.lock().unwrap() = serde_json::to_string(&mismatched).unwrap(); + let result = client.get_app_key(root_epoch, app_epoch).await; + if (message == "root epoch" && root_epoch.is_some()) + || (message == "app epoch" && app_epoch.is_some()) + { + assert!(result.unwrap_err().to_string().contains(message)); + } else { + assert!(result.is_ok()); + } + } + } } *response_body.lock().unwrap() = serde_json::to_string(&response).unwrap(); let production = Arc::new(AttestationVerifier::new_prod(None).unwrap()); @@ -686,15 +787,8 @@ mod tests { .is_err() ); shutdown.notify(); - let server = task.await.unwrap().unwrap(); - assert_eq!( - server - .state::<(Arc>, std::sync::atomic::AtomicUsize)>() - .unwrap() - .1 - .load(std::sync::atomic::Ordering::SeqCst), - 11 - ); + task.await.unwrap().unwrap(); + assert_eq!(requests.lock().unwrap().len(), 38); } #[test] diff --git a/dstack/guest-agent/src/rpc_service.rs b/dstack/guest-agent/src/rpc_service.rs index db0983c83..ee6e90a26 100644 --- a/dstack/guest-agent/src/rpc_service.rs +++ b/dstack/guest-agent/src/rpc_service.rs @@ -454,7 +454,7 @@ impl AppState { let keys = self .inner .cert_client - .get_app_key(epoch, app_epoch) + .get_app_key(Some(epoch), Some(app_epoch)) .await .with_context(|| format!("failed to fetch the app keys of root epoch {epoch}"))?; let root = AppRoot::new(epoch, app_epoch, keys.k256_key, keys.k256_signature)?; From bfd13af57c04f8dbde14628740d5fa06cd5f1886 Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Sun, 8 Nov 2026 01:08:10 -0800 Subject: [PATCH 08/11] style(guest): group key epoch fields in setup initializers Signed-off-by: Kevin Wang --- dstack/cert-client/src/lib.rs | 6 +++--- dstack/dstack-util/src/system_setup.rs | 4 ++-- 2 files changed, 5 insertions(+), 5 deletions(-) diff --git a/dstack/cert-client/src/lib.rs b/dstack/cert-client/src/lib.rs index beccbaf55..0f6d0eba6 100644 --- a/dstack/cert-client/src/lib.rs +++ b/dstack/cert-client/src/lib.rs @@ -183,10 +183,10 @@ impl CertRequestClient { }; let response = client .get_app_key(GetAppKeyRequest { - app_epoch, api_version: 1, vm_config: vm_config.clone(), root_epoch, + app_epoch, }) .await?; if root_epoch.is_some_and(|epoch| epoch != response.root_epoch) { @@ -641,10 +641,10 @@ mod tests { } }; let request = || GetAppKeyRequest { - app_epoch: None, api_version: 1, vm_config: String::new(), root_epoch: None, + app_epoch: None, }; let result = request_app_keys( url.clone(), @@ -663,10 +663,10 @@ mod tests { let other_ca = kms_root().pem_cert; for requested_root_epoch in [None, Some(0), Some(1)] { let request = || GetAppKeyRequest { - app_epoch: None, api_version: 1, vm_config: String::new(), root_epoch: requested_root_epoch, + app_epoch: None, }; let selected = AppKeyResponse { root_epoch: requested_root_epoch.unwrap_or(1), diff --git a/dstack/dstack-util/src/system_setup.rs b/dstack/dstack-util/src/system_setup.rs index 8a961c9eb..652afcb76 100644 --- a/dstack/dstack-util/src/system_setup.rs +++ b/dstack/dstack-util/src/system_setup.rs @@ -2415,10 +2415,10 @@ impl HostShared { let response = request_app_keys( kms_url.clone(), rpc::GetAppKeyRequest { - app_epoch, api_version: 1, vm_config: self.sys_config.vm_config.clone(), root_epoch, + app_epoch, }, cert_pair, attestation_verifier(&self.sys_config)?, @@ -2443,10 +2443,10 @@ impl HostShared { k256_signature: response.k256_signature, gateway_app_id: response.gateway_app_id, key_provider: KeyProvider::Kms { - app_epoch: response.app_epoch, url: kms_url, pubkey: root_pubkey, root_epoch: served, + app_epoch: response.app_epoch, }, }; Ok((keys, response.os_image_hash)) From 57822fe315229ee66c3a22e4011c379f1f5d3300 Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Sun, 8 Nov 2026 01:02:12 -0800 Subject: [PATCH 09/11] docs(kms): clarify rotation features and disk rollback Signed-off-by: Kevin Wang --- CHANGELOG.md | 6 ++++++ docs/app-key-epochs.md | 9 ++++++--- 2 files changed, 12 insertions(+), 3 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 89a44ccd5..01df5d904 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,6 +7,12 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] +### Added + +- KMS root and application key epochs, epoch-scoped authorization, SealedEnv + HPKE encryption, automatic environment migration, and epoch-aware SDK key + selection. + ### Changed - KMS: retain `GetTempCaCert`, temporary CA files and legacy handover fields so diff --git a/docs/app-key-epochs.md b/docs/app-key-epochs.md index 4e13f0652..d1be0d5c3 100644 --- a/docs/app-key-epochs.md +++ b/docs/app-key-epochs.md @@ -49,9 +49,12 @@ regardless of epoch selection; historical responses cannot replace that pin. The guest migrates disk headers and environment ciphertext if either dimension changes and rejects future epochs in either dimension. LUKS tokens record both epochs (missing app epoch means zero). Rekeying replaces the keyslot wrapping key; -it does not rotate the volume encryption key or erase host snapshots. App -rotation isolates future app keys from an exposed old derived key, but cannot -recover from compromise of the KMS root or revoke already disclosed secrets. +it does not rotate the volume encryption key or erase host snapshots. Once a disk +header carries an epoch token, older guest images reject that header as malformed; +reverting the VM image therefore requires migrating the disk with a compatible +image or restoring a compatible header backup. App rotation isolates future app +keys from an exposed old derived key, but cannot recover from compromise of the +KMS root or revoke already disclosed secrets. ## SDK selection From ddf13defed0cbe16a2a9e82c2536b4ce233e07a1 Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Sat, 7 Nov 2026 18:49:03 -0800 Subject: [PATCH 10/11] test(auth): carry operations through application epoch checks Signed-off-by: Kevin Wang --- dstack/crates/dstack-auth/src/main.rs | 1 + dstack/kms/auth-eth/test/DstackKms.t.sol | 14 ++++++------- dstack/kms/src/main_service.rs | 10 ++-------- .../kms/src/main_service/upgrade_authority.rs | 20 +++++++++++++++---- 4 files changed, 26 insertions(+), 19 deletions(-) diff --git a/dstack/crates/dstack-auth/src/main.rs b/dstack/crates/dstack-auth/src/main.rs index 3d2b4ac27..56534fe2c 100644 --- a/dstack/crates/dstack-auth/src/main.rs +++ b/dstack/crates/dstack-auth/src/main.rs @@ -638,6 +638,7 @@ mod tests { let allowed = selectors["requestedAppEpoch"] != 3 && selectors["requestedRootEpoch"] != 2; let mut request = selectors; + request["operation"] = "app_get_key".into(); request["appId"] = id.clone().into(); request["composeHash"] = "ee".into(); let response: serde_json::Value = client diff --git a/dstack/kms/auth-eth/test/DstackKms.t.sol b/dstack/kms/auth-eth/test/DstackKms.t.sol index 1be41706a..c4d4b8f15 100644 --- a/dstack/kms/auth-eth/test/DstackKms.t.sol +++ b/dstack/kms/auth-eth/test/DstackKms.t.sol @@ -618,21 +618,21 @@ contract DstackKmsTest is Test { (uint256 rootEpoch,) = kms.currentRootEpoch(); assertEq(rootEpoch, 0); DstackKms.AuthorizationRequest memory request = DstackKms.AuthorizationRequest( - info, DstackKms.OptionalEpoch(false, 0), DstackKms.OptionalEpoch(false, 0) + info, keccak256("app_get_key"), DstackKms.OptionalUint256(false, 0), DstackKms.OptionalUint256(false, 0) ); - DstackKms.AuthorizationResult memory result = kms.checkAppAuthorization(request); + DstackKms.AuthorizationResult memory result = kms.authorizeApp(request); assertTrue(result.isAllowed); assertEq(result.currentAppEpoch, 1); - request.requestedAppEpoch = DstackKms.OptionalEpoch(true, 0); - assertTrue(kms.checkAppAuthorization(request).isAllowed); + request.requestedAppEpoch = DstackKms.OptionalUint256(true, 0); + assertTrue(kms.authorizeApp(request).isAllowed); request.requestedAppEpoch.value = 2; - result = kms.checkAppAuthorization(request); + result = kms.authorizeApp(request); assertFalse(result.isAllowed); assertEq(result.currentAppEpoch, 1); assertEq(result.rejectionReason, "Requested app epoch is not active"); - request.requestedAppEpoch = DstackKms.OptionalEpoch(false, 0); + request.requestedAppEpoch = DstackKms.OptionalUint256(false, 0); request.appInfo.composeHash = bytes32(0); - result = kms.checkAppAuthorization(request); + result = kms.authorizeApp(request); assertFalse(result.isAllowed); assertEq(result.currentAppEpoch, 1); vm.prank(user); diff --git a/dstack/kms/src/main_service.rs b/dstack/kms/src/main_service.rs index 1c5a851a4..4358e4ad0 100644 --- a/dstack/kms/src/main_service.rs +++ b/dstack/kms/src/main_service.rs @@ -588,14 +588,8 @@ impl RpcHandler { ) -> Result { let att = self.ensure_attested()?; let boot_info = build_boot_info_for_attestation(att, false, vm_config)?; - self.ensure_app_attestation_allowed( - att, - boot_info, - vm_config, - Operation::AppGetKey, - epochs, - ) - .await + self.ensure_app_attestation_allowed(att, boot_info, vm_config, Operation::AppGetKey, epochs) + .await } async fn ensure_app_attestation_allowed( diff --git a/dstack/kms/src/main_service/upgrade_authority.rs b/dstack/kms/src/main_service/upgrade_authority.rs index f464a4e85..c75dc9f25 100644 --- a/dstack/kms/src/main_service/upgrade_authority.rs +++ b/dstack/kms/src/main_service/upgrade_authority.rs @@ -266,8 +266,14 @@ impl AuthApi { if !operation.is_kms() { bail!("invalid KMS authorization operation"); } - self.check_authorization(app_info, "authorizeKms", operation, requested_root_epoch, None) - .await + self.check_authorization( + app_info, + "authorizeKms", + operation, + requested_root_epoch, + None, + ) + .await } pub async fn authorize_app( @@ -280,8 +286,14 @@ impl AuthApi { if operation.is_kms() { bail!("invalid app authorization operation"); } - self.check_authorization(app_info, "authorizeApp", operation, requested_root_epoch, requested_app_epoch) - .await + self.check_authorization( + app_info, + "authorizeApp", + operation, + requested_root_epoch, + requested_app_epoch, + ) + .await } async fn check_authorization( From 48d151db9d6097bc5de83c3ef0186f4aee68766f Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Sun, 8 Nov 2026 06:12:50 -0800 Subject: [PATCH 11/11] test(sdk): retain paired epoch fixture after app epoch rebase Signed-off-by: Kevin Wang --- dstack/cert-client/src/lib.rs | 20 +++++++++----------- 1 file changed, 9 insertions(+), 11 deletions(-) diff --git a/dstack/cert-client/src/lib.rs b/dstack/cert-client/src/lib.rs index 0f6d0eba6..3082ee9ce 100644 --- a/dstack/cert-client/src/lib.rs +++ b/dstack/cert-client/src/lib.rs @@ -711,17 +711,15 @@ mod tests { .unwrap(); assert_eq!(result.ca_cert, other_ca); // First-contact validation still requires the response CA to sign the RPC cert. - assert!( - request_app_keys( - url.clone(), - request(), - client_cert(), - verifier.clone(), - None - ) - .await - .is_err() - ); + assert!(request_app_keys( + url.clone(), + request(), + client_cert(), + verifier.clone(), + None + ) + .await + .is_err()); } assert_eq!(requests.lock().unwrap().len(), 10);