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 new file mode 100644 index 000000000..d1be0d5c3 --- /dev/null +++ b/docs/app-key-epochs.md @@ -0,0 +1,73 @@ +# 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. 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 + +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/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/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/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 9e222bb04..3082ee9ce 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"))?; @@ -168,9 +172,12 @@ 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 { + /// 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"); }; @@ -178,14 +185,15 @@ impl CertRequestClient { .get_app_key(GetAppKeyRequest { api_version: 1, vm_config: vm_config.clone(), - root_epoch: Some(root_epoch), + root_epoch, + app_epoch, }) .await?; - if response.root_epoch != root_epoch { - bail!( - "KMS returned keys of root epoch {} instead of {root_epoch}", - response.root_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) } @@ -463,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()) } @@ -543,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() { @@ -597,6 +644,7 @@ mod tests { api_version: 1, vm_config: String::new(), root_epoch: None, + app_epoch: None, }; let result = request_app_keys( url.clone(), @@ -609,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)] { @@ -616,6 +666,7 @@ mod tests { 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), @@ -670,6 +721,62 @@ mod tests { .await .is_err()); } + assert_eq!(requests.lock().unwrap().len(), 10); + + let cert = client_cert(); + let client = CertRequestClient::Kms { + client: kms_client( + url.clone(), + root.pem_cert.clone(), + cert.cert_pem, + cert.key_pem, + verifier.clone(), + ) + .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()); assert!( @@ -678,15 +785,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/crates/dstack-auth/src/main.rs b/dstack/crates/dstack-auth/src/main.rs index a55f21b13..56534fe2c 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,83 @@ 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["operation"] = "app_get_key".into(); + 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#"{ 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..652afcb76 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}"); @@ -2415,6 +2418,7 @@ impl HostShared { api_version: 1, vm_config: self.sys_config.vm_config.clone(), root_epoch, + app_epoch, }, cert_pair, attestation_verifier(&self.sys_config)?, @@ -2442,6 +2446,7 @@ impl HostShared { url: kms_url, pubkey: root_pubkey, root_epoch: served, + app_epoch: response.app_epoch, }, }; Ok((keys, response.os_image_hash)) @@ -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..ee6e90a26 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(Some(epoch), Some(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..e38cc0ea2 100644 --- a/dstack/kms/README.md +++ b/dstack/kms/README.md @@ -257,16 +257,16 @@ 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 `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..." } ``` 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..c4d4b8f15 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, keccak256("app_get_key"), DstackKms.OptionalUint256(false, 0), DstackKms.OptionalUint256(false, 0) + ); + DstackKms.AuthorizationResult memory result = kms.authorizeApp(request); + assertTrue(result.isAllowed); + assertEq(result.currentAppEpoch, 1); + request.requestedAppEpoch = DstackKms.OptionalUint256(true, 0); + assertTrue(kms.authorizeApp(request).isAllowed); + request.requestedAppEpoch.value = 2; + result = kms.authorizeApp(request); + assertFalse(result.isAllowed); + assertEq(result.currentAppEpoch, 1); + assertEq(result.rejectionReason, "Requested app epoch is not active"); + request.requestedAppEpoch = DstackKms.OptionalUint256(false, 0); + request.appInfo.composeHash = bytes32(0); + result = kms.authorizeApp(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..914ae90d3 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({ @@ -549,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 7444a38ab..0ad01af9a 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) @@ -80,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, '')); @@ -136,6 +138,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 { @@ -159,11 +166,16 @@ class ConfigBackend { } } + async getAppEpoch(appId: string): Promise { + return findApp(this.loadConfig(), appId)?.currentAppEpoch ?? 0; + } + 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: 0, + currentAppEpoch: appConfig?.currentAppEpoch ?? 0, currentRootEpoch: current.rootEpoch ?? 0, currentRootEpochHash: current.rootEpochHash ?? '0x' + '00'.repeat(32), }; @@ -177,23 +189,29 @@ 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); + 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(); - const response = this.evaluateBoot(bootInfo, isKms, config); + 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, 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); @@ -230,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); } } @@ -266,15 +284,12 @@ class ConfigBackend { }; } - private checkAppBoot(bootInfo: BootInfo, config: AuthConfig, deviceId: string): BootResponse { - const appId = normalizeHex(bootInfo.appId); + private checkAppBoot( + bootInfo: BootInfo, config: AuthConfig, deviceId: string, appConfig: AppConfig | undefined + ): BootResponse { const composeHash = normalizeHex(bootInfo.composeHash); // check app exists - const appConfig = Object.entries(config.apps).find( - ([id]) => normalizeHex(id) === appId - )?.[1]; - if (!appConfig) { return { isAllowed: false, @@ -364,6 +379,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..4358e4ad0 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,18 +584,12 @@ 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)?; - self.ensure_app_attestation_allowed( - att, - boot_info, - vm_config, - Operation::AppGetKey, - requested_root_epoch, - ) - .await + self.ensure_app_attestation_allowed(att, boot_info, vm_config, Operation::AppGetKey, epochs) + .await } async fn ensure_app_attestation_allowed( @@ -585,23 +598,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 +647,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 +677,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 +703,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 +714,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 +766,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 +781,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 +791,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 +885,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 +1626,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..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) - .await + self.check_authorization( + app_info, + "authorizeKms", + operation, + requested_root_epoch, + None, + ) + .await } pub async fn authorize_app( @@ -275,12 +281,19 @@ 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) - .await + self.check_authorization( + app_info, + "authorizeApp", + operation, + requested_root_epoch, + requested_app_epoch, + ) + .await } async fn check_authorization( @@ -289,6 +302,7 @@ impl AuthApi { method: &str, operation: Operation, requested_root_epoch: Option, + requested_app_epoch: Option, ) -> Result { #[derive(Serialize)] #[serde(rename_all = "camelCase")] @@ -296,7 +310,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 +332,7 @@ impl AuthApi { &AuthorizationRequest { app_info, operation, - requested_app_epoch: 0, + requested_app_epoch, requested_root_epoch, }, webhook.timeout, @@ -327,6 +342,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 +766,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 +805,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 +829,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 +842,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..29f17fc94 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,11 @@ 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 +267,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 +378,11 @@ 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]],