From 4b05277ce42db66c23ae70243018ea8b3d63db42 Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Thu, 5 Nov 2026 21:43:56 -0800 Subject: [PATCH 001/166] docs: add KMS emergency key rotation protocol design Describe recovery from KMS root-key exposure using a fleet-wide maintenance window: per-epoch roots registered on DstackKms, a single GetAppKey epoch selector (absent = epoch 0, -1 = latest), LUKS2 keyslot rotation with an epoch token, in-CVM env re-encryption, full-history KMS onboarding with a rotation mode, and a clean gateway reset. Signed-off-by: Kevin Wang --- docs/kms-emergency-key-rotation-protocol.md | 303 ++++++++++++++++++++ 1 file changed, 303 insertions(+) create mode 100644 docs/kms-emergency-key-rotation-protocol.md diff --git a/docs/kms-emergency-key-rotation-protocol.md b/docs/kms-emergency-key-rotation-protocol.md new file mode 100644 index 000000000..afae06160 --- /dev/null +++ b/docs/kms-emergency-key-rotation-protocol.md @@ -0,0 +1,303 @@ +# KMS emergency key rotation protocol + +**Status:** Design draft; not implemented. Proposed fields require protobuf definitions. + +## 1. Scope and assumptions + +- Threat model: KMS root private keys are exposed, while disk data, LUKS headers, + and provider/host access are not. Repair the KMS environment before rotation. +- Stop all application CVMs and gateways; upgrade their OS and restart together. +- Reuse DstackKms governance and its OS allowlist. Assume configured chain RPCs + are trusted; RPC rollback and malicious-host/data anti-rollback are out of scope. +- Keep existing root-public-key pins. Affected deployments are handled separately. +- Rotate LUKS keyslots only; the underlying volume key remains unchanged. +- Every current KMS stores and hands over the complete root history, epochs 0..N. +- KMS stores no individual migration progress. No JWT, live rotation, or gateway + admission-lease protocol is introduced. + +**Storage boundary:** old root + old LUKS header reveals the unchanged volume +key. If ciphertext is later exposed, this can affect post-rotation data too. +Treat historical headers/snapshots as sensitive and inventory/isolate/dispose of +copies operationally. Full reencryption is future work if host/storage exposure +enters the threat model; slot replacement is not equivalent protection. + +## 2. Behavior summary + +Historical GetAppKey requests return complete epoch material and stay available +without a deadline, so late or restored instances can still migrate. Automatic recovery covers +disk keyslots and environment ciphertext with cloud-platform assistance; +application-owned GetKey data and wallets remain application responsibilities. +Onboarding supports ordinary full-history import and rotation import followed +by local generation of the next epoch. Sections 6-7 specify these flows. + +## 3. Roots, policy, and compatibility + +`api_version` selects protocol behavior; `root_epoch` identifies root material. +Epoch 0 is the existing roots, with existing derivation formulas and outputs +unchanged. Every epoch selector is optional; unset means the current epoch. +The name leaves room for a separate per-app epoch later. Generate +independent fresh CA and secp256k1 roots for each new epoch; existing derivation +formulas can remain unchanged. Never reuse an epoch for different roots. + +DstackKms keeps an append-only registry of each epoch's root public keys, an +`evidenceUrl` locating the attestation that proves where they were generated, and +a chain hash `h_n = keccak256(h_{n-1} || keccak256(caPubkey) || keccak256(k256Pubkey))` +(zero `h_{-1}`) committing to every key up to `n`; the position in the chain +fixes `n`, and hashing each key first keeps the preimage fixed-length. The +contract accepts only the encodings the KMS produces: a 91-byte P-256 DER +SubjectPublicKeyInfo with an uncompressed point, and a 33-byte compressed +secp256k1 key, so one key cannot be registered under two encodings. The +evidence stays off chain (stored on chain it would cost about 720 gas per byte) +and needs no hash: its report data binds the keys, so only availability depends +on the URL, which the owner can update (`setRootEpochEvidenceUrl`) when the +evidence moves. Registering an epoch and +making it current are separate: `registerRootEpoch` appends (epoch 0 becomes +current at once) and `activateRootEpoch(n)` moves the current epoch forward one +step. `popRootEpoch(n, chainHash)` drops the last epoch while it is still pending, for when +the nodes holding it are lost before activation; a re-rotation then registers a +new `n`. A current epoch is never popped. `kmsInfo` follows the current epoch: +registering epoch 0 and each activation copy its keys there and clear the legacy +quote and event log, and `setKmsInfo` is refused once epochs exist. Epoch 0 must be the `kmsInfo` keys +when those are set; a legacy `kmsInfo` quote and event log are converted off +chain into the standard attestation encoding (`Attestation::from_tdx_quote`) +before publishing. The KMS sends the epoch it serves and its chain hash (`rootEpoch`, +`rootEpochHash`) in every boot check, and the auth backend owns the policy: it +denies a KMS that does not serve the current epoch with the registered keys. +Info responses report the current epoch for KMS startup and onboarding; boot +responses are unchanged, and so is the on-chain `AppBootInfo`. +Governance credentials must remain independent and uncompromised. + +The KMS keyring and backups contain every epoch the node holds, not just the +latest two. Keeping old roots does not authorize issuing old derived keys. There +is no local epoch configuration: a KMS serves the current epoch on chain, +provided its keyring holds it and reproduces its chain hash, and waits +otherwise. When the current epoch moves, the KMS restarts to follow it and +re-signs its RPC certificate under the new root CA. A `root-epoch` marker records +the newest stored epoch. + +Historical epochs are not time-limited. Their roots are already exposed, so +withholding derived keys from authorized CVMs would not reduce what an attacker +can derive, while it would strand instances or backups that miss a deadline. +Old OS images are cut off through the OS allowlist instead. If host-side disk +rollback enters the threat model, a release deadline can be added later. + +Compatibility: keep accepting GetAppKey versions 0/1. An absent selector means +the current epoch, which before any rotation is epoch 0, so old OS images see no +change. After rotation they cannot open migrated disks; remove old OS images +through the existing allowlist and retain OS-image verification. Historical +epoch access never overrides OS/application authorization. + +## 4. GetAppKey additive extension + +Reuse the existing [GetAppKey transport and authorization path](../dstack/kms/src/main_service.rs), +without claiming a new KMS measurement-allowlist check on the guest. Preserve +root-public-key checking through `key_provider_id`; there is no separate guest +epoch pin or chain query. Unpinned deployments retain their existing KMS identity +limitations: this proposal does not create independent replacement-root trust +for them. Changing a pin can change compose/config hashes and related allowlists +or verifier expectations; affected deployments need separate coordinated updates. + +Keep `api_version = 1`; add one epoch selector: + +```protobuf +message GetAppKeyRequest { + uint32 api_version = 1; + string vm_config = 2; + // Unset = the current epoch, N = epoch N. + optional uint32 root_epoch = 3; +} +// Add to AppKeyResponse; preserve fields 1-8 and their types. +uint32 root_epoch = 9; +``` + +Unset resolves to current once for this request; unknown epochs are rejected. +Response `root_epoch` and all epoch-dependent material in fields 1-8 belong to +the resolved epoch. There is no separate source disk field. Every epoch up to +current uses the same normal authorization. A guest asks for an older epoch by +number only when it needs one, so a boot with nothing to migrate makes a single +request. + +An allowed historical request returns the complete historical AppKeyResponse, +including env and application keys, to preserve legacy compatibility. This +replaces the earlier disk-only server response proposal: the new OS extracts +only disk_crypt_key and, when needed, env_crypt_key for migration, discarding +other historical material. It retains its latest bundle for app startup. This +accepted tradeoff restricts recipients, not cryptographic uses of returned keys. + +The guest checks explicit epoch responses against the requested number. An old +KMS ignores the selector and serves epoch 0, reporting 0; that is +indistinguishable from a not-yet-rotated KMS using these fields alone. +Existing root-public-key pins and removal of old KMS endpoints remain essential; +unpinned clients gain no new downgrade protection. Do not rely on cached/disabled +KMS self-authorization for immediate revocation. Test JSON compatibility separately. + +Authenticate the KMS connection against current configured trust. Historical +ca_cert in the response is historical data, not a replacement transport trust +anchor: do not install it or compare it as though it were the current server pin. + +## 5. OS slot migration and LUKS header metadata + +Store the epoch in one LUKS2 JSON token: + +```json +{"type":"dstack-kms","keyslots":["0"],"root_epoch":2} +``` + +The current `validate_single_luks2_header` rejects tokens, nonzero subsystem, +unexpected labels, and any keyslot layout other than slot 0 alone. The new OS +allows at most this one namespaced token, validates its epoch and slot-0 binding, +and preserves the other checks. Validate both header copies consistently. Keep +the existing initialization label; do not repurpose it or the subsystem field. +Disks without the token are epoch 0. + +```text +GetAppKey(root_epoch unset) -> copy the header to tmpfs -> validate the copy + token_or_0 == current -> open with the copy + token_or_0 < current -> old keys of epoch token_or_0 -> use only disk key + -> luksChangeKey + token import on the copy -> validate it again + -> luksHeaderRestore to the disk -> open with the copy + token_or_0 > current -> refuse + -> continue boot using only latest application material +``` + +Use `luksChangeKey` instead of a deliberate add/remove two-slot transition. A +successful change retains slot 0. Never format on unlock failure. Invalid or +unsupported metadata requires recovery. Migrating the validated in-memory copy +keeps the boot from acting on a header the host swaps after validation, and the +key change and epoch token reach the disk in one header restore, so the token +always names the epoch of the key in slot 0. + +No extra reboot is required. Historical material is confined to the migration +step, not persisted or passed to applications. +Old OS implementations with the existing validator reject a token-bearing disk; +this is an additional defense, not a substitute for the OS allowlist. Because of +that, an epoch-0 disk is never tagged: before any rotation the new OS leaves the +header untouched, so rolling back to an older OS image still works. Disks +formatted at a later epoch are tagged immediately. + +Preliminary test (Linux 6.8.0-1028-intel, cryptsetup 2.7.0, 32 MiB file image, +LUKS2/PBKDF2, no dm mapping): format, change-key, token import, and +test-passphrase succeeded; the completed change retained slot 0, rejected old +credentials, and accepted new ones. +Release blockers: shipped-tool compatibility, interruption recovery across both +headers, and old key-area wiping. Completed-operation tests do not prove atomicity. + +## 6. Environment migration, application keys, and TLS + +First decrypt env using latest. On authentication failure, try historical epochs +from current-1 down to 0, at most once each per boot; no platform metadata is +required. Stop on authorization failure. Corrupt ciphertext fails closed; env epoch is independent of +the disk token. Empty env requires no migration. + +Decrypt/re-encrypt only inside the CVM, using current env public key, existing +format, and fresh randomness. Locally verify replacement decryption. +The guest pushes the result to the VMM with `HostApi.UpdateEncryptedEnv` on +every boot, before any application code runs, sending an empty env when nothing +changed. The host API is reached over the VM's own vsock, so the VMM knows which +VM sent it. The VMM records the env each boot starts with and accepts one update +per boot, only until the boot reports init scripts or containers starting, and +only if the stored env is still the boot env, so an owner's concurrent update is +never overwritten. Application containers and scripts therefore cannot replace +the env. Rocket's per-method limit bounds the request (1 MiB by default). A +missed update is retried on the next boot. Plaintext and private keys never +leave the CVM. + +`GetAppEnvEncryptPubKey` keeps its interface. It always derives and signs with +the current epoch's roots and never serves historical epochs: env public keys +encrypt new data, so serving an older epoch would encrypt new secrets to a leaked +key. The signer public key identifies the epoch; no response field is added. +SDKs/frontends refresh cached env public keys before uploads; old-key uploads +are still migrated at boot, but put new secrets under an exposed key. Prompt owners to replace exposed +external credentials: reencryption alone does not revoke passwords/tokens. + +Guest API v1 `GetKey` takes an optional `root_epoch` (unset = current) and reports the epoch it +derived from; the SDKs expose it as an optional argument. The default is the +current epoch so new keys never come from an exposed root by accident. GetKey +outputs therefore change after a rotation: owners migrate app-encrypted records, +signing trust, wallets, and permissions by reading old keys with an explicit +`root_epoch` before releasing business services. The guest agent fetches an +older epoch's app root from the KMS the first time GetKey names it and keeps it +in memory; old roots are not stored in the keys file. The frozen v0 surface +always uses the current epoch. A per-app epoch, if added later, follows the same convention. SignCert +retains existing checks and signs under current roots; GetMeta is not a trust +anchor. Change CVM/gateway certificate trust in both directions. Public ACME +certificates remain separate. Orchestration gates all business activity, including +outbound workers, until readiness. + +## 7. Full-history KMS onboarding + +Onboarding imports the source's whole keyring, 0..S, and checks it by the +current epoch's chain hash; it never selects a subset. With `rotate`, it then +generates epoch S+1. + +| Mode | Source requirement | Result | +| --- | --- | --- | +| Ordinary join | Holds the current epoch, possibly one pending activation | Keyring 0..S | +| Rotation join (`rotate=true`) | Holds exactly up to the current epoch | Keyring 0..S plus fresh epoch S+1 and attestation evidence | + +Use admin authorization and recipient KMS attestation. Disable public legacy +handover and reject incompatible legacy receivers without reducing the payload. +A rotation also requires the source's certificate to carry its OS image hash. +Deployments that never registered epochs can still join at epoch 0. + +For rotation, generate fresh roots and attest them: report data is +`keccak256("dstack-kms-generated-keys-v2:{N};{ca_pubkey_hex};{k256_pubkey_hex};")` +(epoch 0 keeps the original v1 bootstrap encoding). The owner validates the +evidence, approved code, and binding, then registers N; the cluster keeps +serving N-1. After the other nodes onboard and hold N, the owner activates N and +every KMS restarts onto it. An old-root signature is not authorization. +Competing generators cannot both be registered as N; a losing node's keyring no +longer reproduces the chain hash, so it never serves and must onboard again. + +Every authorization request (app checks and the KMS self check that gates all +RPCs) carries the epoch the KMS serves, so a backend refuses a stale or forked +KMS. On a denial the KMS reads the current epoch from the info endpoint and +restarts if it moved. A KMS predating root epochs sends none; backends treat it +as epoch 0 and deny its app checks once a later epoch is current, so old KMS +software is cut off too. In-flight operations are not cancelled. + +Backups retain full history; closing migration does not delete roots. Replacement +KMS boot cannot depend on its own unavailable service or stopped gateway ingress. +Protect root storage and independent admin access. + +## 8. Cutover runbook + +1. Inventory app/key dependencies, env ciphertext, pins, and historical snapshots. + Prepare trusted OS images, application recovery, and provider-pin adjustments. +2. Stop business writes, CVMs, gateways, and compromised KMS paths. Repair/replace + KMS execution and rotate independently exposed admin credentials. +3. Rotation-onboard one node, register its new epoch, then ordinary-onboard peers. + Authorize new KMS images before removing old OS images; do not rotate again + during the cutover. +4. Boot new OS with business services held by orchestration. Migrate slots and + complete platform env replacement; confirm app readiness. Ingress closure alone does + not stop workers, cron jobs, queues, or outbound database/wallet writes. +5. Start gateways with fresh data directories/KV namespaces and clean WireGuard + peers. Rebuild required non-admission configuration from trusted sources; + never restore old registration databases or admit old sync nodes. +6. Re-register CVMs with new certificates and release approved business + services/ingress. Keep failed instances isolated for operator follow-up. +7. Notify env public-key verifiers (vmm-cli signer whitelists, frontends, SDK + callers) to replace the old k256 signer with the new one, preferably by reading + it from DstackKms. A stale whitelist entry does not break migration, but lets + anyone able to intercept client-to-KMS traffic forge env public keys. + +## 9. Invariants and acceptance + +Authenticate actual secret-bearing connections; old-CA signatures alone do not +replace evidence checks. Do not downgrade on failure or expose secrets in logs, +command lines, swap, app key files, or host-shared storage. Released secrets +cannot be recalled. Historical key possession, +current unlock success, and production readiness are three different facts. + +Test legacy compatibility, OS-allowlist cutoff, selector boundaries and epoch-0 compatibility, non-adjacent epochs, +both onboarding modes, full-history atomic persistence, pending registration +restart, competing roots, live stale-node rejection, and epoch mismatch, slot/token crash recovery, +new-key reboot success, old-key unlock failure against the working header, and +absence of residual old paths and recovery from missing/stale epoch tokens. Confirm the volume key intentionally stays the +same. Test env corruption, independent env epochs, concurrent platform updates, +response-source authentication, bounded epoch search, stale-key uploads, +replacement/reboot recovery, application readiness, +and clean gateway reauthentication/sync. +Reports must state OS/TEE, kernel/cryptsetup, filesystem/layout, builds, topology, +policy, and failure injections. The preliminary test above is limited evidence, not release validation. From dac94a18d732434af062a90336ca0d69d479c572 Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Fri, 6 Nov 2026 04:36:24 -0800 Subject: [PATCH 002/166] kms: register root key epochs on chain and gate auth on the current one DstackKms keeps an append-only list of root epochs (CA and k256 public keys, the URL of the attestation proving their generation, and a keccak256 chain hash over all epochs so far) and a separately activated current epoch. Registration and activation are distinct steps so every KMS instance can import an epoch before it goes live. A pending epoch can be popped when the nodes holding it are lost before it is activated. Epoch 0 must be the kmsInfo keys when those are set. The attestation stays off chain: its report data binds the keys, so it needs no on-chain hash, and storing it would cost about 720 gas per byte. The KMS sends the root epoch it serves and its chain hash as rootEpoch/rootEpochHash in the boot info; every auth backend (auth-eth, auth-eth-bun, auth-simple, auth-mock, dstack-auth) denies a KMS that does not serve the current epoch with the registered keys. A KMS that predates root epochs reports none and is cut off once a later epoch is current. Info responses report the current epoch so a KMS can pick it at startup. dstack-types gains the shared chain hash computation and test vectors. Signed-off-by: Kevin Wang --- dstack/Cargo.lock | 2 + dstack/crates/dstack-auth/Cargo.toml | 2 + dstack/crates/dstack-auth/src/main.rs | 157 +++++++++++++- dstack/dstack-types/src/lib.rs | 1 + dstack/dstack-types/src/root_epoch.rs | 57 +++++ dstack/kms/auth-eth-bun/index.test.ts | 67 ++++++ dstack/kms/auth-eth-bun/index.ts | 84 +++++++- dstack/kms/auth-eth-bun/openapi.json | 16 ++ dstack/kms/auth-eth/contracts/DstackKms.sol | 140 ++++++++++++- dstack/kms/auth-eth/script/Manage.s.sol | 52 +++++ dstack/kms/auth-eth/script/Query.s.sol | 16 ++ dstack/kms/auth-eth/script/README.md | 41 ++++ dstack/kms/auth-eth/src/ethereum.ts | 53 ++++- dstack/kms/auth-eth/src/main.test.ts | 51 +++++ dstack/kms/auth-eth/src/server.ts | 4 + dstack/kms/auth-eth/src/types.ts | 11 + dstack/kms/auth-eth/test/DstackKms.t.sol | 220 ++++++++++++++++++++ dstack/kms/auth-mock/README.md | 1 + dstack/kms/auth-mock/index.test.ts | 28 +++ dstack/kms/auth-mock/index.ts | 46 +++- dstack/kms/auth-mock/openapi.json | 10 + dstack/kms/auth-simple/README.md | 2 + dstack/kms/auth-simple/bun.lock | 3 + dstack/kms/auth-simple/index.test.ts | 50 +++++ dstack/kms/auth-simple/index.ts | 78 ++++++- dstack/kms/auth-simple/package.json | 1 + 26 files changed, 1168 insertions(+), 25 deletions(-) create mode 100644 dstack/dstack-types/src/root_epoch.rs diff --git a/dstack/Cargo.lock b/dstack/Cargo.lock index 877a7a98c..550291383 100644 --- a/dstack/Cargo.lock +++ b/dstack/Cargo.lock @@ -1899,6 +1899,8 @@ version = "0.6.3" dependencies = [ "anyhow", "clap", + "dstack-types", + "hex", "rocket", "serde", "serde_json", diff --git a/dstack/crates/dstack-auth/Cargo.toml b/dstack/crates/dstack-auth/Cargo.toml index d9fb1879b..9a374b988 100644 --- a/dstack/crates/dstack-auth/Cargo.toml +++ b/dstack/crates/dstack-auth/Cargo.toml @@ -14,6 +14,8 @@ path = "src/main.rs" [dependencies] anyhow.workspace = true +dstack-types.workspace = true +hex.workspace = true clap.workspace = true rocket = { workspace = true, features = ["json"] } serde.workspace = true diff --git a/dstack/crates/dstack-auth/src/main.rs b/dstack/crates/dstack-auth/src/main.rs index 2a752e4b0..75dc30bff 100644 --- a/dstack/crates/dstack-auth/src/main.rs +++ b/dstack/crates/dstack-auth/src/main.rs @@ -58,6 +58,10 @@ struct BootInfo { app_id: String, compose_hash: String, device_id: String, + /// the root epoch the asking KMS serves and its chain hash; absent from a + /// KMS that predates root epochs, and from KMS checks made while onboarding. + root_epoch: Option, + root_epoch_hash: Option, } #[derive(Serialize)] @@ -68,6 +72,13 @@ struct BootResponse { reason: String, } +/// the current KMS root epoch and its chain hash, which commits to the root +/// public keys of epochs 0..=number. +struct EpochInfo { + number: usize, + hash: String, +} + #[derive(Deserialize, Default)] #[serde(rename_all = "camelCase", default)] struct Allowlist { @@ -75,6 +86,41 @@ struct Allowlist { gateway_app_id: String, kms: KmsRules, apps: HashMap, + /// root public keys of each KMS root epoch, indexed by epoch; append only. + root_epochs: Vec, + /// the root epoch KMS instances serve; raise it once every node holds it. + current_root_epoch: usize, +} + +#[derive(Deserialize, Default, Clone)] +#[serde(rename_all = "camelCase", default)] +struct RootEpoch { + ca_pubkey: String, + k256_pubkey: String, +} + +impl Allowlist { + /// `None` when no epoch is configured, or the configuration is invalid, + /// which a KMS past epoch 0 treats as fail-closed. + fn current_root_epoch(&self) -> Option { + let epochs = self.root_epochs.get(..=self.current_root_epoch)?; + let keys = epochs + .iter() + .map(|e| { + Some(( + hex::decode(norm(&e.ca_pubkey)).ok()?, + hex::decode(norm(&e.k256_pubkey)).ok()?, + )) + }) + .collect::>>()?; + let hashes = dstack_types::root_epoch::chain_hashes( + keys.iter().map(|(ca, k256)| (&ca[..], &k256[..])), + ); + Some(EpochInfo { + number: self.current_root_epoch, + hash: format!("0x{}", hex::encode(hashes.last()?)), + }) + } } #[derive(Deserialize, Default)] @@ -125,11 +171,36 @@ fn deny(al: &Allowlist, reason: &str) -> BootResponse { } } -fn allow(al: &Allowlist) -> BootResponse { - BootResponse { - is_allowed: true, - gateway_app_id: al.gateway_app_id.clone(), - reason: "ok".to_string(), +/// allow, provided the asking KMS serves the current root epoch with the +/// registered keys up to it. A KMS that predates root epochs reports none and +/// holds only epoch 0; KMS checks made while onboarding carry none. +fn allow(info: &BootInfo, is_kms: bool, al: &Allowlist) -> BootResponse { + let current = al.current_root_epoch(); + let current_epoch = current.as_ref().map_or(0, |c| c.number); + let denial = match (info.root_epoch, ¤t) { + (None, _) if is_kms || current_epoch == 0 => None, + (None, _) => Some(format!( + "root epoch {current_epoch} is current, but the KMS reports none" + )), + (Some(0), None) => None, + (Some(epoch), None) => Some(format!("root epoch {epoch} is not registered")), + (Some(epoch), Some(_)) if epoch != current_epoch => Some(format!( + "the KMS serves root epoch {epoch}, but {current_epoch} is current" + )), + (Some(_), Some(c)) if info.root_epoch_hash.as_deref().map(norm) != Some(norm(&c.hash)) => { + Some(format!( + "the KMS root keys up to epoch {current_epoch} are not the registered ones" + )) + } + _ => None, + }; + match denial { + Some(reason) => deny(al, &reason), + None => BootResponse { + is_allowed: true, + gateway_app_id: al.gateway_app_id.clone(), + reason: "ok".to_string(), + }, } } @@ -152,7 +223,7 @@ fn check_app(info: &BootInfo, al: &Allowlist) -> BootResponse { if !device_ok(app.allow_any_device, &app.devices, &info.device_id) { return deny(al, "device not allowed"); } - allow(al) + allow(info, false, al) } fn check_kms(info: &BootInfo, al: &Allowlist) -> BootResponse { @@ -162,7 +233,7 @@ fn check_kms(info: &BootInfo, al: &Allowlist) -> BootResponse { if !device_ok(al.kms.allow_any_device, &al.kms.devices, &info.device_id) { return deny(al, "device not allowed"); } - allow(al) + allow(info, true, al) } /// load the allowlist, failing closed (deny-all) if it's missing or invalid. @@ -206,15 +277,24 @@ fn boot_kms(info: Json, cli: &State) -> Json { /// info endpoint the KMS GETs to populate its metadata. Single-node: no chain. #[get("/")] -fn info() -> Json { - Json(json!({ +fn info(cli: &State) -> Json { + Json(info_json(&load(&cli.config))) +} + +fn info_json(al: &Allowlist) -> serde_json::Value { + let mut info = json!({ "status": "ok", "kmsContractAddr": "", "ethRpcUrl": "", "gatewayAppId": "", "chainId": 0, - "appImplementation": "" - })) + "appImplementation": "", + }); + if let Some(epoch) = al.current_root_epoch() { + info["rootEpoch"] = epoch.number.into(); + info["rootEpochHash"] = epoch.hash.into(); + } + info } #[rocket::main] @@ -283,6 +363,61 @@ mod tests { assert!(!check_kms(&info, &Allowlist::default()).is_allowed); } + #[test] + fn allows_only_a_kms_serving_the_current_root_epoch() { + // Shared with the DstackKms contract tests. + let h0 = "0x10ab6b6c8e7601901b0f1767985bf40d35bf3129e0866fe12dcc7c8f04a298f2"; + let h1 = "0x50152a82023f2219a31723b476c13bc5f608ab4147946dabdec15d4f3d648ed4"; + let mut al: Allowlist = serde_json::from_str( + r#"{"kms":{"mrAggregated":["0xMR"]}, + "rootEpochs":[{"caPubkey":"0x3059301306072a8648ce3d020106082a8648ce3d03010703420004a0a0a0a0a0a0a0a0a0a0a0a0a0a0a0a0a0a0a0a0a0a0a0a0a0a0a0a0a0a0a0a0a0a0a0a0a0a0a0a0a0a0a0a0a0a0a0a0a0a0a0a0a0a0a0a0a0a0a0a0a0a0a0a0","k256Pubkey":"0x02b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0"}, + {"caPubkey":"0x3059301306072a8648ce3d020106082a8648ce3d03010703420004a1a1a1a1a1a1a1a1a1a1a1a1a1a1a1a1a1a1a1a1a1a1a1a1a1a1a1a1a1a1a1a1a1a1a1a1a1a1a1a1a1a1a1a1a1a1a1a1a1a1a1a1a1a1a1a1a1a1a1a1a1a1a1a1","k256Pubkey":"0x02b1b1b1b1b1b1b1b1b1b1b1b1b1b1b1b1b1b1b1b1b1b1b1b1b1b1b1b1b1b1b1b1"}]}"#, + ) + .unwrap(); + let info = |epoch: Option, hash: &str| BootInfo { + mr_aggregated: "0xMR".into(), + root_epoch: epoch, + root_epoch_hash: Some(hash.into()), + ..Default::default() + }; + assert_eq!(info_json(&al)["rootEpochHash"], h0); + assert!(check_kms(&info(Some(0), h0), &al).is_allowed); + + al.current_root_epoch = 1; + assert_eq!(info_json(&al)["rootEpoch"], 1); + assert_eq!(info_json(&al)["rootEpochHash"], h1); + assert!(check_kms(&info(Some(1), &h1.to_uppercase()), &al).is_allowed); + assert!(!check_kms(&info(Some(0), h0), &al).is_allowed); + assert!(!check_kms(&info(Some(1), h0), &al).is_allowed); + // onboarding checks carry no epoch + assert!(check_kms(&info(None, ""), &al).is_allowed); + + al.current_root_epoch = 2; + assert!(al.current_root_epoch().is_none()); + assert!(info_json(&al).get("rootEpoch").is_none()); + assert!(!check_kms(&info(Some(1), h1), &al).is_allowed); + } + + #[test] + fn a_kms_predating_root_epochs_is_cut_off_once_one_is_current() { + let mut al = allowlist(); + let info = boot("0xapp1", "0xhash", "0ximg"); + assert!(check_app(&info, &al).is_allowed); + al.root_epochs = vec![ + RootEpoch { + ca_pubkey: "0xa0".into(), + k256_pubkey: "0xb0".into(), + }, + RootEpoch { + ca_pubkey: "0xa1".into(), + k256_pubkey: "0xb1".into(), + }, + ]; + assert!(check_app(&info, &al).is_allowed); + al.current_root_epoch = 1; + assert!(!check_app(&info, &al).is_allowed); + } + // wire-contract snapshot: BootInfo as the KMS serializes it (camelCase). // Keep these field names in sync with the kms BootInfo. `#[serde(default)]` // means extra fields are ignored AND a renamed field deserializes to "" — diff --git a/dstack/dstack-types/src/lib.rs b/dstack/dstack-types/src/lib.rs index 452a2ce59..cc1bdf530 100644 --- a/dstack/dstack-types/src/lib.rs +++ b/dstack/dstack-types/src/lib.rs @@ -2541,6 +2541,7 @@ pub struct ImageInfo { } pub mod mr_config; +pub mod root_epoch; pub mod sha256sum; pub mod shared_filenames; pub mod version; diff --git a/dstack/dstack-types/src/root_epoch.rs b/dstack/dstack-types/src/root_epoch.rs new file mode 100644 index 000000000..feb194ce5 --- /dev/null +++ b/dstack/dstack-types/src/root_epoch.rs @@ -0,0 +1,57 @@ +// SPDX-FileCopyrightText: © 2026 Phala Network +// +// SPDX-License-Identifier: Apache-2.0 + +//! KMS root epoch chain hashes, computed as the DstackKms contract does. + +use sha3::{Digest, Keccak256}; + +/// Chain hashes of root epochs 0, 1, ... given each epoch's +/// `(ca_pubkey, k256_pubkey)`. Epoch `n`'s hash is +/// `keccak256(previous || keccak256(ca_pubkey) || keccak256(k256_pubkey))` +/// with a zero previous hash for epoch 0, so it commits to every key up to `n`; +/// the position in the chain fixes the epoch number. +pub fn chain_hashes<'a>(epochs: impl IntoIterator) -> Vec<[u8; 32]> { + let mut previous = [0u8; 32]; + epochs + .into_iter() + .map(|(ca_pubkey, k256_pubkey)| { + previous = Keccak256::new() + .chain_update(previous) + .chain_update(Keccak256::digest(ca_pubkey)) + .chain_update(Keccak256::digest(k256_pubkey)) + .finalize() + .into(); + previous + }) + .collect() +} + +#[cfg(test)] +mod tests { + use super::*; + + /// Well-formed root public keys whose point bytes are all `fill`. + fn test_keys(ca_fill: u8, k256_fill: u8) -> (Vec, Vec) { + let mut ca = hex::decode("3059301306072a8648ce3d020106082a8648ce3d03010703420004").unwrap(); + ca.resize(91, ca_fill); + let mut k256 = vec![0x02]; + k256.resize(33, k256_fill); + (ca, k256) + } + + /// Shared with the DstackKms contract and auth-simple tests. + #[test] + fn matches_the_contract_vectors() { + let epochs = [test_keys(0xa0, 0xb0), test_keys(0xa1, 0xb1)]; + let hashes = chain_hashes(epochs.iter().map(|(ca, k256)| (&ca[..], &k256[..]))); + assert_eq!( + hex::encode(hashes[0]), + "10ab6b6c8e7601901b0f1767985bf40d35bf3129e0866fe12dcc7c8f04a298f2" + ); + assert_eq!( + hex::encode(hashes[1]), + "50152a82023f2219a31723b476c13bc5f608ab4147946dabdec15d4f3d648ed4" + ); + } +} diff --git a/dstack/kms/auth-eth-bun/index.test.ts b/dstack/kms/auth-eth-bun/index.test.ts index 4d8a85f45..5a2645400 100644 --- a/dstack/kms/auth-eth-bun/index.test.ts +++ b/dstack/kms/auth-eth-bun/index.test.ts @@ -230,6 +230,66 @@ describe('API Compatibility Tests', () => { }); }); + describe('root epochs', () => { + const validBootInfo = { + mrAggregated: '0x' + '12'.repeat(32), + osImageHash: '0x' + 'ab'.repeat(32), + appId: '0x' + '11'.repeat(20), + composeHash: '0x' + '22'.repeat(32), + instanceId: '0x' + '33'.repeat(20), + deviceId: '0x' + '44'.repeat(32), + }; + + const hash = '0x' + '11'.repeat(32); + const mockChain = (currentRootEpoch: () => unknown) => + mockReadContract.mockImplementation((params) => { + switch (params.functionName) { + case 'isAppAllowed': return [true, '']; + case 'isKmsAllowed': return [true, '']; + case 'gatewayAppId': return '0xgateway'; + case 'appImplementation': return '0x9876543210987654321098765432109876543210'; + case 'currentRootEpoch': return currentRootEpoch(); + } + }); + const boot = async (path: string, fields: object) => { + const response = await appFetch(new Request(`http://localhost:3001/bootAuth/${path}`, { + method: 'POST', + headers: { 'Content-Type': 'application/json' }, + body: JSON.stringify({ ...validBootInfo, ...fields }), + })); + return response.json(); + }; + + it('allows only a KMS serving the current root epoch', async () => { + mockChain(() => [1n, hash]); + + expect((await boot('app', { rootEpoch: 1, rootEpochHash: hash })).isAllowed).toBe(true); + const stale = await boot('app', { rootEpoch: 0, rootEpochHash: hash }); + expect(stale.isAllowed).toBe(false); + expect(stale.reason).toMatch(/1 is current/); + const forked = await boot('app', { rootEpoch: 1, rootEpochHash: '0x' + '22'.repeat(32) }); + expect(forked.reason).toMatch(/not the registered/); + // A KMS predating root epochs is cut off; onboarding checks carry none. + expect((await boot('app', {})).reason).toMatch(/reports none/); + expect((await boot('kms', {})).isAllowed).toBe(true); + + const info = await appFetch(new Request('http://localhost:3001/')); + expect(await info.json()).toMatchObject({ rootEpoch: 1, rootEpochHash: hash }); + }); + + it('treats a contract without the registry as having no epochs', async () => { + mockChain(() => { + throw Object.assign(new Error('call failed'), { + cause: { name: 'ContractFunctionZeroDataError' }, + }); + }); + + expect((await boot('app', {})).isAllowed).toBe(true); + expect((await boot('app', { rootEpoch: 0, rootEpochHash: hash })).isAllowed).toBe(true); + expect((await boot('app', { rootEpoch: 1, rootEpochHash: hash })).reason).toMatch(/not registered/); + }); + }); + describe('POST /bootAuth/kms', () => { const validBootInfo = { mrAggregated: '0x1234567890123456789012345678901234567890123456789012345678901234', @@ -449,6 +509,7 @@ describe('Authorization freshness and domain binding', () => { return decisions === 1 ? [true, 'initial allow'] : [false, 'policy changed']; } if (params.functionName === 'gatewayAppId') return 'gateway-app'; + if (params.functionName === 'currentRootEpoch') return undefined; throw new Error(`unexpected function ${params.functionName}`); }); @@ -468,6 +529,7 @@ describe('Authorization freshness and domain binding', () => { return [true, 'allowed']; } if (params.functionName === 'gatewayAppId') return 'gateway-app'; + if (params.functionName === 'currentRootEpoch') return undefined; throw new Error(`unexpected function ${params.functionName}`); }); @@ -492,6 +554,7 @@ describe('Authorization freshness and domain binding', () => { mockReadContract.mockImplementation((params) => { if (params.functionName === 'isAppAllowed') return [true, 'recovered']; if (params.functionName === 'gatewayAppId') return 'gateway-app'; + if (params.functionName === 'currentRootEpoch') return undefined; throw new Error(`unexpected function ${params.functionName}`); }); const recovered = await postApp(); @@ -522,6 +585,7 @@ describe('Ethereum finalized snapshot authorization', () => { expect(params.blockNumber).toBe(98n); if (params.functionName === 'isAppAllowed') return [true, 'finalized allow']; if (params.functionName === 'gatewayAppId') return 'gateway-app'; + if (params.functionName === 'currentRootEpoch') return undefined; throw new Error(`unexpected function ${params.functionName}`); }); @@ -541,6 +605,7 @@ describe('Ethereum finalized snapshot authorization', () => { return decisions === 1 ? [true, 'old canonical allow'] : [false, 'new canonical deny']; } if (params.functionName === 'gatewayAppId') return 'gateway-app'; + if (params.functionName === 'currentRootEpoch') return undefined; throw new Error(`unexpected function ${params.functionName}`); }); @@ -569,6 +634,7 @@ describe('Ethereum finalized snapshot authorization', () => { mockReadContract.mockImplementation((params) => { if (params.functionName === 'isAppAllowed') return [true, 'recovered']; if (params.functionName === 'gatewayAppId') return 'gateway-app'; + if (params.functionName === 'currentRootEpoch') return undefined; throw new Error(`unexpected function ${params.functionName}`); }); const recovered = await authorize(); @@ -611,6 +677,7 @@ describe('dstack-kms 0.6 bootAuth payload', () => { return [true, '']; } if (params.functionName === 'gatewayAppId') return 'gateway-app'; + if (params.functionName === 'currentRootEpoch') return undefined; throw new Error(`unexpected function ${params.functionName}`); }); }); diff --git a/dstack/kms/auth-eth-bun/index.ts b/dstack/kms/auth-eth-bun/index.ts index b192d9f0f..d274fb009 100644 --- a/dstack/kms/auth-eth-bun/index.ts +++ b/dstack/kms/auth-eth-bun/index.ts @@ -29,7 +29,11 @@ const BootInfoSchema = z.object({ deviceId: boundedHex(32, 'device ID'), tcbStatus: z.string().max(128).optional().default(''), advisoryIds: z.array(z.string().max(256)).max(128).optional().default([]), - mrSystem: boundedHex(32, 'system MR measurement').optional().default('') + mrSystem: boundedHex(32, 'system MR measurement').optional().default(''), + // The root epoch the asking KMS serves and its chain hash; absent from a KMS + // that predates root epochs, and from KMS checks made while onboarding. + rootEpoch: z.number().int().min(0).optional(), + rootEpochHash: boundedHex(32, 'root epoch chain hash').optional() }); const BootResponseSchema = z.object({ @@ -38,9 +42,35 @@ const BootResponseSchema = z.object({ gatewayAppId: z.string() }); +// The current root epoch and its chain hash, which commits to the root public +// keys of epochs 0..=rootEpoch; both absent when none is registered. +type RootEpochInfo = { rootEpoch?: number; rootEpochHash?: Hex }; + type BootInfo = z.infer; type BootResponse = z.infer; +const normHash = (hash?: string) => hash?.toLowerCase().replace(/^0x/, ''); + +// Only a KMS serving the current root epoch, with the registered keys up to it, +// may authorize. A KMS that predates root epochs reports none and holds only +// epoch 0; KMS checks made while onboarding carry none and are not limited. +function rootEpochDenial(bootInfo: BootInfo, isKms: boolean, current: RootEpochInfo): string | undefined { + const currentEpoch = current.rootEpoch ?? 0; + if (bootInfo.rootEpoch === undefined) { + return isKms || currentEpoch === 0 ? undefined : `root epoch ${currentEpoch} is current, but the KMS reports none`; + } + if (current.rootEpoch === undefined) { + return bootInfo.rootEpoch === 0 ? undefined : `root epoch ${bootInfo.rootEpoch} is not registered`; + } + if (bootInfo.rootEpoch !== currentEpoch) { + return `the KMS serves root epoch ${bootInfo.rootEpoch}, but ${currentEpoch} is current`; + } + if (normHash(bootInfo.rootEpochHash) !== normHash(current.rootEpochHash)) { + return `the KMS root keys up to epoch ${currentEpoch} are not the registered ones`; + } + return undefined; +} + // DstackKms contract ABI (minimal required functions) const DSTACK_KMS_ABI = [ { @@ -108,9 +138,30 @@ const DSTACK_KMS_ABI = [ stateMutability: 'view', inputs: [], outputs: [{ name: '', type: 'address' }] + }, + { + name: 'currentRootEpoch', + type: 'function', + stateMutability: 'view', + inputs: [], + outputs: [ + { name: 'epoch', type: 'uint256' }, + { name: 'chainHash', type: 'bytes32' } + ] } ] as const; +// Contracts deployed before root epochs lack the registry getters, and a +// registry with no epoch reverts. +const isMissingFunction = (error: unknown): boolean => { + for (let cause: any = error; cause; cause = cause.cause) { + if (cause.name === 'ContractFunctionZeroDataError' || cause.name === 'ContractFunctionRevertedError') { + return true; + } + } + return false; +}; + // ethereum backend class class EthereumBackend { private client: ReturnType; @@ -185,7 +236,7 @@ class EthereumBackend { }); } - const [isAllowed, reason] = response; + const [allowed, contractReason] = response; const gatewayAppId = await this.client.readContract({ address: this.kmsContractAddr, abi: DSTACK_KMS_ABI, @@ -193,13 +244,36 @@ class EthereumBackend { blockNumber }); + const epochDenial = allowed + ? rootEpochDenial(bootInfo, isKms, await this.getRootEpoch(blockNumber)) + : undefined; return { - isAllowed, - reason, + isAllowed: allowed && !epochDenial, + reason: epochDenial ?? contractReason, gatewayAppId: gatewayAppId as string, }; } + // The current root epoch and its chain hash; empty when none is registered + // or the contract predates root epochs. + async getRootEpoch(blockNumber?: bigint): Promise { + blockNumber ??= await this.finalizedBlockNumber(); + try { + const current = await this.client.readContract({ + address: this.kmsContractAddr, + abi: DSTACK_KMS_ABI, + functionName: 'currentRootEpoch', + blockNumber + }); + if (!current) return {}; + const [epoch, chainHash] = current; + return { rootEpoch: Number(epoch), rootEpochHash: chainHash }; + } catch (error) { + if (isMissingFunction(error)) return {}; + throw error; + } + } + async getGatewayAppId(): Promise { const result = await this.client.readContract({ address: this.kmsContractAddr, @@ -278,6 +352,7 @@ app.get('/', async (c) => { ethereum.getGatewayAppId(), ethereum.getChainId(), ethereum.getAppImplementation(), + ethereum.getRootEpoch(), ]); console.log('batch', batch); @@ -289,6 +364,7 @@ app.get('/', async (c) => { chainId: batch[1], appAuthImplementation: batch[2], // NOTE: for backward compatibility appImplementation: batch[2], + ...batch[3], }); } catch { console.error('authorization backend health check failed'); diff --git a/dstack/kms/auth-eth-bun/openapi.json b/dstack/kms/auth-eth-bun/openapi.json index e87c03611..a9619eb0a 100644 --- a/dstack/kms/auth-eth-bun/openapi.json +++ b/dstack/kms/auth-eth-bun/openapi.json @@ -151,6 +151,14 @@ "mrSystem": { "type": "string", "description": "MR System measurement (optional for compatibility)" + }, + "rootEpoch": { + "type": "integer", + "description": "KMS root epoch the asking KMS serves; absent from a KMS that predates root epochs and from KMS checks made while onboarding" + }, + "rootEpochHash": { + "type": "string", + "description": "Chain hash committing to the root public keys of epochs 0..=rootEpoch (hex)" } } }, @@ -216,6 +224,14 @@ "appImplementation": { "type": "string", "description": "App implementation address" + }, + "rootEpoch": { + "type": "integer", + "description": "Current KMS root epoch; absent when none is registered" + }, + "rootEpochHash": { + "type": "string", + "description": "Chain hash committing to the root public keys of epochs 0..=rootEpoch (hex)" } } }, diff --git a/dstack/kms/auth-eth/contracts/DstackKms.sol b/dstack/kms/auth-eth/contracts/DstackKms.sol index f79325b89..8330ece1f 100644 --- a/dstack/kms/auth-eth/contracts/DstackKms.sol +++ b/dstack/kms/auth-eth/contracts/DstackKms.sol @@ -44,6 +44,31 @@ contract DstackKms is Initializable, Ownable2StepUpgradeable, UUPSUpgradeable, E // DstackApp implementation contract address for factory deployment address public appImplementation; + // DER SubjectPublicKeyInfo prefix of a P-256 key with an uncompressed point; + // the 64 bytes of the point follow. + bytes27 private constant P256_SPKI_PREFIX = 0x3059301306072a8648ce3d020106082a8648ce3d03010703420004; + + struct RootEpoch { + bytes caPubkey; + bytes k256Pubkey; + // Where to fetch the generating KMS's attestation. Its report data binds + // the epoch and both public keys, so the evidence proves itself: anyone + // can check the keys came from an approved KMS in a TEE. + string evidenceUrl; + // Commits to the public keys of this epoch and every earlier one: + // keccak256(previous chainHash || keccak256(caPubkey) || + // keccak256(k256Pubkey)), with a zero previous hash for epoch 0. + bytes32 chainHash; + } + + // Root keys of each KMS root epoch, indexed by epoch; epoch 0 is the + // original roots. Append-only: a registered epoch is never changed. + RootEpoch[] private _rootEpochs; + + // The epoch KMS instances serve. Registering epoch 0 makes it current; + // later epochs take effect only when activated, one at a time. + uint256 private _currentRootEpoch; + // Events // appId is unindexed for backward compatibility with deployed log indexers. // slither-disable-next-line unindexed-event-address @@ -58,6 +83,12 @@ contract DstackKms is Initializable, Ownable2StepUpgradeable, UUPSUpgradeable, E event GatewayAppIdSet(string gatewayAppId); // slither-disable-next-line unindexed-event-address event AppImplementationSet(address implementation); + event RootEpochRegistered( + uint256 indexed epoch, bytes32 chainHash, bytes caPubkey, bytes k256Pubkey, string evidenceUrl + ); + event RootEpochActivated(uint256 indexed epoch); + event RootEpochPopped(uint256 indexed epoch, bytes32 chainHash); + event RootEpochEvidenceUrlSet(uint256 indexed epoch, string evidenceUrl); event AppDeployedViaFactory(address indexed appId, address indexed deployer); /// @notice Additive audit event for reconstructing authorization policy. /// @dev `value` is the affected bytes32 value, address, or hash of dynamic public data. @@ -108,20 +139,125 @@ contract DstackKms is Initializable, Ownable2StepUpgradeable, UUPSUpgradeable, E emit PolicyChanged(msg.sender, keccak256(bytes(policy)), value, enabled); } - // Function to set KMS information + // Function to set KMS information. Once root epochs are registered, + // kmsInfo follows the current one instead. function setKmsInfo(KmsInfo memory info) external onlyOwner { + require(_rootEpochs.length == 0, "kmsInfo follows the current root epoch"); kmsInfo = info; emit KmsInfoSet(info.k256Pubkey); _emitPolicy("kms-info", keccak256(info.k256Pubkey), true); } + // Register the root public keys of the next KMS root epoch, with where to + // fetch the attestation proving where they were generated. It is not served until + // activated, except epoch 0, which is current from the start and must be + // the kmsInfo keys when those are set. + function registerRootEpoch( + // P-256 DER SubjectPublicKeyInfo, uncompressed (91 bytes) + bytes calldata caPubkey, + // secp256k1 SEC1 compressed (33 bytes) + bytes calldata k256Pubkey, + string calldata evidenceUrl + ) + external + onlyOwner + returns (uint256 epoch) + { + require(caPubkey.length == 91 && bytes27(caPubkey[:27]) == P256_SPKI_PREFIX, "Invalid epoch CA public key"); + require( + k256Pubkey.length == 33 && (k256Pubkey[0] == 0x02 || k256Pubkey[0] == 0x03), "Invalid epoch k256 public key" + ); + require(bytes(evidenceUrl).length > 0, "Empty epoch evidence URL"); + epoch = _rootEpochs.length; + if (epoch == 0 && kmsInfo.k256Pubkey.length > 0) { + require( + keccak256(caPubkey) == keccak256(kmsInfo.caPubkey) + && keccak256(k256Pubkey) == keccak256(kmsInfo.k256Pubkey), + "Epoch 0 keys differ from kmsInfo" + ); + } + bytes32 previous = epoch == 0 ? bytes32(0) : _rootEpochs[epoch - 1].chainHash; + bytes32 chainHash = keccak256(abi.encodePacked(previous, keccak256(caPubkey), keccak256(k256Pubkey))); + _rootEpochs.push(RootEpoch(caPubkey, k256Pubkey, evidenceUrl, chainHash)); + emit RootEpochRegistered(epoch, chainHash, caPubkey, k256Pubkey, evidenceUrl); + _emitPolicy("root-epoch", chainHash, true); + if (epoch == 0) { + _useRootEpoch(0); + } + } + + // Make the next registered root epoch the one KMS instances serve + function activateRootEpoch(uint256 epoch) external onlyOwner { + require(epoch == _currentRootEpoch + 1 && epoch < _rootEpochs.length, "Invalid root epoch activation"); + _currentRootEpoch = epoch; + _useRootEpoch(epoch); + _emitPolicy("root-epoch-activated", bytes32(epoch), true); + } + + // Point kmsInfo at the keys of the now current epoch, for readers of the + // legacy KMS identity. Its quote and event log described the previous keys; + // the epoch's evidence is at its evidenceUrl. + function _useRootEpoch(uint256 epoch) private { + RootEpoch storage keys = _rootEpochs[epoch]; + kmsInfo.caPubkey = keys.caPubkey; + kmsInfo.k256Pubkey = keys.k256Pubkey; + delete kmsInfo.quote; + delete kmsInfo.eventlog; + emit KmsInfoSet(keys.k256Pubkey); + emit RootEpochActivated(epoch); + } + + // Point an epoch at a new copy of its evidence, for when the old location + // goes away. The evidence proves itself, so only its location changes. + function setRootEpochEvidenceUrl(uint256 epoch, string calldata evidenceUrl) external onlyOwner { + require(epoch < _rootEpochs.length, "Unknown root epoch"); + require(bytes(evidenceUrl).length > 0, "Empty epoch evidence URL"); + _rootEpochs[epoch].evidenceUrl = evidenceUrl; + emit RootEpochEvidenceUrlSet(epoch, evidenceUrl); + } + + // Drop the last registered epoch while it is still pending, for when the + // nodes holding it are lost before enough others have onboarded it. The + // epoch is named by number and chain hash, so a stale call cannot drop a + // later registration of the same number. + function popRootEpoch(uint256 epoch, bytes32 chainHash) external onlyOwner { + require(epoch + 1 == _rootEpochs.length && epoch > _currentRootEpoch, "Not the pending root epoch"); + require(_rootEpochs[epoch].chainHash == chainHash, "Root epoch chain hash mismatch"); + _rootEpochs.pop(); + emit RootEpochPopped(epoch, chainHash); + _emitPolicy("root-epoch", chainHash, false); + } + + function rootEpochCount() external view returns (uint256) { + return _rootEpochs.length; + } + + // The root epoch KMS instances serve and its chain hash + function currentRootEpoch() external view returns (uint256 epoch, bytes32 chainHash) { + require(_rootEpochs.length > 0, "No root epoch registered"); + epoch = _currentRootEpoch; + chainHash = _rootEpochs[epoch].chainHash; + } + + function rootEpoch(uint256 epoch) + external + view + returns (bytes memory caPubkey, bytes memory k256Pubkey, string memory evidenceUrl, bytes32 chainHash) + { + require(epoch < _rootEpochs.length, "Unknown root epoch"); + RootEpoch storage keys = _rootEpochs[epoch]; + return (keys.caPubkey, keys.k256Pubkey, keys.evidenceUrl, keys.chainHash); + } + // Function to set KMS quote function setKmsQuote(bytes memory quote) external onlyOwner { + require(_rootEpochs.length == 0, "kmsInfo follows the current root epoch"); kmsInfo.quote = quote; } // Function to set KMS eventlog function setKmsEventlog(bytes memory eventlog) external onlyOwner { + require(_rootEpochs.length == 0, "kmsInfo follows the current root epoch"); kmsInfo.eventlog = eventlog; } @@ -306,7 +442,7 @@ contract DstackKms is Initializable, Ownable2StepUpgradeable, UUPSUpgradeable, E } // Add storage gap for upgradeable contracts - uint256[50] private __gap; + uint256[48] private __gap; } function isContract(address addr) view returns (bool) { diff --git a/dstack/kms/auth-eth/script/Manage.s.sol b/dstack/kms/auth-eth/script/Manage.s.sol index 110612b7e..9d87e1f33 100644 --- a/dstack/kms/auth-eth/script/Manage.s.sol +++ b/dstack/kms/auth-eth/script/Manage.s.sol @@ -109,6 +109,58 @@ contract SetKmsInfo is BaseScript { } } +contract RegisterRootEpoch is BaseScript { + function run() external { + bytes memory caPubkey = vm.envBytes("CA_PUBKEY"); + bytes memory k256Pubkey = vm.envBytes("K256_PUBKEY"); + string memory evidenceUrl = vm.envString("EVIDENCE_URL"); + + vm.startBroadcast(vm.envUint("PRIVATE_KEY")); + uint256 epoch = kms.registerRootEpoch(caPubkey, k256Pubkey, evidenceUrl); + vm.stopBroadcast(); + + console.log("Registered KMS root epoch", epoch); + } +} + +contract ActivateRootEpoch is BaseScript { + function run() external { + uint256 epoch = vm.envUint("ROOT_EPOCH"); + + vm.startBroadcast(vm.envUint("PRIVATE_KEY")); + kms.activateRootEpoch(epoch); + vm.stopBroadcast(); + + console.log("Activated KMS root epoch", epoch); + } +} + +contract SetRootEpochEvidenceUrl is BaseScript { + function run() external { + uint256 epoch = vm.envUint("ROOT_EPOCH"); + string memory evidenceUrl = vm.envString("EVIDENCE_URL"); + + vm.startBroadcast(vm.envUint("PRIVATE_KEY")); + kms.setRootEpochEvidenceUrl(epoch, evidenceUrl); + vm.stopBroadcast(); + + console.log("Updated the evidence URL of KMS root epoch", epoch); + } +} + +contract PopRootEpoch is BaseScript { + function run() external { + uint256 epoch = vm.envUint("ROOT_EPOCH"); + bytes32 chainHash = vm.envBytes32("CHAIN_HASH"); + + vm.startBroadcast(vm.envUint("PRIVATE_KEY")); + kms.popRootEpoch(epoch, chainHash); + vm.stopBroadcast(); + + console.log("Popped pending KMS root epoch", epoch); + } +} + // App Management Scripts contract AddComposeHash is BaseScript { function run() external { diff --git a/dstack/kms/auth-eth/script/Query.s.sol b/dstack/kms/auth-eth/script/Query.s.sol index 4a9a01496..32faedec9 100644 --- a/dstack/kms/auth-eth/script/Query.s.sol +++ b/dstack/kms/auth-eth/script/Query.s.sol @@ -211,6 +211,22 @@ contract GetKmsSettings is BaseQueryScript { console.log(" Quote length:", quote.length); console.log(" Eventlog length:", eventlog.length); + uint256 epochCount = kms.rootEpochCount(); + console.log("\nRegistered root epochs:", epochCount); + if (epochCount > 0) { + (uint256 current,) = kms.currentRootEpoch(); + console.log("Current root epoch:", current); + } + for (uint256 epoch = 0; epoch < epochCount; epoch++) { + (bytes memory epochCa, bytes memory epochK256, string memory evidenceUrl, bytes32 chainHash) = + kms.rootEpoch(epoch); + console.log(" Epoch", epoch); + console.log(" CA Pubkey:", vm.toString(epochCa)); + console.log(" K256 Pubkey:", vm.toString(epochK256)); + console.log(" Evidence URL:", evidenceUrl); + console.log(" Chain hash:", vm.toString(chainHash)); + } + // Check implementation via storage bytes32 implSlot = 0x360894a13ba1a3210667c828492db98dca3e2076cc3735a920a3ca505d382bbc; address implementation = address(uint160(uint256(vm.load(address(kms), implSlot)))); diff --git a/dstack/kms/auth-eth/script/README.md b/dstack/kms/auth-eth/script/README.md index 8ff68a00e..0c706f50a 100644 --- a/dstack/kms/auth-eth/script/README.md +++ b/dstack/kms/auth-eth/script/README.md @@ -53,11 +53,52 @@ forge script script/Manage.s.sol:SetGatewayAppId --broadcast --rpc-url $RPC_URL ``` #### Set KMS Info +Only before root epochs are registered; afterwards `kmsInfo` holds the current +epoch's keys, with no quote or event log. ```bash KMS_CONTRACT_ADDR=0x... K256_PUBKEY=0x... CA_PUBKEY=0x... QUOTE=0x... EVENTLOG=0x... \ forge script script/Manage.s.sol:SetKmsInfo --broadcast --rpc-url $RPC_URL ``` +#### Register a KMS Root Epoch +Appends the next root epoch (the original roots are epoch 0). Epoch 0 takes +effect when registered; later epochs take effect when activated, and each takes +over `kmsInfo` then, clearing its quote and event log. `EVIDENCE_URL` +points to the attestation the KMS returned when it generated the keys (Bootstrap, +or Onboard with `rotate`), published where verifiers can fetch it; verify it +before registering. Registered keys cannot be changed, only a pending epoch popped; +the evidence URL can be updated when the evidence moves. +```bash +KMS_CONTRACT_ADDR=0x... CA_PUBKEY=0x... K256_PUBKEY=0x... EVIDENCE_URL=ipfs://... \ +forge script script/Manage.s.sol:RegisterRootEpoch --broadcast --rpc-url $RPC_URL +``` + +#### Activate a KMS Root Epoch +Makes the next registered root epoch the one KMS instances serve. Activate it +only once every KMS node holds it; nodes still at the previous epoch stop +serving, and nodes holding it restart to switch. Epochs only move forward. +```bash +KMS_CONTRACT_ADDR=0x... ROOT_EPOCH=1 \ +forge script script/Manage.s.sol:ActivateRootEpoch --broadcast --rpc-url $RPC_URL +``` + +#### Update a Root Epoch's Evidence URL +Points an epoch at a new copy of its attestation, e.g. when the old host goes away. +```bash +KMS_CONTRACT_ADDR=0x... ROOT_EPOCH=0 EVIDENCE_URL=ipfs://... \ +forge script script/Manage.s.sol:SetRootEpochEvidenceUrl --broadcast --rpc-url $RPC_URL +``` + +#### Pop a Pending KMS Root Epoch +Drops the last registered epoch while it is not yet current, e.g. when the nodes +holding it are lost before it could be activated. Re-rotate to register a new one. +`CHAIN_HASH` is the epoch's chain hash (see `Query.s.sol`), so only that exact +registration is dropped. +```bash +KMS_CONTRACT_ADDR=0x... ROOT_EPOCH=1 CHAIN_HASH=0x... \ +forge script script/Manage.s.sol:PopRootEpoch --broadcast --rpc-url $RPC_URL +``` + #### Remove KMS Aggregated MR ```bash KMS_CONTRACT_ADDR=0x... MR_AGGREGATED=0x1234... \ diff --git a/dstack/kms/auth-eth/src/ethereum.ts b/dstack/kms/auth-eth/src/ethereum.ts index ebd655140..45ff7283e 100644 --- a/dstack/kms/auth-eth/src/ethereum.ts +++ b/dstack/kms/auth-eth/src/ethereum.ts @@ -3,16 +3,44 @@ // SPDX-License-Identifier: Apache-2.0 import { ethers } from 'ethers'; -import { BootInfo, BootResponse } from './types'; +import { BootInfo, BootResponse, RootEpochInfo } from './types'; // Minimal ABI for DstackKms contract const DSTACK_KMS_ABI = [ "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)", "function isKmsAllowed((address appId,bytes32 composeHash,address instanceId,bytes32 deviceId,bytes32 mrAggregated,bytes32 mrSystem,bytes32 osImageHash,string tcbStatus,string[] advisoryIds) bootInfo) view returns (bool, string)", "function gatewayAppId() view returns (string)", - "function appImplementation() view returns (address)" + "function appImplementation() view returns (address)", + "function currentRootEpoch() view returns (uint256 epoch, bytes32 chainHash)" ]; +const normHash = (hash?: string) => hash?.toLowerCase().replace(/^0x/, ''); + +// Only a KMS serving the current root epoch, with the registered keys up to it, +// may authorize. A KMS that predates root epochs reports none and holds only +// epoch 0; KMS checks made while onboarding carry none and are not limited. +export function rootEpochDenial(bootInfo: BootInfo, isKms: boolean, current: RootEpochInfo): string | undefined { + const currentEpoch = current.rootEpoch ?? 0; + if (bootInfo.rootEpoch === undefined) { + return isKms || currentEpoch === 0 ? undefined : `root epoch ${currentEpoch} is current, but the KMS reports none`; + } + if (current.rootEpoch === undefined) { + return bootInfo.rootEpoch === 0 ? undefined : `root epoch ${bootInfo.rootEpoch} is not registered`; + } + if (bootInfo.rootEpoch !== currentEpoch) { + return `the KMS serves root epoch ${bootInfo.rootEpoch}, but ${currentEpoch} is current`; + } + if (normHash(bootInfo.rootEpochHash) !== normHash(current.rootEpochHash)) { + return `the KMS root keys up to epoch ${currentEpoch} are not the registered ones`; + } + return undefined; +} + +// Contracts deployed before key epochs existed lack the registry getters. +function isMissingFunction(error: unknown): boolean { + return ethers.isError(error, 'CALL_EXCEPTION') || ethers.isError(error, 'BAD_DATA'); +} + export class EthereumBackend { private provider: ethers.JsonRpcProvider; private kmsContract: ethers.Contract; @@ -58,15 +86,30 @@ export class EthereumBackend { } else { response = await this.kmsContract.isAppAllowed(bootInfoStruct, { blockTag }); } - const [isAllowed, reason] = response; + const [allowed, contractReason] = response; const gatewayAppId = await this.kmsContract.gatewayAppId({ blockTag }); + const epochDenial = allowed + ? rootEpochDenial(bootInfo, isKms, await this.getRootEpoch(blockTag)) + : undefined; return { - isAllowed, - reason, + isAllowed: allowed && !epochDenial, + reason: epochDenial ?? contractReason, gatewayAppId, } } + // The current root epoch and its chain hash; empty when none is registered + // or the contract predates root epochs. + async getRootEpoch(blockTag?: number): Promise { + try { + const [epoch, chainHash] = await this.kmsContract.currentRootEpoch({ blockTag }); + return { rootEpoch: Number(epoch), rootEpochHash: chainHash }; + } catch (error) { + if (isMissingFunction(error)) return {}; + throw error; + } + } + async getGatewayAppId(): Promise { return await this.kmsContract.gatewayAppId(); } diff --git a/dstack/kms/auth-eth/src/main.test.ts b/dstack/kms/auth-eth/src/main.test.ts index 842f1e36b..239a43847 100644 --- a/dstack/kms/auth-eth/src/main.test.ts +++ b/dstack/kms/auth-eth/src/main.test.ts @@ -4,11 +4,13 @@ import { FastifyInstance } from 'fastify'; import { build } from './server'; +import { rootEpochDenial } from './ethereum'; import { BootInfo } from './types'; // Mock EthereumBackend jest.mock('./ethereum', () => { return { + ...jest.requireActual('./ethereum'), EthereumBackend: jest.fn().mockImplementation(() => ({ checkBoot: jest.fn() })) @@ -88,6 +90,55 @@ describe('Server', () => { }); }); + describe('key epochs', () => { + it('should pass the KMS root epoch through to the backend', async () => { + const mockCheckBoot = jest.fn().mockResolvedValue({ isAllowed: true, reason: '', gatewayAppId: '' }); + app.ethereum.checkBoot = mockCheckBoot; + const rootEpoch = { rootEpoch: 1, rootEpochHash: '0x' + '11'.repeat(32) }; + + const response = await app.inject({ + method: 'POST', + url: '/bootAuth/app', + payload: { ...mockBootInfo, ...rootEpoch } + }); + + expect(response.statusCode).toBe(200); + expect(mockCheckBoot).toHaveBeenCalledWith(expect.objectContaining(rootEpoch), false); + }); + + it('should allow only a KMS serving the current root epoch', () => { + const hash = '0x' + '11'.repeat(32); + const current = { rootEpoch: 1, rootEpochHash: hash }; + const info = (fields: object) => ({ ...mockBootInfo, ...fields }); + + expect(rootEpochDenial(info({ rootEpoch: 1, rootEpochHash: hash.toUpperCase().replace('0X', '') }), false, current)).toBeUndefined(); + expect(rootEpochDenial(info({ rootEpoch: 0, rootEpochHash: hash }), false, current)).toMatch(/1 is current/); + expect(rootEpochDenial(info({ rootEpoch: 1, rootEpochHash: '0x' + '22'.repeat(32) }), false, current)).toMatch(/not the registered/); + // A KMS predating root epochs is cut off once a rotation is current. + expect(rootEpochDenial(info({}), false, current)).toMatch(/reports none/); + expect(rootEpochDenial(info({}), false, { rootEpoch: 0, rootEpochHash: hash })).toBeUndefined(); + expect(rootEpochDenial(info({}), true, current)).toBeUndefined(); + // Nothing registered: only epoch 0 serves. + expect(rootEpochDenial(info({ rootEpoch: 0, rootEpochHash: hash }), false, {})).toBeUndefined(); + expect(rootEpochDenial(info({ rootEpoch: 1, rootEpochHash: hash }), false, {})).toMatch(/not registered/); + }); + + it('should list registered key epochs in the info response', async () => { + const rootEpoch = { rootEpoch: 1, rootEpochHash: '0x' + '11'.repeat(32) }; + Object.assign(app.ethereum, { + getGatewayAppId: jest.fn().mockResolvedValue('gw'), + getChainId: jest.fn().mockResolvedValue(1), + getAppImplementation: jest.fn().mockResolvedValue('0x00'), + getRootEpoch: jest.fn().mockResolvedValue(rootEpoch), + }); + + const response = await app.inject({ method: 'GET', url: '/' }); + + expect(response.statusCode).toBe(200); + expect(JSON.parse(response.payload)).toMatchObject(rootEpoch); + }); + }); + describe('POST /bootAuth/kms', () => { it('should return 200 when backend check passes', async () => { // Mock successful response diff --git a/dstack/kms/auth-eth/src/server.ts b/dstack/kms/auth-eth/src/server.ts index 54dd14346..33c80d1d5 100644 --- a/dstack/kms/auth-eth/src/server.ts +++ b/dstack/kms/auth-eth/src/server.ts @@ -39,6 +39,8 @@ export async function build(): Promise { tcbStatus: { type: 'string', maxLength: 128, default: '' }, advisoryIds: { type: 'array', maxItems: 128, items: { type: 'string', maxLength: 256 }, default: [] }, mrSystem: { ...hex(32, 'System MR measurement'), default: '' }, + rootEpoch: { type: 'integer', minimum: 0 }, + rootEpochHash: hex(32, 'Chain hash of the root epochs up to rootEpoch'), } }); @@ -84,6 +86,7 @@ export async function build(): Promise { server.ethereum.getGatewayAppId(), server.ethereum.getChainId(), server.ethereum.getAppImplementation(), + server.ethereum.getRootEpoch(), ]); return { status: 'ok', @@ -93,6 +96,7 @@ export async function build(): Promise { chainId: batch[1], appAuthImplementation: batch[2], // NOTE: for backward compatibility appImplementation: batch[2], + ...batch[3], }; } catch { _request.log.error('authorization backend health check failed'); diff --git a/dstack/kms/auth-eth/src/types.ts b/dstack/kms/auth-eth/src/types.ts index fd1a89143..c92edcfe0 100644 --- a/dstack/kms/auth-eth/src/types.ts +++ b/dstack/kms/auth-eth/src/types.ts @@ -12,6 +12,10 @@ export interface BootInfo { composeHash: string; instanceId: string; deviceId: string; + // The root epoch the asking KMS serves and its chain hash; absent from a KMS + // that predates root epochs, and from KMS checks made while onboarding. + rootEpoch?: number; + rootEpochHash?: string; } export interface BootResponse { @@ -20,4 +24,11 @@ export interface BootResponse { reason: string; } +// The current KMS root epoch and its chain hash, which commits to the root +// public keys of epochs 0..=rootEpoch. Both absent when none is registered. +export interface RootEpochInfo { + rootEpoch?: number; + rootEpochHash?: string; +} + // Removed KMS_CONTRACT_ABI and APP_CONTRACT_ABI since we're using typechain types now diff --git a/dstack/kms/auth-eth/test/DstackKms.t.sol b/dstack/kms/auth-eth/test/DstackKms.t.sol index 899b2b7b9..847a9153f 100644 --- a/dstack/kms/auth-eth/test/DstackKms.t.sol +++ b/dstack/kms/auth-eth/test/DstackKms.t.sol @@ -19,6 +19,12 @@ contract DstackKmsTest is Test { address public user; event KmsInfoSet(bytes k256Pubkey); + event RootEpochRegistered( + uint256 indexed epoch, bytes32 chainHash, bytes caPubkey, bytes k256Pubkey, string evidenceUrl + ); + event RootEpochActivated(uint256 indexed epoch); + event RootEpochPopped(uint256 indexed epoch, bytes32 chainHash); + event RootEpochEvidenceUrlSet(uint256 indexed epoch, string evidenceUrl); event AppRegistered(address appId); @@ -66,6 +72,220 @@ contract DstackKmsTest is Test { assertEq(storedEventlog, eventlog); } + // Well-formed root public keys whose point bytes are all `fill`; shared with + // the off-chain implementations' tests. + function ca(uint8 fill) internal pure returns (bytes memory key) { + key = bytes.concat(hex"3059301306072a8648ce3d020106082a8648ce3d03010703420004", new bytes(64)); + for (uint256 i = 27; i < 91; i++) { + key[i] = bytes1(fill); + } + } + + function k256(uint8 fill) internal pure returns (bytes memory key) { + key = new bytes(33); + key[0] = 0x02; + for (uint256 i = 1; i < 33; i++) { + key[i] = bytes1(fill); + } + } + + function chainHash(bytes32 previous, bytes memory caKey, bytes memory k256Key) internal pure returns (bytes32) { + return keccak256(abi.encodePacked(previous, keccak256(caKey), keccak256(k256Key))); + } + + function test_RegisterRootEpochAppendsAndChains() public { + assertEq(kms.rootEpochCount(), 0); + vm.expectRevert("No root epoch registered"); + kms.currentRootEpoch(); + + bytes32 h0 = chainHash(bytes32(0), ca(0xa0), k256(0xb0)); + bytes32 h1 = chainHash(h0, ca(0xa1), k256(0xb1)); + // Shared with the off-chain implementations' tests. + assertEq(h0, 0x10ab6b6c8e7601901b0f1767985bf40d35bf3129e0866fe12dcc7c8f04a298f2); + assertEq(h1, 0x50152a82023f2219a31723b476c13bc5f608ab4147946dabdec15d4f3d648ed4); + vm.startPrank(owner); + vm.expectEmit(true, true, true, true); + emit RootEpochRegistered(0, h0, ca(0xa0), k256(0xb0), "ipfs://c0"); + vm.expectEmit(true, true, true, true); + emit RootEpochActivated(0); + assertEq(kms.registerRootEpoch(ca(0xa0), k256(0xb0), "ipfs://c0"), 0); + assertEq(kms.registerRootEpoch(ca(0xa1), k256(0xb1), "ipfs://c1"), 1); + vm.stopPrank(); + + assertEq(kms.rootEpochCount(), 2); + (bytes memory caKey, bytes memory k256Key, string memory evidenceUrl, bytes32 h) = kms.rootEpoch(1); + assertEq(caKey, ca(0xa1)); + assertEq(k256Key, k256(0xb1)); + assertEq(evidenceUrl, "ipfs://c1"); + assertEq(h, h1); + vm.expectRevert("Unknown root epoch"); + kms.rootEpoch(2); + + // Epoch 1 is registered but epoch 0 is still served. + (uint256 current, bytes32 currentHash) = kms.currentRootEpoch(); + assertEq(current, 0); + assertEq(currentHash, h0); + } + + function test_ActivateRootEpochMovesForwardOneAtATime() public { + vm.startPrank(owner); + kms.registerRootEpoch(ca(0xa0), k256(0xb0), "ipfs://c0"); + vm.expectRevert("Invalid root epoch activation"); + kms.activateRootEpoch(1); + kms.registerRootEpoch(ca(0xa1), k256(0xb1), "ipfs://c1"); + kms.registerRootEpoch(ca(0xa2), k256(0xb2), "ipfs://c2"); + vm.expectRevert("Invalid root epoch activation"); + kms.activateRootEpoch(2); + vm.expectEmit(true, true, true, true); + emit RootEpochActivated(1); + kms.activateRootEpoch(1); + vm.expectRevert("Invalid root epoch activation"); + kms.activateRootEpoch(1); + vm.expectRevert("Invalid root epoch activation"); + kms.activateRootEpoch(0); + vm.stopPrank(); + + (uint256 current,) = kms.currentRootEpoch(); + assertEq(current, 1); + + vm.prank(user); + vm.expectRevert(abi.encodeWithSelector(OwnableUpgradeable.OwnableUnauthorizedAccount.selector, user)); + kms.activateRootEpoch(2); + } + + function test_RegisterRootEpochZeroMustMatchKmsInfo() public { + vm.startPrank(owner); + kms.setKmsInfo(DstackKms.KmsInfo(k256(0xb0), ca(0xa0), hex"0d", hex"0e")); + vm.expectRevert("Epoch 0 keys differ from kmsInfo"); + kms.registerRootEpoch(ca(0xa1), k256(0xb0), "ipfs://c0"); + vm.expectRevert("Epoch 0 keys differ from kmsInfo"); + kms.registerRootEpoch(ca(0xa0), k256(0xb1), "ipfs://c0"); + assertEq(kms.registerRootEpoch(ca(0xa0), k256(0xb0), "ipfs://c0"), 0); + // Later epochs are new keys by definition. + assertEq(kms.registerRootEpoch(ca(0xa1), k256(0xb1), "ipfs://c1"), 1); + vm.stopPrank(); + } + + function test_PopRootEpochDropsOnlyThePendingOne() public { + bytes32 h0 = chainHash(bytes32(0), ca(0xa0), k256(0xb0)); + bytes32 h1 = chainHash(h0, ca(0xa1), k256(0xb1)); + bytes32 h1b = chainHash(h0, ca(0xa2), k256(0xb2)); + vm.startPrank(owner); + vm.expectRevert("Not the pending root epoch"); + kms.popRootEpoch(0, h0); + kms.registerRootEpoch(ca(0xa0), k256(0xb0), "ipfs://c0"); + vm.expectRevert("Not the pending root epoch"); + kms.popRootEpoch(0, h0); + kms.registerRootEpoch(ca(0xa1), k256(0xb1), "ipfs://c1"); + vm.expectRevert("Not the pending root epoch"); + kms.popRootEpoch(2, h1); + vm.expectRevert("Root epoch chain hash mismatch"); + kms.popRootEpoch(1, h1b); + + vm.expectEmit(true, true, true, true); + emit RootEpochPopped(1, h1); + kms.popRootEpoch(1, h1); + assertEq(kms.rootEpochCount(), 1); + + // A new epoch 1 replaces the dropped one; a stale pop of the old one + // cannot drop it. + assertEq(kms.registerRootEpoch(ca(0xa2), k256(0xb2), "ipfs://c2"), 1); + vm.expectRevert("Root epoch chain hash mismatch"); + kms.popRootEpoch(1, h1); + kms.activateRootEpoch(1); + vm.expectRevert("Not the pending root epoch"); + kms.popRootEpoch(1, h1b); + vm.stopPrank(); + + (uint256 epoch, bytes32 h) = kms.currentRootEpoch(); + assertEq(epoch, 1); + assertEq(h, h1b); + + vm.prank(user); + vm.expectRevert(abi.encodeWithSelector(OwnableUpgradeable.OwnableUnauthorizedAccount.selector, user)); + kms.popRootEpoch(1, h1b); + } + + function test_SetRootEpochEvidenceUrl() public { + vm.startPrank(owner); + vm.expectRevert("Unknown root epoch"); + kms.setRootEpochEvidenceUrl(0, "ipfs://moved"); + kms.registerRootEpoch(ca(0xa0), k256(0xb0), "ipfs://c0"); + vm.expectRevert("Empty epoch evidence URL"); + kms.setRootEpochEvidenceUrl(0, ""); + vm.expectEmit(true, true, true, true); + emit RootEpochEvidenceUrlSet(0, "https://mirror/c0"); + kms.setRootEpochEvidenceUrl(0, "https://mirror/c0"); + vm.stopPrank(); + + (,, string memory evidenceUrl, bytes32 h) = kms.rootEpoch(0); + assertEq(evidenceUrl, "https://mirror/c0"); + assertEq(h, chainHash(bytes32(0), ca(0xa0), k256(0xb0))); + + vm.prank(user); + vm.expectRevert(abi.encodeWithSelector(OwnableUpgradeable.OwnableUnauthorizedAccount.selector, user)); + kms.setRootEpochEvidenceUrl(0, "ipfs://evil"); + } + + function test_KmsInfoFollowsTheCurrentRootEpoch() public { + vm.startPrank(owner); + kms.setKmsInfo(DstackKms.KmsInfo(k256(0xb0), ca(0xa0), hex"0d", hex"0e")); + kms.registerRootEpoch(ca(0xa0), k256(0xb0), "ipfs://c0"); + (bytes memory k256Key, bytes memory caKey, bytes memory quote, bytes memory eventlog) = kms.kmsInfo(); + assertEq(caKey, ca(0xa0)); + assertEq(k256Key, k256(0xb0)); + assertEq(quote.length + eventlog.length, 0); + + // A pending epoch is not the KMS identity yet. + kms.registerRootEpoch(ca(0xa1), k256(0xb1), "ipfs://c1"); + (k256Key, caKey,,) = kms.kmsInfo(); + assertEq(caKey, ca(0xa0)); + kms.activateRootEpoch(1); + (k256Key, caKey,,) = kms.kmsInfo(); + assertEq(caKey, ca(0xa1)); + assertEq(k256Key, k256(0xb1)); + + vm.expectRevert("kmsInfo follows the current root epoch"); + kms.setKmsInfo(DstackKms.KmsInfo(k256(0xb0), ca(0xa0), "", "")); + vm.expectRevert("kmsInfo follows the current root epoch"); + kms.setKmsQuote(hex"0d"); + vm.expectRevert("kmsInfo follows the current root epoch"); + kms.setKmsEventlog(hex"0e"); + vm.stopPrank(); + } + + function test_RegisterRootEpochOnlyOwner() public { + vm.prank(user); + vm.expectRevert(abi.encodeWithSelector(OwnableUpgradeable.OwnableUnauthorizedAccount.selector, user)); + kms.registerRootEpoch(ca(0xa0), k256(0xb0), "ipfs://c0"); + } + + function test_RegisterRootEpochRejectsMalformedKeys() public { + bytes memory compressedCa = ca(0xa0); + compressedCa[26] = 0x02; + bytes memory otherCurve = ca(0xa0); + otherCurve[25] = 0x0a; + bytes memory uncompressedK256 = bytes.concat(hex"04", new bytes(64)); + vm.startPrank(owner); + vm.expectRevert("Invalid epoch CA public key"); + kms.registerRootEpoch("", k256(0xb0), "ipfs://c0"); + vm.expectRevert("Invalid epoch CA public key"); + kms.registerRootEpoch(bytes.concat(ca(0xa0), hex"00"), k256(0xb0), "ipfs://c0"); + vm.expectRevert("Invalid epoch CA public key"); + kms.registerRootEpoch(compressedCa, k256(0xb0), "ipfs://c0"); + vm.expectRevert("Invalid epoch CA public key"); + kms.registerRootEpoch(otherCurve, k256(0xb0), "ipfs://c0"); + vm.expectRevert("Invalid epoch k256 public key"); + kms.registerRootEpoch(ca(0xa0), "", "ipfs://c0"); + vm.expectRevert("Invalid epoch k256 public key"); + kms.registerRootEpoch(ca(0xa0), uncompressedK256, "ipfs://c0"); + vm.expectRevert("Invalid epoch k256 public key"); + kms.registerRootEpoch(ca(0xa0), bytes.concat(hex"04", new bytes(32)), "ipfs://c0"); + vm.expectRevert("Empty epoch evidence URL"); + kms.registerRootEpoch(ca(0xa0), k256(0xb0), ""); + vm.stopPrank(); + } + function test_SetKmsInfoOnlyOwner() public { bytes memory k256Pubkey = hex"123456"; bytes memory caPubkey = hex"789abc"; diff --git a/dstack/kms/auth-mock/README.md b/dstack/kms/auth-mock/README.md index c26ed0e19..c44c207e8 100644 --- a/dstack/kms/auth-mock/README.md +++ b/dstack/kms/auth-mock/README.md @@ -72,6 +72,7 @@ bun run check - `MOCK_GATEWAY_APP_ID` - mock gateway app ID (default: 0xmockgateway1234567890123456789012345678) - `MOCK_CHAIN_ID` - mock blockchain chain ID (default: 1337) - `MOCK_APP_IMPLEMENTATION` - mock app implementation address (default: 0xmockapp9876543210987654321098765432109) +- `MOCK_ROOT_EPOCH`, `MOCK_ROOT_EPOCH_HASH` - current KMS root epoch and its chain hash, reported by the info endpoint as `rootEpoch`/`rootEpochHash`; boot checks from a KMS serving another epoch are denied (default: none registered, only epoch 0 serves) ## API endpoints diff --git a/dstack/kms/auth-mock/index.test.ts b/dstack/kms/auth-mock/index.test.ts index 493745613..68a3cfbf4 100644 --- a/dstack/kms/auth-mock/index.test.ts +++ b/dstack/kms/auth-mock/index.test.ts @@ -160,6 +160,34 @@ describe('Mock Backend Tests', () => { }); }); + describe('root epochs', () => { + it('allows only a KMS serving MOCK_ROOT_EPOCH', async () => { + const hash = '0x' + '11'.repeat(32); + process.env.MOCK_ROOT_EPOCH = '1'; + process.env.MOCK_ROOT_EPOCH_HASH = hash; + const boot = async (path: string, fields: object) => + (await appFetch(new Request(`http://localhost:3002/bootAuth/${path}`, { + method: 'POST', + headers: { 'Content-Type': 'application/json' }, + body: JSON.stringify({ + mrAggregated: '0x01', osImageHash: '0x02', appId: '0x03', + composeHash: '0x04', instanceId: '0x05', deviceId: '0x06', ...fields, + }), + }))).json(); + try { + expect((await boot('app', { rootEpoch: 1, rootEpochHash: hash })).isAllowed).toBe(true); + expect((await boot('app', { rootEpoch: 0, rootEpochHash: hash })).reason).toMatch(/1 is current/); + expect((await boot('app', {})).reason).toMatch(/reports none/); + expect((await boot('kms', {})).isAllowed).toBe(true); + const info = await appFetch(new Request('http://localhost:3002/')); + expect(await info.json()).toMatchObject({ rootEpoch: 1, rootEpochHash: hash }); + } finally { + delete process.env.MOCK_ROOT_EPOCH; + delete process.env.MOCK_ROOT_EPOCH_HASH; + } + }); + }); + describe('POST /bootAuth/kms', () => { const validBootInfo = { mrAggregated: '0x1234567890123456789012345678901234567890123456789012345678901234', diff --git a/dstack/kms/auth-mock/index.ts b/dstack/kms/auth-mock/index.ts index 577cae583..6d01a947f 100644 --- a/dstack/kms/auth-mock/index.ts +++ b/dstack/kms/auth-mock/index.ts @@ -22,7 +22,11 @@ const BootInfoSchema = z.object({ // optional fields (for full compatibility with BootInfo interface) tcbStatus: z.string().optional().default(''), advisoryIds: z.array(z.string()).optional().default([]), - mrSystem: z.string().optional().default('') + mrSystem: z.string().optional().default(''), + // The root epoch the asking KMS serves and its chain hash; absent from a KMS + // that predates root epochs, and from KMS checks made while onboarding. + rootEpoch: z.number().int().min(0).optional(), + rootEpochHash: z.string().optional() }); const BootResponseSchema = z.object({ @@ -31,9 +35,42 @@ const BootResponseSchema = z.object({ gatewayAppId: z.string() }); +// MOCK_ROOT_EPOCH and MOCK_ROOT_EPOCH_HASH: the current KMS root epoch and its +// chain hash; both unset means none is registered +function getRootEpoch(): { rootEpoch?: number; rootEpochHash?: string } { + const epoch = process.env.MOCK_ROOT_EPOCH; + const hash = process.env.MOCK_ROOT_EPOCH_HASH; + if (!epoch && !hash) return {}; + if (!epoch || !hash) throw new Error('set both MOCK_ROOT_EPOCH and MOCK_ROOT_EPOCH_HASH'); + return { rootEpoch: z.coerce.number().int().min(0).parse(epoch), rootEpochHash: hash }; +} + type BootInfo = z.infer; type BootResponse = z.infer; +const normHash = (hash?: string) => hash?.toLowerCase().replace(/^0x/, ''); + +// Only a KMS serving the current root epoch, with the registered keys up to it, +// may authorize. A KMS that predates root epochs reports none and holds only +// epoch 0; KMS checks made while onboarding carry none and are not limited. +function rootEpochDenial(bootInfo: BootInfo, isKms: boolean): string | undefined { + const current = getRootEpoch(); + const currentEpoch = current.rootEpoch ?? 0; + if (bootInfo.rootEpoch === undefined) { + return isKms || currentEpoch === 0 ? undefined : `root epoch ${currentEpoch} is current, but the KMS reports none`; + } + if (current.rootEpoch === undefined) { + return bootInfo.rootEpoch === 0 ? undefined : `root epoch ${bootInfo.rootEpoch} is not registered`; + } + if (bootInfo.rootEpoch !== currentEpoch) { + return `the KMS serves root epoch ${bootInfo.rootEpoch}, but ${currentEpoch} is current`; + } + if (normHash(bootInfo.rootEpochHash) !== normHash(current.rootEpochHash)) { + return `the KMS root keys up to epoch ${currentEpoch} are not the registered ones`; + } + return undefined; +} + // authorization policy - configurable via environment variables // MOCK_POLICY: "allow-all" (default), "deny-kms", "deny-app", "deny-all", // "allowlist-device", "allowlist-mr" @@ -70,6 +107,12 @@ class MockBackend { } async checkBoot(bootInfo: BootInfo, isKms: boolean): Promise { + const response = this.evaluateBoot(bootInfo, isKms); + const epochDenial = response.isAllowed ? rootEpochDenial(bootInfo, isKms) : undefined; + return epochDenial ? { isAllowed: false, reason: epochDenial, gatewayAppId: '' } : response; + } + + private evaluateBoot(bootInfo: BootInfo, isKms: boolean): BootResponse { const policy = getPolicy(); const deny = (reason: string): BootResponse => ({ isAllowed: false, @@ -147,6 +190,7 @@ app.get('/', async (c) => { chainId: batch[1], appAuthImplementation: batch[2], // NOTE: for backward compatibility appImplementation: batch[2], + ...getRootEpoch(), note: 'this is a mock backend - all authentications will succeed' }); } catch (error) { diff --git a/dstack/kms/auth-mock/openapi.json b/dstack/kms/auth-mock/openapi.json index 1689308d0..6a35cbbc3 100644 --- a/dstack/kms/auth-mock/openapi.json +++ b/dstack/kms/auth-mock/openapi.json @@ -151,6 +151,14 @@ "mrSystem": { "type": "string", "description": "MR System measurement (optional for compatibility)" + }, + "rootEpoch": { + "type": "integer", + "description": "KMS root epoch the asking KMS serves; must equal MOCK_ROOT_EPOCH" + }, + "rootEpochHash": { + "type": "string", + "description": "Its chain hash; must equal MOCK_ROOT_EPOCH_HASH" } } }, @@ -221,6 +229,8 @@ "description": "Mock app implementation address", "example": "0xmockapp9876543210987654321098765432109" }, + "rootEpoch": { "type": "integer", "description": "MOCK_ROOT_EPOCH" }, + "rootEpochHash": { "type": "string", "description": "MOCK_ROOT_EPOCH_HASH" }, "note": { "type": "string", "description": "Note about mock backend behavior", diff --git a/dstack/kms/auth-simple/README.md b/dstack/kms/auth-simple/README.md index efd2efcb5..dc48e7c50 100644 --- a/dstack/kms/auth-simple/README.md +++ b/dstack/kms/auth-simple/README.md @@ -107,6 +107,8 @@ raw NitroTPM evidence by itself. | `apps..composeHashes` | No | Allowed compose hashes for this app | | `apps..devices` | No | Allowed device IDs for this app. Empty denies all unless `allowAnyDevice` is true. | | `apps..allowAnyDevice` | No | If true, skip device ID check for this app | +| `rootEpochs` | No | KMS root public keys of each root epoch, indexed by epoch, `[{"caPubkey": "0x...", "k256Pubkey": "0x..."}]`. Epoch 0 is the original roots. Append only. The info endpoint reports the current epoch as `rootEpoch` with its chain hash as `rootEpochHash`; boot checks from a KMS serving another epoch are denied. | +| `currentRootEpoch` | No | Root epoch KMS instances serve (default 0). Raise it only once every KMS node holds that epoch. | For experimental AMD SEV-SNP dry-run authorization, keep the default fail-closed TCB policy unless you intentionally want the auth webhook to accept non-up-to-date verifier-derived SNP `BootInfo`. To exercise the dry-run path without enabling key release, allowlist the recomputed SNP `mrAggregated`, `osImageHash`, app/compose identity, device/chip identity, and any non-default `allowedTcbStatuses`/`allowedAdvisoryIds` values explicitly. KMS still rejects SNP before returning app keys, KMS keys, or app certificates. diff --git a/dstack/kms/auth-simple/bun.lock b/dstack/kms/auth-simple/bun.lock index daaf50f92..6350fe8d5 100644 --- a/dstack/kms/auth-simple/bun.lock +++ b/dstack/kms/auth-simple/bun.lock @@ -6,6 +6,7 @@ "name": "auth-simple", "dependencies": { "@hono/zod-validator": "0.9.1", + "@noble/hashes": "2.4.0", "hono": "4.13.14", "zod": "4.6.7", }, @@ -29,6 +30,8 @@ "@jridgewell/trace-mapping": ["@jridgewell/trace-mapping@0.3.31", "", { "dependencies": { "@jridgewell/resolve-uri": "^3.1.0", "@jridgewell/sourcemap-codec": "^1.4.14" } }, "sha512-zzNR+SdQSDJzc8joaeP8QQoCQr8NuYx2dIIytl1QeBEZHJ9uW6hebsrYgbz8hJwUQao3TWCMtmfV8Nu1twOLAw=="], + "@noble/hashes": ["@noble/hashes@2.4.0", "", {}, "sha512-X5XaVWZIBCT7HHZGm5I7ZQXDwLG+bGXuSrMQAW+7Zvl87h1kmc1ZB1VSRJcpUfoUrGQp4Fkoxm5kZ+Ms+aW+eA=="], + "@oxc-project/types": ["@oxc-project/types@0.155.0", "", {}, "sha512-SeE8JPwW1xs1tHeNycjUtUBdhorRa8o4jlvjwk49R7/Z291i+4DQa7JJIgOm/TaYQZx+QrPQbVJ/9c8dMIZN8Q=="], "@oxlint/darwin-arm64": ["@oxlint/darwin-arm64@0.9.10", "", { "os": "darwin", "cpu": "arm64" }, "sha512-eOXKZYq5bnCSgDefgM5bzAg+4Fc//Rc4yjgKN8iDWUARweCaChiQXb6TXX8MfEfs6qayEMy6yVj0pqoFz0B1aw=="], diff --git a/dstack/kms/auth-simple/index.test.ts b/dstack/kms/auth-simple/index.test.ts index fee441712..65fc7583a 100644 --- a/dstack/kms/auth-simple/index.test.ts +++ b/dstack/kms/auth-simple/index.test.ts @@ -446,6 +446,56 @@ describe('auth-simple', () => { }); }); + describe('root epochs', () => { + const rootEpochs = [ + // Well-formed keys whose point bytes repeat one value. + { caPubkey: '0x3059301306072a8648ce3d020106082a8648ce3d03010703420004' + 'a0'.repeat(64), k256Pubkey: '0x02' + 'b0'.repeat(32) }, + { caPubkey: '0x3059301306072a8648ce3d020106082a8648ce3d03010703420004' + 'a1'.repeat(64), k256Pubkey: '0x02' + 'b1'.repeat(32) } + ]; + // Shared with the DstackKms contract tests. + const h0 = '0x10ab6b6c8e7601901b0f1767985bf40d35bf3129e0866fe12dcc7c8f04a298f2'; + const h1 = '0x50152a82023f2219a31723b476c13bc5f608ab4147946dabdec15d4f3d648ed4'; + + const config = { + gatewayAppId: '0xgateway', + osImages: [baseBootInfo.osImageHash], + kms: { mrAggregated: [baseBootInfo.mrAggregated], allowAnyDevice: true }, + apps: { '0xapp123': { composeHashes: ['0xcompose456'], allowAnyDevice: true } }, + rootEpochs + }; + const boot = async (path: string, fields: object) => + (await app.fetch(new Request(`http://localhost/bootAuth/${path}`, { + method: 'POST', + headers: { 'Content-Type': 'application/json' }, + body: JSON.stringify({ ...baseBootInfo, ...fields }) + }))).json(); + const info = async () => + (await app.fetch(new Request('http://localhost/', { method: 'GET' }))).json(); + + it('allows only a KMS serving the current root epoch', async () => { + writeTestConfig(config); + expect(await info()).toMatchObject({ rootEpoch: 0, rootEpochHash: h0 }); + expect((await boot('app', { rootEpoch: 0, rootEpochHash: h0 })).isAllowed).toBe(true); + expect((await boot('app', {})).isAllowed).toBe(true); + + writeTestConfig({ ...config, currentRootEpoch: 1 }); + expect(await info()).toMatchObject({ rootEpoch: 1, rootEpochHash: h1 }); + expect((await boot('app', { rootEpoch: 1, rootEpochHash: h1 })).isAllowed).toBe(true); + expect((await boot('app', { rootEpoch: 0, rootEpochHash: h0 })).reason).toMatch(/1 is current/); + expect((await boot('app', { rootEpoch: 1, rootEpochHash: h0 })).reason).toMatch(/not the registered/); + // A KMS predating root epochs is cut off; onboarding checks carry none. + expect((await boot('app', {})).reason).toMatch(/reports none/); + expect((await boot('kms', {})).isAllowed).toBe(true); + }); + + it('serves only epoch 0 when none is configured', async () => { + writeTestConfig({ ...config, rootEpochs: [] }); + expect((await info()).rootEpoch).toBeUndefined(); + expect((await boot('app', { rootEpoch: 0, rootEpochHash: h0 })).isAllowed).toBe(true); + expect((await boot('app', { rootEpoch: 1, rootEpochHash: h1 })).reason).toMatch(/not registered/); + }); + }); + describe('hex normalization', () => { it('handles uppercase hex', async () => { writeTestConfig({ diff --git a/dstack/kms/auth-simple/index.ts b/dstack/kms/auth-simple/index.ts index 9da149014..371bc625d 100644 --- a/dstack/kms/auth-simple/index.ts +++ b/dstack/kms/auth-simple/index.ts @@ -6,6 +6,8 @@ import { Hono } from 'hono'; import { zValidator } from '@hono/zod-validator'; import { z } from 'zod'; import { readFileSync, existsSync } from 'fs'; +import { keccak_256 } from '@noble/hashes/sha3.js'; +import { bytesToHex, concatBytes, hexToBytes } from '@noble/hashes/utils.js'; // zod schemas for validation - compatible with auth-eth implementation const BootInfoSchema = z.object({ @@ -18,7 +20,11 @@ const BootInfoSchema = z.object({ deviceId: z.string().describe('device ID'), tcbStatus: z.string().optional().default(''), advisoryIds: z.array(z.string()).optional().default([]), - mrSystem: z.string().optional().default('') + mrSystem: z.string().optional().default(''), + // The root epoch the asking KMS serves and its chain hash; absent from a KMS + // that predates root epochs, and from KMS checks made while onboarding. + rootEpoch: z.number().int().min(0).optional(), + rootEpochHash: z.string().optional() }); const BootResponseSchema = z.object({ @@ -27,6 +33,11 @@ const BootResponseSchema = z.object({ gatewayAppId: z.string() }); +const RootEpochSchema = z.object({ + caPubkey: z.string(), + k256Pubkey: z.string() +}); + // config schema const AppConfigSchema = z.object({ composeHashes: z.array(z.string()).default([]), @@ -52,13 +63,63 @@ const AuthConfigSchema = z.object({ allowedTcbStatuses: z.array(z.string()).default(['UpToDate']), allowedAdvisoryIds: z.array(z.string()).default([]), kms: KmsConfigSchema.default({}), - apps: z.record(z.string(), AppConfigSchema).default({}) + apps: z.record(z.string(), AppConfigSchema).default({}), + // Root public keys of each KMS root epoch, indexed by epoch; append only. + rootEpochs: z.array(RootEpochSchema).default([]), + // The root epoch KMS instances serve; raise it once every node holds it. + currentRootEpoch: z.number().int().min(0).default(0) }); type BootInfo = z.infer; type BootResponse = z.infer; type AuthConfig = z.infer; +const bytesOf = (hex: string) => hexToBytes(hex.replace(/^0x/i, '')); + +// The current root epoch and its chain hash, as DstackKms computes it: +// keccak256(previous || keccak256(caPubkey) || keccak256(k256Pubkey)), +// with a zero previous hash for epoch 0. Empty when no epoch is configured. +function currentRootEpoch(config: AuthConfig): { rootEpoch?: number; rootEpochHash?: string } { + const current = config.currentRootEpoch; + if (config.rootEpochs.length === 0) return {}; + if (current >= config.rootEpochs.length) { + throw new Error('currentRootEpoch is not a configured root epoch'); + } + let hash = new Uint8Array(32); + for (const keys of config.rootEpochs.slice(0, current + 1)) { + hash = keccak_256(concatBytes( + hash, keccak_256(bytesOf(keys.caPubkey)), keccak_256(bytesOf(keys.k256Pubkey)) + )); + } + return { rootEpoch: current, rootEpochHash: '0x' + bytesToHex(hash) }; +} + +const normHash = (hash?: string) => hash?.toLowerCase().replace(/^0x/, ''); + +// Only a KMS serving the current root epoch, with the registered keys up to it, +// may authorize. A KMS that predates root epochs reports none and holds only +// epoch 0; KMS checks made while onboarding carry none and are not limited. +function rootEpochDenial( + bootInfo: BootInfo, + isKms: boolean, + current: { rootEpoch?: number; rootEpochHash?: string } +): string | undefined { + const currentEpoch = current.rootEpoch ?? 0; + if (bootInfo.rootEpoch === undefined) { + return isKms || currentEpoch === 0 ? undefined : `root epoch ${currentEpoch} is current, but the KMS reports none`; + } + if (current.rootEpoch === undefined) { + return bootInfo.rootEpoch === 0 ? undefined : `root epoch ${bootInfo.rootEpoch} is not registered`; + } + if (bootInfo.rootEpoch !== currentEpoch) { + return `the KMS serves root epoch ${bootInfo.rootEpoch}, but ${currentEpoch} is current`; + } + if (normHash(bootInfo.rootEpochHash) !== normHash(current.rootEpochHash)) { + return `the KMS root keys up to epoch ${currentEpoch} are not the registered ones`; + } + return undefined; +} + // normalize hex string to lowercase with 0x prefix function normalizeHex(hex: string): string { hex = hex.toLowerCase(); @@ -93,6 +154,14 @@ class ConfigBackend { async checkBoot(bootInfo: BootInfo, isKms: boolean): Promise { const config = this.loadConfig(); + const response = this.evaluateBoot(bootInfo, isKms, config); + 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 { const osImageHash = normalizeHex(bootInfo.osImageHash); const deviceId = normalizeHex(bootInfo.deviceId); @@ -220,6 +289,8 @@ class ConfigBackend { kmsContractAddr: string; chainId: number; appImplementation: string; + rootEpoch?: number; + rootEpochHash?: string; }> { const config = this.loadConfig(); return { @@ -227,6 +298,7 @@ class ConfigBackend { kmsContractAddr: config.kmsContractAddr, chainId: config.chainId, appImplementation: config.appImplementation, + ...currentRootEpoch(config), }; } } @@ -248,6 +320,8 @@ app.get('/', async (c) => { chainId: info.chainId, appAuthImplementation: info.appImplementation, // backward compat appImplementation: info.appImplementation, + rootEpoch: info.rootEpoch, + rootEpochHash: info.rootEpochHash, }); } catch (error) { console.error('error in health check:', error); diff --git a/dstack/kms/auth-simple/package.json b/dstack/kms/auth-simple/package.json index bb0bbb7e6..eba28e617 100644 --- a/dstack/kms/auth-simple/package.json +++ b/dstack/kms/auth-simple/package.json @@ -16,6 +16,7 @@ }, "dependencies": { "@hono/zod-validator": "0.9.1", + "@noble/hashes": "2.4.0", "hono": "4.13.14", "zod": "4.6.7" }, From 274e36dbdea5528f1087e69c3e05cb5293593c50 Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Fri, 6 Nov 2026 07:38:54 -0800 Subject: [PATCH 003/166] fix(kms): reject invalid root epoch configuration Signed-off-by: Kevin Wang --- dstack/crates/dstack-auth/src/main.rs | 87 ++++++++++++++++++++------- dstack/kms/auth-simple/index.test.ts | 17 ++++++ dstack/kms/auth-simple/index.ts | 2 +- 3 files changed, 82 insertions(+), 24 deletions(-) diff --git a/dstack/crates/dstack-auth/src/main.rs b/dstack/crates/dstack-auth/src/main.rs index 75dc30bff..00e5b72ee 100644 --- a/dstack/crates/dstack-auth/src/main.rs +++ b/dstack/crates/dstack-auth/src/main.rs @@ -20,8 +20,9 @@ //! (capture `tcbStatus`, deny unless `UpToDate`) if this grows into a //! multi-tenant / hosted deployment. -use anyhow::Result; +use anyhow::{Context, Result}; use clap::Parser; +use rocket::http::Status; use rocket::serde::json::Json; use rocket::{get, post, routes, State}; use serde::{Deserialize, Serialize}; @@ -100,26 +101,32 @@ struct RootEpoch { } impl Allowlist { - /// `None` when no epoch is configured, or the configuration is invalid, - /// which a KMS past epoch 0 treats as fail-closed. - fn current_root_epoch(&self) -> Option { - let epochs = self.root_epochs.get(..=self.current_root_epoch)?; + /// `None` only when no epoch is configured; invalid configuration is an error. + fn current_root_epoch(&self) -> Result> { + if self.root_epochs.is_empty() && self.current_root_epoch == 0 { + return Ok(None); + } + let epochs = self + .root_epochs + .get(..=self.current_root_epoch) + .context("currentRootEpoch is not a configured root epoch")?; let keys = epochs .iter() .map(|e| { - Some(( - hex::decode(norm(&e.ca_pubkey)).ok()?, - hex::decode(norm(&e.k256_pubkey)).ok()?, + Ok(( + hex::decode(norm(&e.ca_pubkey)).context("invalid root epoch CA public key")?, + hex::decode(norm(&e.k256_pubkey)) + .context("invalid root epoch k256 public key")?, )) }) - .collect::>>()?; + .collect::>>()?; let hashes = dstack_types::root_epoch::chain_hashes( keys.iter().map(|(ca, k256)| (&ca[..], &k256[..])), ); - Some(EpochInfo { + Ok(Some(EpochInfo { number: self.current_root_epoch, - hash: format!("0x{}", hex::encode(hashes.last()?)), - }) + hash: format!("0x{}", hex::encode(hashes[self.current_root_epoch])), + })) } } @@ -175,7 +182,10 @@ fn deny(al: &Allowlist, reason: &str) -> BootResponse { /// registered keys up to it. A KMS that predates root epochs reports none and /// holds only epoch 0; KMS checks made while onboarding carry none. fn allow(info: &BootInfo, is_kms: bool, al: &Allowlist) -> BootResponse { - let current = al.current_root_epoch(); + let current = match al.current_root_epoch() { + Ok(current) => current, + Err(error) => return deny(al, &format!("invalid root epoch configuration: {error}")), + }; let current_epoch = current.as_ref().map_or(0, |c| c.number); let denial = match (info.root_epoch, ¤t) { (None, _) if is_kms || current_epoch == 0 => None, @@ -277,11 +287,14 @@ fn boot_kms(info: Json, cli: &State) -> Json { /// info endpoint the KMS GETs to populate its metadata. Single-node: no chain. #[get("/")] -fn info(cli: &State) -> Json { - Json(info_json(&load(&cli.config))) +fn info(cli: &State) -> Result, Status> { + info_json(&load(&cli.config)).map(Json).map_err(|error| { + rocket::warn!("invalid root epoch configuration: {error}"); + Status::ServiceUnavailable + }) } -fn info_json(al: &Allowlist) -> serde_json::Value { +fn info_json(al: &Allowlist) -> Result { let mut info = json!({ "status": "ok", "kmsContractAddr": "", @@ -290,11 +303,11 @@ fn info_json(al: &Allowlist) -> serde_json::Value { "chainId": 0, "appImplementation": "", }); - if let Some(epoch) = al.current_root_epoch() { + if let Some(epoch) = al.current_root_epoch()? { info["rootEpoch"] = epoch.number.into(); info["rootEpochHash"] = epoch.hash.into(); } - info + Ok(info) } #[rocket::main] @@ -380,12 +393,12 @@ mod tests { root_epoch_hash: Some(hash.into()), ..Default::default() }; - assert_eq!(info_json(&al)["rootEpochHash"], h0); + assert_eq!(info_json(&al).unwrap()["rootEpochHash"], h0); assert!(check_kms(&info(Some(0), h0), &al).is_allowed); al.current_root_epoch = 1; - assert_eq!(info_json(&al)["rootEpoch"], 1); - assert_eq!(info_json(&al)["rootEpochHash"], h1); + assert_eq!(info_json(&al).unwrap()["rootEpoch"], 1); + assert_eq!(info_json(&al).unwrap()["rootEpochHash"], h1); assert!(check_kms(&info(Some(1), &h1.to_uppercase()), &al).is_allowed); assert!(!check_kms(&info(Some(0), h0), &al).is_allowed); assert!(!check_kms(&info(Some(1), h0), &al).is_allowed); @@ -393,8 +406,8 @@ mod tests { assert!(check_kms(&info(None, ""), &al).is_allowed); al.current_root_epoch = 2; - assert!(al.current_root_epoch().is_none()); - assert!(info_json(&al).get("rootEpoch").is_none()); + assert!(al.current_root_epoch().is_err()); + assert!(info_json(&al).is_err()); assert!(!check_kms(&info(Some(1), h1), &al).is_allowed); } @@ -418,6 +431,34 @@ mod tests { assert!(!check_app(&info, &al).is_allowed); } + #[test] + fn invalid_root_epochs_deny_legacy_and_epoch_aware_kms() { + let invalid_configs = [ + json!({"currentRootEpoch": 1}), + json!({"currentRootEpoch": 1, "rootEpochs": [{"caPubkey": "a0", "k256Pubkey": "b0"}]}), + json!({"rootEpochs": [{"caPubkey": "invalid", "k256Pubkey": "b0"}]}), + json!({"rootEpochs": [{"caPubkey": "a0", "k256Pubkey": "invalid"}]}), + ]; + for config in invalid_configs { + let epochs: Allowlist = serde_json::from_value(config).unwrap(); + let al = Allowlist { + root_epochs: epochs.root_epochs, + current_root_epoch: epochs.current_root_epoch, + ..allowlist() + }; + assert!(info_json(&al).is_err()); + for epoch in [None, Some(0), Some(1)] { + let info = BootInfo { + mr_aggregated: "0xMR".into(), + root_epoch: epoch, + ..boot("0xapp1", "0xhash", "0ximg") + }; + assert!(!check_app(&info, &al).is_allowed); + assert!(!check_kms(&info, &al).is_allowed); + } + } + } + // wire-contract snapshot: BootInfo as the KMS serializes it (camelCase). // Keep these field names in sync with the kms BootInfo. `#[serde(default)]` // means extra fields are ignored AND a renamed field deserializes to "" — diff --git a/dstack/kms/auth-simple/index.test.ts b/dstack/kms/auth-simple/index.test.ts index 65fc7583a..5e7e37238 100644 --- a/dstack/kms/auth-simple/index.test.ts +++ b/dstack/kms/auth-simple/index.test.ts @@ -488,6 +488,23 @@ describe('auth-simple', () => { expect((await boot('kms', {})).isAllowed).toBe(true); }); + it.each([ + { currentRootEpoch: 1, rootEpochs: [] }, + { currentRootEpoch: 1, rootEpochs: undefined }, + { currentRootEpoch: 2, rootEpochs }, + { rootEpochs: [{ ...rootEpochs[0], caPubkey: '0xINVALID' }] }, + { rootEpochs: [{ ...rootEpochs[0], k256Pubkey: '0xINVALID' }] } + ])('fails closed for invalid root epoch config %j', async (epochs) => { + writeTestConfig({ ...config, ...epochs }); + const health = await app.fetch(new Request('http://localhost/')); + expect(health.status).toBe(500); + for (const path of ['app', 'kms']) { + for (const fields of [{}, { rootEpoch: 0 }, { rootEpoch: 1 }]) { + expect((await boot(path, fields)).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 371bc625d..4a242ac7e 100644 --- a/dstack/kms/auth-simple/index.ts +++ b/dstack/kms/auth-simple/index.ts @@ -81,7 +81,7 @@ const bytesOf = (hex: string) => hexToBytes(hex.replace(/^0x/i, '')); // with a zero previous hash for epoch 0. Empty when no epoch is configured. function currentRootEpoch(config: AuthConfig): { rootEpoch?: number; rootEpochHash?: string } { const current = config.currentRootEpoch; - if (config.rootEpochs.length === 0) return {}; + if (config.rootEpochs.length === 0 && current === 0) return {}; if (current >= config.rootEpochs.length) { throw new Error('currentRootEpoch is not a configured root epoch'); } From 064f7593b4f055ea43bb89c5ebd241f2740131f7 Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Fri, 6 Nov 2026 07:40:19 -0800 Subject: [PATCH 004/166] fix(kms): bind root epoch activation to its chain hash Signed-off-by: Kevin Wang --- dstack/kms/auth-eth/contracts/DstackKms.sol | 6 ++- dstack/kms/auth-eth/script/Manage.s.sol | 3 +- dstack/kms/auth-eth/script/README.md | 4 +- dstack/kms/auth-eth/test/DstackKms.t.sol | 49 +++++++++++++++++---- 4 files changed, 50 insertions(+), 12 deletions(-) diff --git a/dstack/kms/auth-eth/contracts/DstackKms.sol b/dstack/kms/auth-eth/contracts/DstackKms.sol index 8330ece1f..21afb93eb 100644 --- a/dstack/kms/auth-eth/contracts/DstackKms.sol +++ b/dstack/kms/auth-eth/contracts/DstackKms.sol @@ -186,9 +186,11 @@ contract DstackKms is Initializable, Ownable2StepUpgradeable, UUPSUpgradeable, E } } - // Make the next registered root epoch the one KMS instances serve - function activateRootEpoch(uint256 epoch) external onlyOwner { + // Make the next registered root epoch the one KMS instances serve. Bind + // activation to its chain hash in case a pending epoch has been replaced. + function activateRootEpoch(uint256 epoch, bytes32 chainHash) external onlyOwner { require(epoch == _currentRootEpoch + 1 && epoch < _rootEpochs.length, "Invalid root epoch activation"); + require(_rootEpochs[epoch].chainHash == chainHash, "root epoch chain hash mismatch"); _currentRootEpoch = epoch; _useRootEpoch(epoch); _emitPolicy("root-epoch-activated", bytes32(epoch), true); diff --git a/dstack/kms/auth-eth/script/Manage.s.sol b/dstack/kms/auth-eth/script/Manage.s.sol index 9d87e1f33..0dcedc110 100644 --- a/dstack/kms/auth-eth/script/Manage.s.sol +++ b/dstack/kms/auth-eth/script/Manage.s.sol @@ -126,9 +126,10 @@ contract RegisterRootEpoch is BaseScript { contract ActivateRootEpoch is BaseScript { function run() external { uint256 epoch = vm.envUint("ROOT_EPOCH"); + bytes32 chainHash = vm.envBytes32("CHAIN_HASH"); vm.startBroadcast(vm.envUint("PRIVATE_KEY")); - kms.activateRootEpoch(epoch); + kms.activateRootEpoch(epoch, chainHash); vm.stopBroadcast(); console.log("Activated KMS root epoch", epoch); diff --git a/dstack/kms/auth-eth/script/README.md b/dstack/kms/auth-eth/script/README.md index 0c706f50a..9c5630c20 100644 --- a/dstack/kms/auth-eth/script/README.md +++ b/dstack/kms/auth-eth/script/README.md @@ -77,8 +77,10 @@ forge script script/Manage.s.sol:RegisterRootEpoch --broadcast --rpc-url $RPC_UR Makes the next registered root epoch the one KMS instances serve. Activate it only once every KMS node holds it; nodes still at the previous epoch stop serving, and nodes holding it restart to switch. Epochs only move forward. +`CHAIN_HASH` is the expected epoch chain hash (see `Query.s.sol`), +so a stale activation cannot activate replacement keys at the same epoch. ```bash -KMS_CONTRACT_ADDR=0x... ROOT_EPOCH=1 \ +KMS_CONTRACT_ADDR=0x... ROOT_EPOCH=1 CHAIN_HASH=0x... \ forge script script/Manage.s.sol:ActivateRootEpoch --broadcast --rpc-url $RPC_URL ``` diff --git a/dstack/kms/auth-eth/test/DstackKms.t.sol b/dstack/kms/auth-eth/test/DstackKms.t.sol index 847a9153f..ebb886b8d 100644 --- a/dstack/kms/auth-eth/test/DstackKms.t.sol +++ b/dstack/kms/auth-eth/test/DstackKms.t.sol @@ -128,21 +128,24 @@ contract DstackKmsTest is Test { } function test_ActivateRootEpochMovesForwardOneAtATime() public { + bytes32 h0 = chainHash(bytes32(0), ca(0xa0), k256(0xb0)); + bytes32 h1 = chainHash(h0, ca(0xa1), k256(0xb1)); + bytes32 h2 = chainHash(h1, ca(0xa2), k256(0xb2)); vm.startPrank(owner); kms.registerRootEpoch(ca(0xa0), k256(0xb0), "ipfs://c0"); vm.expectRevert("Invalid root epoch activation"); - kms.activateRootEpoch(1); + kms.activateRootEpoch(1, h1); kms.registerRootEpoch(ca(0xa1), k256(0xb1), "ipfs://c1"); kms.registerRootEpoch(ca(0xa2), k256(0xb2), "ipfs://c2"); vm.expectRevert("Invalid root epoch activation"); - kms.activateRootEpoch(2); + kms.activateRootEpoch(2, h2); vm.expectEmit(true, true, true, true); emit RootEpochActivated(1); - kms.activateRootEpoch(1); + kms.activateRootEpoch(1, h1); vm.expectRevert("Invalid root epoch activation"); - kms.activateRootEpoch(1); + kms.activateRootEpoch(1, h1); vm.expectRevert("Invalid root epoch activation"); - kms.activateRootEpoch(0); + kms.activateRootEpoch(0, h0); vm.stopPrank(); (uint256 current,) = kms.currentRootEpoch(); @@ -150,7 +153,36 @@ contract DstackKmsTest is Test { vm.prank(user); vm.expectRevert(abi.encodeWithSelector(OwnableUpgradeable.OwnableUnauthorizedAccount.selector, user)); - kms.activateRootEpoch(2); + kms.activateRootEpoch(2, h2); + } + + function test_activate_root_epoch_rejects_replaced_pending_epoch() public { + bytes32 h0 = chainHash(bytes32(0), ca(0xa0), k256(0xb0)); + bytes32 h1 = chainHash(h0, ca(0xa1), k256(0xb1)); + bytes32 replacementHash = chainHash(h0, ca(0xa2), k256(0xb2)); + vm.startPrank(owner); + kms.registerRootEpoch(ca(0xa0), k256(0xb0), "ipfs://c0"); + kms.registerRootEpoch(ca(0xa1), k256(0xb1), "ipfs://c1"); + kms.popRootEpoch(1, h1); + kms.registerRootEpoch(ca(0xa2), k256(0xb2), "ipfs://replacement"); + + vm.expectRevert("root epoch chain hash mismatch"); + kms.activateRootEpoch(1, h1); + (uint256 current, bytes32 currentHash) = kms.currentRootEpoch(); + assertEq(current, 0); + assertEq(currentHash, h0); + (bytes memory k256Key, bytes memory caKey,,) = kms.kmsInfo(); + assertEq(caKey, ca(0xa0)); + assertEq(k256Key, k256(0xb0)); + + kms.activateRootEpoch(1, replacementHash); + (current, currentHash) = kms.currentRootEpoch(); + assertEq(current, 1); + assertEq(currentHash, replacementHash); + (k256Key, caKey,,) = kms.kmsInfo(); + assertEq(caKey, ca(0xa2)); + assertEq(k256Key, k256(0xb2)); + vm.stopPrank(); } function test_RegisterRootEpochZeroMustMatchKmsInfo() public { @@ -192,7 +224,7 @@ contract DstackKmsTest is Test { assertEq(kms.registerRootEpoch(ca(0xa2), k256(0xb2), "ipfs://c2"), 1); vm.expectRevert("Root epoch chain hash mismatch"); kms.popRootEpoch(1, h1); - kms.activateRootEpoch(1); + kms.activateRootEpoch(1, h1b); vm.expectRevert("Not the pending root epoch"); kms.popRootEpoch(1, h1b); vm.stopPrank(); @@ -240,7 +272,8 @@ contract DstackKmsTest is Test { kms.registerRootEpoch(ca(0xa1), k256(0xb1), "ipfs://c1"); (k256Key, caKey,,) = kms.kmsInfo(); assertEq(caKey, ca(0xa0)); - kms.activateRootEpoch(1); + (,,, bytes32 h1) = kms.rootEpoch(1); + kms.activateRootEpoch(1, h1); (k256Key, caKey,,) = kms.kmsInfo(); assertEq(caKey, ca(0xa1)); assertEq(k256Key, k256(0xb1)); From 220480b63a71dcd2e4254acf0d3f8899501b2b1f Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Fri, 6 Nov 2026 07:43:54 -0800 Subject: [PATCH 005/166] fix(kms): require root epoch registry support in auth backends Signed-off-by: Kevin Wang --- dstack/kms/auth-eth-bun/README.md | 8 ++ dstack/kms/auth-eth-bun/index.test.ts | 93 +++++++++++++++++--- dstack/kms/auth-eth-bun/index.ts | 49 +++++------ dstack/kms/auth-eth/README.md | 8 ++ dstack/kms/auth-eth/src/ethereum.ts | 21 ++--- dstack/kms/auth-eth/src/main.test.ts | 118 ++++++++++++++++++++++++++ 6 files changed, 242 insertions(+), 55 deletions(-) diff --git a/dstack/kms/auth-eth-bun/README.md b/dstack/kms/auth-eth-bun/README.md index 9dcd4c955..8fcafb0bd 100644 --- a/dstack/kms/auth-eth-bun/README.md +++ b/dstack/kms/auth-eth-bun/README.md @@ -17,6 +17,14 @@ a single-file implementation of the dstack KMS ethereum backend using bun + hono bun install ``` +## Upgrade order + +Upgrade the `DstackKms` contract before deploying the new auth backend or KMS. +An empty root epoch registry is supported on the upgraded contract; there is no +fallback for legacy contracts. Old KMS instances that omit the root epoch can +receive app authorization only while no epoch is registered or epoch 0 is current, +subject to the existing authorization policy. + ## usage ### development diff --git a/dstack/kms/auth-eth-bun/index.test.ts b/dstack/kms/auth-eth-bun/index.test.ts index 5a2645400..22af72fc8 100644 --- a/dstack/kms/auth-eth-bun/index.test.ts +++ b/dstack/kms/auth-eth-bun/index.test.ts @@ -58,6 +58,7 @@ describe('API Compatibility Tests', () => { it('should return system info matching OpenAPI spec', async () => { // Mock contract calls mockReadContract.mockImplementation((params) => { + if (params.functionName === 'rootEpochCount') return 0n; if (params.functionName === 'gatewayAppId') { return '0xabcdefabcdefabcdefabcdefabcdefabcdefabcd'; } @@ -119,6 +120,7 @@ describe('API Compatibility Tests', () => { it('should validate app boot with required fields only', async () => { // Mock successful contract response mockReadContract.mockImplementation((params) => { + if (params.functionName === 'rootEpochCount') return 0n; if (params.functionName === 'isAppAllowed') { return [true, 'success']; } @@ -160,6 +162,7 @@ describe('API Compatibility Tests', () => { }; mockReadContract.mockImplementation((params) => { + if (params.functionName === 'rootEpochCount') return 0n; if (params.functionName === 'isAppAllowed') { return [true, 'success with full info']; } @@ -241,13 +244,14 @@ describe('API Compatibility Tests', () => { }; const hash = '0x' + '11'.repeat(32); - const mockChain = (currentRootEpoch: () => unknown) => + const mockChain = (currentRootEpoch: () => unknown, rootEpochCount = () => 2n) => mockReadContract.mockImplementation((params) => { switch (params.functionName) { case 'isAppAllowed': return [true, '']; case 'isKmsAllowed': return [true, '']; case 'gatewayAppId': return '0xgateway'; case 'appImplementation': return '0x9876543210987654321098765432109876543210'; + case 'rootEpochCount': return rootEpochCount(); case 'currentRootEpoch': return currentRootEpoch(); } }); @@ -277,16 +281,74 @@ describe('API Compatibility Tests', () => { expect(await info.json()).toMatchObject({ rootEpoch: 1, rootEpochHash: hash }); }); - it('treats a contract without the registry as having no epochs', async () => { - mockChain(() => { - throw Object.assign(new Error('call failed'), { - cause: { name: 'ContractFunctionZeroDataError' }, - }); - }); + it('supports an upgraded empty registry without querying the current epoch', async () => { + const current = vi.fn(() => { throw new Error('no root epoch registered'); }); + mockChain(current, () => 0n); expect((await boot('app', {})).isAllowed).toBe(true); + expect((await boot('kms', {})).isAllowed).toBe(true); expect((await boot('app', { rootEpoch: 0, rootEpochHash: hash })).isAllowed).toBe(true); expect((await boot('app', { rootEpoch: 1, rootEpochHash: hash })).reason).toMatch(/not registered/); + const info = await appFetch(new Request('http://localhost:3001/')); + expect(info.status).toBe(200); + expect(await info.json()).not.toHaveProperty('rootEpoch'); + expect(current).not.toHaveBeenCalled(); + }); + + it('allows old KMS app requests while epoch 0 is current', async () => { + mockChain(() => [0n, hash], () => 1n); + expect((await boot('app', {})).isAllowed).toBe(true); + }); + + it.each(['app', 'info'])('reads the registry at the same finalized block for %s', async (endpoint) => { + mockChain(() => [1n, hash]); + if (endpoint === 'app') { + expect((await boot('app', { rootEpoch: 1, rootEpochHash: hash })).isAllowed).toBe(true); + } else { + const response = await appFetch(new Request('http://localhost:3001/')); + expect(response.status).toBe(200); + } + const calls = mockReadContract.mock.calls.map(([params]) => params) + .filter(({ functionName }) => ['rootEpochCount', 'currentRootEpoch'].includes(functionName)); + expect(calls.map(({ functionName }) => functionName)).toEqual(['rootEpochCount', 'currentRootEpoch']); + expect(calls.every(({ blockNumber }) => blockNumber === 98n)).toBe(true); + expect(mockGetBlockNumber).toHaveBeenCalledTimes(1); + }); + + describe.each(['rootEpochCount', 'currentRootEpoch'])('%s failures', (getter) => { + it.each([ + ['missing getter', { name: 'ContractFunctionZeroDataError' }], + ['revert', { name: 'ContractFunctionRevertedError' }], + ['RPC failure', new Error('RPC unavailable')], + ])('fails closed on %s', async (_name, cause) => { + const failure = () => { throw Object.assign(new Error('query failed'), { cause }); }; + mockChain(getter === 'currentRootEpoch' ? failure : () => [1n, hash], + getter === 'rootEpochCount' ? failure : () => 2n); + + for (const path of ['app', 'kms']) { + expect(await boot(path, {})).toEqual({ + isAllowed: false, reason: 'authorization backend unavailable', gatewayAppId: '', + }); + expect((await boot(path, { rootEpoch: 0, rootEpochHash: hash })).isAllowed).toBe(false); + } + const info = await appFetch(new Request('http://localhost:3001/')); + expect(info.status).toBe(500); + }); + + it.each(['0x', '0x1234'] as const)('fails closed on malformed ABI data %s', async (data) => { + const { decodeFunctionResult } = await vi.importActual('viem'); + mockReadContract.mockImplementation((params) => { + if (params.functionName === getter) return decodeFunctionResult({ ...params, data }); + if (params.functionName === 'rootEpochCount') return 2n; + if (params.functionName === 'isAppAllowed') return [true, '']; + if (params.functionName === 'gatewayAppId') return 'gateway-app'; + if (params.functionName === 'appImplementation') return '0x' + '11'.repeat(20); + throw new Error(`unexpected function ${params.functionName}`); + }); + expect((await boot('app', {})).isAllowed).toBe(false); + const info = await appFetch(new Request('http://localhost:3001/')); + expect(info.status).toBe(500); + }); }); }); @@ -303,6 +365,7 @@ describe('API Compatibility Tests', () => { it('should validate KMS boot successfully', async () => { // Mock successful contract response mockReadContract.mockImplementation((params) => { + if (params.functionName === 'rootEpochCount') return 0n; if (params.functionName === 'isKmsAllowed') { return [true, 'KMS allowed']; } @@ -329,6 +392,7 @@ describe('API Compatibility Tests', () => { it('should handle KMS rejection', async () => { mockReadContract.mockImplementation((params) => { + if (params.functionName === 'rootEpochCount') return 0n; if (params.functionName === 'isKmsAllowed') { return [false, 'KMS not authorized']; } @@ -459,6 +523,7 @@ describe('Hex Decoding Compatibility', () => { }; mockReadContract.mockImplementation((params) => { + if (params.functionName === 'rootEpochCount') return 0n; if (params.functionName === 'isAppAllowed') { // Verify that hex values are properly formatted const [bootInfoStruct] = params.args; @@ -504,12 +569,12 @@ describe('Authorization freshness and domain binding', () => { let decisions = 0; mockReadContract.mockImplementation((params) => { expect(params.address).toBe('0x1234567890123456789012345678901234567890'); + if (params.functionName === 'rootEpochCount') return 0n; if (params.functionName === 'isAppAllowed') { decisions += 1; return decisions === 1 ? [true, 'initial allow'] : [false, 'policy changed']; } if (params.functionName === 'gatewayAppId') return 'gateway-app'; - if (params.functionName === 'currentRootEpoch') return undefined; throw new Error(`unexpected function ${params.functionName}`); }); @@ -524,12 +589,12 @@ describe('Authorization freshness and domain binding', () => { it('binds changed measurements and identities into distinct contract arguments', async () => { const calls: unknown[] = []; mockReadContract.mockImplementation((params) => { + if (params.functionName === 'rootEpochCount') return 0n; if (params.functionName === 'isAppAllowed') { calls.push(params.args[0]); return [true, 'allowed']; } if (params.functionName === 'gatewayAppId') return 'gateway-app'; - if (params.functionName === 'currentRootEpoch') return undefined; throw new Error(`unexpected function ${params.functionName}`); }); @@ -552,9 +617,9 @@ describe('Authorization freshness and domain binding', () => { }); mockReadContract.mockImplementation((params) => { + if (params.functionName === 'rootEpochCount') return 0n; if (params.functionName === 'isAppAllowed') return [true, 'recovered']; if (params.functionName === 'gatewayAppId') return 'gateway-app'; - if (params.functionName === 'currentRootEpoch') return undefined; throw new Error(`unexpected function ${params.functionName}`); }); const recovered = await postApp(); @@ -583,9 +648,9 @@ describe('Ethereum finalized snapshot authorization', () => { mockGetBlockNumber.mockResolvedValue(100n); mockReadContract.mockImplementation((params) => { expect(params.blockNumber).toBe(98n); + if (params.functionName === 'rootEpochCount') return 0n; if (params.functionName === 'isAppAllowed') return [true, 'finalized allow']; if (params.functionName === 'gatewayAppId') return 'gateway-app'; - if (params.functionName === 'currentRootEpoch') return undefined; throw new Error(`unexpected function ${params.functionName}`); }); @@ -599,13 +664,13 @@ describe('Ethereum finalized snapshot authorization', () => { let decisions = 0; const observedBlocks: bigint[] = []; mockReadContract.mockImplementation((params) => { + if (params.functionName === 'rootEpochCount') return 0n; if (params.functionName === 'isAppAllowed') { observedBlocks.push(params.blockNumber); decisions += 1; return decisions === 1 ? [true, 'old canonical allow'] : [false, 'new canonical deny']; } if (params.functionName === 'gatewayAppId') return 'gateway-app'; - if (params.functionName === 'currentRootEpoch') return undefined; throw new Error(`unexpected function ${params.functionName}`); }); @@ -632,9 +697,9 @@ describe('Ethereum finalized snapshot authorization', () => { mockGetChainId.mockResolvedValue(1337); mockGetBlockNumber.mockResolvedValue(102n); mockReadContract.mockImplementation((params) => { + if (params.functionName === 'rootEpochCount') return 0n; if (params.functionName === 'isAppAllowed') return [true, 'recovered']; if (params.functionName === 'gatewayAppId') return 'gateway-app'; - if (params.functionName === 'currentRootEpoch') return undefined; throw new Error(`unexpected function ${params.functionName}`); }); const recovered = await authorize(); @@ -673,11 +738,11 @@ describe('dstack-kms 0.6 bootAuth payload', () => { beforeEach(() => { mockReadContract.mockImplementation((params) => { + if (params.functionName === 'rootEpochCount') return 0n; if (params.functionName === 'isKmsAllowed' || params.functionName === 'isAppAllowed') { return [true, '']; } if (params.functionName === 'gatewayAppId') return 'gateway-app'; - if (params.functionName === 'currentRootEpoch') return undefined; throw new Error(`unexpected function ${params.functionName}`); }); }); diff --git a/dstack/kms/auth-eth-bun/index.ts b/dstack/kms/auth-eth-bun/index.ts index d274fb009..f7d726377 100644 --- a/dstack/kms/auth-eth-bun/index.ts +++ b/dstack/kms/auth-eth-bun/index.ts @@ -139,6 +139,13 @@ const DSTACK_KMS_ABI = [ inputs: [], outputs: [{ name: '', type: 'address' }] }, + { + name: 'rootEpochCount', + type: 'function', + stateMutability: 'view', + inputs: [], + outputs: [{ name: '', type: 'uint256' }] + }, { name: 'currentRootEpoch', type: 'function', @@ -151,17 +158,6 @@ const DSTACK_KMS_ABI = [ } ] as const; -// Contracts deployed before root epochs lack the registry getters, and a -// registry with no epoch reverts. -const isMissingFunction = (error: unknown): boolean => { - for (let cause: any = error; cause; cause = cause.cause) { - if (cause.name === 'ContractFunctionZeroDataError' || cause.name === 'ContractFunctionRevertedError') { - return true; - } - } - return false; -}; - // ethereum backend class class EthereumBackend { private client: ReturnType; @@ -254,24 +250,23 @@ class EthereumBackend { }; } - // The current root epoch and its chain hash; empty when none is registered - // or the contract predates root epochs. + // The current root epoch and its chain hash; empty only when no epoch is registered. async getRootEpoch(blockNumber?: bigint): Promise { blockNumber ??= await this.finalizedBlockNumber(); - try { - const current = await this.client.readContract({ - address: this.kmsContractAddr, - abi: DSTACK_KMS_ABI, - functionName: 'currentRootEpoch', - blockNumber - }); - if (!current) return {}; - const [epoch, chainHash] = current; - return { rootEpoch: Number(epoch), rootEpochHash: chainHash }; - } catch (error) { - if (isMissingFunction(error)) return {}; - throw error; - } + const count = await this.client.readContract({ + address: this.kmsContractAddr, + abi: DSTACK_KMS_ABI, + functionName: 'rootEpochCount', + blockNumber + }); + if (count === 0n) return {}; + const [epoch, chainHash] = await this.client.readContract({ + address: this.kmsContractAddr, + abi: DSTACK_KMS_ABI, + functionName: 'currentRootEpoch', + blockNumber + }); + return { rootEpoch: Number(epoch), rootEpochHash: chainHash }; } async getGatewayAppId(): Promise { diff --git a/dstack/kms/auth-eth/README.md b/dstack/kms/auth-eth/README.md index 3d3b88f36..bb543b6dd 100644 --- a/dstack/kms/auth-eth/README.md +++ b/dstack/kms/auth-eth/README.md @@ -13,6 +13,14 @@ This project contains upgradeable smart contracts for: - [Foundry](https://book.getfoundry.sh/getting-started/installation) - For smart contract development and testing - Node.js - For the bootAuth server and TypeScript development +## Upgrade order + +Upgrade the `DstackKms` contract before deploying the new auth backend or KMS. +An empty root epoch registry is supported on the upgraded contract; there is no +fallback for legacy contracts. Old KMS instances that omit the root epoch can +receive app authorization only while no epoch is registered or epoch 0 is current, +subject to the existing authorization policy. + ## Setup 1. Install dependencies: diff --git a/dstack/kms/auth-eth/src/ethereum.ts b/dstack/kms/auth-eth/src/ethereum.ts index 45ff7283e..6eab5d874 100644 --- a/dstack/kms/auth-eth/src/ethereum.ts +++ b/dstack/kms/auth-eth/src/ethereum.ts @@ -11,6 +11,7 @@ const DSTACK_KMS_ABI = [ "function isKmsAllowed((address appId,bytes32 composeHash,address instanceId,bytes32 deviceId,bytes32 mrAggregated,bytes32 mrSystem,bytes32 osImageHash,string tcbStatus,string[] advisoryIds) bootInfo) view returns (bool, string)", "function gatewayAppId() view returns (string)", "function appImplementation() view returns (address)", + "function rootEpochCount() view returns (uint256)", "function currentRootEpoch() view returns (uint256 epoch, bytes32 chainHash)" ]; @@ -36,11 +37,6 @@ export function rootEpochDenial(bootInfo: BootInfo, isKms: boolean, current: Roo return undefined; } -// Contracts deployed before key epochs existed lack the registry getters. -function isMissingFunction(error: unknown): boolean { - return ethers.isError(error, 'CALL_EXCEPTION') || ethers.isError(error, 'BAD_DATA'); -} - export class EthereumBackend { private provider: ethers.JsonRpcProvider; private kmsContract: ethers.Contract; @@ -98,16 +94,13 @@ export class EthereumBackend { } } - // The current root epoch and its chain hash; empty when none is registered - // or the contract predates root epochs. + // The current root epoch and its chain hash; empty only when no epoch is registered. async getRootEpoch(blockTag?: number): Promise { - try { - const [epoch, chainHash] = await this.kmsContract.currentRootEpoch({ blockTag }); - return { rootEpoch: Number(epoch), rootEpochHash: chainHash }; - } catch (error) { - if (isMissingFunction(error)) return {}; - throw error; - } + blockTag ??= await this.provider.getBlockNumber(); + const count = await this.kmsContract.rootEpochCount({ blockTag }); + if (count === 0n) return {}; + const [epoch, chainHash] = await this.kmsContract.currentRootEpoch({ blockTag }); + return { rootEpoch: Number(epoch), rootEpochHash: chainHash }; } async getGatewayAppId(): Promise { diff --git a/dstack/kms/auth-eth/src/main.test.ts b/dstack/kms/auth-eth/src/main.test.ts index 239a43847..e7ad49739 100644 --- a/dstack/kms/auth-eth/src/main.test.ts +++ b/dstack/kms/auth-eth/src/main.test.ts @@ -3,6 +3,7 @@ // SPDX-License-Identifier: Apache-2.0 import { FastifyInstance } from 'fastify'; +import { ethers } from 'ethers'; import { build } from './server'; import { rootEpochDenial } from './ethereum'; import { BootInfo } from './types'; @@ -180,3 +181,120 @@ describe('Server', () => { }); }); }); + +describe('Ethereum root epoch registry', () => { + const { EthereumBackend } = jest.requireActual('./ethereum'); + const abi = new ethers.Interface([ + 'function rootEpochCount() view returns (uint256)', + 'function currentRootEpoch() view returns (uint256, bytes32)', + 'function gatewayAppId() view returns (string)', + 'function appImplementation() view returns (address)', + ]); + const hash = '0x' + '11'.repeat(32); + const bootInfo = { + mrAggregated: '11'.repeat(32), osImageHash: '22'.repeat(32), + appId: '33'.repeat(20), composeHash: '44'.repeat(32), + instanceId: '55'.repeat(20), deviceId: '66'.repeat(32), + }; + const getBlockNumber = jest.fn(); + const call = jest.fn(); + let app: FastifyInstance; + let count: bigint; + let epoch: bigint; + + beforeAll(async () => { + app = await build(); + app.ethereum = new EthereumBackend({ + getBlockNumber, call, getNetwork: async () => ({ chainId: 1337n }), + } as unknown as ethers.JsonRpcProvider, + '0x' + '11'.repeat(20)); + }); + afterAll(async () => { await app.close(); }); + beforeEach(() => { + jest.clearAllMocks(); + count = 2n; + epoch = 1n; + getBlockNumber.mockResolvedValue(100); + call.mockImplementation(async (tx: ethers.TransactionRequest) => { + const fn = abi.getFunction(tx.data!.slice(0, 10)); + switch (fn?.name) { + case 'rootEpochCount': return abi.encodeFunctionResult(fn, [count]); + case 'currentRootEpoch': return abi.encodeFunctionResult(fn, [epoch, hash]); + case 'gatewayAppId': return abi.encodeFunctionResult(fn, ['gateway-app']); + case 'appImplementation': return abi.encodeFunctionResult(fn, ['0x' + '11'.repeat(20)]); + default: return ethers.AbiCoder.defaultAbiCoder().encode(['bool', 'string'], [true, '']); + } + }); + }); + const boot = async (path: string, fields = {}) => { + const response = await app.inject({ + method: 'POST', url: `/bootAuth/${path}`, payload: { ...bootInfo, ...fields }, + }); + expect(response.statusCode).toBe(200); + return response.json(); + }; + const registryCalls = () => call.mock.calls.map(([tx]) => ({ + name: abi.getFunction(tx.data.slice(0, 10))?.name, blockTag: tx.blockTag, + })).filter(({ name }) => name === 'rootEpochCount' || name === 'currentRootEpoch'); + + it('supports an upgraded empty registry without querying the current epoch', async () => { + count = 0n; + expect((await boot('app')).isAllowed).toBe(true); + expect((await boot('kms')).isAllowed).toBe(true); + expect((await boot('app', { rootEpoch: 0, rootEpochHash: hash })).isAllowed).toBe(true); + expect((await boot('app', { rootEpoch: 1, rootEpochHash: hash })).reason).toMatch(/not registered/); + expect(await app.ethereum.getRootEpoch()).toEqual({}); + expect(registryCalls().every(({ name }) => name === 'rootEpochCount')).toBe(true); + }); + + it('allows old KMS app requests only through epoch 0 while preserving onboarding', async () => { + count = 1n; + epoch = 0n; + expect((await boot('app')).isAllowed).toBe(true); + count = 2n; + epoch = 1n; + expect((await boot('app')).reason).toMatch(/reports none/); + expect((await boot('kms')).isAllowed).toBe(true); + }); + + it.each(['app', 'info'])('pins both registry reads to the same block for %s', async (endpoint) => { + if (endpoint === 'app') { + expect((await boot('app', { rootEpoch: 1, rootEpochHash: hash })).isAllowed).toBe(true); + expect(call.mock.calls.every(([tx]) => tx.blockTag === 100)).toBe(true); + } else { + const response = await app.inject({ method: 'GET', url: '/' }); + expect(response.statusCode).toBe(200); + expect(response.json()).toMatchObject({ rootEpoch: 1, rootEpochHash: hash }); + } + expect(registryCalls()).toEqual([ + { name: 'rootEpochCount', blockTag: 100 }, + { name: 'currentRootEpoch', blockTag: 100 }, + ]); + expect(getBlockNumber).toHaveBeenCalledTimes(1); + }); + + describe.each(['rootEpochCount', 'currentRootEpoch'])('%s failures', (getter) => { + it.each(['missing getter', 'revert', 'RPC failure', 'malformed ABI data'])('fails closed on %s', async (failure) => { + const normalCall = call.getMockImplementation()!; + call.mockImplementation(async (tx: ethers.TransactionRequest) => { + if (abi.getFunction(tx.data!.slice(0, 10))?.name !== getter) return normalCall(tx); + if (failure === 'missing getter') return '0x'; + if (failure === 'malformed ABI data') return '0x1234'; + if (failure === 'revert') { + throw ethers.AbiCoder.getBuiltinCallException('call', { + to: '0x' + '11'.repeat(20), data: tx.data!, + }, '0x'); + } + throw new Error('RPC unavailable'); + }); + for (const path of ['app', 'kms']) { + expect(await boot(path)).toEqual({ + isAllowed: false, reason: 'authorization backend unavailable', gatewayAppId: '', + }); + expect((await boot(path, { rootEpoch: 0, rootEpochHash: hash })).isAllowed).toBe(false); + } + const info = await app.inject({ method: 'GET', url: '/' }); + expect(info.statusCode).toBe(500); + }); + }); +}); From 7b76414739c2a0a3a1c170d282e3d3e6d8623d42 Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Fri, 6 Nov 2026 07:59:43 -0800 Subject: [PATCH 006/166] fix(kms): bind epoch evidence updates to their chain hash Signed-off-by: Kevin Wang --- dstack/kms/auth-eth/contracts/DstackKms.sol | 5 ++- dstack/kms/auth-eth/script/Manage.s.sol | 3 +- dstack/kms/auth-eth/script/README.md | 4 +- dstack/kms/auth-eth/test/DstackKms.t.sol | 41 ++++++++++++++++++--- 4 files changed, 43 insertions(+), 10 deletions(-) diff --git a/dstack/kms/auth-eth/contracts/DstackKms.sol b/dstack/kms/auth-eth/contracts/DstackKms.sol index 21afb93eb..76f0fda2b 100644 --- a/dstack/kms/auth-eth/contracts/DstackKms.sol +++ b/dstack/kms/auth-eth/contracts/DstackKms.sol @@ -210,9 +210,10 @@ contract DstackKms is Initializable, Ownable2StepUpgradeable, UUPSUpgradeable, E } // Point an epoch at a new copy of its evidence, for when the old location - // goes away. The evidence proves itself, so only its location changes. - function setRootEpochEvidenceUrl(uint256 epoch, string calldata evidenceUrl) external onlyOwner { + // goes away. Bind the update to the keys in case a pending epoch was replaced. + function setRootEpochEvidenceUrl(uint256 epoch, bytes32 chainHash, string calldata evidenceUrl) external onlyOwner { require(epoch < _rootEpochs.length, "Unknown root epoch"); + require(_rootEpochs[epoch].chainHash == chainHash, "root epoch chain hash mismatch"); require(bytes(evidenceUrl).length > 0, "Empty epoch evidence URL"); _rootEpochs[epoch].evidenceUrl = evidenceUrl; emit RootEpochEvidenceUrlSet(epoch, evidenceUrl); diff --git a/dstack/kms/auth-eth/script/Manage.s.sol b/dstack/kms/auth-eth/script/Manage.s.sol index 0dcedc110..a942207e0 100644 --- a/dstack/kms/auth-eth/script/Manage.s.sol +++ b/dstack/kms/auth-eth/script/Manage.s.sol @@ -139,10 +139,11 @@ contract ActivateRootEpoch is BaseScript { contract SetRootEpochEvidenceUrl is BaseScript { function run() external { uint256 epoch = vm.envUint("ROOT_EPOCH"); + bytes32 chainHash = vm.envBytes32("CHAIN_HASH"); string memory evidenceUrl = vm.envString("EVIDENCE_URL"); vm.startBroadcast(vm.envUint("PRIVATE_KEY")); - kms.setRootEpochEvidenceUrl(epoch, evidenceUrl); + kms.setRootEpochEvidenceUrl(epoch, chainHash, evidenceUrl); vm.stopBroadcast(); console.log("Updated the evidence URL of KMS root epoch", epoch); diff --git a/dstack/kms/auth-eth/script/README.md b/dstack/kms/auth-eth/script/README.md index 9c5630c20..b07283feb 100644 --- a/dstack/kms/auth-eth/script/README.md +++ b/dstack/kms/auth-eth/script/README.md @@ -86,8 +86,10 @@ forge script script/Manage.s.sol:ActivateRootEpoch --broadcast --rpc-url $RPC_UR #### Update a Root Epoch's Evidence URL Points an epoch at a new copy of its attestation, e.g. when the old host goes away. +`CHAIN_HASH` is the expected epoch chain hash (see `Query.s.sol`), so a stale +update cannot overwrite the evidence URL of replacement keys at the same epoch. ```bash -KMS_CONTRACT_ADDR=0x... ROOT_EPOCH=0 EVIDENCE_URL=ipfs://... \ +KMS_CONTRACT_ADDR=0x... ROOT_EPOCH=0 CHAIN_HASH=0x... EVIDENCE_URL=ipfs://... \ forge script script/Manage.s.sol:SetRootEpochEvidenceUrl --broadcast --rpc-url $RPC_URL ``` diff --git a/dstack/kms/auth-eth/test/DstackKms.t.sol b/dstack/kms/auth-eth/test/DstackKms.t.sol index ebb886b8d..4661db858 100644 --- a/dstack/kms/auth-eth/test/DstackKms.t.sol +++ b/dstack/kms/auth-eth/test/DstackKms.t.sol @@ -239,24 +239,53 @@ contract DstackKmsTest is Test { } function test_SetRootEpochEvidenceUrl() public { + bytes32 h0 = chainHash(bytes32(0), ca(0xa0), k256(0xb0)); vm.startPrank(owner); vm.expectRevert("Unknown root epoch"); - kms.setRootEpochEvidenceUrl(0, "ipfs://moved"); + kms.setRootEpochEvidenceUrl(0, h0, "ipfs://moved"); kms.registerRootEpoch(ca(0xa0), k256(0xb0), "ipfs://c0"); vm.expectRevert("Empty epoch evidence URL"); - kms.setRootEpochEvidenceUrl(0, ""); + kms.setRootEpochEvidenceUrl(0, h0, ""); vm.expectEmit(true, true, true, true); emit RootEpochEvidenceUrlSet(0, "https://mirror/c0"); - kms.setRootEpochEvidenceUrl(0, "https://mirror/c0"); + kms.setRootEpochEvidenceUrl(0, h0, "https://mirror/c0"); vm.stopPrank(); - (,, string memory evidenceUrl, bytes32 h) = kms.rootEpoch(0); + (bytes memory caKey, bytes memory k256Key, string memory evidenceUrl, bytes32 h) = kms.rootEpoch(0); assertEq(evidenceUrl, "https://mirror/c0"); - assertEq(h, chainHash(bytes32(0), ca(0xa0), k256(0xb0))); + assertEq(caKey, ca(0xa0)); + assertEq(k256Key, k256(0xb0)); + assertEq(h, h0); vm.prank(user); vm.expectRevert(abi.encodeWithSelector(OwnableUpgradeable.OwnableUnauthorizedAccount.selector, user)); - kms.setRootEpochEvidenceUrl(0, "ipfs://evil"); + kms.setRootEpochEvidenceUrl(0, h0, "ipfs://evil"); + } + + function test_set_root_epoch_evidence_url_rejects_replaced_epoch() public { + vm.startPrank(owner); + kms.registerRootEpoch(ca(0xa0), k256(0xb0), "ipfs://c0"); + kms.registerRootEpoch(ca(0xa1), k256(0xb1), "ipfs://c1"); + (,,, bytes32 oldHash) = kms.rootEpoch(1); + kms.popRootEpoch(1, oldHash); + kms.registerRootEpoch(ca(0xa2), k256(0xb2), "ipfs://replacement"); + (,,, bytes32 replacementHash) = kms.rootEpoch(1); + + vm.expectRevert("root epoch chain hash mismatch"); + kms.setRootEpochEvidenceUrl(1, oldHash, "ipfs://old-evidence"); + kms.activateRootEpoch(1, replacementHash); + vm.expectRevert("root epoch chain hash mismatch"); + kms.setRootEpochEvidenceUrl(1, oldHash, "ipfs://old-evidence"); + (,, string memory evidenceUrl,) = kms.rootEpoch(1); + assertEq(evidenceUrl, "ipfs://replacement"); + + kms.setRootEpochEvidenceUrl(1, replacementHash, "ipfs://replacement-mirror"); + vm.stopPrank(); + (bytes memory caKey, bytes memory k256Key, string memory updatedUrl, bytes32 h) = kms.rootEpoch(1); + assertEq(updatedUrl, "ipfs://replacement-mirror"); + assertEq(caKey, ca(0xa2)); + assertEq(k256Key, k256(0xb2)); + assertEq(h, replacementHash); } function test_KmsInfoFollowsTheCurrentRootEpoch() public { From c9427b5d70ef44f94654648431438a2f42fd71fe Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Fri, 6 Nov 2026 18:56:48 -0800 Subject: [PATCH 007/166] fix(kms): keep evidence URL updates independent of chain hashes Evidence updates are recoverable and mismatched evidence fails verification. Revert the unnecessary chain hash requirement for URL updates; retain it for activation and removal. Signed-off-by: Kevin Wang --- dstack/kms/auth-eth/contracts/DstackKms.sol | 5 +-- dstack/kms/auth-eth/script/Manage.s.sol | 3 +- dstack/kms/auth-eth/script/README.md | 4 +- dstack/kms/auth-eth/test/DstackKms.t.sol | 41 +++------------------ 4 files changed, 10 insertions(+), 43 deletions(-) diff --git a/dstack/kms/auth-eth/contracts/DstackKms.sol b/dstack/kms/auth-eth/contracts/DstackKms.sol index 76f0fda2b..21afb93eb 100644 --- a/dstack/kms/auth-eth/contracts/DstackKms.sol +++ b/dstack/kms/auth-eth/contracts/DstackKms.sol @@ -210,10 +210,9 @@ contract DstackKms is Initializable, Ownable2StepUpgradeable, UUPSUpgradeable, E } // Point an epoch at a new copy of its evidence, for when the old location - // goes away. Bind the update to the keys in case a pending epoch was replaced. - function setRootEpochEvidenceUrl(uint256 epoch, bytes32 chainHash, string calldata evidenceUrl) external onlyOwner { + // goes away. The evidence proves itself, so only its location changes. + function setRootEpochEvidenceUrl(uint256 epoch, string calldata evidenceUrl) external onlyOwner { require(epoch < _rootEpochs.length, "Unknown root epoch"); - require(_rootEpochs[epoch].chainHash == chainHash, "root epoch chain hash mismatch"); require(bytes(evidenceUrl).length > 0, "Empty epoch evidence URL"); _rootEpochs[epoch].evidenceUrl = evidenceUrl; emit RootEpochEvidenceUrlSet(epoch, evidenceUrl); diff --git a/dstack/kms/auth-eth/script/Manage.s.sol b/dstack/kms/auth-eth/script/Manage.s.sol index a942207e0..0dcedc110 100644 --- a/dstack/kms/auth-eth/script/Manage.s.sol +++ b/dstack/kms/auth-eth/script/Manage.s.sol @@ -139,11 +139,10 @@ contract ActivateRootEpoch is BaseScript { contract SetRootEpochEvidenceUrl is BaseScript { function run() external { uint256 epoch = vm.envUint("ROOT_EPOCH"); - bytes32 chainHash = vm.envBytes32("CHAIN_HASH"); string memory evidenceUrl = vm.envString("EVIDENCE_URL"); vm.startBroadcast(vm.envUint("PRIVATE_KEY")); - kms.setRootEpochEvidenceUrl(epoch, chainHash, evidenceUrl); + kms.setRootEpochEvidenceUrl(epoch, evidenceUrl); vm.stopBroadcast(); console.log("Updated the evidence URL of KMS root epoch", epoch); diff --git a/dstack/kms/auth-eth/script/README.md b/dstack/kms/auth-eth/script/README.md index b07283feb..9c5630c20 100644 --- a/dstack/kms/auth-eth/script/README.md +++ b/dstack/kms/auth-eth/script/README.md @@ -86,10 +86,8 @@ forge script script/Manage.s.sol:ActivateRootEpoch --broadcast --rpc-url $RPC_UR #### Update a Root Epoch's Evidence URL Points an epoch at a new copy of its attestation, e.g. when the old host goes away. -`CHAIN_HASH` is the expected epoch chain hash (see `Query.s.sol`), so a stale -update cannot overwrite the evidence URL of replacement keys at the same epoch. ```bash -KMS_CONTRACT_ADDR=0x... ROOT_EPOCH=0 CHAIN_HASH=0x... EVIDENCE_URL=ipfs://... \ +KMS_CONTRACT_ADDR=0x... ROOT_EPOCH=0 EVIDENCE_URL=ipfs://... \ forge script script/Manage.s.sol:SetRootEpochEvidenceUrl --broadcast --rpc-url $RPC_URL ``` diff --git a/dstack/kms/auth-eth/test/DstackKms.t.sol b/dstack/kms/auth-eth/test/DstackKms.t.sol index 4661db858..ebb886b8d 100644 --- a/dstack/kms/auth-eth/test/DstackKms.t.sol +++ b/dstack/kms/auth-eth/test/DstackKms.t.sol @@ -239,53 +239,24 @@ contract DstackKmsTest is Test { } function test_SetRootEpochEvidenceUrl() public { - bytes32 h0 = chainHash(bytes32(0), ca(0xa0), k256(0xb0)); vm.startPrank(owner); vm.expectRevert("Unknown root epoch"); - kms.setRootEpochEvidenceUrl(0, h0, "ipfs://moved"); + kms.setRootEpochEvidenceUrl(0, "ipfs://moved"); kms.registerRootEpoch(ca(0xa0), k256(0xb0), "ipfs://c0"); vm.expectRevert("Empty epoch evidence URL"); - kms.setRootEpochEvidenceUrl(0, h0, ""); + kms.setRootEpochEvidenceUrl(0, ""); vm.expectEmit(true, true, true, true); emit RootEpochEvidenceUrlSet(0, "https://mirror/c0"); - kms.setRootEpochEvidenceUrl(0, h0, "https://mirror/c0"); + kms.setRootEpochEvidenceUrl(0, "https://mirror/c0"); vm.stopPrank(); - (bytes memory caKey, bytes memory k256Key, string memory evidenceUrl, bytes32 h) = kms.rootEpoch(0); + (,, string memory evidenceUrl, bytes32 h) = kms.rootEpoch(0); assertEq(evidenceUrl, "https://mirror/c0"); - assertEq(caKey, ca(0xa0)); - assertEq(k256Key, k256(0xb0)); - assertEq(h, h0); + assertEq(h, chainHash(bytes32(0), ca(0xa0), k256(0xb0))); vm.prank(user); vm.expectRevert(abi.encodeWithSelector(OwnableUpgradeable.OwnableUnauthorizedAccount.selector, user)); - kms.setRootEpochEvidenceUrl(0, h0, "ipfs://evil"); - } - - function test_set_root_epoch_evidence_url_rejects_replaced_epoch() public { - vm.startPrank(owner); - kms.registerRootEpoch(ca(0xa0), k256(0xb0), "ipfs://c0"); - kms.registerRootEpoch(ca(0xa1), k256(0xb1), "ipfs://c1"); - (,,, bytes32 oldHash) = kms.rootEpoch(1); - kms.popRootEpoch(1, oldHash); - kms.registerRootEpoch(ca(0xa2), k256(0xb2), "ipfs://replacement"); - (,,, bytes32 replacementHash) = kms.rootEpoch(1); - - vm.expectRevert("root epoch chain hash mismatch"); - kms.setRootEpochEvidenceUrl(1, oldHash, "ipfs://old-evidence"); - kms.activateRootEpoch(1, replacementHash); - vm.expectRevert("root epoch chain hash mismatch"); - kms.setRootEpochEvidenceUrl(1, oldHash, "ipfs://old-evidence"); - (,, string memory evidenceUrl,) = kms.rootEpoch(1); - assertEq(evidenceUrl, "ipfs://replacement"); - - kms.setRootEpochEvidenceUrl(1, replacementHash, "ipfs://replacement-mirror"); - vm.stopPrank(); - (bytes memory caKey, bytes memory k256Key, string memory updatedUrl, bytes32 h) = kms.rootEpoch(1); - assertEq(updatedUrl, "ipfs://replacement-mirror"); - assertEq(caKey, ca(0xa2)); - assertEq(k256Key, k256(0xb2)); - assertEq(h, replacementHash); + kms.setRootEpochEvidenceUrl(0, "ipfs://evil"); } function test_KmsInfoFollowsTheCurrentRootEpoch() public { From ef26d468a1a56524422114dda0c759836a878029 Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Sun, 8 Nov 2026 01:06:44 -0800 Subject: [PATCH 008/166] fix(kms): remove legacy identity setters Signed-off-by: Kevin Wang --- dstack/kms/auth-eth/contracts/DstackKms.sol | 24 +--- .../contracts/test-utils/LegacyDstackKms.sol | 49 ++++++++ dstack/kms/auth-eth/docs/specification.md | 9 +- dstack/kms/auth-eth/script/Manage.s.sol | 21 ---- dstack/kms/auth-eth/script/README.md | 13 +-- .../auth-eth/test/DstackKms.symbolic.t.sol | 4 - dstack/kms/auth-eth/test/DstackKms.t.sol | 108 ++++++++++-------- 7 files changed, 121 insertions(+), 107 deletions(-) create mode 100644 dstack/kms/auth-eth/contracts/test-utils/LegacyDstackKms.sol diff --git a/dstack/kms/auth-eth/contracts/DstackKms.sol b/dstack/kms/auth-eth/contracts/DstackKms.sol index 21afb93eb..7168572e1 100644 --- a/dstack/kms/auth-eth/contracts/DstackKms.sol +++ b/dstack/kms/auth-eth/contracts/DstackKms.sol @@ -22,7 +22,8 @@ contract DstackKms is Initializable, Ownable2StepUpgradeable, UUPSUpgradeable, E bytes eventlog; } - // KMS information + // Retained for proxy storage compatibility and legacy readers. Only root + // epoch registration and activation may update this identity. KmsInfo public kmsInfo; // The dstack-gateway App ID @@ -139,15 +140,6 @@ contract DstackKms is Initializable, Ownable2StepUpgradeable, UUPSUpgradeable, E emit PolicyChanged(msg.sender, keccak256(bytes(policy)), value, enabled); } - // Function to set KMS information. Once root epochs are registered, - // kmsInfo follows the current one instead. - function setKmsInfo(KmsInfo memory info) external onlyOwner { - require(_rootEpochs.length == 0, "kmsInfo follows the current root epoch"); - kmsInfo = info; - emit KmsInfoSet(info.k256Pubkey); - _emitPolicy("kms-info", keccak256(info.k256Pubkey), true); - } - // Register the root public keys of the next KMS root epoch, with where to // fetch the attestation proving where they were generated. It is not served until // activated, except epoch 0, which is current from the start and must be @@ -251,18 +243,6 @@ contract DstackKms is Initializable, Ownable2StepUpgradeable, UUPSUpgradeable, E return (keys.caPubkey, keys.k256Pubkey, keys.evidenceUrl, keys.chainHash); } - // Function to set KMS quote - function setKmsQuote(bytes memory quote) external onlyOwner { - require(_rootEpochs.length == 0, "kmsInfo follows the current root epoch"); - kmsInfo.quote = quote; - } - - // Function to set KMS eventlog - function setKmsEventlog(bytes memory eventlog) external onlyOwner { - require(_rootEpochs.length == 0, "kmsInfo follows the current root epoch"); - kmsInfo.eventlog = eventlog; - } - // Function to set trusted Gateway App ID function setGatewayAppId(string memory appId) external onlyOwner { gatewayAppId = appId; diff --git a/dstack/kms/auth-eth/contracts/test-utils/LegacyDstackKms.sol b/dstack/kms/auth-eth/contracts/test-utils/LegacyDstackKms.sol new file mode 100644 index 000000000..eac29466e --- /dev/null +++ b/dstack/kms/auth-eth/contracts/test-utils/LegacyDstackKms.sol @@ -0,0 +1,49 @@ +/* + * SPDX-FileCopyrightText: © 2026 Phala Network + * + * SPDX-License-Identifier: Apache-2.0 + */ + +pragma solidity ^0.8.24; + +import "@openzeppelin/contracts-upgradeable/access/Ownable2StepUpgradeable.sol"; +import "@openzeppelin/contracts-upgradeable/proxy/utils/UUPSUpgradeable.sol"; +import "@openzeppelin/contracts-upgradeable/utils/introspection/ERC165Upgradeable.sol"; + +// Pre-epoch KMS storage and identity setter, for migration tests. Keep the +// original storage layout independent of DstackKms so upgrades validate it. +contract LegacyDstackKms is Ownable2StepUpgradeable, UUPSUpgradeable, ERC165Upgradeable { + struct KmsInfo { + bytes k256Pubkey; + bytes caPubkey; + bytes quote; + bytes eventlog; + } + + KmsInfo public kmsInfo; + string public tproxyAppId; + mapping(address => bool) public registeredApps; + mapping(bytes32 => bool) public kmsAllowedAggregatedMrs; + mapping(bytes32 => bool) public kmsAllowedDeviceIds; + mapping(bytes32 => bool) public allowedOsImages; + address public appImplementation; + uint256[50] private __gap; + + /// @custom:oz-upgrades-unsafe-allow constructor + constructor() { + _disableInitializers(); + } + + function initialize(address initialOwner, address implementation) public initializer { + __Ownable_init(initialOwner); + __UUPSUpgradeable_init(); + __ERC165_init(); + appImplementation = implementation; + } + + function setKmsInfo(KmsInfo memory info) external onlyOwner { + kmsInfo = info; + } + + function _authorizeUpgrade(address) internal override onlyOwner { } +} diff --git a/dstack/kms/auth-eth/docs/specification.md b/dstack/kms/auth-eth/docs/specification.md index 95a72e506..41cd2b885 100644 --- a/dstack/kms/auth-eth/docs/specification.md +++ b/dstack/kms/auth-eth/docs/specification.md @@ -52,7 +52,7 @@ return-value spoofing are still possible — see §6. | Slot | Name | Type | Writable by | Notes | |---|---|---|---|---| -| 0–3 | `kmsInfo` | `KmsInfo` (struct, 4 fields) | `setKmsInfo`, `setKmsQuote`, `setKmsEventlog` (owner) | | +| 0–3 | `kmsInfo` | `KmsInfo` (struct, 4 fields) | `registerRootEpoch` (epoch 0), `activateRootEpoch` (owner) | Preserved for proxy upgrades and legacy readers; quote and event log are cleared on activation | | 4 | `gatewayAppId` | `string` | `setGatewayAppId` (owner) | Renamed from `tproxyAppId`; carries `@custom:oz-renamed-from` | | 5 | `registeredApps` | `mapping(address => bool)` | `registerApp` (anyone), `deployAndRegisterApp` (anyone) | Permissionless — see §6.1 | | 6 | `kmsAllowedAggregatedMrs` | `mapping(bytes32 => bool)` | `addKmsAggregatedMr`, `removeKmsAggregatedMr` (owner) | | @@ -167,7 +167,12 @@ Decision (in order): ### 3.7 Owner-only KMS mutations (uniform pattern) -For each of `setKmsInfo`, `setKmsQuote`, `setKmsEventlog`, `setGatewayAppId`, +The legacy `setKmsInfo`, `setKmsQuote`, and `setKmsEventlog` entry points are +removed: calls always revert, including before root epoch registration. Initial +setup uses `registerRootEpoch`. On an upgraded proxy with existing legacy roots, +epoch 0 must match the stored `kmsInfo` keys; the upgrade preserves that storage. + +For each of `setGatewayAppId`, `setAppImplementation`, `addKmsAggregatedMr`, `removeKmsAggregatedMr`, `addKmsDevice`, `removeKmsDevice`, `addOsImageHash`, `removeOsImageHash`: diff --git a/dstack/kms/auth-eth/script/Manage.s.sol b/dstack/kms/auth-eth/script/Manage.s.sol index 0dcedc110..7f57f1d62 100644 --- a/dstack/kms/auth-eth/script/Manage.s.sol +++ b/dstack/kms/auth-eth/script/Manage.s.sol @@ -88,27 +88,6 @@ contract SetGatewayAppId is BaseScript { } } -contract SetKmsInfo is BaseScript { - function run() external { - bytes memory k256Pubkey = vm.envBytes("K256_PUBKEY"); - bytes memory caPubkey = vm.envBytes("CA_PUBKEY"); - bytes memory quote = vm.envBytes("QUOTE"); - bytes memory eventlog = vm.envBytes("EVENTLOG"); - - vm.startBroadcast(vm.envUint("PRIVATE_KEY")); - kms.setKmsInfo( - DstackKms.KmsInfo({ k256Pubkey: k256Pubkey, caPubkey: caPubkey, quote: quote, eventlog: eventlog }) - ); - vm.stopBroadcast(); - - console.log("KMS info set successfully"); - console.log(" K256 Pubkey length:", k256Pubkey.length); - console.log(" CA Pubkey length:", caPubkey.length); - console.log(" Quote length:", quote.length); - console.log(" Eventlog length:", eventlog.length); - } -} - contract RegisterRootEpoch is BaseScript { function run() external { bytes memory caPubkey = vm.envBytes("CA_PUBKEY"); diff --git a/dstack/kms/auth-eth/script/README.md b/dstack/kms/auth-eth/script/README.md index 9c5630c20..f993aa951 100644 --- a/dstack/kms/auth-eth/script/README.md +++ b/dstack/kms/auth-eth/script/README.md @@ -52,15 +52,12 @@ KMS_CONTRACT_ADDR=0x... GATEWAY_APP_ID="my-gateway" \ forge script script/Manage.s.sol:SetGatewayAppId --broadcast --rpc-url $RPC_URL ``` -#### Set KMS Info -Only before root epochs are registered; afterwards `kmsInfo` holds the current -epoch's keys, with no quote or event log. -```bash -KMS_CONTRACT_ADDR=0x... K256_PUBKEY=0x... CA_PUBKEY=0x... QUOTE=0x... EVENTLOG=0x... \ -forge script script/Manage.s.sol:SetKmsInfo --broadcast --rpc-url $RPC_URL -``` - #### Register a KMS Root Epoch +Use this for initial setup as well as subsequent root rotations. The legacy +`setKmsInfo`, `setKmsQuote`, and `setKmsEventlog` entry points have been removed; +calls to them revert even before epoch 0 is registered. For an upgraded proxy +with legacy roots in `kmsInfo`, epoch 0 must register those same public keys. + Appends the next root epoch (the original roots are epoch 0). Epoch 0 takes effect when registered; later epochs take effect when activated, and each takes over `kmsInfo` then, clearing its quote and event log. `EVIDENCE_URL` diff --git a/dstack/kms/auth-eth/test/DstackKms.symbolic.t.sol b/dstack/kms/auth-eth/test/DstackKms.symbolic.t.sol index 573aa8cc2..47b79af2c 100644 --- a/dstack/kms/auth-eth/test/DstackKms.symbolic.t.sol +++ b/dstack/kms/auth-eth/test/DstackKms.symbolic.t.sol @@ -369,10 +369,6 @@ contract DstackKmsSymbolicTest is Test { word ); } else if (which == 11) { - data = abi.encodeWithSelector(DstackKms.setKmsQuote.selector, ""); - } else if (which == 12) { - data = abi.encodeWithSelector(DstackKms.setKmsEventlog.selector, ""); - } else if (which == 13) { data = abi.encodeWithSelector(UUPSUpgradeable.upgradeToAndCall.selector, addr, ""); } else { return; // outside the enumerated KMS-own mutating surface diff --git a/dstack/kms/auth-eth/test/DstackKms.t.sol b/dstack/kms/auth-eth/test/DstackKms.t.sol index ebb886b8d..037374aec 100644 --- a/dstack/kms/auth-eth/test/DstackKms.t.sol +++ b/dstack/kms/auth-eth/test/DstackKms.t.sol @@ -11,6 +11,7 @@ import "openzeppelin-foundry-upgrades/Upgrades.sol"; import "../contracts/DstackKms.sol"; import "../contracts/DstackApp.sol"; import "../contracts/IAppAuth.sol"; +import "../contracts/test-utils/LegacyDstackKms.sol"; contract DstackKmsTest is Test { DstackKms public kms; @@ -50,28 +51,66 @@ contract DstackKmsTest is Test { assertEq(kms.appImplementation(), address(appImpl)); } - function test_SetKmsInfo() public { - bytes memory k256Pubkey = hex"123456"; - bytes memory caPubkey = hex"789abc"; - bytes memory quote = hex"defdef"; - bytes memory eventlog = hex"abcabc"; - - vm.prank(owner); - vm.expectEmit(true, true, true, true); - emit KmsInfoSet(k256Pubkey); - - kms.setKmsInfo( - DstackKms.KmsInfo({ k256Pubkey: k256Pubkey, caPubkey: caPubkey, quote: quote, eventlog: eventlog }) + function upgradeLegacyKms() internal { + vm.startPrank(owner); + address proxy = Upgrades.deployUUPSProxy( + "LegacyDstackKms.sol", abi.encodeCall(LegacyDstackKms.initialize, (owner, address(appImpl))) ); + LegacyDstackKms(proxy).setKmsInfo(LegacyDstackKms.KmsInfo(k256(0xb0), ca(0xa0), hex"0d", hex"0e")); + Options memory opts; + opts.referenceContract = "LegacyDstackKms.sol"; + Upgrades.upgradeProxy(proxy, "DstackKms.sol", "", opts); + vm.stopPrank(); + kms = DstackKms(proxy); - (bytes memory storedK256, bytes memory storedCa, bytes memory storedQuote, bytes memory storedEventlog) = + assertEq(kms.owner(), owner); + assertEq(kms.appImplementation(), address(appImpl)); + assertEq(kms.rootEpochCount(), 0); + (bytes memory storedK256, bytes memory storedCa, bytes memory quote, bytes memory eventlog) = kms.kmsInfo(); + assertEq(storedK256, k256(0xb0)); + assertEq(storedCa, ca(0xa0)); + assertEq(quote, hex"0d"); + assertEq(eventlog, hex"0e"); + } + + function assertLegacySettersDisabled(address caller) internal { + (bytes memory key, bytes memory caKey, bytes memory quote, bytes memory eventlog) = kms.kmsInfo(); + vm.startPrank(caller); + (bool ok,) = address(kms) + .call( + abi.encodeWithSignature( + "setKmsInfo((bytes,bytes,bytes,bytes))", DstackKms.KmsInfo(k256(0xb1), ca(0xa1), hex"ff", hex"ff") + ) + ); + assertFalse(ok); + (ok,) = address(kms).call(abi.encodeWithSignature("setKmsQuote(bytes)", hex"ff")); + assertFalse(ok); + (ok,) = address(kms).call(abi.encodeWithSignature("setKmsEventlog(bytes)", hex"ff")); + assertFalse(ok); + vm.stopPrank(); + (bytes memory storedKey, bytes memory storedCa, bytes memory storedQuote, bytes memory storedEventlog) = kms.kmsInfo(); - assertEq(storedK256, k256Pubkey); - assertEq(storedCa, caPubkey); + assertEq(storedKey, key); + assertEq(storedCa, caKey); assertEq(storedQuote, quote); assertEq(storedEventlog, eventlog); } + function test_legacy_setters_disabled_before_and_after_epoch_registration() public { + assertLegacySettersDisabled(owner); + assertLegacySettersDisabled(user); + vm.prank(owner); + kms.registerRootEpoch(ca(0xa0), k256(0xb0), "ipfs://c0"); + assertLegacySettersDisabled(owner); + assertLegacySettersDisabled(user); + } + + function test_legacy_setters_disabled_after_upgrade() public { + upgradeLegacyKms(); + assertLegacySettersDisabled(owner); + assertLegacySettersDisabled(user); + } + // Well-formed root public keys whose point bytes are all `fill`; shared with // the off-chain implementations' tests. function ca(uint8 fill) internal pure returns (bytes memory key) { @@ -186,8 +225,8 @@ contract DstackKmsTest is Test { } function test_RegisterRootEpochZeroMustMatchKmsInfo() public { + upgradeLegacyKms(); vm.startPrank(owner); - kms.setKmsInfo(DstackKms.KmsInfo(k256(0xb0), ca(0xa0), hex"0d", hex"0e")); vm.expectRevert("Epoch 0 keys differ from kmsInfo"); kms.registerRootEpoch(ca(0xa1), k256(0xb0), "ipfs://c0"); vm.expectRevert("Epoch 0 keys differ from kmsInfo"); @@ -260,8 +299,8 @@ contract DstackKmsTest is Test { } function test_KmsInfoFollowsTheCurrentRootEpoch() public { + upgradeLegacyKms(); vm.startPrank(owner); - kms.setKmsInfo(DstackKms.KmsInfo(k256(0xb0), ca(0xa0), hex"0d", hex"0e")); kms.registerRootEpoch(ca(0xa0), k256(0xb0), "ipfs://c0"); (bytes memory k256Key, bytes memory caKey, bytes memory quote, bytes memory eventlog) = kms.kmsInfo(); assertEq(caKey, ca(0xa0)); @@ -278,13 +317,9 @@ contract DstackKmsTest is Test { assertEq(caKey, ca(0xa1)); assertEq(k256Key, k256(0xb1)); - vm.expectRevert("kmsInfo follows the current root epoch"); - kms.setKmsInfo(DstackKms.KmsInfo(k256(0xb0), ca(0xa0), "", "")); - vm.expectRevert("kmsInfo follows the current root epoch"); - kms.setKmsQuote(hex"0d"); - vm.expectRevert("kmsInfo follows the current root epoch"); - kms.setKmsEventlog(hex"0e"); vm.stopPrank(); + assertLegacySettersDisabled(owner); + assertLegacySettersDisabled(user); } function test_RegisterRootEpochOnlyOwner() public { @@ -319,19 +354,6 @@ contract DstackKmsTest is Test { vm.stopPrank(); } - function test_SetKmsInfoOnlyOwner() public { - bytes memory k256Pubkey = hex"123456"; - bytes memory caPubkey = hex"789abc"; - bytes memory quote = hex"defdef"; - bytes memory eventlog = hex"abcabc"; - - vm.prank(user); - vm.expectRevert(); - kms.setKmsInfo( - DstackKms.KmsInfo({ k256Pubkey: k256Pubkey, caPubkey: caPubkey, quote: quote, eventlog: eventlog }) - ); - } - function test_AddKmsAggregatedMr() public { bytes32 mr = bytes32(uint256(123)); @@ -420,20 +442,6 @@ contract DstackKmsTest is Test { vm.stopPrank(); } - function test_SetKmsQuoteAndEventlog() public { - bytes memory newQuote = hex"deadbeef"; - bytes memory newEventlog = hex"cafebabe"; - - vm.startPrank(owner); - kms.setKmsQuote(newQuote); - kms.setKmsEventlog(newEventlog); - vm.stopPrank(); - - (,, bytes memory storedQuote, bytes memory storedEventlog) = kms.kmsInfo(); - assertEq(storedQuote, newQuote); - assertEq(storedEventlog, newEventlog); - } - function test_IsKmsAllowed() public { bytes32 deviceId = bytes32(uint256(123)); bytes32 mrAggregated = bytes32(uint256(456)); From 92525da89e6539f2afc466c9d8b5f1e48f8d710e Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Sun, 8 Nov 2026 07:08:23 -0800 Subject: [PATCH 009/166] fix(ci): pin SDK tests to supported Python 3.11 Signed-off-by: Kevin Wang --- .github/workflows/sdk.yaml | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/.github/workflows/sdk.yaml b/.github/workflows/sdk.yaml index 550806ff9..c17c8269f 100644 --- a/.github/workflows/sdk.yaml +++ b/.github/workflows/sdk.yaml @@ -84,6 +84,12 @@ jobs: # the entries that are actually shared. save-if: ${{ github.event_name == 'push' }} + - name: Set up Python + if: steps.changes.outputs.relevant == 'true' + uses: actions/setup-python@v5 + with: + python-version: '3.11' + - name: SDK tests if: steps.changes.outputs.relevant == 'true' run: cd sdk && ./run-tests.sh From 21f356f7af85794bc2b9c62121d43f9b9d22ef55 Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Fri, 6 Nov 2026 04:37:42 -0800 Subject: [PATCH 010/166] kms: keep a keyring of all root epochs and serve the current one The KMS stores every root epoch it holds (the newest at the existing paths, older ones under epochs//) and at startup serves the epoch the auth backend reports as current, provided its own keys produce the registered chain hash. Every boot check carries the served epoch and its chain hash so the backend can deny a stale or forked KMS; on a denial the KMS re-reads the current epoch and exits to restart onto it if it moved. GetAppKey takes an optional root_epoch (unset = current) so guests can fetch the keys of an older epoch to migrate, and responses carry the epoch served. Onboarding hands over the whole keyring and checks it against the chain. With rotate, the new KMS also generates the next epoch and attests it; the source must then be at the current epoch and pass the strict KMS authorization check. Signed-off-by: Kevin Wang --- docs/deployment.md | 40 ++ dstack/dstack-util/src/main.rs | 1 + dstack/dstack-util/src/system_setup.rs | 1 + dstack/kms/rpc/proto/kms_rpc.proto | 15 +- dstack/kms/src/config.rs | 30 ++ dstack/kms/src/main.rs | 94 +++- dstack/kms/src/main_service.rs | 316 +++++++++--- .../kms/src/main_service/upgrade_authority.rs | 178 ++++++- dstack/kms/src/onboard_service.rs | 450 ++++++++++++++++-- dstack/kms/src/www/onboard.html | 12 +- 10 files changed, 1006 insertions(+), 131 deletions(-) diff --git a/docs/deployment.md b/docs/deployment.md index ec67cdedd..29ec3334b 100644 --- a/docs/deployment.md +++ b/docs/deployment.md @@ -513,6 +513,46 @@ curl http://:9203/finish > > **Key handover.** It is served as `Admin.GetKmsKey` on the `[core.admin]` listener, which needs `[core.admin.tls]` so it can verify the onboarding KMS's RA-TLS certificate. Pass the source admin URL and its token (`source_token`) in the onboarding request. `core.onboard.public_key_handover = true` keeps the legacy public `KMS.GetKmsKey`. +#### Rotating KMS Root Keys + +KMS root keys are versioned by epoch; epoch 0 is the original roots. Every KMS +keeps all epochs so CVMs can migrate from any of them, and serves the epoch +governance made current. See the [rotation design](./kms-emergency-key-rotation-protocol.md) +for the full procedure; the KMS steps are: + +1. If epochs were never registered, publish the attestation the KMS produced at + bootstrap (kept in `bootstrap-info.json` on the bootstrapped node) and + register the current root public keys as epoch 0: + `DstackKms.registerRootEpoch(caPubkey, k256Pubkey, evidenceUrl)`, via + `script/Manage.s.sol:RegisterRootEpoch`, or a `rootEpochs` entry in the + auth-simple config. Epoch 0 is current once registered, and must be the + `kmsInfo` keys when those are set. A deployment whose `kmsInfo` holds a + legacy `quote` and `eventlog` instead of an attestation converts them with + `dstack-kms legacy-attestation --quote --eventlog ` and publishes + the result first: from then on `kmsInfo` follows the current epoch's keys and + its quote and event log are cleared. +2. Onboard a fresh KMS with `"rotate": true` (the **Rotate** checkbox in the + onboard UI) from a KMS at the current epoch N-1. It imports epochs 0..N-1, + generates epoch N, and returns the new public keys with an attestation whose + report data is + `keccak256("dstack-kms-generated-keys-v2:{N};{ca_pubkey_hex};{k256_pubkey_hex};")`. +3. Verify the attestation, publish it, and register epoch N with its URL. The + cluster keeps serving N-1. +4. Onboard the other KMS nodes from the rotated one; they import 0..N and keep + serving N-1. +5. Activate epoch N (`ActivateRootEpoch`, or `currentRootEpoch` for + auth-simple). Every KMS restarts onto N; a KMS that does not hold N waits + until it is onboarded again. + If the nodes holding N are lost before N is activated, drop it with + `PopRootEpoch` (or remove the `rootEpochs` entry) and rotate again; nodes + that imported the dropped N must be onboarded again. + +Each boot check the KMS sends carries the epoch it serves and the chain hash +committing to the root public keys of epochs 0..epoch; the auth backend denies +it unless that is the current epoch with the registered keys. On a denial the +KMS reads the current epoch from the backend's info endpoint and restarts if it +moved. + --- ## Deploying Apps diff --git a/dstack/dstack-util/src/main.rs b/dstack/dstack-util/src/main.rs index 1046d35d1..f34cb0690 100644 --- a/dstack/dstack-util/src/main.rs +++ b/dstack/dstack-util/src/main.rs @@ -697,6 +697,7 @@ async fn cmd_get_keys(args: GetKeysArgs) -> Result<()> { .get_app_key(dstack_kms_rpc::GetAppKeyRequest { api_version: 1, vm_config: "".to_string(), + root_epoch: None, }) .await .context( diff --git a/dstack/dstack-util/src/system_setup.rs b/dstack/dstack-util/src/system_setup.rs index 0d604076b..03056aa4f 100644 --- a/dstack/dstack-util/src/system_setup.rs +++ b/dstack/dstack-util/src/system_setup.rs @@ -2453,6 +2453,7 @@ impl<'a> Stage0<'a> { .get_app_key(rpc::GetAppKeyRequest { api_version: 1, vm_config: self.shared.sys_config.vm_config.clone(), + root_epoch: None, }) .await .context( diff --git a/dstack/kms/rpc/proto/kms_rpc.proto b/dstack/kms/rpc/proto/kms_rpc.proto index 6853df2be..3bc4af8b4 100644 --- a/dstack/kms/rpc/proto/kms_rpc.proto +++ b/dstack/kms/rpc/proto/kms_rpc.proto @@ -11,6 +11,9 @@ package kms; message GetAppKeyRequest { uint32 api_version = 1; string vm_config = 2; + // 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; } message AppId { @@ -46,6 +49,8 @@ message AppKeyResponse { string gateway_app_id = 7; // OS Image hash bytes os_image_hash = 8; + // Root key epoch all epoch-dependent fields above were derived from. + uint32 root_epoch = 9; } message GetMetaResponse { @@ -74,6 +79,7 @@ message KmsKeys { message KmsKeyResponse { string temp_ca_key = 1; + // Root keys of every epoch the KMS holds; the index is the epoch. repeated KmsKeys keys = 2; } @@ -158,11 +164,18 @@ message OnboardRequest { string domain = 2; // Optional bearer token for the source KMS admin listener. string source_token = 3; + // Import the source's epochs 0..N-1 and generate epoch N locally. The source + // must be at the current epoch N-1. + bool rotate = 4; } message OnboardResponse { - // k256 public key (secp256k1) inherited from source KMS + // k256 public key (secp256k1) of the current epoch bytes k256_pubkey = 1; + // Set only by a rotation: the generated CA public key and the attestation + // binding the new epoch's public keys. + bytes ca_pubkey = 2; + bytes attestation = 3; } // Attestation info needed for on-chain KMS authorization. diff --git a/dstack/kms/src/config.rs b/dstack/kms/src/config.rs index fd2200fb8..f90bdf9ce 100644 --- a/dstack/kms/src/config.rs +++ b/dstack/kms/src/config.rs @@ -21,6 +21,8 @@ const RPC_CERT: &str = "rpc.crt"; const RPC_KEY: &str = "rpc.key"; const RPC_DOMAIN: &str = "rpc-domain"; const K256_KEY: &str = "root-k256.key"; +const ROOT_EPOCH: &str = "root-epoch"; +const EPOCHS_DIR: &str = "epochs"; const BOOTSTRAP_INFO: &str = "bootstrap-info.json"; #[derive(Debug, Clone, Deserialize)] @@ -171,6 +173,28 @@ impl KmsConfig { self.cert_dir.join(K256_KEY) } + /// Records the newest epoch this KMS holds, whose roots are the top-level + /// root files. Absent means 0. + pub fn root_epoch(&self) -> PathBuf { + self.cert_dir.join(ROOT_EPOCH) + } + + /// Root key paths of `epoch` in a keyring whose newest epoch is `latest`: + /// the newest roots are the top-level root files, so paths that predate + /// epochs are unchanged. Older epochs live in `epochs//`. + pub fn root_key_paths(&self, epoch: u32, latest: u32) -> RootKeyPaths { + let dir = if epoch == latest { + self.cert_dir.clone() + } else { + self.cert_dir.join(EPOCHS_DIR).join(epoch.to_string()) + }; + RootKeyPaths { + ca_cert: dir.join(ROOT_CA_CERT), + ca_key: dir.join(ROOT_CA_KEY), + k256_key: dir.join(K256_KEY), + } + } + pub fn bootstrap_info(&self) -> PathBuf { self.cert_dir.join(BOOTSTRAP_INFO) } @@ -209,6 +233,12 @@ pub(crate) struct Dev { pub gateway_app_id: String, } +pub(crate) struct RootKeyPaths { + pub ca_cert: PathBuf, + pub ca_key: PathBuf, + pub k256_key: PathBuf, +} + #[derive(Debug, Clone, Deserialize)] pub(crate) struct OnboardConfig { pub enabled: bool, diff --git a/dstack/kms/src/main.rs b/dstack/kms/src/main.rs index f4bf428a8..14796ad3c 100644 --- a/dstack/kms/src/main.rs +++ b/dstack/kms/src/main.rs @@ -3,7 +3,7 @@ // SPDX-License-Identifier: Apache-2.0 use anyhow::{anyhow, bail, Context, Result}; -use clap::Parser; +use clap::{Parser, Subcommand}; use config::KmsConfig; use main_service::{KmsState, RpcHandler}; use ra_rpc::ratls_client_verifier::RaTlsClientAuth; @@ -35,6 +35,32 @@ struct Args { /// Path to the configuration file #[arg(short, long)] config: Option, + #[command(subcommand)] + command: Option, +} + +#[derive(Subcommand)] +enum Command { + /// Print, as hex, the attestation `registerRootEpoch` takes for epoch 0, + /// converted from the legacy `kmsInfo` quote and event log + LegacyAttestation { + /// TDX quote, hex + #[arg(long)] + quote: String, + /// Event log JSON, hex-encoded as `kmsInfo()` returns it + #[arg(long)] + eventlog: String, + }, +} + +fn legacy_attestation(quote: &str, eventlog: &str) -> Result { + let decode = |hex_str: &str| hex::decode(hex_str.trim().trim_start_matches("0x")); + let quote = decode(quote).context("invalid quote hex")?; + let eventlog = decode(eventlog).context("invalid event log hex")?; + let attestation = dstack_attest::attestation::Attestation::from_tdx_quote(quote, &eventlog)? + .into_versioned() + .to_bytes()?; + Ok(format!("0x{}", hex::encode(attestation))) } async fn run_onboard_service(kms_config: KmsConfig, figment: Figment) -> Result<()> { @@ -123,6 +149,10 @@ async fn main() -> Result<()> { fmt().with_env_filter(filter).with_ansi(false).init(); } let args = Args::parse(); + if let Some(Command::LegacyAttestation { quote, eventlog }) = args.command { + println!("{}", legacy_attestation("e, &eventlog)?); + return Ok(()); + } let figment = config::load_config_figment(args.config.as_deref()); let config: KmsConfig = figment.focus("core").extract()?; @@ -143,8 +173,13 @@ async fn main() -> Result<()> { } } + let keyring = main_service::Keyring::load(&config).context("Failed to load the root keys")?; + let serving_epoch = main_service::wait_for_servable_epoch(&config, &keyring).await; + info!("Serving root epoch {serving_epoch}"); + info!("Updating certs"); - if let Err(err) = onboard_service::update_certs(&config).await { + if let Err(err) = onboard_service::update_certs(&config, serving_epoch, keyring.latest()).await + { if config.attest_rpc_cert { return Err(err).context( "Failed to reissue the attested KMS RPC certificate; refusing to start with a \ @@ -176,7 +211,8 @@ async fn main() -> Result<()> { } else { None }; - let state = main_service::KmsState::new(config).context("Failed to initialize KMS state")?; + let state = main_service::KmsState::new(config, keyring, serving_epoch) + .context("Failed to initialize KMS state")?; let quote_verifier = QuoteVerifier::new(state.attestation_verifier()); let admin_quote_verifier = QuoteVerifier::new(state.attestation_verifier()); let figment = figment @@ -213,6 +249,12 @@ async fn main() -> Result<()> { rocket = rocket.manage(quote_verifier); let main_srv = rocket.launch(); + // A move of the current root epoch on chain ends the process so it can be + // restarted onto the new epoch. + let restart = { + let state = state.clone(); + async move { state.restart_requested().await } + }; match admin_setup { Some((admin_figment, admin_fairing)) => { if admin_config.insecure_no_auth { @@ -233,14 +275,46 @@ async fn main() -> Result<()> { .manage(state) .manage(admin_quote_verifier) .launch(); - tokio::try_join!( - async { main_srv.await.map_err(|err| anyhow!(err.to_string())) }, - async { admin_srv.await.map_err(|err| anyhow!(err.to_string())) }, - )?; - } - None => { - main_srv.await.map_err(|err| anyhow!(err.to_string()))?; + let servers = async { + tokio::try_join!( + async { main_srv.await.map_err(|err| anyhow!(err.to_string())) }, + async { admin_srv.await.map_err(|err| anyhow!(err.to_string())) }, + ) + .map(drop) + }; + tokio::select! { + result = servers => result?, + () = restart => bail!("the current root epoch changed; restarting to follow it"), + } } + None => tokio::select! { + result = main_srv => result.map(drop).map_err(|err| anyhow!(err.to_string()))?, + () = restart => bail!("the current root epoch changed; restarting to follow it"), + }, } Ok(()) } + +#[cfg(test)] +mod tests { + use dstack_attest::attestation::VersionedAttestation; + + #[test] + fn legacy_attestation_converts_a_get_quote_response() { + let fixture: serde_json::Value = serde_json::from_str(include_str!( + "../../verifier/fixtures/tdx-lite-getquote.json" + )) + .unwrap(); + let quote = fixture["quote"].as_str().unwrap(); + let eventlog = hex::encode(fixture["event_log"].as_str().unwrap()); + + let attestation = super::legacy_attestation(quote, &format!("0x{eventlog}")).unwrap(); + let attestation = + VersionedAttestation::from_bytes(&hex::decode(&attestation[2..]).unwrap()).unwrap(); + let attestation = attestation.into_v1(); + assert_eq!( + hex::encode(attestation.report_data().unwrap()), + fixture["report_data"].as_str().unwrap() + ); + } +} diff --git a/dstack/kms/src/main_service.rs b/dstack/kms/src/main_service.rs index d845a82c6..083ba991c 100644 --- a/dstack/kms/src/main_service.rs +++ b/dstack/kms/src/main_service.rs @@ -26,15 +26,17 @@ use ra_tls::{ attestation::{resolve_vm_config, AttestationVerifier, TeeVariant, VerifiedAttestation}, cert::{CaCert, CertRequest, CertSigningRequestV1, CertSigningRequestV2, Csr}, kdf, + rcgen::KeyPair, }; use scale::Decode; use tracing::warn; use upgrade_authority::{ - build_boot_info, ensure_app_id_len, local_kms_boot_info, BootInfo, GetInfoResponse, + build_boot_info, ensure_app_id_len, local_kms_boot_info, BootInfo, CurrentRootEpoch, + GetInfoResponse, RootEpochRegistry, }; use crate::{ - config::KmsConfig, + config::{KmsConfig, RootKeyPaths}, crypto::{derive_k256_key, sign_message, sign_message_with_timestamp}, }; @@ -67,10 +69,105 @@ impl std::ops::Deref for KmsState { } } +/// The root key pair of one epoch. +struct RootKeys { + ca: CaCert, + k256: SigningKey, +} + +impl RootKeys { + fn load(paths: &RootKeyPaths) -> Result { + let ca = CaCert::load(&paths.ca_cert, &paths.ca_key) + .context("failed to load root CA certificate")?; + let key_bytes = fs::read(&paths.k256_key).context("failed to read ECDSA root key")?; + let k256 = SigningKey::from_slice(&key_bytes).context("failed to load ECDSA root key")?; + Ok(Self { ca, k256 }) + } +} + +/// The root keys of epochs 0..=latest this KMS holds, and their chain hashes. +pub(crate) struct Keyring { + roots: Vec, + chain_hashes: Vec<[u8; 32]>, +} + +impl Keyring { + pub(crate) fn load(config: &KmsConfig) -> Result { + let latest = read_root_epoch(config)?; + let roots = (0..=latest) + .map(|epoch| { + RootKeys::load(&config.root_key_paths(epoch, latest)) + .with_context(|| format!("failed to load root keys of epoch {epoch}")) + }) + .collect::>>()?; + let chain_hashes = root_chain_hashes(roots.iter().map(|root| (&root.ca.key, &root.k256))); + Ok(Self { + roots, + chain_hashes, + }) + } + + pub(crate) fn latest(&self) -> u32 { + self.roots.len() as u32 - 1 + } + + pub(crate) fn servable_epoch(&self, registry: RootEpochRegistry) -> Result { + servable_epoch(&self.chain_hashes, registry) + } +} + +/// Chain hashes of the root keys of epochs 0, 1, ... +pub(crate) fn root_chain_hashes<'a>( + roots: impl Iterator, +) -> Vec<[u8; 32]> { + let pubkeys: Vec<_> = roots + .map(|(ca, k256)| { + ( + ca.public_key_der(), + k256.verifying_key().to_sec1_bytes().to_vec(), + ) + }) + .collect(); + dstack_types::root_epoch::chain_hashes(pubkeys.iter().map(|(ca, k256)| (&ca[..], &k256[..]))) +} + +/// The epoch a keyring with `chain_hashes` serves under `registry`: the +/// current epoch, which the keyring must hold with a matching chain hash. A +/// keyring nobody registered is servable only at epoch 0, as before root +/// epochs existed. +pub(crate) fn servable_epoch( + chain_hashes: &[[u8; 32]], + registry: RootEpochRegistry, +) -> Result { + let latest = chain_hashes.len() as u32 - 1; + let current = match registry { + RootEpochRegistry::Unchecked => return Ok(latest), + RootEpochRegistry::Empty if latest == 0 => return Ok(0), + RootEpochRegistry::Empty => { + bail!("the keyring holds root epoch {latest}, but no root epoch is registered") + } + RootEpochRegistry::Current(current) => current, + }; + match chain_hashes.get(current.epoch as usize) { + Some(hash) if *hash == current.chain_hash => Ok(current.epoch), + Some(_) => bail!( + "the root keys up to epoch {} are not the registered ones", + current.epoch + ), + None => bail!( + "root epoch {} is current, but the keyring holds only up to {latest}", + current.epoch + ), + } +} + pub struct KmsStateInner { config: KmsConfig, - root_ca: CaCert, - k256_key: SigningKey, + keyring: Keyring, + /// The epoch served for this process's lifetime; a newer current epoch on + /// chain makes the process restart. + serving_epoch: u32, + restart: tokio::sync::Notify, temp_ca_cert: String, temp_ca_key: String, verifier: CvmVerifier, @@ -127,12 +224,10 @@ impl KmsState { self.verifier.clear_cache(image_hash, config_hash) } - pub fn new(config: KmsConfig) -> Result { - let root_ca = CaCert::load(config.root_ca_cert(), config.root_ca_key()) - .context("Failed to load root CA certificate")?; - let key_bytes = fs::read(config.k256_key()).context("Failed to read ECDSA root key")?; - let k256_key = - SigningKey::from_slice(&key_bytes).context("Failed to load ECDSA root key")?; + pub fn new(config: KmsConfig, keyring: Keyring, serving_epoch: u32) -> Result { + if serving_epoch > keyring.latest() { + bail!("cannot serve root epoch {serving_epoch}, which this KMS does not hold"); + } let temp_ca_key = fs::read_to_string(config.tmp_ca_key()).context("Faeild to read temp ca key")?; let temp_ca_cert = @@ -155,8 +250,9 @@ impl KmsState { Ok(Self { inner: Arc::new(KmsStateInner { config, - root_ca, - k256_key, + keyring, + serving_epoch, + restart: tokio::sync::Notify::new(), temp_ca_cert, temp_ca_key, verifier, @@ -179,6 +275,122 @@ impl KmsState { pub(crate) fn attestation_verifier(&self) -> Arc { self.inner.attestation_verifier.clone() } + + /// The served epoch's root keys. + fn root(&self) -> &RootKeys { + self.epoch_root(self.serving_epoch) + } + + fn epoch_root(&self, epoch: u32) -> &RootKeys { + &self.keyring.roots[epoch as usize] + } + + /// Resolve a root epoch selector: unset is the served epoch. Epochs held + /// but not yet current are not served. + fn resolve_epoch(&self, requested: Option) -> Result { + let serving = self.serving_epoch; + match requested { + None => Ok(serving), + Some(epoch) if epoch <= serving => Ok(epoch), + Some(epoch) => bail!("root epoch {epoch} is not served; the current one is {serving}"), + } + } + + /// Resolves once governance moves the current root epoch away from the + /// served one, so the process can restart to follow it. + pub async fn restart_requested(&self) { + self.restart.notified().await + } + + /// The served epoch and its chain hash, sent with every boot check so + /// the auth backend can refuse a KMS that does not serve the current one. + fn serving_root_epoch(&self) -> CurrentRootEpoch { + CurrentRootEpoch { + epoch: self.serving_epoch, + chain_hash: self.keyring.chain_hashes[self.serving_epoch as usize], + } + } + + /// A denial may mean governance moved the current root epoch. If it did, + /// request a restart, which re-checks the keyring and re-signs the RPC + /// certificate. + pub(crate) async fn restart_if_epoch_moved(&self) { + let info = match self + .auth_api_info + .try_get_with((), self.config.auth_api.get_info()) + .await + { + Ok(info) => info, + Err(err) => { + warn!("failed to read the current root epoch: {err:#}"); + return; + } + }; + let serving = self.serving_epoch; + if self.keyring.servable_epoch(info.root_epochs).ok() != Some(serving) { + warn!( + "governance moved off root epoch {serving} ({:?}); restarting", + info.root_epochs + ); + self.restart.notify_one(); + } + } +} + +/// Wait until governance lets `keyring` serve, and return the epoch to serve. +/// A keyring that predates root epochs keeps serving epoch 0 when the auth API +/// is unreachable, as before. +pub(crate) async fn wait_for_servable_epoch(config: &KmsConfig, keyring: &Keyring) -> u32 { + const RETRY: Duration = Duration::from_secs(10); + loop { + let result = match config.auth_api.get_info().await { + Ok(info) => keyring.servable_epoch(info.root_epochs), + Err(err) if keyring.latest() == 0 => { + warn!("skipping the root epoch check: {err:#}"); + return 0; + } + Err(err) => Err(err.context("failed to read the current root epoch")), + }; + match result { + Ok(epoch) => return epoch, + Err(err) => { + warn!("not serving yet: {err:#}; retrying in {}s", RETRY.as_secs()); + tokio::time::sleep(RETRY).await; + } + } + } +} + +struct AppKeys { + disk_crypt_key: Vec, + env_crypt_key: Vec, + k256_key: Vec, + k256_signature: Vec, +} + +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")?; + 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")?; + Ok(AppKeys { + disk_crypt_key: disk_crypt_key.to_vec(), + env_crypt_key: env_crypt_key.to_vec(), + k256_key: k256_key.to_bytes().to_vec(), + k256_signature, + }) +} + +pub(crate) fn read_root_epoch(config: &KmsConfig) -> Result { + match fs::read_to_string(config.root_epoch()) { + Ok(epoch) => epoch.trim().parse().context("invalid root epoch file"), + Err(err) if err.kind() == std::io::ErrorKind::NotFound => Ok(0), + Err(err) => Err(err).context("failed to read root epoch file"), + } } pub struct RpcHandler { @@ -202,10 +414,16 @@ impl RpcHandler { ensure_key_release_allowed(&info, (&self.state.config).into())?; Ok(KmsKeyResponse { temp_ca_key: self.state.inner.temp_ca_key.clone(), - keys: vec![KmsKeys { - ca_key: self.state.inner.root_ca.key.serialize_pem(), - k256_key: self.state.inner.k256_key.to_bytes().to_vec(), - }], + keys: self + .state + .keyring + .roots + .iter() + .map(|root| KmsKeys { + ca_key: root.ca.key.serialize_pem(), + k256_key: root.k256.to_bytes().to_vec(), + }) + .collect(), }) } } @@ -297,10 +515,11 @@ impl RpcHandler { .state .config .auth_api - .is_app_allowed(&boot_info, true) + .is_app_allowed(&boot_info, true, Some(self.state.serving_root_epoch())) .await .context("Failed to call self KMS auth check")?; if !response.is_allowed { + self.state.restart_if_epoch_moved().await; bail!("KMS is not allowed: {}", response.reason); } Ok(()) @@ -342,9 +561,10 @@ impl RpcHandler { .state .config .auth_api - .is_app_allowed(&boot_info, is_kms) + .is_app_allowed(&boot_info, is_kms, Some(self.state.serving_root_epoch())) .await?; if !response.is_allowed { + self.state.restart_if_epoch_moved().await; bail!("Boot denied: {}", response.reason); } // SNP rootfs/app/config binding is handled by the SNP launch-measurement @@ -369,7 +589,7 @@ impl RpcHandler { fn derive_app_ca(&self, app_id: &[u8]) -> Result { let context_data = vec![app_id, b"app-ca"]; - let app_key = kdf::derive_p256_key_pair(&self.state.root_ca.key, &context_data) + let app_key = kdf::derive_p256_key_pair(&self.state.root().ca.key, &context_data) .context("Failed to derive app disk key")?; let req = CertRequest::builder() .key(&app_key) @@ -381,7 +601,8 @@ impl RpcHandler { .build(); let app_ca = self .state - .root_ca + .root() + .ca .sign(req) .context("Failed to sign App CA")?; Ok(CaCert::from_parts(app_key, app_ca)) @@ -404,36 +625,20 @@ impl KmsRpc for RpcHandler { .await .context("App not allowed")?; ensure_key_release_allowed(&boot_info, (&self.state.config).into())?; - let app_id = boot_info.app_id; - let instance_id = boot_info.instance_id; - let os_image_hash = boot_info.os_image_hash; - - let context_data = vec![&app_id[..], &instance_id[..], b"app-disk-crypt-key"]; - let app_disk_key = kdf::derive_dh_secret(&self.state.root_ca.key, &context_data) - .context("Failed to derive app disk key")?; - let env_crypt_key = { - let secret = - kdf::derive_dh_secret(&self.state.root_ca.key, &[&app_id[..], b"env-encrypt-key"]) - .context("Failed to derive env encrypt key")?; - let secret = x25519_dalek::StaticSecret::from(secret); - secret.to_bytes() - }; - - let (k256_key, k256_signature) = { - let (k256_app_key, signature) = derive_k256_key(&self.state.k256_key, &app_id) - .context("Failed to derive app ecdsa key")?; - (k256_app_key.to_bytes().to_vec(), signature) - }; + let root_epoch = self.state.resolve_epoch(request.root_epoch)?; + let root = self.state.epoch_root(root_epoch); + let keys = derive_app_keys(root, &boot_info.app_id, &boot_info.instance_id)?; Ok(AppKeyResponse { - ca_cert: self.state.root_ca.pem_cert.clone(), - disk_crypt_key: app_disk_key.to_vec(), - env_crypt_key: env_crypt_key.to_vec(), - k256_key, - k256_signature, + ca_cert: root.ca.pem_cert.clone(), + disk_crypt_key: keys.disk_crypt_key, + env_crypt_key: keys.env_crypt_key, + k256_key: keys.k256_key, + k256_signature: keys.k256_signature, tproxy_app_id: gateway_app_id.clone(), gateway_app_id, - os_image_hash, + os_image_hash: boot_info.os_image_hash, + root_epoch, }) } @@ -443,7 +648,7 @@ impl KmsRpc for RpcHandler { .context("KMS self authorization failed")?; ensure_app_id_len(&request.app_id)?; let secret = kdf::derive_dh_secret( - &self.state.root_ca.key, + &self.state.root().ca.key, &[&request.app_id[..], "env-encrypt-key".as_bytes()], ) .context("Failed to derive env encrypt key")?; @@ -458,7 +663,7 @@ impl KmsRpc for RpcHandler { // Legacy signature (without timestamp) for backward compatibility let signature = sign_message( - &self.state.k256_key, + &self.state.root().k256, b"dstack-env-encrypt-pubkey", &request.app_id, &public_key, @@ -467,7 +672,7 @@ impl KmsRpc for RpcHandler { // New signature with timestamp to prevent replay attacks let signature_v1 = sign_message_with_timestamp( - &self.state.k256_key, + &self.state.root().k256, b"dstack-env-encrypt-pubkey", &request.app_id, timestamp, @@ -494,12 +699,12 @@ impl KmsRpc for RpcHandler { .await .map_err(|err| anyhow!("{err:#}"))?; Ok(GetMetaResponse { - ca_cert: self.state.inner.root_ca.pem_cert.clone(), + ca_cert: self.state.root().ca.pem_cert.clone(), allow_any_upgrade: self.state.inner.config.auth_api.is_dev(), k256_pubkey: self .state - .inner - .k256_key + .root() + .k256 .verifying_key() .to_sec1_bytes() .to_vec(), @@ -534,7 +739,7 @@ impl KmsRpc for RpcHandler { Ok(GetTempCaCertResponse { temp_ca_cert: self.state.inner.temp_ca_cert.clone(), temp_ca_key: self.state.inner.temp_ca_key.clone(), - ca_cert: self.state.inner.root_ca.pem_cert.clone(), + ca_cert: self.state.root().ca.pem_cert.clone(), }) } @@ -583,7 +788,7 @@ impl KmsRpc for RpcHandler { certificate_chain: vec![ cert.pem(), app_ca.pem_cert.clone(), - self.state.root_ca.pem_cert.clone(), + self.state.root().ca.pem_cert.clone(), ], }) } @@ -1324,7 +1529,8 @@ mod tests { .focus("core") .extract() .unwrap(); - KmsState::new(config).unwrap() + let keyring = Keyring::load(&config).unwrap(); + KmsState::new(config, keyring, 0).unwrap() } const AUTH_API_INFO: &str = r#"{"status":"ok","kmsContractAddr":"0xkms","ethRpcUrl":"https://rpc.example","gatewayAppId":"0xgateway","chainId":1,"appImplementation":"0ximpl"}"#; diff --git a/dstack/kms/src/main_service/upgrade_authority.rs b/dstack/kms/src/main_service/upgrade_authority.rs index 79f4a7c23..be4c294d5 100644 --- a/dstack/kms/src/main_service/upgrade_authority.rs +++ b/dstack/kms/src/main_service/upgrade_authority.rs @@ -86,6 +86,63 @@ pub(crate) struct BootResponse { pub reason: String, } +/// The root epoch fields: the KMS sends the epoch it serves in boot checks, +/// and auth backends report the current one in info responses. +#[derive(Debug, Clone, Default, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub(crate) struct RootEpochFields { + #[serde(default, skip_serializing_if = "Option::is_none")] + pub root_epoch: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub root_epoch_hash: Option, +} + +impl From for RootEpochFields { + fn from(current: CurrentRootEpoch) -> Self { + Self { + root_epoch: Some(current.epoch), + root_epoch_hash: Some(format!("0x{}", hex::encode(current.chain_hash))), + } + } +} + +impl RootEpochFields { + pub fn registry(&self) -> Result { + match (self.root_epoch, &self.root_epoch_hash) { + (None, None) => Ok(RootEpochRegistry::Empty), + (Some(epoch), Some(hash)) => { + let chain_hash = hex::decode(hash.trim_start_matches("0x")) + .ok() + .and_then(|hash| hash.try_into().ok()) + .context("invalid root epoch chain hash")?; + Ok(RootEpochRegistry::Current(CurrentRootEpoch { + epoch, + chain_hash, + })) + } + _ => bail!("the auth API reports a root epoch without its chain hash"), + } + } +} + +/// The root epoch governance made current, with the chain hash committing to +/// the root public keys of epochs 0..=epoch. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +pub(crate) struct CurrentRootEpoch { + pub epoch: u32, + pub chain_hash: [u8; 32], +} + +/// What governance says about root epochs. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +pub(crate) enum RootEpochRegistry { + /// The dev backend: there is nothing to check against. + Unchecked, + /// No epoch is registered, or the backend predates root epochs. + Empty, + Current(CurrentRootEpoch), +} + #[derive(Debug, Serialize, Deserialize)] #[serde(rename_all = "camelCase")] pub(crate) struct AuthApiInfoResponse { @@ -96,6 +153,8 @@ pub(crate) struct AuthApiInfoResponse { pub gateway_app_id: String, pub chain_id: u64, pub app_implementation: String, + #[serde(flatten)] + pub root_epoch: RootEpochFields, } #[derive(Debug, Clone, Serialize, Deserialize)] @@ -107,6 +166,7 @@ pub(crate) struct GetInfoResponse { pub eth_rpc_url: Option, pub chain_id: Option, pub app_implementation: Option, + pub root_epochs: RootEpochRegistry, } const AUTH_API_CONNECT_TIMEOUT: Duration = Duration::from_secs(5); @@ -167,7 +227,22 @@ async fn send_request(req: reqwest::RequestBuilder, url: &s } impl AuthApi { - pub async fn is_app_allowed(&self, boot_info: &BootInfo, is_kms: bool) -> Result { + /// `root_epoch` is the epoch this KMS serves, which the backend checks is + /// current; unset while onboarding, before the KMS serves any. + pub async fn is_app_allowed( + &self, + boot_info: &BootInfo, + is_kms: bool, + root_epoch: Option, + ) -> Result { + #[derive(Serialize)] + struct BootAuthRequest<'a> { + #[serde(flatten)] + boot_info: &'a BootInfo, + #[serde(flatten)] + root_epoch: RootEpochFields, + } + ensure_identity_widths(boot_info)?; match self { AuthApi::Dev { dev } => Ok(BootResponse { @@ -182,7 +257,11 @@ impl AuthApi { "bootAuth/app" }; let url = url_join(&webhook.url, path); - http_post(&url, &boot_info, webhook.timeout).await + let request = BootAuthRequest { + boot_info, + root_epoch: root_epoch.map(Into::into).unwrap_or_default(), + }; + http_post(&url, &request, webhook.timeout).await } } } @@ -196,6 +275,7 @@ impl AuthApi { gateway_app_id: Some(dev.gateway_app_id.clone()), chain_id: None, app_implementation: None, + root_epochs: RootEpochRegistry::Unchecked, }), AuthApi::Webhook { webhook } => { let info: AuthApiInfoResponse = http_get(&webhook.url, webhook.timeout).await?; @@ -211,6 +291,7 @@ impl AuthApi { chain_id: Some(info.chain_id), gateway_app_id: Some(info.gateway_app_id.clone()), app_implementation: Some(info.app_implementation.clone()), + root_epochs: info.root_epoch.registry()?, }) } } @@ -255,7 +336,7 @@ pub(crate) async fn ensure_self_kms_allowed( .context("failed to build local KMS boot info")?; let response = cfg .auth_api - .is_app_allowed(&boot_info, true) + .is_app_allowed(&boot_info, true, None) .await .context("failed to call KMS auth check")?; if !response.is_allowed { @@ -264,10 +345,13 @@ pub(crate) async fn ensure_self_kms_allowed( Ok(()) } +/// `allow_legacy_cert` accepts a source whose certificate predates the unified +/// attestation format, judging its image by mrAggregated alone. pub(crate) async fn ensure_kms_allowed( cfg: &KmsConfig, attestation: &VerifiedAttestation, verifier: &AttestationVerifier, + allow_legacy_cert: bool, ) -> Result<()> { let mut boot_info = build_boot_info_for_attestation(attestation, false, "") .context("failed to build KMS boot info from attestation")?; @@ -277,6 +361,9 @@ pub(crate) async fn ensure_kms_allowed( // A source that presents a config but no os_image_hash is not legacy and gets no fallback. // TODO: remove once all source KMS instances use the unified PHALA_RATLS_ATTESTATION format. if boot_info.os_image_hash.is_empty() && attestation.config.is_empty() { + if !allow_legacy_cert { + bail!("the source KMS certificate does not carry its OS image hash"); + } warn!("source KMS uses the legacy cert format, falling back to the local os_image_hash"); let local_info = local_kms_boot_info(verifier) .await @@ -285,7 +372,7 @@ pub(crate) async fn ensure_kms_allowed( } let response = cfg .auth_api - .is_app_allowed(&boot_info, true) + .is_app_allowed(&boot_info, true, None) .await .context("failed to call KMS auth check")?; if !response.is_allowed { @@ -295,7 +382,7 @@ pub(crate) async fn ensure_kms_allowed( } #[cfg(test)] -mod tests { +pub(crate) mod tests { use super::*; use crate::config::Webhook; use ra_tls::attestation::TeeVariant; @@ -320,7 +407,7 @@ mod tests { } } - fn serve>( + pub(crate) fn serve>( responses: Vec, ) -> (String, thread::JoinHandle>) { let responses: Vec = responses.into_iter().map(Into::into).collect(); @@ -364,9 +451,12 @@ mod tests { .nth(1) .unwrap() .to_string(); - let body = - serde_json::from_slice(&request[header_end..header_end + content_length]) - .unwrap(); + let body = &request[header_end..header_end + content_length]; + let body = if body.is_empty() { + Value::Null + } else { + serde_json::from_slice(body).unwrap() + }; requests.push((path, body)); write!( stream, @@ -381,7 +471,7 @@ mod tests { (format!("http://{address}"), handle) } - fn webhook(url: String) -> AuthApi { + pub(crate) fn webhook(url: String) -> AuthApi { AuthApi::Webhook { webhook: Webhook { url, @@ -436,7 +526,7 @@ mod tests { let mut aliasing_compose_hash = boot_info(1); aliasing_compose_hash.compose_hash.truncate(31); let err = auth - .is_app_allowed(&aliasing_compose_hash, false) + .is_app_allowed(&aliasing_compose_hash, false, None) .await .unwrap_err(); assert!(err.to_string().contains("compose_hash must be 32 bytes")); @@ -444,7 +534,7 @@ mod tests { let mut aliasing_instance_id = boot_info(1); aliasing_instance_id.instance_id.truncate(19); let err = auth - .is_app_allowed(&aliasing_instance_id, false) + .is_app_allowed(&aliasing_instance_id, false, None) .await .unwrap_err(); assert!(err @@ -455,7 +545,7 @@ mod tests { let mut no_instance_id = boot_info(1); no_instance_id.instance_id.clear(); assert!( - auth.is_app_allowed(&no_instance_id, false) + auth.is_app_allowed(&no_instance_id, false, None) .await .unwrap() .is_allowed @@ -472,8 +562,8 @@ mod tests { let auth = webhook(url); let info = boot_info(1); - let first = auth.is_app_allowed(&info, false).await.unwrap(); - let repeated = auth.is_app_allowed(&info, false).await.unwrap(); + let first = auth.is_app_allowed(&info, false, None).await.unwrap(); + let repeated = auth.is_app_allowed(&info, false, None).await.unwrap(); assert!(first.is_allowed); assert!(!repeated.is_allowed); assert_eq!(repeated.reason, "revoked"); @@ -495,8 +585,18 @@ mod tests { let app = boot_info(10); let kms = boot_info(20); - assert!(auth.is_app_allowed(&app, false).await.unwrap().is_allowed); - assert!(auth.is_app_allowed(&kms, true).await.unwrap().is_allowed); + assert!( + auth.is_app_allowed(&app, false, None) + .await + .unwrap() + .is_allowed + ); + assert!( + auth.is_app_allowed(&kms, true, None) + .await + .unwrap() + .is_allowed + ); let requests = server.join().unwrap(); assert_eq!(requests[0].0, "/bootAuth/app"); @@ -509,6 +609,44 @@ mod tests { assert_ne!(requests[0].1["mrAggregated"], requests[1].1["mrAggregated"]); } + #[rocket::async_test] + async fn root_epochs_are_read_from_info_and_sent_in_boot_checks() { + let hash = "0x".to_string() + &"11".repeat(32); + let (url, server) = serve(vec![ + format!( + r#"{{"status":"ok","kmsContractAddr":"0x00","gatewayAppId":"gw","chainId":1, + "appImplementation":"0x00","rootEpoch":3,"rootEpochHash":"{hash}"}}"# + ), + r#"{"status":"ok","kmsContractAddr":"0x00","gatewayAppId":"gw","chainId":1, + "appImplementation":"0x00","rootEpoch":3}"# + .into(), + r#"{"isAllowed":true,"gatewayAppId":"gw","reason":""}"#.into(), + r#"{"isAllowed":true,"gatewayAppId":"gw","reason":""}"#.into(), + ]); + let auth = webhook(url); + let current = CurrentRootEpoch { + epoch: 3, + chain_hash: [0x11; 32], + }; + + assert_eq!( + auth.get_info().await.unwrap().root_epochs, + RootEpochRegistry::Current(current) + ); + assert!(auth.get_info().await.is_err()); + let info = boot_info(50); + auth.is_app_allowed(&info, false, Some(current)) + .await + .unwrap(); + auth.is_app_allowed(&info, true, None).await.unwrap(); + + let requests = server.join().unwrap(); + assert_eq!(requests[2].1["rootEpoch"], 3); + assert_eq!(requests[2].1["rootEpochHash"], hash); + assert_eq!(requests[2].1["appId"], requests[3].1["appId"]); + assert!(requests[3].1.get("rootEpoch").is_none()); + } + #[rocket::async_test] async fn malformed_backend_response_fails_closed_and_next_request_recovers() { let (url, server) = serve(vec![ @@ -518,9 +656,9 @@ mod tests { let auth = webhook(url); let info = boot_info(30); - let error = auth.is_app_allowed(&info, false).await.unwrap_err(); + let error = auth.is_app_allowed(&info, false, None).await.unwrap_err(); assert!(error.to_string().contains("failed to decode response")); - let recovered = auth.is_app_allowed(&info, false).await.unwrap(); + let recovered = auth.is_app_allowed(&info, false, None).await.unwrap(); assert!(recovered.is_allowed); assert_eq!(recovered.reason, "recovered"); assert_eq!(server.join().unwrap().len(), 2); @@ -537,7 +675,7 @@ mod tests { let (url, server) = serve(vec![body]); let error = webhook(url) - .is_app_allowed(&boot_info(40), false) + .is_app_allowed(&boot_info(40), false, None) .await .unwrap_err(); diff --git a/dstack/kms/src/onboard_service.rs b/dstack/kms/src/onboard_service.rs index 480f9b62a..43a7dde8b 100644 --- a/dstack/kms/src/onboard_service.rs +++ b/dstack/kms/src/onboard_service.rs @@ -36,9 +36,10 @@ use tracing::info; use crate::{ config::KmsConfig, main_service::{ - build_boot_info_for_attestation, + build_boot_info_for_attestation, root_chain_hashes, servable_epoch, upgrade_authority::{ app_attest, dstack_client, ensure_kms_allowed, ensure_self_kms_allowed, pad64, + RootEpochRegistry, }, }, }; @@ -146,16 +147,7 @@ impl OnboardRpc for OnboardHandler { let keys = Keys::generate(&request.domain, self.state.config.attest_rpc_cert) .await .context("Failed to generate keys")?; - - let k256_pubkey = keys.k256_key.verifying_key().to_sec1_bytes().to_vec(); - let ca_pubkey = keys.ca_key.public_key_der(); - let attestation = attest_keys(&ca_pubkey, &k256_pubkey).await?; - - let response = BootstrapResponse { - ca_pubkey, - k256_pubkey, - attestation, - }; + let response = keys.attest().await?; // Store the bootstrap info safe_write(cfg.bootstrap_info(), serde_json::to_vec(&response)?)?; keys.store(cfg)?; @@ -180,13 +172,29 @@ impl OnboardRpc for OnboardHandler { &source_url, &request.source_token, &request.domain, + request.rotate, self.state.attestation_verifier.clone(), ) .await .context("Failed to onboard")?; - let k256_pubkey = keys.k256_key.verifying_key().to_sec1_bytes().to_vec(); + if !request.rotate { + let k256_pubkey = keys.k256_key.verifying_key().to_sec1_bytes().to_vec(); + keys.store(cfg).context("Failed to store keys")?; + return Ok(OnboardResponse { + k256_pubkey, + ..Default::default() + }); + } + // The attestation over the new epoch is what governance registers, so + // it is persisted with the keys and stays readable via GetSetupStatus. + let info = keys.attest().await?; + safe_write(cfg.bootstrap_info(), serde_json::to_vec(&info)?)?; keys.store(cfg).context("Failed to store keys")?; - Ok(OnboardResponse { k256_pubkey }) + Ok(OnboardResponse { + k256_pubkey: info.k256_pubkey, + ca_pubkey: info.ca_pubkey, + attestation: info.attestation, + }) } async fn get_attestation_info(self) -> Result { @@ -287,6 +295,8 @@ mod tests { use crate::main_service::amd_attest::{ compute_expected_measurement, MeasurementInput, OvmfSectionParam, }; + use crate::main_service::upgrade_authority::CurrentRootEpoch; + use dstack_kms_rpc::KmsKeys; use sha2::Digest; fn hex_of(byte: u8, len: usize) -> String { @@ -536,6 +546,234 @@ mod tests { assert!(status.attestation.is_empty()); } + fn handover(epochs: u32) -> KmsKeyResponse { + KmsKeyResponse { + temp_ca_key: KeyPair::generate_for(&PKCS_ECDSA_P256_SHA256) + .unwrap() + .serialize_pem(), + keys: (0..epochs) + .map(|_| KmsKeys { + ca_key: KeyPair::generate_for(&PKCS_ECDSA_P256_SHA256) + .unwrap() + .serialize_pem(), + k256_key: SigningKey::random(&mut rand::rngs::OsRng) + .to_bytes() + .to_vec(), + }) + .collect(), + } + } + + /// The registry a correct handover of `res` would be checked against. + fn current_of(res: &KmsKeyResponse, epoch: u32) -> RootEpochRegistry { + let roots: Vec<_> = res + .keys + .iter() + .map(|k| { + ( + KeyPair::from_pem(&k.ca_key).unwrap().public_key_der(), + SigningKey::from_slice(&k.k256_key) + .unwrap() + .verifying_key() + .to_sec1_bytes() + .to_vec(), + ) + }) + .collect(); + let hashes = dstack_types::root_epoch::chain_hashes( + roots.iter().map(|(ca, k256)| (&ca[..], &k256[..])), + ); + RootEpochRegistry::Current(CurrentRootEpoch { + epoch, + chain_hash: hashes[epoch as usize], + }) + } + + async fn onboard( + cfg: &KmsConfig, + res: KmsKeyResponse, + rotate: bool, + registry: RootEpochRegistry, + ) -> Result { + Keys::from_key_response(cfg, res, "kms.example.com", rotate, registry).await + } + + #[rocket::async_test] + async fn join_imports_and_stores_every_epoch() { + let cert_dir = tempfile::tempdir().unwrap(); + let cfg = auto_bootstrap_config(cert_dir.path()); + let res = handover(3); + let sent: Vec<_> = res.keys.iter().map(|k| k.k256_key.clone()).collect(); + + let keys = onboard(&cfg, res, false, RootEpochRegistry::Unchecked) + .await + .unwrap(); + keys.store(&cfg).unwrap(); + + assert_eq!(crate::main_service::read_root_epoch(&cfg).unwrap(), 2); + for (epoch, k256) in (0..).zip(&sent) { + assert_eq!( + &fs::read(cfg.root_key_paths(epoch, 2).k256_key).unwrap(), + k256 + ); + } + assert_eq!(&fs::read(cfg.k256_key()).unwrap(), &sent[2]); + } + + #[rocket::async_test] + async fn rotation_imports_the_keyring_and_generates_the_next_epoch() { + let cert_dir = tempfile::tempdir().unwrap(); + let cfg = auto_bootstrap_config(cert_dir.path()); + let res = handover(2); + let registry = current_of(&res, 1); + let sent: Vec<_> = res.keys.iter().map(|k| k.k256_key.clone()).collect(); + + let keys = onboard(&cfg, res, true, registry).await.unwrap(); + assert_eq!(keys.epoch, 2); + keys.store(&cfg).unwrap(); + + assert_eq!( + fs::read(cfg.root_key_paths(1, 2).k256_key).unwrap(), + sent[1] + ); + let current = fs::read(cfg.k256_key()).unwrap(); + assert!(!sent.contains(¤t)); + } + + #[rocket::async_test] + async fn onboarding_checks_the_handover_against_the_current_epoch() { + let cert_dir = tempfile::tempdir().unwrap(); + let cfg = auto_bootstrap_config(cert_dir.path()); + let res = handover(2); + let other = handover(2); + + // A join may hold an epoch pending activation beyond the current one. + assert!(onboard(&cfg, res.clone(), false, current_of(&res, 1)) + .await + .is_ok()); + assert!(onboard(&cfg, res.clone(), false, current_of(&res, 0)) + .await + .is_ok()); + // A rotation must start at the current epoch. + assert!(onboard(&cfg, res.clone(), true, current_of(&res, 0)) + .await + .is_err()); + // A keyring other than the registered one. + assert!(onboard(&cfg, res.clone(), false, current_of(&other, 1)) + .await + .is_err()); + assert!(onboard(&cfg, res.clone(), true, current_of(&other, 1)) + .await + .is_err()); + // A source behind the current epoch. + assert!(onboard(&cfg, handover(1), false, current_of(&res, 1)) + .await + .is_err()); + // A missing epoch shifts the rest, so the chain hash no longer matches. + let mut gap = res.clone(); + gap.keys.remove(0); + assert!(onboard(&cfg, gap, false, current_of(&res, 0)) + .await + .is_err()); + assert!( + onboard(&cfg, handover(0), false, RootEpochRegistry::Unchecked) + .await + .is_err() + ); + } + + #[rocket::async_test] + async fn an_unregistered_deployment_can_only_join_at_epoch_0() { + let cert_dir = tempfile::tempdir().unwrap(); + let cfg = auto_bootstrap_config(cert_dir.path()); + assert!(onboard(&cfg, handover(1), false, RootEpochRegistry::Empty) + .await + .is_ok()); + assert!(onboard(&cfg, handover(2), false, RootEpochRegistry::Empty) + .await + .is_err()); + assert!(onboard(&cfg, handover(1), true, RootEpochRegistry::Empty) + .await + .is_err()); + } + + /// A stored keyring at `latest`, loaded back. + async fn stored_keyring(epochs: u32) -> (tempfile::TempDir, KmsConfig, KmsKeyResponse) { + let cert_dir = tempfile::tempdir().unwrap(); + let cfg = auto_bootstrap_config(cert_dir.path()); + let res = handover(epochs); + onboard(&cfg, res.clone(), false, RootEpochRegistry::Unchecked) + .await + .unwrap() + .store(&cfg) + .unwrap(); + (cert_dir, cfg, res) + } + + #[rocket::async_test] + async fn a_keyring_serves_the_current_epoch_it_holds() { + use crate::main_service::Keyring; + + let (_dir, cfg, res) = stored_keyring(3).await; + let keyring = Keyring::load(&cfg).unwrap(); + assert_eq!(keyring.latest(), 2); + assert_eq!(keyring.servable_epoch(current_of(&res, 1)).unwrap(), 1); + assert_eq!(keyring.servable_epoch(current_of(&res, 2)).unwrap(), 2); + assert_eq!( + keyring + .servable_epoch(RootEpochRegistry::Unchecked) + .unwrap(), + 2 + ); + assert!(keyring.servable_epoch(RootEpochRegistry::Empty).is_err()); + assert!(keyring.servable_epoch(current_of(&handover(3), 1)).is_err()); + let mut ahead = handover(4); + ahead.keys[..3].clone_from_slice(&res.keys); + assert!(keyring.servable_epoch(current_of(&ahead, 3)).is_err()); + + let (_dir, cfg, _) = stored_keyring(1).await; + let keyring = Keyring::load(&cfg).unwrap(); + assert_eq!(keyring.servable_epoch(RootEpochRegistry::Empty).unwrap(), 0); + } + + #[rocket::async_test] + async fn a_denial_restarts_the_kms_if_the_current_epoch_moved() { + use crate::main_service::upgrade_authority::tests::{serve, webhook}; + use crate::main_service::{Keyring, KmsState}; + + let (_dir, cfg, res) = stored_keyring(3).await; + let restarts = |registry: RootEpochRegistry| { + let mut cfg = cfg.clone(); + async move { + let epoch = match registry { + RootEpochRegistry::Current(current) => format!( + r#","rootEpoch":{},"rootEpochHash":"{}""#, + current.epoch, + hex::encode(current.chain_hash) + ), + _ => String::new(), + }; + let (url, server) = serve(vec![format!( + r#"{{"status":"ok","kmsContractAddr":"0x00","gatewayAppId":"gw","chainId":1, + "appImplementation":"0x00"{epoch}}}"# + )]); + cfg.auth_api = webhook(url); + let keyring = Keyring::load(&cfg).unwrap(); + let state = KmsState::new(cfg, keyring, 1).unwrap(); + state.restart_if_epoch_moved().await; + server.join().unwrap(); + let restart = + tokio::time::timeout(Duration::from_millis(100), state.restart_requested()); + restart.await.is_ok() + } + }; + + assert!(!restarts(current_of(&res, 1)).await); + assert!(restarts(current_of(&res, 2)).await); + assert!(restarts(RootEpochRegistry::Empty).await); + assert!(restarts(current_of(&handover(3), 1)).await); + } + #[rocket::async_test] async fn setup_status_includes_bootstrap_info_when_present() { let cert_dir = tempfile::tempdir().unwrap(); @@ -607,7 +845,28 @@ mod tests { } } +/// Root keys of an epoch older than the current one. +struct HistoricalRoot { + ca_key: KeyPair, + ca_cert: Certificate, + k256_key: SigningKey, +} + +impl HistoricalRoot { + fn new(ca_key: KeyPair, k256_key: SigningKey) -> Result { + let ca_cert = self_signed_root_ca(&ca_key)?; + Ok(Self { + ca_key, + ca_cert, + k256_key, + }) + } +} + struct Keys { + /// Epoch of the root keys below; `history` holds epochs `0..epoch`. + epoch: u32, + history: Vec, k256_key: SigningKey, tmp_ca_key: KeyPair, tmp_ca_cert: Certificate, @@ -635,6 +894,19 @@ impl Keys { .await } + /// Attest the current root public keys. Epoch 0 keeps the original + /// bootstrap encoding; later epochs also bind the epoch number. + async fn attest(&self) -> Result { + let ca_pubkey = self.ca_key.public_key_der(); + let k256_pubkey = self.k256_key.verifying_key().to_sec1_bytes().to_vec(); + let attestation = attest_keys(self.epoch, &ca_pubkey, &k256_pubkey).await?; + Ok(BootstrapResponse { + ca_pubkey, + k256_pubkey, + attestation, + }) + } + async fn from_keys( tmp_ca_key: KeyPair, ca_key: KeyPair, @@ -651,14 +923,7 @@ impl Keys { .build() .self_signed()?; - // Create self-signed KMS cert - let ca_cert = CertRequest::builder() - .org_name("Dstack") - .subject("Dstack KMS CA") - .ca_level(1) - .key(&ca_key) - .build() - .self_signed()?; + let ca_cert = self_signed_root_ca(&ca_key)?; // The only place the KMS embeds its own attestation. Skipping it lets // the KMS run outside a TEE for development; it does not affect the // verification of quotes presented *to* the KMS, which is a separate @@ -689,6 +954,8 @@ impl Keys { .build() .signed_by(&ca_cert, &ca_key)?; Ok(Keys { + epoch: 0, + history: Vec::new(), k256_key, tmp_ca_key, tmp_ca_cert, @@ -705,6 +972,7 @@ impl Keys { other_kms_url: &str, source_token: &str, domain: &str, + rotate: bool, attestation_verifier: Arc, ) -> Result { let attestation_slot = Arc::new(Mutex::new(None::)); @@ -763,46 +1031,101 @@ impl Keys { }; // The attestation is captured only on the connection that returned the keys; // nothing received is used before the source passes this check. - ensure_kms_allowed(cfg, &source_attestation()?, &attestation_verifier) + ensure_kms_allowed(cfg, &source_attestation()?, &attestation_verifier, !rotate) .await .context("Source KMS is not allowed for onboarding")?; - Self::from_key_response(cfg, keys_res, domain).await + Self::from_registered_key_response(cfg, keys_res, domain, rotate).await + } + + /// [`Self::from_key_response`] after reading the registered epochs, so a + /// source holding any keyring other than the registered one is refused. + async fn from_registered_key_response( + cfg: &KmsConfig, + keys_res: KmsKeyResponse, + domain: &str, + rotate: bool, + ) -> Result { + let registry = cfg + .auth_api + .get_info() + .await + .context("failed to read the current root epoch")? + .root_epochs; + Self::from_key_response(cfg, keys_res, domain, rotate, registry).await } + /// Import the source's whole keyring, epochs 0..=latest; a rotation then + /// generates epoch latest+1. async fn from_key_response( cfg: &KmsConfig, keys_res: KmsKeyResponse, domain: &str, + rotate: bool, + registry: RootEpochRegistry, ) -> Result { - if keys_res.keys.len() != 1 { - return Err(anyhow::anyhow!("Invalid keys")); - } - let keys = keys_res.keys[0].clone(); - let tmp_ca_key_pem = keys_res.temp_ca_key; - let root_ca_key_pem = keys.ca_key; - let root_k256_key = keys.k256_key; + let latest = keys_res + .keys + .len() + .checked_sub(1) + .context("source KMS handed over no keys")? as u32; + let mut roots = keys_res + .keys + .into_iter() + .map(|keys| { + let ca_key = KeyPair::from_pem(&keys.ca_key).context("Failed to parse CA key")?; + let k256_key = + SigningKey::from_slice(&keys.k256_key).context("Failed to parse ECDSA key")?; + Ok((ca_key, k256_key)) + }) + .collect::>>()?; + ensure_handover_registered(&roots, registry, rotate)?; + let epoch = if rotate { latest + 1 } else { latest }; + let (ca_key, k256_key) = if rotate { + ( + KeyPair::generate_for(&PKCS_ECDSA_P256_SHA256)?, + SigningKey::random(&mut rand::rngs::OsRng), + ) + } else { + roots.pop().context("source KMS handed over no keys")? + }; + let history = roots + .into_iter() + .map(|(ca_key, k256_key)| HistoricalRoot::new(ca_key, k256_key)) + .collect::>>()?; let rpc_key = KeyPair::generate_for(&PKCS_ECDSA_P256_SHA256)?; - let ca_key = KeyPair::from_pem(&root_ca_key_pem).context("Failed to parse CA key")?; let tmp_ca_key = - KeyPair::from_pem(&tmp_ca_key_pem).context("Failed to parse tmp CA key")?; - let ecdsa_key = - SigningKey::from_slice(&root_k256_key).context("Failed to parse ECDSA key")?; - Self::from_keys( + KeyPair::from_pem(&keys_res.temp_ca_key).context("Failed to parse tmp CA key")?; + let keys = Self::from_keys( tmp_ca_key, ca_key, rpc_key, - ecdsa_key, + k256_key, domain, cfg.attest_rpc_cert, ) - .await + .await?; + Ok(Self { + epoch, + history, + ..keys + }) } /// The root keys go last: see [`KmsConfig::root_keys_exist`]. fn store(&self, cfg: &KmsConfig) -> Result<()> { + for (epoch, root) in (0..).zip(&self.history) { + let paths = cfg.root_key_paths(epoch, self.epoch); + if let Some(dir) = paths.ca_key.parent() { + fs::create_dir_all(dir)?; + } + safe_write(&paths.ca_cert, root.ca_cert.pem())?; + safe_write_with_mode(&paths.ca_key, root.ca_key.serialize_pem(), 0o600)?; + safe_write_with_mode(&paths.k256_key, root.k256_key.to_bytes(), 0o600)?; + } self.store_certs(cfg)?; safe_write(cfg.rpc_domain(), self.rpc_domain.as_bytes())?; + safe_write(cfg.root_epoch(), self.epoch.to_string())?; self.store_keys(cfg)?; Ok(()) } @@ -823,14 +1146,17 @@ impl Keys { } } -pub(crate) async fn update_certs(cfg: &KmsConfig) -> Result<()> { +/// Reissue the RPC certificate under the root CA of the served `epoch`, so +/// guests reach this KMS through the CA of the keys it serves. +pub(crate) async fn update_certs(cfg: &KmsConfig, epoch: u32, latest: u32) -> Result<()> { + let roots = cfg.root_key_paths(epoch, latest); // Read existing keys let tmp_ca_key = KeyPair::from_pem(&fs::read_to_string(cfg.tmp_ca_key())?)?; - let ca_key = KeyPair::from_pem(&fs::read_to_string(cfg.root_ca_key())?)?; + let ca_key = KeyPair::from_pem(&fs::read_to_string(&roots.ca_key)?)?; let rpc_key = KeyPair::from_pem(&fs::read_to_string(cfg.rpc_key())?)?; // Read k256 key - let k256_key_bytes = fs::read(cfg.k256_key())?; + let k256_key_bytes = fs::read(&roots.k256_key)?; let k256_key = SigningKey::from_slice(&k256_key_bytes)?; let domain = if cfg.onboard.auto_bootstrap_domain.is_empty() { @@ -853,7 +1179,7 @@ pub(crate) async fn update_certs(cfg: &KmsConfig) -> Result<()> { .context("Failed to regenerate certificates")?; renew_ca_cert_if_expiring( - cfg.root_ca_cert(), + &roots.ca_cert, keys.ca_cert.pem(), "KMS root CA certificate", )?; @@ -922,10 +1248,48 @@ pub(crate) async fn bootstrap_keys(cfg: &KmsConfig, verifier: &AttestationVerifi Ok(()) } -async fn attest_keys(p256_pubkey: &[u8], k256_pubkey: &[u8]) -> Result> { +/// Check the handed-over roots against governance: they must include the +/// current epoch with a matching chain hash. A rotation must start from the +/// current epoch with nothing pending activation, or it would collide with the +/// next registered epoch. +fn ensure_handover_registered( + roots: &[(KeyPair, SigningKey)], + registry: RootEpochRegistry, + rotate: bool, +) -> Result<()> { + if rotate && registry == RootEpochRegistry::Empty { + bail!("register root epoch 0 before rotating"); + } + let hashes = root_chain_hashes(roots.iter().map(|(ca, k256)| (ca, k256))); + let current = servable_epoch(&hashes, registry)?; + let latest = roots.len() as u32 - 1; + if rotate && current != latest { + bail!( + "the source holds root epochs up to {latest}, but {current} is current; rotate from \ + a source at the current epoch with no epoch pending activation" + ); + } + Ok(()) +} + +fn self_signed_root_ca(ca_key: &KeyPair) -> Result { + CertRequest::builder() + .org_name("Dstack") + .subject("Dstack KMS CA") + .ca_level(1) + .key(ca_key) + .build() + .self_signed() +} + +async fn attest_keys(epoch: u32, p256_pubkey: &[u8], k256_pubkey: &[u8]) -> Result> { let p256_hex = hex::encode(p256_pubkey); let k256_hex = hex::encode(k256_pubkey); - let content_to_quote = format!("dstack-kms-genereted-keys-v1:{p256_hex};{k256_hex};"); + let content_to_quote = if epoch == 0 { + format!("dstack-kms-genereted-keys-v1:{p256_hex};{k256_hex};") + } else { + format!("dstack-kms-generated-keys-v2:{epoch};{p256_hex};{k256_hex};") + }; let hash = keccak256(content_to_quote.as_bytes()); let report_data = pad64(hash); let res = app_attest(report_data).await?; diff --git a/dstack/kms/src/www/onboard.html b/dstack/kms/src/www/onboard.html index 4794cf724..d476e64e8 100644 --- a/dstack/kms/src/www/onboard.html +++ b/dstack/kms/src/www/onboard.html @@ -261,6 +261,12 @@

Onboard from an Existing KMS Instance

+
+ +
Onboard from an Existing KMS Instance sourceUrl: '', sourceToken: '', onboardDomain: '', + rotate: false, error: '', success: '', result: '', @@ -350,7 +357,7 @@

Onboard from an Existing KMS Instance

return; } if (data.attestation) { - this.success = 'Bootstrap successful!'; + this.success = 'New root keys generated!'; this.result = JSON.stringify({ caPubkey: '0x' + data.ca_pubkey, k256Pubkey: '0x' + data.k256_pubkey, @@ -391,7 +398,8 @@

Onboard from an Existing KMS Instance

const data = await rpcCall('Onboard', { source_url: this.sourceUrl + "/prpc", source_token: this.sourceToken, - domain: this.onboardDomain + domain: this.onboardDomain, + rotate: this.rotate }); if (data.error) throw new Error(data.error); From 533d18de283421746331c0d85db9b44284ab6d09 Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Fri, 6 Nov 2026 08:50:40 -0800 Subject: [PATCH 011/166] feat(kms): unify authorization and return the current root epoch Signed-off-by: Kevin Wang --- docs/deployment.md | 16 +- dstack/kms/README.md | 44 +++++ dstack/kms/auth-eth-bun/README.md | 10 ++ dstack/kms/auth-eth-bun/index.test.ts | 35 ++++ dstack/kms/auth-eth-bun/index.ts | 90 +++++++--- dstack/kms/auth-eth-bun/openapi.json | 78 +++++++- dstack/kms/auth-eth/README.md | 10 ++ dstack/kms/auth-eth/contracts/DstackKms.sol | 56 ++++++ dstack/kms/auth-eth/src/ethereum.ts | 22 ++- dstack/kms/auth-eth/src/main.test.ts | 22 +++ dstack/kms/auth-eth/src/server.ts | 21 ++- dstack/kms/auth-eth/src/types.ts | 8 + dstack/kms/auth-eth/test/DstackKms.t.sol | 68 +++++++ dstack/kms/auth-mock/README.md | 10 ++ dstack/kms/auth-mock/index.test.ts | 29 +++ dstack/kms/auth-mock/index.ts | 23 +++ dstack/kms/auth-mock/openapi.json | 88 ++++++++- dstack/kms/auth-simple/README.md | 10 ++ dstack/kms/auth-simple/index.test.ts | 28 +++ dstack/kms/auth-simple/index.ts | 24 +++ dstack/kms/src/main_service.rs | 170 ++++++++++++++---- .../kms/src/main_service/upgrade_authority.rs | 87 ++++++++- dstack/kms/src/onboard_service.rs | 38 ---- 23 files changed, 871 insertions(+), 116 deletions(-) diff --git a/docs/deployment.md b/docs/deployment.md index 29ec3334b..4ec2d5aaf 100644 --- a/docs/deployment.md +++ b/docs/deployment.md @@ -541,17 +541,21 @@ for the full procedure; the KMS steps are: 4. Onboard the other KMS nodes from the rotated one; they import 0..N and keep serving N-1. 5. Activate epoch N (`ActivateRootEpoch`, or `currentRootEpoch` for - auth-simple). Every KMS restarts onto N; a KMS that does not hold N waits + auth-simple). Each KMS restarts onto N on its next uncached authorization + check; a KMS that does not hold N waits until it is onboarded again. If the nodes holding N are lost before N is activated, drop it with `PopRootEpoch` (or remove the `rootEpochs` entry) and rotate again; nodes that imported the dropped N must be onboarded again. -Each boot check the KMS sends carries the epoch it serves and the chain hash -committing to the root public keys of epochs 0..epoch; the auth backend denies -it unless that is the current epoch with the registered keys. On a denial the -KMS reads the current epoch from the backend's info endpoint and restarts if it -moved. +Each `checkAuthorization` result carries the current epoch and its chain hash, +including when authorization is denied. The KMS compares these directly with +its serving roots and requests a restart on a mismatch; it does not read the +cached info endpoint to decide. No registered epoch returns a rejection with +`rejectReason: "Root epoch not registered"` and a zero hash, which does not +trigger a restart. Register epoch 0 before normal key release. Bootstrap keeps +using scoped KMS authorization so it can generate the first roots before +registration. See the [authorization protocol](../dstack/kms/README.md#authorization-protocol). --- diff --git a/dstack/kms/README.md b/dstack/kms/README.md index 49d15e618..aff3f93aa 100644 --- a/dstack/kms/README.md +++ b/dstack/kms/README.md @@ -242,3 +242,47 @@ The `SignCert` RPC is used by the dstack app to sign a TLS certificate. In this - Verify the CSR signature - Query the smart contract to check if the app is authorized - If authorized, sign the CSR with the CA root key and return the certificate chain to the app + + +## Authorization protocol + +Normal KMS operations use `POST /checkAuthorization`. The request carries the +verified attestation identity (the same identity fields as the legacy +`/bootAuth/*` requests). The Ethereum backends call +`DstackKms.checkAuthorization(AppInfo)` once; the contract selects KMS policy when +`appId == address(this)` (the proxy address), and app policy otherwise. KMS-only +operations also check this identity before accepting the result. + +The response is an `AuthorizationResult`: + +```json +{ + "isAllowed": true, + "rejectReason": "", + "gatewayAppId": "0x...", + "currentRootEpoch": 0, + "currentRootEpochHash": "0x..." +} +``` + +Both successful and rejected decisions include the current epoch and chain hash +from the same contract call or auth-simple configuration snapshot. With no +registered epoch, the result is `isAllowed: false`, +`rejectReason: "Root epoch not registered"`, epoch 0 and a zero hash. A zero hash +never authorizes key release or requests an epoch switch. + +The KMS compares each result with its serving epoch and chain hash before +releasing keys or signing certificates. A mismatch requests a process restart, +including when policy also denies authorization. This does not consult the +cached info endpoint. Rotation is detected on authorization requests, not by +background polling; self-authorization retains its one-second cache. + +Upgrade the contract and auth backend and register epoch 0 before using this +protocol for key release. Runtime authorization never falls back to the legacy +API if the new endpoint or contract method is unavailable. The old contract +methods and `/bootAuth/app` and `/bootAuth/kms` endpoints remain available for +older clients. Bootstrap and onboarding retain scoped KMS checks so the first +root keys can be generated and attested before registering epoch 0. For +`auth-simple`, set `kmsContractAddr` to the attested KMS app ID and configure +`rootEpochs`; for `auth-mock`, set `KMS_CONTRACT_ADDR`, `MOCK_ROOT_EPOCH`, and +`MOCK_ROOT_EPOCH_HASH`. diff --git a/dstack/kms/auth-eth-bun/README.md b/dstack/kms/auth-eth-bun/README.md index 8fcafb0bd..92f6842e4 100644 --- a/dstack/kms/auth-eth-bun/README.md +++ b/dstack/kms/auth-eth-bun/README.md @@ -199,3 +199,13 @@ the test suite (`index.test.ts`) validates: - ✅ error handling behavior - ✅ hex encoding/decoding compatibility - ✅ optional field handling + + +### Unified authorization + +`POST /checkAuthorization` returns `AuthorizationResult`, including +`rejectReason`, `gatewayAppId`, `currentRootEpoch`, and `currentRootEpochHash`. +It selects KMS or app policy from the verified `appId`. See the +[KMS authorization protocol](../README.md#authorization-protocol) for registration, +bootstrap, and compatibility requirements. The legacy `/bootAuth/*` endpoints +remain available. diff --git a/dstack/kms/auth-eth-bun/index.test.ts b/dstack/kms/auth-eth-bun/index.test.ts index 22af72fc8..073b8a849 100644 --- a/dstack/kms/auth-eth-bun/index.test.ts +++ b/dstack/kms/auth-eth-bun/index.test.ts @@ -832,3 +832,38 @@ describe('request hardening', () => { expect(response.headers.get('x-frame-options')).toBe('SAMEORIGIN'); }); }); + + +describe('POST /checkAuthorization', () => { + const appInfo = { + appId: '0x' + '11'.repeat(20), composeHash: '0x' + '22'.repeat(32), + instanceId: '', deviceId: '0x' + '33'.repeat(32), + mrAggregated: '0x' + '44'.repeat(32), mrSystem: '0x' + '55'.repeat(32), + osImageHash: '0x' + '66'.repeat(32), tcbStatus: 'UpToDate', advisoryIds: [], + }; + const authorize = () => appFetch(new Request('http://localhost/checkAuthorization', { + method: 'POST', headers: { 'Content-Type': 'application/json' }, body: JSON.stringify(appInfo), + })); + + it('reads the complete result with one contract call at configured finality', async () => { + const result = { + isAllowed: false, rejectReason: 'App not registered', gatewayAppId: 'gateway', + currentRootEpoch: 2n, currentRootEpochHash: '0x' + '77'.repeat(32), + }; + mockReadContract.mockResolvedValueOnce(result); + const response = await authorize(); + expect(response.status).toBe(200); + expect(await response.json()).toEqual({ ...result, currentRootEpoch: 2 }); + expect(mockReadContract).toHaveBeenCalledTimes(1); + expect(mockReadContract).toHaveBeenCalledWith(expect.objectContaining({ + functionName: 'checkAuthorization', blockNumber: 98n, + args: [expect.objectContaining({ appId: appInfo.appId })], + })); + }); + + it('does not fall back to legacy authorization on a failed contract call', async () => { + mockReadContract.mockRejectedValueOnce(new Error('contract unavailable')); + expect((await authorize()).status).toBe(503); + expect(mockReadContract).toHaveBeenCalledTimes(1); + }); +}); diff --git a/dstack/kms/auth-eth-bun/index.ts b/dstack/kms/auth-eth-bun/index.ts index f7d726377..796d83b06 100644 --- a/dstack/kms/auth-eth-bun/index.ts +++ b/dstack/kms/auth-eth-bun/index.ts @@ -71,8 +71,37 @@ function rootEpochDenial(bootInfo: BootInfo, isKms: boolean, current: RootEpochI return undefined; } +const APP_INFO_COMPONENTS = [ + { name: 'appId', type: 'address' }, + { name: 'composeHash', type: 'bytes32' }, + { name: 'instanceId', type: 'address' }, + { name: 'deviceId', type: 'bytes32' }, + { name: 'mrAggregated', type: 'bytes32' }, + { name: 'mrSystem', type: 'bytes32' }, + { name: 'osImageHash', type: 'bytes32' }, + { name: 'tcbStatus', type: 'string' }, + { name: 'advisoryIds', type: 'string[]' } +] as const; + // DstackKms contract ABI (minimal required functions) const DSTACK_KMS_ABI = [ + { + name: 'checkAuthorization', + type: 'function', + stateMutability: 'view', + inputs: [{ name: 'appInfo', type: 'tuple', components: APP_INFO_COMPONENTS }], + outputs: [{ + name: 'result', + type: 'tuple', + components: [ + { name: 'isAllowed', type: 'bool' }, + { name: 'rejectReason', type: 'string' }, + { name: 'gatewayAppId', type: 'string' }, + { name: 'currentRootEpoch', type: 'uint256' }, + { name: 'currentRootEpochHash', type: 'bytes32' }, + ], + }], + }, { name: 'isKmsAllowed', type: 'function', @@ -81,17 +110,7 @@ const DSTACK_KMS_ABI = [ { name: 'bootInfo', type: 'tuple', - components: [ - { name: 'appId', type: 'address' }, - { name: 'composeHash', type: 'bytes32' }, - { name: 'instanceId', type: 'address' }, - { name: 'deviceId', type: 'bytes32' }, - { name: 'mrAggregated', type: 'bytes32' }, - { name: 'mrSystem', type: 'bytes32' }, - { name: 'osImageHash', type: 'bytes32' }, - { name: 'tcbStatus', type: 'string' }, - { name: 'advisoryIds', type: 'string[]' } - ] + components: APP_INFO_COMPONENTS } ], outputs: [ @@ -107,17 +126,7 @@ const DSTACK_KMS_ABI = [ { name: 'bootInfo', type: 'tuple', - components: [ - { name: 'appId', type: 'address' }, - { name: 'composeHash', type: 'bytes32' }, - { name: 'instanceId', type: 'address' }, - { name: 'deviceId', type: 'bytes32' }, - { name: 'mrAggregated', type: 'bytes32' }, - { name: 'mrSystem', type: 'bytes32' }, - { name: 'osImageHash', type: 'bytes32' }, - { name: 'tcbStatus', type: 'string' }, - { name: 'advisoryIds', type: 'string[]' } - ] + components: APP_INFO_COMPONENTS } ], outputs: [ @@ -198,9 +207,9 @@ class EthereumBackend { return `0x${hex}` as Hex; } - async checkBoot(bootInfo: BootInfo, isKms: boolean): Promise { + private contractAppInfo(bootInfo: BootInfo) { // create boot info struct for contract call - const bootInfoStruct = { + return { appId: this.decodeHex(bootInfo.appId, 20) as Address, composeHash: this.decodeHex(bootInfo.composeHash, 32), instanceId: this.decodeHex(bootInfo.instanceId, 20) as Address, @@ -212,6 +221,26 @@ class EthereumBackend { advisoryIds: bootInfo.advisoryIds || [] }; + } + + async checkAuthorization(appInfo: BootInfo) { + const blockNumber = await this.finalizedBlockNumber(); + const result = await this.client.readContract({ + address: this.kmsContractAddr, + abi: DSTACK_KMS_ABI, + functionName: 'checkAuthorization', + args: [this.contractAppInfo(appInfo)], + blockNumber, + }); + const currentRootEpoch = Number(result.currentRootEpoch); + if (!Number.isSafeInteger(currentRootEpoch)) { + throw new Error('current root epoch exceeds the supported range'); + } + return { ...result, currentRootEpoch }; + } + + async checkBoot(bootInfo: BootInfo, isKms: boolean): Promise { + const bootInfoStruct = this.contractAppInfo(bootInfo); const blockNumber = await this.finalizedBlockNumber(); let response; if (isKms) { @@ -370,6 +399,19 @@ app.get('/', async (c) => { } }); +app.post('/checkAuthorization', + bodyLimit({ maxSize: MAX_BOOT_INFO_BYTES }), + zValidator('json', BootInfoSchema), + async (c) => { + try { + return c.json(await ethereum.checkAuthorization(c.req.valid('json'))); + } catch { + console.error('authorization backend failed'); + return c.json({ status: 'error', message: backendUnavailable }, 503); + } + } +); + // app boot authentication app.post('/bootAuth/app', zValidator('json', BootInfoSchema, (result, c) => { diff --git a/dstack/kms/auth-eth-bun/openapi.json b/dstack/kms/auth-eth-bun/openapi.json index a9619eb0a..f7e222054 100644 --- a/dstack/kms/auth-eth-bun/openapi.json +++ b/dstack/kms/auth-eth-bun/openapi.json @@ -98,6 +98,37 @@ } } } + }, + "/checkAuthorization": { + "post": { + "summary": "Authorize an attested identity and return the current root epoch", + "operationId": "checkAuthorization", + "requestBody": { + "required": true, + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/AppInfo" + } + } + } + }, + "responses": { + "200": { + "description": "Authorization decision, including the epoch on rejection.", + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/AuthorizationResult" + } + } + } + }, + "503": { + "description": "Authorization backend unavailable; no legacy fallback." + } + } + } } }, "components": { @@ -198,7 +229,9 @@ "properties": { "status": { "type": "string", - "enum": ["ok"], + "enum": [ + "ok" + ], "description": "System status" }, "kmsContractAddr": { @@ -244,13 +277,54 @@ "properties": { "status": { "type": "string", - "enum": ["error"] + "enum": [ + "error" + ] }, "message": { "type": "string", "description": "Error message" } } + }, + "AppInfo": { + "allOf": [ + { + "$ref": "#/components/schemas/BootInfo" + } + ], + "description": "Verified attestation identity. The KMS app ID is the KMS contract address." + }, + "AuthorizationResult": { + "type": "object", + "required": [ + "isAllowed", + "rejectReason", + "gatewayAppId", + "currentRootEpoch", + "currentRootEpochHash" + ], + "properties": { + "isAllowed": { + "type": "boolean" + }, + "rejectReason": { + "type": "string", + "description": "Empty when allowed." + }, + "gatewayAppId": { + "type": "string" + }, + "currentRootEpoch": { + "type": "integer", + "minimum": 0 + }, + "currentRootEpochHash": { + "type": "string", + "pattern": "^0x[0-9a-fA-F]{64}$", + "description": "Zero when no epoch is registered; authorization is denied." + } + } } } } diff --git a/dstack/kms/auth-eth/README.md b/dstack/kms/auth-eth/README.md index bb543b6dd..d8a91fad4 100644 --- a/dstack/kms/auth-eth/README.md +++ b/dstack/kms/auth-eth/README.md @@ -170,3 +170,13 @@ anvil ``` Documentation: https://book.getfoundry.sh/ + + +### Unified authorization + +`POST /checkAuthorization` returns `AuthorizationResult`, including +`rejectReason`, `gatewayAppId`, `currentRootEpoch`, and `currentRootEpochHash`. +It selects KMS or app policy from the verified `appId`. See the +[KMS authorization protocol](../README.md#authorization-protocol) for registration, +bootstrap, and compatibility requirements. The legacy `/bootAuth/*` endpoints +remain available. diff --git a/dstack/kms/auth-eth/contracts/DstackKms.sol b/dstack/kms/auth-eth/contracts/DstackKms.sol index 7168572e1..7ce8a53f7 100644 --- a/dstack/kms/auth-eth/contracts/DstackKms.sol +++ b/dstack/kms/auth-eth/contracts/DstackKms.sol @@ -14,6 +14,26 @@ import "@openzeppelin/contracts-upgradeable/utils/introspection/ERC165Upgradeabl import "@openzeppelin/contracts/proxy/ERC1967/ERC1967Proxy.sol"; contract DstackKms is Initializable, Ownable2StepUpgradeable, UUPSUpgradeable, ERC165Upgradeable, IAppAuth { + struct AppInfo { + address appId; + bytes32 composeHash; + address instanceId; + bytes32 deviceId; + bytes32 mrAggregated; + bytes32 mrSystem; + bytes32 osImageHash; + string tcbStatus; + string[] advisoryIds; + } + + struct AuthorizationResult { + bool isAllowed; + string rejectReason; + string gatewayAppId; + uint256 currentRootEpoch; + bytes32 currentRootEpochHash; + } + // Struct for KMS information struct KmsInfo { bytes k256Pubkey; @@ -370,8 +390,40 @@ contract DstackKms is Initializable, Ownable2StepUpgradeable, UUPSUpgradeable, E _emitPolicy("os-image", osImageHash, false); } + /// Authorize an attested identity and return the root epoch from the same state. + /// The KMS app ID is this contract's proxy address; all other IDs use app policy. + function checkAuthorization(AppInfo calldata appInfo) external view returns (AuthorizationResult memory result) { + result.gatewayAppId = gatewayAppId; + if (_rootEpochs.length == 0) { + result.rejectReason = "Root epoch not registered"; + return result; + } + result.currentRootEpoch = _currentRootEpoch; + result.currentRootEpochHash = _rootEpochs[_currentRootEpoch].chainHash; + AppBootInfo memory bootInfo = AppBootInfo({ + appId: appInfo.appId, + composeHash: appInfo.composeHash, + instanceId: appInfo.instanceId, + deviceId: appInfo.deviceId, + mrAggregated: appInfo.mrAggregated, + mrSystem: appInfo.mrSystem, + osImageHash: appInfo.osImageHash, + tcbStatus: appInfo.tcbStatus, + advisoryIds: appInfo.advisoryIds + }); + if (appInfo.appId == address(this)) { + (result.isAllowed, result.rejectReason) = _isKmsAllowed(bootInfo); + } else { + (result.isAllowed, result.rejectReason) = _isAppAllowed(bootInfo); + } + } + // Function to check if KMS is allowed to boot function isKmsAllowed(AppBootInfo calldata bootInfo) external view returns (bool isAllowed, string memory reason) { + return _isKmsAllowed(bootInfo); + } + + function _isKmsAllowed(AppBootInfo memory bootInfo) private view returns (bool, string memory) { // Check if the TCB status is up to date if (keccak256(abi.encodePacked(bootInfo.tcbStatus)) != keccak256(abi.encodePacked("UpToDate"))) { return (false, "TCB status is not up to date"); @@ -402,6 +454,10 @@ contract DstackKms is Initializable, Ownable2StepUpgradeable, UUPSUpgradeable, E override returns (bool isAllowed, string memory reason) { + return _isAppAllowed(bootInfo); + } + + function _isAppAllowed(AppBootInfo memory bootInfo) private view returns (bool, string memory) { // Check if app is registered if (!registeredApps[bootInfo.appId]) { return (false, "App not registered"); diff --git a/dstack/kms/auth-eth/src/ethereum.ts b/dstack/kms/auth-eth/src/ethereum.ts index 6eab5d874..a89d2fb9b 100644 --- a/dstack/kms/auth-eth/src/ethereum.ts +++ b/dstack/kms/auth-eth/src/ethereum.ts @@ -3,10 +3,11 @@ // SPDX-License-Identifier: Apache-2.0 import { ethers } from 'ethers'; -import { BootInfo, BootResponse, RootEpochInfo } from './types'; +import { BootInfo, BootResponse, RootEpochInfo, AuthorizationResult } from './types'; // Minimal ABI for DstackKms contract const DSTACK_KMS_ABI = [ + "function checkAuthorization((address appId,bytes32 composeHash,address instanceId,bytes32 deviceId,bytes32 mrAggregated,bytes32 mrSystem,bytes32 osImageHash,string tcbStatus,string[] advisoryIds) appInfo) view returns ((bool isAllowed,string rejectReason,string gatewayAppId,uint256 currentRootEpoch,bytes32 currentRootEpochHash))", "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)", "function isKmsAllowed((address appId,bytes32 composeHash,address instanceId,bytes32 deviceId,bytes32 mrAggregated,bytes32 mrSystem,bytes32 osImageHash,string tcbStatus,string[] advisoryIds) bootInfo) view returns (bool, string)", "function gatewayAppId() view returns (string)", @@ -61,9 +62,9 @@ export class EthereumBackend { return '0x' + hex; } - async checkBoot(bootInfo: BootInfo, isKms: boolean): Promise { + private contractAppInfo(bootInfo: BootInfo) { // Create boot info struct for contract call - const bootInfoStruct = { + return { appId: this.decodeHex(bootInfo.appId, 20), instanceId: this.decodeHex(bootInfo.instanceId, 20), composeHash: this.decodeHex(bootInfo.composeHash, 32), @@ -74,6 +75,21 @@ export class EthereumBackend { tcbStatus: bootInfo.tcbStatus, advisoryIds: bootInfo.advisoryIds }; + } + + async checkAuthorization(appInfo: BootInfo): Promise { + const result = await this.kmsContract.checkAuthorization(this.contractAppInfo(appInfo)); + return { + isAllowed: result.isAllowed, + rejectReason: result.rejectReason, + gatewayAppId: result.gatewayAppId, + currentRootEpoch: ethers.toNumber(result.currentRootEpoch), + currentRootEpochHash: result.currentRootEpochHash, + }; + } + + async checkBoot(bootInfo: BootInfo, isKms: boolean): Promise { + const bootInfoStruct = this.contractAppInfo(bootInfo); // Read the whole decision from one block. const blockTag = await this.provider.getBlockNumber(); let response; diff --git a/dstack/kms/auth-eth/src/main.test.ts b/dstack/kms/auth-eth/src/main.test.ts index e7ad49739..bcb44bd3c 100644 --- a/dstack/kms/auth-eth/src/main.test.ts +++ b/dstack/kms/auth-eth/src/main.test.ts @@ -91,6 +91,28 @@ describe('Server', () => { }); }); + describe('POST /checkAuthorization', () => { + it('preserves the complete decision, including the epoch on rejection', async () => { + const decision = { + isAllowed: false, rejectReason: 'App not registered', gatewayAppId: 'gateway', + currentRootEpoch: 1, currentRootEpochHash: '0x' + '11'.repeat(32), + }; + app.ethereum.checkAuthorization = jest.fn().mockResolvedValue(decision); + const response = await app.inject({ method: 'POST', url: '/checkAuthorization', payload: mockBootInfo }); + expect(response.statusCode).toBe(200); + expect(JSON.parse(response.payload)).toEqual(decision); + expect(app.ethereum.checkAuthorization).toHaveBeenCalledWith(mockBootInfo); + }); + + it('fails closed when the new contract method is unavailable', async () => { + jest.clearAllMocks(); + app.ethereum.checkAuthorization = jest.fn().mockRejectedValue(new Error('missing contract method')); + const response = await app.inject({ method: 'POST', url: '/checkAuthorization', payload: mockBootInfo }); + expect(response.statusCode).toBe(500); + expect(app.ethereum.checkBoot).not.toHaveBeenCalled(); + }); + }); + describe('key epochs', () => { it('should pass the KMS root epoch through to the backend', async () => { const mockCheckBoot = jest.fn().mockResolvedValue({ isAllowed: true, reason: '', gatewayAppId: '' }); diff --git a/dstack/kms/auth-eth/src/server.ts b/dstack/kms/auth-eth/src/server.ts index 33c80d1d5..225607f80 100644 --- a/dstack/kms/auth-eth/src/server.ts +++ b/dstack/kms/auth-eth/src/server.ts @@ -4,7 +4,7 @@ import fastify, { FastifyInstance } from 'fastify'; import { EthereumBackend } from './ethereum'; -import { BootInfo, BootResponse } from './types'; +import { BootInfo, BootResponse, AuthorizationResult } from './types'; import { ethers } from 'ethers'; declare module 'fastify' { @@ -104,6 +104,25 @@ export async function build(): Promise { } }); + server.post<{ Body: BootInfo; Reply: AuthorizationResult }>('/checkAuthorization', { + schema: { + body: { $ref: 'bootInfo#' }, + response: { + 200: { + type: 'object', + required: ['isAllowed', 'rejectReason', 'gatewayAppId', 'currentRootEpoch', 'currentRootEpochHash'], + properties: { + isAllowed: { type: 'boolean' }, + rejectReason: { type: 'string' }, + gatewayAppId: { type: 'string' }, + currentRootEpoch: { type: 'integer', minimum: 0 }, + currentRootEpochHash: hex(32, 'Current root epoch chain hash'), + }, + }, + }, + }, + }, async (request) => server.ethereum.checkAuthorization(request.body)); + // Define routes server.post<{ Body: BootInfo; diff --git a/dstack/kms/auth-eth/src/types.ts b/dstack/kms/auth-eth/src/types.ts index c92edcfe0..7fbd6aed5 100644 --- a/dstack/kms/auth-eth/src/types.ts +++ b/dstack/kms/auth-eth/src/types.ts @@ -32,3 +32,11 @@ export interface RootEpochInfo { } // Removed KMS_CONTRACT_ABI and APP_CONTRACT_ABI since we're using typechain types now + +export interface AuthorizationResult { + isAllowed: boolean; + rejectReason: string; + gatewayAppId: string; + currentRootEpoch: number; + currentRootEpochHash: string; +} diff --git a/dstack/kms/auth-eth/test/DstackKms.t.sol b/dstack/kms/auth-eth/test/DstackKms.t.sol index 037374aec..bfaa3cf87 100644 --- a/dstack/kms/auth-eth/test/DstackKms.t.sol +++ b/dstack/kms/auth-eth/test/DstackKms.t.sol @@ -442,6 +442,74 @@ contract DstackKmsTest is Test { vm.stopPrank(); } + function authorizationInfo(address appId) internal pure returns (DstackKms.AppInfo memory) { + return DstackKms.AppInfo({ + appId: appId, + composeHash: bytes32(uint256(789)), + instanceId: address(0), + deviceId: bytes32(uint256(123)), + mrAggregated: bytes32(uint256(456)), + mrSystem: bytes32(0), + osImageHash: bytes32(uint256(111)), + tcbStatus: "UpToDate", + advisoryIds: new string[](0) + }); + } + + function test_CheckAuthorization_UnregisteredEpochReturnsRejection() public view { + DstackKms.AuthorizationResult memory result = kms.checkAuthorization(authorizationInfo(address(kms))); + assertFalse(result.isAllowed); + assertEq(result.rejectReason, "Root epoch not registered"); + assertEq(result.currentRootEpoch, 0); + assertEq(result.currentRootEpochHash, bytes32(0)); + } + + function test_CheckAuthorization_SelectsPolicyByAppId() public { + DstackKms.AppInfo memory info = authorizationInfo(address(kms)); + vm.startPrank(owner); + kms.registerRootEpoch(ca(0xa0), k256(0xb0), "https://example.com/epoch-0"); + kms.setGatewayAppId("gateway"); + kms.addOsImageHash(info.osImageHash); + kms.addKmsDevice(info.deviceId); + kms.addKmsAggregatedMr(info.mrAggregated); + address appId = kms.deployAndRegisterApp(user, false, false, info.deviceId, info.composeHash); + vm.stopPrank(); + + DstackKms.AuthorizationResult memory result = kms.checkAuthorization(info); + assertTrue(result.isAllowed); + assertEq(result.rejectReason, ""); + assertEq(result.gatewayAppId, "gateway"); + assertEq(result.currentRootEpochHash, chainHash(bytes32(0), ca(0xa0), k256(0xb0))); + + info.appId = user; + result = kms.checkAuthorization(info); + assertFalse(result.isAllowed); + assertEq(result.rejectReason, "App not registered"); + + info.appId = appId; + info.mrAggregated = bytes32(0); + result = kms.checkAuthorization(info); + assertTrue(result.isAllowed); + info.appId = address(kms); + result = kms.checkAuthorization(info); + assertFalse(result.isAllowed); + assertEq(result.rejectReason, "Aggregated MR not allowed"); + } + + function test_CheckAuthorization_RejectionIncludesActivatedEpoch() public { + vm.startPrank(owner); + kms.registerRootEpoch(ca(0xa0), k256(0xb0), "https://example.com/epoch-0"); + kms.registerRootEpoch(ca(0xa1), k256(0xb1), "https://example.com/epoch-1"); + kms.activateRootEpoch(1); + vm.stopPrank(); + + DstackKms.AuthorizationResult memory result = kms.checkAuthorization(authorizationInfo(user)); + assertFalse(result.isAllowed); + assertEq(result.rejectReason, "App not registered"); + assertEq(result.currentRootEpoch, 1); + assertEq(result.currentRootEpochHash, chainHash(chainHash(bytes32(0), ca(0xa0), k256(0xb0)), ca(0xa1), k256(0xb1))); + } + function test_IsKmsAllowed() public { bytes32 deviceId = bytes32(uint256(123)); bytes32 mrAggregated = bytes32(uint256(456)); diff --git a/dstack/kms/auth-mock/README.md b/dstack/kms/auth-mock/README.md index c44c207e8..961acf432 100644 --- a/dstack/kms/auth-mock/README.md +++ b/dstack/kms/auth-mock/README.md @@ -167,3 +167,13 @@ to switch from mock to real backend: 4. update environment variables the API remains identical, so no client code changes are needed. + + +### Unified authorization + +`POST /checkAuthorization` returns `AuthorizationResult`, including +`rejectReason`, `gatewayAppId`, `currentRootEpoch`, and `currentRootEpochHash`. +It selects KMS or app policy from the verified `appId`. See the +[KMS authorization protocol](../README.md#authorization-protocol) for registration, +bootstrap, and compatibility requirements. The legacy `/bootAuth/*` endpoints +remain available. diff --git a/dstack/kms/auth-mock/index.test.ts b/dstack/kms/auth-mock/index.test.ts index 68a3cfbf4..7169a7e50 100644 --- a/dstack/kms/auth-mock/index.test.ts +++ b/dstack/kms/auth-mock/index.test.ts @@ -382,3 +382,32 @@ describe('Mock Behavior Verification', () => { } }); }); + + +describe('POST /checkAuthorization', () => { + const authorize = (appId: string) => appFetch(new Request('http://localhost/checkAuthorization', { + method: 'POST', headers: { 'Content-Type': 'application/json' }, + body: JSON.stringify({ appId, mrAggregated: '', osImageHash: '', composeHash: '', instanceId: '', deviceId: '' }), + })); + + it('requires registered roots and selects policy by app ID', async () => { + vi.stubEnv('MOCK_ROOT_EPOCH', ''); + vi.stubEnv('MOCK_ROOT_EPOCH_HASH', ''); + vi.stubEnv('MOCK_POLICY', 'deny-kms'); + try { + expect(await (await authorize('app')).json()).toMatchObject({ + isAllowed: false, rejectReason: 'Root epoch not registered', + currentRootEpoch: 0, currentRootEpochHash: '0x' + '00'.repeat(32), + }); + vi.stubEnv('MOCK_ROOT_EPOCH', '1'); + vi.stubEnv('MOCK_ROOT_EPOCH_HASH', '0x' + '11'.repeat(32)); + expect(await (await authorize('app')).json()).toMatchObject({ isAllowed: true, rejectReason: '', currentRootEpoch: 1 }); + expect(await (await authorize(process.env.KMS_CONTRACT_ADDR!)).json()).toMatchObject({ + isAllowed: false, rejectReason: 'mock policy: deny-kms', currentRootEpoch: 1, + currentRootEpochHash: '0x' + '11'.repeat(32), + }); + } finally { + vi.unstubAllEnvs(); + } + }); +}); diff --git a/dstack/kms/auth-mock/index.ts b/dstack/kms/auth-mock/index.ts index 6d01a947f..b8d0b586d 100644 --- a/dstack/kms/auth-mock/index.ts +++ b/dstack/kms/auth-mock/index.ts @@ -106,6 +106,21 @@ class MockBackend { this.mockAppImplementation = process.env.MOCK_APP_IMPLEMENTATION || '0xmockapp9876543210987654321098765432109'; } + async checkAuthorization(appInfo: BootInfo) { + const current = getRootEpoch(); + const epoch = { + currentRootEpoch: current.rootEpoch ?? 0, + currentRootEpochHash: current.rootEpochHash ?? '0x' + '00'.repeat(32), + }; + if (current.rootEpoch === undefined) { + return { isAllowed: false, rejectReason: 'Root epoch not registered', gatewayAppId: this.mockGatewayAppId, ...epoch }; + } + const kmsAppId = process.env.KMS_CONTRACT_ADDR || '0xmockcontract1234567890123456789012345678'; + const result = this.evaluateBoot(appInfo, normHash(appInfo.appId) === normHash(kmsAppId)); + return { isAllowed: result.isAllowed, rejectReason: result.isAllowed ? '' : result.reason, + gatewayAppId: this.mockGatewayAppId, ...epoch }; + } + async checkBoot(bootInfo: BootInfo, isKms: boolean): Promise { const response = this.evaluateBoot(bootInfo, isKms); const epochDenial = response.isAllowed ? rootEpochDenial(bootInfo, isKms) : undefined; @@ -202,6 +217,14 @@ app.get('/', async (c) => { } }); +app.post('/checkAuthorization', zValidator('json', BootInfoSchema), async (c) => { + try { + return c.json(await mockBackend.checkAuthorization(c.req.valid('json'))); + } catch { + return c.json({ status: 'error', message: 'authorization backend unavailable' }, 503); + } +}); + // app boot authentication app.post('/bootAuth/app', zValidator('json', BootInfoSchema), diff --git a/dstack/kms/auth-mock/openapi.json b/dstack/kms/auth-mock/openapi.json index 6a35cbbc3..6679d20a8 100644 --- a/dstack/kms/auth-mock/openapi.json +++ b/dstack/kms/auth-mock/openapi.json @@ -98,6 +98,37 @@ } } } + }, + "/checkAuthorization": { + "post": { + "summary": "Authorize an attested identity and return the current root epoch", + "operationId": "checkAuthorization", + "requestBody": { + "required": true, + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/AppInfo" + } + } + } + }, + "responses": { + "200": { + "description": "Authorization decision, including the epoch on rejection.", + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/AuthorizationResult" + } + } + } + }, + "503": { + "description": "Authorization backend unavailable; no legacy fallback." + } + } + } } }, "components": { @@ -201,7 +232,9 @@ "properties": { "status": { "type": "string", - "enum": ["ok"], + "enum": [ + "ok" + ], "description": "System status" }, "kmsContractAddr": { @@ -229,8 +262,14 @@ "description": "Mock app implementation address", "example": "0xmockapp9876543210987654321098765432109" }, - "rootEpoch": { "type": "integer", "description": "MOCK_ROOT_EPOCH" }, - "rootEpochHash": { "type": "string", "description": "MOCK_ROOT_EPOCH_HASH" }, + "rootEpoch": { + "type": "integer", + "description": "MOCK_ROOT_EPOCH" + }, + "rootEpochHash": { + "type": "string", + "description": "MOCK_ROOT_EPOCH_HASH" + }, "note": { "type": "string", "description": "Note about mock backend behavior", @@ -247,13 +286,54 @@ "properties": { "status": { "type": "string", - "enum": ["error"] + "enum": [ + "error" + ] }, "message": { "type": "string", "description": "Error message" } } + }, + "AppInfo": { + "allOf": [ + { + "$ref": "#/components/schemas/BootInfo" + } + ], + "description": "Verified attestation identity. The KMS app ID is the KMS contract address." + }, + "AuthorizationResult": { + "type": "object", + "required": [ + "isAllowed", + "rejectReason", + "gatewayAppId", + "currentRootEpoch", + "currentRootEpochHash" + ], + "properties": { + "isAllowed": { + "type": "boolean" + }, + "rejectReason": { + "type": "string", + "description": "Empty when allowed." + }, + "gatewayAppId": { + "type": "string" + }, + "currentRootEpoch": { + "type": "integer", + "minimum": 0 + }, + "currentRootEpochHash": { + "type": "string", + "pattern": "^0x[0-9a-fA-F]{64}$", + "description": "Zero when no epoch is registered; authorization is denied." + } + } } } } diff --git a/dstack/kms/auth-simple/README.md b/dstack/kms/auth-simple/README.md index dc48e7c50..14f930cf9 100644 --- a/dstack/kms/auth-simple/README.md +++ b/dstack/kms/auth-simple/README.md @@ -252,3 +252,13 @@ bun run test:run - [auth-eth](../auth-eth/) - On-chain governance auth server - [auth-mock](../auth-mock/) - Development/testing auth server (always allows) - [Deployment Guide](../../../docs/deployment.md) - Full deployment instructions + + +### Unified authorization + +`POST /checkAuthorization` returns `AuthorizationResult`, including +`rejectReason`, `gatewayAppId`, `currentRootEpoch`, and `currentRootEpochHash`. +It selects KMS or app policy from the verified `appId`. See the +[KMS authorization protocol](../README.md#authorization-protocol) for registration, +bootstrap, and compatibility requirements. The legacy `/bootAuth/*` endpoints +remain available. diff --git a/dstack/kms/auth-simple/index.test.ts b/dstack/kms/auth-simple/index.test.ts index 5e7e37238..f45087f10 100644 --- a/dstack/kms/auth-simple/index.test.ts +++ b/dstack/kms/auth-simple/index.test.ts @@ -472,6 +472,34 @@ describe('auth-simple', () => { const info = async () => (await app.fetch(new Request('http://localhost/', { method: 'GET' }))).json(); + const authorize = async (fields: object = {}) => + (await app.fetch(new Request('http://localhost/checkAuthorization', { + method: 'POST', headers: { 'Content-Type': 'application/json' }, + body: JSON.stringify({ ...baseBootInfo, ...fields }), + }))).json(); + + it('selects KMS policy by app ID and returns the epoch on rejection', async () => { + const kmsContractAddr = '0x' + '11'.repeat(20); + writeTestConfig({ ...config, kmsContractAddr, currentRootEpoch: 1 }); + expect(await authorize()).toMatchObject({ + isAllowed: true, rejectReason: '', currentRootEpoch: 1, currentRootEpochHash: h1, + }); + expect(await authorize({ appId: kmsContractAddr, composeHash: 'not-an-app-compose' })).toMatchObject({ isAllowed: true }); + expect(await authorize({ appId: kmsContractAddr, mrAggregated: '0xbad' })).toMatchObject({ + isAllowed: false, currentRootEpoch: 1, currentRootEpochHash: h1, + }); + expect(await authorize({ appId: '0x' + '22'.repeat(20) })).toMatchObject({ isAllowed: false }); + }); + + it('denies authorization without registered roots, while bootstrap remains available', async () => { + writeTestConfig({ ...config, rootEpochs: [] }); + expect(await authorize()).toMatchObject({ + isAllowed: false, rejectReason: 'Root epoch not registered', + currentRootEpoch: 0, currentRootEpochHash: '0x' + '00'.repeat(32), + }); + expect((await boot('kms', {})).isAllowed).toBe(true); + }); + it('allows only a KMS serving the current root epoch', async () => { writeTestConfig(config); expect(await info()).toMatchObject({ rootEpoch: 0, rootEpochHash: h0 }); diff --git a/dstack/kms/auth-simple/index.ts b/dstack/kms/auth-simple/index.ts index 4a242ac7e..a1da5814e 100644 --- a/dstack/kms/auth-simple/index.ts +++ b/dstack/kms/auth-simple/index.ts @@ -152,6 +152,21 @@ class ConfigBackend { } } + async checkAuthorization(appInfo: BootInfo) { + const config = this.loadConfig(); + const current = currentRootEpoch(config); + const epoch = { + currentRootEpoch: current.rootEpoch ?? 0, + currentRootEpochHash: current.rootEpochHash ?? '0x' + '00'.repeat(32), + }; + if (current.rootEpoch === undefined) { + return { isAllowed: false, rejectReason: 'Root epoch not registered', gatewayAppId: config.gatewayAppId, ...epoch }; + } + const isKms = normalizeHex(appInfo.appId) === normalizeHex(config.kmsContractAddr); + const result = this.evaluateBoot(appInfo, isKms, config); + return { isAllowed: result.isAllowed, rejectReason: result.reason, gatewayAppId: result.gatewayAppId, ...epoch }; + } + async checkBoot(bootInfo: BootInfo, isKms: boolean): Promise { const config = this.loadConfig(); const response = this.evaluateBoot(bootInfo, isKms, config); @@ -332,6 +347,15 @@ app.get('/', async (c) => { } }); +app.post('/checkAuthorization', zValidator('json', BootInfoSchema), async (c) => { + try { + return c.json(await backend.checkAuthorization(c.req.valid('json'))); + } catch { + console.error('authorization backend failed'); + return c.json({ status: 'error', message: 'authorization backend unavailable' }, 503); + } +}); + // app boot authentication app.post('/bootAuth/app', zValidator('json', BootInfoSchema), diff --git a/dstack/kms/src/main_service.rs b/dstack/kms/src/main_service.rs index 083ba991c..72fe7428b 100644 --- a/dstack/kms/src/main_service.rs +++ b/dstack/kms/src/main_service.rs @@ -31,8 +31,8 @@ use ra_tls::{ use scale::Decode; use tracing::warn; use upgrade_authority::{ - build_boot_info, ensure_app_id_len, local_kms_boot_info, BootInfo, CurrentRootEpoch, - GetInfoResponse, RootEpochRegistry, + build_boot_info, ensure_app_id_len, local_kms_boot_info, AuthorizationResult, BootInfo, + CurrentRootEpoch, GetInfoResponse, RootEpochRegistry, }; use crate::{ @@ -302,8 +302,7 @@ impl KmsState { self.restart.notified().await } - /// The served epoch and its chain hash, sent with every boot check so - /// the auth backend can refuse a KMS that does not serve the current one. + /// The served epoch and chain hash to compare with each authorization result. fn serving_root_epoch(&self) -> CurrentRootEpoch { CurrentRootEpoch { epoch: self.serving_epoch, @@ -311,35 +310,57 @@ impl KmsState { } } - /// A denial may mean governance moved the current root epoch. If it did, - /// request a restart, which re-checks the keyring and re-signs the RPC - /// certificate. - pub(crate) async fn restart_if_epoch_moved(&self) { - let info = match self + /// KMS-only operations must not accept an app that passes ordinary app policy. + async fn ensure_kms_identity(&self, app_info: &BootInfo) -> Result<()> { + if self.config.auth_api.is_dev() { + return Ok(()); + } + let info = self .auth_api_info .try_get_with((), self.config.auth_api.get_info()) .await - { - Ok(info) => info, - Err(err) => { - warn!("failed to read the current root epoch: {err:#}"); - return; + .map_err(|err| anyhow!("{err:#}"))?; + let address = info + .kms_contract_address + .as_deref() + .context("missing KMS contract address")?; + let app_id = hex::decode(address.trim_start_matches("0x")) + .context("invalid KMS contract address")?; + if app_info.app_id != app_id { + bail!("KMS app ID does not match the KMS contract address"); + } + Ok(()) + } + + /// Compare against the epoch returned with this decision, never a cached info response. + fn ensure_authorized(&self, result: &AuthorizationResult) -> Result<()> { + if !self.config.auth_api.is_dev() { + if let Some(current) = result.current_root()? { + if current != self.serving_root_epoch() { + self.restart.notify_one(); + bail!("the authorized root epoch or chain hash changed; restarting"); + } + } else { + bail!( + "Authorization denied: {}", + if result.reject_reason.is_empty() { + "Root epoch not registered" + } else { + &result.reject_reason + } + ); } - }; - let serving = self.serving_epoch; - if self.keyring.servable_epoch(info.root_epochs).ok() != Some(serving) { - warn!( - "governance moved off root epoch {serving} ({:?}); restarting", - info.root_epochs - ); - self.restart.notify_one(); } + if !result.is_allowed { + bail!("Authorization denied: {}", result.reject_reason); + } + Ok(()) } } /// Wait until governance lets `keyring` serve, and return the epoch to serve. -/// A keyring that predates root epochs keeps serving epoch 0 when the auth API -/// is unreachable, as before. +/// An epoch-0 keyring may start while the info endpoint is unreachable, but +/// key release still requires an authorization result with registered roots. pub(crate) async fn wait_for_servable_epoch(config: &KmsConfig, keyring: &Keyring) -> u32 { const RETRY: Duration = Duration::from_secs(10); loop { @@ -508,6 +529,7 @@ impl RpcHandler { .await .map_err(|err| anyhow!("{err:#}")) .context("Failed to load cached self boot info")?; + self.state.ensure_kms_identity(&boot_info).await?; self.state .self_allowed .try_get_with((), async { @@ -515,14 +537,10 @@ impl RpcHandler { .state .config .auth_api - .is_app_allowed(&boot_info, true, Some(self.state.serving_root_epoch())) + .check_authorization(&boot_info) .await .context("Failed to call self KMS auth check")?; - if !response.is_allowed { - self.state.restart_if_epoch_moved().await; - bail!("KMS is not allowed: {}", response.reason); - } - Ok(()) + self.state.ensure_authorized(&response) }) .await .map_err(|err| anyhow!("{err:#}"))?; @@ -557,16 +575,16 @@ impl RpcHandler { vm_config_str: &str, ) -> Result { let boot_info = build_boot_info_for_attestation(att, use_boottime_mr, vm_config_str)?; + if is_kms { + self.state.ensure_kms_identity(&boot_info).await?; + } let response = self .state .config .auth_api - .is_app_allowed(&boot_info, is_kms, Some(self.state.serving_root_epoch())) + .check_authorization(&boot_info) .await?; - if !response.is_allowed { - self.state.restart_if_epoch_moved().await; - bail!("Boot denied: {}", response.reason); - } + self.state.ensure_authorized(&response)?; // SNP rootfs/app/config binding is handled by the SNP launch-measurement // helper above. The legacy OS-image verifier is TDX-oriented and still // rejects SNP quotes; keep SNP on the explicit fail-closed helper path. @@ -1533,6 +1551,88 @@ mod tests { KmsState::new(config, keyring, 0).unwrap() } + #[rocket::async_test] + async fn authorization_epoch_is_checked_without_reading_cached_info() { + for (changed, allowed, registered, restart) in [ + (false, true, true, false), + (false, false, true, false), + (true, true, true, true), + (true, false, true, true), + (false, false, false, false), + (false, true, false, false), + ] { + let dir = tempfile::tempdir().unwrap(); + let state = kms_state(dir.path(), "http://127.0.0.1:1"); + let mut current = state.serving_root_epoch(); + if changed { + current.epoch = u32::MAX; + } + let result = AuthorizationResult { + is_allowed: allowed, + reject_reason: if allowed { "" } else { "policy denied" }.into(), + gateway_app_id: "gateway".into(), + current_root_epoch: current.epoch, + current_root_epoch_hash: hex::encode(if registered { + current.chain_hash + } else { + [0; 32] + }), + }; + assert_eq!( + state.ensure_authorized(&result).is_ok(), + allowed && !changed && registered + ); + assert_eq!( + tokio::time::timeout(Duration::from_millis(10), state.restart_requested()) + .await + .is_ok(), + restart, + ); + } + } + + #[rocket::async_test] + async fn authorization_hash_mismatch_restarts_and_malformed_hash_is_rejected() { + let dir = tempfile::tempdir().unwrap(); + let state = kms_state(dir.path(), "http://127.0.0.1:1"); + let mut result = AuthorizationResult { + is_allowed: true, + reject_reason: String::new(), + gateway_app_id: String::new(), + current_root_epoch: 0, + current_root_epoch_hash: "invalid".into(), + }; + assert!(state.ensure_authorized(&result).is_err()); + assert!( + tokio::time::timeout(Duration::from_millis(10), state.restart_requested()) + .await + .is_err() + ); + result.current_root_epoch_hash = hex::encode([0x55; 32]); + assert!(state.ensure_authorized(&result).is_err()); + assert!( + tokio::time::timeout(Duration::from_millis(10), state.restart_requested()) + .await + .is_ok() + ); + } + + #[rocket::async_test] + async fn kms_identity_cannot_be_replaced_with_an_authorized_app() { + use super::upgrade_authority::tests::serve; + let app_id = vec![0x11; 20]; + let response = AUTH_API_INFO.replace("0xkms", &format!("0x{}", hex::encode(&app_id))); + let (url, server) = serve(vec![response]); + let dir = tempfile::tempdir().unwrap(); + let state = kms_state(dir.path(), &url); + let mut info = super::upgrade_authority::tests::boot_info(1); + info.app_id = app_id; + state.ensure_kms_identity(&info).await.unwrap(); + info.app_id[0] ^= 1; + assert!(state.ensure_kms_identity(&info).await.is_err()); + assert_eq!(server.join().unwrap().len(), 1); + } + const AUTH_API_INFO: &str = r#"{"status":"ok","kmsContractAddr":"0xkms","ethRpcUrl":"https://rpc.example","gatewayAppId":"0xgateway","chainId":1,"appImplementation":"0ximpl"}"#; #[rocket::async_test] diff --git a/dstack/kms/src/main_service/upgrade_authority.rs b/dstack/kms/src/main_service/upgrade_authority.rs index be4c294d5..40de002ec 100644 --- a/dstack/kms/src/main_service/upgrade_authority.rs +++ b/dstack/kms/src/main_service/upgrade_authority.rs @@ -86,6 +86,33 @@ pub(crate) struct BootResponse { pub reason: String, } +#[derive(Debug, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub(crate) struct AuthorizationResult { + pub is_allowed: bool, + pub reject_reason: String, + pub gateway_app_id: String, + pub current_root_epoch: u32, + pub current_root_epoch_hash: String, +} + +impl AuthorizationResult { + pub fn current_root(&self) -> Result> { + let chain_hash: [u8; 32] = + hex::decode(self.current_root_epoch_hash.trim_start_matches("0x")) + .ok() + .and_then(|hash| hash.try_into().ok()) + .context("invalid authorization root epoch chain hash")?; + if chain_hash == [0; 32] { + return Ok(None); + } + Ok(Some(CurrentRootEpoch { + epoch: self.current_root_epoch, + chain_hash, + })) + } +} + /// The root epoch fields: the KMS sends the epoch it serves in boot checks, /// and auth backends report the current one in info responses. #[derive(Debug, Clone, Default, Serialize, Deserialize)] @@ -227,8 +254,25 @@ async fn send_request(req: reqwest::RequestBuilder, url: &s } impl AuthApi { - /// `root_epoch` is the epoch this KMS serves, which the backend checks is - /// current; unset while onboarding, before the KMS serves any. + pub async fn check_authorization(&self, app_info: &BootInfo) -> Result { + ensure_identity_widths(app_info)?; + match self { + AuthApi::Dev { dev } => Ok(AuthorizationResult { + is_allowed: true, + reject_reason: String::new(), + gateway_app_id: dev.gateway_app_id.clone(), + current_root_epoch: 0, + current_root_epoch_hash: hex::encode([0; 32]), + }), + AuthApi::Webhook { webhook } => { + let url = url_join(&webhook.url, "checkAuthorization"); + http_post(&url, app_info, webhook.timeout).await + } + } + } + + /// Legacy scoped checks remain available during bootstrap and onboarding, + /// before the first root epoch can be registered. pub async fn is_app_allowed( &self, boot_info: &BootInfo, @@ -391,7 +435,7 @@ pub(crate) mod tests { use std::net::TcpListener; use std::thread; - fn boot_info(identity: u8) -> BootInfo { + pub(crate) fn boot_info(identity: u8) -> BootInfo { BootInfo { tee_variant: TeeVariant::DstackTdx, mr_aggregated: vec![identity; 32], @@ -609,6 +653,43 @@ pub(crate) mod tests { assert_ne!(requests[0].1["mrAggregated"], requests[1].1["mrAggregated"]); } + #[rocket::async_test] + async fn authorization_returns_the_epoch_with_each_decision() { + let response = |allowed: bool, epoch: u32| { + format!( + r#"{{"isAllowed":{allowed},"rejectReason":"","gatewayAppId":"gw", "currentRootEpoch":{epoch},"currentRootEpochHash":"0x{}"}}"#, + "11".repeat(32) + ) + }; + let (url, server) = serve(vec![response(true, 0), response(false, 1)]); + let auth = webhook(url); + let info = boot_info(50); + let first = auth.check_authorization(&info).await.unwrap(); + let second = auth.check_authorization(&info).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); + let requests = server.join().unwrap(); + assert_eq!(requests.len(), 2); + assert!(requests + .iter() + .all(|(path, _)| path == "/checkAuthorization")); + assert_eq!(requests[0].1, serde_json::to_value(info).unwrap()); + } + + #[rocket::async_test] + async fn authorization_without_epoch_fields_is_not_accepted_as_legacy() { + let (url, server) = serve(vec![ + r#"{"isAllowed":true,"reason":"","gatewayAppId":"gw"}"#, + ]); + assert!(webhook(url) + .check_authorization(&boot_info(50)) + .await + .is_err()); + assert_eq!(server.join().unwrap().len(), 1); + } + #[rocket::async_test] async fn root_epochs_are_read_from_info_and_sent_in_boot_checks() { let hash = "0x".to_string() + &"11".repeat(32); diff --git a/dstack/kms/src/onboard_service.rs b/dstack/kms/src/onboard_service.rs index 43a7dde8b..6d4a857c8 100644 --- a/dstack/kms/src/onboard_service.rs +++ b/dstack/kms/src/onboard_service.rs @@ -736,44 +736,6 @@ mod tests { assert_eq!(keyring.servable_epoch(RootEpochRegistry::Empty).unwrap(), 0); } - #[rocket::async_test] - async fn a_denial_restarts_the_kms_if_the_current_epoch_moved() { - use crate::main_service::upgrade_authority::tests::{serve, webhook}; - use crate::main_service::{Keyring, KmsState}; - - let (_dir, cfg, res) = stored_keyring(3).await; - let restarts = |registry: RootEpochRegistry| { - let mut cfg = cfg.clone(); - async move { - let epoch = match registry { - RootEpochRegistry::Current(current) => format!( - r#","rootEpoch":{},"rootEpochHash":"{}""#, - current.epoch, - hex::encode(current.chain_hash) - ), - _ => String::new(), - }; - let (url, server) = serve(vec![format!( - r#"{{"status":"ok","kmsContractAddr":"0x00","gatewayAppId":"gw","chainId":1, - "appImplementation":"0x00"{epoch}}}"# - )]); - cfg.auth_api = webhook(url); - let keyring = Keyring::load(&cfg).unwrap(); - let state = KmsState::new(cfg, keyring, 1).unwrap(); - state.restart_if_epoch_moved().await; - server.join().unwrap(); - let restart = - tokio::time::timeout(Duration::from_millis(100), state.restart_requested()); - restart.await.is_ok() - } - }; - - assert!(!restarts(current_of(&res, 1)).await); - assert!(restarts(current_of(&res, 2)).await); - assert!(restarts(RootEpochRegistry::Empty).await); - assert!(restarts(current_of(&handover(3), 1)).await); - } - #[rocket::async_test] async fn setup_status_includes_bootstrap_info_when_present() { let cert_dir = tempfile::tempdir().unwrap(); From 1a778c987302e6b1f00afe9010edb35a980373a2 Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Fri, 6 Nov 2026 08:52:33 -0800 Subject: [PATCH 012/166] style(kms): format authorization regression test Signed-off-by: Kevin Wang --- dstack/kms/auth-eth/test/DstackKms.t.sol | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/dstack/kms/auth-eth/test/DstackKms.t.sol b/dstack/kms/auth-eth/test/DstackKms.t.sol index bfaa3cf87..777e7fdbf 100644 --- a/dstack/kms/auth-eth/test/DstackKms.t.sol +++ b/dstack/kms/auth-eth/test/DstackKms.t.sol @@ -507,7 +507,9 @@ contract DstackKmsTest is Test { assertFalse(result.isAllowed); assertEq(result.rejectReason, "App not registered"); assertEq(result.currentRootEpoch, 1); - assertEq(result.currentRootEpochHash, chainHash(chainHash(bytes32(0), ca(0xa0), k256(0xb0)), ca(0xa1), k256(0xb1))); + assertEq( + result.currentRootEpochHash, chainHash(chainHash(bytes32(0), ca(0xa0), k256(0xb0)), ca(0xa1), k256(0xb1)) + ); } function test_IsKmsAllowed() public { From 61ccda299b070a97397485b44a29290e54ce5611 Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Fri, 6 Nov 2026 09:43:29 -0800 Subject: [PATCH 013/166] feat(kms): select requested epochs from authorization results Signed-off-by: Kevin Wang --- docs/deployment.md | 18 +- dstack/kms/README.md | 26 +- dstack/kms/auth-eth-bun/README.md | 4 +- dstack/kms/auth-eth-bun/index.test.ts | 17 +- dstack/kms/auth-eth-bun/index.ts | 21 +- dstack/kms/auth-eth-bun/openapi.json | 71 +++- dstack/kms/auth-eth/README.md | 4 +- dstack/kms/auth-eth/contracts/DstackKms.sol | 21 +- dstack/kms/auth-eth/src/ethereum.ts | 11 +- dstack/kms/auth-eth/src/main.test.ts | 2 + dstack/kms/auth-eth/src/server.ts | 27 +- dstack/kms/auth-eth/src/types.ts | 4 + dstack/kms/auth-eth/test/DstackKms.t.sol | 46 ++- dstack/kms/auth-mock/README.md | 4 +- dstack/kms/auth-mock/index.ts | 12 +- dstack/kms/auth-mock/openapi.json | 71 +++- dstack/kms/auth-simple/README.md | 4 +- dstack/kms/auth-simple/index.test.ts | 5 + dstack/kms/auth-simple/index.ts | 12 +- dstack/kms/src/main.rs | 20 +- dstack/kms/src/main_service.rs | 309 +++++++++--------- .../kms/src/main_service/upgrade_authority.rs | 43 ++- 22 files changed, 531 insertions(+), 221 deletions(-) diff --git a/docs/deployment.md b/docs/deployment.md index 4ec2d5aaf..56c3fe7f8 100644 --- a/docs/deployment.md +++ b/docs/deployment.md @@ -549,13 +549,17 @@ for the full procedure; the KMS steps are: that imported the dropped N must be onboarded again. Each `checkAuthorization` result carries the current epoch and its chain hash, -including when authorization is denied. The KMS compares these directly with -its serving roots and requests a restart on a mismatch; it does not read the -cached info endpoint to decide. No registered epoch returns a rejection with -`rejectReason: "Root epoch not registered"` and a zero hash, which does not -trigger a restart. Register epoch 0 before normal key release. Bootstrap keeps -using scoped KMS authorization so it can generate the first roots before -registration. See the [authorization protocol](../dstack/kms/README.md#authorization-protocol). +including when authorization is denied. KMS verifies these against its keyring; +there is no cached serving epoch. An optional `requestedEpoch` selects historical +keys, and omission selects the returned current epoch. Requests for future +epochs are denied. No registered epoch returns a rejection with +`rejectReason: "Root epoch not registered"` and a zero hash. + +RPC certificates still require a listener restart when their root CA changes; +callers retry after the certificate is refreshed. Register epoch 0 before normal +key release. Bootstrap keeps using scoped KMS authorization so it can generate +the first roots before registration. See the +[authorization protocol](../dstack/kms/README.md#authorization-protocol). --- diff --git a/dstack/kms/README.md b/dstack/kms/README.md index aff3f93aa..8f69cf08e 100644 --- a/dstack/kms/README.md +++ b/dstack/kms/README.md @@ -249,10 +249,18 @@ The `SignCert` RPC is used by the dstack app to sign a TLS certificate. In this Normal KMS operations use `POST /checkAuthorization`. The request carries the verified attestation identity (the same identity fields as the legacy `/bootAuth/*` requests). The Ethereum backends call -`DstackKms.checkAuthorization(AppInfo)` once; the contract selects KMS policy when +`DstackKms.checkAuthorization(AuthorizationRequest)` once; the contract selects KMS policy when `appId == address(this)` (the proxy address), and app policy otherwise. KMS-only operations also check this identity before accepting the result. +`requestedEpoch` 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 +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 omits this selector and still transfers the complete keyring, +including pending epochs. + The response is an `AuthorizationResult`: ```json @@ -271,11 +279,17 @@ registered epoch, the result is `isAllowed: false`, `rejectReason: "Root epoch not registered"`, epoch 0 and a zero hash. A zero hash never authorizes key release or requests an epoch switch. -The KMS compares each result with its serving epoch and chain hash before -releasing keys or signing certificates. A mismatch requests a process restart, -including when policy also denies authorization. This does not consult the -cached info endpoint. Rotation is detected on authorization requests, not by -background polling; self-authorization retains its one-second cache. +The KMS holds no serving-epoch snapshot. It verifies the returned current epoch +and chain hash against its keyring, then derives keys from `requestedEpoch` when +specified or the returned current epoch otherwise. Missing or mismatched roots +fail closed. Self-authorization retains its one-second cache of complete results. + +The TLS listener has a separate lifetime: when an authorized current CA differs +from the running RPC certificate's issuer, the KMS requests a restart to refresh +that certificate and the caller must retry. The RPC CA does not select key +epochs. `GetTempCaCert` reports the actual TLS issuer; application key responses +and signed certificate chains report the selected epoch's CA. CA certificates +are reloaded after renewal at startup. Upgrade the contract and auth backend and register epoch 0 before using this protocol for key release. Runtime authorization never falls back to the legacy diff --git a/dstack/kms/auth-eth-bun/README.md b/dstack/kms/auth-eth-bun/README.md index 92f6842e4..3e202644b 100644 --- a/dstack/kms/auth-eth-bun/README.md +++ b/dstack/kms/auth-eth-bun/README.md @@ -205,7 +205,9 @@ the test suite (`index.test.ts`) validates: `POST /checkAuthorization` returns `AuthorizationResult`, including `rejectReason`, `gatewayAppId`, `currentRootEpoch`, and `currentRootEpochHash`. -It selects KMS or app policy from the verified `appId`. See the +It selects KMS or app policy from the verified `appId`. The optional +`requestedEpoch` 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, bootstrap, and compatibility requirements. The legacy `/bootAuth/*` endpoints remain available. diff --git a/dstack/kms/auth-eth-bun/index.test.ts b/dstack/kms/auth-eth-bun/index.test.ts index 073b8a849..3f07d5c17 100644 --- a/dstack/kms/auth-eth-bun/index.test.ts +++ b/dstack/kms/auth-eth-bun/index.test.ts @@ -841,8 +841,8 @@ describe('POST /checkAuthorization', () => { mrAggregated: '0x' + '44'.repeat(32), mrSystem: '0x' + '55'.repeat(32), osImageHash: '0x' + '66'.repeat(32), tcbStatus: 'UpToDate', advisoryIds: [], }; - const authorize = () => appFetch(new Request('http://localhost/checkAuthorization', { - method: 'POST', headers: { 'Content-Type': 'application/json' }, body: JSON.stringify(appInfo), + const authorize = (requestedEpoch?: number) => appFetch(new Request('http://localhost/checkAuthorization', { + method: 'POST', headers: { 'Content-Type': 'application/json' }, body: JSON.stringify({ ...appInfo, requestedEpoch }), })); it('reads the complete result with one contract call at configured finality', async () => { @@ -857,7 +857,18 @@ describe('POST /checkAuthorization', () => { expect(mockReadContract).toHaveBeenCalledTimes(1); expect(mockReadContract).toHaveBeenCalledWith(expect.objectContaining({ functionName: 'checkAuthorization', blockNumber: 98n, - args: [expect.objectContaining({ appId: appInfo.appId })], + args: [expect.objectContaining({ appInfo: expect.objectContaining({ appId: appInfo.appId }), requestedEpoch: { isSet: false, value: 0n } })], + })); + }); + + it('encodes an explicit epoch zero as present', async () => { + mockReadContract.mockResolvedValueOnce({ + isAllowed: true, rejectReason: '', gatewayAppId: 'gateway', + currentRootEpoch: 1n, currentRootEpochHash: '0x' + '77'.repeat(32), + }); + expect((await authorize(0)).status).toBe(200); + expect(mockReadContract).toHaveBeenCalledWith(expect.objectContaining({ + args: [expect.objectContaining({ requestedEpoch: { isSet: true, value: 0n } })], })); }); diff --git a/dstack/kms/auth-eth-bun/index.ts b/dstack/kms/auth-eth-bun/index.ts index 796d83b06..933e5923c 100644 --- a/dstack/kms/auth-eth-bun/index.ts +++ b/dstack/kms/auth-eth-bun/index.ts @@ -36,6 +36,11 @@ const BootInfoSchema = z.object({ rootEpochHash: boundedHex(32, 'root epoch chain hash').optional() }); +const AuthorizationRequestSchema = BootInfoSchema.omit({ rootEpoch: true, rootEpochHash: true }).extend({ + requestedEpoch: z.number().int().min(0).max(0xffffffff).optional(), +}); +type AuthorizationRequest = z.infer; + const BootResponseSchema = z.object({ isAllowed: z.boolean(), reason: z.string(), @@ -89,7 +94,12 @@ const DSTACK_KMS_ABI = [ name: 'checkAuthorization', type: 'function', stateMutability: 'view', - inputs: [{ name: 'appInfo', type: 'tuple', components: APP_INFO_COMPONENTS }], + inputs: [{ name: 'request', type: 'tuple', components: [ + { name: 'appInfo', type: 'tuple', components: APP_INFO_COMPONENTS }, + { name: 'requestedEpoch', type: 'tuple', components: [ + { name: 'isSet', type: 'bool' }, { name: 'value', type: 'uint256' }, + ] }, + ] }], outputs: [{ name: 'result', type: 'tuple', @@ -223,13 +233,16 @@ class EthereumBackend { } - async checkAuthorization(appInfo: BootInfo) { + async checkAuthorization(appInfo: AuthorizationRequest) { const blockNumber = await this.finalizedBlockNumber(); const result = await this.client.readContract({ address: this.kmsContractAddr, abi: DSTACK_KMS_ABI, functionName: 'checkAuthorization', - args: [this.contractAppInfo(appInfo)], + args: [{ + appInfo: this.contractAppInfo(appInfo), + requestedEpoch: { isSet: appInfo.requestedEpoch !== undefined, value: BigInt(appInfo.requestedEpoch ?? 0) }, + }], blockNumber, }); const currentRootEpoch = Number(result.currentRootEpoch); @@ -401,7 +414,7 @@ app.get('/', async (c) => { app.post('/checkAuthorization', bodyLimit({ maxSize: MAX_BOOT_INFO_BYTES }), - zValidator('json', BootInfoSchema), + zValidator('json', AuthorizationRequestSchema), async (c) => { try { return c.json(await ethereum.checkAuthorization(c.req.valid('json'))); diff --git a/dstack/kms/auth-eth-bun/openapi.json b/dstack/kms/auth-eth-bun/openapi.json index f7e222054..2660158d3 100644 --- a/dstack/kms/auth-eth-bun/openapi.json +++ b/dstack/kms/auth-eth-bun/openapi.json @@ -108,7 +108,7 @@ "content": { "application/json": { "schema": { - "$ref": "#/components/schemas/AppInfo" + "$ref": "#/components/schemas/AuthorizationRequest" } } } @@ -288,11 +288,56 @@ } }, "AppInfo": { - "allOf": [ - { - "$ref": "#/components/schemas/BootInfo" - } + "type": "object", + "required": [ + "mrAggregated", + "osImageHash", + "appId", + "composeHash", + "instanceId", + "deviceId" ], + "properties": { + "mrAggregated": { + "type": "string", + "description": "Aggregated MR measurement" + }, + "osImageHash": { + "type": "string", + "description": "OS Image hash" + }, + "appId": { + "type": "string", + "description": "Application ID" + }, + "composeHash": { + "type": "string", + "description": "Compose hash" + }, + "instanceId": { + "type": "string", + "description": "Instance ID" + }, + "deviceId": { + "type": "string", + "description": "Device ID" + }, + "tcbStatus": { + "type": "string", + "description": "TCB status (optional for compatibility)" + }, + "advisoryIds": { + "type": "array", + "items": { + "type": "string" + }, + "description": "Advisory IDs (optional for compatibility)" + }, + "mrSystem": { + "type": "string", + "description": "MR System measurement (optional for compatibility)" + } + }, "description": "Verified attestation identity. The KMS app ID is the KMS contract address." }, "AuthorizationResult": { @@ -325,6 +370,22 @@ "description": "Zero when no epoch is registered; authorization is denied." } } + }, + "AuthorizationRequest": { + "allOf": [ + { + "$ref": "#/components/schemas/AppInfo" + } + ], + "type": "object", + "properties": { + "requestedEpoch": { + "type": "integer", + "minimum": 0, + "maximum": 4294967295, + "description": "Omit to use the current epoch. An explicit zero selects epoch 0; future epochs are denied." + } + } } } } diff --git a/dstack/kms/auth-eth/README.md b/dstack/kms/auth-eth/README.md index d8a91fad4..397064ac7 100644 --- a/dstack/kms/auth-eth/README.md +++ b/dstack/kms/auth-eth/README.md @@ -176,7 +176,9 @@ Documentation: https://book.getfoundry.sh/ `POST /checkAuthorization` returns `AuthorizationResult`, including `rejectReason`, `gatewayAppId`, `currentRootEpoch`, and `currentRootEpochHash`. -It selects KMS or app policy from the verified `appId`. See the +It selects KMS or app policy from the verified `appId`. The optional +`requestedEpoch` 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, bootstrap, and compatibility requirements. The legacy `/bootAuth/*` endpoints remain available. diff --git a/dstack/kms/auth-eth/contracts/DstackKms.sol b/dstack/kms/auth-eth/contracts/DstackKms.sol index 7ce8a53f7..2b7b2e343 100644 --- a/dstack/kms/auth-eth/contracts/DstackKms.sol +++ b/dstack/kms/auth-eth/contracts/DstackKms.sol @@ -26,6 +26,16 @@ contract DstackKms is Initializable, Ownable2StepUpgradeable, UUPSUpgradeable, E string[] advisoryIds; } + struct OptionalEpoch { + bool isSet; + uint256 value; + } + + struct AuthorizationRequest { + AppInfo appInfo; + OptionalEpoch requestedEpoch; + } + struct AuthorizationResult { bool isAllowed; string rejectReason; @@ -392,7 +402,11 @@ contract DstackKms is Initializable, Ownable2StepUpgradeable, UUPSUpgradeable, E /// Authorize an attested identity and return the root epoch from the same state. /// The KMS app ID is this contract's proxy address; all other IDs use app policy. - function checkAuthorization(AppInfo calldata appInfo) external view returns (AuthorizationResult memory result) { + function checkAuthorization(AuthorizationRequest calldata request) + external + view + returns (AuthorizationResult memory result) + { result.gatewayAppId = gatewayAppId; if (_rootEpochs.length == 0) { result.rejectReason = "Root epoch not registered"; @@ -400,6 +414,11 @@ contract DstackKms is Initializable, Ownable2StepUpgradeable, UUPSUpgradeable, E } result.currentRootEpoch = _currentRootEpoch; result.currentRootEpochHash = _rootEpochs[_currentRootEpoch].chainHash; + if (request.requestedEpoch.isSet && request.requestedEpoch.value > _currentRootEpoch) { + result.rejectReason = "Requested root epoch is not active"; + return result; + } + AppInfo calldata appInfo = request.appInfo; AppBootInfo memory bootInfo = AppBootInfo({ appId: appInfo.appId, composeHash: appInfo.composeHash, diff --git a/dstack/kms/auth-eth/src/ethereum.ts b/dstack/kms/auth-eth/src/ethereum.ts index a89d2fb9b..6af9cf360 100644 --- a/dstack/kms/auth-eth/src/ethereum.ts +++ b/dstack/kms/auth-eth/src/ethereum.ts @@ -3,11 +3,11 @@ // SPDX-License-Identifier: Apache-2.0 import { ethers } from 'ethers'; -import { BootInfo, BootResponse, RootEpochInfo, AuthorizationResult } from './types'; +import { BootInfo, BootResponse, RootEpochInfo, AuthorizationResult, AuthorizationRequest } from './types'; // Minimal ABI for DstackKms contract const DSTACK_KMS_ABI = [ - "function checkAuthorization((address appId,bytes32 composeHash,address instanceId,bytes32 deviceId,bytes32 mrAggregated,bytes32 mrSystem,bytes32 osImageHash,string tcbStatus,string[] advisoryIds) appInfo) view returns ((bool isAllowed,string rejectReason,string gatewayAppId,uint256 currentRootEpoch,bytes32 currentRootEpochHash))", + "function checkAuthorization(((address appId,bytes32 composeHash,address instanceId,bytes32 deviceId,bytes32 mrAggregated,bytes32 mrSystem,bytes32 osImageHash,string tcbStatus,string[] advisoryIds) appInfo,(bool isSet,uint256 value) requestedEpoch) request) view returns ((bool isAllowed,string rejectReason,string gatewayAppId,uint256 currentRootEpoch,bytes32 currentRootEpochHash))", "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)", "function isKmsAllowed((address appId,bytes32 composeHash,address instanceId,bytes32 deviceId,bytes32 mrAggregated,bytes32 mrSystem,bytes32 osImageHash,string tcbStatus,string[] advisoryIds) bootInfo) view returns (bool, string)", "function gatewayAppId() view returns (string)", @@ -77,8 +77,11 @@ export class EthereumBackend { }; } - async checkAuthorization(appInfo: BootInfo): Promise { - const result = await this.kmsContract.checkAuthorization(this.contractAppInfo(appInfo)); + async checkAuthorization(appInfo: AuthorizationRequest): Promise { + const result = await this.kmsContract.checkAuthorization({ + appInfo: this.contractAppInfo(appInfo), + requestedEpoch: { isSet: appInfo.requestedEpoch !== undefined, value: appInfo.requestedEpoch ?? 0 }, + }); return { isAllowed: result.isAllowed, rejectReason: result.rejectReason, diff --git a/dstack/kms/auth-eth/src/main.test.ts b/dstack/kms/auth-eth/src/main.test.ts index bcb44bd3c..1b15f500f 100644 --- a/dstack/kms/auth-eth/src/main.test.ts +++ b/dstack/kms/auth-eth/src/main.test.ts @@ -102,6 +102,8 @@ describe('Server', () => { expect(response.statusCode).toBe(200); expect(JSON.parse(response.payload)).toEqual(decision); expect(app.ethereum.checkAuthorization).toHaveBeenCalledWith(mockBootInfo); + await app.inject({ method: 'POST', url: '/checkAuthorization', payload: { ...mockBootInfo, requestedEpoch: 0 } }); + expect(app.ethereum.checkAuthorization).toHaveBeenLastCalledWith({ ...mockBootInfo, requestedEpoch: 0 }); }); it('fails closed when the new contract method is unavailable', async () => { diff --git a/dstack/kms/auth-eth/src/server.ts b/dstack/kms/auth-eth/src/server.ts index 225607f80..c702e8749 100644 --- a/dstack/kms/auth-eth/src/server.ts +++ b/dstack/kms/auth-eth/src/server.ts @@ -4,7 +4,7 @@ import fastify, { FastifyInstance } from 'fastify'; import { EthereumBackend } from './ethereum'; -import { BootInfo, BootResponse, AuthorizationResult } from './types'; +import { BootInfo, BootResponse, AuthorizationResult, AuthorizationRequest } from './types'; import { ethers } from 'ethers'; declare module 'fastify' { @@ -25,8 +25,7 @@ export async function build(): Promise { description, }); - server.addSchema({ - $id: 'bootInfo', + const appInfoSchema = { type: 'object', required: ['mrAggregated', 'osImageHash', 'appId', 'composeHash', 'instanceId', 'deviceId'], properties: { @@ -39,9 +38,25 @@ export async function build(): Promise { tcbStatus: { type: 'string', maxLength: 128, default: '' }, advisoryIds: { type: 'array', maxItems: 128, items: { type: 'string', maxLength: 256 }, default: [] }, mrSystem: { ...hex(32, 'System MR measurement'), default: '' }, + } + }; + + server.addSchema({ + ...appInfoSchema, + $id: 'bootInfo', + properties: { + ...appInfoSchema.properties, rootEpoch: { type: 'integer', minimum: 0 }, rootEpochHash: hex(32, 'Chain hash of the root epochs up to rootEpoch'), - } + }, + }); + server.addSchema({ + ...appInfoSchema, + $id: 'authorizationRequest', + properties: { + ...appInfoSchema.properties, + requestedEpoch: { type: 'integer', minimum: 0, maximum: 4294967295 }, + }, }); server.addSchema({ @@ -104,9 +119,9 @@ export async function build(): Promise { } }); - server.post<{ Body: BootInfo; Reply: AuthorizationResult }>('/checkAuthorization', { + server.post<{ Body: AuthorizationRequest; Reply: AuthorizationResult }>('/checkAuthorization', { schema: { - body: { $ref: 'bootInfo#' }, + body: { $ref: 'authorizationRequest#' }, response: { 200: { type: 'object', diff --git a/dstack/kms/auth-eth/src/types.ts b/dstack/kms/auth-eth/src/types.ts index 7fbd6aed5..f59adf624 100644 --- a/dstack/kms/auth-eth/src/types.ts +++ b/dstack/kms/auth-eth/src/types.ts @@ -40,3 +40,7 @@ export interface AuthorizationResult { currentRootEpoch: number; currentRootEpochHash: string; } + +export interface AuthorizationRequest extends Omit { + requestedEpoch?: number; +} diff --git a/dstack/kms/auth-eth/test/DstackKms.t.sol b/dstack/kms/auth-eth/test/DstackKms.t.sol index 777e7fdbf..e6f358d64 100644 --- a/dstack/kms/auth-eth/test/DstackKms.t.sol +++ b/dstack/kms/auth-eth/test/DstackKms.t.sol @@ -456,8 +456,42 @@ contract DstackKmsTest is Test { }); } + function authorize(DstackKms.AppInfo memory info) internal view returns (DstackKms.AuthorizationResult memory) { + return kms.checkAuthorization(DstackKms.AuthorizationRequest(info, DstackKms.OptionalEpoch(false, 0))); + } + + function test_CheckAuthorization_OptionalEpochDistinguishesZeroFromDefault() public { + DstackKms.AppInfo memory info = authorizationInfo(address(kms)); + vm.startPrank(owner); + kms.registerRootEpoch(ca(0xa0), k256(0xb0), "https://example.com/epoch-0"); + kms.registerRootEpoch(ca(0xa1), k256(0xb1), "https://example.com/epoch-1"); + kms.addOsImageHash(info.osImageHash); + kms.addKmsDevice(info.deviceId); + kms.addKmsAggregatedMr(info.mrAggregated); + vm.stopPrank(); + + DstackKms.AuthorizationRequest memory request = + DstackKms.AuthorizationRequest(info, DstackKms.OptionalEpoch(true, 1)); + DstackKms.AuthorizationResult memory result = kms.checkAuthorization(request); + assertFalse(result.isAllowed); + assertEq(result.rejectReason, "Requested root epoch is not active"); + bytes32 h1 = chainHash(chainHash(bytes32(0), ca(0xa0), k256(0xb0)), ca(0xa1), k256(0xb1)); + vm.prank(owner); + kms.activateRootEpoch(1, h1); + assertTrue(kms.checkAuthorization(request).isAllowed); + request.requestedEpoch.value = 0; + assertTrue(kms.checkAuthorization(request).isAllowed); + request.requestedEpoch.value = type(uint256).max; + assertFalse(kms.checkAuthorization(request).isAllowed); + request.requestedEpoch.isSet = false; + result = kms.checkAuthorization(request); + assertTrue(result.isAllowed); + assertEq(result.currentRootEpoch, 1); + assertEq(result.currentRootEpochHash, h1); + } + function test_CheckAuthorization_UnregisteredEpochReturnsRejection() public view { - DstackKms.AuthorizationResult memory result = kms.checkAuthorization(authorizationInfo(address(kms))); + DstackKms.AuthorizationResult memory result = authorize(authorizationInfo(address(kms))); assertFalse(result.isAllowed); assertEq(result.rejectReason, "Root epoch not registered"); assertEq(result.currentRootEpoch, 0); @@ -475,23 +509,23 @@ contract DstackKmsTest is Test { address appId = kms.deployAndRegisterApp(user, false, false, info.deviceId, info.composeHash); vm.stopPrank(); - DstackKms.AuthorizationResult memory result = kms.checkAuthorization(info); + DstackKms.AuthorizationResult memory result = authorize(info); assertTrue(result.isAllowed); assertEq(result.rejectReason, ""); assertEq(result.gatewayAppId, "gateway"); assertEq(result.currentRootEpochHash, chainHash(bytes32(0), ca(0xa0), k256(0xb0))); info.appId = user; - result = kms.checkAuthorization(info); + result = authorize(info); assertFalse(result.isAllowed); assertEq(result.rejectReason, "App not registered"); info.appId = appId; info.mrAggregated = bytes32(0); - result = kms.checkAuthorization(info); + result = authorize(info); assertTrue(result.isAllowed); info.appId = address(kms); - result = kms.checkAuthorization(info); + result = authorize(info); assertFalse(result.isAllowed); assertEq(result.rejectReason, "Aggregated MR not allowed"); } @@ -503,7 +537,7 @@ contract DstackKmsTest is Test { kms.activateRootEpoch(1); vm.stopPrank(); - DstackKms.AuthorizationResult memory result = kms.checkAuthorization(authorizationInfo(user)); + DstackKms.AuthorizationResult memory result = authorize(authorizationInfo(user)); assertFalse(result.isAllowed); assertEq(result.rejectReason, "App not registered"); assertEq(result.currentRootEpoch, 1); diff --git a/dstack/kms/auth-mock/README.md b/dstack/kms/auth-mock/README.md index 961acf432..665695c98 100644 --- a/dstack/kms/auth-mock/README.md +++ b/dstack/kms/auth-mock/README.md @@ -173,7 +173,9 @@ the API remains identical, so no client code changes are needed. `POST /checkAuthorization` returns `AuthorizationResult`, including `rejectReason`, `gatewayAppId`, `currentRootEpoch`, and `currentRootEpochHash`. -It selects KMS or app policy from the verified `appId`. See the +It selects KMS or app policy from the verified `appId`. The optional +`requestedEpoch` 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, bootstrap, and compatibility requirements. The legacy `/bootAuth/*` endpoints remain available. diff --git a/dstack/kms/auth-mock/index.ts b/dstack/kms/auth-mock/index.ts index b8d0b586d..408c813fb 100644 --- a/dstack/kms/auth-mock/index.ts +++ b/dstack/kms/auth-mock/index.ts @@ -29,6 +29,11 @@ const BootInfoSchema = z.object({ rootEpochHash: z.string().optional() }); +const AuthorizationRequestSchema = BootInfoSchema.omit({ rootEpoch: true, rootEpochHash: true }).extend({ + requestedEpoch: z.number().int().min(0).max(0xffffffff).optional(), +}); +type AuthorizationRequest = z.infer; + const BootResponseSchema = z.object({ isAllowed: z.boolean(), reason: z.string(), @@ -106,7 +111,7 @@ class MockBackend { this.mockAppImplementation = process.env.MOCK_APP_IMPLEMENTATION || '0xmockapp9876543210987654321098765432109'; } - async checkAuthorization(appInfo: BootInfo) { + async checkAuthorization(appInfo: AuthorizationRequest) { const current = getRootEpoch(); const epoch = { currentRootEpoch: current.rootEpoch ?? 0, @@ -115,6 +120,9 @@ class MockBackend { if (current.rootEpoch === undefined) { return { isAllowed: false, rejectReason: 'Root epoch not registered', gatewayAppId: this.mockGatewayAppId, ...epoch }; } + if (appInfo.requestedEpoch !== undefined && appInfo.requestedEpoch > current.rootEpoch) { + return { isAllowed: false, rejectReason: 'Requested root epoch is not active', gatewayAppId: this.mockGatewayAppId, ...epoch }; + } const kmsAppId = process.env.KMS_CONTRACT_ADDR || '0xmockcontract1234567890123456789012345678'; const result = this.evaluateBoot(appInfo, normHash(appInfo.appId) === normHash(kmsAppId)); return { isAllowed: result.isAllowed, rejectReason: result.isAllowed ? '' : result.reason, @@ -217,7 +225,7 @@ app.get('/', async (c) => { } }); -app.post('/checkAuthorization', zValidator('json', BootInfoSchema), async (c) => { +app.post('/checkAuthorization', zValidator('json', AuthorizationRequestSchema), async (c) => { try { return c.json(await mockBackend.checkAuthorization(c.req.valid('json'))); } catch { diff --git a/dstack/kms/auth-mock/openapi.json b/dstack/kms/auth-mock/openapi.json index 6679d20a8..d9c75f5e3 100644 --- a/dstack/kms/auth-mock/openapi.json +++ b/dstack/kms/auth-mock/openapi.json @@ -108,7 +108,7 @@ "content": { "application/json": { "schema": { - "$ref": "#/components/schemas/AppInfo" + "$ref": "#/components/schemas/AuthorizationRequest" } } } @@ -297,11 +297,56 @@ } }, "AppInfo": { - "allOf": [ - { - "$ref": "#/components/schemas/BootInfo" - } + "type": "object", + "required": [ + "mrAggregated", + "osImageHash", + "appId", + "composeHash", + "instanceId", + "deviceId" ], + "properties": { + "mrAggregated": { + "type": "string", + "description": "Aggregated MR measurement" + }, + "osImageHash": { + "type": "string", + "description": "OS Image hash" + }, + "appId": { + "type": "string", + "description": "Application ID" + }, + "composeHash": { + "type": "string", + "description": "Compose hash" + }, + "instanceId": { + "type": "string", + "description": "Instance ID" + }, + "deviceId": { + "type": "string", + "description": "Device ID" + }, + "tcbStatus": { + "type": "string", + "description": "TCB status (optional for compatibility)" + }, + "advisoryIds": { + "type": "array", + "items": { + "type": "string" + }, + "description": "Advisory IDs (optional for compatibility)" + }, + "mrSystem": { + "type": "string", + "description": "MR System measurement (optional for compatibility)" + } + }, "description": "Verified attestation identity. The KMS app ID is the KMS contract address." }, "AuthorizationResult": { @@ -334,6 +379,22 @@ "description": "Zero when no epoch is registered; authorization is denied." } } + }, + "AuthorizationRequest": { + "allOf": [ + { + "$ref": "#/components/schemas/AppInfo" + } + ], + "type": "object", + "properties": { + "requestedEpoch": { + "type": "integer", + "minimum": 0, + "maximum": 4294967295, + "description": "Omit to use the current epoch. An explicit zero selects epoch 0; future epochs are denied." + } + } } } } diff --git a/dstack/kms/auth-simple/README.md b/dstack/kms/auth-simple/README.md index 14f930cf9..d0f7f00fc 100644 --- a/dstack/kms/auth-simple/README.md +++ b/dstack/kms/auth-simple/README.md @@ -258,7 +258,9 @@ bun run test:run `POST /checkAuthorization` returns `AuthorizationResult`, including `rejectReason`, `gatewayAppId`, `currentRootEpoch`, and `currentRootEpochHash`. -It selects KMS or app policy from the verified `appId`. See the +It selects KMS or app policy from the verified `appId`. The optional +`requestedEpoch` 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, bootstrap, and compatibility requirements. The legacy `/bootAuth/*` endpoints remain available. diff --git a/dstack/kms/auth-simple/index.test.ts b/dstack/kms/auth-simple/index.test.ts index f45087f10..191137cd3 100644 --- a/dstack/kms/auth-simple/index.test.ts +++ b/dstack/kms/auth-simple/index.test.ts @@ -484,6 +484,11 @@ describe('auth-simple', () => { expect(await authorize()).toMatchObject({ isAllowed: true, rejectReason: '', currentRootEpoch: 1, currentRootEpochHash: h1, }); + expect(await authorize({ requestedEpoch: 0 })).toMatchObject({ isAllowed: true, currentRootEpoch: 1 }); + expect(await authorize({ requestedEpoch: 1 })).toMatchObject({ isAllowed: true }); + expect(await authorize({ requestedEpoch: 2 })).toMatchObject({ + isAllowed: false, rejectReason: 'Requested root epoch is not active', currentRootEpoch: 1, currentRootEpochHash: h1, + }); expect(await authorize({ appId: kmsContractAddr, composeHash: 'not-an-app-compose' })).toMatchObject({ isAllowed: true }); expect(await authorize({ appId: kmsContractAddr, mrAggregated: '0xbad' })).toMatchObject({ isAllowed: false, currentRootEpoch: 1, currentRootEpochHash: h1, diff --git a/dstack/kms/auth-simple/index.ts b/dstack/kms/auth-simple/index.ts index a1da5814e..c639d5c23 100644 --- a/dstack/kms/auth-simple/index.ts +++ b/dstack/kms/auth-simple/index.ts @@ -27,6 +27,11 @@ const BootInfoSchema = z.object({ rootEpochHash: z.string().optional() }); +const AuthorizationRequestSchema = BootInfoSchema.omit({ rootEpoch: true, rootEpochHash: true }).extend({ + requestedEpoch: z.number().int().min(0).max(0xffffffff).optional(), +}); +type AuthorizationRequest = z.infer; + const BootResponseSchema = z.object({ isAllowed: z.boolean(), reason: z.string(), @@ -152,7 +157,7 @@ class ConfigBackend { } } - async checkAuthorization(appInfo: BootInfo) { + async checkAuthorization(appInfo: AuthorizationRequest) { const config = this.loadConfig(); const current = currentRootEpoch(config); const epoch = { @@ -162,6 +167,9 @@ class ConfigBackend { if (current.rootEpoch === undefined) { return { isAllowed: false, rejectReason: 'Root epoch not registered', gatewayAppId: config.gatewayAppId, ...epoch }; } + if (appInfo.requestedEpoch !== undefined && appInfo.requestedEpoch > current.rootEpoch) { + return { isAllowed: false, rejectReason: 'Requested root epoch is not active', gatewayAppId: config.gatewayAppId, ...epoch }; + } const isKms = normalizeHex(appInfo.appId) === normalizeHex(config.kmsContractAddr); const result = this.evaluateBoot(appInfo, isKms, config); return { isAllowed: result.isAllowed, rejectReason: result.reason, gatewayAppId: result.gatewayAppId, ...epoch }; @@ -347,7 +355,7 @@ app.get('/', async (c) => { } }); -app.post('/checkAuthorization', zValidator('json', BootInfoSchema), async (c) => { +app.post('/checkAuthorization', zValidator('json', AuthorizationRequestSchema), async (c) => { try { return c.json(await backend.checkAuthorization(c.req.valid('json'))); } catch { diff --git a/dstack/kms/src/main.rs b/dstack/kms/src/main.rs index 14796ad3c..205790d79 100644 --- a/dstack/kms/src/main.rs +++ b/dstack/kms/src/main.rs @@ -174,12 +174,11 @@ async fn main() -> Result<()> { } let keyring = main_service::Keyring::load(&config).context("Failed to load the root keys")?; - let serving_epoch = main_service::wait_for_servable_epoch(&config, &keyring).await; - info!("Serving root epoch {serving_epoch}"); + let rpc_epoch = main_service::wait_for_servable_epoch(&config, &keyring).await; + info!("Issuing RPC certificate under root epoch {rpc_epoch}"); info!("Updating certs"); - if let Err(err) = onboard_service::update_certs(&config, serving_epoch, keyring.latest()).await - { + if let Err(err) = onboard_service::update_certs(&config, rpc_epoch, keyring.latest()).await { if config.attest_rpc_cert { return Err(err).context( "Failed to reissue the attested KMS RPC certificate; refusing to start with a \ @@ -211,7 +210,11 @@ async fn main() -> Result<()> { } else { None }; - let state = main_service::KmsState::new(config, keyring, serving_epoch) + // Reload certificates after renewal before exposing them through RPCs. + let rpc_ca_cert = + fs_err::read_to_string(config.root_key_paths(rpc_epoch, keyring.latest()).ca_cert)?; + let keyring = main_service::Keyring::load(&config).context("Failed to reload the root keys")?; + let state = main_service::KmsState::new(config, keyring, rpc_ca_cert) .context("Failed to initialize KMS state")?; let quote_verifier = QuoteVerifier::new(state.attestation_verifier()); let admin_quote_verifier = QuoteVerifier::new(state.attestation_verifier()); @@ -249,8 +252,7 @@ async fn main() -> Result<()> { rocket = rocket.manage(quote_verifier); let main_srv = rocket.launch(); - // A move of the current root epoch on chain ends the process so it can be - // restarted onto the new epoch. + // Refresh the TLS listener when the authorized root CA changes. let restart = { let state = state.clone(); async move { state.restart_requested().await } @@ -284,12 +286,12 @@ async fn main() -> Result<()> { }; tokio::select! { result = servers => result?, - () = restart => bail!("the current root epoch changed; restarting to follow it"), + () = restart => bail!("the RPC root CA changed; restarting to refresh the TLS certificate"), } } None => tokio::select! { result = main_srv => result.map(drop).map_err(|err| anyhow!(err.to_string()))?, - () = restart => bail!("the current root epoch changed; restarting to follow it"), + () = restart => bail!("the RPC root CA changed; restarting to refresh the TLS certificate"), }, } Ok(()) diff --git a/dstack/kms/src/main_service.rs b/dstack/kms/src/main_service.rs index 72fe7428b..574f6312d 100644 --- a/dstack/kms/src/main_service.rs +++ b/dstack/kms/src/main_service.rs @@ -32,7 +32,7 @@ use scale::Decode; use tracing::warn; use upgrade_authority::{ build_boot_info, ensure_app_id_len, local_kms_boot_info, AuthorizationResult, BootInfo, - CurrentRootEpoch, GetInfoResponse, RootEpochRegistry, + GetInfoResponse, RootEpochRegistry, }; use crate::{ @@ -164,16 +164,16 @@ pub(crate) fn servable_epoch( pub struct KmsStateInner { config: KmsConfig, keyring: Keyring, - /// The epoch served for this process's lifetime; a newer current epoch on - /// chain makes the process restart. - serving_epoch: u32, + /// The CA of the running TLS listener, independent of key selection. + rpc_ca_cert: String, + rpc_ca_pubkey: Vec, restart: tokio::sync::Notify, temp_ca_cert: String, temp_ca_key: String, verifier: CvmVerifier, attestation_verifier: Arc, self_boot_info: Cache<(), BootInfo>, - self_allowed: Cache<(), ()>, + self_allowed: Cache<(), AuthorizationResult>, auth_api_info: Cache<(), GetInfoResponse>, metrics: KmsMetrics, } @@ -224,10 +224,15 @@ impl KmsState { self.verifier.clear_cache(image_hash, config_hash) } - pub fn new(config: KmsConfig, keyring: Keyring, serving_epoch: u32) -> Result { - if serving_epoch > keyring.latest() { - bail!("cannot serve root epoch {serving_epoch}, which this KMS does not hold"); - } + pub fn new(config: KmsConfig, keyring: Keyring, rpc_ca_cert: String) -> Result { + let (_, pem) = x509_parser::pem::parse_x509_pem(rpc_ca_cert.as_bytes()) + .context("invalid RPC CA certificate")?; + let rpc_ca_pubkey = pem + .parse_x509() + .context("invalid RPC CA certificate")? + .public_key() + .raw + .to_vec(); let temp_ca_key = fs::read_to_string(config.tmp_ca_key()).context("Faeild to read temp ca key")?; let temp_ca_cert = @@ -251,7 +256,8 @@ impl KmsState { inner: Arc::new(KmsStateInner { config, keyring, - serving_epoch, + rpc_ca_cert, + rpc_ca_pubkey, restart: tokio::sync::Notify::new(), temp_ca_cert, temp_ca_key, @@ -276,38 +282,24 @@ impl KmsState { self.inner.attestation_verifier.clone() } - /// The served epoch's root keys. - fn root(&self) -> &RootKeys { - self.epoch_root(self.serving_epoch) - } - - fn epoch_root(&self, epoch: u32) -> &RootKeys { - &self.keyring.roots[epoch as usize] + fn epoch_root(&self, epoch: u32) -> Result<&RootKeys> { + self.keyring + .roots + .get(epoch as usize) + .with_context(|| format!("root epoch {epoch} is not held by this KMS")) } - /// Resolve a root epoch selector: unset is the served epoch. Epochs held - /// but not yet current are not served. - fn resolve_epoch(&self, requested: Option) -> Result { - let serving = self.serving_epoch; - match requested { - None => Ok(serving), - Some(epoch) if epoch <= serving => Ok(epoch), - Some(epoch) => bail!("root epoch {epoch} is not served; the current one is {serving}"), - } - } - - /// Resolves once governance moves the current root epoch away from the - /// served one, so the process can restart to follow it. pub async fn restart_requested(&self) { self.restart.notified().await } - /// The served epoch and chain hash to compare with each authorization result. - fn serving_root_epoch(&self) -> CurrentRootEpoch { - CurrentRootEpoch { - epoch: self.serving_epoch, - chain_hash: self.keyring.chain_hashes[self.serving_epoch as usize], - } + async fn current_epoch_from_info(&self) -> Result { + let info = self + .auth_api_info + .try_get_with((), self.config.auth_api.get_info()) + .await + .map_err(|err| anyhow!("{err:#}"))?; + self.keyring.servable_epoch(info.root_epochs) } /// KMS-only operations must not accept an app that passes ordinary app policy. @@ -332,33 +324,43 @@ impl KmsState { Ok(()) } - /// Compare against the epoch returned with this decision, never a cached info response. - fn ensure_authorized(&self, result: &AuthorizationResult) -> Result<()> { - if !self.config.auth_api.is_dev() { - if let Some(current) = result.current_root()? { - if current != self.serving_root_epoch() { - self.restart.notify_one(); - bail!("the authorized root epoch or chain hash changed; restarting"); - } - } else { - bail!( - "Authorization denied: {}", - if result.reject_reason.is_empty() { - "Root epoch not registered" - } else { - &result.reject_reason - } - ); - } - } + /// Resolve this request from its authorization result; no serving epoch is cached. + fn authorized_epoch( + &self, + result: &AuthorizationResult, + requested: Option, + ) -> Result { + let current = if self.config.auth_api.is_dev() { + self.keyring.latest() + } else { + let current = result + .current_root()? + .context("Root epoch not registered")?; + self.keyring + .servable_epoch(RootEpochRegistry::Current(current))? + }; if !result.is_allowed { bail!("Authorization denied: {}", result.reject_reason); } + let epoch = requested.unwrap_or(current); + if epoch > current { + bail!("root epoch {epoch} is not active; the current one is {current}"); + } + self.epoch_root(epoch)?; + Ok(epoch) + } + + /// Keep the running TLS issuer consistent with default key responses until reload. + fn ensure_rpc_ca(&self, current_epoch: u32) -> Result<()> { + if self.epoch_root(current_epoch)?.ca.key.public_key_der() != self.rpc_ca_pubkey { + self.restart.notify_one(); + bail!("the RPC root CA changed; retry after the TLS certificate is refreshed"); + } Ok(()) } } -/// Wait until governance lets `keyring` serve, and return the epoch to serve. +/// Wait for the registered current roots to issue the RPC certificate at startup. /// An epoch-0 keyring may start while the info endpoint is unreachable, but /// key release still requires an authorization result with registered roots. pub(crate) async fn wait_for_servable_epoch(config: &KmsConfig, keyring: &Keyring) -> u32 { @@ -452,6 +454,7 @@ impl RpcHandler { struct BootConfig { boot_info: BootInfo, gateway_app_id: String, + root_epoch: u32, } pub(crate) fn build_boot_info_for_attestation( @@ -518,9 +521,11 @@ fn ensure_self_key_release_allowed( } impl RpcHandler { - async fn ensure_self_allowed(&self) -> Result> { + async fn ensure_self_allowed(&self) -> Result<(Option, u32)> { if !self.state.config.enforce_self_authorization { - return Ok(None); + let epoch = self.state.current_epoch_from_info().await?; + self.state.ensure_rpc_ca(epoch)?; + return Ok((None, epoch)); } let boot_info = self .state @@ -530,21 +535,25 @@ impl RpcHandler { .map_err(|err| anyhow!("{err:#}")) .context("Failed to load cached self boot info")?; self.state.ensure_kms_identity(&boot_info).await?; - self.state + let response = self + .state .self_allowed .try_get_with((), async { let response = self .state .config .auth_api - .check_authorization(&boot_info) + .check_authorization(&boot_info, None) .await .context("Failed to call self KMS auth check")?; - self.state.ensure_authorized(&response) + let epoch = self.state.authorized_epoch(&response, None)?; + self.state.ensure_rpc_ca(epoch)?; + Ok::<_, anyhow::Error>(response) }) .await .map_err(|err| anyhow!("{err:#}"))?; - Ok(Some(boot_info)) + let epoch = self.state.authorized_epoch(&response, None)?; + Ok((Some(boot_info), epoch)) } fn ensure_attested(&self) -> Result<&VerifiedAttestation> { @@ -556,14 +565,18 @@ impl RpcHandler { async fn ensure_kms_allowed(&self, vm_config: &str) -> Result { let att = self.ensure_attested()?; - self.ensure_app_attestation_allowed(att, true, false, vm_config) + self.ensure_app_attestation_allowed(att, true, false, vm_config, None) .await .map(|c| c.boot_info) } - async fn ensure_app_boot_allowed(&self, vm_config: &str) -> Result { + async fn ensure_app_boot_allowed( + &self, + vm_config: &str, + requested_epoch: Option, + ) -> Result { let att = self.ensure_attested()?; - self.ensure_app_attestation_allowed(att, false, false, vm_config) + self.ensure_app_attestation_allowed(att, false, false, vm_config, requested_epoch) .await } @@ -573,6 +586,7 @@ impl RpcHandler { is_kms: bool, use_boottime_mr: bool, vm_config_str: &str, + requested_epoch: Option, ) -> Result { let boot_info = build_boot_info_for_attestation(att, use_boottime_mr, vm_config_str)?; if is_kms { @@ -582,9 +596,15 @@ impl RpcHandler { .state .config .auth_api - .check_authorization(&boot_info) + .check_authorization(&boot_info, requested_epoch) .await?; - self.state.ensure_authorized(&response)?; + let root_epoch = self.state.authorized_epoch(&response, requested_epoch)?; + let current_epoch = if self.state.config.auth_api.is_dev() { + self.state.keyring.latest() + } else { + response.current_root_epoch + }; + self.state.ensure_rpc_ca(current_epoch)?; // SNP rootfs/app/config binding is handled by the SNP launch-measurement // helper above. The legacy OS-image verifier is TDX-oriented and still // rejects SNP quotes; keep SNP on the explicit fail-closed helper path. @@ -602,12 +622,13 @@ impl RpcHandler { Ok(BootConfig { boot_info, gateway_app_id: response.gateway_app_id, + root_epoch, }) } - fn derive_app_ca(&self, app_id: &[u8]) -> Result { + fn derive_app_ca(&self, app_id: &[u8], root: &RootKeys) -> Result { let context_data = vec![app_id, b"app-ca"]; - let app_key = kdf::derive_p256_key_pair(&self.state.root().ca.key, &context_data) + let app_key = kdf::derive_p256_key_pair(&root.ca.key, &context_data) .context("Failed to derive app disk key")?; let req = CertRequest::builder() .key(&app_key) @@ -617,12 +638,7 @@ impl RpcHandler { .app_id(app_id) .special_usage("app:ca") .build(); - let app_ca = self - .state - .root() - .ca - .sign(req) - .context("Failed to sign App CA")?; + let app_ca = root.ca.sign(req).context("Failed to sign App CA")?; Ok(CaCert::from_parts(app_key, app_ca)) } } @@ -638,13 +654,13 @@ impl KmsRpc for RpcHandler { let BootConfig { boot_info, gateway_app_id, + root_epoch, } = self - .ensure_app_boot_allowed(&request.vm_config) + .ensure_app_boot_allowed(&request.vm_config, request.root_epoch) .await .context("App not allowed")?; ensure_key_release_allowed(&boot_info, (&self.state.config).into())?; - let root_epoch = self.state.resolve_epoch(request.root_epoch)?; - let root = self.state.epoch_root(root_epoch); + let root = self.state.epoch_root(root_epoch)?; let keys = derive_app_keys(root, &boot_info.app_id, &boot_info.instance_id)?; Ok(AppKeyResponse { @@ -661,12 +677,14 @@ impl KmsRpc for RpcHandler { } async fn get_app_env_encrypt_pub_key(self, request: AppId) -> Result { - self.ensure_self_allowed() + let (_, epoch) = self + .ensure_self_allowed() .await .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( - &self.state.root().ca.key, + &root.ca.key, &[&request.app_id[..], "env-encrypt-key".as_bytes()], ) .context("Failed to derive env encrypt key")?; @@ -681,7 +699,7 @@ impl KmsRpc for RpcHandler { // Legacy signature (without timestamp) for backward compatibility let signature = sign_message( - &self.state.root().k256, + &root.k256, b"dstack-env-encrypt-pubkey", &request.app_id, &public_key, @@ -690,7 +708,7 @@ impl KmsRpc for RpcHandler { // New signature with timestamp to prevent replay attacks let signature_v1 = sign_message_with_timestamp( - &self.state.root().k256, + &root.k256, b"dstack-env-encrypt-pubkey", &request.app_id, timestamp, @@ -716,16 +734,12 @@ impl KmsRpc for RpcHandler { .try_get_with((), self.state.config.auth_api.get_info()) .await .map_err(|err| anyhow!("{err:#}"))?; + let epoch = self.state.keyring.servable_epoch(info.root_epochs)?; + let root = self.state.epoch_root(epoch)?; Ok(GetMetaResponse { - ca_cert: self.state.root().ca.pem_cert.clone(), + ca_cert: root.ca.pem_cert.clone(), allow_any_upgrade: self.state.inner.config.auth_api.is_dev(), - k256_pubkey: self - .state - .root() - .k256 - .verifying_key() - .to_sec1_bytes() - .to_vec(), + k256_pubkey: root.k256.verifying_key().to_sec1_bytes().to_vec(), bootstrap_info, is_dev: self.state.config.auth_api.is_dev(), kms_contract_address: info.kms_contract_address, @@ -749,7 +763,7 @@ impl KmsRpc for RpcHandler { /// guests use self-issued RA-TLS certificates and fetch the root CA via GetMeta. /// The returned key authenticates nobody; client identity comes from attestation. async fn get_temp_ca_cert(self) -> Result { - let self_boot_info = self + let (self_boot_info, _) = self .ensure_self_allowed() .await .context("KMS self authorization failed")?; @@ -757,7 +771,7 @@ impl KmsRpc for RpcHandler { Ok(GetTempCaCertResponse { temp_ca_cert: self.state.inner.temp_ca_cert.clone(), temp_ca_key: self.state.inner.temp_ca_key.clone(), - ca_cert: self.state.root().ca.pem_cert.clone(), + ca_cert: self.state.rpc_ca_cert.clone(), }) } @@ -790,10 +804,11 @@ impl KmsRpc for RpcHandler { .await .context("Quote verification failed")?; let app_info = self - .ensure_app_attestation_allowed(&attestation, false, true, &request.vm_config) + .ensure_app_attestation_allowed(&attestation, false, true, &request.vm_config, None) .await?; ensure_key_release_allowed(&app_info.boot_info, (&self.state.config).into())?; - let app_ca = self.derive_app_ca(&app_info.boot_info.app_id)?; + let root = self.state.epoch_root(app_info.root_epoch)?; + let app_ca = self.derive_app_ca(&app_info.boot_info.app_id, root)?; // 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 @@ -806,7 +821,7 @@ impl KmsRpc for RpcHandler { certificate_chain: vec![ cert.pem(), app_ca.pem_cert.clone(), - self.state.root().ca.pem_cert.clone(), + root.ca.pem_cert.clone(), ], }) } @@ -1548,73 +1563,69 @@ mod tests { .extract() .unwrap(); let keyring = Keyring::load(&config).unwrap(); - KmsState::new(config, keyring, 0).unwrap() - } - - #[rocket::async_test] - async fn authorization_epoch_is_checked_without_reading_cached_info() { - for (changed, allowed, registered, restart) in [ - (false, true, true, false), - (false, false, true, false), - (true, true, true, true), - (true, false, true, true), - (false, false, false, false), - (false, true, false, false), - ] { - let dir = tempfile::tempdir().unwrap(); - let state = kms_state(dir.path(), "http://127.0.0.1:1"); - let mut current = state.serving_root_epoch(); - if changed { - current.epoch = u32::MAX; - } - let result = AuthorizationResult { - is_allowed: allowed, - reject_reason: if allowed { "" } else { "policy denied" }.into(), - gateway_app_id: "gateway".into(), - current_root_epoch: current.epoch, - current_root_epoch_hash: hex::encode(if registered { - current.chain_hash - } else { - [0; 32] - }), - }; - assert_eq!( - state.ensure_authorized(&result).is_ok(), - allowed && !changed && registered - ); - assert_eq!( - tokio::time::timeout(Duration::from_millis(10), state.restart_requested()) - .await - .is_ok(), - restart, - ); - } + KmsState::new(config, keyring, ca_cert.pem()).unwrap() } #[rocket::async_test] - async fn authorization_hash_mismatch_restarts_and_malformed_hash_is_rejected() { + async fn authorization_selects_epochs_without_a_serving_epoch_snapshot() { let dir = tempfile::tempdir().unwrap(); - let state = kms_state(dir.path(), "http://127.0.0.1:1"); - let mut result = AuthorizationResult { + let initial = kms_state(dir.path(), "http://127.0.0.1:1"); + let mut keyring = Keyring::load(&initial.config).unwrap(); + let key = KeyPair::generate_for(&ra_tls::rcgen::PKCS_ECDSA_P256_SHA256).unwrap(); + let cert = CertRequest::builder() + .key(&key) + .subject("Root epoch 1") + .ca_level(1) + .build() + .self_signed() + .unwrap(); + keyring.roots.push(RootKeys { + ca: CaCert::from_parts(key, cert), + k256: SigningKey::random(&mut rand::rngs::OsRng), + }); + keyring.chain_hashes = + root_chain_hashes(keyring.roots.iter().map(|root| (&root.ca.key, &root.k256))); + let state = + KmsState::new(initial.config.clone(), keyring, initial.rpc_ca_cert.clone()).unwrap(); + let result = |epoch: u32| AuthorizationResult { is_allowed: true, reject_reason: String::new(), gateway_app_id: String::new(), - current_root_epoch: 0, - current_root_epoch_hash: "invalid".into(), + current_root_epoch: epoch, + current_root_epoch_hash: hex::encode(state.keyring.chain_hashes[epoch as usize]), }; - assert!(state.ensure_authorized(&result).is_err()); - assert!( - tokio::time::timeout(Duration::from_millis(10), state.restart_requested()) - .await - .is_err() - ); - result.current_root_epoch_hash = hex::encode([0x55; 32]); - assert!(state.ensure_authorized(&result).is_err()); + assert_eq!(state.authorized_epoch(&result(0), None).unwrap(), 0); + assert!(state.authorized_epoch(&result(0), Some(1)).is_err()); + assert_eq!(state.authorized_epoch(&result(1), None).unwrap(), 1); + assert_eq!(state.authorized_epoch(&result(1), Some(0)).unwrap(), 0); + assert_eq!(state.authorized_epoch(&result(1), Some(1)).unwrap(), 1); + assert!(state.authorized_epoch(&result(1), Some(u32::MAX)).is_err()); + // Key selection already follows epoch 1; only the TLS listener needs refreshing. + assert!(state.ensure_rpc_ca(1).is_err()); assert!( tokio::time::timeout(Duration::from_millis(10), state.restart_requested()) .await .is_ok() ); + assert_eq!(state.rpc_ca_cert, initial.rpc_ca_cert); + + let mut rejected = result(1); + rejected.is_allowed = false; + rejected.reject_reason = "Policy denied".into(); + assert!(state.authorized_epoch(&rejected, Some(0)).is_err()); + rejected.is_allowed = true; + for hash in [ + "invalid".into(), + hex::encode([0; 32]), + hex::encode([0x55; 32]), + ] { + rejected.current_root_epoch_hash = hash; + assert!(state.authorized_epoch(&rejected, None).is_err()); + } + let mut missing = result(1); + missing.current_root_epoch = u32::MAX; + assert!(state.authorized_epoch(&missing, Some(0)).is_err()); + assert!(state.epoch_root(u32::MAX).is_err()); } #[rocket::async_test] diff --git a/dstack/kms/src/main_service/upgrade_authority.rs b/dstack/kms/src/main_service/upgrade_authority.rs index 40de002ec..a49d84066 100644 --- a/dstack/kms/src/main_service/upgrade_authority.rs +++ b/dstack/kms/src/main_service/upgrade_authority.rs @@ -86,7 +86,7 @@ pub(crate) struct BootResponse { pub reason: String, } -#[derive(Debug, Serialize, Deserialize)] +#[derive(Debug, Clone, Serialize, Deserialize)] #[serde(rename_all = "camelCase")] pub(crate) struct AuthorizationResult { pub is_allowed: bool, @@ -254,7 +254,19 @@ async fn send_request(req: reqwest::RequestBuilder, url: &s } impl AuthApi { - pub async fn check_authorization(&self, app_info: &BootInfo) -> Result { + pub async fn check_authorization( + &self, + app_info: &BootInfo, + requested_epoch: Option, + ) -> Result { + #[derive(Serialize)] + #[serde(rename_all = "camelCase")] + struct AuthorizationRequest<'a> { + #[serde(flatten)] + app_info: &'a BootInfo, + #[serde(skip_serializing_if = "Option::is_none")] + requested_epoch: Option, + } ensure_identity_widths(app_info)?; match self { AuthApi::Dev { dev } => Ok(AuthorizationResult { @@ -266,7 +278,15 @@ impl AuthApi { }), AuthApi::Webhook { webhook } => { let url = url_join(&webhook.url, "checkAuthorization"); - http_post(&url, app_info, webhook.timeout).await + http_post( + &url, + &AuthorizationRequest { + app_info, + requested_epoch, + }, + webhook.timeout, + ) + .await } } } @@ -661,21 +681,28 @@ pub(crate) mod tests { "11".repeat(32) ) }; - let (url, server) = serve(vec![response(true, 0), response(false, 1)]); + let (url, server) = serve(vec![ + response(true, 0), + response(false, 1), + response(true, 1), + ]); let auth = webhook(url); let info = boot_info(50); - let first = auth.check_authorization(&info).await.unwrap(); - let second = auth.check_authorization(&info).await.unwrap(); + let first = auth.check_authorization(&info, None).await.unwrap(); + let second = auth.check_authorization(&info, 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.check_authorization(&info, Some(0)).await.unwrap(); let requests = server.join().unwrap(); - assert_eq!(requests.len(), 2); + assert_eq!(requests.len(), 3); assert!(requests .iter() .all(|(path, _)| path == "/checkAuthorization")); assert_eq!(requests[0].1, serde_json::to_value(info).unwrap()); + assert!(requests[0].1.get("requestedEpoch").is_none()); + assert_eq!(requests[2].1["requestedEpoch"], 0); } #[rocket::async_test] @@ -684,7 +711,7 @@ pub(crate) mod tests { r#"{"isAllowed":true,"reason":"","gatewayAppId":"gw"}"#, ]); assert!(webhook(url) - .check_authorization(&boot_info(50)) + .check_authorization(&boot_info(50), None) .await .is_err()); assert_eq!(server.join().unwrap().len(), 1); From c3e9319abc61788b31e8d228cec6812741d9138d Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Fri, 6 Nov 2026 09:53:18 -0800 Subject: [PATCH 014/166] refactor(kms): rename rejectionReason consistently Signed-off-by: Kevin Wang --- docs/deployment.md | 2 +- dstack/kms/README.md | 4 ++-- dstack/kms/auth-eth-bun/README.md | 2 +- dstack/kms/auth-eth-bun/index.test.ts | 4 ++-- dstack/kms/auth-eth-bun/index.ts | 2 +- dstack/kms/auth-eth-bun/openapi.json | 4 ++-- dstack/kms/auth-eth/README.md | 2 +- dstack/kms/auth-eth/contracts/DstackKms.sol | 10 +++++----- dstack/kms/auth-eth/src/ethereum.ts | 4 ++-- dstack/kms/auth-eth/src/main.test.ts | 2 +- dstack/kms/auth-eth/src/server.ts | 4 ++-- dstack/kms/auth-eth/src/types.ts | 2 +- dstack/kms/auth-eth/test/DstackKms.t.sol | 12 ++++++------ dstack/kms/auth-mock/README.md | 2 +- dstack/kms/auth-mock/index.test.ts | 6 +++--- dstack/kms/auth-mock/index.ts | 6 +++--- dstack/kms/auth-mock/openapi.json | 4 ++-- dstack/kms/auth-simple/README.md | 2 +- dstack/kms/auth-simple/index.test.ts | 6 +++--- dstack/kms/auth-simple/index.ts | 6 +++--- dstack/kms/src/main_service.rs | 6 +++--- dstack/kms/src/main_service/upgrade_authority.rs | 6 +++--- 22 files changed, 49 insertions(+), 49 deletions(-) diff --git a/docs/deployment.md b/docs/deployment.md index 56c3fe7f8..946700e4d 100644 --- a/docs/deployment.md +++ b/docs/deployment.md @@ -553,7 +553,7 @@ including when authorization is denied. KMS verifies these against its keyring; there is no cached serving epoch. An optional `requestedEpoch` selects historical keys, and omission selects the returned current epoch. Requests for future epochs are denied. No registered epoch returns a rejection with -`rejectReason: "Root epoch not registered"` and a zero hash. +`rejectionReason: "Root epoch not registered"` and a zero hash. RPC certificates still require a listener restart when their root CA changes; callers retry after the certificate is refreshed. Register epoch 0 before normal diff --git a/dstack/kms/README.md b/dstack/kms/README.md index 8f69cf08e..5a516823a 100644 --- a/dstack/kms/README.md +++ b/dstack/kms/README.md @@ -266,7 +266,7 @@ The response is an `AuthorizationResult`: ```json { "isAllowed": true, - "rejectReason": "", + "rejectionReason": "", "gatewayAppId": "0x...", "currentRootEpoch": 0, "currentRootEpochHash": "0x..." @@ -276,7 +276,7 @@ The response is an `AuthorizationResult`: Both successful and rejected decisions include the current epoch and chain hash from the same contract call or auth-simple configuration snapshot. With no registered epoch, the result is `isAllowed: false`, -`rejectReason: "Root epoch not registered"`, epoch 0 and a zero hash. A zero hash +`rejectionReason: "Root epoch not registered"`, epoch 0 and a zero hash. A zero hash never authorizes key release or requests an epoch switch. The KMS holds no serving-epoch snapshot. It verifies the returned current epoch diff --git a/dstack/kms/auth-eth-bun/README.md b/dstack/kms/auth-eth-bun/README.md index 3e202644b..61345133f 100644 --- a/dstack/kms/auth-eth-bun/README.md +++ b/dstack/kms/auth-eth-bun/README.md @@ -204,7 +204,7 @@ the test suite (`index.test.ts`) validates: ### Unified authorization `POST /checkAuthorization` returns `AuthorizationResult`, including -`rejectReason`, `gatewayAppId`, `currentRootEpoch`, and `currentRootEpochHash`. +`rejectionReason`, `gatewayAppId`, `currentRootEpoch`, and `currentRootEpochHash`. It selects KMS or app policy from the verified `appId`. The optional `requestedEpoch` selects historical keys; omission uses the returned current epoch, and an explicit `0` selects epoch 0. See the diff --git a/dstack/kms/auth-eth-bun/index.test.ts b/dstack/kms/auth-eth-bun/index.test.ts index 3f07d5c17..fbbd64841 100644 --- a/dstack/kms/auth-eth-bun/index.test.ts +++ b/dstack/kms/auth-eth-bun/index.test.ts @@ -847,7 +847,7 @@ describe('POST /checkAuthorization', () => { it('reads the complete result with one contract call at configured finality', async () => { const result = { - isAllowed: false, rejectReason: 'App not registered', gatewayAppId: 'gateway', + isAllowed: false, rejectionReason: 'App not registered', gatewayAppId: 'gateway', currentRootEpoch: 2n, currentRootEpochHash: '0x' + '77'.repeat(32), }; mockReadContract.mockResolvedValueOnce(result); @@ -863,7 +863,7 @@ describe('POST /checkAuthorization', () => { it('encodes an explicit epoch zero as present', async () => { mockReadContract.mockResolvedValueOnce({ - isAllowed: true, rejectReason: '', gatewayAppId: 'gateway', + isAllowed: true, rejectionReason: '', gatewayAppId: 'gateway', currentRootEpoch: 1n, currentRootEpochHash: '0x' + '77'.repeat(32), }); expect((await authorize(0)).status).toBe(200); diff --git a/dstack/kms/auth-eth-bun/index.ts b/dstack/kms/auth-eth-bun/index.ts index 933e5923c..8355ba255 100644 --- a/dstack/kms/auth-eth-bun/index.ts +++ b/dstack/kms/auth-eth-bun/index.ts @@ -105,7 +105,7 @@ const DSTACK_KMS_ABI = [ type: 'tuple', components: [ { name: 'isAllowed', type: 'bool' }, - { name: 'rejectReason', type: 'string' }, + { name: 'rejectionReason', type: 'string' }, { name: 'gatewayAppId', type: 'string' }, { name: 'currentRootEpoch', type: 'uint256' }, { name: 'currentRootEpochHash', type: 'bytes32' }, diff --git a/dstack/kms/auth-eth-bun/openapi.json b/dstack/kms/auth-eth-bun/openapi.json index 2660158d3..5425de88f 100644 --- a/dstack/kms/auth-eth-bun/openapi.json +++ b/dstack/kms/auth-eth-bun/openapi.json @@ -344,7 +344,7 @@ "type": "object", "required": [ "isAllowed", - "rejectReason", + "rejectionReason", "gatewayAppId", "currentRootEpoch", "currentRootEpochHash" @@ -353,7 +353,7 @@ "isAllowed": { "type": "boolean" }, - "rejectReason": { + "rejectionReason": { "type": "string", "description": "Empty when allowed." }, diff --git a/dstack/kms/auth-eth/README.md b/dstack/kms/auth-eth/README.md index 397064ac7..af26e2455 100644 --- a/dstack/kms/auth-eth/README.md +++ b/dstack/kms/auth-eth/README.md @@ -175,7 +175,7 @@ Documentation: https://book.getfoundry.sh/ ### Unified authorization `POST /checkAuthorization` returns `AuthorizationResult`, including -`rejectReason`, `gatewayAppId`, `currentRootEpoch`, and `currentRootEpochHash`. +`rejectionReason`, `gatewayAppId`, `currentRootEpoch`, and `currentRootEpochHash`. It selects KMS or app policy from the verified `appId`. The optional `requestedEpoch` selects historical keys; omission uses the returned current epoch, and an explicit `0` selects epoch 0. See the diff --git a/dstack/kms/auth-eth/contracts/DstackKms.sol b/dstack/kms/auth-eth/contracts/DstackKms.sol index 2b7b2e343..d750f71b4 100644 --- a/dstack/kms/auth-eth/contracts/DstackKms.sol +++ b/dstack/kms/auth-eth/contracts/DstackKms.sol @@ -38,7 +38,7 @@ contract DstackKms is Initializable, Ownable2StepUpgradeable, UUPSUpgradeable, E struct AuthorizationResult { bool isAllowed; - string rejectReason; + string rejectionReason; string gatewayAppId; uint256 currentRootEpoch; bytes32 currentRootEpochHash; @@ -409,13 +409,13 @@ contract DstackKms is Initializable, Ownable2StepUpgradeable, UUPSUpgradeable, E { result.gatewayAppId = gatewayAppId; if (_rootEpochs.length == 0) { - result.rejectReason = "Root epoch not registered"; + result.rejectionReason = "Root epoch not registered"; return result; } result.currentRootEpoch = _currentRootEpoch; result.currentRootEpochHash = _rootEpochs[_currentRootEpoch].chainHash; if (request.requestedEpoch.isSet && request.requestedEpoch.value > _currentRootEpoch) { - result.rejectReason = "Requested root epoch is not active"; + result.rejectionReason = "Requested root epoch is not active"; return result; } AppInfo calldata appInfo = request.appInfo; @@ -431,9 +431,9 @@ contract DstackKms is Initializable, Ownable2StepUpgradeable, UUPSUpgradeable, E advisoryIds: appInfo.advisoryIds }); if (appInfo.appId == address(this)) { - (result.isAllowed, result.rejectReason) = _isKmsAllowed(bootInfo); + (result.isAllowed, result.rejectionReason) = _isKmsAllowed(bootInfo); } else { - (result.isAllowed, result.rejectReason) = _isAppAllowed(bootInfo); + (result.isAllowed, result.rejectionReason) = _isAppAllowed(bootInfo); } } diff --git a/dstack/kms/auth-eth/src/ethereum.ts b/dstack/kms/auth-eth/src/ethereum.ts index 6af9cf360..c8451b824 100644 --- a/dstack/kms/auth-eth/src/ethereum.ts +++ b/dstack/kms/auth-eth/src/ethereum.ts @@ -7,7 +7,7 @@ import { BootInfo, BootResponse, RootEpochInfo, AuthorizationResult, Authorizati // Minimal ABI for DstackKms contract const DSTACK_KMS_ABI = [ - "function checkAuthorization(((address appId,bytes32 composeHash,address instanceId,bytes32 deviceId,bytes32 mrAggregated,bytes32 mrSystem,bytes32 osImageHash,string tcbStatus,string[] advisoryIds) appInfo,(bool isSet,uint256 value) requestedEpoch) request) view returns ((bool isAllowed,string rejectReason,string gatewayAppId,uint256 currentRootEpoch,bytes32 currentRootEpochHash))", + "function checkAuthorization(((address appId,bytes32 composeHash,address instanceId,bytes32 deviceId,bytes32 mrAggregated,bytes32 mrSystem,bytes32 osImageHash,string tcbStatus,string[] advisoryIds) appInfo,(bool isSet,uint256 value) requestedEpoch) request) view returns ((bool isAllowed,string rejectionReason,string gatewayAppId,uint256 currentRootEpoch,bytes32 currentRootEpochHash))", "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)", "function isKmsAllowed((address appId,bytes32 composeHash,address instanceId,bytes32 deviceId,bytes32 mrAggregated,bytes32 mrSystem,bytes32 osImageHash,string tcbStatus,string[] advisoryIds) bootInfo) view returns (bool, string)", "function gatewayAppId() view returns (string)", @@ -84,7 +84,7 @@ export class EthereumBackend { }); return { isAllowed: result.isAllowed, - rejectReason: result.rejectReason, + rejectionReason: result.rejectionReason, gatewayAppId: result.gatewayAppId, currentRootEpoch: ethers.toNumber(result.currentRootEpoch), currentRootEpochHash: result.currentRootEpochHash, diff --git a/dstack/kms/auth-eth/src/main.test.ts b/dstack/kms/auth-eth/src/main.test.ts index 1b15f500f..a9e044e99 100644 --- a/dstack/kms/auth-eth/src/main.test.ts +++ b/dstack/kms/auth-eth/src/main.test.ts @@ -94,7 +94,7 @@ describe('Server', () => { describe('POST /checkAuthorization', () => { it('preserves the complete decision, including the epoch on rejection', async () => { const decision = { - isAllowed: false, rejectReason: 'App not registered', gatewayAppId: 'gateway', + isAllowed: false, rejectionReason: 'App not registered', gatewayAppId: 'gateway', currentRootEpoch: 1, currentRootEpochHash: '0x' + '11'.repeat(32), }; app.ethereum.checkAuthorization = jest.fn().mockResolvedValue(decision); diff --git a/dstack/kms/auth-eth/src/server.ts b/dstack/kms/auth-eth/src/server.ts index c702e8749..17deb15b5 100644 --- a/dstack/kms/auth-eth/src/server.ts +++ b/dstack/kms/auth-eth/src/server.ts @@ -125,10 +125,10 @@ export async function build(): Promise { response: { 200: { type: 'object', - required: ['isAllowed', 'rejectReason', 'gatewayAppId', 'currentRootEpoch', 'currentRootEpochHash'], + required: ['isAllowed', 'rejectionReason', 'gatewayAppId', 'currentRootEpoch', 'currentRootEpochHash'], properties: { isAllowed: { type: 'boolean' }, - rejectReason: { type: 'string' }, + rejectionReason: { type: 'string' }, gatewayAppId: { type: 'string' }, currentRootEpoch: { type: 'integer', minimum: 0 }, currentRootEpochHash: hex(32, 'Current root epoch chain hash'), diff --git a/dstack/kms/auth-eth/src/types.ts b/dstack/kms/auth-eth/src/types.ts index f59adf624..1b43e3fac 100644 --- a/dstack/kms/auth-eth/src/types.ts +++ b/dstack/kms/auth-eth/src/types.ts @@ -35,7 +35,7 @@ export interface RootEpochInfo { export interface AuthorizationResult { isAllowed: boolean; - rejectReason: string; + rejectionReason: string; gatewayAppId: string; currentRootEpoch: number; currentRootEpochHash: string; diff --git a/dstack/kms/auth-eth/test/DstackKms.t.sol b/dstack/kms/auth-eth/test/DstackKms.t.sol index e6f358d64..c2b005250 100644 --- a/dstack/kms/auth-eth/test/DstackKms.t.sol +++ b/dstack/kms/auth-eth/test/DstackKms.t.sol @@ -474,7 +474,7 @@ contract DstackKmsTest is Test { DstackKms.AuthorizationRequest(info, DstackKms.OptionalEpoch(true, 1)); DstackKms.AuthorizationResult memory result = kms.checkAuthorization(request); assertFalse(result.isAllowed); - assertEq(result.rejectReason, "Requested root epoch is not active"); + assertEq(result.rejectionReason, "Requested root epoch is not active"); bytes32 h1 = chainHash(chainHash(bytes32(0), ca(0xa0), k256(0xb0)), ca(0xa1), k256(0xb1)); vm.prank(owner); kms.activateRootEpoch(1, h1); @@ -493,7 +493,7 @@ contract DstackKmsTest is Test { function test_CheckAuthorization_UnregisteredEpochReturnsRejection() public view { DstackKms.AuthorizationResult memory result = authorize(authorizationInfo(address(kms))); assertFalse(result.isAllowed); - assertEq(result.rejectReason, "Root epoch not registered"); + assertEq(result.rejectionReason, "Root epoch not registered"); assertEq(result.currentRootEpoch, 0); assertEq(result.currentRootEpochHash, bytes32(0)); } @@ -511,14 +511,14 @@ contract DstackKmsTest is Test { DstackKms.AuthorizationResult memory result = authorize(info); assertTrue(result.isAllowed); - assertEq(result.rejectReason, ""); + assertEq(result.rejectionReason, ""); assertEq(result.gatewayAppId, "gateway"); assertEq(result.currentRootEpochHash, chainHash(bytes32(0), ca(0xa0), k256(0xb0))); info.appId = user; result = authorize(info); assertFalse(result.isAllowed); - assertEq(result.rejectReason, "App not registered"); + assertEq(result.rejectionReason, "App not registered"); info.appId = appId; info.mrAggregated = bytes32(0); @@ -527,7 +527,7 @@ contract DstackKmsTest is Test { info.appId = address(kms); result = authorize(info); assertFalse(result.isAllowed); - assertEq(result.rejectReason, "Aggregated MR not allowed"); + assertEq(result.rejectionReason, "Aggregated MR not allowed"); } function test_CheckAuthorization_RejectionIncludesActivatedEpoch() public { @@ -539,7 +539,7 @@ contract DstackKmsTest is Test { DstackKms.AuthorizationResult memory result = authorize(authorizationInfo(user)); assertFalse(result.isAllowed); - assertEq(result.rejectReason, "App not registered"); + assertEq(result.rejectionReason, "App not registered"); assertEq(result.currentRootEpoch, 1); assertEq( result.currentRootEpochHash, chainHash(chainHash(bytes32(0), ca(0xa0), k256(0xb0)), ca(0xa1), k256(0xb1)) diff --git a/dstack/kms/auth-mock/README.md b/dstack/kms/auth-mock/README.md index 665695c98..78585d661 100644 --- a/dstack/kms/auth-mock/README.md +++ b/dstack/kms/auth-mock/README.md @@ -172,7 +172,7 @@ the API remains identical, so no client code changes are needed. ### Unified authorization `POST /checkAuthorization` returns `AuthorizationResult`, including -`rejectReason`, `gatewayAppId`, `currentRootEpoch`, and `currentRootEpochHash`. +`rejectionReason`, `gatewayAppId`, `currentRootEpoch`, and `currentRootEpochHash`. It selects KMS or app policy from the verified `appId`. The optional `requestedEpoch` selects historical keys; omission uses the returned current epoch, and an explicit `0` selects epoch 0. See the diff --git a/dstack/kms/auth-mock/index.test.ts b/dstack/kms/auth-mock/index.test.ts index 7169a7e50..3336d4f5f 100644 --- a/dstack/kms/auth-mock/index.test.ts +++ b/dstack/kms/auth-mock/index.test.ts @@ -396,14 +396,14 @@ describe('POST /checkAuthorization', () => { vi.stubEnv('MOCK_POLICY', 'deny-kms'); try { expect(await (await authorize('app')).json()).toMatchObject({ - isAllowed: false, rejectReason: 'Root epoch not registered', + isAllowed: false, rejectionReason: 'Root epoch not registered', currentRootEpoch: 0, currentRootEpochHash: '0x' + '00'.repeat(32), }); vi.stubEnv('MOCK_ROOT_EPOCH', '1'); vi.stubEnv('MOCK_ROOT_EPOCH_HASH', '0x' + '11'.repeat(32)); - expect(await (await authorize('app')).json()).toMatchObject({ isAllowed: true, rejectReason: '', currentRootEpoch: 1 }); + expect(await (await authorize('app')).json()).toMatchObject({ isAllowed: true, rejectionReason: '', currentRootEpoch: 1 }); expect(await (await authorize(process.env.KMS_CONTRACT_ADDR!)).json()).toMatchObject({ - isAllowed: false, rejectReason: 'mock policy: deny-kms', currentRootEpoch: 1, + isAllowed: false, rejectionReason: 'mock policy: deny-kms', currentRootEpoch: 1, currentRootEpochHash: '0x' + '11'.repeat(32), }); } finally { diff --git a/dstack/kms/auth-mock/index.ts b/dstack/kms/auth-mock/index.ts index 408c813fb..dc9f3048f 100644 --- a/dstack/kms/auth-mock/index.ts +++ b/dstack/kms/auth-mock/index.ts @@ -118,14 +118,14 @@ class MockBackend { currentRootEpochHash: current.rootEpochHash ?? '0x' + '00'.repeat(32), }; if (current.rootEpoch === undefined) { - return { isAllowed: false, rejectReason: 'Root epoch not registered', gatewayAppId: this.mockGatewayAppId, ...epoch }; + return { isAllowed: false, rejectionReason: 'Root epoch not registered', gatewayAppId: this.mockGatewayAppId, ...epoch }; } if (appInfo.requestedEpoch !== undefined && appInfo.requestedEpoch > current.rootEpoch) { - return { isAllowed: false, rejectReason: 'Requested root epoch is not active', gatewayAppId: this.mockGatewayAppId, ...epoch }; + return { isAllowed: false, rejectionReason: 'Requested root epoch is not active', gatewayAppId: this.mockGatewayAppId, ...epoch }; } const kmsAppId = process.env.KMS_CONTRACT_ADDR || '0xmockcontract1234567890123456789012345678'; const result = this.evaluateBoot(appInfo, normHash(appInfo.appId) === normHash(kmsAppId)); - return { isAllowed: result.isAllowed, rejectReason: result.isAllowed ? '' : result.reason, + return { isAllowed: result.isAllowed, rejectionReason: result.isAllowed ? '' : result.reason, gatewayAppId: this.mockGatewayAppId, ...epoch }; } diff --git a/dstack/kms/auth-mock/openapi.json b/dstack/kms/auth-mock/openapi.json index d9c75f5e3..03d5c95fa 100644 --- a/dstack/kms/auth-mock/openapi.json +++ b/dstack/kms/auth-mock/openapi.json @@ -353,7 +353,7 @@ "type": "object", "required": [ "isAllowed", - "rejectReason", + "rejectionReason", "gatewayAppId", "currentRootEpoch", "currentRootEpochHash" @@ -362,7 +362,7 @@ "isAllowed": { "type": "boolean" }, - "rejectReason": { + "rejectionReason": { "type": "string", "description": "Empty when allowed." }, diff --git a/dstack/kms/auth-simple/README.md b/dstack/kms/auth-simple/README.md index d0f7f00fc..af0e99419 100644 --- a/dstack/kms/auth-simple/README.md +++ b/dstack/kms/auth-simple/README.md @@ -257,7 +257,7 @@ bun run test:run ### Unified authorization `POST /checkAuthorization` returns `AuthorizationResult`, including -`rejectReason`, `gatewayAppId`, `currentRootEpoch`, and `currentRootEpochHash`. +`rejectionReason`, `gatewayAppId`, `currentRootEpoch`, and `currentRootEpochHash`. It selects KMS or app policy from the verified `appId`. The optional `requestedEpoch` selects historical keys; omission uses the returned current epoch, and an explicit `0` selects epoch 0. See the diff --git a/dstack/kms/auth-simple/index.test.ts b/dstack/kms/auth-simple/index.test.ts index 191137cd3..fc3edea38 100644 --- a/dstack/kms/auth-simple/index.test.ts +++ b/dstack/kms/auth-simple/index.test.ts @@ -482,12 +482,12 @@ describe('auth-simple', () => { const kmsContractAddr = '0x' + '11'.repeat(20); writeTestConfig({ ...config, kmsContractAddr, currentRootEpoch: 1 }); expect(await authorize()).toMatchObject({ - isAllowed: true, rejectReason: '', currentRootEpoch: 1, currentRootEpochHash: h1, + isAllowed: true, rejectionReason: '', currentRootEpoch: 1, currentRootEpochHash: h1, }); expect(await authorize({ requestedEpoch: 0 })).toMatchObject({ isAllowed: true, currentRootEpoch: 1 }); expect(await authorize({ requestedEpoch: 1 })).toMatchObject({ isAllowed: true }); expect(await authorize({ requestedEpoch: 2 })).toMatchObject({ - isAllowed: false, rejectReason: 'Requested root epoch is not active', currentRootEpoch: 1, currentRootEpochHash: h1, + isAllowed: false, rejectionReason: 'Requested root epoch is not active', currentRootEpoch: 1, currentRootEpochHash: h1, }); expect(await authorize({ appId: kmsContractAddr, composeHash: 'not-an-app-compose' })).toMatchObject({ isAllowed: true }); expect(await authorize({ appId: kmsContractAddr, mrAggregated: '0xbad' })).toMatchObject({ @@ -499,7 +499,7 @@ describe('auth-simple', () => { it('denies authorization without registered roots, while bootstrap remains available', async () => { writeTestConfig({ ...config, rootEpochs: [] }); expect(await authorize()).toMatchObject({ - isAllowed: false, rejectReason: 'Root epoch not registered', + isAllowed: false, rejectionReason: 'Root epoch not registered', currentRootEpoch: 0, currentRootEpochHash: '0x' + '00'.repeat(32), }); expect((await boot('kms', {})).isAllowed).toBe(true); diff --git a/dstack/kms/auth-simple/index.ts b/dstack/kms/auth-simple/index.ts index c639d5c23..dbca135ed 100644 --- a/dstack/kms/auth-simple/index.ts +++ b/dstack/kms/auth-simple/index.ts @@ -165,14 +165,14 @@ class ConfigBackend { currentRootEpochHash: current.rootEpochHash ?? '0x' + '00'.repeat(32), }; if (current.rootEpoch === undefined) { - return { isAllowed: false, rejectReason: 'Root epoch not registered', gatewayAppId: config.gatewayAppId, ...epoch }; + return { isAllowed: false, rejectionReason: 'Root epoch not registered', gatewayAppId: config.gatewayAppId, ...epoch }; } if (appInfo.requestedEpoch !== undefined && appInfo.requestedEpoch > current.rootEpoch) { - return { isAllowed: false, rejectReason: 'Requested root epoch is not active', gatewayAppId: config.gatewayAppId, ...epoch }; + return { isAllowed: false, rejectionReason: 'Requested root epoch is not active', gatewayAppId: config.gatewayAppId, ...epoch }; } const isKms = normalizeHex(appInfo.appId) === normalizeHex(config.kmsContractAddr); const result = this.evaluateBoot(appInfo, isKms, config); - return { isAllowed: result.isAllowed, rejectReason: result.reason, gatewayAppId: result.gatewayAppId, ...epoch }; + return { isAllowed: result.isAllowed, rejectionReason: result.reason, gatewayAppId: result.gatewayAppId, ...epoch }; } async checkBoot(bootInfo: BootInfo, isKms: boolean): Promise { diff --git a/dstack/kms/src/main_service.rs b/dstack/kms/src/main_service.rs index 574f6312d..4c5e72545 100644 --- a/dstack/kms/src/main_service.rs +++ b/dstack/kms/src/main_service.rs @@ -340,7 +340,7 @@ impl KmsState { .servable_epoch(RootEpochRegistry::Current(current))? }; if !result.is_allowed { - bail!("Authorization denied: {}", result.reject_reason); + bail!("Authorization denied: {}", result.rejection_reason); } let epoch = requested.unwrap_or(current); if epoch > current { @@ -1589,7 +1589,7 @@ mod tests { KmsState::new(initial.config.clone(), keyring, initial.rpc_ca_cert.clone()).unwrap(); let result = |epoch: u32| AuthorizationResult { is_allowed: true, - reject_reason: String::new(), + rejection_reason: String::new(), gateway_app_id: String::new(), current_root_epoch: epoch, current_root_epoch_hash: hex::encode(state.keyring.chain_hashes[epoch as usize]), @@ -1611,7 +1611,7 @@ mod tests { let mut rejected = result(1); rejected.is_allowed = false; - rejected.reject_reason = "Policy denied".into(); + rejected.rejection_reason = "Policy denied".into(); assert!(state.authorized_epoch(&rejected, Some(0)).is_err()); rejected.is_allowed = true; for hash in [ diff --git a/dstack/kms/src/main_service/upgrade_authority.rs b/dstack/kms/src/main_service/upgrade_authority.rs index a49d84066..d44e91a0d 100644 --- a/dstack/kms/src/main_service/upgrade_authority.rs +++ b/dstack/kms/src/main_service/upgrade_authority.rs @@ -90,7 +90,7 @@ pub(crate) struct BootResponse { #[serde(rename_all = "camelCase")] pub(crate) struct AuthorizationResult { pub is_allowed: bool, - pub reject_reason: String, + pub rejection_reason: String, pub gateway_app_id: String, pub current_root_epoch: u32, pub current_root_epoch_hash: String, @@ -271,7 +271,7 @@ impl AuthApi { match self { AuthApi::Dev { dev } => Ok(AuthorizationResult { is_allowed: true, - reject_reason: String::new(), + rejection_reason: String::new(), gateway_app_id: dev.gateway_app_id.clone(), current_root_epoch: 0, current_root_epoch_hash: hex::encode([0; 32]), @@ -677,7 +677,7 @@ pub(crate) mod tests { async fn authorization_returns_the_epoch_with_each_decision() { let response = |allowed: bool, epoch: u32| { format!( - r#"{{"isAllowed":{allowed},"rejectReason":"","gatewayAppId":"gw", "currentRootEpoch":{epoch},"currentRootEpochHash":"0x{}"}}"#, + r#"{{"isAllowed":{allowed},"rejectionReason":"","gatewayAppId":"gw", "currentRootEpoch":{epoch},"currentRootEpochHash":"0x{}"}}"#, "11".repeat(32) ) }; From 2f3d7858b94079d2a55aa71ea296c53605af342a Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Fri, 6 Nov 2026 19:02:27 -0800 Subject: [PATCH 015/166] test(kms): retain hash-bound activation after rebase Signed-off-by: Kevin Wang --- dstack/kms/auth-eth/test/DstackKms.t.sol | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/dstack/kms/auth-eth/test/DstackKms.t.sol b/dstack/kms/auth-eth/test/DstackKms.t.sol index c2b005250..a6635a5ca 100644 --- a/dstack/kms/auth-eth/test/DstackKms.t.sol +++ b/dstack/kms/auth-eth/test/DstackKms.t.sol @@ -534,7 +534,7 @@ contract DstackKmsTest is Test { vm.startPrank(owner); kms.registerRootEpoch(ca(0xa0), k256(0xb0), "https://example.com/epoch-0"); kms.registerRootEpoch(ca(0xa1), k256(0xb1), "https://example.com/epoch-1"); - kms.activateRootEpoch(1); + kms.activateRootEpoch(1, chainHash(chainHash(bytes32(0), ca(0xa0), k256(0xb0)), ca(0xa1), k256(0xb1))); vm.stopPrank(); DstackKms.AuthorizationResult memory result = authorize(authorizationInfo(user)); From f697616cfd4e75a5b17c6e2f7c33132b3adc21f2 Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Fri, 6 Nov 2026 19:20:22 -0800 Subject: [PATCH 016/166] refactor(kms): remove temporary CA bootstrap and fetch keys directly Signed-off-by: Kevin Wang --- CHANGELOG.md | 7 + docs/deployment.md | 4 +- docs/security/security-best-practices.md | 2 +- docs/security/security-model.md | 2 +- docs/tutorials/kms-cvm-deployment.md | 2 - .../troubleshooting-kms-deployment.md | 2 +- dstack/Cargo.lock | 1 + dstack/cert-client/Cargo.toml | 3 +- dstack/cert-client/src/lib.rs | 277 +++++++++++++++++- dstack/dstack-util/src/main.rs | 65 +--- dstack/dstack-util/src/system_setup.rs | 50 +--- dstack/kms/README.md | 28 +- dstack/kms/rpc/proto/kms_rpc.proto | 28 +- dstack/kms/src/config.rs | 14 +- dstack/kms/src/main_service.rs | 92 +----- dstack/kms/src/onboard_service.rs | 92 +----- dstack/ra-rpc/src/ratls_client_verifier.rs | 2 +- dstack/tests/docs/kms-bootstrap-onboard.md | 39 +-- dstack/tests/docs/kms-self-authorization.md | 18 +- 19 files changed, 405 insertions(+), 323 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 8319bdc58..a222c8c15 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,6 +7,13 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] +### Removed + +- KMS: remove `GetTempCaCert` and temporary CA files and handover material. + Guests fetch keys directly with self-issued RA-TLS certificates and validate + the response connection's attestation and root CA. Older clients and sources + that require the temporary CA protocol must be upgraded. + ### Changed - dstack-util, cert-client: guests mint self-issued RA-TLS client diff --git a/docs/deployment.md b/docs/deployment.md index 946700e4d..27adbe9e8 100644 --- a/docs/deployment.md +++ b/docs/deployment.md @@ -470,8 +470,8 @@ Additional KMS instances can onboard from an existing KMS to share the same root **How it works:** 1. New KMS starts in onboard mode (empty `auto_bootstrap_domain`) -2. New KMS generates an RA-TLS client certificate with its TDX quote: self-issued when given a `source_token`, or minted from the temp CA fetched with `GetTempCaCert` for a source older than 0.6.0 -3. New KMS calls `GetKmsKey` over mTLS: `Admin.GetKmsKey` with the token, or the legacy public `KMS.GetKmsKey` without one +2. New KMS generates a self-issued RA-TLS client certificate with its TDX quote locally; no CA preflight is needed +3. New KMS calls `GetKmsKey` over mTLS: `Admin.GetKmsKey` with the token, or the explicitly enabled public `KMS.GetKmsKey` without one 4. Source KMS verifies attestation via `bootAuth/kms` webhook 5. If approved, source KMS returns root keys 6. New KMS verifies the source's attestation on that same connection and checks it against its own policy before using the keys diff --git a/docs/security/security-best-practices.md b/docs/security/security-best-practices.md index 5948b2c28..31cba2658 100644 --- a/docs/security/security-best-practices.md +++ b/docs/security/security-best-practices.md @@ -111,7 +111,7 @@ Development settings are intentionally easy to audit, but they are not productio - The KMS contract pins a concrete gateway app id. Do not use `gateway_app_id = "any"` for production traffic. - TEE quotes are evaluated by deployment policy, including TCB status and expected OS/application measurements. -The KMS TLS listener verifies client certificates by the attestation they carry rather than by an issuer CA, so it needs no `rpc.tls.mutual` section. It still accepts connections without a client certificate, because bootstrap and public metadata endpoints must be reachable before a client has an RA-TLS certificate. Current guests present self-issued certificates, but `GetTempCaCert` remains in use by older guest images and by KMS-to-KMS onboarding without a `source_token`, which mint their client certificates from that CA; it returns temp CA private material, so treat it as bootstrap-sensitive. +The KMS TLS listener verifies client certificates by the attestation they carry rather than by an issuer CA, so it needs no `rpc.tls.mutual` section. It still accepts connections without a client certificate, because bootstrap and public metadata endpoints must be reachable before a client has an RA-TLS certificate. Guests and onboarding clients use self-issued certificates. The temporary-CA RPC and key material have been removed; older clients that require them must be upgraded. App key release and KMS key handover still require verified caller attestation from the RA-TLS client certificate. Certificate signing verifies the CSR signature and embedded attestation before signing. diff --git a/docs/security/security-model.md b/docs/security/security-model.md index b3d43ed3b..d89e703ef 100644 --- a/docs/security/security-model.md +++ b/docs/security/security-model.md @@ -425,7 +425,7 @@ The KMS Rocket TLS listener permits connections without a client certificate bec App key release and KMS key handover require verified caller attestation from the RA-TLS client certificate. Certificate signing verifies the CSR signature and the attestation embedded in the CSR before signing. -The unauthenticated or non-client-certificate surface includes bootstrap and temp-CA bootstrap material retrieval, env-encryption public-key retrieval, metadata, health, and metrics behavior documented for operators. `GetTempCaCert` returns temp CA private material and remains in use by older guest images and by KMS-to-KMS onboarding without a `source_token`, which mint their client certificates from that CA; operators must treat it as bootstrap-sensitive rather than harmless public metadata. +The unauthenticated or non-client-certificate surface includes bootstrap, env-encryption public-key retrieval, metadata, health, and metrics behavior documented for operators. Clients create self-issued certificates locally. The temporary-CA RPC and its key material have been removed. Guest key retrieval verifies the actual response connection's attestation and checks that the returned CA signed its RPC certificate. Provider pins and launch measurement continue to bind the selected KMS root identity; attestation alone does not authorize an arbitrary CVM as the expected KMS. KMS-to-KMS root-key handover is served by `Admin.GetKmsKey` on the authenticated admin listener, which additionally requires a verified RA-TLS client certificate. diff --git a/docs/tutorials/kms-cvm-deployment.md b/docs/tutorials/kms-cvm-deployment.md index 24c68287e..2b6de47a0 100644 --- a/docs/tutorials/kms-cvm-deployment.md +++ b/docs/tutorials/kms-cvm-deployment.md @@ -73,8 +73,6 @@ Inside the CVM at `/etc/kms/certs/`: | `root-ca.key` | Root CA signing key (P256 ECDSA) | | `rpc.crt` | TLS certificate for RPC server | | `rpc.key` | RPC server private key | -| `tmp-ca.crt` | Temp CA, served by the legacy `GetTempCaCert` bootstrap RPC | -| `tmp-ca.key` | Temp CA private key, served alongside it | | `root-k256.key` | Ethereum signing key (secp256k1) | | `bootstrap-info.json` | Public keys and TDX attestation quote | diff --git a/docs/tutorials/troubleshooting-kms-deployment.md b/docs/tutorials/troubleshooting-kms-deployment.md index 7282db94d..ac38cbb4d 100644 --- a/docs/tutorials/troubleshooting-kms-deployment.md +++ b/docs/tutorials/troubleshooting-kms-deployment.md @@ -303,7 +303,7 @@ sudo systemctl restart qgsd ### GetMeta Returns "Connection refused" on Port 9200 -**Symptom:** KMS responds to `GetTempCaCert` but GetMeta returns: +**Symptom:** The KMS health endpoint responds but GetMeta returns: ```json {"error": "error sending request for url (http://127.0.0.1:9200/): ...Connection refused (os error 111)"} ``` diff --git a/dstack/Cargo.lock b/dstack/Cargo.lock index 550291383..a67cc5543 100644 --- a/dstack/Cargo.lock +++ b/dstack/Cargo.lock @@ -832,6 +832,7 @@ dependencies = [ "dstack-types", "hex", "http-client", + "mock-attestation", "ra-rpc", "ra-tls", "rocket", diff --git a/dstack/cert-client/Cargo.toml b/dstack/cert-client/Cargo.toml index 092824e28..2f1a56e0e 100644 --- a/dstack/cert-client/Cargo.toml +++ b/dstack/cert-client/Cargo.toml @@ -20,10 +20,11 @@ tdx-attest.workspace = true dstack-guest-agent-rpc.workspace = true http-client = { workspace = true, features = ["prpc"] } tokio.workspace = true +x509-parser = { workspace = true, features = ["verify"] } [dev-dependencies] +mock-attestation = { path = "../crates/mock-attestation" } hex.workspace = true rocket.workspace = true tempfile.workspace = true -x509-parser.workspace = true tokio = { workspace = true, features = ["macros", "rt", "rt-multi-thread", "sync"] } diff --git a/dstack/cert-client/src/lib.rs b/dstack/cert-client/src/lib.rs index b7fdda86a..5e45ce24c 100644 --- a/dstack/cert-client/src/lib.rs +++ b/dstack/cert-client/src/lib.rs @@ -2,15 +2,15 @@ // // SPDX-License-Identifier: Apache-2.0 -use std::sync::Arc; +use std::sync::{Arc, Mutex}; use anyhow::{bail, Context, Result}; -use dstack_kms_rpc::{kms_client::KmsClient, SignCertRequest}; +use dstack_kms_rpc::{kms_client::KmsClient, AppKeyResponse, GetAppKeyRequest, SignCertRequest}; use dstack_types::{AppKeys, KeyProvider}; use ra_rpc::client::{CertInfo, RaClient, RaClientConfig}; use ra_tls::{ attestation::{AppInfo, AttestationVerifier, VersionedAttestation}, - cert::{generate_self_signed_ra_cert, CaCert, CertSigningRequestV2}, + cert::{generate_self_signed_ra_cert, CaCert, CertPair, CertSigningRequestV2}, }; pub enum CertRequestClient { @@ -49,6 +49,82 @@ pub fn validate_kms_rpc_cert(cert: Option) -> Result<()> { Ok(()) } +/// Fetch keys directly over the attested connection, without a metadata/CA preflight. +/// Historical requests must pin the already trusted current KMS CA, since their +/// response CA belongs to the requested epoch rather than the TLS listener. +/// Without a supplied CA, the caller must still bind the returned root identity +/// through its provider pin or launch-measurement policy. +pub async fn request_app_keys( + url: String, + request: GetAppKeyRequest, + client_cert: CertPair, + attestation_verifier: Arc, + trusted_ca: Option, +) -> Result { + if request.root_epoch.is_some() && trusted_ca.is_none() { + bail!("historical key requests require the trusted KMS CA"); + } + let server_cert = Arc::new(Mutex::new(None)); + let captured = server_cert.clone(); + let client = RaClientConfig::builder() + .remote_uri(url) + .tls_client_cert(client_cert.cert_pem) + .tls_client_key(client_cert.key_pem) + .tls_no_check(trusted_ca.is_none()) + .tls_built_in_root_certs(false) + .maybe_tls_ca_cert(trusted_ca.clone()) + .attestation_verifier(attestation_verifier) + .cert_validator(Box::new(move |info| { + let info = info.context("missing KMS server certificate")?; + if info.attestation.is_none() { + bail!("KMS server certificate has no verified attestation"); + } + let cert_der = info.cert_der.clone(); + validate_kms_rpc_cert(Some(info))?; + *captured + .lock() + .map_err(|_| anyhow::anyhow!("KMS certificate lock poisoned"))? = Some(cert_der); + Ok(()) + })) + .build() + .into_client()?; + let requested_epoch = request.root_epoch; + let response = KmsClient::new(client).get_app_key(request).await?; + if requested_epoch.is_some_and(|epoch| epoch != response.root_epoch) { + bail!("KMS returned a different root epoch than requested"); + } + let cert = server_cert + .lock() + .map_err(|_| anyhow::anyhow!("KMS certificate lock poisoned"))?; + validate_kms_root( + cert.as_deref().context("missing KMS server certificate")?, + if requested_epoch.is_some() { + trusted_ca + .as_deref() + .context("historical key requests require the trusted KMS CA")? + } else { + &response.ca_cert + }, + )?; + Ok(response) +} + +fn validate_kms_root(server_cert: &[u8], root_ca: &str) -> Result<()> { + let (_, pem) = + x509_parser::pem::parse_x509_pem(root_ca.as_bytes()).context("invalid KMS root CA")?; + let root = pem.parse_x509().context("invalid KMS root CA")?; + let (_, leaf) = + x509_parser::parse_x509_certificate(server_cert).context("invalid KMS RPC certificate")?; + if !root.is_ca() || !root.validity().is_valid() || !leaf.validity().is_valid() { + bail!("invalid or expired KMS certificate"); + } + if leaf.issuer() != root.subject() { + bail!("KMS RPC certificate issuer does not match the root CA"); + } + leaf.verify_signature(Some(root.public_key())) + .context("KMS root CA did not sign the attested RPC certificate") +} + fn kms_client( url: String, ca_cert: String, @@ -374,6 +450,201 @@ mod tests { ) } + #[rocket::post("/GetAppKey", data = "<_request>")] + fn direct_app_keys( + _request: String, + state: &rocket::State<(Arc>, std::sync::atomic::AtomicUsize)>, + ) -> rocket::response::content::RawJson { + state.1.fetch_add(1, std::sync::atomic::Ordering::SeqCst); + rocket::response::content::RawJson(state.0.lock().unwrap().clone()) + } + + #[tokio::test(flavor = "multi_thread")] + async fn direct_key_fetch_verifies_the_response_connection_without_preflight() { + use ra_tls::attestation::{ + AttestationVerifierConfig, DstackNitroQuote, QuoteContentType, RootCaPaths, + }; + let dir = tempfile::tempdir().unwrap(); + let attester = mock_attestation::nsm::NsmGenerator::new().unwrap(); + let attestation_root = dir.path().join("attestation-root.pem"); + std::fs::write(&attestation_root, attester.root_ca_pem()).unwrap(); + let verifier = Arc::new( + AttestationVerifier::load(&AttestationVerifierConfig { + insecure_allow_external_trust_anchors: true, + root_ca: RootCaPaths { + aws_nitro_enclave: Some(attestation_root), + ..Default::default() + }, + ..Default::default() + }) + .unwrap(), + ); + let root = kms_root(); + let rpc_key = p256_key(); + let report_data = QuoteContentType::RaTlsCert.to_report_data(&rpc_key.public_key_der()); + let pcrs = [0, 1, 2, 4] + .into_iter() + .map(|i| (i, vec![i as u8; 48])) + .collect(); + let evidence = attester.attest_with_pcrs(&report_data, pcrs).unwrap(); + let attestation = Attestation { + quote: AttestationQuote::DstackNitroEnclave(DstackNitroQuote { + nsm_quote: evidence, + }), + runtime_events: Vec::new(), + report_data, + config: String::new(), + report: (), + } + .into_versioned(); + let rpc_cert = root + .sign( + CertRequest::builder() + .subject("localhost") + .alt_names(&["127.0.0.1".into()]) + .key(&rpc_key) + .special_usage("kms:rpc") + .usage_server_auth(true) + .attestation(&attestation) + .build(), + ) + .unwrap(); + let cert_path = dir.path().join("rpc.crt"); + let key_path = dir.path().join("rpc.key"); + std::fs::write(&cert_path, rpc_cert.pem()).unwrap(); + std::fs::write(&key_path, rpc_key.serialize_pem()).unwrap(); + let response = AppKeyResponse { + ca_cert: root.pem_cert.clone(), + root_epoch: 1, + ..Default::default() + }; + let figment = rocket::Config::figment() + .merge(("port", 0)) + .merge(("address", "127.0.0.1")) + .merge(("log_level", "off")) + .merge(("shutdown.ctrlc", false)) + .merge(("tls.certs", cert_path)) + .merge(("tls.key", key_path)); + 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 server = rocket::custom(figment) + .manage(( + response_body.clone(), + std::sync::atomic::AtomicUsize::new(0), + )) + .mount("/", rocket::routes![direct_app_keys]) + .attach(rocket::fairing::AdHoc::on_liftoff("port", move |rocket| { + if let Some(tx) = tx.lock().unwrap().take() { + let _ = tx.send(rocket.endpoints().find_map(|e| e.port()).unwrap()); + } + Box::pin(async {}) + })) + .ignite() + .await + .unwrap(); + let shutdown = server.shutdown(); + let task = tokio::spawn(server.launch()); + let url = format!("https://127.0.0.1:{}", rx.await.unwrap()); + let client_cert = || { + let key = p256_key(); + let cert = CertRequest::builder() + .subject("guest") + .key(&key) + .usage_client_auth(true) + .build() + .self_signed() + .unwrap(); + CertPair { + cert_pem: cert.pem(), + key_pem: key.serialize_pem(), + } + }; + let request = || GetAppKeyRequest { + api_version: 1, + vm_config: String::new(), + root_epoch: None, + }; + let result = request_app_keys( + url.clone(), + request(), + client_cert(), + verifier.clone(), + None, + ) + .await + .unwrap(); + assert_eq!(result.ca_cert, response.ca_cert); + assert_eq!(result.root_epoch, 1); + + let wrong_root = AppKeyResponse { + ca_cert: kms_root().pem_cert, + ..response.clone() + }; + *response_body.lock().unwrap() = serde_json::to_string(&wrong_root).unwrap(); + assert!(request_app_keys( + url.clone(), + request(), + client_cert(), + verifier.clone(), + None + ) + .await + .is_err()); + *response_body.lock().unwrap() = serde_json::to_string(&response).unwrap(); + + let mut historical_request = request(); + historical_request.root_epoch = Some(1); + assert!(request_app_keys( + url.clone(), + historical_request, + client_cert(), + verifier.clone(), + Some(root.pem_cert.clone()) + ) + .await + .is_ok()); + + let production = Arc::new(AttestationVerifier::new_prod(None).unwrap()); + assert!( + request_app_keys(url.clone(), request(), client_cert(), production, None) + .await + .is_err() + ); + let mut historical = request(); + historical.root_epoch = Some(0); + let err = request_app_keys(url, historical, client_cert(), verifier, None) + .await + .unwrap_err(); + assert!(err.to_string().contains("require the trusted KMS CA")); + 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), + 4 + ); + } + + #[test] + fn the_returned_root_must_sign_the_actual_rpc_certificate() { + let root = kms_root(); + let (leaf, _) = kms_rpc_cert(&root, "kms.example"); + let (_, leaf) = x509_parser::pem::parse_x509_pem(leaf.as_bytes()).unwrap(); + validate_kms_root(&leaf.contents, &root.pem_cert).unwrap(); + let unrelated = kms_root(); + assert!(validate_kms_root(&leaf.contents, &unrelated.pem_cert).is_err()); + assert!(validate_kms_root(&leaf.contents, "not a certificate").is_err()); + assert!(validate_kms_root(b"not a certificate", &root.pem_cert).is_err()); + + let (app_leaf, _) = app_cert(&root, &[0xa1; 20], "kms.example"); + let (_, app_leaf) = x509_parser::pem::parse_x509_pem(app_leaf.as_bytes()).unwrap(); + assert!(validate_kms_root(&app_leaf.contents, &root.pem_cert).is_err()); + } + #[rocket::post("/GetMeta")] fn get_meta() -> &'static str { r#"{"ca_cert":"served ca cert"}"# diff --git a/dstack/dstack-util/src/main.rs b/dstack/dstack-util/src/main.rs index f34cb0690..e135bdf64 100644 --- a/dstack/dstack-util/src/main.rs +++ b/dstack/dstack-util/src/main.rs @@ -639,9 +639,6 @@ fn cmd_attest_strip(args: AttestStripArgs) -> Result<()> { } async fn cmd_get_keys(args: GetKeysArgs) -> Result<()> { - use dstack_kms_rpc::kms_client::KmsClient; - use ra_rpc::client::RaClientConfig; - let kms_url = normalize_prpc_url(&args.kms_url); // Load root CA if provided for TLS pinning @@ -653,58 +650,24 @@ async fn cmd_get_keys(args: GetKeysArgs) -> Result<()> { None }; - // Step 1: Fetch the KMS root CA - eprintln!("Connecting to KMS: {kms_url}"); - let tls_skip_ca_verification = root_ca_pem.is_none(); - if tls_skip_ca_verification { - eprintln!("Warning: no --root-ca provided, TLS certificate verification is disabled for initial connection"); - } - let root_ca = { - let client = RaClientConfig::builder() - .remote_uri(kms_url.clone()) - .tls_skip_ca_verification(tls_skip_ca_verification) - .tls_built_in_root_certs(false) - .maybe_tls_ca_cert(root_ca_pem.clone()) - .build() - .into_client() - .context("failed to create client")?; - KmsClient::new(client) - .get_meta() - .await - .context("failed to get KMS meta")? - .ca_cert - }; - - // Step 2: Generate a self-issued RA-TLS client certificate. The quote inside it - // is the identity the KMS authenticates, so no CA material is needed. let app_id = decode_app_id(args.app_id.as_deref())?; let cert_pair = generate_self_signed_ra_cert(app_id).context("Failed to generate RA cert")?; - - // Step 3: Create authenticated client and request app keys - let ra_client = RaClientConfig::builder() - .tls_skip_ca_verification(false) - .tls_built_in_root_certs(false) - .remote_uri(kms_url.clone()) - .tls_client_cert(cert_pair.cert_pem) - .tls_client_key(cert_pair.key_pem) - .tls_ca_cert(root_ca.clone()) - .build() - .into_client() - .context("Failed to create RA client")?; - - let kms_client = KmsClient::new(ra_client); - let response = kms_client - .get_app_key(dstack_kms_rpc::GetAppKeyRequest { + let response = cert_client::request_app_keys( + kms_url.clone(), + dstack_kms_rpc::GetAppKeyRequest { api_version: 1, - vm_config: "".to_string(), + vm_config: String::new(), root_epoch: None, - }) - .await - .context( - "failed to get app key (the KMS must be 0.6.0 or later to accept a self-issued RA-TLS client certificate)", - )?; + }, + cert_pair, + std::sync::Arc::new(ra_tls::attestation::AttestationVerifier::new_prod(None)?), + root_ca_pem, + ) + .await + .context("Failed to get app keys from the attested KMS")?; + let root_ca = response.ca_cert; - // Step 4: Build AppKeys structure + // Build AppKeys from the authenticated response. let (_, ca_pem) = x509_parser::pem::parse_x509_pem(root_ca.as_bytes()).context("Failed to parse CA cert")?; let x509 = ca_pem.parse_x509().context("Failed to parse CA cert")?; @@ -723,7 +686,7 @@ async fn cmd_get_keys(args: GetKeysArgs) -> Result<()> { }, }; - // Step 5: Output result + // Output result. let json = serde_json::to_string_pretty(&keys).context("Failed to serialize app keys")?; if let Some(output_path) = args.output { safe_write_with_mode(&output_path, &json, 0o600).context("Failed to write app keys")?; diff --git a/dstack/dstack-util/src/system_setup.rs b/dstack/dstack-util/src/system_setup.rs index 03056aa4f..1beb2b6ad 100644 --- a/dstack/dstack-util/src/system_setup.rs +++ b/dstack/dstack-util/src/system_setup.rs @@ -55,7 +55,7 @@ use crate::{ deserialize_json_file, sha256, sha256_file, AppCompose, AppKeys, KeyProviderKind, SysConfig, }, }; -use cert_client::{validate_kms_rpc_cert, CertRequestClient}; +use cert_client::{request_app_keys, CertRequestClient}; use cmd_lib::run_fun as cmd; use dstack_api_auth::jwt; use dstack_gateway_rpc::{ @@ -2417,48 +2417,22 @@ impl<'a> Stage0<'a> { async fn request_app_keys_from_kms_url(&self, kms_url: String) -> Result<(AppKeys, Vec)> { info!("Requesting app keys from KMS: {kms_url}"); - // Learn the KMS root CA over an unauthenticated connection, as with - // `GetTempCaCert` before. Relying parties bind the KMS through the - // `key-provider` launch event; `verify_key_provider_id` also enforces it - // in-guest when the app pins `key_provider_id`. - let root_ca = { - info!("Getting KMS root CA"); - let client = RaClient::new(kms_url.clone(), true)?; - let kms_client = dstack_kms_rpc::kms_client::KmsClient::new(client); - kms_client - .get_meta() - .await - .context("failed to get KMS meta")? - .ca_cert - }; - // Self-issued: the quote inside the certificate is the identity, so there is - // no CA to fetch and no shared private key to hold. let cert_pair = generate_self_signed_ra_cert(None).context("Failed to generate RA-TLS cert")?; - let attestation_verifier = attestation_verifier(&self.shared.sys_config)?; - let ra_client = RaClientConfig::builder() - .tls_skip_ca_verification(false) - .tls_built_in_root_certs(false) - .remote_uri(kms_url.clone()) - .tls_client_cert(cert_pair.cert_pem) - .tls_client_key(cert_pair.key_pem) - .tls_ca_cert(root_ca.clone()) - .attestation_verifier(attestation_verifier) - .cert_validator(Box::new(validate_kms_rpc_cert)) - .build() - .into_client() - .context("Failed to create client")?; - let kms_client = dstack_kms_rpc::kms_client::KmsClient::new(ra_client); - let response = kms_client - .get_app_key(rpc::GetAppKeyRequest { + let response = request_app_keys( + kms_url.clone(), + rpc::GetAppKeyRequest { api_version: 1, vm_config: self.shared.sys_config.vm_config.clone(), root_epoch: None, - }) - .await - .context( - "failed to get app key (the KMS must be 0.6.0 or later to accept a self-issued RA-TLS client certificate)", - )?; + }, + cert_pair, + attestation_verifier(&self.shared.sys_config)?, + None, + ) + .await + .context("Failed to get app keys from the attested KMS")?; + let root_ca = response.ca_cert; let (_, ca_pem) = x509_parser::pem::parse_x509_pem(root_ca.as_bytes()) .context("Failed to parse ca cert")?; diff --git a/dstack/kms/README.md b/dstack/kms/README.md index 5a516823a..be6cb0d92 100644 --- a/dstack/kms/README.md +++ b/dstack/kms/README.md @@ -287,8 +287,8 @@ fail closed. Self-authorization retains its one-second cache of complete results The TLS listener has a separate lifetime: when an authorized current CA differs from the running RPC certificate's issuer, the KMS requests a restart to refresh that certificate and the caller must retry. The RPC CA does not select key -epochs. `GetTempCaCert` reports the actual TLS issuer; application key responses -and signed certificate chains report the selected epoch's CA. CA certificates +epochs. Application key responses and signed certificate chains report the +selected epoch's CA. CA certificates are reloaded after renewal at startup. Upgrade the contract and auth backend and register epoch 0 before using this @@ -300,3 +300,27 @@ root keys can be generated and attested before registering epoch 0. For `auth-simple`, set `kmsContractAddr` to the attested KMS app ID and configure `rootEpochs`; for `auth-mock`, set `KMS_CONTRACT_ADDR`, `MOCK_ROOT_EPOCH`, and `MOCK_ROOT_EPOCH_HASH`. + + +## Direct attested key retrieval + +`GetTempCaCert` and temporary CA generation, persistence, renewal, and handover +are removed. Its former handover field number and name are reserved in protobuf. +Existing temporary CA files are ignored; they are not needed to start the KMS. +Clients that require the removed RPC, and onboarding sources that require +CA-issued client certificates, are no longer supported. + +Guests and `dstack-util get-keys` generate their self-issued RA-TLS certificate +locally, then call `GetAppKey` directly. They verify the response connection's +attestation and `kms:rpc` usage, and verify that the returned root CA signed that +connection's certificate. There is no `GetMeta` preflight or first unauthenticated +connection. The optional provider pin and the key-provider launch measurement +still bind the root identity; a valid quote alone is not a KMS policy decision. +`--root-ca` additionally pins the CLI connection before sending the request. + +Onboarding uses the same self-issued client flow with or without an admin token, +and checks the source attestation captured from the key-carrying connection +against local KMS policy before importing anything. Retries and reconnects repeat +verification on the response connection; no assumption relies on one TCP session +lasting forever. Historical guest key requests use the already trusted current +KMS CA for TLS, not the historical response CA. diff --git a/dstack/kms/rpc/proto/kms_rpc.proto b/dstack/kms/rpc/proto/kms_rpc.proto index 3bc4af8b4..25e21975e 100644 --- a/dstack/kms/rpc/proto/kms_rpc.proto +++ b/dstack/kms/rpc/proto/kms_rpc.proto @@ -78,22 +78,12 @@ message KmsKeys { } message KmsKeyResponse { - string temp_ca_key = 1; + reserved 1; + reserved "temp_ca_key"; // Root keys of every epoch the KMS holds; the index is the epoch. repeated KmsKeys keys = 2; } -// See the note on `GetTempCaCert`. -// -// `temp_ca_key` is a CA private key handed to any caller, so it authenticates -// nobody. It exists because the KMS pinned this CA for mutual TLS, which left a -// self-issued RA-TLS certificate with nothing to chain to. -message GetTempCaCertResponse { - string temp_ca_cert = 1; - string temp_ca_key = 2; - string ca_cert = 3; -} - message SignCertRequest { uint32 api_version = 1; bytes csr = 2; @@ -115,20 +105,6 @@ service KMS { rpc GetAppEnvEncryptPubKey(AppId) returns (PublicKeyResponse); // Request the KMS instance metadata for use as a probe and health check. rpc GetMeta(google.protobuf.Empty) returns (GetMetaResponse); - // Request the temporary CA certificate and key. - // - // Both current callers - guests at boot, and KMS-to-KMS onboarding - fetch this CA - // and mint their client certificate from it, because the KMS used to pin the CA for - // mutual TLS and a self-issued certificate had nothing to chain to. - // - // That pin is gone: the KMS now verifies client certificates by the attestation - // they carry and ignores the issuer, so a self-issued certificate would be accepted - // on the same terms. Neither caller has been migrated yet, so this RPC still has to - // work; it is scheduled for removal once both are. - // - // Do not build new callers on it. A caller that only needs the KMS root CA can read - // `ca_cert` from `GetMeta` instead. - rpc GetTempCaCert(google.protobuf.Empty) returns (GetTempCaCertResponse); // Sign a certificate rpc SignCert(SignCertRequest) returns (SignCertResponse); } diff --git a/dstack/kms/src/config.rs b/dstack/kms/src/config.rs index f90bdf9ce..25bb10f3d 100644 --- a/dstack/kms/src/config.rs +++ b/dstack/kms/src/config.rs @@ -13,8 +13,6 @@ pub fn load_config_figment(config_file: Option<&str>) -> Figment { load_config("kms", DEFAULT_CONFIG, config_file, false) } -const TEMP_CA_CERT: &str = "tmp-ca.crt"; -const TEMP_CA_KEY: &str = "tmp-ca.key"; const ROOT_CA_CERT: &str = "root-ca.crt"; const ROOT_CA_KEY: &str = "root-ca.key"; const RPC_CERT: &str = "rpc.crt"; @@ -132,23 +130,13 @@ impl KmsConfig { } pub fn keys_exists(&self) -> bool { - self.tmp_ca_cert().exists() - && self.tmp_ca_key().exists() - && self.root_ca_cert().exists() + self.root_ca_cert().exists() && self.root_ca_key().exists() && self.rpc_cert().exists() && self.rpc_key().exists() && self.k256_key().exists() } - pub fn tmp_ca_cert(&self) -> PathBuf { - self.cert_dir.join(TEMP_CA_CERT) - } - - pub fn tmp_ca_key(&self) -> PathBuf { - self.cert_dir.join(TEMP_CA_KEY) - } - pub fn root_ca_cert(&self) -> PathBuf { self.cert_dir.join(ROOT_CA_CERT) } diff --git a/dstack/kms/src/main_service.rs b/dstack/kms/src/main_service.rs index 4c5e72545..188145cc7 100644 --- a/dstack/kms/src/main_service.rs +++ b/dstack/kms/src/main_service.rs @@ -13,9 +13,8 @@ use std::{ use anyhow::{anyhow, bail, Context, Result}; use dstack_kms_rpc::{ kms_server::{KmsRpc, KmsServer}, - AppId, AppKeyResponse, GetAppKeyRequest, GetKmsKeyRequest, GetMetaResponse, - GetTempCaCertResponse, KmsKeyResponse, KmsKeys, PublicKeyResponse, SignCertRequest, - SignCertResponse, + AppId, AppKeyResponse, GetAppKeyRequest, GetKmsKeyRequest, GetMetaResponse, KmsKeyResponse, + KmsKeys, PublicKeyResponse, SignCertRequest, SignCertResponse, }; use dstack_verifier::{CvmVerifier, VerificationDetails}; use fs_err as fs; @@ -45,7 +44,7 @@ pub(crate) mod upgrade_authority; /// How long the auth API's info and this KMS's own authorization are reused. /// -/// `GetMeta`, `GetAppEnvEncryptPubKey` and `GetTempCaCert` need no client certificate, and each +/// `GetMeta` and `GetAppEnvEncryptPubKey` need no client certificate, and each /// would otherwise cost the auth API a round of chain RPC calls. Concurrent callers share one /// in-flight request. App authorization is never cached. const AUTH_API_CACHE_TTL: Duration = Duration::from_secs(1); @@ -165,11 +164,8 @@ pub struct KmsStateInner { config: KmsConfig, keyring: Keyring, /// The CA of the running TLS listener, independent of key selection. - rpc_ca_cert: String, rpc_ca_pubkey: Vec, restart: tokio::sync::Notify, - temp_ca_cert: String, - temp_ca_key: String, verifier: CvmVerifier, attestation_verifier: Arc, self_boot_info: Cache<(), BootInfo>, @@ -233,10 +229,6 @@ impl KmsState { .public_key() .raw .to_vec(); - let temp_ca_key = - fs::read_to_string(config.tmp_ca_key()).context("Faeild to read temp ca key")?; - let temp_ca_cert = - fs::read_to_string(config.tmp_ca_cert()).context("Faeild to read temp ca cert")?; let attestation_verifier = Arc::new( AttestationVerifier::load(&config.attestation) .context("failed to load attestation verifier")?, @@ -256,11 +248,8 @@ impl KmsState { inner: Arc::new(KmsStateInner { config, keyring, - rpc_ca_cert, rpc_ca_pubkey, restart: tokio::sync::Notify::new(), - temp_ca_cert, - temp_ca_key, verifier, attestation_verifier, self_boot_info: Cache::builder() @@ -436,7 +425,6 @@ impl RpcHandler { let info = self.ensure_kms_allowed(&request.vm_config).await?; ensure_key_release_allowed(&info, (&self.state.config).into())?; Ok(KmsKeyResponse { - temp_ca_key: self.state.inner.temp_ca_key.clone(), keys: self .state .keyring @@ -510,22 +498,12 @@ fn ensure_key_release_allowed(boot_info: &BootInfo, policy: KeyReleasePolicy) -> } } -fn ensure_self_key_release_allowed( - self_boot_info: Option<&BootInfo>, - policy: KeyReleasePolicy, -) -> Result<()> { - if let Some(boot_info) = self_boot_info { - ensure_key_release_allowed(boot_info, policy)?; - } - Ok(()) -} - impl RpcHandler { - async fn ensure_self_allowed(&self) -> Result<(Option, u32)> { + async fn ensure_self_allowed(&self) -> Result { if !self.state.config.enforce_self_authorization { let epoch = self.state.current_epoch_from_info().await?; self.state.ensure_rpc_ca(epoch)?; - return Ok((None, epoch)); + return Ok(epoch); } let boot_info = self .state @@ -553,7 +531,7 @@ impl RpcHandler { .await .map_err(|err| anyhow!("{err:#}"))?; let epoch = self.state.authorized_epoch(&response, None)?; - Ok((Some(boot_info), epoch)) + Ok(epoch) } fn ensure_attested(&self) -> Result<&VerifiedAttestation> { @@ -677,7 +655,7 @@ impl KmsRpc for RpcHandler { } async fn get_app_env_encrypt_pub_key(self, request: AppId) -> Result { - let (_, epoch) = self + let epoch = self .ensure_self_allowed() .await .context("KMS self authorization failed")?; @@ -757,24 +735,6 @@ impl KmsRpc for RpcHandler { self.handover_keys(request).await } - /// Serve the temp CA certificate and key. - /// - /// Retained for older guest images and tokenless KMS onboarding. Current - /// guests use self-issued RA-TLS certificates and fetch the root CA via GetMeta. - /// The returned key authenticates nobody; client identity comes from attestation. - async fn get_temp_ca_cert(self) -> Result { - let (self_boot_info, _) = self - .ensure_self_allowed() - .await - .context("KMS self authorization failed")?; - ensure_self_key_release_allowed(self_boot_info.as_ref(), (&self.state.config).into())?; - Ok(GetTempCaCertResponse { - temp_ca_cert: self.state.inner.temp_ca_cert.clone(), - temp_ca_key: self.state.inner.temp_ca_key.clone(), - ca_cert: self.state.rpc_ca_cert.clone(), - }) - } - async fn sign_cert(self, request: SignCertRequest) -> Result { self.ensure_self_allowed() .await @@ -1409,16 +1369,6 @@ mod tests { .expect("TCB/advisory policy should be decided by the auth API, not this local gate"); } - #[test] - fn snp_self_boot_info_uses_same_release_policy_for_temp_ca() { - let boot_info = snp_boot_info(); - - ensure_self_key_release_allowed(Some(&boot_info), KeyReleasePolicy::default()) - .expect_err("disabled SNP self boot info must not receive temp CA key material"); - ensure_self_key_release_allowed(Some(&boot_info), SNP_ENABLED) - .expect("enabled clean SNP self boot info should pass the temp CA release gate"); - } - #[test] fn nitro_enclave_key_release_requires_explicit_enablement() { let mut boot_info = snp_boot_info(); @@ -1436,13 +1386,6 @@ mod tests { // explicitly enabled permits it ensure_key_release_allowed(&boot_info, NITRO_ENCLAVE_ENABLED) .expect("explicitly enabled nitro enclave key release should allow the boot info"); - - // the temp CA path uses the same gate - ensure_self_key_release_allowed(Some(&boot_info), KeyReleasePolicy::default()).expect_err( - "disabled nitro enclave self boot info must not receive temp CA key material", - ); - ensure_self_key_release_allowed(Some(&boot_info), NITRO_ENCLAVE_ENABLED) - .expect("enabled nitro enclave self boot info should pass the temp CA release gate"); } #[test] @@ -1518,19 +1461,8 @@ mod tests { .build() .self_signed() .unwrap(); - let tmp_ca_key = KeyPair::generate_for(&PKCS_ECDSA_P256_SHA256).unwrap(); - let tmp_ca_cert = CertRequest::builder() - .org_name("Dstack") - .subject("Dstack Client Temp CA") - .ca_level(0) - .key(&tmp_ca_key) - .build() - .self_signed() - .unwrap(); fs::write(cert_dir.join("root-ca.key"), ca_key.serialize_pem()).unwrap(); fs::write(cert_dir.join("root-ca.crt"), ca_cert.pem()).unwrap(); - fs::write(cert_dir.join("tmp-ca.key"), tmp_ca_key.serialize_pem()).unwrap(); - fs::write(cert_dir.join("tmp-ca.crt"), tmp_ca_cert.pem()).unwrap(); fs::write( cert_dir.join("root-k256.key"), SigningKey::random(&mut rand::rngs::OsRng).to_bytes(), @@ -1585,8 +1517,12 @@ mod tests { }); keyring.chain_hashes = root_chain_hashes(keyring.roots.iter().map(|root| (&root.ca.key, &root.k256))); - let state = - KmsState::new(initial.config.clone(), keyring, initial.rpc_ca_cert.clone()).unwrap(); + let state = KmsState::new( + initial.config.clone(), + keyring, + initial.keyring.roots[0].ca.pem_cert.clone(), + ) + .unwrap(); let result = |epoch: u32| AuthorizationResult { is_allowed: true, rejection_reason: String::new(), @@ -1607,7 +1543,7 @@ mod tests { .await .is_ok() ); - assert_eq!(state.rpc_ca_cert, initial.rpc_ca_cert); + assert_eq!(state.rpc_ca_pubkey, initial.rpc_ca_pubkey); let mut rejected = result(1); rejected.is_allowed = false; diff --git a/dstack/kms/src/onboard_service.rs b/dstack/kms/src/onboard_service.rs index 6d4a857c8..fb0f8a7d6 100644 --- a/dstack/kms/src/onboard_service.rs +++ b/dstack/kms/src/onboard_service.rs @@ -18,7 +18,7 @@ use dstack_kms_rpc::{ use fs_err as fs; use k256::ecdsa::SigningKey; use ra_rpc::{ - client::{CertInfo, RaClient, RaClientConfig}, + client::{CertInfo, RaClientConfig}, CallContext, RpcCall, }; use ra_tls::{ @@ -26,7 +26,7 @@ use ra_tls::{ AttestationVerifier, GetDeviceId, PlatformEvidence, QuoteContentType, VerifiedAttestation, VersionedAttestation, }, - cert::{CaCert, CertRequest}, + cert::CertRequest, rcgen::{Certificate, KeyPair, PKCS_ECDSA_P256_SHA256}, }; use safe_write::{safe_write, safe_write_with_mode}; @@ -548,9 +548,6 @@ mod tests { fn handover(epochs: u32) -> KmsKeyResponse { KmsKeyResponse { - temp_ca_key: KeyPair::generate_for(&PKCS_ECDSA_P256_SHA256) - .unwrap() - .serialize_pem(), keys: (0..epochs) .map(|_| KmsKeys { ca_key: KeyPair::generate_for(&PKCS_ECDSA_P256_SHA256) @@ -609,6 +606,8 @@ mod tests { .await .unwrap(); keys.store(&cfg).unwrap(); + assert!(!cfg.cert_dir.join("tmp-ca.key").exists()); + assert!(!cfg.cert_dir.join("tmp-ca.crt").exists()); assert_eq!(crate::main_service::read_root_epoch(&cfg).unwrap(), 2); for (epoch, k256) in (0..).zip(&sent) { @@ -830,8 +829,6 @@ struct Keys { epoch: u32, history: Vec, k256_key: SigningKey, - tmp_ca_key: KeyPair, - tmp_ca_cert: Certificate, ca_key: KeyPair, ca_cert: Certificate, rpc_key: KeyPair, @@ -841,19 +838,10 @@ struct Keys { impl Keys { async fn generate(domain: &str, attest_rpc_cert: bool) -> Result { - let tmp_ca_key = KeyPair::generate_for(&PKCS_ECDSA_P256_SHA256)?; let ca_key = KeyPair::generate_for(&PKCS_ECDSA_P256_SHA256)?; let rpc_key = KeyPair::generate_for(&PKCS_ECDSA_P256_SHA256)?; let k256_key = SigningKey::random(&mut rand::rngs::OsRng); - Self::from_keys( - tmp_ca_key, - ca_key, - rpc_key, - k256_key, - domain, - attest_rpc_cert, - ) - .await + Self::from_keys(ca_key, rpc_key, k256_key, domain, attest_rpc_cert).await } /// Attest the current root public keys. Epoch 0 keeps the original @@ -870,21 +858,12 @@ impl Keys { } async fn from_keys( - tmp_ca_key: KeyPair, ca_key: KeyPair, rpc_key: KeyPair, k256_key: SigningKey, domain: &str, attest_rpc_cert: bool, ) -> Result { - let tmp_ca_cert = CertRequest::builder() - .org_name("Dstack") - .subject("Dstack Client Temp CA") - .ca_level(0) - .key(&tmp_ca_key) - .build() - .self_signed()?; - let ca_cert = self_signed_root_ca(&ca_key)?; // The only place the KMS embeds its own attestation. Skipping it lets // the KMS run outside a TEE for development; it does not affect the @@ -919,8 +898,6 @@ impl Keys { epoch: 0, history: Vec::new(), k256_key, - tmp_ca_key, - tmp_ca_cert, ca_key, ca_cert, rpc_key, @@ -960,17 +937,7 @@ impl Keys { .context("Missing source KMS attestation") }; - // With a token, `other_kms_url` is the source's admin listener, which accepts - // a self-issued RA-TLS client certificate. Without one, the source predates - // 0.6.0: it serves `KMS.GetKmsKey` publicly and pins its temp CA. - let tmp_ca = if source_token.is_empty() { - let client = RaClient::new(other_kms_url.to_string(), true)?; - let tmp_ca = KmsClient::new(client).get_temp_ca_cert().await?; - Some((tmp_ca.temp_ca_cert, tmp_ca.temp_ca_key)) - } else { - None - }; - let (ra_cert, ra_key) = gen_ra_cert(tmp_ca).await?; + let (ra_cert, ra_key) = gen_ra_cert().await?; let client = RaClientConfig::builder() .tls_skip_ca_verification(true) .tls_built_in_root_certs(false) @@ -1056,17 +1023,7 @@ impl Keys { .collect::>>()?; let rpc_key = KeyPair::generate_for(&PKCS_ECDSA_P256_SHA256)?; - let tmp_ca_key = - KeyPair::from_pem(&keys_res.temp_ca_key).context("Failed to parse tmp CA key")?; - let keys = Self::from_keys( - tmp_ca_key, - ca_key, - rpc_key, - k256_key, - domain, - cfg.attest_rpc_cert, - ) - .await?; + let keys = Self::from_keys(ca_key, rpc_key, k256_key, domain, cfg.attest_rpc_cert).await?; Ok(Self { epoch, history, @@ -1093,7 +1050,6 @@ impl Keys { } fn store_keys(&self, cfg: &KmsConfig) -> Result<()> { - safe_write_with_mode(cfg.tmp_ca_key(), self.tmp_ca_key.serialize_pem(), 0o600)?; safe_write_with_mode(cfg.rpc_key(), self.rpc_key.serialize_pem(), 0o600)?; safe_write_with_mode(cfg.root_ca_key(), self.ca_key.serialize_pem(), 0o600)?; safe_write_with_mode(cfg.k256_key(), self.k256_key.to_bytes(), 0o600)?; @@ -1101,7 +1057,6 @@ impl Keys { } fn store_certs(&self, cfg: &KmsConfig) -> Result<()> { - safe_write(cfg.tmp_ca_cert(), self.tmp_ca_cert.pem())?; safe_write(cfg.root_ca_cert(), self.ca_cert.pem())?; safe_write(cfg.rpc_cert(), self.rpc_cert.pem())?; Ok(()) @@ -1113,7 +1068,6 @@ impl Keys { pub(crate) async fn update_certs(cfg: &KmsConfig, epoch: u32, latest: u32) -> Result<()> { let roots = cfg.root_key_paths(epoch, latest); // Read existing keys - let tmp_ca_key = KeyPair::from_pem(&fs::read_to_string(cfg.tmp_ca_key())?)?; let ca_key = KeyPair::from_pem(&fs::read_to_string(&roots.ca_key)?)?; let rpc_key = KeyPair::from_pem(&fs::read_to_string(cfg.rpc_key())?)?; @@ -1129,28 +1083,15 @@ pub(crate) async fn update_certs(cfg: &KmsConfig, epoch: u32, latest: u32) -> Re let domain = domain.trim(); // Regenerate certificates using existing keys - let keys = Keys::from_keys( - tmp_ca_key, - ca_key, - rpc_key, - k256_key, - domain, - cfg.attest_rpc_cert, - ) - .await - .context("Failed to regenerate certificates")?; + let keys = Keys::from_keys(ca_key, rpc_key, k256_key, domain, cfg.attest_rpc_cert) + .await + .context("Failed to regenerate certificates")?; renew_ca_cert_if_expiring( &roots.ca_cert, keys.ca_cert.pem(), "KMS root CA certificate", )?; - renew_ca_cert_if_expiring( - cfg.tmp_ca_cert(), - keys.tmp_ca_cert.pem(), - "temporary client CA certificate", - )?; - // The RPC leaf depends on the refreshed domain and platform attestation, so // it is reissued on every startup. safe_write(cfg.rpc_cert(), keys.rpc_cert.pem())?; @@ -1265,8 +1206,8 @@ fn keccak256(msg: &[u8]) -> [u8; 32] { hasher.finalize().into() } -/// Mint an RA-TLS client certificate, signed by `ca` (cert and key PEM) or self-issued. -async fn gen_ra_cert(ca: Option<(String, String)>) -> Result<(String, String)> { +/// Mint a self-issued RA-TLS client certificate. +async fn gen_ra_cert() -> Result<(String, String)> { use ra_tls::cert::CertRequest; use ra_tls::rcgen::{KeyPair, PKCS_ECDSA_P256_SHA256}; @@ -1279,14 +1220,11 @@ async fn gen_ra_cert(ca: Option<(String, String)>) -> Result<(String, String)> { let attestation = VersionedAttestation::from_bytes(&response.attestation).context("Invalid attestation")?; let req = CertRequest::builder() - .subject("RA-TLS TEMP Cert") + .subject("RA-TLS Client Cert") .attestation(&attestation) .key(&key) + .usage_client_auth(true) .build(); - let cert = match ca { - Some((cert_pem, key_pem)) => CaCert::new(cert_pem, key_pem)?.sign(req), - None => req.self_signed(), - } - .context("Failed to sign certificate")?; + let cert = req.self_signed().context("Failed to sign certificate")?; Ok((cert.pem(), key.serialize_pem())) } diff --git a/dstack/ra-rpc/src/ratls_client_verifier.rs b/dstack/ra-rpc/src/ratls_client_verifier.rs index d3c467808..f886aac5b 100644 --- a/dstack/ra-rpc/src/ratls_client_verifier.rs +++ b/dstack/ra-rpc/src/ratls_client_verifier.rs @@ -61,7 +61,7 @@ impl ClientCertVerifier for RaTlsClientVerifier { /// Anonymous connections stay allowed, matching the previous /// `[rpc.tls.mutual] mandatory = false`. Handlers that need an attested caller - /// enforce it themselves; the unauthenticated RPCs (`GetMeta`, `GetTempCaCert`) + /// enforce it themselves; the unauthenticated RPCs (such as `GetMeta`) /// are reachable without a certificate by design. fn client_auth_mandatory(&self) -> bool { false diff --git a/dstack/tests/docs/kms-bootstrap-onboard.md b/dstack/tests/docs/kms-bootstrap-onboard.md index f9f4d3234..ac8b32164 100644 --- a/dstack/tests/docs/kms-bootstrap-onboard.md +++ b/dstack/tests/docs/kms-bootstrap-onboard.md @@ -76,7 +76,7 @@ Operational notes: 2. `Boot Progress: done` does **not** guarantee the onboard endpoint is ready. 3. The onboarding completion endpoint is **GET `/finish`**. 4. On teepod with gateway, onboard mode usually uses the `-8000` URL, while runtime TLS KMS RPC usually uses the `-8000s` URL. **Port forwarding** (`--port tcp:0.0.0.0::8000`) is simpler than gateway for testing, because gateway requires the auth API to return a `gatewayAppId` at boot time. -5. If you use a very small custom webhook instead of the real auth service, `KMS.GetMeta` may fail because `auth_api.get_info()` expects extra chain / contract metadata fields. In that case, use `GetTempCaCert` as the runtime readiness probe. +5. If you use a very small custom webhook instead of the real auth service, `KMS.GetMeta` may fail because `auth_api.get_info()` expects extra chain / contract metadata fields. In that case, use `GetAppEnvEncryptPubKey` as the runtime readiness probe. 6. dstack CVMs use QEMU user-mode networking — the host is reachable at **`10.0.2.2`** from inside the CVM. The `source_url` in `Onboard.Onboard` must use a CVM-reachable address (e.g., `https://10.0.2.2:/prpc`), not `127.0.0.1`. 7. **~~Remote KMS attestation has an empty `osImageHash`.~~** Fixed: RA-TLS certs now use the unified `PHALA_RATLS_ATTESTATION` format which preserves `vm_config`. For old source KMS instances, the receiver-side check fills `osImageHash` from the local KMS's own value automatically. No special `"0x"` entry in `osImages` is needed anymore. @@ -205,7 +205,7 @@ curl -sk "${KMS_SRC_RUNTIME%/}/prpc/KMS.GetMeta?json" \ - `k256_pubkey` - `bootstrap_info` -If `KMS.GetMeta` fails because your minimal webhook does not return chain / contract info, use `GetTempCaCert` below as the runtime readiness probe instead. +If `KMS.GetMeta` fails because your minimal webhook does not return chain / contract info, use `GetAppEnvEncryptPubKey` below as the runtime readiness probe instead. --- @@ -273,24 +273,25 @@ This section folds the runtime trusted-RPC verification into the same flow. | bootstrap deny | source-side auth leaves `kms.mrAggregated` empty or omits the current `kms-src` MR | `Onboard.Bootstrap` on `kms-src` | `KMS is not allowed to bootstrap`, `MR aggregated not allowed` | | onboard deny (receiver-side) | destination-side auth leaves `kms.mrAggregated` empty or omits the current `kms-src` MR | `Onboard.Onboard` on `kms-dst` | source KMS not allowed / onboarding failed | | onboard deny (source-side) | source-side auth leaves `kms.mrAggregated` empty or omits the current `kms-dst` MR | `Onboard.Onboard` on `kms-dst` | source rejected destination caller / `GetKmsKey` authorization failed | -| runtime deny | auth removes the running KMS from `kms.mrAggregated` | `GetTempCaCert` or another trusted RPC | `KMS self authorization failed`, `KMS is not allowed` | +| runtime deny | auth removes the running KMS from `kms.mrAggregated` | `GetAppEnvEncryptPubKey` or another trusted RPC | `KMS self authorization failed`, `KMS is not allowed` | Use the happy-path steps below first, then flip policies one by one and rerun the indicated probe. -### 7.1 Minimum canary: `GetTempCaCert` +### 7.1 Minimum canary: `GetAppEnvEncryptPubKey` ```bash -curl -sk "${KMS_SRC_RUNTIME%/}/prpc/KMS.GetTempCaCert?json" \ - | tee /tmp/kms-bootstrap-onboard/kms-src-get-temp-ca.json | jq . +curl -sk "${KMS_SRC_RUNTIME%/}/prpc/KMS.GetAppEnvEncryptPubKey?json" \ + -H 'Content-Type: application/json' -d '{"app_id":"1111111111111111111111111111111111111111"}' \ + | tee /tmp/kms-bootstrap-onboard/kms-src-get-env-pubkey.json | jq . ``` Expected result: - success - response contains: - - `temp_ca_cert` - - `temp_ca_key` - - `ca_cert` + - `public_key` + - `signature_v1` + - `timestamp` ### 7.2 `GetKmsKey` @@ -332,8 +333,9 @@ Expected result: After a normal happy-path run, flip source-side auth policy to deny `kms-src` itself and retry: ```bash -curl -sk "${KMS_SRC_RUNTIME%/}/prpc/KMS.GetTempCaCert?json" \ - | tee /tmp/kms-bootstrap-onboard/kms-src-get-temp-ca-after-deny.json | jq . +curl -sk "${KMS_SRC_RUNTIME%/}/prpc/KMS.GetAppEnvEncryptPubKey?json" \ + -H 'Content-Type: application/json' -d '{"app_id":"1111111111111111111111111111111111111111"}' \ + | tee /tmp/kms-bootstrap-onboard/kms-src-get-env-pubkey-after-deny.json | jq . ``` Expected result: @@ -405,8 +407,9 @@ Then restore both source-side and destination-side allowlists and rerun the norm After a successful bootstrap or onboard, remove the running KMS's own MR from `kms.mrAggregated` and retry: ```bash -curl -sk "${KMS_SRC_RUNTIME%/}/prpc/KMS.GetTempCaCert?json" \ - | tee /tmp/kms-bootstrap-onboard/kms-src-get-temp-ca-denied.json | jq . +curl -sk "${KMS_SRC_RUNTIME%/}/prpc/KMS.GetAppEnvEncryptPubKey?json" \ + -H 'Content-Type: application/json' -d '{"app_id":"1111111111111111111111111111111111111111"}' \ + | tee /tmp/kms-bootstrap-onboard/kms-src-get-env-pubkey-denied.json | jq . ``` Expected result: @@ -414,7 +417,7 @@ Expected result: - response contains `.error` - error indicates KMS self authorization failed or the KMS is not allowed -You can repeat the same check on `kms-dst` after onboard by removing `kms-dst` from destination-side policy and retrying `KMS.GetTempCaCert`. +You can repeat the same check on `kms-dst` after onboard by removing `kms-dst` from destination-side policy and retrying `KMS.GetAppEnvEncryptPubKey`. --- @@ -427,8 +430,8 @@ For each run, save: 3. onboard response 4. `/finish` responses 5. runtime `KMS.GetMeta` responses -6. trusted RPC responses such as `GetTempCaCert` -7. deny-case responses such as `kms-src-bootstrap-denied.json`, `kms-dst-onboard-denied.json`, `kms-dst-onboard-denied-by-src.json`, and `kms-src-get-temp-ca-denied.json` +6. trusted RPC responses such as `GetAppEnvEncryptPubKey` +7. deny-case responses such as `kms-src-bootstrap-denied.json`, `kms-dst-onboard-denied.json`, `kms-dst-onboard-denied-by-src.json`, and `kms-src-get-env-pubkey-denied.json` 8. auth policy snapshots used during the run Recommended archive: @@ -447,8 +450,8 @@ The flow is considered validated if all of the following are true: 2. `kms-src` transitions to runtime mode successfully 3. `kms-dst` onboard succeeds against `kms-src` 4. `kms-dst` transitions to runtime mode successfully -5. runtime metadata probes succeed on both KMS instances, or `GetTempCaCert` succeeds when `GetMeta` is unavailable with a minimal webhook -6. at least one trusted runtime RPC such as `GetTempCaCert` succeeds +5. runtime metadata probes succeed on both KMS instances, or `GetAppEnvEncryptPubKey` succeeds when `GetMeta` is unavailable with a minimal webhook +6. at least one trusted runtime RPC such as `GetAppEnvEncryptPubKey` succeeds 7. the selected deny cases fail at the expected RPC with an authorization error --- diff --git a/dstack/tests/docs/kms-self-authorization.md b/dstack/tests/docs/kms-self-authorization.md index d9e365fd9..88b33c787 100644 --- a/dstack/tests/docs/kms-self-authorization.md +++ b/dstack/tests/docs/kms-self-authorization.md @@ -6,7 +6,7 @@ The goal is to validate the following behaviors without depending on `kms/e2e/` 1. **Bootstrap self-check**: a KMS must call the auth API and verify that **itself** is allowed before bootstrap succeeds. 2. **Onboard receiver-side source check**: a new KMS must reject onboarding if the **source KMS** is not allowed by the receiver's auth policy. -3. **Trusted RPC self-check**: trusted KMS RPCs such as `GetTempCaCert`, `GetKmsKey`, `GetAppKey`, and `SignCert` must fail when the running KMS is no longer allowed by its auth policy. +3. **Trusted RPC self-check**: trusted KMS RPCs such as `GetAppEnvEncryptPubKey`, `GetKmsKey`, `GetAppKey`, and `SignCert` must fail when the running KMS is no longer allowed by its auth policy. 4. **Attestation requirement**: KMS always requires attestation; for local development without TDX hardware, use `sdk/simulator`. This guide is written as a deployment-and-test runbook so an AI agent can follow it end-to-end. @@ -589,15 +589,16 @@ Verify that a running KMS re-checks its own authorization on trusted RPCs. ### Recommended canary RPC -Use `GetTempCaCert` first. It is simpler than `GetAppKey` because it does not require preparing an attested app client, but it still exercises the new runtime self-check. +Use `GetAppEnvEncryptPubKey` first. It is simpler than `GetAppKey` because it does not require preparing an attested app client, but it still exercises the new runtime self-check. ### Steps 1. While `kms-src` is healthy, confirm the canary RPC works: ```bash -curl -sk "${KMS_SRC_RUNTIME%/}/prpc/KMS.GetTempCaCert?json" \ - | tee /tmp/kms-self-auth/get-temp-ca-allowed.json | jq . +curl -sk "${KMS_SRC_RUNTIME%/}/prpc/KMS.GetAppEnvEncryptPubKey?json" \ + -H 'Content-Type: application/json' -d '{"app_id":"1111111111111111111111111111111111111111"}' \ + | tee /tmp/kms-self-auth/get-env-pubkey-allowed.json | jq . ``` 2. Flip `auth-simple-src` to deny `kms-src` itself: @@ -609,13 +610,14 @@ cp /tmp/kms-self-auth/auth-src-deny-self.json /tmp/kms-self-auth/auth-src.json 3. Retry the same RPC: ```bash -curl -sk "${KMS_SRC_RUNTIME%/}/prpc/KMS.GetTempCaCert?json" \ - | tee /tmp/kms-self-auth/get-temp-ca-denied.json | jq . +curl -sk "${KMS_SRC_RUNTIME%/}/prpc/KMS.GetAppEnvEncryptPubKey?json" \ + -H 'Content-Type: application/json' -d '{"app_id":"1111111111111111111111111111111111111111"}' \ + | tee /tmp/kms-self-auth/get-env-pubkey-denied.json | jq . ``` ### Expected result -- before the policy flip: `GetTempCaCert` succeeds +- before the policy flip: `GetAppEnvEncryptPubKey` succeeds - after the policy flip: the response contains `.error` - the error should indicate KMS self-authorization failed, or that the KMS is not allowed @@ -667,7 +669,7 @@ For each run, save: 2. auth config snapshots used for each step 3. bootstrap/onboard RPC responses 4. `KMS.GetMeta` output after successful boot -5. `GetTempCaCert` allow/deny responses +5. `GetAppEnvEncryptPubKey` allow/deny responses 6. relevant CVM logs if a step fails unexpectedly Recommended archive: From 9c90d549f623e4329ec37084c6d6374ed6540ecd Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Fri, 6 Nov 2026 19:22:37 -0800 Subject: [PATCH 017/166] docs(kms): align unreleased notes with direct key retrieval Signed-off-by: Kevin Wang --- CHANGELOG.md | 19 ++++++++----------- 1 file changed, 8 insertions(+), 11 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index a222c8c15..3d7ae700c 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -19,23 +19,20 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 - dstack-util, cert-client: guests mint self-issued RA-TLS client certificates for KMS calls instead of fetching the temp CA and minting from it; the KMS authenticates the quote inside the certificate, not its issuer. - The KMS root CA now comes from `GetMeta`, over the same unauthenticated - connection `GetTempCaCert` used, so the trust posture is unchanged: relying - parties bind the KMS through the `key-provider` launch event (RTMR3/PCR14) - against the on-chain `kmsInfo` or a pinned key, and `verify_key_provider_id` - also enforces it in-guest when the app pins `key_provider_id`. A guest built - from this change needs KMS 0.6.0 or later; older KMS releases pin the temp CA - and refuse the certificate at the TLS handshake + Guests fetch the root CA with the key response and bind it to the attested + RPC certificate, without a metadata preflight. Provider pins and the + `key-provider` launch event (RTMR3/PCR14) retain the root identity binding. + The new flow requires an attested KMS that accepts self-issued certificates. - dstack-types: `KeyProvider::Kms` drops `tmp_ca_key` and `tmp_ca_cert`, so `.appkeys.json` no longer carries them. Files written by older images still parse, but out-of-tree readers that require those fields fail on new files ### Fixed -- kms: onboarding from a source older than 0.6.0 (no `source_token`) now - verifies the source's attestation on the connection that returns the root - keys, and checks it against local policy before using them. It used to - verify the source only on the earlier `GetTempCaCert` connection +- kms: onboarding with or without `source_token` verifies the source's + attestation on the connection that returns the root keys and checks local + policy before using them. Legacy sources requiring temp-CA-issued clients + must be upgraded. ## [0.6.3] - 2026-11-01 From 8237ca9899d44ca69c529eaac3df8500198e673f Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Fri, 6 Nov 2026 19:28:24 -0800 Subject: [PATCH 018/166] docs(kms): record temporary CA protocol retirement Signed-off-by: Kevin Wang --- docs/security/public-security-reports.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/security/public-security-reports.md b/docs/security/public-security-reports.md index d04b40628..03f342b23 100644 --- a/docs/security/public-security-reports.md +++ b/docs/security/public-security-reports.md @@ -59,7 +59,7 @@ These issues were filed as concrete vulnerability reports, security audit findin | [#616](https://github.com/Dstack-TEE/dstack/issues/616) Host-controlled Docker registry mirror enables image substitution attacks | Closed | Not a production vulnerability | Registry mirrors are untrusted transport. Digest-pinned image references and measured compose configuration protect against substitution. No code fix was applied | | [#617](https://github.com/Dstack-TEE/dstack/issues/617) Guest agent exposes raw private keys to all local processes | Closed | Not a production vulnerability | dstack treats a CVM as one application trust domain. It does not provide per-container key isolation inside the same measured application. No code fix was applied | | [#618](https://github.com/Dstack-TEE/dstack/issues/618) Disk encryption disableable via kernel cmdline, not measured in RTMR | Closed | Not a production vulnerability | The kernel command line is measured into RTMR2, so changing `dstack.storage_encrypted=false` changes attestation evidence. No code fix was applied | -| [#619](https://github.com/Dstack-TEE/dstack/issues/619) KMS `get_temp_ca_cert` returns temp CA private key without authentication | Closed | Duplicate | The report duplicates the private advisory response for the temp CA bootstrap flow. [#1106](https://github.com/Dstack-TEE/dstack/pull/1106) removes the reason the temp CA existed - the KMS no longer pins it - so callers can be migrated off it; the RPC itself is retained until guests and KMS-to-KMS onboarding are | +| [#619](https://github.com/Dstack-TEE/dstack/issues/619) KMS `get_temp_ca_cert` returns temp CA private key without authentication | Closed | Duplicate | The report duplicates the private advisory response for the temp CA bootstrap flow. [#1106](https://github.com/Dstack-TEE/dstack/pull/1106) removes the reason the temp CA existed - the KMS no longer pins it - so callers can be migrated off it; the root-epoch rotation stack removes the RPC and temporary CA material after migrating guests and KMS-to-KMS onboarding | ## Related security roadmap and hardening From 846180f48fdc83f605393515f4b2d55d05eb63e6 Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Fri, 6 Nov 2026 20:02:23 -0800 Subject: [PATCH 019/166] refactor(ra-rpc): name CA bypass explicitly and isolate TLS configuration Signed-off-by: Kevin Wang --- dstack/cert-client/src/lib.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/dstack/cert-client/src/lib.rs b/dstack/cert-client/src/lib.rs index 5e45ce24c..431d94ad9 100644 --- a/dstack/cert-client/src/lib.rs +++ b/dstack/cert-client/src/lib.rs @@ -70,7 +70,7 @@ pub async fn request_app_keys( .remote_uri(url) .tls_client_cert(client_cert.cert_pem) .tls_client_key(client_cert.key_pem) - .tls_no_check(trusted_ca.is_none()) + .tls_skip_ca_verification(trusted_ca.is_none()) .tls_built_in_root_certs(false) .maybe_tls_ca_cert(trusted_ca.clone()) .attestation_verifier(attestation_verifier) From 8fdbd3d15c3f8b3be16a93f5c6aab2232f051ed6 Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Fri, 6 Nov 2026 23:07:01 -0800 Subject: [PATCH 020/166] refactor(kms): separate KMS and app authorization and reserve app epochs Signed-off-by: Kevin Wang --- docs/deployment.md | 2 +- dstack/kms/README.md | 21 +++-- dstack/kms/auth-eth-bun/README.md | 8 +- dstack/kms/auth-eth-bun/index.test.ts | 8 +- dstack/kms/auth-eth-bun/index.ts | 77 ++++++++++--------- dstack/kms/auth-eth-bun/openapi.json | 44 ++++++++++- dstack/kms/auth-eth/README.md | 8 +- dstack/kms/auth-eth/contracts/DstackKms.sol | 34 +++++++- dstack/kms/auth-eth/src/ethereum.ts | 8 +- dstack/kms/auth-eth/src/main.test.ts | 15 ++-- dstack/kms/auth-eth/src/server.ts | 33 ++++---- dstack/kms/auth-eth/src/types.ts | 1 + dstack/kms/auth-eth/test/DstackKms.t.sol | 51 ++++++++---- dstack/kms/auth-mock/README.md | 8 +- dstack/kms/auth-mock/index.test.ts | 8 +- dstack/kms/auth-mock/index.ts | 27 ++++--- dstack/kms/auth-mock/openapi.json | 44 ++++++++++- dstack/kms/auth-simple/README.md | 8 +- dstack/kms/auth-simple/index.test.ts | 14 ++-- dstack/kms/auth-simple/index.ts | 28 ++++--- dstack/kms/src/main_service.rs | 57 ++------------ .../kms/src/main_service/upgrade_authority.rs | 46 ++++++++--- 22 files changed, 357 insertions(+), 193 deletions(-) diff --git a/docs/deployment.md b/docs/deployment.md index 27adbe9e8..223e7d191 100644 --- a/docs/deployment.md +++ b/docs/deployment.md @@ -548,7 +548,7 @@ for the full procedure; the KMS steps are: `PopRootEpoch` (or remove the `rootEpochs` entry) and rotate again; nodes that imported the dropped N must be onboarded again. -Each `checkAuthorization` result carries the current epoch and its chain hash, +Each `checkKmsAuthorization` or `checkAppAuthorization` result carries the current epoch and its chain hash, including when authorization is denied. KMS verifies these against its keyring; there is no cached serving epoch. An optional `requestedEpoch` selects historical keys, and omission selects the returned current epoch. Requests for future diff --git a/dstack/kms/README.md b/dstack/kms/README.md index be6cb0d92..1e0d4f254 100644 --- a/dstack/kms/README.md +++ b/dstack/kms/README.md @@ -246,12 +246,21 @@ The `SignCert` RPC is used by the dstack app to sign a TLS certificate. In this ## Authorization protocol -Normal KMS operations use `POST /checkAuthorization`. The request carries the -verified attestation identity (the same identity fields as the legacy -`/bootAuth/*` requests). The Ethereum backends call -`DstackKms.checkAuthorization(AuthorizationRequest)` once; the contract selects KMS policy when -`appId == address(this)` (the proxy address), and app policy otherwise. KMS-only -operations also check this identity before accepting the result. +Normal operations use two explicit authorization endpoints: + +- `POST /checkKmsAuthorization` for KMS self-authorization and root-key handover. +- `POST /checkAppAuthorization` for application keys and certificates. + +The Ethereum backends call the matching `DstackKms` method. Both accept +`AuthorizationRequest` and return `AuthorizationResult`; no caller-kind flag is +sent. The KMS method requires the attested `appId` to equal the proxy address, +and the app method rejects that address and applies registered app policy. +KMS callers do not fetch contract metadata to decide which policy to use. + +`appEpoch` reserves the application's own key epoch, independently of KMS root +epochs. HTTP requests default it to 0; KMS currently sends 0. Nonzero values are +rejected until app-epoch support is implemented, and rejection responses still +carry the current root epoch and chain hash. `requestedEpoch` is optional in the JSON request. Omit it to use the epoch in this authorization result; an explicit `0` selects epoch 0. Solidity represents diff --git a/dstack/kms/auth-eth-bun/README.md b/dstack/kms/auth-eth-bun/README.md index 61345133f..91178ac8f 100644 --- a/dstack/kms/auth-eth-bun/README.md +++ b/dstack/kms/auth-eth-bun/README.md @@ -201,11 +201,13 @@ the test suite (`index.test.ts`) validates: - ✅ optional field handling -### Unified authorization +### Authorization endpoints -`POST /checkAuthorization` returns `AuthorizationResult`, including +`POST /checkKmsAuthorization` and `POST /checkAppAuthorization` return `AuthorizationResult`, including `rejectionReason`, `gatewayAppId`, `currentRootEpoch`, and `currentRootEpochHash`. -It selects KMS or app policy from the verified `appId`. The optional +The selected endpoint determines the policy; the backend checks that the verified +`appId` matches that scope. `appEpoch` defaults to 0; nonzero values are currently +rejected. The optional `requestedEpoch` 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 fbbd64841..504902f49 100644 --- a/dstack/kms/auth-eth-bun/index.test.ts +++ b/dstack/kms/auth-eth-bun/index.test.ts @@ -834,14 +834,14 @@ describe('request hardening', () => { }); -describe('POST /checkAuthorization', () => { +describe('POST /checkAppAuthorization', () => { const appInfo = { appId: '0x' + '11'.repeat(20), composeHash: '0x' + '22'.repeat(32), instanceId: '', deviceId: '0x' + '33'.repeat(32), mrAggregated: '0x' + '44'.repeat(32), mrSystem: '0x' + '55'.repeat(32), osImageHash: '0x' + '66'.repeat(32), tcbStatus: 'UpToDate', advisoryIds: [], }; - const authorize = (requestedEpoch?: number) => appFetch(new Request('http://localhost/checkAuthorization', { + const authorize = (requestedEpoch?: number) => appFetch(new Request('http://localhost/checkAppAuthorization', { method: 'POST', headers: { 'Content-Type': 'application/json' }, body: JSON.stringify({ ...appInfo, requestedEpoch }), })); @@ -856,8 +856,8 @@ describe('POST /checkAuthorization', () => { expect(await response.json()).toEqual({ ...result, currentRootEpoch: 2 }); expect(mockReadContract).toHaveBeenCalledTimes(1); expect(mockReadContract).toHaveBeenCalledWith(expect.objectContaining({ - functionName: 'checkAuthorization', blockNumber: 98n, - args: [expect.objectContaining({ appInfo: expect.objectContaining({ appId: appInfo.appId }), requestedEpoch: { isSet: false, value: 0n } })], + functionName: 'checkAppAuthorization', blockNumber: 98n, + args: [expect.objectContaining({ appInfo: expect.objectContaining({ appId: appInfo.appId }), appEpoch: 0n, requestedEpoch: { isSet: false, value: 0n } })], })); }); diff --git a/dstack/kms/auth-eth-bun/index.ts b/dstack/kms/auth-eth-bun/index.ts index 8355ba255..dc7763edf 100644 --- a/dstack/kms/auth-eth-bun/index.ts +++ b/dstack/kms/auth-eth-bun/index.ts @@ -37,6 +37,7 @@ const BootInfoSchema = z.object({ }); const AuthorizationRequestSchema = BootInfoSchema.omit({ rootEpoch: true, rootEpochHash: true }).extend({ + appEpoch: z.number().int().min(0).max(0xffffffff).default(0), requestedEpoch: z.number().int().min(0).max(0xffffffff).optional(), }); type AuthorizationRequest = z.infer; @@ -88,30 +89,33 @@ const APP_INFO_COMPONENTS = [ { name: 'advisoryIds', type: 'string[]' } ] as const; +const AUTHORIZATION_ABI = { + type: 'function', + stateMutability: 'view', + inputs: [{ name: 'request', type: 'tuple', components: [ + { name: 'appInfo', type: 'tuple', components: APP_INFO_COMPONENTS }, + { name: 'appEpoch', type: 'uint256' }, + { name: 'requestedEpoch', type: 'tuple', components: [ + { name: 'isSet', type: 'bool' }, { name: 'value', type: 'uint256' }, + ] }, + ] }], + outputs: [{ + name: 'result', + type: 'tuple', + components: [ + { name: 'isAllowed', type: 'bool' }, + { name: 'rejectionReason', type: 'string' }, + { name: 'gatewayAppId', type: 'string' }, + { name: 'currentRootEpoch', type: 'uint256' }, + { name: 'currentRootEpochHash', type: 'bytes32' }, + ], + }], +} as const; + // DstackKms contract ABI (minimal required functions) const DSTACK_KMS_ABI = [ - { - name: 'checkAuthorization', - type: 'function', - stateMutability: 'view', - inputs: [{ name: 'request', type: 'tuple', components: [ - { name: 'appInfo', type: 'tuple', components: APP_INFO_COMPONENTS }, - { name: 'requestedEpoch', type: 'tuple', components: [ - { name: 'isSet', type: 'bool' }, { name: 'value', type: 'uint256' }, - ] }, - ] }], - outputs: [{ - name: 'result', - type: 'tuple', - components: [ - { name: 'isAllowed', type: 'bool' }, - { name: 'rejectionReason', type: 'string' }, - { name: 'gatewayAppId', type: 'string' }, - { name: 'currentRootEpoch', type: 'uint256' }, - { name: 'currentRootEpochHash', type: 'bytes32' }, - ], - }], - }, + { ...AUTHORIZATION_ABI, name: 'checkKmsAuthorization' }, + { ...AUTHORIZATION_ABI, name: 'checkAppAuthorization' }, { name: 'isKmsAllowed', type: 'function', @@ -233,14 +237,15 @@ class EthereumBackend { } - async checkAuthorization(appInfo: AuthorizationRequest) { + async checkAuthorization(appInfo: AuthorizationRequest, method: 'checkKmsAuthorization' | 'checkAppAuthorization') { const blockNumber = await this.finalizedBlockNumber(); const result = await this.client.readContract({ address: this.kmsContractAddr, abi: DSTACK_KMS_ABI, - functionName: 'checkAuthorization', + functionName: method, args: [{ appInfo: this.contractAppInfo(appInfo), + appEpoch: BigInt(appInfo.appEpoch), requestedEpoch: { isSet: appInfo.requestedEpoch !== undefined, value: BigInt(appInfo.requestedEpoch ?? 0) }, }], blockNumber, @@ -412,18 +417,20 @@ app.get('/', async (c) => { } }); -app.post('/checkAuthorization', - bodyLimit({ maxSize: MAX_BOOT_INFO_BYTES }), - zValidator('json', AuthorizationRequestSchema), - async (c) => { - try { - return c.json(await ethereum.checkAuthorization(c.req.valid('json'))); - } catch { - console.error('authorization backend failed'); - return c.json({ status: 'error', message: backendUnavailable }, 503); +for (const method of ['checkKmsAuthorization', 'checkAppAuthorization'] as const) { + app.post(`/${method}`, + bodyLimit({ maxSize: MAX_BOOT_INFO_BYTES }), + zValidator('json', AuthorizationRequestSchema), + async (c) => { + try { + return c.json(await ethereum.checkAuthorization(c.req.valid('json'), method)); + } catch { + console.error('authorization backend failed'); + return c.json({ status: 'error', message: backendUnavailable }, 503); + } } - } -); + ); +} // app boot authentication app.post('/bootAuth/app', diff --git a/dstack/kms/auth-eth-bun/openapi.json b/dstack/kms/auth-eth-bun/openapi.json index 5425de88f..f07781e1c 100644 --- a/dstack/kms/auth-eth-bun/openapi.json +++ b/dstack/kms/auth-eth-bun/openapi.json @@ -99,10 +99,41 @@ } } }, - "/checkAuthorization": { + "/checkKmsAuthorization": { "post": { - "summary": "Authorize an attested identity and return the current root epoch", - "operationId": "checkAuthorization", + "summary": "Authorize a Kms identity and return the current root epoch", + "operationId": "checkKmsAuthorization", + "requestBody": { + "required": true, + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/AuthorizationRequest" + } + } + } + }, + "responses": { + "200": { + "description": "Authorization decision, including the epoch on rejection.", + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/AuthorizationResult" + } + } + } + }, + "503": { + "description": "Authorization backend unavailable; no legacy fallback." + } + } + } + }, + "/checkAppAuthorization": { + "post": { + "summary": "Authorize a App identity and return the current root epoch", + "operationId": "checkAppAuthorization", "requestBody": { "required": true, "content": { @@ -384,6 +415,13 @@ "minimum": 0, "maximum": 4294967295, "description": "Omit to use the current epoch. An explicit zero selects epoch 0; future epochs are denied." + }, + "appEpoch": { + "type": "integer", + "minimum": 0, + "maximum": 4294967295, + "default": 0, + "description": "Application epoch, independent of KMS root epochs. Only zero is currently supported." } } } diff --git a/dstack/kms/auth-eth/README.md b/dstack/kms/auth-eth/README.md index af26e2455..2580aea62 100644 --- a/dstack/kms/auth-eth/README.md +++ b/dstack/kms/auth-eth/README.md @@ -172,11 +172,13 @@ anvil Documentation: https://book.getfoundry.sh/ -### Unified authorization +### Authorization endpoints -`POST /checkAuthorization` returns `AuthorizationResult`, including +`POST /checkKmsAuthorization` and `POST /checkAppAuthorization` return `AuthorizationResult`, including `rejectionReason`, `gatewayAppId`, `currentRootEpoch`, and `currentRootEpochHash`. -It selects KMS or app policy from the verified `appId`. The optional +The selected endpoint determines the policy; the backend checks that the verified +`appId` matches that scope. `appEpoch` defaults to 0; nonzero values are currently +rejected. The optional `requestedEpoch` 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/DstackKms.sol b/dstack/kms/auth-eth/contracts/DstackKms.sol index d750f71b4..806c029de 100644 --- a/dstack/kms/auth-eth/contracts/DstackKms.sol +++ b/dstack/kms/auth-eth/contracts/DstackKms.sol @@ -33,6 +33,7 @@ contract DstackKms is Initializable, Ownable2StepUpgradeable, UUPSUpgradeable, E struct AuthorizationRequest { AppInfo appInfo; + uint256 appEpoch; OptionalEpoch requestedEpoch; } @@ -401,10 +402,29 @@ contract DstackKms is Initializable, Ownable2StepUpgradeable, UUPSUpgradeable, E } /// Authorize an attested identity and return the root epoch from the same state. - /// The KMS app ID is this contract's proxy address; all other IDs use app policy. - function checkAuthorization(AuthorizationRequest calldata request) + /// The caller selects KMS or app policy; KMS authorization requires the proxy app ID. + function checkKmsAuthorization(AuthorizationRequest calldata request) external view + returns (AuthorizationResult memory) + { + return _checkAuthorization(request, true); + } + + function checkAppAuthorization(AuthorizationRequest calldata request) + external + view + returns (AuthorizationResult memory) + { + return _checkAuthorization(request, false); + } + + function _checkAuthorization( + AuthorizationRequest calldata request, + bool isKms + ) + private + view returns (AuthorizationResult memory result) { result.gatewayAppId = gatewayAppId; @@ -418,7 +438,15 @@ contract DstackKms is Initializable, Ownable2StepUpgradeable, UUPSUpgradeable, E result.rejectionReason = "Requested root epoch is not active"; return result; } + if (request.appEpoch != 0) { + result.rejectionReason = "App epoch is not supported"; + return result; + } AppInfo calldata appInfo = request.appInfo; + if (isKms != (appInfo.appId == address(this))) { + result.rejectionReason = "Authorization kind does not match app ID"; + return result; + } AppBootInfo memory bootInfo = AppBootInfo({ appId: appInfo.appId, composeHash: appInfo.composeHash, @@ -430,7 +458,7 @@ contract DstackKms is Initializable, Ownable2StepUpgradeable, UUPSUpgradeable, E tcbStatus: appInfo.tcbStatus, advisoryIds: appInfo.advisoryIds }); - if (appInfo.appId == address(this)) { + if (isKms) { (result.isAllowed, result.rejectionReason) = _isKmsAllowed(bootInfo); } else { (result.isAllowed, result.rejectionReason) = _isAppAllowed(bootInfo); diff --git a/dstack/kms/auth-eth/src/ethereum.ts b/dstack/kms/auth-eth/src/ethereum.ts index c8451b824..20dd03fbe 100644 --- a/dstack/kms/auth-eth/src/ethereum.ts +++ b/dstack/kms/auth-eth/src/ethereum.ts @@ -7,7 +7,8 @@ import { BootInfo, BootResponse, RootEpochInfo, AuthorizationResult, Authorizati // Minimal ABI for DstackKms contract const DSTACK_KMS_ABI = [ - "function checkAuthorization(((address appId,bytes32 composeHash,address instanceId,bytes32 deviceId,bytes32 mrAggregated,bytes32 mrSystem,bytes32 osImageHash,string tcbStatus,string[] advisoryIds) appInfo,(bool isSet,uint256 value) requestedEpoch) request) view returns ((bool isAllowed,string rejectionReason,string gatewayAppId,uint256 currentRootEpoch,bytes32 currentRootEpochHash))", + "function checkKmsAuthorization(((address appId,bytes32 composeHash,address instanceId,bytes32 deviceId,bytes32 mrAggregated,bytes32 mrSystem,bytes32 osImageHash,string tcbStatus,string[] advisoryIds) appInfo,uint256 appEpoch,(bool isSet,uint256 value) requestedEpoch) request) view returns ((bool isAllowed,string rejectionReason,string gatewayAppId,uint256 currentRootEpoch,bytes32 currentRootEpochHash))", + "function checkAppAuthorization(((address appId,bytes32 composeHash,address instanceId,bytes32 deviceId,bytes32 mrAggregated,bytes32 mrSystem,bytes32 osImageHash,string tcbStatus,string[] advisoryIds) appInfo,uint256 appEpoch,(bool isSet,uint256 value) requestedEpoch) request) view returns ((bool isAllowed,string rejectionReason,string gatewayAppId,uint256 currentRootEpoch,bytes32 currentRootEpochHash))", "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)", "function isKmsAllowed((address appId,bytes32 composeHash,address instanceId,bytes32 deviceId,bytes32 mrAggregated,bytes32 mrSystem,bytes32 osImageHash,string tcbStatus,string[] advisoryIds) bootInfo) view returns (bool, string)", "function gatewayAppId() view returns (string)", @@ -77,9 +78,10 @@ export class EthereumBackend { }; } - async checkAuthorization(appInfo: AuthorizationRequest): Promise { - const result = await this.kmsContract.checkAuthorization({ + async checkAuthorization(appInfo: AuthorizationRequest, method: 'checkKmsAuthorization' | 'checkAppAuthorization'): Promise { + const result = await this.kmsContract[method]({ appInfo: this.contractAppInfo(appInfo), + appEpoch: appInfo.appEpoch ?? 0, requestedEpoch: { isSet: appInfo.requestedEpoch !== undefined, value: appInfo.requestedEpoch ?? 0 }, }); return { diff --git a/dstack/kms/auth-eth/src/main.test.ts b/dstack/kms/auth-eth/src/main.test.ts index a9e044e99..679a229f9 100644 --- a/dstack/kms/auth-eth/src/main.test.ts +++ b/dstack/kms/auth-eth/src/main.test.ts @@ -91,25 +91,28 @@ describe('Server', () => { }); }); - describe('POST /checkAuthorization', () => { + describe('POST /checkAppAuthorization', () => { + const appRequest = { ...mockBootInfo, appEpoch: 0 }; it('preserves the complete decision, including the epoch on rejection', async () => { const decision = { isAllowed: false, rejectionReason: 'App not registered', gatewayAppId: 'gateway', currentRootEpoch: 1, currentRootEpochHash: '0x' + '11'.repeat(32), }; app.ethereum.checkAuthorization = jest.fn().mockResolvedValue(decision); - const response = await app.inject({ method: 'POST', url: '/checkAuthorization', payload: mockBootInfo }); + const response = await app.inject({ method: 'POST', url: '/checkAppAuthorization', payload: appRequest }); expect(response.statusCode).toBe(200); expect(JSON.parse(response.payload)).toEqual(decision); - expect(app.ethereum.checkAuthorization).toHaveBeenCalledWith(mockBootInfo); - await app.inject({ method: 'POST', url: '/checkAuthorization', payload: { ...mockBootInfo, requestedEpoch: 0 } }); - expect(app.ethereum.checkAuthorization).toHaveBeenLastCalledWith({ ...mockBootInfo, requestedEpoch: 0 }); + expect(app.ethereum.checkAuthorization).toHaveBeenCalledWith(appRequest, 'checkAppAuthorization'); + await app.inject({ method: 'POST', url: '/checkKmsAuthorization', payload: appRequest }); + expect(app.ethereum.checkAuthorization).toHaveBeenLastCalledWith(appRequest, 'checkKmsAuthorization'); + await app.inject({ method: 'POST', url: '/checkAppAuthorization', payload: { ...appRequest, requestedEpoch: 0 } }); + expect(app.ethereum.checkAuthorization).toHaveBeenLastCalledWith({ ...appRequest, requestedEpoch: 0 }, 'checkAppAuthorization'); }); it('fails closed when the new contract method is unavailable', async () => { jest.clearAllMocks(); app.ethereum.checkAuthorization = jest.fn().mockRejectedValue(new Error('missing contract method')); - const response = await app.inject({ method: 'POST', url: '/checkAuthorization', payload: mockBootInfo }); + const response = await app.inject({ method: 'POST', url: '/checkAppAuthorization', payload: appRequest }); expect(response.statusCode).toBe(500); expect(app.ethereum.checkBoot).not.toHaveBeenCalled(); }); diff --git a/dstack/kms/auth-eth/src/server.ts b/dstack/kms/auth-eth/src/server.ts index 17deb15b5..081c56f0e 100644 --- a/dstack/kms/auth-eth/src/server.ts +++ b/dstack/kms/auth-eth/src/server.ts @@ -55,6 +55,7 @@ export async function build(): Promise { $id: 'authorizationRequest', properties: { ...appInfoSchema.properties, + appEpoch: { type: 'integer', minimum: 0, maximum: 4294967295, default: 0 }, requestedEpoch: { type: 'integer', minimum: 0, maximum: 4294967295 }, }, }); @@ -119,24 +120,26 @@ export async function build(): Promise { } }); - server.post<{ Body: AuthorizationRequest; Reply: AuthorizationResult }>('/checkAuthorization', { - schema: { - body: { $ref: 'authorizationRequest#' }, - response: { - 200: { - type: 'object', - required: ['isAllowed', 'rejectionReason', 'gatewayAppId', 'currentRootEpoch', 'currentRootEpochHash'], - properties: { - isAllowed: { type: 'boolean' }, - rejectionReason: { type: 'string' }, - gatewayAppId: { type: 'string' }, - currentRootEpoch: { type: 'integer', minimum: 0 }, - currentRootEpochHash: hex(32, 'Current root epoch chain hash'), + for (const method of ['checkKmsAuthorization', 'checkAppAuthorization'] as const) { + server.post<{ Body: AuthorizationRequest; Reply: AuthorizationResult }>(`/${method}`, { + schema: { + body: { $ref: 'authorizationRequest#' }, + response: { + 200: { + type: 'object', + required: ['isAllowed', 'rejectionReason', 'gatewayAppId', 'currentRootEpoch', 'currentRootEpochHash'], + properties: { + isAllowed: { type: 'boolean' }, + rejectionReason: { type: 'string' }, + gatewayAppId: { type: 'string' }, + currentRootEpoch: { type: 'integer', minimum: 0 }, + currentRootEpochHash: hex(32, 'Current root epoch chain hash'), + }, }, }, }, - }, - }, async (request) => server.ethereum.checkAuthorization(request.body)); + }, async (request) => server.ethereum.checkAuthorization(request.body, method)); + } // Define routes server.post<{ diff --git a/dstack/kms/auth-eth/src/types.ts b/dstack/kms/auth-eth/src/types.ts index 1b43e3fac..40d781acf 100644 --- a/dstack/kms/auth-eth/src/types.ts +++ b/dstack/kms/auth-eth/src/types.ts @@ -42,5 +42,6 @@ export interface AuthorizationResult { } export interface AuthorizationRequest extends Omit { + appEpoch?: number; requestedEpoch?: number; } diff --git a/dstack/kms/auth-eth/test/DstackKms.t.sol b/dstack/kms/auth-eth/test/DstackKms.t.sol index a6635a5ca..0eb2d3f1e 100644 --- a/dstack/kms/auth-eth/test/DstackKms.t.sol +++ b/dstack/kms/auth-eth/test/DstackKms.t.sol @@ -456,8 +456,17 @@ contract DstackKmsTest is Test { }); } - function authorize(DstackKms.AppInfo memory info) internal view returns (DstackKms.AuthorizationResult memory) { - return kms.checkAuthorization(DstackKms.AuthorizationRequest(info, DstackKms.OptionalEpoch(false, 0))); + function authorize( + DstackKms.AppInfo memory info, + bool isKms + ) + internal + view + returns (DstackKms.AuthorizationResult memory) + { + DstackKms.AuthorizationRequest memory request = + DstackKms.AuthorizationRequest(info, 0, DstackKms.OptionalEpoch(false, 0)); + return isKms ? kms.checkKmsAuthorization(request) : kms.checkAppAuthorization(request); } function test_CheckAuthorization_OptionalEpochDistinguishesZeroFromDefault() public { @@ -471,34 +480,34 @@ contract DstackKmsTest is Test { vm.stopPrank(); DstackKms.AuthorizationRequest memory request = - DstackKms.AuthorizationRequest(info, DstackKms.OptionalEpoch(true, 1)); - DstackKms.AuthorizationResult memory result = kms.checkAuthorization(request); + DstackKms.AuthorizationRequest(info, 0, DstackKms.OptionalEpoch(true, 1)); + DstackKms.AuthorizationResult memory result = kms.checkKmsAuthorization(request); assertFalse(result.isAllowed); assertEq(result.rejectionReason, "Requested root epoch is not active"); bytes32 h1 = chainHash(chainHash(bytes32(0), ca(0xa0), k256(0xb0)), ca(0xa1), k256(0xb1)); vm.prank(owner); kms.activateRootEpoch(1, h1); - assertTrue(kms.checkAuthorization(request).isAllowed); + assertTrue(kms.checkKmsAuthorization(request).isAllowed); request.requestedEpoch.value = 0; - assertTrue(kms.checkAuthorization(request).isAllowed); + assertTrue(kms.checkKmsAuthorization(request).isAllowed); request.requestedEpoch.value = type(uint256).max; - assertFalse(kms.checkAuthorization(request).isAllowed); + assertFalse(kms.checkKmsAuthorization(request).isAllowed); request.requestedEpoch.isSet = false; - result = kms.checkAuthorization(request); + result = kms.checkKmsAuthorization(request); assertTrue(result.isAllowed); assertEq(result.currentRootEpoch, 1); assertEq(result.currentRootEpochHash, h1); } function test_CheckAuthorization_UnregisteredEpochReturnsRejection() public view { - DstackKms.AuthorizationResult memory result = authorize(authorizationInfo(address(kms))); + DstackKms.AuthorizationResult memory result = authorize(authorizationInfo(address(kms)), true); assertFalse(result.isAllowed); assertEq(result.rejectionReason, "Root epoch not registered"); assertEq(result.currentRootEpoch, 0); assertEq(result.currentRootEpochHash, bytes32(0)); } - function test_CheckAuthorization_SelectsPolicyByAppId() public { + function test_CheckAuthorization_ValidatesRequestedPolicy() public { DstackKms.AppInfo memory info = authorizationInfo(address(kms)); vm.startPrank(owner); kms.registerRootEpoch(ca(0xa0), k256(0xb0), "https://example.com/epoch-0"); @@ -509,23 +518,35 @@ contract DstackKmsTest is Test { address appId = kms.deployAndRegisterApp(user, false, false, info.deviceId, info.composeHash); vm.stopPrank(); - DstackKms.AuthorizationResult memory result = authorize(info); + DstackKms.AuthorizationResult memory result = authorize(info, true); assertTrue(result.isAllowed); assertEq(result.rejectionReason, ""); assertEq(result.gatewayAppId, "gateway"); assertEq(result.currentRootEpochHash, chainHash(bytes32(0), ca(0xa0), k256(0xb0))); info.appId = user; - result = authorize(info); + result = authorize(info, false); assertFalse(result.isAllowed); assertEq(result.rejectionReason, "App not registered"); info.appId = appId; info.mrAggregated = bytes32(0); - result = authorize(info); + result = authorize(info, false); assertTrue(result.isAllowed); + result = authorize(info, true); + assertFalse(result.isAllowed); + assertEq(result.rejectionReason, "Authorization kind does not match app ID"); + DstackKms.AuthorizationRequest memory appRequest = + DstackKms.AuthorizationRequest(info, 1, DstackKms.OptionalEpoch(false, 0)); + result = kms.checkAppAuthorization(appRequest); + assertFalse(result.isAllowed); + assertEq(result.rejectionReason, "App epoch is not supported"); + assertEq(result.currentRootEpochHash, chainHash(bytes32(0), ca(0xa0), k256(0xb0))); info.appId = address(kms); - result = authorize(info); + result = authorize(info, false); + assertFalse(result.isAllowed); + assertEq(result.rejectionReason, "Authorization kind does not match app ID"); + result = authorize(info, true); assertFalse(result.isAllowed); assertEq(result.rejectionReason, "Aggregated MR not allowed"); } @@ -537,7 +558,7 @@ contract DstackKmsTest is Test { kms.activateRootEpoch(1, chainHash(chainHash(bytes32(0), ca(0xa0), k256(0xb0)), ca(0xa1), k256(0xb1))); vm.stopPrank(); - DstackKms.AuthorizationResult memory result = authorize(authorizationInfo(user)); + DstackKms.AuthorizationResult memory result = authorize(authorizationInfo(user), false); assertFalse(result.isAllowed); assertEq(result.rejectionReason, "App not registered"); assertEq(result.currentRootEpoch, 1); diff --git a/dstack/kms/auth-mock/README.md b/dstack/kms/auth-mock/README.md index 78585d661..3bd731747 100644 --- a/dstack/kms/auth-mock/README.md +++ b/dstack/kms/auth-mock/README.md @@ -169,11 +169,13 @@ to switch from mock to real backend: the API remains identical, so no client code changes are needed. -### Unified authorization +### Authorization endpoints -`POST /checkAuthorization` returns `AuthorizationResult`, including +`POST /checkKmsAuthorization` and `POST /checkAppAuthorization` return `AuthorizationResult`, including `rejectionReason`, `gatewayAppId`, `currentRootEpoch`, and `currentRootEpochHash`. -It selects KMS or app policy from the verified `appId`. The optional +The selected endpoint determines the policy; the backend checks that the verified +`appId` matches that scope. `appEpoch` defaults to 0; nonzero values are currently +rejected. The optional `requestedEpoch` 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.test.ts b/dstack/kms/auth-mock/index.test.ts index 3336d4f5f..6f6858fe1 100644 --- a/dstack/kms/auth-mock/index.test.ts +++ b/dstack/kms/auth-mock/index.test.ts @@ -384,10 +384,10 @@ describe('Mock Behavior Verification', () => { }); -describe('POST /checkAuthorization', () => { - const authorize = (appId: string) => appFetch(new Request('http://localhost/checkAuthorization', { +describe('authorization endpoints', () => { + const authorize = (appId: string, isKms = false, appEpoch = 0) => appFetch(new Request(`http://localhost/${isKms ? 'checkKmsAuthorization' : 'checkAppAuthorization'}`, { method: 'POST', headers: { 'Content-Type': 'application/json' }, - body: JSON.stringify({ appId, mrAggregated: '', osImageHash: '', composeHash: '', instanceId: '', deviceId: '' }), + body: JSON.stringify({ appId, appEpoch, mrAggregated: '', osImageHash: '', composeHash: '', instanceId: '', deviceId: '' }), })); it('requires registered roots and selects policy by app ID', async () => { @@ -402,7 +402,7 @@ describe('POST /checkAuthorization', () => { vi.stubEnv('MOCK_ROOT_EPOCH', '1'); vi.stubEnv('MOCK_ROOT_EPOCH_HASH', '0x' + '11'.repeat(32)); expect(await (await authorize('app')).json()).toMatchObject({ isAllowed: true, rejectionReason: '', currentRootEpoch: 1 }); - expect(await (await authorize(process.env.KMS_CONTRACT_ADDR!)).json()).toMatchObject({ + expect(await (await authorize(process.env.KMS_CONTRACT_ADDR!, true)).json()).toMatchObject({ isAllowed: false, rejectionReason: 'mock policy: deny-kms', currentRootEpoch: 1, currentRootEpochHash: '0x' + '11'.repeat(32), }); diff --git a/dstack/kms/auth-mock/index.ts b/dstack/kms/auth-mock/index.ts index dc9f3048f..4c067189f 100644 --- a/dstack/kms/auth-mock/index.ts +++ b/dstack/kms/auth-mock/index.ts @@ -30,6 +30,7 @@ const BootInfoSchema = z.object({ }); const AuthorizationRequestSchema = BootInfoSchema.omit({ rootEpoch: true, rootEpochHash: true }).extend({ + appEpoch: z.number().int().min(0).max(0xffffffff).default(0), requestedEpoch: z.number().int().min(0).max(0xffffffff).optional(), }); type AuthorizationRequest = z.infer; @@ -111,7 +112,7 @@ class MockBackend { this.mockAppImplementation = process.env.MOCK_APP_IMPLEMENTATION || '0xmockapp9876543210987654321098765432109'; } - async checkAuthorization(appInfo: AuthorizationRequest) { + async checkAuthorization(appInfo: AuthorizationRequest, isKms: boolean) { const current = getRootEpoch(); const epoch = { currentRootEpoch: current.rootEpoch ?? 0, @@ -123,8 +124,14 @@ class MockBackend { if (appInfo.requestedEpoch !== undefined && appInfo.requestedEpoch > current.rootEpoch) { return { isAllowed: false, rejectionReason: 'Requested root epoch is not active', gatewayAppId: this.mockGatewayAppId, ...epoch }; } + if (appInfo.appEpoch !== 0) { + return { isAllowed: false, rejectionReason: 'App epoch is not supported', gatewayAppId: this.mockGatewayAppId, ...epoch }; + } const kmsAppId = process.env.KMS_CONTRACT_ADDR || '0xmockcontract1234567890123456789012345678'; - const result = this.evaluateBoot(appInfo, normHash(appInfo.appId) === normHash(kmsAppId)); + if (isKms !== (normHash(appInfo.appId) === normHash(kmsAppId))) { + return { isAllowed: false, rejectionReason: 'Authorization kind does not match app ID', gatewayAppId: this.mockGatewayAppId, ...epoch }; + } + const result = this.evaluateBoot(appInfo, isKms); return { isAllowed: result.isAllowed, rejectionReason: result.isAllowed ? '' : result.reason, gatewayAppId: this.mockGatewayAppId, ...epoch }; } @@ -225,13 +232,15 @@ app.get('/', async (c) => { } }); -app.post('/checkAuthorization', zValidator('json', AuthorizationRequestSchema), async (c) => { - try { - return c.json(await mockBackend.checkAuthorization(c.req.valid('json'))); - } catch { - return c.json({ status: 'error', message: 'authorization backend unavailable' }, 503); - } -}); +for (const method of ['checkKmsAuthorization', 'checkAppAuthorization'] as const) { + app.post(`/${method}`, zValidator('json', AuthorizationRequestSchema), async (c) => { + try { + return c.json(await mockBackend.checkAuthorization(c.req.valid('json'), method === 'checkKmsAuthorization')); + } catch { + return c.json({ status: 'error', message: 'authorization backend unavailable' }, 503); + } + }); +} // app boot authentication app.post('/bootAuth/app', diff --git a/dstack/kms/auth-mock/openapi.json b/dstack/kms/auth-mock/openapi.json index 03d5c95fa..6f862157c 100644 --- a/dstack/kms/auth-mock/openapi.json +++ b/dstack/kms/auth-mock/openapi.json @@ -99,10 +99,41 @@ } } }, - "/checkAuthorization": { + "/checkKmsAuthorization": { "post": { - "summary": "Authorize an attested identity and return the current root epoch", - "operationId": "checkAuthorization", + "summary": "Authorize a Kms identity and return the current root epoch", + "operationId": "checkKmsAuthorization", + "requestBody": { + "required": true, + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/AuthorizationRequest" + } + } + } + }, + "responses": { + "200": { + "description": "Authorization decision, including the epoch on rejection.", + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/AuthorizationResult" + } + } + } + }, + "503": { + "description": "Authorization backend unavailable; no legacy fallback." + } + } + } + }, + "/checkAppAuthorization": { + "post": { + "summary": "Authorize a App identity and return the current root epoch", + "operationId": "checkAppAuthorization", "requestBody": { "required": true, "content": { @@ -393,6 +424,13 @@ "minimum": 0, "maximum": 4294967295, "description": "Omit to use the current epoch. An explicit zero selects epoch 0; future epochs are denied." + }, + "appEpoch": { + "type": "integer", + "minimum": 0, + "maximum": 4294967295, + "default": 0, + "description": "Application epoch, independent of KMS root epochs. Only zero is currently supported." } } } diff --git a/dstack/kms/auth-simple/README.md b/dstack/kms/auth-simple/README.md index af0e99419..9caf9387e 100644 --- a/dstack/kms/auth-simple/README.md +++ b/dstack/kms/auth-simple/README.md @@ -254,11 +254,13 @@ bun run test:run - [Deployment Guide](../../../docs/deployment.md) - Full deployment instructions -### Unified authorization +### Authorization endpoints -`POST /checkAuthorization` returns `AuthorizationResult`, including +`POST /checkKmsAuthorization` and `POST /checkAppAuthorization` return `AuthorizationResult`, including `rejectionReason`, `gatewayAppId`, `currentRootEpoch`, and `currentRootEpochHash`. -It selects KMS or app policy from the verified `appId`. The optional +The selected endpoint determines the policy; the backend checks that the verified +`appId` matches that scope. `appEpoch` defaults to 0; nonzero values are currently +rejected. The optional `requestedEpoch` 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 fc3edea38..50574cdb1 100644 --- a/dstack/kms/auth-simple/index.test.ts +++ b/dstack/kms/auth-simple/index.test.ts @@ -472,13 +472,13 @@ describe('auth-simple', () => { const info = async () => (await app.fetch(new Request('http://localhost/', { method: 'GET' }))).json(); - const authorize = async (fields: object = {}) => - (await app.fetch(new Request('http://localhost/checkAuthorization', { + const authorize = async (fields: object = {}, isKms = false) => + (await app.fetch(new Request(`http://localhost/${isKms ? 'checkKmsAuthorization' : 'checkAppAuthorization'}`, { method: 'POST', headers: { 'Content-Type': 'application/json' }, body: JSON.stringify({ ...baseBootInfo, ...fields }), }))).json(); - it('selects KMS policy by app ID and returns the epoch on rejection', async () => { + it('checks the requested policy and returns the epoch on rejection', async () => { const kmsContractAddr = '0x' + '11'.repeat(20); writeTestConfig({ ...config, kmsContractAddr, currentRootEpoch: 1 }); expect(await authorize()).toMatchObject({ @@ -486,11 +486,15 @@ describe('auth-simple', () => { }); expect(await authorize({ requestedEpoch: 0 })).toMatchObject({ isAllowed: true, currentRootEpoch: 1 }); expect(await authorize({ requestedEpoch: 1 })).toMatchObject({ isAllowed: true }); + expect(await authorize({}, true)).toMatchObject({ isAllowed: false, rejectionReason: 'Authorization kind does not match app ID' }); + expect(await authorize({ appId: kmsContractAddr })).toMatchObject({ isAllowed: false, rejectionReason: 'Authorization kind does not match app ID' }); + expect(await authorize({ appEpoch: 1 })).toMatchObject({ isAllowed: false, rejectionReason: 'App epoch is not supported', currentRootEpoch: 1 }); + expect(await authorize({ appEpoch: 0 })).toMatchObject({ isAllowed: true }); expect(await authorize({ requestedEpoch: 2 })).toMatchObject({ isAllowed: false, rejectionReason: 'Requested root epoch is not active', currentRootEpoch: 1, currentRootEpochHash: h1, }); - expect(await authorize({ appId: kmsContractAddr, composeHash: 'not-an-app-compose' })).toMatchObject({ isAllowed: true }); - expect(await authorize({ appId: kmsContractAddr, mrAggregated: '0xbad' })).toMatchObject({ + expect(await authorize({ appId: kmsContractAddr, composeHash: 'not-an-app-compose' }, true)).toMatchObject({ isAllowed: true }); + expect(await authorize({ appId: kmsContractAddr, mrAggregated: '0xbad' }, true)).toMatchObject({ isAllowed: false, currentRootEpoch: 1, currentRootEpochHash: h1, }); expect(await authorize({ appId: '0x' + '22'.repeat(20) })).toMatchObject({ isAllowed: false }); diff --git a/dstack/kms/auth-simple/index.ts b/dstack/kms/auth-simple/index.ts index dbca135ed..d04297565 100644 --- a/dstack/kms/auth-simple/index.ts +++ b/dstack/kms/auth-simple/index.ts @@ -28,6 +28,7 @@ const BootInfoSchema = z.object({ }); const AuthorizationRequestSchema = BootInfoSchema.omit({ rootEpoch: true, rootEpochHash: true }).extend({ + appEpoch: z.number().int().min(0).max(0xffffffff).default(0), requestedEpoch: z.number().int().min(0).max(0xffffffff).optional(), }); type AuthorizationRequest = z.infer; @@ -157,7 +158,7 @@ class ConfigBackend { } } - async checkAuthorization(appInfo: AuthorizationRequest) { + async checkAuthorization(appInfo: AuthorizationRequest, isKms: boolean) { const config = this.loadConfig(); const current = currentRootEpoch(config); const epoch = { @@ -170,7 +171,12 @@ class ConfigBackend { if (appInfo.requestedEpoch !== undefined && appInfo.requestedEpoch > current.rootEpoch) { return { isAllowed: false, rejectionReason: 'Requested root epoch is not active', gatewayAppId: config.gatewayAppId, ...epoch }; } - const isKms = normalizeHex(appInfo.appId) === normalizeHex(config.kmsContractAddr); + if (appInfo.appEpoch !== 0) { + return { isAllowed: false, rejectionReason: 'App epoch is not supported', gatewayAppId: config.gatewayAppId, ...epoch }; + } + if (isKms !== (normalizeHex(appInfo.appId) === normalizeHex(config.kmsContractAddr))) { + return { isAllowed: false, rejectionReason: 'Authorization kind does not match app ID', gatewayAppId: config.gatewayAppId, ...epoch }; + } const result = this.evaluateBoot(appInfo, isKms, config); return { isAllowed: result.isAllowed, rejectionReason: result.reason, gatewayAppId: result.gatewayAppId, ...epoch }; } @@ -355,14 +361,16 @@ app.get('/', async (c) => { } }); -app.post('/checkAuthorization', zValidator('json', AuthorizationRequestSchema), async (c) => { - try { - return c.json(await backend.checkAuthorization(c.req.valid('json'))); - } catch { - console.error('authorization backend failed'); - return c.json({ status: 'error', message: 'authorization backend unavailable' }, 503); - } -}); +for (const method of ['checkKmsAuthorization', 'checkAppAuthorization'] as const) { + app.post(`/${method}`, zValidator('json', AuthorizationRequestSchema), async (c) => { + try { + return c.json(await backend.checkAuthorization(c.req.valid('json'), method === 'checkKmsAuthorization')); + } catch { + console.error('authorization backend failed'); + return c.json({ status: 'error', message: 'authorization backend unavailable' }, 503); + } + }); +} // app boot authentication app.post('/bootAuth/app', diff --git a/dstack/kms/src/main_service.rs b/dstack/kms/src/main_service.rs index 188145cc7..0bf9e73c9 100644 --- a/dstack/kms/src/main_service.rs +++ b/dstack/kms/src/main_service.rs @@ -291,28 +291,6 @@ impl KmsState { self.keyring.servable_epoch(info.root_epochs) } - /// KMS-only operations must not accept an app that passes ordinary app policy. - async fn ensure_kms_identity(&self, app_info: &BootInfo) -> Result<()> { - if self.config.auth_api.is_dev() { - return Ok(()); - } - let info = self - .auth_api_info - .try_get_with((), self.config.auth_api.get_info()) - .await - .map_err(|err| anyhow!("{err:#}"))?; - let address = info - .kms_contract_address - .as_deref() - .context("missing KMS contract address")?; - let app_id = hex::decode(address.trim_start_matches("0x")) - .context("invalid KMS contract address")?; - if app_info.app_id != app_id { - bail!("KMS app ID does not match the KMS contract address"); - } - Ok(()) - } - /// Resolve this request from its authorization result; no serving epoch is cached. fn authorized_epoch( &self, @@ -512,7 +490,6 @@ impl RpcHandler { .await .map_err(|err| anyhow!("{err:#}")) .context("Failed to load cached self boot info")?; - self.state.ensure_kms_identity(&boot_info).await?; let response = self .state .self_allowed @@ -521,7 +498,7 @@ impl RpcHandler { .state .config .auth_api - .check_authorization(&boot_info, None) + .check_kms_authorization(&boot_info) .await .context("Failed to call self KMS auth check")?; let epoch = self.state.authorized_epoch(&response, None)?; @@ -567,15 +544,13 @@ impl RpcHandler { requested_epoch: Option, ) -> Result { let boot_info = build_boot_info_for_attestation(att, use_boottime_mr, vm_config_str)?; - if is_kms { - self.state.ensure_kms_identity(&boot_info).await?; - } - let response = self - .state - .config - .auth_api - .check_authorization(&boot_info, requested_epoch) - .await?; + let auth = &self.state.config.auth_api; + let response = if is_kms { + auth.check_kms_authorization(&boot_info).await? + } else { + auth.check_app_authorization(&boot_info, requested_epoch) + .await? + }; let root_epoch = self.state.authorized_epoch(&response, requested_epoch)?; let current_epoch = if self.state.config.auth_api.is_dev() { self.state.keyring.latest() @@ -1564,22 +1539,6 @@ mod tests { assert!(state.epoch_root(u32::MAX).is_err()); } - #[rocket::async_test] - async fn kms_identity_cannot_be_replaced_with_an_authorized_app() { - use super::upgrade_authority::tests::serve; - let app_id = vec![0x11; 20]; - let response = AUTH_API_INFO.replace("0xkms", &format!("0x{}", hex::encode(&app_id))); - let (url, server) = serve(vec![response]); - let dir = tempfile::tempdir().unwrap(); - let state = kms_state(dir.path(), &url); - let mut info = super::upgrade_authority::tests::boot_info(1); - info.app_id = app_id; - state.ensure_kms_identity(&info).await.unwrap(); - info.app_id[0] ^= 1; - assert!(state.ensure_kms_identity(&info).await.is_err()); - assert_eq!(server.join().unwrap().len(), 1); - } - const AUTH_API_INFO: &str = r#"{"status":"ok","kmsContractAddr":"0xkms","ethRpcUrl":"https://rpc.example","gatewayAppId":"0xgateway","chainId":1,"appImplementation":"0ximpl"}"#; #[rocket::async_test] diff --git a/dstack/kms/src/main_service/upgrade_authority.rs b/dstack/kms/src/main_service/upgrade_authority.rs index d44e91a0d..a960c957f 100644 --- a/dstack/kms/src/main_service/upgrade_authority.rs +++ b/dstack/kms/src/main_service/upgrade_authority.rs @@ -254,9 +254,27 @@ async fn send_request(req: reqwest::RequestBuilder, url: &s } impl AuthApi { - pub async fn check_authorization( + pub async fn check_kms_authorization( &self, app_info: &BootInfo, + ) -> Result { + self.check_authorization(app_info, "checkKmsAuthorization", None) + .await + } + + pub async fn check_app_authorization( + &self, + app_info: &BootInfo, + requested_epoch: Option, + ) -> Result { + self.check_authorization(app_info, "checkAppAuthorization", requested_epoch) + .await + } + + async fn check_authorization( + &self, + app_info: &BootInfo, + method: &str, requested_epoch: Option, ) -> Result { #[derive(Serialize)] @@ -264,6 +282,7 @@ impl AuthApi { struct AuthorizationRequest<'a> { #[serde(flatten)] app_info: &'a BootInfo, + app_epoch: u32, #[serde(skip_serializing_if = "Option::is_none")] requested_epoch: Option, } @@ -277,11 +296,12 @@ impl AuthApi { current_root_epoch_hash: hex::encode([0; 32]), }), AuthApi::Webhook { webhook } => { - let url = url_join(&webhook.url, "checkAuthorization"); + let url = url_join(&webhook.url, method); http_post( &url, &AuthorizationRequest { app_info, + app_epoch: 0, requested_epoch, }, webhook.timeout, @@ -685,22 +705,28 @@ pub(crate) mod tests { response(true, 0), response(false, 1), response(true, 1), + response(true, 1), ]); let auth = webhook(url); let info = boot_info(50); - let first = auth.check_authorization(&info, None).await.unwrap(); - let second = auth.check_authorization(&info, None).await.unwrap(); + let first = auth.check_app_authorization(&info, None).await.unwrap(); + let second = auth.check_app_authorization(&info, 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.check_authorization(&info, Some(0)).await.unwrap(); + auth.check_app_authorization(&info, Some(0)).await.unwrap(); + auth.check_kms_authorization(&info).await.unwrap(); let requests = server.join().unwrap(); - assert_eq!(requests.len(), 3); - assert!(requests + assert_eq!(requests.len(), 4); + assert!(requests[..3] .iter() - .all(|(path, _)| path == "/checkAuthorization")); - assert_eq!(requests[0].1, serde_json::to_value(info).unwrap()); + .all(|(path, _)| path == "/checkAppAuthorization")); + assert_eq!(requests[3].0, "/checkKmsAuthorization"); + let mut expected = serde_json::to_value(info).unwrap(); + expected["appEpoch"] = 0.into(); + assert_eq!(requests[0].1, expected); + assert_eq!(requests[2].1["appEpoch"], 0); assert!(requests[0].1.get("requestedEpoch").is_none()); assert_eq!(requests[2].1["requestedEpoch"], 0); } @@ -711,7 +737,7 @@ pub(crate) mod tests { r#"{"isAllowed":true,"reason":"","gatewayAppId":"gw"}"#, ]); assert!(webhook(url) - .check_authorization(&boot_info(50), None) + .check_app_authorization(&boot_info(50), None) .await .is_err()); assert_eq!(server.join().unwrap().len(), 1); From 91934bda83bd3ed1a0450f7a34c950606b7556dd Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Fri, 6 Nov 2026 23:23:08 -0800 Subject: [PATCH 021/166] fix(cert-client): keep KMS CA pins independent of root epoch selection Signed-off-by: Kevin Wang --- dstack/cert-client/src/lib.rs | 116 +++++++++++++++++++-------------- dstack/dstack-util/src/main.rs | 4 +- dstack/kms/README.md | 8 ++- 3 files changed, 75 insertions(+), 53 deletions(-) diff --git a/dstack/cert-client/src/lib.rs b/dstack/cert-client/src/lib.rs index 431d94ad9..6c52588bf 100644 --- a/dstack/cert-client/src/lib.rs +++ b/dstack/cert-client/src/lib.rs @@ -50,10 +50,10 @@ pub fn validate_kms_rpc_cert(cert: Option) -> Result<()> { } /// Fetch keys directly over the attested connection, without a metadata/CA preflight. -/// Historical requests must pin the already trusted current KMS CA, since their -/// response CA belongs to the requested epoch rather than the TLS listener. -/// Without a supplied CA, the caller must still bind the returned root identity -/// through its provider pin or launch-measurement policy. +/// A supplied CA always pins the connection, independently of the epoch selector. +/// Callers retain the startup KMS CA for the CVM lifetime, including later key +/// requests. Without a supplied CA, bind the response CA to this connection and +/// establish its identity through the provider pin or launch-measurement policy. pub async fn request_app_keys( url: String, request: GetAppKeyRequest, @@ -61,9 +61,6 @@ pub async fn request_app_keys( attestation_verifier: Arc, trusted_ca: Option, ) -> Result { - if request.root_epoch.is_some() && trusted_ca.is_none() { - bail!("historical key requests require the trusted KMS CA"); - } let server_cert = Arc::new(Mutex::new(None)); let captured = server_cert.clone(); let client = RaClientConfig::builder() @@ -98,13 +95,7 @@ pub async fn request_app_keys( .map_err(|_| anyhow::anyhow!("KMS certificate lock poisoned"))?; validate_kms_root( cert.as_deref().context("missing KMS server certificate")?, - if requested_epoch.is_some() { - trusted_ca - .as_deref() - .context("historical key requests require the trusted KMS CA")? - } else { - &response.ca_cert - }, + trusted_ca.as_deref().unwrap_or(&response.ca_cert), )?; Ok(response) } @@ -577,46 +568,73 @@ mod tests { assert_eq!(result.ca_cert, response.ca_cert); assert_eq!(result.root_epoch, 1); - let wrong_root = AppKeyResponse { - ca_cert: kms_root().pem_cert, - ..response.clone() - }; - *response_body.lock().unwrap() = serde_json::to_string(&wrong_root).unwrap(); - assert!(request_app_keys( - url.clone(), - request(), - client_cert(), - verifier.clone(), - None - ) - .await - .is_err()); + let other_ca = kms_root().pem_cert; + for requested_epoch in [None, Some(0), Some(1)] { + let request = || GetAppKeyRequest { + api_version: 1, + vm_config: String::new(), + root_epoch: requested_epoch, + }; + let selected = AppKeyResponse { + root_epoch: requested_epoch.unwrap_or(1), + ..response.clone() + }; + *response_body.lock().unwrap() = serde_json::to_string(&selected).unwrap(); + // Omission and explicit epoch selectors have the same trust semantics. + assert!(request_app_keys( + url.clone(), + request(), + client_cert(), + verifier.clone(), + None + ) + .await + .is_ok()); + // An unrelated pin must fail in TLS, before the RPC handler runs. + assert!(request_app_keys( + url.clone(), + request(), + client_cert(), + verifier.clone(), + Some(other_ca.clone()) + ) + .await + .is_err()); + + let different_epoch_ca = AppKeyResponse { + ca_cert: other_ca.clone(), + ..selected + }; + *response_body.lock().unwrap() = serde_json::to_string(&different_epoch_ca).unwrap(); + // A key response cannot replace the startup CA pin, even for default requests. + let result = request_app_keys( + url.clone(), + request(), + client_cert(), + verifier.clone(), + Some(root.pem_cert.clone()), + ) + .await + .unwrap(); + assert_eq!(result.ca_cert, other_ca); + // First-contact validation still requires the response CA to sign the RPC cert. + assert!(request_app_keys( + url.clone(), + request(), + client_cert(), + verifier.clone(), + None + ) + .await + .is_err()); + } *response_body.lock().unwrap() = serde_json::to_string(&response).unwrap(); - - let mut historical_request = request(); - historical_request.root_epoch = Some(1); - assert!(request_app_keys( - url.clone(), - historical_request, - client_cert(), - verifier.clone(), - Some(root.pem_cert.clone()) - ) - .await - .is_ok()); - let production = Arc::new(AttestationVerifier::new_prod(None).unwrap()); assert!( - request_app_keys(url.clone(), request(), client_cert(), production, None) + request_app_keys(url, request(), client_cert(), production, None) .await .is_err() ); - let mut historical = request(); - historical.root_epoch = Some(0); - let err = request_app_keys(url, historical, client_cert(), verifier, None) - .await - .unwrap_err(); - assert!(err.to_string().contains("require the trusted KMS CA")); shutdown.notify(); let server = task.await.unwrap().unwrap(); assert_eq!( @@ -625,7 +643,7 @@ mod tests { .unwrap() .1 .load(std::sync::atomic::Ordering::SeqCst), - 4 + 11 ); } diff --git a/dstack/dstack-util/src/main.rs b/dstack/dstack-util/src/main.rs index e135bdf64..f081b3235 100644 --- a/dstack/dstack-util/src/main.rs +++ b/dstack/dstack-util/src/main.rs @@ -661,11 +661,11 @@ async fn cmd_get_keys(args: GetKeysArgs) -> Result<()> { }, cert_pair, std::sync::Arc::new(ra_tls::attestation::AttestationVerifier::new_prod(None)?), - root_ca_pem, + root_ca_pem.clone(), ) .await .context("Failed to get app keys from the attested KMS")?; - let root_ca = response.ca_cert; + let root_ca = root_ca_pem.unwrap_or(response.ca_cert); // Build AppKeys from the authenticated response. let (_, ca_pem) = diff --git a/dstack/kms/README.md b/dstack/kms/README.md index 1e0d4f254..e1a857cea 100644 --- a/dstack/kms/README.md +++ b/dstack/kms/README.md @@ -331,5 +331,9 @@ Onboarding uses the same self-issued client flow with or without an admin token, and checks the source attestation captured from the key-carrying connection against local KMS policy before importing anything. Retries and reconnects repeat verification on the response connection; no assumption relies on one TCP session -lasting forever. Historical guest key requests use the already trusted current -KMS CA for TLS, not the historical response CA. +lasting forever. All subsequent key requests pin the KMS CA established at startup, regardless +of whether the epoch selector is omitted or explicit. A supplied `trusted_ca` +always wins over the response CA. The selected epoch's CA never replaces this +lifetime pin; a TLS CA change requires a fresh boot/trust establishment rather +than silently trusting another KMS. Without a supplied pin, the response CA +must sign the actual RPC certificate, even when an epoch is explicitly selected. From f8dff0a9c6d159be08f7a7c6afded06b7c0ac1aa Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Sat, 7 Nov 2026 00:36:43 -0800 Subject: [PATCH 022/166] refactor(kms): name requested root and app epochs explicitly Signed-off-by: Kevin Wang --- docs/deployment.md | 2 +- dstack/cert-client/src/lib.rs | 10 ++++----- dstack/kms/README.md | 6 ++--- dstack/kms/auth-eth-bun/README.md | 4 ++-- dstack/kms/auth-eth-bun/index.test.ts | 8 +++---- dstack/kms/auth-eth-bun/index.ts | 12 +++++----- dstack/kms/auth-eth-bun/openapi.json | 4 ++-- dstack/kms/auth-eth/README.md | 4 ++-- dstack/kms/auth-eth/contracts/DstackKms.sol | 11 +++++----- dstack/kms/auth-eth/src/ethereum.ts | 8 +++---- dstack/kms/auth-eth/src/main.test.ts | 6 ++--- dstack/kms/auth-eth/src/server.ts | 4 ++-- dstack/kms/auth-eth/src/types.ts | 4 ++-- dstack/kms/auth-eth/test/DstackKms.t.sol | 8 +++---- dstack/kms/auth-mock/README.md | 4 ++-- dstack/kms/auth-mock/index.test.ts | 4 ++-- dstack/kms/auth-mock/index.ts | 10 ++++----- dstack/kms/auth-mock/openapi.json | 4 ++-- dstack/kms/auth-simple/README.md | 4 ++-- dstack/kms/auth-simple/index.test.ts | 10 ++++----- dstack/kms/auth-simple/index.ts | 10 ++++----- dstack/kms/src/main_service.rs | 12 +++++----- .../kms/src/main_service/upgrade_authority.rs | 22 +++++++++---------- 23 files changed, 87 insertions(+), 84 deletions(-) diff --git a/docs/deployment.md b/docs/deployment.md index 223e7d191..d99d10933 100644 --- a/docs/deployment.md +++ b/docs/deployment.md @@ -550,7 +550,7 @@ for the full procedure; the KMS steps are: Each `checkKmsAuthorization` or `checkAppAuthorization` result carries the current epoch and its chain hash, including when authorization is denied. KMS verifies these against its keyring; -there is no cached serving epoch. An optional `requestedEpoch` selects historical +there is no cached serving epoch. An optional `requestedRootEpoch` selects historical keys, and omission selects the returned current epoch. Requests for future epochs are denied. No registered epoch returns a rejection with `rejectionReason: "Root epoch not registered"` and a zero hash. diff --git a/dstack/cert-client/src/lib.rs b/dstack/cert-client/src/lib.rs index 6c52588bf..5a2ef732f 100644 --- a/dstack/cert-client/src/lib.rs +++ b/dstack/cert-client/src/lib.rs @@ -85,9 +85,9 @@ pub async fn request_app_keys( })) .build() .into_client()?; - let requested_epoch = request.root_epoch; + let requested_root_epoch = request.root_epoch; let response = KmsClient::new(client).get_app_key(request).await?; - if requested_epoch.is_some_and(|epoch| epoch != response.root_epoch) { + if requested_root_epoch.is_some_and(|epoch| epoch != response.root_epoch) { bail!("KMS returned a different root epoch than requested"); } let cert = server_cert @@ -569,14 +569,14 @@ mod tests { assert_eq!(result.root_epoch, 1); let other_ca = kms_root().pem_cert; - for requested_epoch in [None, Some(0), Some(1)] { + for requested_root_epoch in [None, Some(0), Some(1)] { let request = || GetAppKeyRequest { api_version: 1, vm_config: String::new(), - root_epoch: requested_epoch, + root_epoch: requested_root_epoch, }; let selected = AppKeyResponse { - root_epoch: requested_epoch.unwrap_or(1), + root_epoch: requested_root_epoch.unwrap_or(1), ..response.clone() }; *response_body.lock().unwrap() = serde_json::to_string(&selected).unwrap(); diff --git a/dstack/kms/README.md b/dstack/kms/README.md index e1a857cea..2c0eb64bf 100644 --- a/dstack/kms/README.md +++ b/dstack/kms/README.md @@ -257,12 +257,12 @@ sent. The KMS method requires the attested `appId` to equal the proxy address, and the app method rejects that address and applies registered app policy. KMS callers do not fetch contract metadata to decide which policy to use. -`appEpoch` reserves the application's own key epoch, independently of KMS root +`requestedAppEpoch` reserves the application's own key epoch, independently of KMS root epochs. HTTP requests default it to 0; KMS currently sends 0. Nonzero values are rejected until app-epoch support is implemented, and rejection responses still carry the current root epoch and chain hash. -`requestedEpoch` is optional in the JSON request. Omit it to use the epoch in +`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 contract rejects a requested epoch newer than the current one; historical @@ -289,7 +289,7 @@ registered epoch, the result is `isAllowed: false`, never authorizes key release or requests an epoch switch. The KMS holds no serving-epoch snapshot. It verifies the returned current epoch -and chain hash against its keyring, then derives keys from `requestedEpoch` when +and chain hash against its keyring, then derives keys from `requestedRootEpoch` when specified or the returned current epoch otherwise. Missing or mismatched roots fail closed. Self-authorization retains its one-second cache of complete results. diff --git a/dstack/kms/auth-eth-bun/README.md b/dstack/kms/auth-eth-bun/README.md index 91178ac8f..bb028370f 100644 --- a/dstack/kms/auth-eth-bun/README.md +++ b/dstack/kms/auth-eth-bun/README.md @@ -206,9 +206,9 @@ the test suite (`index.test.ts`) validates: `POST /checkKmsAuthorization` and `POST /checkAppAuthorization` return `AuthorizationResult`, including `rejectionReason`, `gatewayAppId`, `currentRootEpoch`, and `currentRootEpochHash`. The selected endpoint determines the policy; the backend checks that the verified -`appId` matches that scope. `appEpoch` defaults to 0; nonzero values are currently +`appId` matches that scope. `requestedAppEpoch` defaults to 0; nonzero values are currently rejected. The optional -`requestedEpoch` selects historical keys; omission uses the returned current +`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, bootstrap, and compatibility requirements. The legacy `/bootAuth/*` endpoints diff --git a/dstack/kms/auth-eth-bun/index.test.ts b/dstack/kms/auth-eth-bun/index.test.ts index 504902f49..76042671d 100644 --- a/dstack/kms/auth-eth-bun/index.test.ts +++ b/dstack/kms/auth-eth-bun/index.test.ts @@ -841,8 +841,8 @@ describe('POST /checkAppAuthorization', () => { mrAggregated: '0x' + '44'.repeat(32), mrSystem: '0x' + '55'.repeat(32), osImageHash: '0x' + '66'.repeat(32), tcbStatus: 'UpToDate', advisoryIds: [], }; - const authorize = (requestedEpoch?: number) => appFetch(new Request('http://localhost/checkAppAuthorization', { - method: 'POST', headers: { 'Content-Type': 'application/json' }, body: JSON.stringify({ ...appInfo, requestedEpoch }), + const authorize = (requestedRootEpoch?: number) => appFetch(new Request('http://localhost/checkAppAuthorization', { + method: 'POST', headers: { 'Content-Type': 'application/json' }, body: JSON.stringify({ ...appInfo, requestedRootEpoch }), })); it('reads the complete result with one contract call at configured finality', async () => { @@ -857,7 +857,7 @@ describe('POST /checkAppAuthorization', () => { expect(mockReadContract).toHaveBeenCalledTimes(1); expect(mockReadContract).toHaveBeenCalledWith(expect.objectContaining({ functionName: 'checkAppAuthorization', blockNumber: 98n, - args: [expect.objectContaining({ appInfo: expect.objectContaining({ appId: appInfo.appId }), appEpoch: 0n, requestedEpoch: { isSet: false, value: 0n } })], + args: [expect.objectContaining({ appInfo: expect.objectContaining({ appId: appInfo.appId }), requestedAppEpoch: 0n, requestedRootEpoch: { isSet: false, value: 0n } })], })); }); @@ -868,7 +868,7 @@ describe('POST /checkAppAuthorization', () => { }); expect((await authorize(0)).status).toBe(200); expect(mockReadContract).toHaveBeenCalledWith(expect.objectContaining({ - args: [expect.objectContaining({ requestedEpoch: { isSet: true, value: 0n } })], + args: [expect.objectContaining({ requestedRootEpoch: { isSet: true, value: 0n } })], })); }); diff --git a/dstack/kms/auth-eth-bun/index.ts b/dstack/kms/auth-eth-bun/index.ts index dc7763edf..e9892d51a 100644 --- a/dstack/kms/auth-eth-bun/index.ts +++ b/dstack/kms/auth-eth-bun/index.ts @@ -37,8 +37,8 @@ const BootInfoSchema = z.object({ }); const AuthorizationRequestSchema = BootInfoSchema.omit({ rootEpoch: true, rootEpochHash: true }).extend({ - appEpoch: z.number().int().min(0).max(0xffffffff).default(0), - requestedEpoch: z.number().int().min(0).max(0xffffffff).optional(), + requestedAppEpoch: z.number().int().min(0).max(0xffffffff).default(0), + requestedRootEpoch: z.number().int().min(0).max(0xffffffff).optional(), }); type AuthorizationRequest = z.infer; @@ -94,8 +94,8 @@ const AUTHORIZATION_ABI = { stateMutability: 'view', inputs: [{ name: 'request', type: 'tuple', components: [ { name: 'appInfo', type: 'tuple', components: APP_INFO_COMPONENTS }, - { name: 'appEpoch', type: 'uint256' }, - { name: 'requestedEpoch', type: 'tuple', components: [ + { name: 'requestedAppEpoch', type: 'uint256' }, + { name: 'requestedRootEpoch', type: 'tuple', components: [ { name: 'isSet', type: 'bool' }, { name: 'value', type: 'uint256' }, ] }, ] }], @@ -245,8 +245,8 @@ class EthereumBackend { functionName: method, args: [{ appInfo: this.contractAppInfo(appInfo), - appEpoch: BigInt(appInfo.appEpoch), - requestedEpoch: { isSet: appInfo.requestedEpoch !== undefined, value: BigInt(appInfo.requestedEpoch ?? 0) }, + requestedAppEpoch: BigInt(appInfo.requestedAppEpoch), + requestedRootEpoch: { isSet: appInfo.requestedRootEpoch !== undefined, value: BigInt(appInfo.requestedRootEpoch ?? 0) }, }], blockNumber, }); diff --git a/dstack/kms/auth-eth-bun/openapi.json b/dstack/kms/auth-eth-bun/openapi.json index f07781e1c..005f64a86 100644 --- a/dstack/kms/auth-eth-bun/openapi.json +++ b/dstack/kms/auth-eth-bun/openapi.json @@ -410,13 +410,13 @@ ], "type": "object", "properties": { - "requestedEpoch": { + "requestedRootEpoch": { "type": "integer", "minimum": 0, "maximum": 4294967295, "description": "Omit to use the current epoch. An explicit zero selects epoch 0; future epochs are denied." }, - "appEpoch": { + "requestedAppEpoch": { "type": "integer", "minimum": 0, "maximum": 4294967295, diff --git a/dstack/kms/auth-eth/README.md b/dstack/kms/auth-eth/README.md index 2580aea62..02e488b4b 100644 --- a/dstack/kms/auth-eth/README.md +++ b/dstack/kms/auth-eth/README.md @@ -177,9 +177,9 @@ Documentation: https://book.getfoundry.sh/ `POST /checkKmsAuthorization` and `POST /checkAppAuthorization` return `AuthorizationResult`, including `rejectionReason`, `gatewayAppId`, `currentRootEpoch`, and `currentRootEpochHash`. The selected endpoint determines the policy; the backend checks that the verified -`appId` matches that scope. `appEpoch` defaults to 0; nonzero values are currently +`appId` matches that scope. `requestedAppEpoch` defaults to 0; nonzero values are currently rejected. The optional -`requestedEpoch` selects historical keys; omission uses the returned current +`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, bootstrap, and compatibility requirements. The legacy `/bootAuth/*` endpoints diff --git a/dstack/kms/auth-eth/contracts/DstackKms.sol b/dstack/kms/auth-eth/contracts/DstackKms.sol index 806c029de..f52177272 100644 --- a/dstack/kms/auth-eth/contracts/DstackKms.sol +++ b/dstack/kms/auth-eth/contracts/DstackKms.sol @@ -33,8 +33,9 @@ contract DstackKms is Initializable, Ownable2StepUpgradeable, UUPSUpgradeable, E struct AuthorizationRequest { AppInfo appInfo; - uint256 appEpoch; - OptionalEpoch requestedEpoch; + // Reserved for application key rotation; only zero is currently supported. + uint256 requestedAppEpoch; + OptionalEpoch requestedRootEpoch; } struct AuthorizationResult { @@ -434,12 +435,12 @@ contract DstackKms is Initializable, Ownable2StepUpgradeable, UUPSUpgradeable, E } result.currentRootEpoch = _currentRootEpoch; result.currentRootEpochHash = _rootEpochs[_currentRootEpoch].chainHash; - if (request.requestedEpoch.isSet && request.requestedEpoch.value > _currentRootEpoch) { + if (request.requestedRootEpoch.isSet && request.requestedRootEpoch.value > _currentRootEpoch) { result.rejectionReason = "Requested root epoch is not active"; return result; } - if (request.appEpoch != 0) { - result.rejectionReason = "App epoch is not supported"; + if (request.requestedAppEpoch != 0) { + result.rejectionReason = "Requested app epoch must be zero"; return result; } AppInfo calldata appInfo = request.appInfo; diff --git a/dstack/kms/auth-eth/src/ethereum.ts b/dstack/kms/auth-eth/src/ethereum.ts index 20dd03fbe..1ac43cb6b 100644 --- a/dstack/kms/auth-eth/src/ethereum.ts +++ b/dstack/kms/auth-eth/src/ethereum.ts @@ -7,8 +7,8 @@ import { BootInfo, BootResponse, RootEpochInfo, AuthorizationResult, Authorizati // Minimal ABI for DstackKms contract const DSTACK_KMS_ABI = [ - "function checkKmsAuthorization(((address appId,bytes32 composeHash,address instanceId,bytes32 deviceId,bytes32 mrAggregated,bytes32 mrSystem,bytes32 osImageHash,string tcbStatus,string[] advisoryIds) appInfo,uint256 appEpoch,(bool isSet,uint256 value) requestedEpoch) request) view returns ((bool isAllowed,string rejectionReason,string gatewayAppId,uint256 currentRootEpoch,bytes32 currentRootEpochHash))", - "function checkAppAuthorization(((address appId,bytes32 composeHash,address instanceId,bytes32 deviceId,bytes32 mrAggregated,bytes32 mrSystem,bytes32 osImageHash,string tcbStatus,string[] advisoryIds) appInfo,uint256 appEpoch,(bool isSet,uint256 value) requestedEpoch) request) view returns ((bool isAllowed,string rejectionReason,string gatewayAppId,uint256 currentRootEpoch,bytes32 currentRootEpochHash))", + "function checkKmsAuthorization(((address appId,bytes32 composeHash,address instanceId,bytes32 deviceId,bytes32 mrAggregated,bytes32 mrSystem,bytes32 osImageHash,string tcbStatus,string[] advisoryIds) appInfo,uint256 requestedAppEpoch,(bool isSet,uint256 value) requestedRootEpoch) request) view returns ((bool isAllowed,string rejectionReason,string gatewayAppId,uint256 currentRootEpoch,bytes32 currentRootEpochHash))", + "function checkAppAuthorization(((address appId,bytes32 composeHash,address instanceId,bytes32 deviceId,bytes32 mrAggregated,bytes32 mrSystem,bytes32 osImageHash,string tcbStatus,string[] advisoryIds) appInfo,uint256 requestedAppEpoch,(bool isSet,uint256 value) requestedRootEpoch) request) view returns ((bool isAllowed,string rejectionReason,string gatewayAppId,uint256 currentRootEpoch,bytes32 currentRootEpochHash))", "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)", "function isKmsAllowed((address appId,bytes32 composeHash,address instanceId,bytes32 deviceId,bytes32 mrAggregated,bytes32 mrSystem,bytes32 osImageHash,string tcbStatus,string[] advisoryIds) bootInfo) view returns (bool, string)", "function gatewayAppId() view returns (string)", @@ -81,8 +81,8 @@ export class EthereumBackend { async checkAuthorization(appInfo: AuthorizationRequest, method: 'checkKmsAuthorization' | 'checkAppAuthorization'): Promise { const result = await this.kmsContract[method]({ appInfo: this.contractAppInfo(appInfo), - appEpoch: appInfo.appEpoch ?? 0, - requestedEpoch: { isSet: appInfo.requestedEpoch !== undefined, value: appInfo.requestedEpoch ?? 0 }, + requestedAppEpoch: appInfo.requestedAppEpoch ?? 0, + requestedRootEpoch: { isSet: appInfo.requestedRootEpoch !== undefined, value: appInfo.requestedRootEpoch ?? 0 }, }); return { isAllowed: result.isAllowed, diff --git a/dstack/kms/auth-eth/src/main.test.ts b/dstack/kms/auth-eth/src/main.test.ts index 679a229f9..da54da869 100644 --- a/dstack/kms/auth-eth/src/main.test.ts +++ b/dstack/kms/auth-eth/src/main.test.ts @@ -92,7 +92,7 @@ describe('Server', () => { }); describe('POST /checkAppAuthorization', () => { - const appRequest = { ...mockBootInfo, appEpoch: 0 }; + const appRequest = { ...mockBootInfo, requestedAppEpoch: 0 }; it('preserves the complete decision, including the epoch on rejection', async () => { const decision = { isAllowed: false, rejectionReason: 'App not registered', gatewayAppId: 'gateway', @@ -105,8 +105,8 @@ describe('Server', () => { expect(app.ethereum.checkAuthorization).toHaveBeenCalledWith(appRequest, 'checkAppAuthorization'); await app.inject({ method: 'POST', url: '/checkKmsAuthorization', payload: appRequest }); expect(app.ethereum.checkAuthorization).toHaveBeenLastCalledWith(appRequest, 'checkKmsAuthorization'); - await app.inject({ method: 'POST', url: '/checkAppAuthorization', payload: { ...appRequest, requestedEpoch: 0 } }); - expect(app.ethereum.checkAuthorization).toHaveBeenLastCalledWith({ ...appRequest, requestedEpoch: 0 }, 'checkAppAuthorization'); + await app.inject({ method: 'POST', url: '/checkAppAuthorization', payload: { ...appRequest, requestedRootEpoch: 0 } }); + expect(app.ethereum.checkAuthorization).toHaveBeenLastCalledWith({ ...appRequest, requestedRootEpoch: 0 }, 'checkAppAuthorization'); }); it('fails closed when the new contract method is unavailable', async () => { diff --git a/dstack/kms/auth-eth/src/server.ts b/dstack/kms/auth-eth/src/server.ts index 081c56f0e..af4cd74d0 100644 --- a/dstack/kms/auth-eth/src/server.ts +++ b/dstack/kms/auth-eth/src/server.ts @@ -55,8 +55,8 @@ export async function build(): Promise { $id: 'authorizationRequest', properties: { ...appInfoSchema.properties, - appEpoch: { type: 'integer', minimum: 0, maximum: 4294967295, default: 0 }, - requestedEpoch: { type: 'integer', minimum: 0, maximum: 4294967295 }, + requestedAppEpoch: { type: 'integer', minimum: 0, maximum: 4294967295, default: 0 }, + requestedRootEpoch: { type: 'integer', minimum: 0, maximum: 4294967295 }, }, }); diff --git a/dstack/kms/auth-eth/src/types.ts b/dstack/kms/auth-eth/src/types.ts index 40d781acf..3a8ea1cbc 100644 --- a/dstack/kms/auth-eth/src/types.ts +++ b/dstack/kms/auth-eth/src/types.ts @@ -42,6 +42,6 @@ export interface AuthorizationResult { } export interface AuthorizationRequest extends Omit { - appEpoch?: number; - requestedEpoch?: number; + requestedAppEpoch?: number; + requestedRootEpoch?: number; } diff --git a/dstack/kms/auth-eth/test/DstackKms.t.sol b/dstack/kms/auth-eth/test/DstackKms.t.sol index 0eb2d3f1e..adf734784 100644 --- a/dstack/kms/auth-eth/test/DstackKms.t.sol +++ b/dstack/kms/auth-eth/test/DstackKms.t.sol @@ -488,11 +488,11 @@ contract DstackKmsTest is Test { vm.prank(owner); kms.activateRootEpoch(1, h1); assertTrue(kms.checkKmsAuthorization(request).isAllowed); - request.requestedEpoch.value = 0; + request.requestedRootEpoch.value = 0; assertTrue(kms.checkKmsAuthorization(request).isAllowed); - request.requestedEpoch.value = type(uint256).max; + request.requestedRootEpoch.value = type(uint256).max; assertFalse(kms.checkKmsAuthorization(request).isAllowed); - request.requestedEpoch.isSet = false; + request.requestedRootEpoch.isSet = false; result = kms.checkKmsAuthorization(request); assertTrue(result.isAllowed); assertEq(result.currentRootEpoch, 1); @@ -540,7 +540,7 @@ contract DstackKmsTest is Test { DstackKms.AuthorizationRequest(info, 1, DstackKms.OptionalEpoch(false, 0)); result = kms.checkAppAuthorization(appRequest); assertFalse(result.isAllowed); - assertEq(result.rejectionReason, "App epoch is not supported"); + assertEq(result.rejectionReason, "Requested app epoch must be zero"); assertEq(result.currentRootEpochHash, chainHash(bytes32(0), ca(0xa0), k256(0xb0))); info.appId = address(kms); result = authorize(info, false); diff --git a/dstack/kms/auth-mock/README.md b/dstack/kms/auth-mock/README.md index 3bd731747..d32ee9ada 100644 --- a/dstack/kms/auth-mock/README.md +++ b/dstack/kms/auth-mock/README.md @@ -174,9 +174,9 @@ the API remains identical, so no client code changes are needed. `POST /checkKmsAuthorization` and `POST /checkAppAuthorization` return `AuthorizationResult`, including `rejectionReason`, `gatewayAppId`, `currentRootEpoch`, and `currentRootEpochHash`. The selected endpoint determines the policy; the backend checks that the verified -`appId` matches that scope. `appEpoch` defaults to 0; nonzero values are currently +`appId` matches that scope. `requestedAppEpoch` defaults to 0; nonzero values are currently rejected. The optional -`requestedEpoch` selects historical keys; omission uses the returned current +`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, bootstrap, and compatibility requirements. The legacy `/bootAuth/*` endpoints diff --git a/dstack/kms/auth-mock/index.test.ts b/dstack/kms/auth-mock/index.test.ts index 6f6858fe1..d9a22ff3c 100644 --- a/dstack/kms/auth-mock/index.test.ts +++ b/dstack/kms/auth-mock/index.test.ts @@ -385,9 +385,9 @@ describe('Mock Behavior Verification', () => { describe('authorization endpoints', () => { - const authorize = (appId: string, isKms = false, appEpoch = 0) => appFetch(new Request(`http://localhost/${isKms ? 'checkKmsAuthorization' : 'checkAppAuthorization'}`, { + const authorize = (appId: string, isKms = false, requestedAppEpoch = 0) => appFetch(new Request(`http://localhost/${isKms ? 'checkKmsAuthorization' : 'checkAppAuthorization'}`, { method: 'POST', headers: { 'Content-Type': 'application/json' }, - body: JSON.stringify({ appId, appEpoch, mrAggregated: '', osImageHash: '', composeHash: '', instanceId: '', deviceId: '' }), + body: JSON.stringify({ appId, requestedAppEpoch, mrAggregated: '', osImageHash: '', composeHash: '', instanceId: '', deviceId: '' }), })); it('requires registered roots and selects policy by app ID', async () => { diff --git a/dstack/kms/auth-mock/index.ts b/dstack/kms/auth-mock/index.ts index 4c067189f..d5b21c6c2 100644 --- a/dstack/kms/auth-mock/index.ts +++ b/dstack/kms/auth-mock/index.ts @@ -30,8 +30,8 @@ const BootInfoSchema = z.object({ }); const AuthorizationRequestSchema = BootInfoSchema.omit({ rootEpoch: true, rootEpochHash: true }).extend({ - appEpoch: z.number().int().min(0).max(0xffffffff).default(0), - requestedEpoch: z.number().int().min(0).max(0xffffffff).optional(), + requestedAppEpoch: z.number().int().min(0).max(0xffffffff).default(0), + requestedRootEpoch: z.number().int().min(0).max(0xffffffff).optional(), }); type AuthorizationRequest = z.infer; @@ -121,11 +121,11 @@ class MockBackend { if (current.rootEpoch === undefined) { return { isAllowed: false, rejectionReason: 'Root epoch not registered', gatewayAppId: this.mockGatewayAppId, ...epoch }; } - if (appInfo.requestedEpoch !== undefined && appInfo.requestedEpoch > current.rootEpoch) { + if (appInfo.requestedRootEpoch !== undefined && appInfo.requestedRootEpoch > current.rootEpoch) { return { isAllowed: false, rejectionReason: 'Requested root epoch is not active', gatewayAppId: this.mockGatewayAppId, ...epoch }; } - if (appInfo.appEpoch !== 0) { - return { isAllowed: false, rejectionReason: 'App epoch is not supported', gatewayAppId: this.mockGatewayAppId, ...epoch }; + if (appInfo.requestedAppEpoch !== 0) { + return { isAllowed: false, rejectionReason: 'Requested app epoch must be zero', gatewayAppId: this.mockGatewayAppId, ...epoch }; } const kmsAppId = process.env.KMS_CONTRACT_ADDR || '0xmockcontract1234567890123456789012345678'; if (isKms !== (normHash(appInfo.appId) === normHash(kmsAppId))) { diff --git a/dstack/kms/auth-mock/openapi.json b/dstack/kms/auth-mock/openapi.json index 6f862157c..e26c8ead2 100644 --- a/dstack/kms/auth-mock/openapi.json +++ b/dstack/kms/auth-mock/openapi.json @@ -419,13 +419,13 @@ ], "type": "object", "properties": { - "requestedEpoch": { + "requestedRootEpoch": { "type": "integer", "minimum": 0, "maximum": 4294967295, "description": "Omit to use the current epoch. An explicit zero selects epoch 0; future epochs are denied." }, - "appEpoch": { + "requestedAppEpoch": { "type": "integer", "minimum": 0, "maximum": 4294967295, diff --git a/dstack/kms/auth-simple/README.md b/dstack/kms/auth-simple/README.md index 9caf9387e..4918f2568 100644 --- a/dstack/kms/auth-simple/README.md +++ b/dstack/kms/auth-simple/README.md @@ -259,9 +259,9 @@ bun run test:run `POST /checkKmsAuthorization` and `POST /checkAppAuthorization` return `AuthorizationResult`, including `rejectionReason`, `gatewayAppId`, `currentRootEpoch`, and `currentRootEpochHash`. The selected endpoint determines the policy; the backend checks that the verified -`appId` matches that scope. `appEpoch` defaults to 0; nonzero values are currently +`appId` matches that scope. `requestedAppEpoch` defaults to 0; nonzero values are currently rejected. The optional -`requestedEpoch` selects historical keys; omission uses the returned current +`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, bootstrap, and compatibility requirements. The legacy `/bootAuth/*` endpoints diff --git a/dstack/kms/auth-simple/index.test.ts b/dstack/kms/auth-simple/index.test.ts index 50574cdb1..f490d03e7 100644 --- a/dstack/kms/auth-simple/index.test.ts +++ b/dstack/kms/auth-simple/index.test.ts @@ -484,13 +484,13 @@ describe('auth-simple', () => { expect(await authorize()).toMatchObject({ isAllowed: true, rejectionReason: '', currentRootEpoch: 1, currentRootEpochHash: h1, }); - expect(await authorize({ requestedEpoch: 0 })).toMatchObject({ isAllowed: true, currentRootEpoch: 1 }); - expect(await authorize({ requestedEpoch: 1 })).toMatchObject({ isAllowed: true }); + expect(await authorize({ requestedRootEpoch: 0 })).toMatchObject({ isAllowed: true, currentRootEpoch: 1 }); + expect(await authorize({ requestedRootEpoch: 1 })).toMatchObject({ isAllowed: true }); expect(await authorize({}, true)).toMatchObject({ isAllowed: false, rejectionReason: 'Authorization kind does not match app ID' }); expect(await authorize({ appId: kmsContractAddr })).toMatchObject({ isAllowed: false, rejectionReason: 'Authorization kind does not match app ID' }); - expect(await authorize({ appEpoch: 1 })).toMatchObject({ isAllowed: false, rejectionReason: 'App epoch is not supported', currentRootEpoch: 1 }); - expect(await authorize({ appEpoch: 0 })).toMatchObject({ isAllowed: true }); - expect(await authorize({ requestedEpoch: 2 })).toMatchObject({ + expect(await authorize({ requestedAppEpoch: 1 })).toMatchObject({ isAllowed: false, rejectionReason: 'Requested app epoch must be zero', 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, }); expect(await authorize({ appId: kmsContractAddr, composeHash: 'not-an-app-compose' }, true)).toMatchObject({ isAllowed: true }); diff --git a/dstack/kms/auth-simple/index.ts b/dstack/kms/auth-simple/index.ts index d04297565..fa8747628 100644 --- a/dstack/kms/auth-simple/index.ts +++ b/dstack/kms/auth-simple/index.ts @@ -28,8 +28,8 @@ const BootInfoSchema = z.object({ }); const AuthorizationRequestSchema = BootInfoSchema.omit({ rootEpoch: true, rootEpochHash: true }).extend({ - appEpoch: z.number().int().min(0).max(0xffffffff).default(0), - requestedEpoch: z.number().int().min(0).max(0xffffffff).optional(), + requestedAppEpoch: z.number().int().min(0).max(0xffffffff).default(0), + requestedRootEpoch: z.number().int().min(0).max(0xffffffff).optional(), }); type AuthorizationRequest = z.infer; @@ -168,11 +168,11 @@ class ConfigBackend { if (current.rootEpoch === undefined) { return { isAllowed: false, rejectionReason: 'Root epoch not registered', gatewayAppId: config.gatewayAppId, ...epoch }; } - if (appInfo.requestedEpoch !== undefined && appInfo.requestedEpoch > current.rootEpoch) { + if (appInfo.requestedRootEpoch !== undefined && appInfo.requestedRootEpoch > current.rootEpoch) { return { isAllowed: false, rejectionReason: 'Requested root epoch is not active', gatewayAppId: config.gatewayAppId, ...epoch }; } - if (appInfo.appEpoch !== 0) { - return { isAllowed: false, rejectionReason: 'App epoch is not supported', gatewayAppId: config.gatewayAppId, ...epoch }; + if (appInfo.requestedAppEpoch !== 0) { + return { isAllowed: false, rejectionReason: 'Requested app epoch must be zero', gatewayAppId: config.gatewayAppId, ...epoch }; } if (isKms !== (normalizeHex(appInfo.appId) === normalizeHex(config.kmsContractAddr))) { return { isAllowed: false, rejectionReason: 'Authorization kind does not match app ID', gatewayAppId: config.gatewayAppId, ...epoch }; diff --git a/dstack/kms/src/main_service.rs b/dstack/kms/src/main_service.rs index 0bf9e73c9..f5f9c4b82 100644 --- a/dstack/kms/src/main_service.rs +++ b/dstack/kms/src/main_service.rs @@ -528,10 +528,10 @@ impl RpcHandler { async fn ensure_app_boot_allowed( &self, vm_config: &str, - requested_epoch: Option, + requested_root_epoch: Option, ) -> Result { let att = self.ensure_attested()?; - self.ensure_app_attestation_allowed(att, false, false, vm_config, requested_epoch) + self.ensure_app_attestation_allowed(att, false, false, vm_config, requested_root_epoch) .await } @@ -541,17 +541,19 @@ impl RpcHandler { is_kms: bool, use_boottime_mr: bool, vm_config_str: &str, - requested_epoch: Option, + requested_root_epoch: Option, ) -> Result { let boot_info = build_boot_info_for_attestation(att, use_boottime_mr, vm_config_str)?; let auth = &self.state.config.auth_api; let response = if is_kms { auth.check_kms_authorization(&boot_info).await? } else { - auth.check_app_authorization(&boot_info, requested_epoch) + auth.check_app_authorization(&boot_info, requested_root_epoch) .await? }; - let root_epoch = self.state.authorized_epoch(&response, requested_epoch)?; + let root_epoch = self + .state + .authorized_epoch(&response, requested_root_epoch)?; let current_epoch = if self.state.config.auth_api.is_dev() { self.state.keyring.latest() } else { diff --git a/dstack/kms/src/main_service/upgrade_authority.rs b/dstack/kms/src/main_service/upgrade_authority.rs index a960c957f..89407de23 100644 --- a/dstack/kms/src/main_service/upgrade_authority.rs +++ b/dstack/kms/src/main_service/upgrade_authority.rs @@ -265,9 +265,9 @@ impl AuthApi { pub async fn check_app_authorization( &self, app_info: &BootInfo, - requested_epoch: Option, + requested_root_epoch: Option, ) -> Result { - self.check_authorization(app_info, "checkAppAuthorization", requested_epoch) + self.check_authorization(app_info, "checkAppAuthorization", requested_root_epoch) .await } @@ -275,16 +275,16 @@ impl AuthApi { &self, app_info: &BootInfo, method: &str, - requested_epoch: Option, + requested_root_epoch: Option, ) -> Result { #[derive(Serialize)] #[serde(rename_all = "camelCase")] struct AuthorizationRequest<'a> { #[serde(flatten)] app_info: &'a BootInfo, - app_epoch: u32, + requested_app_epoch: u32, #[serde(skip_serializing_if = "Option::is_none")] - requested_epoch: Option, + requested_root_epoch: Option, } ensure_identity_widths(app_info)?; match self { @@ -301,8 +301,8 @@ impl AuthApi { &url, &AuthorizationRequest { app_info, - app_epoch: 0, - requested_epoch, + requested_app_epoch: 0, + requested_root_epoch, }, webhook.timeout, ) @@ -724,11 +724,11 @@ pub(crate) mod tests { .all(|(path, _)| path == "/checkAppAuthorization")); assert_eq!(requests[3].0, "/checkKmsAuthorization"); let mut expected = serde_json::to_value(info).unwrap(); - expected["appEpoch"] = 0.into(); + expected["requestedAppEpoch"] = 0.into(); assert_eq!(requests[0].1, expected); - assert_eq!(requests[2].1["appEpoch"], 0); - assert!(requests[0].1.get("requestedEpoch").is_none()); - assert_eq!(requests[2].1["requestedEpoch"], 0); + assert_eq!(requests[2].1["requestedAppEpoch"], 0); + assert!(requests[0].1.get("requestedRootEpoch").is_none()); + assert_eq!(requests[2].1["requestedRootEpoch"], 0); } #[rocket::async_test] From ee4be5a318123d51a0e0448c49dfeecd69e9e119 Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Sat, 7 Nov 2026 00:42:05 -0800 Subject: [PATCH 023/166] feat(kms): reserve optional app epochs and report currentAppEpoch Signed-off-by: Kevin Wang --- dstack/kms/README.md | 4 +++- dstack/kms/auth-eth-bun/index.test.ts | 8 ++++---- dstack/kms/auth-eth-bun/index.ts | 11 ++++++++--- dstack/kms/auth-eth-bun/openapi.json | 8 +++++++- dstack/kms/auth-eth/contracts/DstackKms.sol | 5 +++-- dstack/kms/auth-eth/src/ethereum.ts | 7 ++++--- dstack/kms/auth-eth/src/main.test.ts | 2 +- dstack/kms/auth-eth/src/server.ts | 3 ++- dstack/kms/auth-eth/src/types.ts | 1 + dstack/kms/auth-eth/test/DstackKms.t.sol | 6 +++--- dstack/kms/auth-mock/index.ts | 1 + dstack/kms/auth-mock/openapi.json | 8 +++++++- dstack/kms/auth-simple/index.ts | 1 + dstack/kms/src/main_service.rs | 4 ++++ dstack/kms/src/main_service/upgrade_authority.rs | 4 +++- 15 files changed, 52 insertions(+), 21 deletions(-) diff --git a/dstack/kms/README.md b/dstack/kms/README.md index 2c0eb64bf..682839967 100644 --- a/dstack/kms/README.md +++ b/dstack/kms/README.md @@ -258,7 +258,9 @@ and the app method rejects that address and applies registered app 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. Nonzero values are +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. diff --git a/dstack/kms/auth-eth-bun/index.test.ts b/dstack/kms/auth-eth-bun/index.test.ts index 76042671d..63f19808d 100644 --- a/dstack/kms/auth-eth-bun/index.test.ts +++ b/dstack/kms/auth-eth-bun/index.test.ts @@ -848,23 +848,23 @@ describe('POST /checkAppAuthorization', () => { it('reads the complete result with one contract call at configured finality', async () => { const result = { isAllowed: false, rejectionReason: 'App not registered', gatewayAppId: 'gateway', - currentRootEpoch: 2n, currentRootEpochHash: '0x' + '77'.repeat(32), + currentAppEpoch: 0n, currentRootEpoch: 2n, currentRootEpochHash: '0x' + '77'.repeat(32), }; mockReadContract.mockResolvedValueOnce(result); const response = await authorize(); expect(response.status).toBe(200); - expect(await response.json()).toEqual({ ...result, currentRootEpoch: 2 }); + expect(await response.json()).toEqual({ ...result, currentRootEpoch: 2, currentAppEpoch: 0 }); expect(mockReadContract).toHaveBeenCalledTimes(1); expect(mockReadContract).toHaveBeenCalledWith(expect.objectContaining({ functionName: 'checkAppAuthorization', blockNumber: 98n, - args: [expect.objectContaining({ appInfo: expect.objectContaining({ appId: appInfo.appId }), requestedAppEpoch: 0n, requestedRootEpoch: { isSet: false, value: 0n } })], + args: [expect.objectContaining({ appInfo: expect.objectContaining({ appId: appInfo.appId }), requestedAppEpoch: { isSet: true, value: 0n }, requestedRootEpoch: { isSet: false, value: 0n } })], })); }); it('encodes an explicit epoch zero as present', async () => { mockReadContract.mockResolvedValueOnce({ isAllowed: true, rejectionReason: '', gatewayAppId: 'gateway', - currentRootEpoch: 1n, currentRootEpochHash: '0x' + '77'.repeat(32), + currentAppEpoch: 0n, currentRootEpoch: 1n, currentRootEpochHash: '0x' + '77'.repeat(32), }); expect((await authorize(0)).status).toBe(200); expect(mockReadContract).toHaveBeenCalledWith(expect.objectContaining({ diff --git a/dstack/kms/auth-eth-bun/index.ts b/dstack/kms/auth-eth-bun/index.ts index e9892d51a..a29bd40ee 100644 --- a/dstack/kms/auth-eth-bun/index.ts +++ b/dstack/kms/auth-eth-bun/index.ts @@ -94,7 +94,9 @@ const AUTHORIZATION_ABI = { stateMutability: 'view', inputs: [{ name: 'request', type: 'tuple', components: [ { name: 'appInfo', type: 'tuple', components: APP_INFO_COMPONENTS }, - { name: 'requestedAppEpoch', type: 'uint256' }, + { name: 'requestedAppEpoch', type: 'tuple', components: [ + { name: 'isSet', type: 'bool' }, { name: 'value', type: 'uint256' }, + ] }, { name: 'requestedRootEpoch', type: 'tuple', components: [ { name: 'isSet', type: 'bool' }, { name: 'value', type: 'uint256' }, ] }, @@ -108,6 +110,7 @@ const AUTHORIZATION_ABI = { { name: 'gatewayAppId', type: 'string' }, { name: 'currentRootEpoch', type: 'uint256' }, { name: 'currentRootEpochHash', type: 'bytes32' }, + { name: 'currentAppEpoch', type: 'uint256' }, ], }], } as const; @@ -245,7 +248,7 @@ class EthereumBackend { functionName: method, args: [{ appInfo: this.contractAppInfo(appInfo), - requestedAppEpoch: BigInt(appInfo.requestedAppEpoch), + requestedAppEpoch: { isSet: appInfo.requestedAppEpoch !== undefined, value: BigInt(appInfo.requestedAppEpoch ?? 0) }, requestedRootEpoch: { isSet: appInfo.requestedRootEpoch !== undefined, value: BigInt(appInfo.requestedRootEpoch ?? 0) }, }], blockNumber, @@ -254,7 +257,9 @@ class EthereumBackend { if (!Number.isSafeInteger(currentRootEpoch)) { throw new Error('current root epoch exceeds the supported range'); } - return { ...result, currentRootEpoch }; + const currentAppEpoch = Number(result.currentAppEpoch); + if (!Number.isSafeInteger(currentAppEpoch)) throw new Error('current app epoch exceeds the supported range'); + return { ...result, currentRootEpoch, currentAppEpoch }; } async checkBoot(bootInfo: BootInfo, isKms: boolean): Promise { diff --git a/dstack/kms/auth-eth-bun/openapi.json b/dstack/kms/auth-eth-bun/openapi.json index 005f64a86..46a99357a 100644 --- a/dstack/kms/auth-eth-bun/openapi.json +++ b/dstack/kms/auth-eth-bun/openapi.json @@ -378,7 +378,8 @@ "rejectionReason", "gatewayAppId", "currentRootEpoch", - "currentRootEpochHash" + "currentRootEpochHash", + "currentAppEpoch" ], "properties": { "isAllowed": { @@ -399,6 +400,11 @@ "type": "string", "pattern": "^0x[0-9a-fA-F]{64}$", "description": "Zero when no epoch is registered; authorization is denied." + }, + "currentAppEpoch": { + "type": "integer", + "minimum": 0, + "description": "Current application epoch; zero until app rotation is supported." } } }, diff --git a/dstack/kms/auth-eth/contracts/DstackKms.sol b/dstack/kms/auth-eth/contracts/DstackKms.sol index f52177272..a65fce682 100644 --- a/dstack/kms/auth-eth/contracts/DstackKms.sol +++ b/dstack/kms/auth-eth/contracts/DstackKms.sol @@ -34,7 +34,7 @@ contract DstackKms is Initializable, Ownable2StepUpgradeable, UUPSUpgradeable, E struct AuthorizationRequest { AppInfo appInfo; // Reserved for application key rotation; only zero is currently supported. - uint256 requestedAppEpoch; + OptionalEpoch requestedAppEpoch; OptionalEpoch requestedRootEpoch; } @@ -44,6 +44,7 @@ contract DstackKms is Initializable, Ownable2StepUpgradeable, UUPSUpgradeable, E string gatewayAppId; uint256 currentRootEpoch; bytes32 currentRootEpochHash; + uint256 currentAppEpoch; } // Struct for KMS information @@ -439,7 +440,7 @@ contract DstackKms is Initializable, Ownable2StepUpgradeable, UUPSUpgradeable, E result.rejectionReason = "Requested root epoch is not active"; return result; } - if (request.requestedAppEpoch != 0) { + if (request.requestedAppEpoch.isSet && request.requestedAppEpoch.value != 0) { result.rejectionReason = "Requested app epoch must be zero"; return result; } diff --git a/dstack/kms/auth-eth/src/ethereum.ts b/dstack/kms/auth-eth/src/ethereum.ts index 1ac43cb6b..99ce15c42 100644 --- a/dstack/kms/auth-eth/src/ethereum.ts +++ b/dstack/kms/auth-eth/src/ethereum.ts @@ -7,8 +7,8 @@ import { BootInfo, BootResponse, RootEpochInfo, AuthorizationResult, Authorizati // Minimal ABI for DstackKms contract const DSTACK_KMS_ABI = [ - "function checkKmsAuthorization(((address appId,bytes32 composeHash,address instanceId,bytes32 deviceId,bytes32 mrAggregated,bytes32 mrSystem,bytes32 osImageHash,string tcbStatus,string[] advisoryIds) appInfo,uint256 requestedAppEpoch,(bool isSet,uint256 value) requestedRootEpoch) request) view returns ((bool isAllowed,string rejectionReason,string gatewayAppId,uint256 currentRootEpoch,bytes32 currentRootEpochHash))", - "function checkAppAuthorization(((address appId,bytes32 composeHash,address instanceId,bytes32 deviceId,bytes32 mrAggregated,bytes32 mrSystem,bytes32 osImageHash,string tcbStatus,string[] advisoryIds) appInfo,uint256 requestedAppEpoch,(bool isSet,uint256 value) requestedRootEpoch) request) view returns ((bool isAllowed,string rejectionReason,string gatewayAppId,uint256 currentRootEpoch,bytes32 currentRootEpochHash))", + "function checkKmsAuthorization(((address appId,bytes32 composeHash,address instanceId,bytes32 deviceId,bytes32 mrAggregated,bytes32 mrSystem,bytes32 osImageHash,string tcbStatus,string[] advisoryIds) appInfo,(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 checkAppAuthorization(((address appId,bytes32 composeHash,address instanceId,bytes32 deviceId,bytes32 mrAggregated,bytes32 mrSystem,bytes32 osImageHash,string tcbStatus,string[] advisoryIds) appInfo,(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)", "function isKmsAllowed((address appId,bytes32 composeHash,address instanceId,bytes32 deviceId,bytes32 mrAggregated,bytes32 mrSystem,bytes32 osImageHash,string tcbStatus,string[] advisoryIds) bootInfo) view returns (bool, string)", "function gatewayAppId() view returns (string)", @@ -81,7 +81,7 @@ export class EthereumBackend { async checkAuthorization(appInfo: AuthorizationRequest, method: 'checkKmsAuthorization' | 'checkAppAuthorization'): Promise { const result = await this.kmsContract[method]({ appInfo: this.contractAppInfo(appInfo), - requestedAppEpoch: appInfo.requestedAppEpoch ?? 0, + requestedAppEpoch: { isSet: appInfo.requestedAppEpoch !== undefined, value: appInfo.requestedAppEpoch ?? 0 }, requestedRootEpoch: { isSet: appInfo.requestedRootEpoch !== undefined, value: appInfo.requestedRootEpoch ?? 0 }, }); return { @@ -90,6 +90,7 @@ export class EthereumBackend { gatewayAppId: result.gatewayAppId, currentRootEpoch: ethers.toNumber(result.currentRootEpoch), currentRootEpochHash: result.currentRootEpochHash, + currentAppEpoch: ethers.toNumber(result.currentAppEpoch), }; } diff --git a/dstack/kms/auth-eth/src/main.test.ts b/dstack/kms/auth-eth/src/main.test.ts index da54da869..24d7c3924 100644 --- a/dstack/kms/auth-eth/src/main.test.ts +++ b/dstack/kms/auth-eth/src/main.test.ts @@ -96,7 +96,7 @@ describe('Server', () => { it('preserves the complete decision, including the epoch on rejection', async () => { const decision = { isAllowed: false, rejectionReason: 'App not registered', gatewayAppId: 'gateway', - currentRootEpoch: 1, currentRootEpochHash: '0x' + '11'.repeat(32), + currentAppEpoch: 0, currentRootEpoch: 1, currentRootEpochHash: '0x' + '11'.repeat(32), }; app.ethereum.checkAuthorization = jest.fn().mockResolvedValue(decision); const response = await app.inject({ method: 'POST', url: '/checkAppAuthorization', payload: appRequest }); diff --git a/dstack/kms/auth-eth/src/server.ts b/dstack/kms/auth-eth/src/server.ts index af4cd74d0..bc996330b 100644 --- a/dstack/kms/auth-eth/src/server.ts +++ b/dstack/kms/auth-eth/src/server.ts @@ -127,12 +127,13 @@ export async function build(): Promise { response: { 200: { type: 'object', - required: ['isAllowed', 'rejectionReason', 'gatewayAppId', 'currentRootEpoch', 'currentRootEpochHash'], + required: ['isAllowed', 'rejectionReason', 'gatewayAppId', 'currentRootEpoch', 'currentRootEpochHash', 'currentAppEpoch'], properties: { isAllowed: { type: 'boolean' }, rejectionReason: { type: 'string' }, gatewayAppId: { type: 'string' }, currentRootEpoch: { type: 'integer', minimum: 0 }, + currentAppEpoch: { type: 'integer', minimum: 0 }, currentRootEpochHash: hex(32, 'Current root epoch chain hash'), }, }, diff --git a/dstack/kms/auth-eth/src/types.ts b/dstack/kms/auth-eth/src/types.ts index 3a8ea1cbc..6171f1b01 100644 --- a/dstack/kms/auth-eth/src/types.ts +++ b/dstack/kms/auth-eth/src/types.ts @@ -39,6 +39,7 @@ export interface AuthorizationResult { gatewayAppId: string; currentRootEpoch: number; currentRootEpochHash: string; + currentAppEpoch: number; } export interface AuthorizationRequest extends Omit { diff --git a/dstack/kms/auth-eth/test/DstackKms.t.sol b/dstack/kms/auth-eth/test/DstackKms.t.sol index adf734784..5f185bab3 100644 --- a/dstack/kms/auth-eth/test/DstackKms.t.sol +++ b/dstack/kms/auth-eth/test/DstackKms.t.sol @@ -465,7 +465,7 @@ contract DstackKmsTest is Test { returns (DstackKms.AuthorizationResult memory) { DstackKms.AuthorizationRequest memory request = - DstackKms.AuthorizationRequest(info, 0, DstackKms.OptionalEpoch(false, 0)); + DstackKms.AuthorizationRequest(info, DstackKms.OptionalEpoch(false, 0), DstackKms.OptionalEpoch(false, 0)); return isKms ? kms.checkKmsAuthorization(request) : kms.checkAppAuthorization(request); } @@ -480,7 +480,7 @@ contract DstackKmsTest is Test { vm.stopPrank(); DstackKms.AuthorizationRequest memory request = - DstackKms.AuthorizationRequest(info, 0, DstackKms.OptionalEpoch(true, 1)); + DstackKms.AuthorizationRequest(info, DstackKms.OptionalEpoch(false, 0), DstackKms.OptionalEpoch(true, 1)); DstackKms.AuthorizationResult memory result = kms.checkKmsAuthorization(request); assertFalse(result.isAllowed); assertEq(result.rejectionReason, "Requested root epoch is not active"); @@ -537,7 +537,7 @@ contract DstackKmsTest is Test { assertFalse(result.isAllowed); assertEq(result.rejectionReason, "Authorization kind does not match app ID"); DstackKms.AuthorizationRequest memory appRequest = - DstackKms.AuthorizationRequest(info, 1, DstackKms.OptionalEpoch(false, 0)); + DstackKms.AuthorizationRequest(info, DstackKms.OptionalEpoch(true, 1), DstackKms.OptionalEpoch(false, 0)); result = kms.checkAppAuthorization(appRequest); assertFalse(result.isAllowed); assertEq(result.rejectionReason, "Requested app epoch must be zero"); diff --git a/dstack/kms/auth-mock/index.ts b/dstack/kms/auth-mock/index.ts index d5b21c6c2..66abd04fb 100644 --- a/dstack/kms/auth-mock/index.ts +++ b/dstack/kms/auth-mock/index.ts @@ -115,6 +115,7 @@ class MockBackend { async checkAuthorization(appInfo: AuthorizationRequest, isKms: boolean) { const current = getRootEpoch(); const epoch = { + currentAppEpoch: 0, currentRootEpoch: current.rootEpoch ?? 0, currentRootEpochHash: current.rootEpochHash ?? '0x' + '00'.repeat(32), }; diff --git a/dstack/kms/auth-mock/openapi.json b/dstack/kms/auth-mock/openapi.json index e26c8ead2..90d833adc 100644 --- a/dstack/kms/auth-mock/openapi.json +++ b/dstack/kms/auth-mock/openapi.json @@ -387,7 +387,8 @@ "rejectionReason", "gatewayAppId", "currentRootEpoch", - "currentRootEpochHash" + "currentRootEpochHash", + "currentAppEpoch" ], "properties": { "isAllowed": { @@ -408,6 +409,11 @@ "type": "string", "pattern": "^0x[0-9a-fA-F]{64}$", "description": "Zero when no epoch is registered; authorization is denied." + }, + "currentAppEpoch": { + "type": "integer", + "minimum": 0, + "description": "Current application epoch; zero until app rotation is supported." } } }, diff --git a/dstack/kms/auth-simple/index.ts b/dstack/kms/auth-simple/index.ts index fa8747628..c2c07ea69 100644 --- a/dstack/kms/auth-simple/index.ts +++ b/dstack/kms/auth-simple/index.ts @@ -162,6 +162,7 @@ class ConfigBackend { const config = this.loadConfig(); const current = currentRootEpoch(config); const epoch = { + currentAppEpoch: 0, currentRootEpoch: current.rootEpoch ?? 0, currentRootEpochHash: current.rootEpochHash ?? '0x' + '00'.repeat(32), }; diff --git a/dstack/kms/src/main_service.rs b/dstack/kms/src/main_service.rs index f5f9c4b82..969b5b05c 100644 --- a/dstack/kms/src/main_service.rs +++ b/dstack/kms/src/main_service.rs @@ -297,6 +297,9 @@ 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 { @@ -1505,6 +1508,7 @@ mod tests { rejection_reason: String::new(), gateway_app_id: String::new(), current_root_epoch: epoch, + current_app_epoch: 0, current_root_epoch_hash: hex::encode(state.keyring.chain_hashes[epoch as usize]), }; assert_eq!(state.authorized_epoch(&result(0), None).unwrap(), 0); diff --git a/dstack/kms/src/main_service/upgrade_authority.rs b/dstack/kms/src/main_service/upgrade_authority.rs index 89407de23..92cdc0e49 100644 --- a/dstack/kms/src/main_service/upgrade_authority.rs +++ b/dstack/kms/src/main_service/upgrade_authority.rs @@ -93,6 +93,7 @@ pub(crate) struct AuthorizationResult { pub rejection_reason: String, pub gateway_app_id: String, pub current_root_epoch: u32, + pub current_app_epoch: u32, pub current_root_epoch_hash: String, } @@ -293,6 +294,7 @@ impl AuthApi { rejection_reason: String::new(), gateway_app_id: dev.gateway_app_id.clone(), current_root_epoch: 0, + current_app_epoch: 0, current_root_epoch_hash: hex::encode([0; 32]), }), AuthApi::Webhook { webhook } => { @@ -697,7 +699,7 @@ pub(crate) mod tests { async fn authorization_returns_the_epoch_with_each_decision() { let response = |allowed: bool, epoch: u32| { format!( - r#"{{"isAllowed":{allowed},"rejectionReason":"","gatewayAppId":"gw", "currentRootEpoch":{epoch},"currentRootEpochHash":"0x{}"}}"#, + r#"{{"isAllowed":{allowed},"rejectionReason":"","gatewayAppId":"gw", "currentRootEpoch":{epoch},"currentAppEpoch":0,"currentRootEpochHash":"0x{}"}}"#, "11".repeat(32) ) }; From 2f0968e33e94a3b35b93a5337d44c70f4c0a8bfb Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Sat, 7 Nov 2026 05:51:10 -0800 Subject: [PATCH 024/166] refactor(kms): split authorization by caller kind Signed-off-by: Kevin Wang --- dstack/kms/auth-eth/README.md | 2 +- dstack/kms/auth-eth/contracts/DstackApp.sol | 2 +- dstack/kms/auth-eth/contracts/DstackKms.sol | 30 ++++------------ dstack/kms/auth-eth/contracts/IAppAuth.sol | 4 +-- .../kms/auth-eth/docs/formal-verification.md | 2 +- dstack/kms/auth-eth/docs/specification.md | 6 ++-- dstack/kms/auth-eth/script/Query.s.sol | 4 +-- dstack/kms/auth-eth/src/ethereum.ts | 6 ++-- dstack/kms/auth-eth/src/main.test.ts | 16 ++++----- dstack/kms/auth-eth/src/server.ts | 2 +- dstack/kms/auth-eth/test/DstackApp.t.sol | 8 ++--- .../auth-eth/test/DstackKms.symbolic.t.sol | 16 ++++----- dstack/kms/auth-eth/test/DstackKms.t.sol | 34 +++++++++---------- .../auth-eth/test/UpgradesWithPlugin.t.sol | 6 ++-- dstack/kms/auth-mock/README.md | 2 +- dstack/kms/auth-mock/index.test.ts | 2 +- dstack/kms/auth-mock/index.ts | 8 ++--- dstack/kms/auth-mock/openapi.json | 8 ++--- dstack/kms/auth-simple/README.md | 2 +- dstack/kms/auth-simple/index.test.ts | 6 ++-- dstack/kms/auth-simple/index.ts | 7 ++-- 21 files changed, 74 insertions(+), 99 deletions(-) diff --git a/dstack/kms/auth-eth/README.md b/dstack/kms/auth-eth/README.md index 02e488b4b..47f880315 100644 --- a/dstack/kms/auth-eth/README.md +++ b/dstack/kms/auth-eth/README.md @@ -174,7 +174,7 @@ Documentation: https://book.getfoundry.sh/ ### Authorization endpoints -`POST /checkKmsAuthorization` and `POST /checkAppAuthorization` return `AuthorizationResult`, including +`POST /authorizeKms` and `POST /authorizeApp` return `AuthorizationResult`, including `rejectionReason`, `gatewayAppId`, `currentRootEpoch`, and `currentRootEpochHash`. The selected endpoint determines the policy; the backend checks that the verified `appId` matches that scope. `requestedAppEpoch` defaults to 0; nonzero values are currently diff --git a/dstack/kms/auth-eth/contracts/DstackApp.sol b/dstack/kms/auth-eth/contracts/DstackApp.sol index 7d90f29ef..d55a6c883 100644 --- a/dstack/kms/auth-eth/contracts/DstackApp.sol +++ b/dstack/kms/auth-eth/contracts/DstackApp.sol @@ -188,7 +188,7 @@ contract DstackApp is } // Check if an app is allowed to boot - function isAppAllowed(IAppAuth.AppBootInfo calldata bootInfo) + function isAppAllowed(IAppAuth.AppInfo calldata bootInfo) external view override diff --git a/dstack/kms/auth-eth/contracts/DstackKms.sol b/dstack/kms/auth-eth/contracts/DstackKms.sol index a65fce682..e14926087 100644 --- a/dstack/kms/auth-eth/contracts/DstackKms.sol +++ b/dstack/kms/auth-eth/contracts/DstackKms.sol @@ -14,18 +14,6 @@ import "@openzeppelin/contracts-upgradeable/utils/introspection/ERC165Upgradeabl import "@openzeppelin/contracts/proxy/ERC1967/ERC1967Proxy.sol"; contract DstackKms is Initializable, Ownable2StepUpgradeable, UUPSUpgradeable, ERC165Upgradeable, IAppAuth { - struct AppInfo { - address appId; - bytes32 composeHash; - address instanceId; - bytes32 deviceId; - bytes32 mrAggregated; - bytes32 mrSystem; - bytes32 osImageHash; - string tcbStatus; - string[] advisoryIds; - } - struct OptionalEpoch { bool isSet; uint256 value; @@ -405,7 +393,7 @@ contract DstackKms is Initializable, Ownable2StepUpgradeable, UUPSUpgradeable, E /// 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 checkKmsAuthorization(AuthorizationRequest calldata request) + function authorizeKms(AuthorizationRequest calldata request) external view returns (AuthorizationResult memory) @@ -413,7 +401,7 @@ contract DstackKms is Initializable, Ownable2StepUpgradeable, UUPSUpgradeable, E return _checkAuthorization(request, true); } - function checkAppAuthorization(AuthorizationRequest calldata request) + function authorizeApp(AuthorizationRequest calldata request) external view returns (AuthorizationResult memory) @@ -445,11 +433,7 @@ contract DstackKms is Initializable, Ownable2StepUpgradeable, UUPSUpgradeable, E return result; } AppInfo calldata appInfo = request.appInfo; - if (isKms != (appInfo.appId == address(this))) { - result.rejectionReason = "Authorization kind does not match app ID"; - return result; - } - AppBootInfo memory bootInfo = AppBootInfo({ + AppInfo memory bootInfo = AppInfo({ appId: appInfo.appId, composeHash: appInfo.composeHash, instanceId: appInfo.instanceId, @@ -468,11 +452,11 @@ contract DstackKms is Initializable, Ownable2StepUpgradeable, UUPSUpgradeable, E } // Function to check if KMS is allowed to boot - function isKmsAllowed(AppBootInfo calldata bootInfo) external view returns (bool isAllowed, string memory reason) { + function isKmsAllowed(AppInfo calldata bootInfo) external view returns (bool isAllowed, string memory reason) { return _isKmsAllowed(bootInfo); } - function _isKmsAllowed(AppBootInfo memory bootInfo) private view returns (bool, string memory) { + function _isKmsAllowed(AppInfo memory bootInfo) private view returns (bool, string memory) { // Check if the TCB status is up to date if (keccak256(abi.encodePacked(bootInfo.tcbStatus)) != keccak256(abi.encodePacked("UpToDate"))) { return (false, "TCB status is not up to date"); @@ -497,7 +481,7 @@ contract DstackKms is Initializable, Ownable2StepUpgradeable, UUPSUpgradeable, E } // Function to check if an app is allowed to boot - function isAppAllowed(AppBootInfo calldata bootInfo) + function isAppAllowed(AppInfo calldata bootInfo) external view override @@ -506,7 +490,7 @@ contract DstackKms is Initializable, Ownable2StepUpgradeable, UUPSUpgradeable, E return _isAppAllowed(bootInfo); } - function _isAppAllowed(AppBootInfo memory bootInfo) private view returns (bool, string memory) { + function _isAppAllowed(AppInfo memory bootInfo) private view returns (bool, string memory) { // Check if app is registered if (!registeredApps[bootInfo.appId]) { return (false, "App not registered"); diff --git a/dstack/kms/auth-eth/contracts/IAppAuth.sol b/dstack/kms/auth-eth/contracts/IAppAuth.sol index 56b768971..896c60954 100644 --- a/dstack/kms/auth-eth/contracts/IAppAuth.sol +++ b/dstack/kms/auth-eth/contracts/IAppAuth.sol @@ -33,7 +33,7 @@ interface IAppAuth is IERC165 { * @param tcbStatus Trusted Computing Base status * @param advisoryIds Array of security advisory identifiers */ - struct AppBootInfo { + struct AppInfo { address appId; bytes32 composeHash; address instanceId; @@ -52,5 +52,5 @@ interface IAppAuth is IERC165 { * @return isAllowed True if the app is authorized to boot, false otherwise * @return reason Human-readable reason for the decision (empty if allowed) */ - function isAppAllowed(AppBootInfo calldata bootInfo) external view returns (bool isAllowed, string memory reason); + function isAppAllowed(AppInfo calldata bootInfo) external view returns (bool isAllowed, string memory reason); } diff --git a/dstack/kms/auth-eth/docs/formal-verification.md b/dstack/kms/auth-eth/docs/formal-verification.md index 67d776c4a..e9987e409 100644 --- a/dstack/kms/auth-eth/docs/formal-verification.md +++ b/dstack/kms/auth-eth/docs/formal-verification.md @@ -124,7 +124,7 @@ Single-call (`test/DstackKms.symbolic.t.sol`): - [x] `check_RegisterApp_AnyCallerCanRegisterNonZeroAddress` + `_RejectsZeroAddress` — codifies the permissionless-by-design behavior; see "Findings" - [x] `check_IsAppAllowed_RejectsUnregisteredApp` / `_RejectsUnknownOsImage` - — fully symbolic `AppBootInfo`; the failing gate produces the + — fully symbolic `AppInfo`; the failing gate produces the right rejection reason without delegating - [x] `check_IsAppAllowed_DelegatesFaithfully` — when both KMS gates pass, the outer return equals the registered `IAppAuth`'s return. diff --git a/dstack/kms/auth-eth/docs/specification.md b/dstack/kms/auth-eth/docs/specification.md index 41cd2b885..01f12bf22 100644 --- a/dstack/kms/auth-eth/docs/specification.md +++ b/dstack/kms/auth-eth/docs/specification.md @@ -125,7 +125,7 @@ Signature: `deployAndRegisterApp(address initialOwner, bool disableUpgrades, boo - Semantics: equivalent to the 6-arg overload with `requireTcbUpToDate = false`. - (verified: `DstackKmsSymbolicTest.check_DeployAndRegisterApp5Arg_DefaultsTcbToFalse`) -### 3.5 `DstackKms.isAppAllowed(AppBootInfo b) external view returns (bool, string)` +### 3.5 `DstackKms.isAppAllowed(AppInfo b) external view returns (bool, string)` Decision (in order): @@ -147,7 +147,7 @@ Decision (in order): behavior of an adversarial registered app; those propagate to the outer call by EVM semantics (§5.1). -### 3.6 `DstackKms.isKmsAllowed(AppBootInfo b) external view returns (bool, string)` +### 3.6 `DstackKms.isKmsAllowed(AppInfo b) external view returns (bool, string)` Decision (in order): @@ -200,7 +200,7 @@ For each of `setGatewayAppId`, - (verified: `DstackAppSymbolicTest.check_Initialize5Arg_DefaultsTcbToFalse`, `check_Initialize6Arg_HonorsTcbFlag`) -### 3.9 `DstackApp.isAppAllowed(AppBootInfo b) external view returns (bool, string)` +### 3.9 `DstackApp.isAppAllowed(AppInfo b) external view returns (bool, string)` Decision (in order, all evaluated under `b.appId == address(this)` by convention): diff --git a/dstack/kms/auth-eth/script/Query.s.sol b/dstack/kms/auth-eth/script/Query.s.sol index 32faedec9..933d13578 100644 --- a/dstack/kms/auth-eth/script/Query.s.sol +++ b/dstack/kms/auth-eth/script/Query.s.sol @@ -124,7 +124,7 @@ contract CheckKmsAllowed is BaseQueryScript { require(address(kms) != address(0), "KMS_CONTRACT_ADDR not set"); // Read boot info from environment - IAppAuth.AppBootInfo memory bootInfo = IAppAuth.AppBootInfo({ + IAppAuth.AppInfo memory bootInfo = IAppAuth.AppInfo({ appId: vm.envAddress("APP_ID"), composeHash: vm.envBytes32("COMPOSE_HASH"), instanceId: vm.envOr("INSTANCE_ID", address(0)), @@ -150,7 +150,7 @@ contract CheckAppAllowed is BaseQueryScript { require(address(kms) != address(0), "KMS_CONTRACT_ADDR not set"); // Read boot info from environment - IAppAuth.AppBootInfo memory bootInfo = IAppAuth.AppBootInfo({ + IAppAuth.AppInfo memory bootInfo = IAppAuth.AppInfo({ appId: vm.envAddress("APP_ID"), composeHash: vm.envBytes32("COMPOSE_HASH"), instanceId: vm.envOr("INSTANCE_ID", address(0)), diff --git a/dstack/kms/auth-eth/src/ethereum.ts b/dstack/kms/auth-eth/src/ethereum.ts index 99ce15c42..8617ffe08 100644 --- a/dstack/kms/auth-eth/src/ethereum.ts +++ b/dstack/kms/auth-eth/src/ethereum.ts @@ -7,8 +7,8 @@ import { BootInfo, BootResponse, RootEpochInfo, AuthorizationResult, Authorizati // Minimal ABI for DstackKms contract const DSTACK_KMS_ABI = [ - "function checkKmsAuthorization(((address appId,bytes32 composeHash,address instanceId,bytes32 deviceId,bytes32 mrAggregated,bytes32 mrSystem,bytes32 osImageHash,string tcbStatus,string[] advisoryIds) appInfo,(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 checkAppAuthorization(((address appId,bytes32 composeHash,address instanceId,bytes32 deviceId,bytes32 mrAggregated,bytes32 mrSystem,bytes32 osImageHash,string tcbStatus,string[] advisoryIds) appInfo,(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 authorizeKms(((address appId,bytes32 composeHash,address instanceId,bytes32 deviceId,bytes32 mrAggregated,bytes32 mrSystem,bytes32 osImageHash,string tcbStatus,string[] advisoryIds) appInfo,(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,(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)", "function isKmsAllowed((address appId,bytes32 composeHash,address instanceId,bytes32 deviceId,bytes32 mrAggregated,bytes32 mrSystem,bytes32 osImageHash,string tcbStatus,string[] advisoryIds) bootInfo) view returns (bool, string)", "function gatewayAppId() view returns (string)", @@ -78,7 +78,7 @@ export class EthereumBackend { }; } - async checkAuthorization(appInfo: AuthorizationRequest, method: 'checkKmsAuthorization' | 'checkAppAuthorization'): Promise { + async checkAuthorization(appInfo: AuthorizationRequest, method: 'authorizeKms' | 'authorizeApp'): Promise { const result = await this.kmsContract[method]({ appInfo: this.contractAppInfo(appInfo), requestedAppEpoch: { isSet: appInfo.requestedAppEpoch !== undefined, value: appInfo.requestedAppEpoch ?? 0 }, diff --git a/dstack/kms/auth-eth/src/main.test.ts b/dstack/kms/auth-eth/src/main.test.ts index 24d7c3924..a8d035749 100644 --- a/dstack/kms/auth-eth/src/main.test.ts +++ b/dstack/kms/auth-eth/src/main.test.ts @@ -91,7 +91,7 @@ describe('Server', () => { }); }); - describe('POST /checkAppAuthorization', () => { + describe('POST /authorizeApp', () => { const appRequest = { ...mockBootInfo, requestedAppEpoch: 0 }; it('preserves the complete decision, including the epoch on rejection', async () => { const decision = { @@ -99,20 +99,20 @@ describe('Server', () => { currentAppEpoch: 0, currentRootEpoch: 1, currentRootEpochHash: '0x' + '11'.repeat(32), }; app.ethereum.checkAuthorization = jest.fn().mockResolvedValue(decision); - const response = await app.inject({ method: 'POST', url: '/checkAppAuthorization', payload: appRequest }); + const response = await app.inject({ method: 'POST', url: '/authorizeApp', payload: appRequest }); expect(response.statusCode).toBe(200); expect(JSON.parse(response.payload)).toEqual(decision); - expect(app.ethereum.checkAuthorization).toHaveBeenCalledWith(appRequest, 'checkAppAuthorization'); - await app.inject({ method: 'POST', url: '/checkKmsAuthorization', payload: appRequest }); - expect(app.ethereum.checkAuthorization).toHaveBeenLastCalledWith(appRequest, 'checkKmsAuthorization'); - await app.inject({ method: 'POST', url: '/checkAppAuthorization', payload: { ...appRequest, requestedRootEpoch: 0 } }); - expect(app.ethereum.checkAuthorization).toHaveBeenLastCalledWith({ ...appRequest, requestedRootEpoch: 0 }, 'checkAppAuthorization'); + expect(app.ethereum.checkAuthorization).toHaveBeenCalledWith(appRequest, 'authorizeApp'); + await app.inject({ method: 'POST', url: '/authorizeKms', payload: appRequest }); + expect(app.ethereum.checkAuthorization).toHaveBeenLastCalledWith(appRequest, 'authorizeKms'); + await app.inject({ method: 'POST', url: '/authorizeApp', payload: { ...appRequest, requestedRootEpoch: 0 } }); + expect(app.ethereum.checkAuthorization).toHaveBeenLastCalledWith({ ...appRequest, requestedRootEpoch: 0 }, 'authorizeApp'); }); it('fails closed when the new contract method is unavailable', async () => { jest.clearAllMocks(); app.ethereum.checkAuthorization = jest.fn().mockRejectedValue(new Error('missing contract method')); - const response = await app.inject({ method: 'POST', url: '/checkAppAuthorization', payload: appRequest }); + const response = await app.inject({ method: 'POST', url: '/authorizeApp', payload: appRequest }); expect(response.statusCode).toBe(500); expect(app.ethereum.checkBoot).not.toHaveBeenCalled(); }); diff --git a/dstack/kms/auth-eth/src/server.ts b/dstack/kms/auth-eth/src/server.ts index bc996330b..c7de71c7b 100644 --- a/dstack/kms/auth-eth/src/server.ts +++ b/dstack/kms/auth-eth/src/server.ts @@ -120,7 +120,7 @@ export async function build(): Promise { } }); - for (const method of ['checkKmsAuthorization', 'checkAppAuthorization'] as const) { + for (const method of ['authorizeKms', 'authorizeApp'] as const) { server.post<{ Body: AuthorizationRequest; Reply: AuthorizationResult }>(`/${method}`, { schema: { body: { $ref: 'authorizationRequest#' }, diff --git a/dstack/kms/auth-eth/test/DstackApp.t.sol b/dstack/kms/auth-eth/test/DstackApp.t.sol index 9f590d2e3..6776e0b2d 100644 --- a/dstack/kms/auth-eth/test/DstackApp.t.sol +++ b/dstack/kms/auth-eth/test/DstackApp.t.sol @@ -170,7 +170,7 @@ contract DstackAppTest is Test { app.addComposeHash(composeHash); vm.stopPrank(); - IAppAuth.AppBootInfo memory bootInfo = IAppAuth.AppBootInfo({ + IAppAuth.AppInfo memory bootInfo = IAppAuth.AppInfo({ appId: address(app), composeHash: composeHash, instanceId: address(0), @@ -212,7 +212,7 @@ contract DstackAppTest is Test { app.addComposeHash(allowedHash); vm.stopPrank(); - IAppAuth.AppBootInfo memory bootInfo = IAppAuth.AppBootInfo({ + IAppAuth.AppInfo memory bootInfo = IAppAuth.AppInfo({ appId: address(app), composeHash: unallowedHash, instanceId: address(0), @@ -270,7 +270,7 @@ contract DstackAppTest is Test { string[] memory advisoryIds = new string[](1); advisoryIds[0] = "INTEL-SA-00614"; - IAppAuth.AppBootInfo memory bootInfo = IAppAuth.AppBootInfo({ + IAppAuth.AppInfo memory bootInfo = IAppAuth.AppInfo({ appId: address(app), composeHash: composeHash, instanceId: address(0), @@ -318,7 +318,7 @@ contract DstackAppTest is Test { assertTrue(tcbApp.allowedDeviceIds(deviceId)); assertTrue(tcbApp.allowedComposeHashes(composeHash)); - IAppAuth.AppBootInfo memory bootInfo = IAppAuth.AppBootInfo({ + IAppAuth.AppInfo memory bootInfo = IAppAuth.AppInfo({ appId: address(tcbApp), composeHash: composeHash, instanceId: address(0), diff --git a/dstack/kms/auth-eth/test/DstackKms.symbolic.t.sol b/dstack/kms/auth-eth/test/DstackKms.symbolic.t.sol index 47b79af2c..ad9f7ac63 100644 --- a/dstack/kms/auth-eth/test/DstackKms.symbolic.t.sol +++ b/dstack/kms/auth-eth/test/DstackKms.symbolic.t.sol @@ -32,7 +32,7 @@ contract MockConfigurableApp { _ret = ret; } - function isAppAllowed(IAppAuth.AppBootInfo calldata) external view returns (bool, string memory) { + function isAppAllowed(IAppAuth.AppInfo calldata) external view returns (bool, string memory) { return (_ret, ""); } } @@ -41,7 +41,7 @@ contract MockConfigurableApp { /// to verify that a malicious registered contract's revert /// propagates out of `kms.isAppAllowed` rather than being swallowed. contract MockRevertingApp { - function isAppAllowed(IAppAuth.AppBootInfo calldata) external pure returns (bool, string memory) { + function isAppAllowed(IAppAuth.AppInfo calldata) external pure returns (bool, string memory) { revert("malicious app revert"); } } @@ -87,14 +87,14 @@ contract DstackKmsSymbolicTest is Test { // the (false, reason) return without delegating. // --------------------------------------------------------------- - function check_IsAppAllowed_RejectsUnregisteredApp(IAppAuth.AppBootInfo calldata bootInfo) external view { + function check_IsAppAllowed_RejectsUnregisteredApp(IAppAuth.AppInfo calldata bootInfo) external view { // No apps registered in setUp(), so any bootInfo.appId is rejected. (bool allowed, string memory reason) = kms.isAppAllowed(bootInfo); assert(!allowed); assert(keccak256(bytes(reason)) == keccak256(bytes("App not registered"))); } - function check_IsAppAllowed_RejectsUnknownOsImage(address appId, IAppAuth.AppBootInfo calldata bootInfo) external { + function check_IsAppAllowed_RejectsUnknownOsImage(address appId, IAppAuth.AppInfo calldata bootInfo) external { vm.assume(appId != address(0)); // Register the appId so the registration check passes. @@ -128,7 +128,7 @@ contract DstackKmsSymbolicTest is Test { kms.addOsImageHash(osImageHash); vm.stopPrank(); - IAppAuth.AppBootInfo memory b = IAppAuth.AppBootInfo({ + IAppAuth.AppInfo memory b = IAppAuth.AppInfo({ appId: address(mock), composeHash: bytes32(0), instanceId: address(0), @@ -158,7 +158,7 @@ contract DstackKmsSymbolicTest is Test { kms.addOsImageHash(osImageHash); vm.stopPrank(); - IAppAuth.AppBootInfo memory b = IAppAuth.AppBootInfo({ + IAppAuth.AppInfo memory b = IAppAuth.AppInfo({ appId: address(mock), composeHash: bytes32(0), instanceId: address(0), @@ -192,7 +192,7 @@ contract DstackKmsSymbolicTest is Test { vm.prank(OWNER); kms.addOsImageHash(osImageHash); - IAppAuth.AppBootInfo memory b = IAppAuth.AppBootInfo({ + IAppAuth.AppInfo memory b = IAppAuth.AppInfo({ appId: address(0), composeHash: bytes32(0), instanceId: address(0), @@ -222,7 +222,7 @@ contract DstackKmsSymbolicTest is Test { vm.stopPrank(); // kmsAllowedDeviceIds is left empty. - IAppAuth.AppBootInfo memory b = IAppAuth.AppBootInfo({ + IAppAuth.AppInfo memory b = IAppAuth.AppInfo({ appId: address(0), composeHash: bytes32(0), instanceId: address(0), diff --git a/dstack/kms/auth-eth/test/DstackKms.t.sol b/dstack/kms/auth-eth/test/DstackKms.t.sol index 5f185bab3..a34094469 100644 --- a/dstack/kms/auth-eth/test/DstackKms.t.sol +++ b/dstack/kms/auth-eth/test/DstackKms.t.sol @@ -442,8 +442,8 @@ contract DstackKmsTest is Test { vm.stopPrank(); } - function authorizationInfo(address appId) internal pure returns (DstackKms.AppInfo memory) { - return DstackKms.AppInfo({ + function authorizationInfo(address appId) internal pure returns (IAppAuth.AppInfo memory) { + return IAppAuth.AppInfo({ appId: appId, composeHash: bytes32(uint256(789)), instanceId: address(0), @@ -457,7 +457,7 @@ contract DstackKmsTest is Test { } function authorize( - DstackKms.AppInfo memory info, + IAppAuth.AppInfo memory info, bool isKms ) internal @@ -466,11 +466,11 @@ contract DstackKmsTest is Test { { DstackKms.AuthorizationRequest memory request = DstackKms.AuthorizationRequest(info, DstackKms.OptionalEpoch(false, 0), DstackKms.OptionalEpoch(false, 0)); - return isKms ? kms.checkKmsAuthorization(request) : kms.checkAppAuthorization(request); + return isKms ? kms.authorizeKms(request) : kms.authorizeApp(request); } function test_CheckAuthorization_OptionalEpochDistinguishesZeroFromDefault() public { - DstackKms.AppInfo memory info = authorizationInfo(address(kms)); + IAppAuth.AppInfo memory info = authorizationInfo(address(kms)); vm.startPrank(owner); kms.registerRootEpoch(ca(0xa0), k256(0xb0), "https://example.com/epoch-0"); kms.registerRootEpoch(ca(0xa1), k256(0xb1), "https://example.com/epoch-1"); @@ -481,19 +481,19 @@ contract DstackKmsTest is Test { DstackKms.AuthorizationRequest memory request = DstackKms.AuthorizationRequest(info, DstackKms.OptionalEpoch(false, 0), DstackKms.OptionalEpoch(true, 1)); - DstackKms.AuthorizationResult memory result = kms.checkKmsAuthorization(request); + DstackKms.AuthorizationResult memory result = kms.authorizeKms(request); assertFalse(result.isAllowed); assertEq(result.rejectionReason, "Requested root epoch is not active"); bytes32 h1 = chainHash(chainHash(bytes32(0), ca(0xa0), k256(0xb0)), ca(0xa1), k256(0xb1)); vm.prank(owner); kms.activateRootEpoch(1, h1); - assertTrue(kms.checkKmsAuthorization(request).isAllowed); + assertTrue(kms.authorizeKms(request).isAllowed); request.requestedRootEpoch.value = 0; - assertTrue(kms.checkKmsAuthorization(request).isAllowed); + assertTrue(kms.authorizeKms(request).isAllowed); request.requestedRootEpoch.value = type(uint256).max; - assertFalse(kms.checkKmsAuthorization(request).isAllowed); + assertFalse(kms.authorizeKms(request).isAllowed); request.requestedRootEpoch.isSet = false; - result = kms.checkKmsAuthorization(request); + result = kms.authorizeKms(request); assertTrue(result.isAllowed); assertEq(result.currentRootEpoch, 1); assertEq(result.currentRootEpochHash, h1); @@ -508,7 +508,7 @@ contract DstackKmsTest is Test { } function test_CheckAuthorization_ValidatesRequestedPolicy() public { - DstackKms.AppInfo memory info = authorizationInfo(address(kms)); + IAppAuth.AppInfo memory info = authorizationInfo(address(kms)); vm.startPrank(owner); kms.registerRootEpoch(ca(0xa0), k256(0xb0), "https://example.com/epoch-0"); kms.setGatewayAppId("gateway"); @@ -535,17 +535,17 @@ contract DstackKmsTest is Test { assertTrue(result.isAllowed); result = authorize(info, true); assertFalse(result.isAllowed); - assertEq(result.rejectionReason, "Authorization kind does not match app ID"); + assertEq(result.rejectionReason, "Aggregated MR not allowed"); DstackKms.AuthorizationRequest memory appRequest = DstackKms.AuthorizationRequest(info, DstackKms.OptionalEpoch(true, 1), DstackKms.OptionalEpoch(false, 0)); - result = kms.checkAppAuthorization(appRequest); + result = kms.authorizeApp(appRequest); assertFalse(result.isAllowed); assertEq(result.rejectionReason, "Requested app epoch must be zero"); assertEq(result.currentRootEpochHash, chainHash(bytes32(0), ca(0xa0), k256(0xb0))); info.appId = address(kms); result = authorize(info, false); assertFalse(result.isAllowed); - assertEq(result.rejectionReason, "Authorization kind does not match app ID"); + assertEq(result.rejectionReason, "App not registered"); result = authorize(info, true); assertFalse(result.isAllowed); assertEq(result.rejectionReason, "Aggregated MR not allowed"); @@ -579,7 +579,7 @@ contract DstackKmsTest is Test { kms.addOsImageHash(osImageHash); vm.stopPrank(); - IAppAuth.AppBootInfo memory bootInfo = IAppAuth.AppBootInfo({ + IAppAuth.AppInfo memory bootInfo = IAppAuth.AppInfo({ appId: address(0), composeHash: bytes32(0), instanceId: address(0), @@ -608,7 +608,7 @@ contract DstackKmsTest is Test { kms.addOsImageHash(osImageHash); vm.stopPrank(); - IAppAuth.AppBootInfo memory bootInfo = IAppAuth.AppBootInfo({ + IAppAuth.AppInfo memory bootInfo = IAppAuth.AppInfo({ appId: address(0), composeHash: bytes32(0), instanceId: address(0), @@ -636,7 +636,7 @@ contract DstackKmsTest is Test { address appId = kms.deployAndRegisterApp(user, false, false, deviceId, composeHash); vm.stopPrank(); - IAppAuth.AppBootInfo memory bootInfo = IAppAuth.AppBootInfo({ + IAppAuth.AppInfo memory bootInfo = IAppAuth.AppInfo({ appId: appId, composeHash: composeHash, instanceId: address(0), diff --git a/dstack/kms/auth-eth/test/UpgradesWithPlugin.t.sol b/dstack/kms/auth-eth/test/UpgradesWithPlugin.t.sol index 0cb05013d..770482fba 100644 --- a/dstack/kms/auth-eth/test/UpgradesWithPlugin.t.sol +++ b/dstack/kms/auth-eth/test/UpgradesWithPlugin.t.sol @@ -284,7 +284,7 @@ contract UpgradesWithPluginTest is Test { assertTrue(app.allowAnyDevice()); assertTrue(app.allowedComposeHashes(composeHash)); - IAppAuth.AppBootInfo memory bootInfo = IAppAuth.AppBootInfo({ + IAppAuth.AppInfo memory bootInfo = IAppAuth.AppInfo({ appId: appId, composeHash: composeHash, instanceId: address(0), @@ -335,7 +335,7 @@ contract UpgradesWithPluginTest is Test { assertFalse(app.requireTcbUpToDate()); // Outdated TCB still allowed when flag is off - IAppAuth.AppBootInfo memory bootInfo = IAppAuth.AppBootInfo({ + IAppAuth.AppInfo memory bootInfo = IAppAuth.AppInfo({ appId: appId, composeHash: composeHash, instanceId: address(0), @@ -399,7 +399,7 @@ contract UpgradesWithPluginTest is Test { assertEq(app.owner(), owner); // Test the full flow still works - IAppAuth.AppBootInfo memory bootInfo = IAppAuth.AppBootInfo({ + IAppAuth.AppInfo memory bootInfo = IAppAuth.AppInfo({ appId: appId, composeHash: composeHash, instanceId: address(0), diff --git a/dstack/kms/auth-mock/README.md b/dstack/kms/auth-mock/README.md index d32ee9ada..8d69c0a06 100644 --- a/dstack/kms/auth-mock/README.md +++ b/dstack/kms/auth-mock/README.md @@ -171,7 +171,7 @@ the API remains identical, so no client code changes are needed. ### Authorization endpoints -`POST /checkKmsAuthorization` and `POST /checkAppAuthorization` return `AuthorizationResult`, including +`POST /authorizeKms` and `POST /authorizeApp` return `AuthorizationResult`, including `rejectionReason`, `gatewayAppId`, `currentRootEpoch`, and `currentRootEpochHash`. The selected endpoint determines the policy; the backend checks that the verified `appId` matches that scope. `requestedAppEpoch` defaults to 0; nonzero values are currently diff --git a/dstack/kms/auth-mock/index.test.ts b/dstack/kms/auth-mock/index.test.ts index d9a22ff3c..dc2ee8445 100644 --- a/dstack/kms/auth-mock/index.test.ts +++ b/dstack/kms/auth-mock/index.test.ts @@ -385,7 +385,7 @@ describe('Mock Behavior Verification', () => { describe('authorization endpoints', () => { - const authorize = (appId: string, isKms = false, requestedAppEpoch = 0) => appFetch(new Request(`http://localhost/${isKms ? 'checkKmsAuthorization' : 'checkAppAuthorization'}`, { + const authorize = (appId: string, isKms = false, requestedAppEpoch = 0) => appFetch(new Request(`http://localhost/${isKms ? 'authorizeKms' : 'authorizeApp'}`, { method: 'POST', headers: { 'Content-Type': 'application/json' }, body: JSON.stringify({ appId, requestedAppEpoch, mrAggregated: '', osImageHash: '', composeHash: '', instanceId: '', deviceId: '' }), })); diff --git a/dstack/kms/auth-mock/index.ts b/dstack/kms/auth-mock/index.ts index 66abd04fb..1ea6e326c 100644 --- a/dstack/kms/auth-mock/index.ts +++ b/dstack/kms/auth-mock/index.ts @@ -128,10 +128,6 @@ class MockBackend { if (appInfo.requestedAppEpoch !== 0) { return { isAllowed: false, rejectionReason: 'Requested app epoch must be zero', gatewayAppId: this.mockGatewayAppId, ...epoch }; } - const kmsAppId = process.env.KMS_CONTRACT_ADDR || '0xmockcontract1234567890123456789012345678'; - if (isKms !== (normHash(appInfo.appId) === normHash(kmsAppId))) { - return { isAllowed: false, rejectionReason: 'Authorization kind does not match app ID', gatewayAppId: this.mockGatewayAppId, ...epoch }; - } const result = this.evaluateBoot(appInfo, isKms); return { isAllowed: result.isAllowed, rejectionReason: result.isAllowed ? '' : result.reason, gatewayAppId: this.mockGatewayAppId, ...epoch }; @@ -233,10 +229,10 @@ app.get('/', async (c) => { } }); -for (const method of ['checkKmsAuthorization', 'checkAppAuthorization'] as const) { +for (const method of ['authorizeKms', 'authorizeApp'] as const) { app.post(`/${method}`, zValidator('json', AuthorizationRequestSchema), async (c) => { try { - return c.json(await mockBackend.checkAuthorization(c.req.valid('json'), method === 'checkKmsAuthorization')); + return c.json(await mockBackend.checkAuthorization(c.req.valid('json'), method === 'authorizeKms')); } catch { return c.json({ status: 'error', message: 'authorization backend unavailable' }, 503); } diff --git a/dstack/kms/auth-mock/openapi.json b/dstack/kms/auth-mock/openapi.json index 90d833adc..c4d1bc05a 100644 --- a/dstack/kms/auth-mock/openapi.json +++ b/dstack/kms/auth-mock/openapi.json @@ -99,10 +99,10 @@ } } }, - "/checkKmsAuthorization": { + "/authorizeKms": { "post": { "summary": "Authorize a Kms identity and return the current root epoch", - "operationId": "checkKmsAuthorization", + "operationId": "authorizeKms", "requestBody": { "required": true, "content": { @@ -130,10 +130,10 @@ } } }, - "/checkAppAuthorization": { + "/authorizeApp": { "post": { "summary": "Authorize a App identity and return the current root epoch", - "operationId": "checkAppAuthorization", + "operationId": "authorizeApp", "requestBody": { "required": true, "content": { diff --git a/dstack/kms/auth-simple/README.md b/dstack/kms/auth-simple/README.md index 4918f2568..494b3fcce 100644 --- a/dstack/kms/auth-simple/README.md +++ b/dstack/kms/auth-simple/README.md @@ -256,7 +256,7 @@ bun run test:run ### Authorization endpoints -`POST /checkKmsAuthorization` and `POST /checkAppAuthorization` return `AuthorizationResult`, including +`POST /authorizeKms` and `POST /authorizeApp` return `AuthorizationResult`, including `rejectionReason`, `gatewayAppId`, `currentRootEpoch`, and `currentRootEpochHash`. The selected endpoint determines the policy; the backend checks that the verified `appId` matches that scope. `requestedAppEpoch` defaults to 0; nonzero values are currently diff --git a/dstack/kms/auth-simple/index.test.ts b/dstack/kms/auth-simple/index.test.ts index f490d03e7..b372091f5 100644 --- a/dstack/kms/auth-simple/index.test.ts +++ b/dstack/kms/auth-simple/index.test.ts @@ -473,7 +473,7 @@ describe('auth-simple', () => { (await app.fetch(new Request('http://localhost/', { method: 'GET' }))).json(); const authorize = async (fields: object = {}, isKms = false) => - (await app.fetch(new Request(`http://localhost/${isKms ? 'checkKmsAuthorization' : 'checkAppAuthorization'}`, { + (await app.fetch(new Request(`http://localhost/${isKms ? 'authorizeKms' : 'authorizeApp'}`, { method: 'POST', headers: { 'Content-Type': 'application/json' }, body: JSON.stringify({ ...baseBootInfo, ...fields }), }))).json(); @@ -486,9 +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({}, true)).toMatchObject({ isAllowed: false, rejectionReason: 'Authorization kind does not match app ID' }); - expect(await authorize({ appId: kmsContractAddr })).toMatchObject({ isAllowed: false, rejectionReason: 'Authorization kind does not match app ID' }); - 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 must be zero', 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, diff --git a/dstack/kms/auth-simple/index.ts b/dstack/kms/auth-simple/index.ts index c2c07ea69..05d75b1d5 100644 --- a/dstack/kms/auth-simple/index.ts +++ b/dstack/kms/auth-simple/index.ts @@ -175,9 +175,6 @@ class ConfigBackend { if (appInfo.requestedAppEpoch !== 0) { return { isAllowed: false, rejectionReason: 'Requested app epoch must be zero', gatewayAppId: config.gatewayAppId, ...epoch }; } - if (isKms !== (normalizeHex(appInfo.appId) === normalizeHex(config.kmsContractAddr))) { - return { isAllowed: false, rejectionReason: 'Authorization kind does not match app ID', gatewayAppId: config.gatewayAppId, ...epoch }; - } const result = this.evaluateBoot(appInfo, isKms, config); return { isAllowed: result.isAllowed, rejectionReason: result.reason, gatewayAppId: result.gatewayAppId, ...epoch }; } @@ -362,10 +359,10 @@ app.get('/', async (c) => { } }); -for (const method of ['checkKmsAuthorization', 'checkAppAuthorization'] as const) { +for (const method of ['authorizeKms', 'authorizeApp'] as const) { app.post(`/${method}`, zValidator('json', AuthorizationRequestSchema), async (c) => { try { - return c.json(await backend.checkAuthorization(c.req.valid('json'), method === 'checkKmsAuthorization')); + return c.json(await backend.checkAuthorization(c.req.valid('json'), method === 'authorizeKms')); } catch { console.error('authorization backend failed'); return c.json({ status: 'error', message: 'authorization backend unavailable' }, 503); From 4a2a8d1ef487dd201ee732b549f98c23165e8c78 Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Sat, 7 Nov 2026 06:11:26 -0800 Subject: [PATCH 025/166] refactor(kms): rename authorization endpoints consistently Signed-off-by: Kevin Wang --- dstack/kms/README.md | 4 ++-- dstack/kms/auth-eth-bun/README.md | 2 +- dstack/kms/auth-eth-bun/index.test.ts | 6 +++--- dstack/kms/auth-eth-bun/index.ts | 8 ++++---- dstack/kms/auth-eth-bun/openapi.json | 8 ++++---- dstack/kms/src/main_service/upgrade_authority.rs | 8 ++++---- 6 files changed, 18 insertions(+), 18 deletions(-) diff --git a/dstack/kms/README.md b/dstack/kms/README.md index 682839967..c21132f3d 100644 --- a/dstack/kms/README.md +++ b/dstack/kms/README.md @@ -248,8 +248,8 @@ The `SignCert` RPC is used by the dstack app to sign a TLS certificate. In this Normal operations use two explicit authorization endpoints: -- `POST /checkKmsAuthorization` for KMS self-authorization and root-key handover. -- `POST /checkAppAuthorization` for application keys and certificates. +- `POST /authorizeKms` for KMS self-authorization and root-key handover. +- `POST /authorizeApp` for application keys and certificates. The Ethereum backends call the matching `DstackKms` method. Both accept `AuthorizationRequest` and return `AuthorizationResult`; no caller-kind flag is diff --git a/dstack/kms/auth-eth-bun/README.md b/dstack/kms/auth-eth-bun/README.md index bb028370f..92d50cced 100644 --- a/dstack/kms/auth-eth-bun/README.md +++ b/dstack/kms/auth-eth-bun/README.md @@ -203,7 +203,7 @@ the test suite (`index.test.ts`) validates: ### Authorization endpoints -`POST /checkKmsAuthorization` and `POST /checkAppAuthorization` return `AuthorizationResult`, including +`POST /authorizeKms` and `POST /authorizeApp` return `AuthorizationResult`, including `rejectionReason`, `gatewayAppId`, `currentRootEpoch`, and `currentRootEpochHash`. The selected endpoint determines the policy; the backend checks that the verified `appId` matches that scope. `requestedAppEpoch` defaults to 0; nonzero values are currently diff --git a/dstack/kms/auth-eth-bun/index.test.ts b/dstack/kms/auth-eth-bun/index.test.ts index 63f19808d..2ce209916 100644 --- a/dstack/kms/auth-eth-bun/index.test.ts +++ b/dstack/kms/auth-eth-bun/index.test.ts @@ -834,14 +834,14 @@ describe('request hardening', () => { }); -describe('POST /checkAppAuthorization', () => { +describe('POST /authorizeApp', () => { const appInfo = { appId: '0x' + '11'.repeat(20), composeHash: '0x' + '22'.repeat(32), instanceId: '', deviceId: '0x' + '33'.repeat(32), mrAggregated: '0x' + '44'.repeat(32), mrSystem: '0x' + '55'.repeat(32), osImageHash: '0x' + '66'.repeat(32), tcbStatus: 'UpToDate', advisoryIds: [], }; - const authorize = (requestedRootEpoch?: number) => appFetch(new Request('http://localhost/checkAppAuthorization', { + const authorize = (requestedRootEpoch?: number) => appFetch(new Request('http://localhost/authorizeApp', { method: 'POST', headers: { 'Content-Type': 'application/json' }, body: JSON.stringify({ ...appInfo, requestedRootEpoch }), })); @@ -856,7 +856,7 @@ describe('POST /checkAppAuthorization', () => { expect(await response.json()).toEqual({ ...result, currentRootEpoch: 2, currentAppEpoch: 0 }); expect(mockReadContract).toHaveBeenCalledTimes(1); expect(mockReadContract).toHaveBeenCalledWith(expect.objectContaining({ - functionName: 'checkAppAuthorization', blockNumber: 98n, + functionName: 'authorizeApp', blockNumber: 98n, args: [expect.objectContaining({ appInfo: expect.objectContaining({ appId: appInfo.appId }), requestedAppEpoch: { isSet: true, 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 a29bd40ee..540ae00d5 100644 --- a/dstack/kms/auth-eth-bun/index.ts +++ b/dstack/kms/auth-eth-bun/index.ts @@ -117,8 +117,8 @@ const AUTHORIZATION_ABI = { // DstackKms contract ABI (minimal required functions) const DSTACK_KMS_ABI = [ - { ...AUTHORIZATION_ABI, name: 'checkKmsAuthorization' }, - { ...AUTHORIZATION_ABI, name: 'checkAppAuthorization' }, + { ...AUTHORIZATION_ABI, name: 'authorizeKms' }, + { ...AUTHORIZATION_ABI, name: 'authorizeApp' }, { name: 'isKmsAllowed', type: 'function', @@ -240,7 +240,7 @@ class EthereumBackend { } - async checkAuthorization(appInfo: AuthorizationRequest, method: 'checkKmsAuthorization' | 'checkAppAuthorization') { + async checkAuthorization(appInfo: AuthorizationRequest, method: 'authorizeKms' | 'authorizeApp') { const blockNumber = await this.finalizedBlockNumber(); const result = await this.client.readContract({ address: this.kmsContractAddr, @@ -422,7 +422,7 @@ app.get('/', async (c) => { } }); -for (const method of ['checkKmsAuthorization', 'checkAppAuthorization'] as const) { +for (const method of ['authorizeKms', 'authorizeApp'] as const) { app.post(`/${method}`, bodyLimit({ maxSize: MAX_BOOT_INFO_BYTES }), zValidator('json', AuthorizationRequestSchema), diff --git a/dstack/kms/auth-eth-bun/openapi.json b/dstack/kms/auth-eth-bun/openapi.json index 46a99357a..b7030e175 100644 --- a/dstack/kms/auth-eth-bun/openapi.json +++ b/dstack/kms/auth-eth-bun/openapi.json @@ -99,10 +99,10 @@ } } }, - "/checkKmsAuthorization": { + "/authorizeKms": { "post": { "summary": "Authorize a Kms identity and return the current root epoch", - "operationId": "checkKmsAuthorization", + "operationId": "authorizeKms", "requestBody": { "required": true, "content": { @@ -130,10 +130,10 @@ } } }, - "/checkAppAuthorization": { + "/authorizeApp": { "post": { "summary": "Authorize a App identity and return the current root epoch", - "operationId": "checkAppAuthorization", + "operationId": "authorizeApp", "requestBody": { "required": true, "content": { diff --git a/dstack/kms/src/main_service/upgrade_authority.rs b/dstack/kms/src/main_service/upgrade_authority.rs index 92cdc0e49..08acbc75f 100644 --- a/dstack/kms/src/main_service/upgrade_authority.rs +++ b/dstack/kms/src/main_service/upgrade_authority.rs @@ -259,7 +259,7 @@ impl AuthApi { &self, app_info: &BootInfo, ) -> Result { - self.check_authorization(app_info, "checkKmsAuthorization", None) + self.check_authorization(app_info, "authorizeKms", None) .await } @@ -268,7 +268,7 @@ impl AuthApi { app_info: &BootInfo, requested_root_epoch: Option, ) -> Result { - self.check_authorization(app_info, "checkAppAuthorization", requested_root_epoch) + self.check_authorization(app_info, "authorizeApp", requested_root_epoch) .await } @@ -723,8 +723,8 @@ pub(crate) mod tests { assert_eq!(requests.len(), 4); assert!(requests[..3] .iter() - .all(|(path, _)| path == "/checkAppAuthorization")); - assert_eq!(requests[3].0, "/checkKmsAuthorization"); + .all(|(path, _)| path == "/authorizeApp")); + assert_eq!(requests[3].0, "/authorizeKms"); let mut expected = serde_json::to_value(info).unwrap(); expected["requestedAppEpoch"] = 0.into(); assert_eq!(requests[0].1, expected); From de25eab5c38157fb74253be98f07d99c82ff40dc Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Sat, 7 Nov 2026 06:28:05 -0800 Subject: [PATCH 026/166] refactor(kms): name optional integer wrapper explicitly Signed-off-by: Kevin Wang --- dstack/kms/auth-eth/contracts/DstackKms.sol | 6 +++--- dstack/kms/auth-eth/test/DstackKms.t.sol | 8 ++++---- 2 files changed, 7 insertions(+), 7 deletions(-) diff --git a/dstack/kms/auth-eth/contracts/DstackKms.sol b/dstack/kms/auth-eth/contracts/DstackKms.sol index e14926087..f87c04b74 100644 --- a/dstack/kms/auth-eth/contracts/DstackKms.sol +++ b/dstack/kms/auth-eth/contracts/DstackKms.sol @@ -14,7 +14,7 @@ import "@openzeppelin/contracts-upgradeable/utils/introspection/ERC165Upgradeabl import "@openzeppelin/contracts/proxy/ERC1967/ERC1967Proxy.sol"; contract DstackKms is Initializable, Ownable2StepUpgradeable, UUPSUpgradeable, ERC165Upgradeable, IAppAuth { - struct OptionalEpoch { + struct OptionalU256 { bool isSet; uint256 value; } @@ -22,8 +22,8 @@ contract DstackKms is Initializable, Ownable2StepUpgradeable, UUPSUpgradeable, E struct AuthorizationRequest { AppInfo appInfo; // Reserved for application key rotation; only zero is currently supported. - OptionalEpoch requestedAppEpoch; - OptionalEpoch requestedRootEpoch; + OptionalU256 requestedAppEpoch; + OptionalU256 requestedRootEpoch; } struct AuthorizationResult { diff --git a/dstack/kms/auth-eth/test/DstackKms.t.sol b/dstack/kms/auth-eth/test/DstackKms.t.sol index a34094469..a35f445f1 100644 --- a/dstack/kms/auth-eth/test/DstackKms.t.sol +++ b/dstack/kms/auth-eth/test/DstackKms.t.sol @@ -465,11 +465,11 @@ contract DstackKmsTest is Test { returns (DstackKms.AuthorizationResult memory) { DstackKms.AuthorizationRequest memory request = - DstackKms.AuthorizationRequest(info, DstackKms.OptionalEpoch(false, 0), DstackKms.OptionalEpoch(false, 0)); + DstackKms.AuthorizationRequest(info, DstackKms.OptionalU256(false, 0), DstackKms.OptionalU256(false, 0)); return isKms ? kms.authorizeKms(request) : kms.authorizeApp(request); } - function test_CheckAuthorization_OptionalEpochDistinguishesZeroFromDefault() public { + function test_CheckAuthorization_OptionalU256DistinguishesZeroFromDefault() public { IAppAuth.AppInfo memory info = authorizationInfo(address(kms)); vm.startPrank(owner); kms.registerRootEpoch(ca(0xa0), k256(0xb0), "https://example.com/epoch-0"); @@ -480,7 +480,7 @@ contract DstackKmsTest is Test { vm.stopPrank(); DstackKms.AuthorizationRequest memory request = - DstackKms.AuthorizationRequest(info, DstackKms.OptionalEpoch(false, 0), DstackKms.OptionalEpoch(true, 1)); + DstackKms.AuthorizationRequest(info, DstackKms.OptionalU256(false, 0), DstackKms.OptionalU256(true, 1)); DstackKms.AuthorizationResult memory result = kms.authorizeKms(request); assertFalse(result.isAllowed); assertEq(result.rejectionReason, "Requested root epoch is not active"); @@ -537,7 +537,7 @@ contract DstackKmsTest is Test { assertFalse(result.isAllowed); assertEq(result.rejectionReason, "Aggregated MR not allowed"); DstackKms.AuthorizationRequest memory appRequest = - DstackKms.AuthorizationRequest(info, DstackKms.OptionalEpoch(true, 1), DstackKms.OptionalEpoch(false, 0)); + DstackKms.AuthorizationRequest(info, DstackKms.OptionalU256(true, 1), DstackKms.OptionalU256(false, 0)); result = kms.authorizeApp(appRequest); assertFalse(result.isAllowed); assertEq(result.rejectionReason, "Requested app epoch must be zero"); From 3eefe36fc7a81eef878632e4d2cfb376cbadbfe5 Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Sat, 7 Nov 2026 06:31:48 -0800 Subject: [PATCH 027/166] refactor(kms): name optional uint256 wrapper precisely Signed-off-by: Kevin Wang --- dstack/kms/auth-eth/contracts/DstackKms.sol | 6 +++--- dstack/kms/auth-eth/test/DstackKms.t.sol | 8 ++++---- 2 files changed, 7 insertions(+), 7 deletions(-) diff --git a/dstack/kms/auth-eth/contracts/DstackKms.sol b/dstack/kms/auth-eth/contracts/DstackKms.sol index f87c04b74..15d253275 100644 --- a/dstack/kms/auth-eth/contracts/DstackKms.sol +++ b/dstack/kms/auth-eth/contracts/DstackKms.sol @@ -14,7 +14,7 @@ import "@openzeppelin/contracts-upgradeable/utils/introspection/ERC165Upgradeabl import "@openzeppelin/contracts/proxy/ERC1967/ERC1967Proxy.sol"; contract DstackKms is Initializable, Ownable2StepUpgradeable, UUPSUpgradeable, ERC165Upgradeable, IAppAuth { - struct OptionalU256 { + struct OptionalUint256 { bool isSet; uint256 value; } @@ -22,8 +22,8 @@ contract DstackKms is Initializable, Ownable2StepUpgradeable, UUPSUpgradeable, E struct AuthorizationRequest { AppInfo appInfo; // Reserved for application key rotation; only zero is currently supported. - OptionalU256 requestedAppEpoch; - OptionalU256 requestedRootEpoch; + OptionalUint256 requestedAppEpoch; + OptionalUint256 requestedRootEpoch; } struct AuthorizationResult { diff --git a/dstack/kms/auth-eth/test/DstackKms.t.sol b/dstack/kms/auth-eth/test/DstackKms.t.sol index a35f445f1..791785354 100644 --- a/dstack/kms/auth-eth/test/DstackKms.t.sol +++ b/dstack/kms/auth-eth/test/DstackKms.t.sol @@ -465,11 +465,11 @@ contract DstackKmsTest is Test { returns (DstackKms.AuthorizationResult memory) { DstackKms.AuthorizationRequest memory request = - DstackKms.AuthorizationRequest(info, DstackKms.OptionalU256(false, 0), DstackKms.OptionalU256(false, 0)); + DstackKms.AuthorizationRequest(info, DstackKms.OptionalUint256(false, 0), DstackKms.OptionalUint256(false, 0)); return isKms ? kms.authorizeKms(request) : kms.authorizeApp(request); } - function test_CheckAuthorization_OptionalU256DistinguishesZeroFromDefault() public { + function test_CheckAuthorization_OptionalUint256DistinguishesZeroFromDefault() public { IAppAuth.AppInfo memory info = authorizationInfo(address(kms)); vm.startPrank(owner); kms.registerRootEpoch(ca(0xa0), k256(0xb0), "https://example.com/epoch-0"); @@ -480,7 +480,7 @@ contract DstackKmsTest is Test { vm.stopPrank(); DstackKms.AuthorizationRequest memory request = - DstackKms.AuthorizationRequest(info, DstackKms.OptionalU256(false, 0), DstackKms.OptionalU256(true, 1)); + DstackKms.AuthorizationRequest(info, DstackKms.OptionalUint256(false, 0), DstackKms.OptionalUint256(true, 1)); DstackKms.AuthorizationResult memory result = kms.authorizeKms(request); assertFalse(result.isAllowed); assertEq(result.rejectionReason, "Requested root epoch is not active"); @@ -537,7 +537,7 @@ contract DstackKmsTest is Test { assertFalse(result.isAllowed); assertEq(result.rejectionReason, "Aggregated MR not allowed"); DstackKms.AuthorizationRequest memory appRequest = - DstackKms.AuthorizationRequest(info, DstackKms.OptionalU256(true, 1), DstackKms.OptionalU256(false, 0)); + DstackKms.AuthorizationRequest(info, DstackKms.OptionalUint256(true, 1), DstackKms.OptionalUint256(false, 0)); result = kms.authorizeApp(appRequest); assertFalse(result.isAllowed); assertEq(result.rejectionReason, "Requested app epoch must be zero"); From 697920c30a4992b62215eb17f6d41b65d90a0c56 Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Sat, 7 Nov 2026 07:25:20 -0800 Subject: [PATCH 028/166] test(cert-client): retain direct retrieval TLS authentication coverage Signed-off-by: Kevin Wang --- dstack/cert-client/Cargo.toml | 1 + dstack/cert-client/src/lib.rs | 20 ++++++++++++++++++++ dstack/kms/README.md | 4 +++- 3 files changed, 24 insertions(+), 1 deletion(-) diff --git a/dstack/cert-client/Cargo.toml b/dstack/cert-client/Cargo.toml index 2f1a56e0e..8a715e379 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] +ra-rpc = { workspace = true, features = ["rocket", "client"] } mock-attestation = { path = "../crates/mock-attestation" } hex.workspace = true rocket.workspace = true diff --git a/dstack/cert-client/src/lib.rs b/dstack/cert-client/src/lib.rs index 5a2ef732f..65498f275 100644 --- a/dstack/cert-client/src/lib.rs +++ b/dstack/cert-client/src/lib.rs @@ -444,6 +444,7 @@ mod tests { #[rocket::post("/GetAppKey", data = "<_request>")] fn direct_app_keys( _request: String, + _cert: rocket::mtls::Certificate<'_>, state: &rocket::State<(Arc>, std::sync::atomic::AtomicUsize)>, ) -> rocket::response::content::RawJson { state.1.fetch_add(1, std::sync::atomic::Ordering::SeqCst); @@ -455,6 +456,7 @@ mod tests { use ra_tls::attestation::{ AttestationVerifierConfig, DstackNitroQuote, QuoteContentType, RootCaPaths, }; + use rocket::tls::Resolver as _; let dir = tempfile::tempdir().unwrap(); let attester = mock_attestation::nsm::NsmGenerator::new().unwrap(); let attestation_root = dir.path().join("attestation-root.pem"); @@ -520,6 +522,7 @@ mod tests { let tx = Mutex::new(Some(tx)); let response_body = Arc::new(Mutex::new(serde_json::to_string(&response).unwrap())); let server = rocket::custom(figment) + .attach(ra_rpc::ratls_client_verifier::RaTlsClientAuth::fairing()) .manage(( response_body.clone(), std::sync::atomic::AtomicUsize::new(0), @@ -539,9 +542,26 @@ mod tests { let url = format!("https://127.0.0.1:{}", rx.await.unwrap()); let client_cert = || { let key = p256_key(); + let report_data = QuoteContentType::RaTlsCert.to_report_data(&key.public_key_der()); + let pcrs = [0, 1, 2, 4] + .into_iter() + .map(|i| (i, vec![i as u8; 48])) + .collect(); + let evidence = attester.attest_with_pcrs(&report_data, pcrs).unwrap(); + let attestation = Attestation { + quote: AttestationQuote::DstackNitroEnclave(DstackNitroQuote { + nsm_quote: evidence, + }), + runtime_events: Vec::new(), + report_data, + config: String::new(), + report: (), + } + .into_versioned(); let cert = CertRequest::builder() .subject("guest") .key(&key) + .attestation(&attestation) .usage_client_auth(true) .build() .self_signed() diff --git a/dstack/kms/README.md b/dstack/kms/README.md index c21132f3d..4e56c057e 100644 --- a/dstack/kms/README.md +++ b/dstack/kms/README.md @@ -324,7 +324,9 @@ CA-issued client certificates, are no longer supported. Guests and `dstack-util get-keys` generate their self-issued RA-TLS certificate locally, then call `GetAppKey` directly. They verify the response connection's attestation and `kms:rpc` usage, and verify that the returned root CA signed that -connection's certificate. There is no `GetMeta` preflight or first unauthenticated +connection's certificate. Skipping CA-chain validation during this first request +still verifies the TLS handshake signature, so replaying an attested certificate +without its private key cannot impersonate the server. There is no `GetMeta` preflight or first unauthenticated connection. The optional provider pin and the key-provider launch measurement still bind the root identity; a valid quote alone is not a KMS policy decision. `--root-ca` additionally pins the CLI connection before sending the request. From 913ce0c4f7fc5d665d3fb54a20d4c987600c2fe2 Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Sat, 7 Nov 2026 17:33:53 -0800 Subject: [PATCH 029/166] fix(auth): implement epoch authorization in the local backend Signed-off-by: Kevin Wang --- docs/onboarding.md | 17 ++ dstack/Cargo.lock | 1 + dstack/crates/dstack-auth/Cargo.toml | 3 + dstack/crates/dstack-auth/src/main.rs | 253 +++++++++++++++++++++++--- 4 files changed, 253 insertions(+), 21 deletions(-) diff --git a/docs/onboarding.md b/docs/onboarding.md index 28b216ff2..ff948d2e3 100644 --- a/docs/onboarding.md +++ b/docs/onboarding.md @@ -312,3 +312,20 @@ dstack logs ``` The KMS VM ID is printed by `dstackup install` when it creates or reuses the KMS CVM. + +### Upgrading the local KMS to root epochs + +Upgrade the bundled `dstack-auth` together with the KMS. Before normal key release, +set `kmsContractAddr` in the host's `auth-allowlist.json` to the KMS's attested app ID +and register the verified bootstrap CA/k256 public keys in `rootEpochs` (epoch zero). +Use the bootstrap evidence to verify these keys; do not trust an unauthenticated +metadata response. Set `currentRootEpoch` to the active registered epoch (initially +zero). See [root registration and rotation](deployment.md#rotating-kms-root-keys). +The local installer does not register these keys automatically; an older installed +KMS can continue using the legacy boot routes while this migration is prepared. + +The Rust backend implements `/authorizeKms` and `/authorizeApp` with the same +registered-root requirement as the other backends. Missing registration denies key +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. diff --git a/dstack/Cargo.lock b/dstack/Cargo.lock index a67cc5543..7f86192db 100644 --- a/dstack/Cargo.lock +++ b/dstack/Cargo.lock @@ -1905,6 +1905,7 @@ dependencies = [ "rocket", "serde", "serde_json", + "tempfile", ] [[package]] diff --git a/dstack/crates/dstack-auth/Cargo.toml b/dstack/crates/dstack-auth/Cargo.toml index 9a374b988..518c712c2 100644 --- a/dstack/crates/dstack-auth/Cargo.toml +++ b/dstack/crates/dstack-auth/Cargo.toml @@ -20,3 +20,6 @@ clap.workspace = true rocket = { workspace = true, features = ["json"] } serde.workspace = true serde_json.workspace = true + +[dev-dependencies] +tempfile.workspace = true diff --git a/dstack/crates/dstack-auth/src/main.rs b/dstack/crates/dstack-auth/src/main.rs index 00e5b72ee..50d8ef937 100644 --- a/dstack/crates/dstack-auth/src/main.rs +++ b/dstack/crates/dstack-auth/src/main.rs @@ -7,8 +7,8 @@ //! //! Runs on the host as `dstack-auth.service`; the KMS-in-CVM reaches it at //! `http://10.0.2.2:` under user-mode networking and POSTs `BootInfo` to -//! `/bootAuth/app` (compose-hash allowlist) and `/bootAuth/kms` (mrAggregated -//! allowlist). The allowlist JSON is re-read on every request, so `dstack run` +//! `/authorizeApp` (compose-hash allowlist) and `/authorizeKms` (mrAggregated +//! allowlist). Legacy `/bootAuth/*` routes remain available for bootstrap. The allowlist JSON is re-read on every request, so `dstack run` //! can add an app without a restart. Fails closed: a missing/invalid allowlist //! denies everything. //! @@ -73,6 +73,26 @@ struct BootResponse { reason: String, } +#[derive(Deserialize)] +#[serde(rename_all = "camelCase")] +struct AuthorizationRequest { + #[serde(flatten)] + boot_info: BootInfo, + requested_root_epoch: Option, + requested_app_epoch: Option, +} + +#[derive(Serialize)] +#[serde(rename_all = "camelCase")] +struct AuthorizationResult { + is_allowed: bool, + rejection_reason: String, + gateway_app_id: String, + current_root_epoch: usize, + current_root_epoch_hash: String, + current_app_epoch: u32, +} + /// the current KMS root epoch and its chain hash, which commits to the root /// public keys of epochs 0..=number. struct EpochInfo { @@ -84,6 +104,7 @@ struct EpochInfo { #[serde(rename_all = "camelCase", default)] struct Allowlist { os_images: Vec, + kms_contract_addr: String, gateway_app_id: String, kms: KmsRules, apps: HashMap, @@ -101,11 +122,20 @@ struct RootEpoch { } impl Allowlist { + fn app(&self, app_id: &str) -> Option<&AppRules> { + let app_id = norm(app_id); + self.apps + .iter() + .find(|(id, _)| norm(id) == app_id) + .map(|(_, app)| app) + } + /// `None` only when no epoch is configured; invalid configuration is an error. fn current_root_epoch(&self) -> Result> { if self.root_epochs.is_empty() && self.current_root_epoch == 0 { return Ok(None); } + u32::try_from(self.current_root_epoch).context("root epoch exceeds uint32")?; let epochs = self .root_epochs .get(..=self.current_root_epoch) @@ -214,17 +244,11 @@ fn allow(info: &BootInfo, is_kms: bool, al: &Allowlist) -> BootResponse { } } -fn check_app(info: &BootInfo, al: &Allowlist) -> BootResponse { +fn app_policy(info: &BootInfo, al: &Allowlist) -> BootResponse { if !al.os_images.is_empty() && !contains(&al.os_images, &info.os_image_hash) { return deny(al, "os image not allowed"); } - let app_id = norm(&info.app_id); - let Some(app) = al - .apps - .iter() - .find(|(k, _)| norm(k) == app_id) - .map(|(_, v)| v) - else { + let Some(app) = al.app(&info.app_id) else { return deny(al, "app not registered"); }; if !contains(&app.compose_hashes, &info.compose_hash) { @@ -233,17 +257,111 @@ fn check_app(info: &BootInfo, al: &Allowlist) -> BootResponse { if !device_ok(app.allow_any_device, &app.devices, &info.device_id) { return deny(al, "device not allowed"); } - allow(info, false, al) + BootResponse { + is_allowed: true, + gateway_app_id: al.gateway_app_id.clone(), + reason: String::new(), + } } -fn check_kms(info: &BootInfo, al: &Allowlist) -> BootResponse { +fn kms_policy(info: &BootInfo, al: &Allowlist) -> BootResponse { if !contains(&al.kms.mr_aggregated, &info.mr_aggregated) { return deny(al, "kms mrAggregated not allowed"); } if !device_ok(al.kms.allow_any_device, &al.kms.devices, &info.device_id) { return deny(al, "device not allowed"); } - allow(info, true, al) + BootResponse { + is_allowed: true, + gateway_app_id: al.gateway_app_id.clone(), + reason: String::new(), + } +} + +fn check_app(info: &BootInfo, al: &Allowlist) -> BootResponse { + let result = app_policy(info, al); + if result.is_allowed { + allow(info, false, al) + } else { + result + } +} + +fn check_kms(info: &BootInfo, al: &Allowlist) -> BootResponse { + let result = kms_policy(info, al); + if result.is_allowed { + allow(info, true, al) + } else { + result + } +} + +fn authorize( + request: &AuthorizationRequest, + is_kms: bool, + al: &Allowlist, +) -> Result { + let current = al.current_root_epoch()?; + let mut result = AuthorizationResult { + is_allowed: false, + rejection_reason: String::new(), + gateway_app_id: al.gateway_app_id.clone(), + current_root_epoch: current.as_ref().map_or(0, |epoch| epoch.number), + current_root_epoch_hash: current.as_ref().map_or_else( + || format!("0x{}", "00".repeat(32)), + |epoch| epoch.hash.clone(), + ), + current_app_epoch: 0, + }; + let info = &request.boot_info; + result.rejection_reason = if current.is_none() { + "root epoch not registered".into() + } else if request + .requested_root_epoch + .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 norm(&al.kms_contract_addr).is_empty() { + "kmsContractAddr is not configured".into() + } else if (norm(&info.app_id) == norm(&al.kms_contract_addr)) != is_kms { + "caller does not match authorization scope".into() + } else { + let policy = if is_kms { + kms_policy(info, al) + } else { + app_policy(info, al) + }; + result.is_allowed = policy.is_allowed; + policy.reason + }; + Ok(result) +} + +fn unavailable(error: anyhow::Error) -> Status { + rocket::warn!("invalid authorization configuration: {error}"); + Status::ServiceUnavailable +} + +#[post("/authorizeApp", data = "")] +fn authorize_app( + request: Json, + cli: &State, +) -> Result, Status> { + authorize(&request, false, &load(&cli.config)) + .map(Json) + .map_err(unavailable) +} + +#[post("/authorizeKms", data = "")] +fn authorize_kms( + request: Json, + cli: &State, +) -> Result, Status> { + authorize(&request, true, &load(&cli.config)) + .map(Json) + .map_err(unavailable) } /// load the allowlist, failing closed (deny-all) if it's missing or invalid. @@ -297,9 +415,9 @@ fn info(cli: &State) -> Result, Status> { fn info_json(al: &Allowlist) -> Result { let mut info = json!({ "status": "ok", - "kmsContractAddr": "", + "kmsContractAddr": al.kms_contract_addr, "ethRpcUrl": "", - "gatewayAppId": "", + "gatewayAppId": al.gateway_app_id, "chainId": 0, "appImplementation": "", }); @@ -310,15 +428,19 @@ fn info_json(al: &Allowlist) -> Result { Ok(info) } -#[rocket::main] -async fn main() -> Result<()> { - let cli = Cli::parse(); +fn server(cli: Cli) -> rocket::Rocket { let figment = rocket::Config::figment() .merge(("address", cli.address.clone())) .merge(("port", cli.port)); - rocket::custom(figment) - .manage(cli) - .mount("/", routes![info, boot_app, boot_kms]) + rocket::custom(figment).manage(cli).mount( + "/", + routes![info, boot_app, boot_kms, authorize_app, authorize_kms], + ) +} + +#[rocket::main] +async fn main() -> Result<()> { + server(Cli::parse()) .launch() .await .map_err(|e| anyhow::anyhow!("auth webhook failed: {e}"))?; @@ -329,6 +451,95 @@ async fn main() -> Result<()> { mod tests { use super::*; + #[rocket::async_test] + async fn authorization_routes_use_registered_roots_and_separate_caller_scopes() { + use rocket::{http::ContentType, local::asynchronous::Client}; + + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("allowlist.json"); + let mut policy = json!({ + "kmsContractAddr": "0x1111111111111111111111111111111111111111", + "rootEpochs": [{"caPubkey": "aa", "k256Pubkey": "bb"}], + "osImages": ["aa"], + "kms": {"mrAggregated": ["bb"]}, + "apps": {"0x2222222222222222222222222222222222222222": {"composeHashes": ["cc"]}} + }); + std::fs::write(&path, policy.to_string()).unwrap(); + let client = Client::tracked(server(Cli { + config: path.clone(), + address: "127.0.0.1".into(), + port: 0, + })) + .await + .unwrap(); + let kms = json!({"appId": "11".repeat(20), "mrAggregated": "bb"}); + let app = json!({"appId": "22".repeat(20), "composeHash": "cc", "osImageHash": "aa"}); + for (route, body, allowed) in [ + ("/authorizeKms", kms.clone(), true), + ("/authorizeApp", app.clone(), true), + ("/authorizeKms", app.clone(), false), + ("/authorizeApp", kms.clone(), false), + ] { + let response = client + .post(route) + .header(ContentType::JSON) + .body(body.to_string()) + .dispatch() + .await; + assert_eq!(response.status(), Status::Ok); + let response: serde_json::Value = response.into_json().await.unwrap(); + assert_eq!(response["isAllowed"], allowed); + assert_eq!(response["currentRootEpoch"], 0); + assert_eq!(response["currentAppEpoch"], 0); + assert_ne!( + response["currentRootEpochHash"], + format!("0x{}", "00".repeat(32)) + ); + } + for field in ["requestedRootEpoch", "requestedAppEpoch"] { + let mut body = app.clone(); + body[field] = 1.into(); + let response: serde_json::Value = client + .post("/authorizeApp") + .header(ContentType::JSON) + .body(body.to_string()) + .dispatch() + .await + .into_json() + .await + .unwrap(); + assert_eq!(response["isAllowed"], false); + } + policy["rootEpochs"] = json!([]); + std::fs::write(&path, policy.to_string()).unwrap(); + let response: serde_json::Value = client + .post("/authorizeKms") + .header(ContentType::JSON) + .body(kms.to_string()) + .dispatch() + .await + .into_json() + .await + .unwrap(); + assert_eq!(response["isAllowed"], false); + assert_eq!( + response["currentRootEpochHash"], + format!("0x{}", "00".repeat(32)) + ); + policy["currentRootEpoch"] = 1.into(); + std::fs::write(&path, policy.to_string()).unwrap(); + assert_eq!( + client + .post("/authorizeApp") + .header(ContentType::JSON) + .body(app.to_string()) + .dispatch() + .await + .status(), + Status::ServiceUnavailable + ); + } + fn allowlist() -> Allowlist { serde_json::from_str( r#"{ From bff2fd91d466b61742017e202ac84156b98ec655 Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Sat, 7 Nov 2026 17:43:28 -0800 Subject: [PATCH 030/166] style(auth-eth): format epoch authorization declarations Signed-off-by: Kevin Wang --- dstack/kms/auth-eth/contracts/DstackKms.sol | 12 ++---------- dstack/kms/auth-eth/test/DstackKms.t.sol | 15 +++++++++------ 2 files changed, 11 insertions(+), 16 deletions(-) diff --git a/dstack/kms/auth-eth/contracts/DstackKms.sol b/dstack/kms/auth-eth/contracts/DstackKms.sol index 15d253275..465bef86c 100644 --- a/dstack/kms/auth-eth/contracts/DstackKms.sol +++ b/dstack/kms/auth-eth/contracts/DstackKms.sol @@ -393,19 +393,11 @@ contract DstackKms is Initializable, Ownable2StepUpgradeable, UUPSUpgradeable, E /// 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) - { + function authorizeKms(AuthorizationRequest calldata request) external view returns (AuthorizationResult memory) { return _checkAuthorization(request, true); } - function authorizeApp(AuthorizationRequest calldata request) - external - view - returns (AuthorizationResult memory) - { + function authorizeApp(AuthorizationRequest calldata request) external view returns (AuthorizationResult memory) { return _checkAuthorization(request, false); } diff --git a/dstack/kms/auth-eth/test/DstackKms.t.sol b/dstack/kms/auth-eth/test/DstackKms.t.sol index 791785354..e076d5b19 100644 --- a/dstack/kms/auth-eth/test/DstackKms.t.sol +++ b/dstack/kms/auth-eth/test/DstackKms.t.sol @@ -464,8 +464,9 @@ contract DstackKmsTest is Test { view returns (DstackKms.AuthorizationResult memory) { - DstackKms.AuthorizationRequest memory request = - DstackKms.AuthorizationRequest(info, DstackKms.OptionalUint256(false, 0), DstackKms.OptionalUint256(false, 0)); + DstackKms.AuthorizationRequest memory request = DstackKms.AuthorizationRequest( + info, DstackKms.OptionalUint256(false, 0), DstackKms.OptionalUint256(false, 0) + ); return isKms ? kms.authorizeKms(request) : kms.authorizeApp(request); } @@ -479,8 +480,9 @@ contract DstackKmsTest is Test { kms.addKmsAggregatedMr(info.mrAggregated); vm.stopPrank(); - DstackKms.AuthorizationRequest memory request = - DstackKms.AuthorizationRequest(info, DstackKms.OptionalUint256(false, 0), DstackKms.OptionalUint256(true, 1)); + DstackKms.AuthorizationRequest memory request = DstackKms.AuthorizationRequest( + info, DstackKms.OptionalUint256(false, 0), DstackKms.OptionalUint256(true, 1) + ); DstackKms.AuthorizationResult memory result = kms.authorizeKms(request); assertFalse(result.isAllowed); assertEq(result.rejectionReason, "Requested root epoch is not active"); @@ -536,8 +538,9 @@ contract DstackKmsTest is Test { result = authorize(info, true); assertFalse(result.isAllowed); assertEq(result.rejectionReason, "Aggregated MR not allowed"); - DstackKms.AuthorizationRequest memory appRequest = - DstackKms.AuthorizationRequest(info, DstackKms.OptionalUint256(true, 1), DstackKms.OptionalUint256(false, 0)); + DstackKms.AuthorizationRequest memory appRequest = DstackKms.AuthorizationRequest( + info, DstackKms.OptionalUint256(true, 1), DstackKms.OptionalUint256(false, 0) + ); result = kms.authorizeApp(appRequest); assertFalse(result.isAllowed); assertEq(result.rejectionReason, "Requested app epoch must be zero"); From 3587d243486f7721f553e7fafde050439d6a5455 Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Sat, 7 Nov 2026 18:11:43 -0800 Subject: [PATCH 031/166] refactor(kms): separate app and handover authorization Signed-off-by: Kevin Wang --- docs/deployment.md | 14 ++++--- dstack/kms/src/main_service.rs | 67 +++++++++++++++++++++------------- 2 files changed, 50 insertions(+), 31 deletions(-) diff --git a/docs/deployment.md b/docs/deployment.md index d99d10933..052375db6 100644 --- a/docs/deployment.md +++ b/docs/deployment.md @@ -131,9 +131,13 @@ Production KMS requires: | Custom | Your own authorization logic | Implement webhook interface | All auth servers implement the same webhook interface: -- `GET /` - Health check -- `POST /bootAuth/app` - App boot authorization -- `POST /bootAuth/kms` - KMS boot authorization + +- `GET /` - Health check and current root epoch +- `POST /authorizeApp` - Application key and certificate authorization +- `POST /authorizeKms` - KMS self-authorization and key handover + +The legacy `/bootAuth/app` and `/bootAuth/kms` routes remain available for older +clients and KMS bootstrap/onboarding checks. #### Using auth-simple (Config-Based) @@ -472,7 +476,7 @@ Additional KMS instances can onboard from an existing KMS to share the same root 1. New KMS starts in onboard mode (empty `auto_bootstrap_domain`) 2. New KMS generates a self-issued RA-TLS client certificate with its TDX quote locally; no CA preflight is needed 3. New KMS calls `GetKmsKey` over mTLS: `Admin.GetKmsKey` with the token, or the explicitly enabled public `KMS.GetKmsKey` without one -4. Source KMS verifies attestation via `bootAuth/kms` webhook +4. Source KMS verifies the attestation and authorizes handover via `/authorizeKms` 5. If approved, source KMS returns root keys 6. New KMS verifies the source's attestation on that same connection and checks it against its own policy before using the keys 7. Both KMS instances now derive identical keys @@ -548,7 +552,7 @@ for the full procedure; the KMS steps are: `PopRootEpoch` (or remove the `rootEpochs` entry) and rotate again; nodes that imported the dropped N must be onboarded again. -Each `checkKmsAuthorization` or `checkAppAuthorization` result carries the current epoch and its chain hash, +Each `authorizeKms` or `authorizeApp` result carries the current epoch and its chain hash, including when authorization is denied. KMS verifies these against its keyring; there is no cached serving epoch. An optional `requestedRootEpoch` selects historical keys, and omission selects the returned current epoch. Requests for future diff --git a/dstack/kms/src/main_service.rs b/dstack/kms/src/main_service.rs index 969b5b05c..fc7871d42 100644 --- a/dstack/kms/src/main_service.rs +++ b/dstack/kms/src/main_service.rs @@ -523,9 +523,17 @@ impl RpcHandler { async fn ensure_kms_allowed(&self, vm_config: &str) -> Result { let att = self.ensure_attested()?; - self.ensure_app_attestation_allowed(att, true, false, vm_config, None) - .await - .map(|c| c.boot_info) + let boot_info = build_boot_info_for_attestation(att, false, vm_config)?; + let response = self + .state + .config + .auth_api + .check_kms_authorization(&boot_info) + .await?; + self.state.authorized_epoch(&response, None)?; + self.ensure_attestation_policy(att, vm_config, &boot_info, &response) + .await?; + Ok(boot_info) } async fn ensure_app_boot_allowed( @@ -534,29 +542,43 @@ impl RpcHandler { requested_root_epoch: Option, ) -> Result { let att = self.ensure_attested()?; - self.ensure_app_attestation_allowed(att, false, false, vm_config, requested_root_epoch) + let boot_info = build_boot_info_for_attestation(att, false, vm_config)?; + self.ensure_app_attestation_allowed(att, boot_info, vm_config, requested_root_epoch) .await } async fn ensure_app_attestation_allowed( &self, att: &VerifiedAttestation, - is_kms: bool, - use_boottime_mr: bool, - vm_config_str: &str, + boot_info: BootInfo, + vm_config: &str, requested_root_epoch: Option, ) -> Result { - let boot_info = build_boot_info_for_attestation(att, use_boottime_mr, vm_config_str)?; - let auth = &self.state.config.auth_api; - let response = if is_kms { - auth.check_kms_authorization(&boot_info).await? - } else { - auth.check_app_authorization(&boot_info, requested_root_epoch) - .await? - }; + let response = self + .state + .config + .auth_api + .check_app_authorization(&boot_info, requested_root_epoch) + .await?; let root_epoch = self .state .authorized_epoch(&response, requested_root_epoch)?; + self.ensure_attestation_policy(att, vm_config, &boot_info, &response) + .await?; + Ok(BootConfig { + boot_info, + gateway_app_id: response.gateway_app_id, + root_epoch, + }) + } + + async fn ensure_attestation_policy( + &self, + att: &VerifiedAttestation, + vm_config: &str, + boot_info: &BootInfo, + response: &AuthorizationResult, + ) -> Result<()> { let current_epoch = if self.state.config.auth_api.is_dev() { self.state.keyring.latest() } else { @@ -569,19 +591,11 @@ impl RpcHandler { if boot_info.tee_variant != TeeVariant::DstackAmdSevSnp { self.state .verifier - .verify_os_image_hash( - vm_config_str.into(), - att, - &mut VerificationDetails::default(), - ) + .verify_os_image_hash(vm_config.into(), att, &mut VerificationDetails::default()) .await .context("Failed to verify os image hash")?; } - Ok(BootConfig { - boot_info, - gateway_app_id: response.gateway_app_id, - root_epoch, - }) + Ok(()) } fn derive_app_ca(&self, app_id: &[u8], root: &RootKeys) -> Result { @@ -743,8 +757,9 @@ impl KmsRpc for RpcHandler { .verify_with_ra_pubkey(&csr.pubkey, &self.state.attestation_verifier) .await .context("Quote verification failed")?; + let boot_info = build_boot_info_for_attestation(&attestation, true, &request.vm_config)?; let app_info = self - .ensure_app_attestation_allowed(&attestation, false, true, &request.vm_config, None) + .ensure_app_attestation_allowed(&attestation, boot_info, &request.vm_config, None) .await?; ensure_key_release_allowed(&app_info.boot_info, (&self.state.config).into())?; let root = self.state.epoch_root(app_info.root_epoch)?; From fdf0e31b6c7b3363ea5f2f374d491f7aca11e205 Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Sat, 7 Nov 2026 18:35:13 -0800 Subject: [PATCH 032/166] fix(kms): retain temporary CA compatibility for existing clients Signed-off-by: Kevin Wang --- CHANGELOG.md | 10 +--- docs/deployment.md | 2 +- docs/security/public-security-reports.md | 2 +- docs/security/security-best-practices.md | 2 +- docs/security/security-model.md | 2 +- dstack/kms/README.md | 20 ++++--- dstack/kms/rpc/proto/kms_rpc.proto | 12 +++- dstack/kms/src/config.rs | 10 ++++ dstack/kms/src/main_service.rs | 44 +++++++++++++- dstack/kms/src/onboard_service.rs | 73 +++++++++++++++++++++--- 10 files changed, 147 insertions(+), 30 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 3d7ae700c..89a44ccd5 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,15 +7,11 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] -### Removed - -- KMS: remove `GetTempCaCert` and temporary CA files and handover material. - Guests fetch keys directly with self-issued RA-TLS certificates and validate - the response connection's attestation and root CA. Older clients and sources - that require the temporary CA protocol must be upgraded. - ### Changed +- KMS: retain `GetTempCaCert`, temporary CA files and legacy handover fields so + KMS nodes can be upgraded before existing OS images. + - dstack-util, cert-client: guests mint self-issued RA-TLS client certificates for KMS calls instead of fetching the temp CA and minting from it; the KMS authenticates the quote inside the certificate, not its issuer. diff --git a/docs/deployment.md b/docs/deployment.md index 052375db6..3621228ed 100644 --- a/docs/deployment.md +++ b/docs/deployment.md @@ -474,7 +474,7 @@ Additional KMS instances can onboard from an existing KMS to share the same root **How it works:** 1. New KMS starts in onboard mode (empty `auto_bootstrap_domain`) -2. New KMS generates a self-issued RA-TLS client certificate with its TDX quote locally; no CA preflight is needed +2. New KMS generates an RA-TLS client certificate with its attestation: self-issued with an admin token, or signed by the source temporary CA for tokenless legacy onboarding 3. New KMS calls `GetKmsKey` over mTLS: `Admin.GetKmsKey` with the token, or the explicitly enabled public `KMS.GetKmsKey` without one 4. Source KMS verifies the attestation and authorizes handover via `/authorizeKms` 5. If approved, source KMS returns root keys diff --git a/docs/security/public-security-reports.md b/docs/security/public-security-reports.md index 03f342b23..373299a7c 100644 --- a/docs/security/public-security-reports.md +++ b/docs/security/public-security-reports.md @@ -59,7 +59,7 @@ These issues were filed as concrete vulnerability reports, security audit findin | [#616](https://github.com/Dstack-TEE/dstack/issues/616) Host-controlled Docker registry mirror enables image substitution attacks | Closed | Not a production vulnerability | Registry mirrors are untrusted transport. Digest-pinned image references and measured compose configuration protect against substitution. No code fix was applied | | [#617](https://github.com/Dstack-TEE/dstack/issues/617) Guest agent exposes raw private keys to all local processes | Closed | Not a production vulnerability | dstack treats a CVM as one application trust domain. It does not provide per-container key isolation inside the same measured application. No code fix was applied | | [#618](https://github.com/Dstack-TEE/dstack/issues/618) Disk encryption disableable via kernel cmdline, not measured in RTMR | Closed | Not a production vulnerability | The kernel command line is measured into RTMR2, so changing `dstack.storage_encrypted=false` changes attestation evidence. No code fix was applied | -| [#619](https://github.com/Dstack-TEE/dstack/issues/619) KMS `get_temp_ca_cert` returns temp CA private key without authentication | Closed | Duplicate | The report duplicates the private advisory response for the temp CA bootstrap flow. [#1106](https://github.com/Dstack-TEE/dstack/pull/1106) removes the reason the temp CA existed - the KMS no longer pins it - so callers can be migrated off it; the root-epoch rotation stack removes the RPC and temporary CA material after migrating guests and KMS-to-KMS onboarding | +| [#619](https://github.com/Dstack-TEE/dstack/issues/619) KMS `get_temp_ca_cert` returns temp CA private key without authentication | Closed | Duplicate | The report duplicates the private advisory response for the temp CA bootstrap flow. [#1106](https://github.com/Dstack-TEE/dstack/pull/1106) removes the reason the temp CA existed - the KMS no longer pins it - so callers can be migrated off it; the RPC and temporary CA material remain available for existing OS images and legacy onboarding while new guests use self-issued certificates | ## Related security roadmap and hardening diff --git a/docs/security/security-best-practices.md b/docs/security/security-best-practices.md index 31cba2658..e3f61a666 100644 --- a/docs/security/security-best-practices.md +++ b/docs/security/security-best-practices.md @@ -111,7 +111,7 @@ Development settings are intentionally easy to audit, but they are not productio - The KMS contract pins a concrete gateway app id. Do not use `gateway_app_id = "any"` for production traffic. - TEE quotes are evaluated by deployment policy, including TCB status and expected OS/application measurements. -The KMS TLS listener verifies client certificates by the attestation they carry rather than by an issuer CA, so it needs no `rpc.tls.mutual` section. It still accepts connections without a client certificate, because bootstrap and public metadata endpoints must be reachable before a client has an RA-TLS certificate. Guests and onboarding clients use self-issued certificates. The temporary-CA RPC and key material have been removed; older clients that require them must be upgraded. +The KMS TLS listener verifies client certificates by the attestation they carry rather than by an issuer CA, so it needs no `rpc.tls.mutual` section. It still accepts connections without a client certificate, because bootstrap and public metadata endpoints must be reachable before a client has an RA-TLS certificate. Guests and onboarding clients use self-issued certificates. The temporary-CA RPC remains available for old OS images and tokenless legacy onboarding; its public key material authenticates no caller. App key release and KMS key handover still require verified caller attestation from the RA-TLS client certificate. Certificate signing verifies the CSR signature and embedded attestation before signing. diff --git a/docs/security/security-model.md b/docs/security/security-model.md index d89e703ef..ccf70629f 100644 --- a/docs/security/security-model.md +++ b/docs/security/security-model.md @@ -425,7 +425,7 @@ The KMS Rocket TLS listener permits connections without a client certificate bec App key release and KMS key handover require verified caller attestation from the RA-TLS client certificate. Certificate signing verifies the CSR signature and the attestation embedded in the CSR before signing. -The unauthenticated or non-client-certificate surface includes bootstrap, env-encryption public-key retrieval, metadata, health, and metrics behavior documented for operators. Clients create self-issued certificates locally. The temporary-CA RPC and its key material have been removed. Guest key retrieval verifies the actual response connection's attestation and checks that the returned CA signed its RPC certificate. Provider pins and launch measurement continue to bind the selected KMS root identity; attestation alone does not authorize an arbitrary CVM as the expected KMS. +The unauthenticated or non-client-certificate surface includes bootstrap, env-encryption public-key retrieval, metadata, health, and metrics behavior documented for operators. Clients create self-issued certificates locally. The temporary-CA RPC and its public key material remain available for legacy clients. Guest key retrieval verifies the actual response connection's attestation and checks that the returned CA signed its RPC certificate. Provider pins and launch measurement continue to bind the selected KMS root identity; attestation alone does not authorize an arbitrary CVM as the expected KMS. KMS-to-KMS root-key handover is served by `Admin.GetKmsKey` on the authenticated admin listener, which additionally requires a verified RA-TLS client certificate. diff --git a/dstack/kms/README.md b/dstack/kms/README.md index 4e56c057e..5deb3378f 100644 --- a/dstack/kms/README.md +++ b/dstack/kms/README.md @@ -315,11 +315,16 @@ root keys can be generated and attested before registering epoch 0. For ## Direct attested key retrieval -`GetTempCaCert` and temporary CA generation, persistence, renewal, and handover -are removed. Its former handover field number and name are reserved in protobuf. -Existing temporary CA files are ignored; they are not needed to start the KMS. -Clients that require the removed RPC, and onboarding sources that require -CA-issued client certificates, are no longer supported. +`GetTempCaCert`, its original response fields, and the `temp_ca_key` handover +field remain available for existing OS images and legacy KMS onboarding. The +compatibility CA is stored in `tmp-ca.key`/`tmp-ca.crt`, independent of all root +epochs; existing files are reused and missing files are recreated. Its public +private key authenticates nobody. Client identity still requires attestation. +The returned `ca_cert` is the currently served root CA. + +Upgrade contracts and auth backends, register the original roots as epoch 0, +then upgrade KMS nodes before guests. Keep root and app epochs at zero during +the mixed-version phase; rotate only after guests support key migration. Guests and `dstack-util get-keys` generate their self-issued RA-TLS certificate locally, then call `GetAppKey` directly. They verify the response connection's @@ -331,8 +336,9 @@ connection. The optional provider pin and the key-provider launch measurement still bind the root identity; a valid quote alone is not a KMS policy decision. `--root-ca` additionally pins the CLI connection before sending the request. -Onboarding uses the same self-issued client flow with or without an admin token, -and checks the source attestation captured from the key-carrying connection +Onboarding with an admin token uses a self-issued client certificate. Tokenless +onboarding obtains the source temp CA first so pre-0.6 sources remain usable. +Both paths check the source attestation captured from the key-carrying connection against local KMS policy before importing anything. Retries and reconnects repeat verification on the response connection; no assumption relies on one TCP session lasting forever. All subsequent key requests pin the KMS CA established at startup, regardless diff --git a/dstack/kms/rpc/proto/kms_rpc.proto b/dstack/kms/rpc/proto/kms_rpc.proto index 25e21975e..13083647b 100644 --- a/dstack/kms/rpc/proto/kms_rpc.proto +++ b/dstack/kms/rpc/proto/kms_rpc.proto @@ -78,12 +78,18 @@ message KmsKeys { } message KmsKeyResponse { - reserved 1; - reserved "temp_ca_key"; + // Retained for legacy KMS onboarding. This public CA authenticates no caller. + string temp_ca_key = 1; // Root keys of every epoch the KMS holds; the index is the epoch. repeated KmsKeys keys = 2; } +message GetTempCaCertResponse { + string temp_ca_cert = 1; + string temp_ca_key = 2; + string ca_cert = 3; +} + message SignCertRequest { uint32 api_version = 1; bytes csr = 2; @@ -105,6 +111,8 @@ service KMS { rpc GetAppEnvEncryptPubKey(AppId) returns (PublicKeyResponse); // Request the KMS instance metadata for use as a probe and health check. rpc GetMeta(google.protobuf.Empty) returns (GetMetaResponse); + // Compatibility for existing OS images and tokenless legacy onboarding. + rpc GetTempCaCert(google.protobuf.Empty) returns (GetTempCaCertResponse); // Sign a certificate rpc SignCert(SignCertRequest) returns (SignCertResponse); } diff --git a/dstack/kms/src/config.rs b/dstack/kms/src/config.rs index 25bb10f3d..507e4f136 100644 --- a/dstack/kms/src/config.rs +++ b/dstack/kms/src/config.rs @@ -13,6 +13,8 @@ pub fn load_config_figment(config_file: Option<&str>) -> Figment { load_config("kms", DEFAULT_CONFIG, config_file, false) } +const TEMP_CA_CERT: &str = "tmp-ca.crt"; +const TEMP_CA_KEY: &str = "tmp-ca.key"; const ROOT_CA_CERT: &str = "root-ca.crt"; const ROOT_CA_KEY: &str = "root-ca.key"; const RPC_CERT: &str = "rpc.crt"; @@ -137,6 +139,14 @@ impl KmsConfig { && self.k256_key().exists() } + pub fn tmp_ca_cert(&self) -> PathBuf { + self.cert_dir.join(TEMP_CA_CERT) + } + + pub fn tmp_ca_key(&self) -> PathBuf { + self.cert_dir.join(TEMP_CA_KEY) + } + pub fn root_ca_cert(&self) -> PathBuf { self.cert_dir.join(ROOT_CA_CERT) } diff --git a/dstack/kms/src/main_service.rs b/dstack/kms/src/main_service.rs index fc7871d42..25f3c0c23 100644 --- a/dstack/kms/src/main_service.rs +++ b/dstack/kms/src/main_service.rs @@ -13,8 +13,9 @@ use std::{ use anyhow::{anyhow, bail, Context, Result}; use dstack_kms_rpc::{ kms_server::{KmsRpc, KmsServer}, - AppId, AppKeyResponse, GetAppKeyRequest, GetKmsKeyRequest, GetMetaResponse, KmsKeyResponse, - KmsKeys, PublicKeyResponse, SignCertRequest, SignCertResponse, + AppId, AppKeyResponse, GetAppKeyRequest, GetKmsKeyRequest, GetMetaResponse, + GetTempCaCertResponse, KmsKeyResponse, KmsKeys, PublicKeyResponse, SignCertRequest, + SignCertResponse, }; use dstack_verifier::{CvmVerifier, VerificationDetails}; use fs_err as fs; @@ -163,6 +164,7 @@ pub(crate) fn servable_epoch( pub struct KmsStateInner { config: KmsConfig, keyring: Keyring, + temp_ca: CaCert, /// The CA of the running TLS listener, independent of key selection. rpc_ca_pubkey: Vec, restart: tokio::sync::Notify, @@ -246,6 +248,7 @@ impl KmsState { } Ok(Self { inner: Arc::new(KmsStateInner { + temp_ca: crate::onboard_service::ensure_temp_ca(&config)?, config, keyring, rpc_ca_pubkey, @@ -406,6 +409,7 @@ impl RpcHandler { let info = self.ensure_kms_allowed(&request.vm_config).await?; ensure_key_release_allowed(&info, (&self.state.config).into())?; Ok(KmsKeyResponse { + temp_ca_key: self.state.temp_ca.key.serialize_pem(), keys: self .state .keyring @@ -729,6 +733,24 @@ impl KmsRpc for RpcHandler { self.handover_keys(request).await } + async fn get_temp_ca_cert(self) -> Result { + let epoch = self.ensure_self_allowed().await?; + if self.state.config.enforce_self_authorization { + let boot_info = self + .state + .self_boot_info + .get(&()) + .await + .context("missing authorized KMS boot info")?; + ensure_key_release_allowed(&boot_info, (&self.state.config).into())?; + } + Ok(GetTempCaCertResponse { + temp_ca_cert: self.state.temp_ca.pem_cert.clone(), + temp_ca_key: self.state.temp_ca.key.serialize_pem(), + ca_cert: self.state.epoch_root(epoch)?.ca.pem_cert.clone(), + }) + } + async fn sign_cert(self, request: SignCertRequest) -> Result { self.ensure_self_allowed() .await @@ -1562,6 +1584,24 @@ mod tests { const AUTH_API_INFO: &str = r#"{"status":"ok","kmsContractAddr":"0xkms","ethRpcUrl":"https://rpc.example","gatewayAppId":"0xgateway","chainId":1,"appImplementation":"0ximpl"}"#; + #[rocket::async_test] + async fn legacy_temp_ca_rpc_returns_a_separate_ca_and_the_served_root() { + assert!(dstack_kms_rpc::kms_server::KMSMethod::from_str("GetTempCaCert").is_some()); + let dir = tempfile::tempdir().unwrap(); + let (url, _) = serve_concurrent_auth_api(AUTH_API_INFO); + let state = kms_state(dir.path(), &url); + let response = RpcHandler { + state: state.clone(), + attestation: None, + } + .get_temp_ca_cert() + .await + .unwrap(); + let temp_ca = CaCert::new(response.temp_ca_cert, response.temp_ca_key).unwrap(); + assert_ne!(temp_ca.key.public_key_der(), state.rpc_ca_pubkey); + assert_eq!(response.ca_cert, state.epoch_root(0).unwrap().ca.pem_cert); + } + #[rocket::async_test] async fn concurrent_get_meta_calls_share_one_auth_api_request() { let dir = tempfile::tempdir().unwrap(); diff --git a/dstack/kms/src/onboard_service.rs b/dstack/kms/src/onboard_service.rs index fb0f8a7d6..31be460a0 100644 --- a/dstack/kms/src/onboard_service.rs +++ b/dstack/kms/src/onboard_service.rs @@ -18,7 +18,7 @@ use dstack_kms_rpc::{ use fs_err as fs; use k256::ecdsa::SigningKey; use ra_rpc::{ - client::{CertInfo, RaClientConfig}, + client::{CertInfo, RaClient, RaClientConfig}, CallContext, RpcCall, }; use ra_tls::{ @@ -26,7 +26,7 @@ use ra_tls::{ AttestationVerifier, GetDeviceId, PlatformEvidence, QuoteContentType, VerifiedAttestation, VersionedAttestation, }, - cert::CertRequest, + cert::{CaCert, CertRequest}, rcgen::{Certificate, KeyPair, PKCS_ECDSA_P256_SHA256}, }; use safe_write::{safe_write, safe_write_with_mode}; @@ -548,6 +548,7 @@ mod tests { fn handover(epochs: u32) -> KmsKeyResponse { KmsKeyResponse { + temp_ca_key: String::new(), keys: (0..epochs) .map(|_| KmsKeys { ca_key: KeyPair::generate_for(&PKCS_ECDSA_P256_SHA256) @@ -595,6 +596,20 @@ mod tests { Keys::from_key_response(cfg, res, "kms.example.com", rotate, registry).await } + #[test] + fn compatibility_ca_survives_restart_and_recovers_a_missing_certificate() { + let dir = tempfile::tempdir().unwrap(); + let cfg = auto_bootstrap_config(dir.path()); + let first = ensure_temp_ca(&cfg).unwrap(); + let second = ensure_temp_ca(&cfg).unwrap(); + assert_eq!(first.key.public_key_der(), second.key.public_key_der()); + assert_eq!(first.pem_cert, second.pem_cert); + fs::remove_file(cfg.tmp_ca_cert()).unwrap(); + let recovered = ensure_temp_ca(&cfg).unwrap(); + assert_eq!(first.key.public_key_der(), recovered.key.public_key_der()); + assert!(!cfg.root_keys_exist()); + } + #[rocket::async_test] async fn join_imports_and_stores_every_epoch() { let cert_dir = tempfile::tempdir().unwrap(); @@ -606,8 +621,8 @@ mod tests { .await .unwrap(); keys.store(&cfg).unwrap(); - assert!(!cfg.cert_dir.join("tmp-ca.key").exists()); - assert!(!cfg.cert_dir.join("tmp-ca.crt").exists()); + assert!(cfg.tmp_ca_key().exists()); + assert!(cfg.tmp_ca_cert().exists()); assert_eq!(crate::main_service::read_root_epoch(&cfg).unwrap(), 2); for (epoch, k256) in (0..).zip(&sent) { @@ -937,7 +952,14 @@ impl Keys { .context("Missing source KMS attestation") }; - let (ra_cert, ra_key) = gen_ra_cert().await?; + let temp_ca = if source_token.is_empty() { + let client = RaClient::new(other_kms_url.to_string(), true)?; + let ca = KmsClient::new(client).get_temp_ca_cert().await?; + Some((ca.temp_ca_cert, ca.temp_ca_key)) + } else { + None + }; + let (ra_cert, ra_key) = gen_ra_cert(temp_ca).await?; let client = RaClientConfig::builder() .tls_skip_ca_verification(true) .tls_built_in_root_certs(false) @@ -1033,6 +1055,7 @@ impl Keys { /// The root keys go last: see [`KmsConfig::root_keys_exist`]. fn store(&self, cfg: &KmsConfig) -> Result<()> { + ensure_temp_ca(cfg)?; for (epoch, root) in (0..).zip(&self.history) { let paths = cfg.root_key_paths(epoch, self.epoch); if let Some(dir) = paths.ca_key.parent() { @@ -1063,9 +1086,39 @@ impl Keys { } } +/// Keep the public compatibility CA independent of all epoch root keys. +pub(crate) fn ensure_temp_ca(cfg: &KmsConfig) -> Result { + let had_key = cfg.tmp_ca_key().exists(); + let key = if had_key { + KeyPair::from_pem(&fs::read_to_string(cfg.tmp_ca_key())?)? + } else { + let key = KeyPair::generate_for(&PKCS_ECDSA_P256_SHA256)?; + safe_write_with_mode(cfg.tmp_ca_key(), key.serialize_pem(), 0o600)?; + key + }; + let cert = CertRequest::builder() + .org_name("Dstack") + .subject("Dstack Client Temp CA") + .ca_level(0) + .key(&key) + .build() + .self_signed()?; + if had_key && cfg.tmp_ca_cert().exists() { + renew_ca_cert_if_expiring( + cfg.tmp_ca_cert(), + cert.pem(), + "temporary client CA certificate", + )?; + } else { + safe_write(cfg.tmp_ca_cert(), cert.pem())?; + } + CaCert::load(cfg.tmp_ca_cert(), cfg.tmp_ca_key()) +} + /// Reissue the RPC certificate under the root CA of the served `epoch`, so /// guests reach this KMS through the CA of the keys it serves. pub(crate) async fn update_certs(cfg: &KmsConfig, epoch: u32, latest: u32) -> Result<()> { + ensure_temp_ca(cfg)?; let roots = cfg.root_key_paths(epoch, latest); // Read existing keys let ca_key = KeyPair::from_pem(&fs::read_to_string(&roots.ca_key)?)?; @@ -1206,8 +1259,8 @@ fn keccak256(msg: &[u8]) -> [u8; 32] { hasher.finalize().into() } -/// Mint a self-issued RA-TLS client certificate. -async fn gen_ra_cert() -> Result<(String, String)> { +/// Legacy sources require a temp-CA-issued client certificate. +async fn gen_ra_cert(ca: Option<(String, String)>) -> Result<(String, String)> { use ra_tls::cert::CertRequest; use ra_tls::rcgen::{KeyPair, PKCS_ECDSA_P256_SHA256}; @@ -1225,6 +1278,10 @@ async fn gen_ra_cert() -> Result<(String, String)> { .key(&key) .usage_client_auth(true) .build(); - let cert = req.self_signed().context("Failed to sign certificate")?; + let cert = match ca { + Some((cert, key)) => CaCert::new(cert, key)?.sign(req), + None => req.self_signed(), + } + .context("failed to sign client certificate")?; Ok((cert.pem(), key.serialize_pem())) } From 8608761c7054955b00c8509bbfe1ca286ae920ec Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Sat, 7 Nov 2026 18:42:49 -0800 Subject: [PATCH 033/166] feat(kms): authorize explicit operations on both sides of handover Signed-off-by: Kevin Wang --- dstack/crates/dstack-auth/src/main.rs | 30 ++++- dstack/dstack-types/src/authorization.rs | 27 +++++ dstack/dstack-types/src/lib.rs | 2 + dstack/kms/README.md | 36 ++++-- dstack/kms/auth-eth-bun/README.md | 7 ++ dstack/kms/auth-eth-bun/index.test.ts | 5 +- dstack/kms/auth-eth-bun/index.ts | 5 +- dstack/kms/auth-eth-bun/openapi.json | 16 ++- dstack/kms/auth-eth/README.md | 7 ++ .../contracts/AuthorizationOperation.sol | 24 ++++ dstack/kms/auth-eth/contracts/DstackKms.sol | 19 +++- dstack/kms/auth-eth/src/ethereum.ts | 5 +- dstack/kms/auth-eth/src/main.test.ts | 2 +- dstack/kms/auth-eth/src/server.ts | 2 + dstack/kms/auth-eth/src/types.ts | 1 + dstack/kms/auth-eth/test/DstackKms.t.sol | 36 +++++- dstack/kms/auth-mock/README.md | 7 ++ dstack/kms/auth-mock/index.test.ts | 2 +- dstack/kms/auth-mock/index.ts | 9 +- dstack/kms/auth-mock/openapi.json | 16 ++- dstack/kms/auth-simple/README.md | 7 ++ dstack/kms/auth-simple/index.test.ts | 11 +- dstack/kms/auth-simple/index.ts | 9 +- dstack/kms/src/main_service.rs | 101 +++++++++++++---- .../kms/src/main_service/upgrade_authority.rs | 106 +++++++++++++++--- dstack/kms/src/onboard_service.rs | 35 +++--- 26 files changed, 438 insertions(+), 89 deletions(-) create mode 100644 dstack/dstack-types/src/authorization.rs create mode 100644 dstack/kms/auth-eth/contracts/AuthorizationOperation.sol diff --git a/dstack/crates/dstack-auth/src/main.rs b/dstack/crates/dstack-auth/src/main.rs index 50d8ef937..062228be7 100644 --- a/dstack/crates/dstack-auth/src/main.rs +++ b/dstack/crates/dstack-auth/src/main.rs @@ -22,6 +22,7 @@ use anyhow::{Context, Result}; use clap::Parser; +use dstack_types::authorization::AuthorizationOperation as Operation; use rocket::http::Status; use rocket::serde::json::Json; use rocket::{get, post, routes, State}; @@ -76,6 +77,7 @@ struct BootResponse { #[derive(Deserialize)] #[serde(rename_all = "camelCase")] struct AuthorizationRequest { + operation: Operation, #[serde(flatten)] boot_info: BootInfo, requested_root_epoch: Option, @@ -314,7 +316,9 @@ fn authorize( current_app_epoch: 0, }; let info = &request.boot_info; - result.rejection_reason = if current.is_none() { + result.rejection_reason = if request.operation.is_kms() != is_kms { + "invalid authorization operation".into() + } else if current.is_none() && request.operation != Operation::Bootstrap { "root epoch not registered".into() } else if request .requested_root_epoch @@ -323,9 +327,11 @@ fn authorize( "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 norm(&al.kms_contract_addr).is_empty() { + } else if request.operation != Operation::Bootstrap && norm(&al.kms_contract_addr).is_empty() { "kmsContractAddr is not configured".into() - } else if (norm(&info.app_id) == norm(&al.kms_contract_addr)) != is_kms { + } else if request.operation != Operation::Bootstrap + && (norm(&info.app_id) == norm(&al.kms_contract_addr)) != is_kms + { "caller does not match authorization scope".into() } else { let policy = if is_kms { @@ -472,8 +478,8 @@ mod tests { })) .await .unwrap(); - let kms = json!({"appId": "11".repeat(20), "mrAggregated": "bb"}); - let app = json!({"appId": "22".repeat(20), "composeHash": "cc", "osImageHash": "aa"}); + let kms = json!({"operation": "serve", "appId": "11".repeat(20), "mrAggregated": "bb"}); + let app = json!({"operation": "get_app_key", "appId": "22".repeat(20), "composeHash": "cc", "osImageHash": "aa"}); for (route, body, allowed) in [ ("/authorizeKms", kms.clone(), true), ("/authorizeApp", app.clone(), true), @@ -496,6 +502,20 @@ mod tests { format!("0x{}", "00".repeat(32)) ); } + for operation in [serde_json::Value::Null, json!("unknown")] { + let mut body = app.clone(); + body["operation"] = operation; + let response = client + .post("/authorizeApp") + .header(ContentType::JSON) + .body(body.to_string()) + .dispatch() + .await; + assert!(matches!( + response.status(), + Status::BadRequest | Status::UnprocessableEntity + )); + } for field in ["requestedRootEpoch", "requestedAppEpoch"] { let mut body = app.clone(); body[field] = 1.into(); diff --git a/dstack/dstack-types/src/authorization.rs b/dstack/dstack-types/src/authorization.rs new file mode 100644 index 000000000..b3f5ecec9 --- /dev/null +++ b/dstack/dstack-types/src/authorization.rs @@ -0,0 +1,27 @@ +// SPDX-FileCopyrightText: © 2026 Phala Network +// +// SPDX-License-Identifier: Apache-2.0 + +//! Operations performed by the attested subject of an authorization request. + +use serde::{Deserialize, Serialize}; + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum AuthorizationOperation { + Bootstrap, + Serve, + ExportKmsKeys, + ImportKmsKeys, + GetAppKey, + SignCert, +} + +impl AuthorizationOperation { + pub fn is_kms(self) -> bool { + matches!( + self, + Self::Bootstrap | Self::Serve | Self::ExportKmsKeys | Self::ImportKmsKeys + ) + } +} diff --git a/dstack/dstack-types/src/lib.rs b/dstack/dstack-types/src/lib.rs index cc1bdf530..c9b276f6e 100644 --- a/dstack/dstack-types/src/lib.rs +++ b/dstack/dstack-types/src/lib.rs @@ -14,6 +14,8 @@ use size_parser::human_size; /// infrastructure-provider initialization stages. pub const MAX_INIT_SCRIPTS: usize = 5; +pub mod authorization; + pub mod init_script_hashes { use serde::{Deserialize, Deserializer, Serialize, Serializer}; use serde_human_bytes::ByteBuf; diff --git a/dstack/kms/README.md b/dstack/kms/README.md index 5deb3378f..0fd0b0803 100644 --- a/dstack/kms/README.md +++ b/dstack/kms/README.md @@ -246,15 +246,15 @@ The `SignCert` RPC is used by the dstack app to sign a TLS certificate. In this ## Authorization protocol -Normal operations use two explicit authorization endpoints: +Authorization uses two explicit endpoints: - `POST /authorizeKms` for KMS self-authorization and root-key handover. - `POST /authorizeApp` for application keys and certificates. The Ethereum backends call the matching `DstackKms` method. Both accept `AuthorizationRequest` and return `AuthorizationResult`; no caller-kind flag is -sent. The KMS method requires the attested `appId` to equal the proxy address, -and the app method rejects that address and applies registered app policy. +sent. The KMS method checks the KMS policy; the app method applies registered app +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 @@ -269,8 +269,29 @@ this authorization result; an explicit `0` selects epoch 0. Solidity represents it as `OptionalEpoch { 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 omits this selector and still transfers the complete keyring, -including pending epochs. +KMS-to-KMS handover transfers the complete keyring. The receiver explicitly +authorizes the source to export the highest received epoch before importing it. + +Every request includes a required `operation` describing what the attested +subject is authorized to do. JSON uses `snake_case`; the contract receives +`keccak256(bytes(operation))` as `bytes32`. Rust uses the shared serde-derived +`AuthorizationOperation` enum. Unknown operations and operations on the wrong +endpoint are rejected. The supported operations currently share the existing +material policy, leaving room for independent policies later: + +| Endpoint | Operation | Subject | +| --- | --- | --- | +| `authorizeKms` | `bootstrap` | KMS generating its first roots | +| `authorizeKms` | `serve` | KMS serving app keys, public keys or certificates | +| `authorizeKms` | `export_kms_keys` | Source KMS handing over its keyring | +| `authorizeKms` | `import_kms_keys` | Recipient KMS receiving a keyring | +| `authorizeApp` | `get_app_key` | App receiving derived keys | +| `authorizeApp` | `sign_cert` | App requesting a certificate | + +The KMS chooses the operation from its code path, not an untrusted RPC field. +Self-authorization caches are keyed by operation. Legacy authorization endpoints +retain their existing semantics for old KMS binaries; fine-grained operation +policies cannot protect calls still served by those old binaries. The response is an `AuthorizationResult`: @@ -306,8 +327,9 @@ Upgrade the contract and auth backend and register epoch 0 before using this protocol for key release. Runtime authorization never falls back to the legacy API if the new endpoint or contract method is unavailable. The old contract methods and `/bootAuth/app` and `/bootAuth/kms` endpoints remain available for -older clients. Bootstrap and onboarding retain scoped KMS checks so the first -root keys can be generated and attested before registering epoch 0. For +older clients. Bootstrap uses the `bootstrap` operation, which checks KMS policy at epoch zero +even before roots are registered. Other operations require registered roots. +Onboarding checks its source through `authorizeKms`, without a legacy fallback. For `auth-simple`, set `kmsContractAddr` to the attested KMS app ID and configure `rootEpochs`; for `auth-mock`, set `KMS_CONTRACT_ADDR`, `MOCK_ROOT_EPOCH`, and `MOCK_ROOT_EPOCH_HASH`. diff --git a/dstack/kms/auth-eth-bun/README.md b/dstack/kms/auth-eth-bun/README.md index 92d50cced..bc9c72e19 100644 --- a/dstack/kms/auth-eth-bun/README.md +++ b/dstack/kms/auth-eth-bun/README.md @@ -213,3 +213,10 @@ epoch, and an explicit `0` selects epoch 0. See the [KMS authorization protocol](../README.md#authorization-protocol) for registration, bootstrap, and compatibility requirements. The legacy `/bootAuth/*` endpoints remain available. + +### Authorization operations + +`/authorizeKms` and `/authorizeApp` require an `operation` string. See the +[KMS authorization protocol](../README.md#authorization-protocol) for the names +and subject semantics. Bootstrap alone may authorize before root registration; +unknown operations are rejected. Legacy `/bootAuth/*` routes are unchanged. diff --git a/dstack/kms/auth-eth-bun/index.test.ts b/dstack/kms/auth-eth-bun/index.test.ts index 2ce209916..132fbf67f 100644 --- a/dstack/kms/auth-eth-bun/index.test.ts +++ b/dstack/kms/auth-eth-bun/index.test.ts @@ -10,7 +10,8 @@ const mockReadContract = vi.fn(); const mockGetChainId = vi.fn(); const mockGetBlockNumber = vi.fn(); -vi.mock('viem', () => ({ +vi.mock('viem', async (importOriginal) => ({ + ...await importOriginal(), createPublicClient: vi.fn(() => ({ readContract: mockReadContract, getChainId: mockGetChainId, @@ -842,7 +843,7 @@ describe('POST /authorizeApp', () => { osImageHash: '0x' + '66'.repeat(32), tcbStatus: 'UpToDate', advisoryIds: [], }; const authorize = (requestedRootEpoch?: number) => appFetch(new Request('http://localhost/authorizeApp', { - method: 'POST', headers: { 'Content-Type': 'application/json' }, body: JSON.stringify({ ...appInfo, requestedRootEpoch }), + method: 'POST', headers: { 'Content-Type': 'application/json' }, body: JSON.stringify({ ...appInfo, operation: 'get_app_key', requestedRootEpoch }), })); it('reads the complete result with one contract call at configured finality', async () => { diff --git a/dstack/kms/auth-eth-bun/index.ts b/dstack/kms/auth-eth-bun/index.ts index 540ae00d5..8361d5a90 100644 --- a/dstack/kms/auth-eth-bun/index.ts +++ b/dstack/kms/auth-eth-bun/index.ts @@ -7,7 +7,7 @@ import { bodyLimit } from 'hono/body-limit'; import { secureHeaders } from 'hono/secure-headers'; import { zValidator } from '@hono/zod-validator'; import { z } from 'zod'; -import { createPublicClient, http, type Address, type Hex } from 'viem'; +import { createPublicClient, http, keccak256, toBytes, type Address, type Hex } from 'viem'; // zod schemas for validation - compatible with original fastify implementation const boundedHex = (bytes: number, description: string) => @@ -37,6 +37,7 @@ const BootInfoSchema = z.object({ }); const AuthorizationRequestSchema = BootInfoSchema.omit({ rootEpoch: true, rootEpochHash: true }).extend({ + operation: z.enum(['bootstrap', 'serve', 'export_kms_keys', 'import_kms_keys', 'get_app_key', 'sign_cert']), requestedAppEpoch: z.number().int().min(0).max(0xffffffff).default(0), requestedRootEpoch: z.number().int().min(0).max(0xffffffff).optional(), }); @@ -94,6 +95,7 @@ const AUTHORIZATION_ABI = { stateMutability: 'view', inputs: [{ name: 'request', type: 'tuple', components: [ { name: 'appInfo', type: 'tuple', components: APP_INFO_COMPONENTS }, + { name: 'operation', type: 'bytes32' }, { name: 'requestedAppEpoch', type: 'tuple', components: [ { name: 'isSet', type: 'bool' }, { name: 'value', type: 'uint256' }, ] }, @@ -248,6 +250,7 @@ class EthereumBackend { functionName: method, args: [{ appInfo: this.contractAppInfo(appInfo), + operation: keccak256(toBytes(appInfo.operation)), requestedAppEpoch: { isSet: appInfo.requestedAppEpoch !== undefined, value: BigInt(appInfo.requestedAppEpoch ?? 0) }, requestedRootEpoch: { isSet: appInfo.requestedRootEpoch !== undefined, value: BigInt(appInfo.requestedRootEpoch ?? 0) }, }], diff --git a/dstack/kms/auth-eth-bun/openapi.json b/dstack/kms/auth-eth-bun/openapi.json index b7030e175..cab64187a 100644 --- a/dstack/kms/auth-eth-bun/openapi.json +++ b/dstack/kms/auth-eth-bun/openapi.json @@ -428,8 +428,22 @@ "maximum": 4294967295, "default": 0, "description": "Application epoch, independent of KMS root epochs. Only zero is currently supported." + }, + "operation": { + "type": "string", + "enum": [ + "bootstrap", + "serve", + "export_kms_keys", + "import_kms_keys", + "get_app_key", + "sign_cert" + ] } - } + }, + "required": [ + "operation" + ] } } } diff --git a/dstack/kms/auth-eth/README.md b/dstack/kms/auth-eth/README.md index 47f880315..c7a8d6e7b 100644 --- a/dstack/kms/auth-eth/README.md +++ b/dstack/kms/auth-eth/README.md @@ -184,3 +184,10 @@ epoch, and an explicit `0` selects epoch 0. See the [KMS authorization protocol](../README.md#authorization-protocol) for registration, bootstrap, and compatibility requirements. The legacy `/bootAuth/*` endpoints remain available. + +### Authorization operations + +`/authorizeKms` and `/authorizeApp` require an `operation` string. See the +[KMS authorization protocol](../README.md#authorization-protocol) for the names +and subject semantics. Bootstrap alone may authorize before root registration; +unknown operations are rejected. Legacy `/bootAuth/*` routes are unchanged. diff --git a/dstack/kms/auth-eth/contracts/AuthorizationOperation.sol b/dstack/kms/auth-eth/contracts/AuthorizationOperation.sol new file mode 100644 index 000000000..2f60eed4a --- /dev/null +++ b/dstack/kms/auth-eth/contracts/AuthorizationOperation.sol @@ -0,0 +1,24 @@ +// SPDX-FileCopyrightText: © 2026 Phala Network +// +// SPDX-License-Identifier: Apache-2.0 + +pragma solidity ^0.8.24; + +/// The operation is performed by the attested subject, not the policy caller. +library AuthorizationOperation { + bytes32 internal constant BOOTSTRAP = keccak256("bootstrap"); + bytes32 internal constant SERVE = keccak256("serve"); + bytes32 internal constant EXPORT_KMS_KEYS = keccak256("export_kms_keys"); + bytes32 internal constant IMPORT_KMS_KEYS = keccak256("import_kms_keys"); + bytes32 internal constant GET_APP_KEY = keccak256("get_app_key"); + bytes32 internal constant SIGN_CERT = keccak256("sign_cert"); + + function isKms(bytes32 operation) internal pure returns (bool) { + return + operation == BOOTSTRAP || operation == SERVE || operation == EXPORT_KMS_KEYS || operation == IMPORT_KMS_KEYS; + } + + function isApp(bytes32 operation) internal pure returns (bool) { + return operation == GET_APP_KEY || operation == SIGN_CERT; + } +} diff --git a/dstack/kms/auth-eth/contracts/DstackKms.sol b/dstack/kms/auth-eth/contracts/DstackKms.sol index 465bef86c..9fa11980e 100644 --- a/dstack/kms/auth-eth/contracts/DstackKms.sol +++ b/dstack/kms/auth-eth/contracts/DstackKms.sol @@ -7,6 +7,7 @@ pragma solidity ^0.8.24; import "./IAppAuth.sol"; +import "./AuthorizationOperation.sol"; import "@openzeppelin/contracts-upgradeable/proxy/utils/Initializable.sol"; import "@openzeppelin/contracts-upgradeable/access/Ownable2StepUpgradeable.sol"; import "@openzeppelin/contracts-upgradeable/proxy/utils/UUPSUpgradeable.sol"; @@ -21,6 +22,7 @@ contract DstackKms is Initializable, Ownable2StepUpgradeable, UUPSUpgradeable, E struct AuthorizationRequest { AppInfo appInfo; + bytes32 operation; // Reserved for application key rotation; only zero is currently supported. OptionalUint256 requestedAppEpoch; OptionalUint256 requestedRootEpoch; @@ -410,12 +412,21 @@ contract DstackKms is Initializable, Ownable2StepUpgradeable, UUPSUpgradeable, E returns (AuthorizationResult memory result) { result.gatewayAppId = gatewayAppId; - if (_rootEpochs.length == 0) { - result.rejectionReason = "Root epoch not registered"; + if (!(isKms + ? AuthorizationOperation.isKms(request.operation) + : AuthorizationOperation.isApp(request.operation))) { + result.rejectionReason = "Invalid authorization operation"; return result; } - result.currentRootEpoch = _currentRootEpoch; - result.currentRootEpochHash = _rootEpochs[_currentRootEpoch].chainHash; + if (_rootEpochs.length == 0) { + if (request.operation != AuthorizationOperation.BOOTSTRAP) { + result.rejectionReason = "Root epoch not registered"; + return result; + } + } else { + result.currentRootEpoch = _currentRootEpoch; + result.currentRootEpochHash = _rootEpochs[_currentRootEpoch].chainHash; + } if (request.requestedRootEpoch.isSet && request.requestedRootEpoch.value > _currentRootEpoch) { result.rejectionReason = "Requested root epoch is not active"; return result; diff --git a/dstack/kms/auth-eth/src/ethereum.ts b/dstack/kms/auth-eth/src/ethereum.ts index 8617ffe08..2aedc8b54 100644 --- a/dstack/kms/auth-eth/src/ethereum.ts +++ b/dstack/kms/auth-eth/src/ethereum.ts @@ -7,8 +7,8 @@ import { BootInfo, BootResponse, RootEpochInfo, AuthorizationResult, Authorizati // Minimal ABI for DstackKms contract const DSTACK_KMS_ABI = [ - "function authorizeKms(((address appId,bytes32 composeHash,address instanceId,bytes32 deviceId,bytes32 mrAggregated,bytes32 mrSystem,bytes32 osImageHash,string tcbStatus,string[] advisoryIds) appInfo,(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,(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 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)", "function isKmsAllowed((address appId,bytes32 composeHash,address instanceId,bytes32 deviceId,bytes32 mrAggregated,bytes32 mrSystem,bytes32 osImageHash,string tcbStatus,string[] advisoryIds) bootInfo) view returns (bool, string)", "function gatewayAppId() view returns (string)", @@ -81,6 +81,7 @@ export class EthereumBackend { async checkAuthorization(appInfo: AuthorizationRequest, method: 'authorizeKms' | 'authorizeApp'): Promise { const result = await this.kmsContract[method]({ appInfo: this.contractAppInfo(appInfo), + operation: ethers.id(appInfo.operation), requestedAppEpoch: { isSet: appInfo.requestedAppEpoch !== undefined, value: appInfo.requestedAppEpoch ?? 0 }, requestedRootEpoch: { isSet: appInfo.requestedRootEpoch !== undefined, value: appInfo.requestedRootEpoch ?? 0 }, }); diff --git a/dstack/kms/auth-eth/src/main.test.ts b/dstack/kms/auth-eth/src/main.test.ts index a8d035749..93567235a 100644 --- a/dstack/kms/auth-eth/src/main.test.ts +++ b/dstack/kms/auth-eth/src/main.test.ts @@ -92,7 +92,7 @@ describe('Server', () => { }); describe('POST /authorizeApp', () => { - const appRequest = { ...mockBootInfo, requestedAppEpoch: 0 }; + const appRequest = { ...mockBootInfo, operation: 'get_app_key', requestedAppEpoch: 0 }; it('preserves the complete decision, including the epoch on rejection', async () => { const decision = { isAllowed: false, rejectionReason: 'App not registered', gatewayAppId: 'gateway', diff --git a/dstack/kms/auth-eth/src/server.ts b/dstack/kms/auth-eth/src/server.ts index c7de71c7b..afbc49aa8 100644 --- a/dstack/kms/auth-eth/src/server.ts +++ b/dstack/kms/auth-eth/src/server.ts @@ -53,8 +53,10 @@ export async function build(): Promise { server.addSchema({ ...appInfoSchema, $id: 'authorizationRequest', + required: [...appInfoSchema.required, 'operation'], properties: { ...appInfoSchema.properties, + operation: { type: 'string', enum: ['bootstrap', 'serve', 'export_kms_keys', 'import_kms_keys', 'get_app_key', 'sign_cert'] }, requestedAppEpoch: { type: 'integer', minimum: 0, maximum: 4294967295, default: 0 }, requestedRootEpoch: { type: 'integer', minimum: 0, maximum: 4294967295 }, }, diff --git a/dstack/kms/auth-eth/src/types.ts b/dstack/kms/auth-eth/src/types.ts index 6171f1b01..79db187cc 100644 --- a/dstack/kms/auth-eth/src/types.ts +++ b/dstack/kms/auth-eth/src/types.ts @@ -43,6 +43,7 @@ export interface AuthorizationResult { } export interface AuthorizationRequest extends Omit { + operation: 'bootstrap' | 'serve' | 'export_kms_keys' | 'import_kms_keys' | 'get_app_key' | 'sign_cert'; requestedAppEpoch?: number; requestedRootEpoch?: number; } diff --git a/dstack/kms/auth-eth/test/DstackKms.t.sol b/dstack/kms/auth-eth/test/DstackKms.t.sol index e076d5b19..3f6b9809c 100644 --- a/dstack/kms/auth-eth/test/DstackKms.t.sol +++ b/dstack/kms/auth-eth/test/DstackKms.t.sol @@ -465,11 +465,41 @@ contract DstackKmsTest is Test { returns (DstackKms.AuthorizationResult memory) { DstackKms.AuthorizationRequest memory request = DstackKms.AuthorizationRequest( - info, DstackKms.OptionalUint256(false, 0), DstackKms.OptionalUint256(false, 0) + info, + isKms ? keccak256("serve") : keccak256("get_app_key"), + DstackKms.OptionalUint256(false, 0), + DstackKms.OptionalUint256(false, 0) ); return isKms ? kms.authorizeKms(request) : kms.authorizeApp(request); } + function test_AuthorizationOperationsFailClosedAndBootstrapBeforeRegistration() public { + IAppAuth.AppInfo memory info = authorizationInfo(address(kms)); + vm.startPrank(owner); + kms.addOsImageHash(info.osImageHash); + kms.addKmsDevice(info.deviceId); + kms.addKmsAggregatedMr(info.mrAggregated); + vm.stopPrank(); + DstackKms.AuthorizationRequest memory request = DstackKms.AuthorizationRequest( + info, keccak256("bootstrap"), DstackKms.OptionalUint256(false, 0), DstackKms.OptionalUint256(false, 0) + ); + assertTrue(kms.authorizeKms(request).isAllowed); + assertEq(kms.authorizeKms(request).currentRootEpochHash, bytes32(0)); + request.operation = keccak256("export_kms_keys"); + assertEq(kms.authorizeKms(request).rejectionReason, "Root epoch not registered"); + vm.prank(owner); + kms.registerRootEpoch(ca(0xa0), k256(0xb0), "https://example.com/epoch-0"); + assertTrue(kms.authorizeKms(request).isAllowed); + request.operation = keccak256("import_kms_keys"); + assertTrue(kms.authorizeKms(request).isAllowed); + assertEq(kms.authorizeApp(request).rejectionReason, "Invalid authorization operation"); + request.operation = keccak256("sign_cert"); + assertEq(kms.authorizeKms(request).rejectionReason, "Invalid authorization operation"); + request.operation = bytes32(0); + assertEq(kms.authorizeKms(request).rejectionReason, "Invalid authorization operation"); + assertEq(kms.authorizeApp(request).rejectionReason, "Invalid authorization operation"); + } + function test_CheckAuthorization_OptionalUint256DistinguishesZeroFromDefault() public { IAppAuth.AppInfo memory info = authorizationInfo(address(kms)); vm.startPrank(owner); @@ -481,7 +511,7 @@ contract DstackKmsTest is Test { vm.stopPrank(); DstackKms.AuthorizationRequest memory request = DstackKms.AuthorizationRequest( - info, DstackKms.OptionalUint256(false, 0), DstackKms.OptionalUint256(true, 1) + info, keccak256("serve"), DstackKms.OptionalUint256(false, 0), DstackKms.OptionalUint256(true, 1) ); DstackKms.AuthorizationResult memory result = kms.authorizeKms(request); assertFalse(result.isAllowed); @@ -539,7 +569,7 @@ contract DstackKmsTest is Test { assertFalse(result.isAllowed); assertEq(result.rejectionReason, "Aggregated MR not allowed"); DstackKms.AuthorizationRequest memory appRequest = DstackKms.AuthorizationRequest( - info, DstackKms.OptionalUint256(true, 1), DstackKms.OptionalUint256(false, 0) + info, keccak256("get_app_key"), DstackKms.OptionalUint256(true, 1), DstackKms.OptionalUint256(false, 0) ); result = kms.authorizeApp(appRequest); assertFalse(result.isAllowed); diff --git a/dstack/kms/auth-mock/README.md b/dstack/kms/auth-mock/README.md index 8d69c0a06..4ddb11646 100644 --- a/dstack/kms/auth-mock/README.md +++ b/dstack/kms/auth-mock/README.md @@ -181,3 +181,10 @@ epoch, and an explicit `0` selects epoch 0. See the [KMS authorization protocol](../README.md#authorization-protocol) for registration, bootstrap, and compatibility requirements. The legacy `/bootAuth/*` endpoints remain available. + +### Authorization operations + +`/authorizeKms` and `/authorizeApp` require an `operation` string. See the +[KMS authorization protocol](../README.md#authorization-protocol) for the names +and subject semantics. Bootstrap alone may authorize before root registration; +unknown operations are rejected. Legacy `/bootAuth/*` routes are unchanged. diff --git a/dstack/kms/auth-mock/index.test.ts b/dstack/kms/auth-mock/index.test.ts index dc2ee8445..f067e8add 100644 --- a/dstack/kms/auth-mock/index.test.ts +++ b/dstack/kms/auth-mock/index.test.ts @@ -387,7 +387,7 @@ describe('Mock Behavior Verification', () => { describe('authorization endpoints', () => { const authorize = (appId: string, isKms = false, requestedAppEpoch = 0) => appFetch(new Request(`http://localhost/${isKms ? 'authorizeKms' : 'authorizeApp'}`, { method: 'POST', headers: { 'Content-Type': 'application/json' }, - body: JSON.stringify({ appId, requestedAppEpoch, mrAggregated: '', osImageHash: '', composeHash: '', instanceId: '', deviceId: '' }), + body: JSON.stringify({ appId, operation: isKms ? 'serve' : 'get_app_key', requestedAppEpoch, mrAggregated: '', osImageHash: '', composeHash: '', instanceId: '', deviceId: '' }), })); it('requires registered roots and selects policy by app ID', async () => { diff --git a/dstack/kms/auth-mock/index.ts b/dstack/kms/auth-mock/index.ts index 1ea6e326c..fe8805c50 100644 --- a/dstack/kms/auth-mock/index.ts +++ b/dstack/kms/auth-mock/index.ts @@ -30,6 +30,7 @@ const BootInfoSchema = z.object({ }); const AuthorizationRequestSchema = BootInfoSchema.omit({ rootEpoch: true, rootEpochHash: true }).extend({ + operation: z.enum(['bootstrap', 'serve', 'export_kms_keys', 'import_kms_keys', 'get_app_key', 'sign_cert']), requestedAppEpoch: z.number().int().min(0).max(0xffffffff).default(0), requestedRootEpoch: z.number().int().min(0).max(0xffffffff).optional(), }); @@ -119,10 +120,14 @@ class MockBackend { currentRootEpoch: current.rootEpoch ?? 0, currentRootEpochHash: current.rootEpochHash ?? '0x' + '00'.repeat(32), }; - if (current.rootEpoch === undefined) { + const kmsOperation = ['bootstrap', 'serve', 'export_kms_keys', 'import_kms_keys'].includes(appInfo.operation); + if (kmsOperation !== isKms) { + return { isAllowed: false, rejectionReason: 'Invalid authorization operation', gatewayAppId: this.mockGatewayAppId, ...epoch }; + } + if (current.rootEpoch === undefined && appInfo.operation !== 'bootstrap') { return { isAllowed: false, rejectionReason: 'Root epoch not registered', gatewayAppId: this.mockGatewayAppId, ...epoch }; } - if (appInfo.requestedRootEpoch !== undefined && appInfo.requestedRootEpoch > current.rootEpoch) { + 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) { diff --git a/dstack/kms/auth-mock/openapi.json b/dstack/kms/auth-mock/openapi.json index c4d1bc05a..307db98dc 100644 --- a/dstack/kms/auth-mock/openapi.json +++ b/dstack/kms/auth-mock/openapi.json @@ -437,8 +437,22 @@ "maximum": 4294967295, "default": 0, "description": "Application epoch, independent of KMS root epochs. Only zero is currently supported." + }, + "operation": { + "type": "string", + "enum": [ + "bootstrap", + "serve", + "export_kms_keys", + "import_kms_keys", + "get_app_key", + "sign_cert" + ] } - } + }, + "required": [ + "operation" + ] } } } diff --git a/dstack/kms/auth-simple/README.md b/dstack/kms/auth-simple/README.md index 494b3fcce..bfa0951b1 100644 --- a/dstack/kms/auth-simple/README.md +++ b/dstack/kms/auth-simple/README.md @@ -266,3 +266,10 @@ epoch, and an explicit `0` selects epoch 0. See the [KMS authorization protocol](../README.md#authorization-protocol) for registration, bootstrap, and compatibility requirements. The legacy `/bootAuth/*` endpoints remain available. + +### Authorization operations + +`/authorizeKms` and `/authorizeApp` require an `operation` string. See the +[KMS authorization protocol](../README.md#authorization-protocol) for the names +and subject semantics. Bootstrap alone may authorize before root registration; +unknown operations are rejected. Legacy `/bootAuth/*` routes are unchanged. diff --git a/dstack/kms/auth-simple/index.test.ts b/dstack/kms/auth-simple/index.test.ts index b372091f5..80a075b6f 100644 --- a/dstack/kms/auth-simple/index.test.ts +++ b/dstack/kms/auth-simple/index.test.ts @@ -475,7 +475,7 @@ describe('auth-simple', () => { const authorize = async (fields: object = {}, isKms = false) => (await app.fetch(new Request(`http://localhost/${isKms ? 'authorizeKms' : 'authorizeApp'}`, { method: 'POST', headers: { 'Content-Type': 'application/json' }, - body: JSON.stringify({ ...baseBootInfo, ...fields }), + body: JSON.stringify({ ...baseBootInfo, operation: isKms ? 'serve' : 'get_app_key', ...fields }), }))).json(); it('checks the requested policy and returns the epoch on rejection', async () => { @@ -505,6 +505,15 @@ describe('auth-simple', () => { currentRootEpoch: 0, currentRootEpochHash: '0x' + '00'.repeat(32), }); expect((await boot('kms', {})).isAllowed).toBe(true); + expect(await authorize({ operation: 'bootstrap' }, true)).toMatchObject({ isAllowed: true }); + expect(await authorize({ operation: 'export_kms_keys' }, true)).toMatchObject({ isAllowed: false }); + }); + + it('rejects operations belonging to the other authorization scope', async () => { + writeTestConfig(config); + expect(await authorize({ operation: 'serve' })).toMatchObject({ isAllowed: false, rejectionReason: 'Invalid authorization operation' }); + expect(await authorize({ operation: 'sign_cert' }, true)).toMatchObject({ isAllowed: false, rejectionReason: 'Invalid authorization operation' }); + expect(await authorize({ operation: 'sign_cert' })).toMatchObject({ isAllowed: true }); }); it('allows only a KMS serving the current root epoch', async () => { diff --git a/dstack/kms/auth-simple/index.ts b/dstack/kms/auth-simple/index.ts index 05d75b1d5..eda249642 100644 --- a/dstack/kms/auth-simple/index.ts +++ b/dstack/kms/auth-simple/index.ts @@ -28,6 +28,7 @@ const BootInfoSchema = z.object({ }); const AuthorizationRequestSchema = BootInfoSchema.omit({ rootEpoch: true, rootEpochHash: true }).extend({ + operation: z.enum(['bootstrap', 'serve', 'export_kms_keys', 'import_kms_keys', 'get_app_key', 'sign_cert']), requestedAppEpoch: z.number().int().min(0).max(0xffffffff).default(0), requestedRootEpoch: z.number().int().min(0).max(0xffffffff).optional(), }); @@ -166,10 +167,14 @@ class ConfigBackend { currentRootEpoch: current.rootEpoch ?? 0, currentRootEpochHash: current.rootEpochHash ?? '0x' + '00'.repeat(32), }; - if (current.rootEpoch === undefined) { + const kmsOperation = ['bootstrap', 'serve', 'export_kms_keys', 'import_kms_keys'].includes(appInfo.operation); + if (kmsOperation !== isKms) { + return { isAllowed: false, rejectionReason: 'Invalid authorization operation', gatewayAppId: config.gatewayAppId, ...epoch }; + } + if (current.rootEpoch === undefined && appInfo.operation !== 'bootstrap') { return { isAllowed: false, rejectionReason: 'Root epoch not registered', gatewayAppId: config.gatewayAppId, ...epoch }; } - if (appInfo.requestedRootEpoch !== undefined && appInfo.requestedRootEpoch > current.rootEpoch) { + 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) { diff --git a/dstack/kms/src/main_service.rs b/dstack/kms/src/main_service.rs index 25f3c0c23..86e542ba4 100644 --- a/dstack/kms/src/main_service.rs +++ b/dstack/kms/src/main_service.rs @@ -32,7 +32,7 @@ use scale::Decode; use tracing::warn; use upgrade_authority::{ build_boot_info, ensure_app_id_len, local_kms_boot_info, AuthorizationResult, BootInfo, - GetInfoResponse, RootEpochRegistry, + GetInfoResponse, Operation, RootEpochRegistry, }; use crate::{ @@ -171,7 +171,7 @@ pub struct KmsStateInner { verifier: CvmVerifier, attestation_verifier: Arc, self_boot_info: Cache<(), BootInfo>, - self_allowed: Cache<(), AuthorizationResult>, + self_allowed: Cache, auth_api_info: Cache<(), GetInfoResponse>, metrics: KmsMetrics, } @@ -207,9 +207,13 @@ impl KmsMetrics { } } -fn short_lived_cache() -> Cache<(), V> { +fn short_lived_cache(capacity: u64) -> Cache +where + K: Eq + std::hash::Hash + Send + Sync + 'static, + V: Clone + Send + Sync + 'static, +{ Cache::builder() - .max_capacity(1) + .max_capacity(capacity) .time_to_live(AUTH_API_CACHE_TTL) .build() } @@ -259,8 +263,8 @@ impl KmsState { .max_capacity(1) .time_to_live(SELF_BOOT_INFO_TTL) .build(), - self_allowed: short_lived_cache(), - auth_api_info: short_lived_cache(), + self_allowed: short_lived_cache(4), + auth_api_info: short_lived_cache(1), metrics: KmsMetrics::default(), }), }) @@ -403,7 +407,7 @@ impl RpcHandler { } pub(crate) async fn handover_keys(self, request: GetKmsKeyRequest) -> Result { - self.ensure_self_allowed() + self.ensure_self_allowed(Operation::ExportKmsKeys) .await .context("KMS self authorization failed")?; let info = self.ensure_kms_allowed(&request.vm_config).await?; @@ -484,7 +488,7 @@ fn ensure_key_release_allowed(boot_info: &BootInfo, policy: KeyReleasePolicy) -> } impl RpcHandler { - async fn ensure_self_allowed(&self) -> Result { + async fn ensure_self_allowed(&self, operation: Operation) -> Result { if !self.state.config.enforce_self_authorization { let epoch = self.state.current_epoch_from_info().await?; self.state.ensure_rpc_ca(epoch)?; @@ -500,12 +504,12 @@ impl RpcHandler { let response = self .state .self_allowed - .try_get_with((), async { + .try_get_with(operation, async { let response = self .state .config .auth_api - .check_kms_authorization(&boot_info) + .check_kms_authorization(&boot_info, operation, None) .await .context("Failed to call self KMS auth check")?; let epoch = self.state.authorized_epoch(&response, None)?; @@ -532,7 +536,7 @@ impl RpcHandler { .state .config .auth_api - .check_kms_authorization(&boot_info) + .check_kms_authorization(&boot_info, Operation::ImportKmsKeys, None) .await?; self.state.authorized_epoch(&response, None)?; self.ensure_attestation_policy(att, vm_config, &boot_info, &response) @@ -547,8 +551,14 @@ impl RpcHandler { ) -> Result { let att = self.ensure_attested()?; let boot_info = build_boot_info_for_attestation(att, false, vm_config)?; - self.ensure_app_attestation_allowed(att, boot_info, vm_config, requested_root_epoch) - .await + self.ensure_app_attestation_allowed( + att, + boot_info, + vm_config, + Operation::GetAppKey, + requested_root_epoch, + ) + .await } async fn ensure_app_attestation_allowed( @@ -556,13 +566,14 @@ impl RpcHandler { att: &VerifiedAttestation, boot_info: BootInfo, vm_config: &str, + operation: Operation, requested_root_epoch: Option, ) -> Result { let response = self .state .config .auth_api - .check_app_authorization(&boot_info, requested_root_epoch) + .check_app_authorization(&boot_info, operation, requested_root_epoch) .await?; let root_epoch = self .state @@ -624,7 +635,7 @@ impl KmsRpc for RpcHandler { if request.api_version > 1 { bail!("Unsupported API version: {}", request.api_version); } - self.ensure_self_allowed() + self.ensure_self_allowed(Operation::Serve) .await .context("KMS self authorization failed")?; let BootConfig { @@ -654,7 +665,7 @@ impl KmsRpc for RpcHandler { async fn get_app_env_encrypt_pub_key(self, request: AppId) -> Result { let epoch = self - .ensure_self_allowed() + .ensure_self_allowed(Operation::Serve) .await .context("KMS self authorization failed")?; let root = self.state.epoch_root(epoch)?; @@ -734,7 +745,7 @@ impl KmsRpc for RpcHandler { } async fn get_temp_ca_cert(self) -> Result { - let epoch = self.ensure_self_allowed().await?; + let epoch = self.ensure_self_allowed(Operation::Serve).await?; if self.state.config.enforce_self_authorization { let boot_info = self .state @@ -752,7 +763,7 @@ impl KmsRpc for RpcHandler { } async fn sign_cert(self, request: SignCertRequest) -> Result { - self.ensure_self_allowed() + self.ensure_self_allowed(Operation::Serve) .await .context("KMS self authorization failed")?; let csr = match request.api_version { @@ -781,7 +792,13 @@ impl KmsRpc for RpcHandler { .context("Quote verification failed")?; let boot_info = build_boot_info_for_attestation(&attestation, true, &request.vm_config)?; let app_info = self - .ensure_app_attestation_allowed(&attestation, boot_info, &request.vm_config, None) + .ensure_app_attestation_allowed( + &attestation, + boot_info, + &request.vm_config, + Operation::SignCert, + None, + ) .await?; ensure_key_release_allowed(&app_info.boot_info, (&self.state.config).into())?; let root = self.state.epoch_root(app_info.root_epoch)?; @@ -1584,6 +1601,52 @@ mod tests { const AUTH_API_INFO: &str = r#"{"status":"ok","kmsContractAddr":"0xkms","ethRpcUrl":"https://rpc.example","gatewayAppId":"0xgateway","chainId":1,"appImplementation":"0ximpl"}"#; + #[rocket::async_test] + async fn self_authorization_cache_cannot_authorize_a_different_operation() { + use upgrade_authority::tests::{boot_info, serve}; + let dir = tempfile::tempdir().unwrap(); + let initial = kms_state(dir.path(), "http://127.0.0.1:1"); + let response = |allowed| { + serde_json::json!({ + "isAllowed": allowed, "rejectionReason": "export denied", "gatewayAppId": "", + "currentRootEpoch": 0, "currentAppEpoch": 0, + "currentRootEpochHash": hex::encode(initial.keyring.chain_hashes[0]), + }) + .to_string() + }; + let (url, server) = serve(vec![response(true), response(false)]); + let mut config = initial.config.clone(); + config.enforce_self_authorization = true; + config.auth_api = upgrade_authority::tests::webhook(url); + let keyring = Keyring::load(&config).unwrap(); + let state = KmsState::new( + config, + keyring, + initial.epoch_root(0).unwrap().ca.pem_cert.clone(), + ) + .unwrap(); + state.self_boot_info.insert((), boot_info(1)).await; + let handler = RpcHandler { + state, + attestation: None, + }; + assert_eq!( + handler.ensure_self_allowed(Operation::Serve).await.unwrap(), + 0 + ); + assert!(handler + .ensure_self_allowed(Operation::ExportKmsKeys) + .await + .is_err()); + assert_eq!( + handler.ensure_self_allowed(Operation::Serve).await.unwrap(), + 0 + ); + let requests = server.join().unwrap(); + assert_eq!(requests[0].1["operation"], "serve"); + assert_eq!(requests[1].1["operation"], "export_kms_keys"); + } + #[rocket::async_test] async fn legacy_temp_ca_rpc_returns_a_separate_ca_and_the_served_root() { assert!(dstack_kms_rpc::kms_server::KMSMethod::from_str("GetTempCaCert").is_some()); diff --git a/dstack/kms/src/main_service/upgrade_authority.rs b/dstack/kms/src/main_service/upgrade_authority.rs index 08acbc75f..4b6c48391 100644 --- a/dstack/kms/src/main_service/upgrade_authority.rs +++ b/dstack/kms/src/main_service/upgrade_authority.rs @@ -8,6 +8,7 @@ use anyhow::{bail, Context, Result}; use dstack_guest_agent_rpc::v0::{ dstack_guest_client::DstackGuestClient, AttestResponse, RawQuoteArgs, }; +pub(crate) use dstack_types::authorization::AuthorizationOperation as Operation; use http_client::prpc::PrpcClient; use ra_tls::attestation::{AttestationVerifier, VerifiedAttestation, VersionedAttestation}; use serde::de::DeserializeOwned; @@ -78,6 +79,7 @@ pub(crate) async fn local_kms_boot_info(verifier: &AttestationVerifier) -> Resul build_boot_info_for_attestation(&verified, false, "") } +#[cfg(test)] #[derive(Debug, Serialize, Deserialize)] #[serde(rename_all = "camelCase")] pub(crate) struct BootResponse { @@ -258,17 +260,26 @@ impl AuthApi { pub async fn check_kms_authorization( &self, app_info: &BootInfo, + operation: Operation, + requested_root_epoch: Option, ) -> Result { - self.check_authorization(app_info, "authorizeKms", None) + if !operation.is_kms() { + bail!("invalid KMS authorization operation"); + } + self.check_authorization(app_info, "authorizeKms", operation, requested_root_epoch) .await } pub async fn check_app_authorization( &self, app_info: &BootInfo, + operation: Operation, requested_root_epoch: Option, ) -> Result { - self.check_authorization(app_info, "authorizeApp", requested_root_epoch) + if operation.is_kms() { + bail!("invalid app authorization operation"); + } + self.check_authorization(app_info, "authorizeApp", operation, requested_root_epoch) .await } @@ -276,6 +287,7 @@ impl AuthApi { &self, app_info: &BootInfo, method: &str, + operation: Operation, requested_root_epoch: Option, ) -> Result { #[derive(Serialize)] @@ -283,6 +295,7 @@ impl AuthApi { struct AuthorizationRequest<'a> { #[serde(flatten)] app_info: &'a BootInfo, + operation: Operation, requested_app_epoch: u32, #[serde(skip_serializing_if = "Option::is_none")] requested_root_epoch: Option, @@ -303,6 +316,7 @@ impl AuthApi { &url, &AuthorizationRequest { app_info, + operation, requested_app_epoch: 0, requested_root_epoch, }, @@ -313,8 +327,7 @@ impl AuthApi { } } - /// Legacy scoped checks remain available during bootstrap and onboarding, - /// before the first root epoch can be registered. + #[cfg(test)] pub async fn is_app_allowed( &self, boot_info: &BootInfo, @@ -422,11 +435,11 @@ pub(crate) async fn ensure_self_kms_allowed( .context("failed to build local KMS boot info")?; let response = cfg .auth_api - .is_app_allowed(&boot_info, true, None) + .check_kms_authorization(&boot_info, Operation::Bootstrap, None) .await - .context("failed to call KMS auth check")?; + .context("failed to call KMS bootstrap authorization")?; if !response.is_allowed { - bail!("boot denied: {}", response.reason); + bail!("boot denied: {}", response.rejection_reason); } Ok(()) } @@ -438,7 +451,8 @@ pub(crate) async fn ensure_kms_allowed( attestation: &VerifiedAttestation, verifier: &AttestationVerifier, allow_legacy_cert: bool, -) -> Result<()> { + root_epoch: u32, +) -> Result { let mut boot_info = build_boot_info_for_attestation(attestation, false, "") .context("failed to build KMS boot info from attestation")?; // Workaround: old source KMS instances use the legacy cert format (separate TDX_QUOTE + @@ -458,13 +472,19 @@ pub(crate) async fn ensure_kms_allowed( } let response = cfg .auth_api - .is_app_allowed(&boot_info, true, None) + .check_kms_authorization(&boot_info, Operation::ExportKmsKeys, Some(root_epoch)) .await - .context("failed to call KMS auth check")?; + .context("failed to authorize source KMS key export")?; if !response.is_allowed { - bail!("boot denied: {}", response.reason); + bail!("boot denied: {}", response.rejection_reason); } - Ok(()) + if cfg.auth_api.is_dev() { + return Ok(RootEpochRegistry::Unchecked); + } + let current = response + .current_root()? + .context("root epoch not registered")?; + Ok(RootEpochRegistry::Current(current)) } #[cfg(test)] @@ -695,6 +715,47 @@ pub(crate) mod tests { assert_ne!(requests[0].1["mrAggregated"], requests[1].1["mrAggregated"]); } + #[rocket::async_test] + async fn authorization_operations_are_typed_and_sent_without_legacy_fallback() { + let operations = [ + Operation::Bootstrap, + Operation::Serve, + Operation::ExportKmsKeys, + Operation::ImportKmsKeys, + Operation::GetAppKey, + Operation::SignCert, + ]; + let response = r#"{"isAllowed":true,"rejectionReason":"","gatewayAppId":"","currentRootEpoch":0,"currentAppEpoch":0,"currentRootEpochHash":"0000000000000000000000000000000000000000000000000000000000000000"}"#; + let (url, server) = serve(vec![response; operations.len()]); + let auth = webhook(url); + for operation in operations { + if operation.is_kms() { + auth.check_kms_authorization(&boot_info(1), operation, Some(0)) + .await + .unwrap(); + } else { + auth.check_app_authorization(&boot_info(1), operation, None) + .await + .unwrap(); + } + } + for ((route, body), operation) in server.join().unwrap().into_iter().zip(operations) { + assert_eq!( + route, + if operation.is_kms() { + "/authorizeKms" + } else { + "/authorizeApp" + } + ); + assert_eq!( + serde_json::from_value::(body["operation"].clone()).unwrap(), + operation + ); + } + assert!(serde_json::from_str::(r#""unknown""#).is_err()); + } + #[rocket::async_test] async fn authorization_returns_the_epoch_with_each_decision() { let response = |allowed: bool, epoch: u32| { @@ -711,14 +772,24 @@ pub(crate) mod tests { ]); let auth = webhook(url); let info = boot_info(50); - let first = auth.check_app_authorization(&info, None).await.unwrap(); - let second = auth.check_app_authorization(&info, None).await.unwrap(); + let first = auth + .check_app_authorization(&info, Operation::GetAppKey, None) + .await + .unwrap(); + let second = auth + .check_app_authorization(&info, Operation::GetAppKey, 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.check_app_authorization(&info, Some(0)).await.unwrap(); - auth.check_kms_authorization(&info).await.unwrap(); + auth.check_app_authorization(&info, Operation::GetAppKey, Some(0)) + .await + .unwrap(); + auth.check_kms_authorization(&info, Operation::Serve, None) + .await + .unwrap(); let requests = server.join().unwrap(); assert_eq!(requests.len(), 4); assert!(requests[..3] @@ -727,6 +798,7 @@ pub(crate) mod tests { assert_eq!(requests[3].0, "/authorizeKms"); let mut expected = serde_json::to_value(info).unwrap(); expected["requestedAppEpoch"] = 0.into(); + expected["operation"] = "get_app_key".into(); assert_eq!(requests[0].1, expected); assert_eq!(requests[2].1["requestedAppEpoch"], 0); assert!(requests[0].1.get("requestedRootEpoch").is_none()); @@ -739,7 +811,7 @@ pub(crate) mod tests { r#"{"isAllowed":true,"reason":"","gatewayAppId":"gw"}"#, ]); assert!(webhook(url) - .check_app_authorization(&boot_info(50), None) + .check_app_authorization(&boot_info(50), Operation::GetAppKey, None) .await .is_err()); assert_eq!(server.join().unwrap().len(), 1); diff --git a/dstack/kms/src/onboard_service.rs b/dstack/kms/src/onboard_service.rs index 31be460a0..e7b944433 100644 --- a/dstack/kms/src/onboard_service.rs +++ b/dstack/kms/src/onboard_service.rs @@ -982,26 +982,21 @@ impl Keys { }; // The attestation is captured only on the connection that returned the keys; // nothing received is used before the source passes this check. - ensure_kms_allowed(cfg, &source_attestation()?, &attestation_verifier, !rotate) - .await - .context("Source KMS is not allowed for onboarding")?; - Self::from_registered_key_response(cfg, keys_res, domain, rotate).await - } - - /// [`Self::from_key_response`] after reading the registered epochs, so a - /// source holding any keyring other than the registered one is refused. - async fn from_registered_key_response( - cfg: &KmsConfig, - keys_res: KmsKeyResponse, - domain: &str, - rotate: bool, - ) -> Result { - let registry = cfg - .auth_api - .get_info() - .await - .context("failed to read the current root epoch")? - .root_epochs; + let highest = keys_res + .keys + .len() + .checked_sub(1) + .context("source KMS handed over no keys")?; + let highest = u32::try_from(highest).context("too many root epochs")?; + let registry = ensure_kms_allowed( + cfg, + &source_attestation()?, + &attestation_verifier, + !rotate, + highest, + ) + .await + .context("source KMS is not allowed for onboarding")?; Self::from_key_response(cfg, keys_res, domain, rotate, registry).await } From 35772d8e9cabf3f7063dbde08c23187d5771f0a7 Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Sat, 7 Nov 2026 19:30:07 -0800 Subject: [PATCH 034/166] refactor(auth): name app certificate retrieval explicitly Signed-off-by: Kevin Wang --- dstack/dstack-types/src/authorization.rs | 2 +- dstack/kms/README.md | 2 +- dstack/kms/auth-eth-bun/index.ts | 2 +- dstack/kms/auth-eth-bun/openapi.json | 2 +- dstack/kms/auth-eth/contracts/AuthorizationOperation.sol | 4 ++-- dstack/kms/auth-eth/src/server.ts | 2 +- dstack/kms/auth-eth/src/types.ts | 2 +- dstack/kms/auth-eth/test/DstackKms.t.sol | 2 +- dstack/kms/auth-mock/index.ts | 2 +- dstack/kms/auth-mock/openapi.json | 2 +- dstack/kms/auth-simple/index.test.ts | 4 ++-- dstack/kms/auth-simple/index.ts | 2 +- dstack/kms/src/main_service.rs | 2 +- dstack/kms/src/main_service/upgrade_authority.rs | 2 +- 14 files changed, 16 insertions(+), 16 deletions(-) diff --git a/dstack/dstack-types/src/authorization.rs b/dstack/dstack-types/src/authorization.rs index b3f5ecec9..caea4e21a 100644 --- a/dstack/dstack-types/src/authorization.rs +++ b/dstack/dstack-types/src/authorization.rs @@ -14,7 +14,7 @@ pub enum AuthorizationOperation { ExportKmsKeys, ImportKmsKeys, GetAppKey, - SignCert, + GetAppCert, } impl AuthorizationOperation { diff --git a/dstack/kms/README.md b/dstack/kms/README.md index 0fd0b0803..208d4d68f 100644 --- a/dstack/kms/README.md +++ b/dstack/kms/README.md @@ -286,7 +286,7 @@ material policy, leaving room for independent policies later: | `authorizeKms` | `export_kms_keys` | Source KMS handing over its keyring | | `authorizeKms` | `import_kms_keys` | Recipient KMS receiving a keyring | | `authorizeApp` | `get_app_key` | App receiving derived keys | -| `authorizeApp` | `sign_cert` | App requesting a certificate | +| `authorizeApp` | `get_app_cert` | App requesting a certificate | The KMS chooses the operation from its code path, not an untrusted RPC field. Self-authorization caches are keyed by operation. Legacy authorization endpoints diff --git a/dstack/kms/auth-eth-bun/index.ts b/dstack/kms/auth-eth-bun/index.ts index 8361d5a90..03863a8a6 100644 --- a/dstack/kms/auth-eth-bun/index.ts +++ b/dstack/kms/auth-eth-bun/index.ts @@ -37,7 +37,7 @@ const BootInfoSchema = z.object({ }); const AuthorizationRequestSchema = BootInfoSchema.omit({ rootEpoch: true, rootEpochHash: true }).extend({ - operation: z.enum(['bootstrap', 'serve', 'export_kms_keys', 'import_kms_keys', 'get_app_key', 'sign_cert']), + operation: z.enum(['bootstrap', 'serve', 'export_kms_keys', 'import_kms_keys', 'get_app_key', 'get_app_cert']), requestedAppEpoch: z.number().int().min(0).max(0xffffffff).default(0), requestedRootEpoch: z.number().int().min(0).max(0xffffffff).optional(), }); diff --git a/dstack/kms/auth-eth-bun/openapi.json b/dstack/kms/auth-eth-bun/openapi.json index cab64187a..75b31079b 100644 --- a/dstack/kms/auth-eth-bun/openapi.json +++ b/dstack/kms/auth-eth-bun/openapi.json @@ -437,7 +437,7 @@ "export_kms_keys", "import_kms_keys", "get_app_key", - "sign_cert" + "get_app_cert" ] } }, diff --git a/dstack/kms/auth-eth/contracts/AuthorizationOperation.sol b/dstack/kms/auth-eth/contracts/AuthorizationOperation.sol index 2f60eed4a..7a7ef0551 100644 --- a/dstack/kms/auth-eth/contracts/AuthorizationOperation.sol +++ b/dstack/kms/auth-eth/contracts/AuthorizationOperation.sol @@ -11,7 +11,7 @@ library AuthorizationOperation { bytes32 internal constant EXPORT_KMS_KEYS = keccak256("export_kms_keys"); bytes32 internal constant IMPORT_KMS_KEYS = keccak256("import_kms_keys"); bytes32 internal constant GET_APP_KEY = keccak256("get_app_key"); - bytes32 internal constant SIGN_CERT = keccak256("sign_cert"); + bytes32 internal constant GET_APP_CERT = keccak256("get_app_cert"); function isKms(bytes32 operation) internal pure returns (bool) { return @@ -19,6 +19,6 @@ library AuthorizationOperation { } function isApp(bytes32 operation) internal pure returns (bool) { - return operation == GET_APP_KEY || operation == SIGN_CERT; + return operation == GET_APP_KEY || operation == GET_APP_CERT; } } diff --git a/dstack/kms/auth-eth/src/server.ts b/dstack/kms/auth-eth/src/server.ts index afbc49aa8..c3ac99d10 100644 --- a/dstack/kms/auth-eth/src/server.ts +++ b/dstack/kms/auth-eth/src/server.ts @@ -56,7 +56,7 @@ export async function build(): Promise { required: [...appInfoSchema.required, 'operation'], properties: { ...appInfoSchema.properties, - operation: { type: 'string', enum: ['bootstrap', 'serve', 'export_kms_keys', 'import_kms_keys', 'get_app_key', 'sign_cert'] }, + operation: { type: 'string', enum: ['bootstrap', 'serve', 'export_kms_keys', 'import_kms_keys', 'get_app_key', 'get_app_cert'] }, requestedAppEpoch: { type: 'integer', minimum: 0, maximum: 4294967295, default: 0 }, requestedRootEpoch: { type: 'integer', minimum: 0, maximum: 4294967295 }, }, diff --git a/dstack/kms/auth-eth/src/types.ts b/dstack/kms/auth-eth/src/types.ts index 79db187cc..307ebfbb8 100644 --- a/dstack/kms/auth-eth/src/types.ts +++ b/dstack/kms/auth-eth/src/types.ts @@ -43,7 +43,7 @@ export interface AuthorizationResult { } export interface AuthorizationRequest extends Omit { - operation: 'bootstrap' | 'serve' | 'export_kms_keys' | 'import_kms_keys' | 'get_app_key' | 'sign_cert'; + operation: 'bootstrap' | 'serve' | 'export_kms_keys' | 'import_kms_keys' | 'get_app_key' | 'get_app_cert'; requestedAppEpoch?: number; requestedRootEpoch?: number; } diff --git a/dstack/kms/auth-eth/test/DstackKms.t.sol b/dstack/kms/auth-eth/test/DstackKms.t.sol index 3f6b9809c..a489b14f4 100644 --- a/dstack/kms/auth-eth/test/DstackKms.t.sol +++ b/dstack/kms/auth-eth/test/DstackKms.t.sol @@ -493,7 +493,7 @@ contract DstackKmsTest is Test { request.operation = keccak256("import_kms_keys"); assertTrue(kms.authorizeKms(request).isAllowed); assertEq(kms.authorizeApp(request).rejectionReason, "Invalid authorization operation"); - request.operation = keccak256("sign_cert"); + request.operation = keccak256("get_app_cert"); assertEq(kms.authorizeKms(request).rejectionReason, "Invalid authorization operation"); request.operation = bytes32(0); assertEq(kms.authorizeKms(request).rejectionReason, "Invalid authorization operation"); diff --git a/dstack/kms/auth-mock/index.ts b/dstack/kms/auth-mock/index.ts index fe8805c50..827b5c703 100644 --- a/dstack/kms/auth-mock/index.ts +++ b/dstack/kms/auth-mock/index.ts @@ -30,7 +30,7 @@ const BootInfoSchema = z.object({ }); const AuthorizationRequestSchema = BootInfoSchema.omit({ rootEpoch: true, rootEpochHash: true }).extend({ - operation: z.enum(['bootstrap', 'serve', 'export_kms_keys', 'import_kms_keys', 'get_app_key', 'sign_cert']), + operation: z.enum(['bootstrap', 'serve', 'export_kms_keys', 'import_kms_keys', 'get_app_key', 'get_app_cert']), requestedAppEpoch: z.number().int().min(0).max(0xffffffff).default(0), requestedRootEpoch: z.number().int().min(0).max(0xffffffff).optional(), }); diff --git a/dstack/kms/auth-mock/openapi.json b/dstack/kms/auth-mock/openapi.json index 307db98dc..804c111c7 100644 --- a/dstack/kms/auth-mock/openapi.json +++ b/dstack/kms/auth-mock/openapi.json @@ -446,7 +446,7 @@ "export_kms_keys", "import_kms_keys", "get_app_key", - "sign_cert" + "get_app_cert" ] } }, diff --git a/dstack/kms/auth-simple/index.test.ts b/dstack/kms/auth-simple/index.test.ts index 80a075b6f..da516fef0 100644 --- a/dstack/kms/auth-simple/index.test.ts +++ b/dstack/kms/auth-simple/index.test.ts @@ -512,8 +512,8 @@ describe('auth-simple', () => { it('rejects operations belonging to the other authorization scope', async () => { writeTestConfig(config); expect(await authorize({ operation: 'serve' })).toMatchObject({ isAllowed: false, rejectionReason: 'Invalid authorization operation' }); - expect(await authorize({ operation: 'sign_cert' }, true)).toMatchObject({ isAllowed: false, rejectionReason: 'Invalid authorization operation' }); - expect(await authorize({ operation: 'sign_cert' })).toMatchObject({ isAllowed: true }); + expect(await authorize({ operation: 'get_app_cert' }, true)).toMatchObject({ isAllowed: false, rejectionReason: 'Invalid authorization operation' }); + expect(await authorize({ operation: 'get_app_cert' })).toMatchObject({ isAllowed: true }); }); it('allows only a KMS serving the current root epoch', async () => { diff --git a/dstack/kms/auth-simple/index.ts b/dstack/kms/auth-simple/index.ts index eda249642..1773bec12 100644 --- a/dstack/kms/auth-simple/index.ts +++ b/dstack/kms/auth-simple/index.ts @@ -28,7 +28,7 @@ const BootInfoSchema = z.object({ }); const AuthorizationRequestSchema = BootInfoSchema.omit({ rootEpoch: true, rootEpochHash: true }).extend({ - operation: z.enum(['bootstrap', 'serve', 'export_kms_keys', 'import_kms_keys', 'get_app_key', 'sign_cert']), + operation: z.enum(['bootstrap', 'serve', 'export_kms_keys', 'import_kms_keys', 'get_app_key', 'get_app_cert']), requestedAppEpoch: z.number().int().min(0).max(0xffffffff).default(0), requestedRootEpoch: z.number().int().min(0).max(0xffffffff).optional(), }); diff --git a/dstack/kms/src/main_service.rs b/dstack/kms/src/main_service.rs index 86e542ba4..7d3ded7ca 100644 --- a/dstack/kms/src/main_service.rs +++ b/dstack/kms/src/main_service.rs @@ -796,7 +796,7 @@ impl KmsRpc for RpcHandler { &attestation, boot_info, &request.vm_config, - Operation::SignCert, + Operation::GetAppCert, None, ) .await?; diff --git a/dstack/kms/src/main_service/upgrade_authority.rs b/dstack/kms/src/main_service/upgrade_authority.rs index 4b6c48391..0161765e1 100644 --- a/dstack/kms/src/main_service/upgrade_authority.rs +++ b/dstack/kms/src/main_service/upgrade_authority.rs @@ -723,7 +723,7 @@ pub(crate) mod tests { Operation::ExportKmsKeys, Operation::ImportKmsKeys, Operation::GetAppKey, - Operation::SignCert, + Operation::GetAppCert, ]; let response = r#"{"isAllowed":true,"rejectionReason":"","gatewayAppId":"","currentRootEpoch":0,"currentAppEpoch":0,"currentRootEpochHash":"0000000000000000000000000000000000000000000000000000000000000000"}"#; let (url, server) = serve(vec![response; operations.len()]); From 27bf9668e0047a6362aebffac8b1504dcb8021c5 Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Sun, 8 Nov 2026 01:07:03 -0800 Subject: [PATCH 035/166] refactor(auth): prefix operations with their subject scope Signed-off-by: Kevin Wang --- dstack/crates/dstack-auth/src/main.rs | 10 +++---- dstack/dstack-types/src/authorization.rs | 14 +++++----- dstack/kms/README.md | 14 +++++----- dstack/kms/auth-eth-bun/index.test.ts | 2 +- dstack/kms/auth-eth-bun/index.ts | 2 +- dstack/kms/auth-eth-bun/openapi.json | 12 ++++---- .../contracts/AuthorizationOperation.sol | 16 +++++------ dstack/kms/auth-eth/contracts/DstackKms.sol | 2 +- dstack/kms/auth-eth/src/main.test.ts | 2 +- dstack/kms/auth-eth/src/server.ts | 2 +- dstack/kms/auth-eth/src/types.ts | 2 +- dstack/kms/auth-eth/test/DstackKms.t.sol | 14 +++++----- dstack/kms/auth-mock/index.test.ts | 2 +- dstack/kms/auth-mock/index.ts | 6 ++-- dstack/kms/auth-mock/openapi.json | 12 ++++---- dstack/kms/auth-simple/index.test.ts | 12 ++++---- dstack/kms/auth-simple/index.ts | 6 ++-- dstack/kms/src/main_service.rs | 26 ++++++++--------- .../kms/src/main_service/upgrade_authority.rs | 28 +++++++++---------- 19 files changed, 92 insertions(+), 92 deletions(-) diff --git a/dstack/crates/dstack-auth/src/main.rs b/dstack/crates/dstack-auth/src/main.rs index 062228be7..06f3f1186 100644 --- a/dstack/crates/dstack-auth/src/main.rs +++ b/dstack/crates/dstack-auth/src/main.rs @@ -318,7 +318,7 @@ fn authorize( let info = &request.boot_info; result.rejection_reason = if request.operation.is_kms() != is_kms { "invalid authorization operation".into() - } else if current.is_none() && request.operation != Operation::Bootstrap { + } else if current.is_none() && request.operation != Operation::KmsBootstrap { "root epoch not registered".into() } else if request .requested_root_epoch @@ -327,9 +327,9 @@ fn authorize( "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.operation != Operation::Bootstrap && norm(&al.kms_contract_addr).is_empty() { + } else if request.operation != Operation::KmsBootstrap && norm(&al.kms_contract_addr).is_empty() { "kmsContractAddr is not configured".into() - } else if request.operation != Operation::Bootstrap + } else if request.operation != Operation::KmsBootstrap && (norm(&info.app_id) == norm(&al.kms_contract_addr)) != is_kms { "caller does not match authorization scope".into() @@ -478,8 +478,8 @@ mod tests { })) .await .unwrap(); - let kms = json!({"operation": "serve", "appId": "11".repeat(20), "mrAggregated": "bb"}); - let app = json!({"operation": "get_app_key", "appId": "22".repeat(20), "composeHash": "cc", "osImageHash": "aa"}); + let kms = json!({"operation": "kms_serve", "appId": "11".repeat(20), "mrAggregated": "bb"}); + let app = json!({"operation": "app_get_key", "appId": "22".repeat(20), "composeHash": "cc", "osImageHash": "aa"}); for (route, body, allowed) in [ ("/authorizeKms", kms.clone(), true), ("/authorizeApp", app.clone(), true), diff --git a/dstack/dstack-types/src/authorization.rs b/dstack/dstack-types/src/authorization.rs index caea4e21a..1fdb48228 100644 --- a/dstack/dstack-types/src/authorization.rs +++ b/dstack/dstack-types/src/authorization.rs @@ -9,19 +9,19 @@ use serde::{Deserialize, Serialize}; #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)] #[serde(rename_all = "snake_case")] pub enum AuthorizationOperation { - Bootstrap, - Serve, - ExportKmsKeys, - ImportKmsKeys, - GetAppKey, - GetAppCert, + KmsBootstrap, + KmsServe, + KmsExportKeys, + KmsImportKeys, + AppGetKey, + AppGetCert, } impl AuthorizationOperation { pub fn is_kms(self) -> bool { matches!( self, - Self::Bootstrap | Self::Serve | Self::ExportKmsKeys | Self::ImportKmsKeys + Self::KmsBootstrap | Self::KmsServe | Self::KmsExportKeys | Self::KmsImportKeys ) } } diff --git a/dstack/kms/README.md b/dstack/kms/README.md index 208d4d68f..7e5d4b710 100644 --- a/dstack/kms/README.md +++ b/dstack/kms/README.md @@ -281,12 +281,12 @@ material policy, leaving room for independent policies later: | Endpoint | Operation | Subject | | --- | --- | --- | -| `authorizeKms` | `bootstrap` | KMS generating its first roots | -| `authorizeKms` | `serve` | KMS serving app keys, public keys or certificates | -| `authorizeKms` | `export_kms_keys` | Source KMS handing over its keyring | -| `authorizeKms` | `import_kms_keys` | Recipient KMS receiving a keyring | -| `authorizeApp` | `get_app_key` | App receiving derived keys | -| `authorizeApp` | `get_app_cert` | App requesting a certificate | +| `authorizeKms` | `kms_bootstrap` | KMS generating its first roots | +| `authorizeKms` | `kms_serve` | KMS serving app keys, public keys or certificates | +| `authorizeKms` | `kms_export_keys` | Source KMS handing over its keyring | +| `authorizeKms` | `kms_import_keys` | Recipient KMS receiving a keyring | +| `authorizeApp` | `app_get_key` | App receiving derived keys | +| `authorizeApp` | `app_get_cert` | App requesting a certificate | The KMS chooses the operation from its code path, not an untrusted RPC field. Self-authorization caches are keyed by operation. Legacy authorization endpoints @@ -327,7 +327,7 @@ Upgrade the contract and auth backend and register epoch 0 before using this protocol for key release. Runtime authorization never falls back to the legacy API if the new endpoint or contract method is unavailable. The old contract methods and `/bootAuth/app` and `/bootAuth/kms` endpoints remain available for -older clients. Bootstrap uses the `bootstrap` operation, which checks KMS policy at epoch zero +older clients. Bootstrap uses the `kms_bootstrap` operation, which checks KMS policy at epoch zero even before roots are registered. Other operations require registered roots. Onboarding checks its source through `authorizeKms`, without a legacy fallback. For `auth-simple`, set `kmsContractAddr` to the attested KMS app ID and configure diff --git a/dstack/kms/auth-eth-bun/index.test.ts b/dstack/kms/auth-eth-bun/index.test.ts index 132fbf67f..dda821bba 100644 --- a/dstack/kms/auth-eth-bun/index.test.ts +++ b/dstack/kms/auth-eth-bun/index.test.ts @@ -843,7 +843,7 @@ describe('POST /authorizeApp', () => { osImageHash: '0x' + '66'.repeat(32), tcbStatus: 'UpToDate', advisoryIds: [], }; const authorize = (requestedRootEpoch?: number) => appFetch(new Request('http://localhost/authorizeApp', { - method: 'POST', headers: { 'Content-Type': 'application/json' }, body: JSON.stringify({ ...appInfo, operation: 'get_app_key', requestedRootEpoch }), + method: 'POST', headers: { 'Content-Type': 'application/json' }, body: JSON.stringify({ ...appInfo, operation: 'app_get_key', requestedRootEpoch }), })); it('reads the complete result with one contract call at configured finality', async () => { diff --git a/dstack/kms/auth-eth-bun/index.ts b/dstack/kms/auth-eth-bun/index.ts index 03863a8a6..f4e12ab3a 100644 --- a/dstack/kms/auth-eth-bun/index.ts +++ b/dstack/kms/auth-eth-bun/index.ts @@ -37,7 +37,7 @@ const BootInfoSchema = z.object({ }); const AuthorizationRequestSchema = BootInfoSchema.omit({ rootEpoch: true, rootEpochHash: true }).extend({ - operation: z.enum(['bootstrap', 'serve', 'export_kms_keys', 'import_kms_keys', 'get_app_key', 'get_app_cert']), + 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), requestedRootEpoch: z.number().int().min(0).max(0xffffffff).optional(), }); diff --git a/dstack/kms/auth-eth-bun/openapi.json b/dstack/kms/auth-eth-bun/openapi.json index 75b31079b..e431225a9 100644 --- a/dstack/kms/auth-eth-bun/openapi.json +++ b/dstack/kms/auth-eth-bun/openapi.json @@ -432,12 +432,12 @@ "operation": { "type": "string", "enum": [ - "bootstrap", - "serve", - "export_kms_keys", - "import_kms_keys", - "get_app_key", - "get_app_cert" + "kms_bootstrap", + "kms_serve", + "kms_export_keys", + "kms_import_keys", + "app_get_key", + "app_get_cert" ] } }, diff --git a/dstack/kms/auth-eth/contracts/AuthorizationOperation.sol b/dstack/kms/auth-eth/contracts/AuthorizationOperation.sol index 7a7ef0551..3d217976e 100644 --- a/dstack/kms/auth-eth/contracts/AuthorizationOperation.sol +++ b/dstack/kms/auth-eth/contracts/AuthorizationOperation.sol @@ -6,19 +6,19 @@ pragma solidity ^0.8.24; /// The operation is performed by the attested subject, not the policy caller. library AuthorizationOperation { - bytes32 internal constant BOOTSTRAP = keccak256("bootstrap"); - bytes32 internal constant SERVE = keccak256("serve"); - bytes32 internal constant EXPORT_KMS_KEYS = keccak256("export_kms_keys"); - bytes32 internal constant IMPORT_KMS_KEYS = keccak256("import_kms_keys"); - bytes32 internal constant GET_APP_KEY = keccak256("get_app_key"); - bytes32 internal constant GET_APP_CERT = keccak256("get_app_cert"); + bytes32 internal constant KMS_BOOTSTRAP = keccak256("kms_bootstrap"); + bytes32 internal constant KMS_SERVE = keccak256("kms_serve"); + bytes32 internal constant KMS_EXPORT_KEYS = keccak256("kms_export_keys"); + bytes32 internal constant KMS_IMPORT_KEYS = keccak256("kms_import_keys"); + bytes32 internal constant APP_GET_KEY = keccak256("app_get_key"); + bytes32 internal constant APP_GET_CERT = keccak256("app_get_cert"); function isKms(bytes32 operation) internal pure returns (bool) { return - operation == BOOTSTRAP || operation == SERVE || operation == EXPORT_KMS_KEYS || operation == IMPORT_KMS_KEYS; + operation == KMS_BOOTSTRAP || operation == KMS_SERVE || operation == KMS_EXPORT_KEYS || operation == KMS_IMPORT_KEYS; } function isApp(bytes32 operation) internal pure returns (bool) { - return operation == GET_APP_KEY || operation == GET_APP_CERT; + return operation == APP_GET_KEY || operation == APP_GET_CERT; } } diff --git a/dstack/kms/auth-eth/contracts/DstackKms.sol b/dstack/kms/auth-eth/contracts/DstackKms.sol index 9fa11980e..a31483e60 100644 --- a/dstack/kms/auth-eth/contracts/DstackKms.sol +++ b/dstack/kms/auth-eth/contracts/DstackKms.sol @@ -419,7 +419,7 @@ contract DstackKms is Initializable, Ownable2StepUpgradeable, UUPSUpgradeable, E return result; } if (_rootEpochs.length == 0) { - if (request.operation != AuthorizationOperation.BOOTSTRAP) { + if (request.operation != AuthorizationOperation.KMS_BOOTSTRAP) { result.rejectionReason = "Root epoch not registered"; return result; } diff --git a/dstack/kms/auth-eth/src/main.test.ts b/dstack/kms/auth-eth/src/main.test.ts index 93567235a..06884976e 100644 --- a/dstack/kms/auth-eth/src/main.test.ts +++ b/dstack/kms/auth-eth/src/main.test.ts @@ -92,7 +92,7 @@ describe('Server', () => { }); describe('POST /authorizeApp', () => { - const appRequest = { ...mockBootInfo, operation: 'get_app_key', requestedAppEpoch: 0 }; + const appRequest = { ...mockBootInfo, operation: 'app_get_key', requestedAppEpoch: 0 }; it('preserves the complete decision, including the epoch on rejection', async () => { const decision = { isAllowed: false, rejectionReason: 'App not registered', gatewayAppId: 'gateway', diff --git a/dstack/kms/auth-eth/src/server.ts b/dstack/kms/auth-eth/src/server.ts index c3ac99d10..e3d81424a 100644 --- a/dstack/kms/auth-eth/src/server.ts +++ b/dstack/kms/auth-eth/src/server.ts @@ -56,7 +56,7 @@ export async function build(): Promise { required: [...appInfoSchema.required, 'operation'], properties: { ...appInfoSchema.properties, - operation: { type: 'string', enum: ['bootstrap', 'serve', 'export_kms_keys', 'import_kms_keys', 'get_app_key', 'get_app_cert'] }, + 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 }, requestedRootEpoch: { type: 'integer', minimum: 0, maximum: 4294967295 }, }, diff --git a/dstack/kms/auth-eth/src/types.ts b/dstack/kms/auth-eth/src/types.ts index 307ebfbb8..8f7223c6a 100644 --- a/dstack/kms/auth-eth/src/types.ts +++ b/dstack/kms/auth-eth/src/types.ts @@ -43,7 +43,7 @@ export interface AuthorizationResult { } export interface AuthorizationRequest extends Omit { - operation: 'bootstrap' | 'serve' | 'export_kms_keys' | 'import_kms_keys' | 'get_app_key' | 'get_app_cert'; + operation: 'kms_bootstrap' | 'kms_serve' | 'kms_export_keys' | 'kms_import_keys' | 'app_get_key' | 'app_get_cert'; requestedAppEpoch?: number; requestedRootEpoch?: number; } diff --git a/dstack/kms/auth-eth/test/DstackKms.t.sol b/dstack/kms/auth-eth/test/DstackKms.t.sol index a489b14f4..138759983 100644 --- a/dstack/kms/auth-eth/test/DstackKms.t.sol +++ b/dstack/kms/auth-eth/test/DstackKms.t.sol @@ -466,7 +466,7 @@ contract DstackKmsTest is Test { { DstackKms.AuthorizationRequest memory request = DstackKms.AuthorizationRequest( info, - isKms ? keccak256("serve") : keccak256("get_app_key"), + isKms ? keccak256("kms_serve") : keccak256("app_get_key"), DstackKms.OptionalUint256(false, 0), DstackKms.OptionalUint256(false, 0) ); @@ -481,19 +481,19 @@ contract DstackKmsTest is Test { kms.addKmsAggregatedMr(info.mrAggregated); vm.stopPrank(); DstackKms.AuthorizationRequest memory request = DstackKms.AuthorizationRequest( - info, keccak256("bootstrap"), DstackKms.OptionalUint256(false, 0), DstackKms.OptionalUint256(false, 0) + info, keccak256("kms_bootstrap"), DstackKms.OptionalUint256(false, 0), DstackKms.OptionalUint256(false, 0) ); assertTrue(kms.authorizeKms(request).isAllowed); assertEq(kms.authorizeKms(request).currentRootEpochHash, bytes32(0)); - request.operation = keccak256("export_kms_keys"); + request.operation = keccak256("kms_export_keys"); assertEq(kms.authorizeKms(request).rejectionReason, "Root epoch not registered"); vm.prank(owner); kms.registerRootEpoch(ca(0xa0), k256(0xb0), "https://example.com/epoch-0"); assertTrue(kms.authorizeKms(request).isAllowed); - request.operation = keccak256("import_kms_keys"); + request.operation = keccak256("kms_import_keys"); assertTrue(kms.authorizeKms(request).isAllowed); assertEq(kms.authorizeApp(request).rejectionReason, "Invalid authorization operation"); - request.operation = keccak256("get_app_cert"); + request.operation = keccak256("app_get_cert"); assertEq(kms.authorizeKms(request).rejectionReason, "Invalid authorization operation"); request.operation = bytes32(0); assertEq(kms.authorizeKms(request).rejectionReason, "Invalid authorization operation"); @@ -511,7 +511,7 @@ contract DstackKmsTest is Test { vm.stopPrank(); DstackKms.AuthorizationRequest memory request = DstackKms.AuthorizationRequest( - info, keccak256("serve"), DstackKms.OptionalUint256(false, 0), DstackKms.OptionalUint256(true, 1) + info, keccak256("kms_serve"), DstackKms.OptionalUint256(false, 0), DstackKms.OptionalUint256(true, 1) ); DstackKms.AuthorizationResult memory result = kms.authorizeKms(request); assertFalse(result.isAllowed); @@ -569,7 +569,7 @@ contract DstackKmsTest is Test { assertFalse(result.isAllowed); assertEq(result.rejectionReason, "Aggregated MR not allowed"); DstackKms.AuthorizationRequest memory appRequest = DstackKms.AuthorizationRequest( - info, keccak256("get_app_key"), DstackKms.OptionalUint256(true, 1), DstackKms.OptionalUint256(false, 0) + info, keccak256("app_get_key"), DstackKms.OptionalUint256(true, 1), DstackKms.OptionalUint256(false, 0) ); result = kms.authorizeApp(appRequest); assertFalse(result.isAllowed); diff --git a/dstack/kms/auth-mock/index.test.ts b/dstack/kms/auth-mock/index.test.ts index f067e8add..d244b938f 100644 --- a/dstack/kms/auth-mock/index.test.ts +++ b/dstack/kms/auth-mock/index.test.ts @@ -387,7 +387,7 @@ describe('Mock Behavior Verification', () => { describe('authorization endpoints', () => { const authorize = (appId: string, isKms = false, requestedAppEpoch = 0) => appFetch(new Request(`http://localhost/${isKms ? 'authorizeKms' : 'authorizeApp'}`, { method: 'POST', headers: { 'Content-Type': 'application/json' }, - body: JSON.stringify({ appId, operation: isKms ? 'serve' : 'get_app_key', requestedAppEpoch, mrAggregated: '', osImageHash: '', composeHash: '', instanceId: '', deviceId: '' }), + body: JSON.stringify({ appId, operation: isKms ? 'kms_serve' : 'app_get_key', requestedAppEpoch, mrAggregated: '', osImageHash: '', composeHash: '', instanceId: '', deviceId: '' }), })); it('requires registered roots and selects policy by app ID', async () => { diff --git a/dstack/kms/auth-mock/index.ts b/dstack/kms/auth-mock/index.ts index 827b5c703..518b3e334 100644 --- a/dstack/kms/auth-mock/index.ts +++ b/dstack/kms/auth-mock/index.ts @@ -30,7 +30,7 @@ const BootInfoSchema = z.object({ }); const AuthorizationRequestSchema = BootInfoSchema.omit({ rootEpoch: true, rootEpochHash: true }).extend({ - operation: z.enum(['bootstrap', 'serve', 'export_kms_keys', 'import_kms_keys', 'get_app_key', 'get_app_cert']), + 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), requestedRootEpoch: z.number().int().min(0).max(0xffffffff).optional(), }); @@ -120,11 +120,11 @@ class MockBackend { currentRootEpoch: current.rootEpoch ?? 0, currentRootEpochHash: current.rootEpochHash ?? '0x' + '00'.repeat(32), }; - const kmsOperation = ['bootstrap', 'serve', 'export_kms_keys', 'import_kms_keys'].includes(appInfo.operation); + const kmsOperation = ['kms_bootstrap', 'kms_serve', 'kms_export_keys', 'kms_import_keys'].includes(appInfo.operation); if (kmsOperation !== isKms) { return { isAllowed: false, rejectionReason: 'Invalid authorization operation', gatewayAppId: this.mockGatewayAppId, ...epoch }; } - if (current.rootEpoch === undefined && appInfo.operation !== 'bootstrap') { + if (current.rootEpoch === undefined && appInfo.operation !== 'kms_bootstrap') { return { isAllowed: false, rejectionReason: 'Root epoch not registered', gatewayAppId: this.mockGatewayAppId, ...epoch }; } if (appInfo.requestedRootEpoch !== undefined && appInfo.requestedRootEpoch > (current.rootEpoch ?? 0)) { diff --git a/dstack/kms/auth-mock/openapi.json b/dstack/kms/auth-mock/openapi.json index 804c111c7..73b1b16c5 100644 --- a/dstack/kms/auth-mock/openapi.json +++ b/dstack/kms/auth-mock/openapi.json @@ -441,12 +441,12 @@ "operation": { "type": "string", "enum": [ - "bootstrap", - "serve", - "export_kms_keys", - "import_kms_keys", - "get_app_key", - "get_app_cert" + "kms_bootstrap", + "kms_serve", + "kms_export_keys", + "kms_import_keys", + "app_get_key", + "app_get_cert" ] } }, diff --git a/dstack/kms/auth-simple/index.test.ts b/dstack/kms/auth-simple/index.test.ts index da516fef0..2a08a5877 100644 --- a/dstack/kms/auth-simple/index.test.ts +++ b/dstack/kms/auth-simple/index.test.ts @@ -475,7 +475,7 @@ describe('auth-simple', () => { const authorize = async (fields: object = {}, isKms = false) => (await app.fetch(new Request(`http://localhost/${isKms ? 'authorizeKms' : 'authorizeApp'}`, { method: 'POST', headers: { 'Content-Type': 'application/json' }, - body: JSON.stringify({ ...baseBootInfo, operation: isKms ? 'serve' : 'get_app_key', ...fields }), + body: JSON.stringify({ ...baseBootInfo, operation: isKms ? 'kms_serve' : 'app_get_key', ...fields }), }))).json(); it('checks the requested policy and returns the epoch on rejection', async () => { @@ -505,15 +505,15 @@ describe('auth-simple', () => { currentRootEpoch: 0, currentRootEpochHash: '0x' + '00'.repeat(32), }); expect((await boot('kms', {})).isAllowed).toBe(true); - expect(await authorize({ operation: 'bootstrap' }, true)).toMatchObject({ isAllowed: true }); - expect(await authorize({ operation: 'export_kms_keys' }, true)).toMatchObject({ isAllowed: false }); + expect(await authorize({ operation: 'kms_bootstrap' }, true)).toMatchObject({ isAllowed: true }); + expect(await authorize({ operation: 'kms_export_keys' }, true)).toMatchObject({ isAllowed: false }); }); it('rejects operations belonging to the other authorization scope', async () => { writeTestConfig(config); - expect(await authorize({ operation: 'serve' })).toMatchObject({ isAllowed: false, rejectionReason: 'Invalid authorization operation' }); - expect(await authorize({ operation: 'get_app_cert' }, true)).toMatchObject({ isAllowed: false, rejectionReason: 'Invalid authorization operation' }); - expect(await authorize({ operation: 'get_app_cert' })).toMatchObject({ isAllowed: true }); + expect(await authorize({ operation: 'kms_serve' })).toMatchObject({ isAllowed: false, rejectionReason: 'Invalid authorization operation' }); + expect(await authorize({ operation: 'app_get_cert' }, true)).toMatchObject({ isAllowed: false, rejectionReason: 'Invalid authorization operation' }); + expect(await authorize({ operation: 'app_get_cert' })).toMatchObject({ isAllowed: true }); }); it('allows only a KMS serving the current root epoch', async () => { diff --git a/dstack/kms/auth-simple/index.ts b/dstack/kms/auth-simple/index.ts index 1773bec12..7444a38ab 100644 --- a/dstack/kms/auth-simple/index.ts +++ b/dstack/kms/auth-simple/index.ts @@ -28,7 +28,7 @@ const BootInfoSchema = z.object({ }); const AuthorizationRequestSchema = BootInfoSchema.omit({ rootEpoch: true, rootEpochHash: true }).extend({ - operation: z.enum(['bootstrap', 'serve', 'export_kms_keys', 'import_kms_keys', 'get_app_key', 'get_app_cert']), + 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), requestedRootEpoch: z.number().int().min(0).max(0xffffffff).optional(), }); @@ -167,11 +167,11 @@ class ConfigBackend { currentRootEpoch: current.rootEpoch ?? 0, currentRootEpochHash: current.rootEpochHash ?? '0x' + '00'.repeat(32), }; - const kmsOperation = ['bootstrap', 'serve', 'export_kms_keys', 'import_kms_keys'].includes(appInfo.operation); + const kmsOperation = ['kms_bootstrap', 'kms_serve', 'kms_export_keys', 'kms_import_keys'].includes(appInfo.operation); if (kmsOperation !== isKms) { return { isAllowed: false, rejectionReason: 'Invalid authorization operation', gatewayAppId: config.gatewayAppId, ...epoch }; } - if (current.rootEpoch === undefined && appInfo.operation !== 'bootstrap') { + if (current.rootEpoch === undefined && appInfo.operation !== 'kms_bootstrap') { return { isAllowed: false, rejectionReason: 'Root epoch not registered', gatewayAppId: config.gatewayAppId, ...epoch }; } if (appInfo.requestedRootEpoch !== undefined && appInfo.requestedRootEpoch > (current.rootEpoch ?? 0)) { diff --git a/dstack/kms/src/main_service.rs b/dstack/kms/src/main_service.rs index 7d3ded7ca..4a95b63dc 100644 --- a/dstack/kms/src/main_service.rs +++ b/dstack/kms/src/main_service.rs @@ -407,7 +407,7 @@ impl RpcHandler { } pub(crate) async fn handover_keys(self, request: GetKmsKeyRequest) -> Result { - self.ensure_self_allowed(Operation::ExportKmsKeys) + self.ensure_self_allowed(Operation::KmsExportKeys) .await .context("KMS self authorization failed")?; let info = self.ensure_kms_allowed(&request.vm_config).await?; @@ -536,7 +536,7 @@ impl RpcHandler { .state .config .auth_api - .check_kms_authorization(&boot_info, Operation::ImportKmsKeys, None) + .check_kms_authorization(&boot_info, Operation::KmsImportKeys, None) .await?; self.state.authorized_epoch(&response, None)?; self.ensure_attestation_policy(att, vm_config, &boot_info, &response) @@ -555,7 +555,7 @@ impl RpcHandler { att, boot_info, vm_config, - Operation::GetAppKey, + Operation::AppGetKey, requested_root_epoch, ) .await @@ -635,7 +635,7 @@ impl KmsRpc for RpcHandler { if request.api_version > 1 { bail!("Unsupported API version: {}", request.api_version); } - self.ensure_self_allowed(Operation::Serve) + self.ensure_self_allowed(Operation::KmsServe) .await .context("KMS self authorization failed")?; let BootConfig { @@ -665,7 +665,7 @@ impl KmsRpc for RpcHandler { async fn get_app_env_encrypt_pub_key(self, request: AppId) -> Result { let epoch = self - .ensure_self_allowed(Operation::Serve) + .ensure_self_allowed(Operation::KmsServe) .await .context("KMS self authorization failed")?; let root = self.state.epoch_root(epoch)?; @@ -745,7 +745,7 @@ impl KmsRpc for RpcHandler { } async fn get_temp_ca_cert(self) -> Result { - let epoch = self.ensure_self_allowed(Operation::Serve).await?; + let epoch = self.ensure_self_allowed(Operation::KmsServe).await?; if self.state.config.enforce_self_authorization { let boot_info = self .state @@ -763,7 +763,7 @@ impl KmsRpc for RpcHandler { } async fn sign_cert(self, request: SignCertRequest) -> Result { - self.ensure_self_allowed(Operation::Serve) + self.ensure_self_allowed(Operation::KmsServe) .await .context("KMS self authorization failed")?; let csr = match request.api_version { @@ -796,7 +796,7 @@ impl KmsRpc for RpcHandler { &attestation, boot_info, &request.vm_config, - Operation::GetAppCert, + Operation::AppGetCert, None, ) .await?; @@ -1631,20 +1631,20 @@ mod tests { attestation: None, }; assert_eq!( - handler.ensure_self_allowed(Operation::Serve).await.unwrap(), + handler.ensure_self_allowed(Operation::KmsServe).await.unwrap(), 0 ); assert!(handler - .ensure_self_allowed(Operation::ExportKmsKeys) + .ensure_self_allowed(Operation::KmsExportKeys) .await .is_err()); assert_eq!( - handler.ensure_self_allowed(Operation::Serve).await.unwrap(), + handler.ensure_self_allowed(Operation::KmsServe).await.unwrap(), 0 ); let requests = server.join().unwrap(); - assert_eq!(requests[0].1["operation"], "serve"); - assert_eq!(requests[1].1["operation"], "export_kms_keys"); + assert_eq!(requests[0].1["operation"], "kms_serve"); + assert_eq!(requests[1].1["operation"], "kms_export_keys"); } #[rocket::async_test] diff --git a/dstack/kms/src/main_service/upgrade_authority.rs b/dstack/kms/src/main_service/upgrade_authority.rs index 0161765e1..cbdf8eeac 100644 --- a/dstack/kms/src/main_service/upgrade_authority.rs +++ b/dstack/kms/src/main_service/upgrade_authority.rs @@ -435,7 +435,7 @@ pub(crate) async fn ensure_self_kms_allowed( .context("failed to build local KMS boot info")?; let response = cfg .auth_api - .check_kms_authorization(&boot_info, Operation::Bootstrap, None) + .check_kms_authorization(&boot_info, Operation::KmsBootstrap, None) .await .context("failed to call KMS bootstrap authorization")?; if !response.is_allowed { @@ -472,7 +472,7 @@ pub(crate) async fn ensure_kms_allowed( } let response = cfg .auth_api - .check_kms_authorization(&boot_info, Operation::ExportKmsKeys, Some(root_epoch)) + .check_kms_authorization(&boot_info, Operation::KmsExportKeys, Some(root_epoch)) .await .context("failed to authorize source KMS key export")?; if !response.is_allowed { @@ -718,12 +718,12 @@ pub(crate) mod tests { #[rocket::async_test] async fn authorization_operations_are_typed_and_sent_without_legacy_fallback() { let operations = [ - Operation::Bootstrap, - Operation::Serve, - Operation::ExportKmsKeys, - Operation::ImportKmsKeys, - Operation::GetAppKey, - Operation::GetAppCert, + Operation::KmsBootstrap, + Operation::KmsServe, + Operation::KmsExportKeys, + Operation::KmsImportKeys, + Operation::AppGetKey, + Operation::AppGetCert, ]; let response = r#"{"isAllowed":true,"rejectionReason":"","gatewayAppId":"","currentRootEpoch":0,"currentAppEpoch":0,"currentRootEpochHash":"0000000000000000000000000000000000000000000000000000000000000000"}"#; let (url, server) = serve(vec![response; operations.len()]); @@ -773,21 +773,21 @@ pub(crate) mod tests { let auth = webhook(url); let info = boot_info(50); let first = auth - .check_app_authorization(&info, Operation::GetAppKey, None) + .check_app_authorization(&info, Operation::AppGetKey, None) .await .unwrap(); let second = auth - .check_app_authorization(&info, Operation::GetAppKey, None) + .check_app_authorization(&info, Operation::AppGetKey, 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.check_app_authorization(&info, Operation::GetAppKey, Some(0)) + auth.check_app_authorization(&info, Operation::AppGetKey, Some(0)) .await .unwrap(); - auth.check_kms_authorization(&info, Operation::Serve, None) + auth.check_kms_authorization(&info, Operation::KmsServe, None) .await .unwrap(); let requests = server.join().unwrap(); @@ -798,7 +798,7 @@ pub(crate) mod tests { assert_eq!(requests[3].0, "/authorizeKms"); let mut expected = serde_json::to_value(info).unwrap(); expected["requestedAppEpoch"] = 0.into(); - expected["operation"] = "get_app_key".into(); + expected["operation"] = "app_get_key".into(); assert_eq!(requests[0].1, expected); assert_eq!(requests[2].1["requestedAppEpoch"], 0); assert!(requests[0].1.get("requestedRootEpoch").is_none()); @@ -811,7 +811,7 @@ pub(crate) mod tests { r#"{"isAllowed":true,"reason":"","gatewayAppId":"gw"}"#, ]); assert!(webhook(url) - .check_app_authorization(&boot_info(50), Operation::GetAppKey, None) + .check_app_authorization(&boot_info(50), Operation::AppGetKey, None) .await .is_err()); assert_eq!(server.join().unwrap().len(), 1); From 65180eac681fa35a38566f8c371661f56e418e35 Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Sun, 8 Nov 2026 01:07:03 -0800 Subject: [PATCH 036/166] refactor(kms): align Rust authorization method names Signed-off-by: Kevin Wang --- dstack/kms/src/main_service.rs | 6 ++--- .../kms/src/main_service/upgrade_authority.rs | 22 +++++++++---------- 2 files changed, 14 insertions(+), 14 deletions(-) diff --git a/dstack/kms/src/main_service.rs b/dstack/kms/src/main_service.rs index 4a95b63dc..c0bdafbcb 100644 --- a/dstack/kms/src/main_service.rs +++ b/dstack/kms/src/main_service.rs @@ -509,7 +509,7 @@ impl RpcHandler { .state .config .auth_api - .check_kms_authorization(&boot_info, operation, None) + .authorize_kms(&boot_info, operation, None) .await .context("Failed to call self KMS auth check")?; let epoch = self.state.authorized_epoch(&response, None)?; @@ -536,7 +536,7 @@ impl RpcHandler { .state .config .auth_api - .check_kms_authorization(&boot_info, Operation::KmsImportKeys, None) + .authorize_kms(&boot_info, Operation::KmsImportKeys, None) .await?; self.state.authorized_epoch(&response, None)?; self.ensure_attestation_policy(att, vm_config, &boot_info, &response) @@ -573,7 +573,7 @@ impl RpcHandler { .state .config .auth_api - .check_app_authorization(&boot_info, operation, requested_root_epoch) + .authorize_app(&boot_info, operation, requested_root_epoch) .await?; let root_epoch = self .state diff --git a/dstack/kms/src/main_service/upgrade_authority.rs b/dstack/kms/src/main_service/upgrade_authority.rs index cbdf8eeac..c265a555b 100644 --- a/dstack/kms/src/main_service/upgrade_authority.rs +++ b/dstack/kms/src/main_service/upgrade_authority.rs @@ -257,7 +257,7 @@ async fn send_request(req: reqwest::RequestBuilder, url: &s } impl AuthApi { - pub async fn check_kms_authorization( + pub async fn authorize_kms( &self, app_info: &BootInfo, operation: Operation, @@ -270,7 +270,7 @@ impl AuthApi { .await } - pub async fn check_app_authorization( + pub async fn authorize_app( &self, app_info: &BootInfo, operation: Operation, @@ -435,7 +435,7 @@ pub(crate) async fn ensure_self_kms_allowed( .context("failed to build local KMS boot info")?; let response = cfg .auth_api - .check_kms_authorization(&boot_info, Operation::KmsBootstrap, None) + .authorize_kms(&boot_info, Operation::KmsBootstrap, None) .await .context("failed to call KMS bootstrap authorization")?; if !response.is_allowed { @@ -472,7 +472,7 @@ pub(crate) async fn ensure_kms_allowed( } let response = cfg .auth_api - .check_kms_authorization(&boot_info, Operation::KmsExportKeys, Some(root_epoch)) + .authorize_kms(&boot_info, Operation::KmsExportKeys, Some(root_epoch)) .await .context("failed to authorize source KMS key export")?; if !response.is_allowed { @@ -730,11 +730,11 @@ pub(crate) mod tests { let auth = webhook(url); for operation in operations { if operation.is_kms() { - auth.check_kms_authorization(&boot_info(1), operation, Some(0)) + auth.authorize_kms(&boot_info(1), operation, Some(0)) .await .unwrap(); } else { - auth.check_app_authorization(&boot_info(1), operation, None) + auth.authorize_app(&boot_info(1), operation, None) .await .unwrap(); } @@ -773,21 +773,21 @@ pub(crate) mod tests { let auth = webhook(url); let info = boot_info(50); let first = auth - .check_app_authorization(&info, Operation::AppGetKey, None) + .authorize_app(&info, Operation::AppGetKey, None) .await .unwrap(); let second = auth - .check_app_authorization(&info, Operation::AppGetKey, None) + .authorize_app(&info, Operation::AppGetKey, 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.check_app_authorization(&info, Operation::AppGetKey, Some(0)) + auth.authorize_app(&info, Operation::AppGetKey, Some(0)) .await .unwrap(); - auth.check_kms_authorization(&info, Operation::KmsServe, None) + auth.authorize_kms(&info, Operation::KmsServe, None) .await .unwrap(); let requests = server.join().unwrap(); @@ -811,7 +811,7 @@ pub(crate) mod tests { r#"{"isAllowed":true,"reason":"","gatewayAppId":"gw"}"#, ]); assert!(webhook(url) - .check_app_authorization(&boot_info(50), Operation::AppGetKey, None) + .authorize_app(&boot_info(50), Operation::AppGetKey, None) .await .is_err()); assert_eq!(server.join().unwrap().len(), 1); From 0a8f311e2f1a38473f0d2a6a1e08ed47fb14105e Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Sun, 8 Nov 2026 06:09:23 -0800 Subject: [PATCH 037/166] style(auth): format scoped operations in the base PR Signed-off-by: Kevin Wang --- dstack/crates/dstack-auth/src/main.rs | 3 ++- dstack/kms/src/main_service.rs | 10 ++++++++-- 2 files changed, 10 insertions(+), 3 deletions(-) diff --git a/dstack/crates/dstack-auth/src/main.rs b/dstack/crates/dstack-auth/src/main.rs index 06f3f1186..a55f21b13 100644 --- a/dstack/crates/dstack-auth/src/main.rs +++ b/dstack/crates/dstack-auth/src/main.rs @@ -327,7 +327,8 @@ fn authorize( "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.operation != Operation::KmsBootstrap && norm(&al.kms_contract_addr).is_empty() { + } else if request.operation != Operation::KmsBootstrap && norm(&al.kms_contract_addr).is_empty() + { "kmsContractAddr is not configured".into() } else if request.operation != Operation::KmsBootstrap && (norm(&info.app_id) == norm(&al.kms_contract_addr)) != is_kms diff --git a/dstack/kms/src/main_service.rs b/dstack/kms/src/main_service.rs index c0bdafbcb..605092339 100644 --- a/dstack/kms/src/main_service.rs +++ b/dstack/kms/src/main_service.rs @@ -1631,7 +1631,10 @@ mod tests { attestation: None, }; assert_eq!( - handler.ensure_self_allowed(Operation::KmsServe).await.unwrap(), + handler + .ensure_self_allowed(Operation::KmsServe) + .await + .unwrap(), 0 ); assert!(handler @@ -1639,7 +1642,10 @@ mod tests { .await .is_err()); assert_eq!( - handler.ensure_self_allowed(Operation::KmsServe).await.unwrap(), + handler + .ensure_self_allowed(Operation::KmsServe) + .await + .unwrap(), 0 ); let requests = server.join().unwrap(); From 46c5ec2aa952acf1c2c465f4b1d1bbfd91f9b26a Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Sun, 8 Nov 2026 06:14:13 -0800 Subject: [PATCH 038/166] style(auth): format scoped Solidity operation predicates Signed-off-by: Kevin Wang --- dstack/kms/auth-eth/contracts/AuthorizationOperation.sol | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/dstack/kms/auth-eth/contracts/AuthorizationOperation.sol b/dstack/kms/auth-eth/contracts/AuthorizationOperation.sol index 3d217976e..d5f926616 100644 --- a/dstack/kms/auth-eth/contracts/AuthorizationOperation.sol +++ b/dstack/kms/auth-eth/contracts/AuthorizationOperation.sol @@ -14,8 +14,8 @@ library AuthorizationOperation { bytes32 internal constant APP_GET_CERT = keccak256("app_get_cert"); function isKms(bytes32 operation) internal pure returns (bool) { - return - operation == KMS_BOOTSTRAP || operation == KMS_SERVE || operation == KMS_EXPORT_KEYS || operation == KMS_IMPORT_KEYS; + return operation == KMS_BOOTSTRAP || operation == KMS_SERVE || operation == KMS_EXPORT_KEYS + || operation == KMS_IMPORT_KEYS; } function isApp(bytes32 operation) internal pure returns (bool) { From 2ca1679040b344124a6bc6aa6cff141599a5f541 Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Sat, 7 Nov 2026 07:14:53 -0800 Subject: [PATCH 039/166] refactor(dstack-util): move app key retrieval to shared setup state Signed-off-by: Kevin Wang --- dstack/dstack-util/src/system_setup.rs | 86 ++++++++++++++------------ 1 file changed, 46 insertions(+), 40 deletions(-) diff --git a/dstack/dstack-util/src/system_setup.rs b/dstack/dstack-util/src/system_setup.rs index 1beb2b6ad..5b2ade00f 100644 --- a/dstack/dstack-util/src/system_setup.rs +++ b/dstack/dstack-util/src/system_setup.rs @@ -2353,6 +2353,47 @@ struct AppInfo { gpu_attestation_event: Option>, } +impl HostShared { + async fn request_app_keys_from_kms_url(&self, kms_url: String) -> Result<(AppKeys, Vec)> { + info!("Requesting app keys from KMS: {kms_url}"); + let cert_pair = + generate_self_signed_ra_cert(None).context("Failed to generate RA-TLS cert")?; + let response = request_app_keys( + kms_url.clone(), + rpc::GetAppKeyRequest { + api_version: 1, + vm_config: self.sys_config.vm_config.clone(), + root_epoch: None, + }, + cert_pair, + attestation_verifier(&self.sys_config)?, + None, + ) + .await + .context("Failed to get app keys from the attested KMS")?; + let root_ca = response.ca_cert; + + let (_, ca_pem) = x509_parser::pem::parse_x509_pem(root_ca.as_bytes()) + .context("Failed to parse ca cert")?; + let x509 = ca_pem.parse_x509().context("Failed to parse ca cert")?; + let root_pubkey = x509.public_key().raw.to_vec(); + + let keys = AppKeys { + ca_cert: root_ca, + disk_crypt_key: response.disk_crypt_key, + env_crypt_key: response.env_crypt_key, + k256_key: response.k256_key, + k256_signature: response.k256_signature, + gateway_app_id: response.gateway_app_id, + key_provider: KeyProvider::Kms { + url: kms_url, + pubkey: root_pubkey, + }, + }; + Ok((keys, response.os_image_hash)) + } +} + struct Stage0<'a> { args: &'a SetupArgs, shared: HostShared, @@ -2415,45 +2456,6 @@ impl<'a> Stage0<'a> { self.shared.dir.join(APP_KEYS) } - async fn request_app_keys_from_kms_url(&self, kms_url: String) -> Result<(AppKeys, Vec)> { - info!("Requesting app keys from KMS: {kms_url}"); - let cert_pair = - generate_self_signed_ra_cert(None).context("Failed to generate RA-TLS cert")?; - let response = request_app_keys( - kms_url.clone(), - rpc::GetAppKeyRequest { - api_version: 1, - vm_config: self.shared.sys_config.vm_config.clone(), - root_epoch: None, - }, - cert_pair, - attestation_verifier(&self.shared.sys_config)?, - None, - ) - .await - .context("Failed to get app keys from the attested KMS")?; - let root_ca = response.ca_cert; - - let (_, ca_pem) = x509_parser::pem::parse_x509_pem(root_ca.as_bytes()) - .context("Failed to parse ca cert")?; - let x509 = ca_pem.parse_x509().context("Failed to parse ca cert")?; - let root_pubkey = x509.public_key().raw.to_vec(); - - let keys = AppKeys { - ca_cert: root_ca, - disk_crypt_key: response.disk_crypt_key, - env_crypt_key: response.env_crypt_key, - k256_key: response.k256_key, - k256_signature: response.k256_signature, - gateway_app_id: response.gateway_app_id, - key_provider: KeyProvider::Kms { - url: kms_url, - pubkey: root_pubkey, - }, - }; - Ok((keys, response.os_image_hash)) - } - async fn request_app_keys_from_kms(&self) -> Result { if self.shared.sys_config.kms_urls.is_empty() { bail!("No KMS URLs are set"); @@ -2461,7 +2463,11 @@ impl<'a> Stage0<'a> { let mut errors = vec![]; for kms_url in &self.shared.sys_config.kms_urls { let kms_url = kms_rpc_url(kms_url); - match self.request_app_keys_from_kms_url(kms_url.clone()).await { + match self + .shared + .request_app_keys_from_kms_url(kms_url.clone()) + .await + { Ok((keys, os_image_hash)) => { // Measured here, outside the per-URL request, so that a // failover can never extend os-image-hash twice. From e7bacf8fe528e643f3244e6419ddf2f9757fe3fb Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Sat, 7 Nov 2026 07:14:53 -0800 Subject: [PATCH 040/166] feat(dstack-util): re-key data disks left on an older KMS root epoch Signed-off-by: Kevin Wang --- docs/deployment.md | 2 + dstack/dstack-types/src/lib.rs | 17 +- dstack/dstack-util/src/main.rs | 1 + dstack/dstack-util/src/system_setup.rs | 352 ++++++++++++++++++-- dstack/dstack-util/src/system_setup/luks.rs | 14 +- 5 files changed, 349 insertions(+), 37 deletions(-) diff --git a/docs/deployment.md b/docs/deployment.md index 3621228ed..d68fdea44 100644 --- a/docs/deployment.md +++ b/docs/deployment.md @@ -551,6 +551,8 @@ for the full procedure; the KMS steps are: If the nodes holding N are lost before N is activated, drop it with `PopRootEpoch` (or remove the `rootEpochs` entry) and rotate again; nodes that imported the dropped N must be onboarded again. +6. Boot CVMs on an OS that supports root epochs. At boot they re-key their data + disk from the epoch recorded in its LUKS header to the current one. Each `authorizeKms` or `authorizeApp` result carries the current epoch and its chain hash, including when authorization is denied. KMS verifies these against its keyring; diff --git a/dstack/dstack-types/src/lib.rs b/dstack/dstack-types/src/lib.rs index c9b276f6e..d4a51710f 100644 --- a/dstack/dstack-types/src/lib.rs +++ b/dstack/dstack-types/src/lib.rs @@ -2409,6 +2409,17 @@ pub struct AppKeys { pub key_provider: KeyProvider, } +impl AppKeys { + /// KMS root epoch the keys were derived from; 0 for other key providers, + /// which have a single root. + pub fn root_epoch(&self) -> u32 { + match self.key_provider { + KeyProvider::Kms { root_epoch, .. } => root_epoch, + _ => 0, + } + } +} + /// Maximum length of a framed message exchanged with the local key provider. pub const KEY_PROVIDER_MAX_FRAME_SIZE: usize = 8 * 1024 * 1024; @@ -2431,6 +2442,9 @@ pub enum KeyProvider { url: String, #[serde(with = "hex_bytes")] pubkey: Vec, + /// KMS root key epoch the keys were derived from. + #[serde(default)] + root_epoch: u32, }, } @@ -2493,7 +2507,7 @@ mod key_provider_tests { assert_eq!( serde_json::to_value(provider).unwrap(), serde_json::json!({ - "Kms": { "url": "https://kms.example/prpc", "pubkey": "abcd" } + "Kms": { "url": "https://kms.example/prpc", "pubkey": "abcd", "root_epoch": 0 } }) ); } @@ -2512,6 +2526,7 @@ mod key_provider_tests { let kms = KeyProvider::Kms { url: "https://kms.example".into(), pubkey: vec![0xab; 32], + root_epoch: 0, }; assert_eq!(kms.id(), &[0xab; 32]); } diff --git a/dstack/dstack-util/src/main.rs b/dstack/dstack-util/src/main.rs index f081b3235..ab1f6837b 100644 --- a/dstack/dstack-util/src/main.rs +++ b/dstack/dstack-util/src/main.rs @@ -683,6 +683,7 @@ async fn cmd_get_keys(args: GetKeysArgs) -> Result<()> { key_provider: KeyProvider::Kms { 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 5b2ade00f..e58ae61e4 100644 --- a/dstack/dstack-util/src/system_setup.rs +++ b/dstack/dstack-util/src/system_setup.rs @@ -29,7 +29,7 @@ use dstack_types::{ use fs_err as fs; use luks::{ LuksAf, LuksArea, LuksConfig, LuksDigest, LuksHeader, LuksJson, LuksKdf, LuksKeyslot, - LuksSegment, BINARY_HEADER_SIZE, + LuksSegment, BINARY_HEADER_SIZE, ROOT_EPOCH_TOKEN_TYPE, }; use ra_rpc::{ client::{CertInfo, RaClient, RaClientConfig}, @@ -2354,8 +2354,14 @@ struct AppInfo { } impl HostShared { - async fn request_app_keys_from_kms_url(&self, kms_url: String) -> Result<(AppKeys, Vec)> { - info!("Requesting app keys from KMS: {kms_url}"); + /// Request app keys of `root_epoch` (`None` = current) from one KMS. + async fn request_app_keys_from_kms_url( + &self, + kms_url: String, + root_epoch: Option, + trusted_ca: Option, + ) -> Result<(AppKeys, Vec)> { + info!("Requesting app keys of root epoch {root_epoch:?} from KMS: {kms_url}"); let cert_pair = generate_self_signed_ra_cert(None).context("Failed to generate RA-TLS cert")?; let response = request_app_keys( @@ -2363,15 +2369,16 @@ impl HostShared { rpc::GetAppKeyRequest { api_version: 1, vm_config: self.sys_config.vm_config.clone(), - root_epoch: None, + root_epoch, }, cert_pair, attestation_verifier(&self.sys_config)?, - None, + trusted_ca.clone(), ) .await .context("Failed to get app keys from the attested KMS")?; - let root_ca = response.ca_cert; + let served = response.root_epoch; + let root_ca = trusted_ca.unwrap_or(response.ca_cert); let (_, ca_pem) = x509_parser::pem::parse_x509_pem(root_ca.as_bytes()) .context("Failed to parse ca cert")?; @@ -2388,10 +2395,23 @@ impl HostShared { key_provider: KeyProvider::Kms { url: kms_url, pubkey: root_pubkey, + root_epoch: served, }, }; Ok((keys, response.os_image_hash)) } + + /// 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 { + 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(url.clone(), Some(epoch), Some(keys.ca_cert.clone())) + .await? + .0) + } } struct Stage0<'a> { @@ -2465,7 +2485,7 @@ impl<'a> Stage0<'a> { let kms_url = kms_rpc_url(kms_url); match self .shared - .request_app_keys_from_kms_url(kms_url.clone()) + .request_app_keys_from_kms_url(kms_url.clone(), None, None) .await { Ok((keys, os_image_hash)) => { @@ -2710,7 +2730,7 @@ impl<'a> Stage0<'a> { initialized } - async fn mount_data_disk(&self, disk_crypt_key: &str, opts: &DstackOptions) -> Result<()> { + async fn mount_data_disk(&self, keys: &AppKeys, opts: &DstackOptions) -> Result<()> { let name = "dstack_data_disk"; let mount_point = &self.args.mount_point; @@ -2731,8 +2751,9 @@ impl<'a> Stage0<'a> { let has_filesystem = if opts.storage_encrypted && disk_initialized { info!("Mounting encrypted data disk"); - let initializing = - self.open_encrypted_volume(disk_crypt_key, name, opts.storage_discard)?; + let initializing = self + .open_encrypted_volume(keys, name, opts.storage_discard) + .await?; if initializing { warn!("data disk initialization was interrupted, initializing it again"); } @@ -2750,7 +2771,7 @@ impl<'a> Stage0<'a> { info!("Skipping disk encryption as requested by kernel cmdline"); } else if !disk_initialized { info!("Setting up disk encryption"); - self.luks_setup(disk_crypt_key, name, opts.storage_discard)?; + self.luks_setup(keys, name, opts.storage_discard).await?; } match opts.storage_fs { @@ -2879,7 +2900,8 @@ impl<'a> Stage0<'a> { Ok(()) } - fn luks_setup(&self, disk_crypt_key: &str, name: &str, discard: bool) -> Result<()> { + async fn luks_setup(&self, keys: &AppKeys, name: &str, discard: bool) -> Result<()> { + let disk_crypt_key = hex::encode(&keys.disk_crypt_key); let root_hd = &self.args.device; let sector_offset = PAYLOAD_OFFSET / 512; info!("Formatting encrypted disk"); @@ -2917,19 +2939,49 @@ impl<'a> Stage0<'a> { { bail!("Failed to setup luks volume"); } - self.open_encrypted_volume(disk_crypt_key, name, discard)?; + if let Some(epoch) = kms_root_epoch(keys).filter(|epoch| *epoch > 0) { + write_luks_root_epoch(root_hd, epoch)?; + } + self.open_encrypted_volume(keys, name, discard).await?; + Ok(()) + } + + /// Re-key the in-memory `header` of a data disk still protected by an + /// older KMS epoch, and write it back to the disk. + /// + /// Only the keyslot passphrase changes; the volume key stays. The epoch + /// is recorded in the header so the right old key is used after the next + /// rotation. An unrotated epoch-0 disk is left untouched, so it stays + /// readable by older OS images. + async fn migrate_disk_key( + &self, + keys: &AppKeys, + header: &Path, + from: u32, + to: 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( + header, + &hex::encode(&old.disk_crypt_key), + &hex::encode(&keys.disk_crypt_key), + to, + )?; + let root_hd = &self.args.device; + cmd!(cryptsetup luksHeaderRestore --batch-mode --header-backup-file=$header $root_hd) + .context("failed to write the re-keyed LUKS header")?; Ok(()) } - /// Returns whether the header still carries the initializing label. - fn open_encrypted_volume( + async fn open_encrypted_volume( &self, - disk_crypt_key: &str, + keys: &AppKeys, name: &str, discard: bool, ) -> Result { let root_hd = &self.args.device; - let disk_crypt_key = disk_crypt_key.trim(); + let disk_crypt_key = hex::encode(&keys.disk_crypt_key); // Create a private tmpfs mount to ensure the header stays in-memory. let tmp_hdr_dir = "/tmp/dstack-luks-header"; let in_mem_hdr = format!("{tmp_hdr_dir}/luks-header"); @@ -2951,9 +3003,27 @@ impl<'a> Stage0<'a> { } .context("Failed to load LUKS2 header")?; - let hdr_file = fs::File::open(&in_mem_hdr).context("Failed to open LUKS2 header")?; - let initializing = - validate_luks2_headers(hdr_file).context("Failed to validate LUKS2 header")?; + let validate = || { + let hdr_file = fs::File::open(&in_mem_hdr).context("Failed to open LUKS2 header")?; + validate_luks2_headers(hdr_file).context("Failed to validate LUKS2 header") + }; + 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}"); + } + if stored < current { + self.migrate_disk_key(keys, Path::new(&in_mem_hdr), stored, current) + .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}"); + } + } + } + let initializing = header.initializing; info!("Opening the device"); let mut command = Command::new("cryptsetup"); @@ -3182,8 +3252,7 @@ impl<'a> Stage0<'a> { opts.storage_encrypted, opts.storage_fs ); - self.mount_data_disk(&hex::encode(&app_keys.disk_crypt_key), &opts) - .await?; + self.mount_data_disk(&app_keys, &opts).await?; self.setup_swap(self.shared.app_compose.swap_size, &opts) .await?; self.vmm @@ -3376,18 +3445,107 @@ macro_rules! const_pad { const PAYLOAD_OFFSET: u64 = 16777216; +fn kms_root_epoch(keys: &AppKeys) -> Option { + match keys.key_provider { + KeyProvider::Kms { root_epoch, .. } => Some(root_epoch), + _ => None, + } +} + +fn cryptsetup(args: &[&str], device: &Path) -> Command { + let mut command = Command::new("cryptsetup"); + command.args(args).arg(device); + command +} + +/// Run `command` with `input` on stdin and return its exit code. +fn run_with_input(mut command: Command, input: &[u8]) -> Result> { + let mut child = command + .stdin(Stdio::piped()) + .spawn() + .context("failed to start cryptsetup")?; + child + .stdin + .take() + .context("cryptsetup stdin is unavailable")? + .write_all(input) + .context("failed to write to cryptsetup")?; + Ok(child + .wait() + .context("failed to wait for cryptsetup")? + .code()) +} + +fn write_luks_root_epoch(device: &Path, epoch: u32) -> Result<()> { + let token = serde_json::json!({ + "type": ROOT_EPOCH_TOKEN_TYPE, + "keyslots": ["0"], + "root_epoch": epoch, + }); + let args = ["token", "import", "--token-id", "0", "--token-replace"]; + let code = run_with_input(cryptsetup(&args, device), token.to_string().as_bytes())?; + if code != Some(0) { + bail!("failed to record root epoch {epoch} in the LUKS header: {code:?}"); + } + Ok(()) +} + +/// 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. +fn rekey_luks_header(header: &Path, old_key: &str, new_key: &str, epoch: u32) -> Result<()> { + use std::os::unix::fs::PermissionsExt; + + // luksHeaderBackup leaves the copy read-only. + fs::set_permissions(header, std::fs::Permissions::from_mode(0o600))?; + let old_file = header.with_file_name("old-key"); + defer! { + fs::remove_file(&old_file).ok(); + } + safe_write_with_mode(&old_file, old_key, 0o600)?; + let mut command = cryptsetup( + &[ + "luksChangeKey", + "--batch-mode", + "--pbkdf", + "pbkdf2", + "--key-slot", + "0", + ], + header, + ); + command.arg("--key-file").arg(&old_file); + if run_with_input(command, new_key.as_bytes())? != Some(0) { + bail!("cryptsetup luksChangeKey failed"); + } + write_luks_root_epoch(header, epoch) +} + /// Set by luksFormat and cleared once the filesystem exists, so a disk still /// carrying it was interrupted mid-initialization and holds no data. const LUKS_INITIALIZING_LABEL: &str = "dstack-initializing"; -/// Returns whether the primary header carries [`LUKS_INITIALIZING_LABEL`]. -fn validate_luks2_headers(mut reader: impl std::io::Read) -> Result { - let initializing = validate_single_luks2_header(&mut reader, 0)?; - validate_single_luks2_header(&mut reader, 1)?; - Ok(initializing) +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +struct LuksHeaderState { + /// The header carries [`LUKS_INITIALIZING_LABEL`]. + initializing: bool, + /// Root epoch recorded in the dstack token; `None` means epoch 0. + root_epoch: Option, } -fn validate_single_luks2_header(mut reader: impl std::io::Read, hdr_ind: u64) -> Result { +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 { + bail!("LUKS header copies disagree on the key epoch"); + } + Ok(primary) +} + +fn validate_single_luks2_header( + mut reader: impl std::io::Read, + hdr_ind: u64, +) -> Result { let mut hdr_data = vec![0u8; BINARY_HEADER_SIZE as usize]; reader .read_exact(&mut hdr_data) @@ -3458,9 +3616,17 @@ fn validate_single_luks2_header(mut reader: impl std::io::Read, hdr_ind: u64) -> if keyslots.len() != 1 { bail!("Invalid LUKS keyslots"); } - if !tokens.is_empty() { - bail!("Invalid LUKS tokens"); - } + let root_epoch = match tokens.len() { + 0 => None, + 1 => { + let token = tokens.get(&0).context("Invalid LUKS token id")?; + if token.kind != ROOT_EPOCH_TOKEN_TYPE || token.keyslots != ["0"] { + bail!("Invalid LUKS token"); + } + Some(token.root_epoch.context("LUKS token has no root epoch")?) + } + _ => bail!("Invalid LUKS tokens"), + }; if segments.len() != 1 { bail!("Invalid LUKS segments"); } @@ -3576,18 +3742,31 @@ fn validate_single_luks2_header(mut reader: impl std::io::Read, hdr_ind: u64) -> bail!("Invalid LUKS digest segments: {segments:?}"); } } - Ok(initializing) + Ok(LuksHeaderState { + initializing, + root_epoch, + }) } #[test] fn test_validate_luks2_header() { let mut header_data = include_bytes!("../tests/fixtures/luks_header_good").to_vec(); - assert!( - !validate_luks2_headers(&mut &header_data[..]).expect("Failed to validate LUKS2 header") + let state = + validate_luks2_headers(&mut &header_data[..]).expect("Failed to validate LUKS2 header"); + assert_eq!( + state, + LuksHeaderState { + initializing: false, + root_epoch: None + } ); header_data[24..24 + LUKS_INITIALIZING_LABEL.len()] .copy_from_slice(LUKS_INITIALIZING_LABEL.as_bytes()); - assert!(validate_luks2_headers(&mut &header_data[..]).expect("Failed to validate LUKS2 header")); + assert!( + validate_luks2_headers(&mut &header_data[..]) + .expect("Failed to validate LUKS2 header") + .initializing + ); let header_data = include_bytes!("../tests/fixtures/luks_header_cipher_null").to_vec(); let error = validate_luks2_headers(&mut &header_data[..]).unwrap_err(); assert!(error @@ -3631,6 +3810,109 @@ fn patch_luks_json(header: &mut [u8], needle: &[u8], replacement: &[u8]) -> usiz patched } +#[test] +fn test_disk_key_migration_with_cryptsetup() { + if Command::new("cryptsetup") + .arg("--version") + .output() + .is_err() + { + eprintln!("cryptsetup is unavailable; skipping"); + return; + } + let dir = tempfile::tempdir().unwrap(); + let image = dir.path().join("disk.img"); + fs::File::create(&image) + .unwrap() + .set_len(2 * PAYLOAD_OFFSET) + .unwrap(); + let offset = (PAYLOAD_OFFSET / 512).to_string(); + let format = cryptsetup( + &[ + "luksFormat", + "--batch-mode", + "--type", + "luks2", + "--offset", + &offset, + "--cipher", + "aes-xts-plain64", + "--pbkdf", + "pbkdf2", + "--sector-size", + "512", + "-d-", + ], + &image, + ); + assert_eq!(run_with_input(format, b"old").unwrap(), Some(0)); + let unlocks = |key: &str| { + let test = cryptsetup( + &["open", "--test-passphrase", "--key-slot", "0", "-d-"], + &image, + ); + run_with_input(test, key.as_bytes()).unwrap() == Some(0) + }; + + // Re-key a detached copy of the header, as the boot does, and restore it. + let header = dir.path().join("header"); + let backup = cryptsetup(&["luksHeaderBackup", "--header-backup-file"], &header) + .arg(&image) + .status() + .unwrap(); + assert!(backup.success()); + rekey_luks_header(&header, "old", "new", 2).unwrap(); + assert!(!dir.path().join("old-key").exists()); + assert!( + unlocks("old"), + "the disk is unchanged until the header is restored" + ); + let state = validate_luks2_headers(fs::File::open(&header).unwrap()).unwrap(); + assert_eq!(state.root_epoch, Some(2)); + + let restore = cryptsetup( + &["luksHeaderRestore", "--batch-mode", "--header-backup-file"], + &header, + ) + .arg(&image) + .status() + .unwrap(); + assert!(restore.success()); + assert!(unlocks("new")); + assert!(!unlocks("old")); + + // A later rotation replaces the recorded epoch. + write_luks_root_epoch(&image, 3).unwrap(); + fs::remove_file(&header).unwrap(); + let backup = cryptsetup(&["luksHeaderBackup", "--header-backup-file"], &header) + .arg(&image) + .status() + .unwrap(); + assert!(backup.success()); + let state = validate_luks2_headers(fs::File::open(&header).unwrap()).unwrap(); + assert_eq!(state.root_epoch, Some(3)); +} + +#[test] +fn test_validate_luks2_header_root_epoch_token() { + let with_token = |token: &str| { + let mut header = include_bytes!("../tests/fixtures/luks_header_good").to_vec(); + let replacement = format!(r#""tokens":{{"0":{token}}}"#); + let patched = patch_luks_json(&mut header, br#""tokens":{}"#, replacement.as_bytes()); + assert_eq!(patched, 2, "expected to patch both header copies"); + validate_luks2_headers(&mut &header[..]) + }; + let state = with_token(r#"{"type":"dstack-kms","keyslots":["0"],"root_epoch":2}"#).unwrap(); + assert_eq!(state.root_epoch, Some(2)); + for token in [ + r#"{"type":"luks2-keyring","keyslots":["0"],"root_epoch":2}"#, + r#"{"type":"dstack-kms","keyslots":["1"],"root_epoch":2}"#, + r#"{"type":"dstack-kms","keyslots":["0"]}"#, + ] { + assert!(with_token(token).is_err(), "{token}"); + } +} + #[test] fn test_validate_luks2_header_rejects_keyslot_area_that_overflows() { let mut header = include_bytes!("../tests/fixtures/luks_header_good").to_vec(); diff --git a/dstack/dstack-util/src/system_setup/luks.rs b/dstack/dstack-util/src/system_setup/luks.rs index 6d0381b9e..aaa184a1c 100644 --- a/dstack/dstack-util/src/system_setup/luks.rs +++ b/dstack/dstack-util/src/system_setup/luks.rs @@ -49,12 +49,24 @@ impl LuksHeader { #[derive(Deserialize)] pub(super) struct LuksJson { pub keyslots: HashMap, - pub tokens: HashMap, + pub tokens: HashMap, pub segments: HashMap, pub digests: HashMap, pub config: LuksConfig, } +/// Type of the only token dstack writes: the KMS key epoch of keyslot 0. +pub(super) const ROOT_EPOCH_TOKEN_TYPE: &str = "dstack-kms"; + +#[derive(Deserialize)] +pub(super) struct LuksToken { + #[serde(rename = "type")] + pub kind: String, + pub keyslots: Vec, + #[serde(default)] + pub root_epoch: Option, +} + #[derive(Deserialize)] #[serde(tag = "type", rename_all = "lowercase")] pub(super) enum LuksKeyslot { From 0455c68763cf06d2ab82aeeca5b94ff249300ec1 Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Sat, 7 Nov 2026 07:19:53 -0800 Subject: [PATCH 041/166] refactor(dstack-util): simplify shared app key request name Signed-off-by: Kevin Wang --- dstack/dstack-util/src/system_setup.rs | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/dstack/dstack-util/src/system_setup.rs b/dstack/dstack-util/src/system_setup.rs index e58ae61e4..0321f17c1 100644 --- a/dstack/dstack-util/src/system_setup.rs +++ b/dstack/dstack-util/src/system_setup.rs @@ -2355,7 +2355,7 @@ struct AppInfo { impl HostShared { /// Request app keys of `root_epoch` (`None` = current) from one KMS. - async fn request_app_keys_from_kms_url( + async fn request_app_keys_from_kms( &self, kms_url: String, root_epoch: Option, @@ -2408,7 +2408,7 @@ impl HostShared { bail!("only KMS-provided keys have older root epochs"); }; Ok(self - .request_app_keys_from_kms_url(url.clone(), Some(epoch), Some(keys.ca_cert.clone())) + .request_app_keys_from_kms(url.clone(), Some(epoch), Some(keys.ca_cert.clone())) .await? .0) } @@ -2485,7 +2485,7 @@ impl<'a> Stage0<'a> { let kms_url = kms_rpc_url(kms_url); match self .shared - .request_app_keys_from_kms_url(kms_url.clone(), None, None) + .request_app_keys_from_kms(kms_url.clone(), None, None) .await { Ok((keys, os_image_hash)) => { From 37b9e0891c9caa451da7acaac77a83c9bdb81df6 Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Sat, 7 Nov 2026 18:12:21 -0800 Subject: [PATCH 042/166] refactor(dstack-util): keep disk header operations in luks module Signed-off-by: Kevin Wang --- dstack/dstack-util/src/system_setup.rs | 500 +----------------- dstack/dstack-util/src/system_setup/luks.rs | 541 +++++++++++++++++++- 2 files changed, 527 insertions(+), 514 deletions(-) diff --git a/dstack/dstack-util/src/system_setup.rs b/dstack/dstack-util/src/system_setup.rs index 0321f17c1..0699befb9 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::{ - LuksAf, LuksArea, LuksConfig, LuksDigest, LuksHeader, LuksJson, LuksKdf, LuksKeyslot, - LuksSegment, BINARY_HEADER_SIZE, ROOT_EPOCH_TOKEN_TYPE, + rekey_luks_header, validate_luks2_headers, write_luks_root_epoch, LUKS_INITIALIZING_LABEL, + PAYLOAD_OFFSET, }; use ra_rpc::{ client::{CertInfo, RaClient, RaClientConfig}, @@ -3428,23 +3428,6 @@ impl Stage1<'_> { } } -macro_rules! const_pad { - ($s:expr, $len:expr) => { - const { - assert!($s.len() <= $len, "The s is too long"); - let mut padded: [u8; $len] = [0; $len]; - let mut i = 0; - while i < $s.len() { - padded[i] = $s[i]; - i += 1; - } - padded - } - }; -} - -const PAYLOAD_OFFSET: u64 = 16777216; - fn kms_root_epoch(keys: &AppKeys) -> Option { match keys.key_provider { KeyProvider::Kms { root_epoch, .. } => Some(root_epoch), @@ -3452,485 +3435,6 @@ fn kms_root_epoch(keys: &AppKeys) -> Option { } } -fn cryptsetup(args: &[&str], device: &Path) -> Command { - let mut command = Command::new("cryptsetup"); - command.args(args).arg(device); - command -} - -/// Run `command` with `input` on stdin and return its exit code. -fn run_with_input(mut command: Command, input: &[u8]) -> Result> { - let mut child = command - .stdin(Stdio::piped()) - .spawn() - .context("failed to start cryptsetup")?; - child - .stdin - .take() - .context("cryptsetup stdin is unavailable")? - .write_all(input) - .context("failed to write to cryptsetup")?; - Ok(child - .wait() - .context("failed to wait for cryptsetup")? - .code()) -} - -fn write_luks_root_epoch(device: &Path, epoch: u32) -> Result<()> { - let token = serde_json::json!({ - "type": ROOT_EPOCH_TOKEN_TYPE, - "keyslots": ["0"], - "root_epoch": epoch, - }); - let args = ["token", "import", "--token-id", "0", "--token-replace"]; - let code = run_with_input(cryptsetup(&args, device), token.to_string().as_bytes())?; - if code != Some(0) { - bail!("failed to record root epoch {epoch} in the LUKS header: {code:?}"); - } - Ok(()) -} - -/// 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. -fn rekey_luks_header(header: &Path, old_key: &str, new_key: &str, epoch: u32) -> Result<()> { - use std::os::unix::fs::PermissionsExt; - - // luksHeaderBackup leaves the copy read-only. - fs::set_permissions(header, std::fs::Permissions::from_mode(0o600))?; - let old_file = header.with_file_name("old-key"); - defer! { - fs::remove_file(&old_file).ok(); - } - safe_write_with_mode(&old_file, old_key, 0o600)?; - let mut command = cryptsetup( - &[ - "luksChangeKey", - "--batch-mode", - "--pbkdf", - "pbkdf2", - "--key-slot", - "0", - ], - header, - ); - command.arg("--key-file").arg(&old_file); - if run_with_input(command, new_key.as_bytes())? != Some(0) { - bail!("cryptsetup luksChangeKey failed"); - } - write_luks_root_epoch(header, epoch) -} - -/// Set by luksFormat and cleared once the filesystem exists, so a disk still -/// carrying it was interrupted mid-initialization and holds no data. -const LUKS_INITIALIZING_LABEL: &str = "dstack-initializing"; - -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -struct LuksHeaderState { - /// The header carries [`LUKS_INITIALIZING_LABEL`]. - initializing: bool, - /// Root epoch recorded in the dstack token; `None` means epoch 0. - root_epoch: Option, -} - -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 { - bail!("LUKS header copies disagree on the key epoch"); - } - Ok(primary) -} - -fn validate_single_luks2_header( - mut reader: impl std::io::Read, - hdr_ind: u64, -) -> Result { - let mut hdr_data = vec![0u8; BINARY_HEADER_SIZE as usize]; - reader - .read_exact(&mut hdr_data) - .context("Failed to read LUKS header")?; - let LuksHeader { - magic, - version, - hdr_size, - label, - csum_alg, - subsystem, - hdr_offset, - } = LuksHeader::parse(&hdr_data).context("Failed to decode LUKS header")?; - - let expected_magic = match hdr_ind { - 0 => [76, 85, 75, 83, 186, 190], - 1 => [83, 75, 85, 76, 186, 190], - _ => bail!("Invalid LUKS header index: {hdr_ind}"), - }; - if magic != expected_magic { - bail!("Invalid LUKS magic: {magic:?}"); - } - if version != 2 { - bail!("Invalid LUKS version: {version}"); - } - let initializing = label == const_pad!(LUKS_INITIALIZING_LABEL.as_bytes(), 48); - if !initializing && label != [0; 48] { - bail!("Invalid LUKS label: {:?}", label); - } - if csum_alg != const_pad!(b"sha256", 32) { - bail!("Invalid LUKS checksum algorithm"); - } - if subsystem != [0; 48] { - bail!("Invalid LUKS subsystem"); - } - if hdr_offset != hdr_ind * hdr_size { - bail!("Invalid LUKS header offset: {hdr_offset}"); - } - if !(4096..=1024 * 1024 * 16).contains(&hdr_size) { - bail!("Invalid LUKS header size: {hdr_size}"); - } - - // Check JSON - let json_size = hdr_size - BINARY_HEADER_SIZE; - let mut jsn_data = vec![0u8; json_size as usize]; - reader - .read_exact(&mut jsn_data) - .context("Failed to read LUKS JSON")?; - let json_end = jsn_data - .iter() - .position(|&b| b == 0) - .unwrap_or(jsn_data.len()); - jsn_data.truncate(json_end); - - let LuksJson { - keyslots, - tokens, - segments, - digests, - config: - LuksConfig { - keyslots_size, - flags, - requirements, - }, - } = serde_json::from_slice(&jsn_data).context("Failed to decode LUKS JSON")?; - - if keyslots.len() != 1 { - bail!("Invalid LUKS keyslots"); - } - let root_epoch = match tokens.len() { - 0 => None, - 1 => { - let token = tokens.get(&0).context("Invalid LUKS token id")?; - if token.kind != ROOT_EPOCH_TOKEN_TYPE || token.keyslots != ["0"] { - bail!("Invalid LUKS token"); - } - Some(token.root_epoch.context("LUKS token has no root epoch")?) - } - _ => bail!("Invalid LUKS tokens"), - }; - if segments.len() != 1 { - bail!("Invalid LUKS segments"); - } - if digests.len() != 1 { - bail!("Invalid LUKS digests"); - } - if keyslots_size % 4096 != 0 { - bail!("Invalid LUKS keyslots size: {keyslots_size}"); - } - if flags.is_some() { - bail!("Invalid LUKS flags"); - } - if requirements.is_some() { - bail!("Invalid LUKS requirements"); - } - - { - let LuksKeyslot::Luks2 { - key_size, - area: - LuksArea::Raw { - encryption, - key_size: area_key_size, - offset, - size, - }, - kdf, - af: - LuksAf::Luks1 { - hash: af_hash, - stripes, - }, - priority, - } = keyslots.get(&0).context("no LUKS keyslot")?; - if encryption != "aes-xts-plain64" { - bail!("Invalid LUKS keyslot encryption: {encryption}"); - } - // Pin where the encrypted key material is read from. The binary area - // must sit between the two header copies and the encrypted payload; - // otherwise a host with raw disk access could redirect it elsewhere. - let area_end = offset.checked_add(*size); - if *offset < 2 * hdr_size || area_end.is_none_or(|end| end > PAYLOAD_OFFSET) { - bail!("Invalid LUKS keyslot area: offset={offset} size={size}"); - } - if *key_size != 64 { - bail!("Invalid LUKS keyslot key size: {key_size}"); - } - if *area_key_size != 64 { - bail!("Invalid LUKS keyslot key size: {area_key_size}"); - } - let LuksKdf::Pbkdf2 { hash: kdf_hash } = kdf else { - bail!("Invalid LUKS keyslot KDF"); - }; - if kdf_hash != "sha256" { - bail!("Invalid LUKS keyslot hash: {kdf_hash}"); - } - if af_hash != "sha256" { - bail!("Invalid LUKS keyslot hash: {af_hash}"); - } - if *stripes != 4000 { - bail!("Invalid LUKS keyslot stripes: {stripes}"); - } - if priority.is_some() { - bail!("Invalid LUKS keyslot priority"); - } - } - - { - let LuksSegment::Crypt { - offset, - size, - iv_tweak, - encryption, - sector_size, - integrity, - flags, - } = segments.get(&0).context("no LUKS segment")?; - if *offset != PAYLOAD_OFFSET { - bail!("Invalid LUKS segment offset"); - } - if size != "dynamic" { - bail!("Invalid LUKS segment size"); - } - if *iv_tweak != 0 { - bail!("Invalid LUKS segment IV tweak"); - } - if encryption != "aes-xts-plain64" { - bail!("Invalid LUKS segment encryption"); - } - if *sector_size != 512 { - bail!("Invalid LUKS segment sector size"); - } - if integrity.is_some() { - bail!("Invalid LUKS segment integrity"); - } - if flags.is_some() { - bail!("Invalid LUKS segment flags"); - } - } - { - let LuksDigest::Pbkdf2 { - keyslots, - segments, - hash, - } = digests.get(&0).context("no LUKS digest")?; - if hash != "sha256" { - bail!("Invalid LUKS digest hash: {hash}"); - } - if keyslots != &[0] { - bail!("Invalid LUKS digest keyslots: {keyslots:?}"); - } - if segments != &[0] { - bail!("Invalid LUKS digest segments: {segments:?}"); - } - } - Ok(LuksHeaderState { - initializing, - root_epoch, - }) -} - -#[test] -fn test_validate_luks2_header() { - let mut header_data = include_bytes!("../tests/fixtures/luks_header_good").to_vec(); - let state = - validate_luks2_headers(&mut &header_data[..]).expect("Failed to validate LUKS2 header"); - assert_eq!( - state, - LuksHeaderState { - initializing: false, - root_epoch: None - } - ); - header_data[24..24 + LUKS_INITIALIZING_LABEL.len()] - .copy_from_slice(LUKS_INITIALIZING_LABEL.as_bytes()); - assert!( - validate_luks2_headers(&mut &header_data[..]) - .expect("Failed to validate LUKS2 header") - .initializing - ); - let header_data = include_bytes!("../tests/fixtures/luks_header_cipher_null").to_vec(); - let error = validate_luks2_headers(&mut &header_data[..]).unwrap_err(); - assert!(error - .to_string() - .contains("Invalid LUKS keyslot encryption")); -} - -#[test] -fn test_validate_luks2_header_rejects_out_of_range_keyslot_area() { - // Redirect the keyslot binary area below the header region. Same length - // so the surrounding header stays intact; "00768" parses to 768, which is - // inside the header copies (< 2 * hdr_size) rather than the metadata gap. - let mut header = include_bytes!("../tests/fixtures/luks_header_good").to_vec(); - let patched = patch_luks_json(&mut header, br#""offset":"32768""#, br#""offset":"00768""#); - assert_eq!(patched, 2, "expected to patch both header copies"); - let error = validate_luks2_headers(&mut &header[..]).unwrap_err(); - assert!(error.to_string().contains("Invalid LUKS keyslot area")); -} - -/// Replace `needle` with `replacement` inside every LUKS JSON region, keeping -/// each region the same length by giving back trailing NUL padding. -#[cfg(test)] -fn patch_luks_json(header: &mut [u8], needle: &[u8], replacement: &[u8]) -> usize { - assert!(replacement.len() >= needle.len()); - let grow = replacement.len() - needle.len(); - let mut patched = 0; - let mut i = 0; - while i + needle.len() <= header.len() { - if &header[i..i + needle.len()] != needle { - i += 1; - continue; - } - let tail = &mut header[i..]; - let end = tail.iter().position(|b| *b == 0).expect("NUL padding"); - assert!(end + grow < tail.len(), "not enough NUL padding"); - tail.copy_within(needle.len()..end, replacement.len()); - tail[..replacement.len()].copy_from_slice(replacement); - patched += 1; - i += replacement.len(); - } - patched -} - -#[test] -fn test_disk_key_migration_with_cryptsetup() { - if Command::new("cryptsetup") - .arg("--version") - .output() - .is_err() - { - eprintln!("cryptsetup is unavailable; skipping"); - return; - } - let dir = tempfile::tempdir().unwrap(); - let image = dir.path().join("disk.img"); - fs::File::create(&image) - .unwrap() - .set_len(2 * PAYLOAD_OFFSET) - .unwrap(); - let offset = (PAYLOAD_OFFSET / 512).to_string(); - let format = cryptsetup( - &[ - "luksFormat", - "--batch-mode", - "--type", - "luks2", - "--offset", - &offset, - "--cipher", - "aes-xts-plain64", - "--pbkdf", - "pbkdf2", - "--sector-size", - "512", - "-d-", - ], - &image, - ); - assert_eq!(run_with_input(format, b"old").unwrap(), Some(0)); - let unlocks = |key: &str| { - let test = cryptsetup( - &["open", "--test-passphrase", "--key-slot", "0", "-d-"], - &image, - ); - run_with_input(test, key.as_bytes()).unwrap() == Some(0) - }; - - // Re-key a detached copy of the header, as the boot does, and restore it. - let header = dir.path().join("header"); - let backup = cryptsetup(&["luksHeaderBackup", "--header-backup-file"], &header) - .arg(&image) - .status() - .unwrap(); - assert!(backup.success()); - rekey_luks_header(&header, "old", "new", 2).unwrap(); - assert!(!dir.path().join("old-key").exists()); - assert!( - unlocks("old"), - "the disk is unchanged until the header is restored" - ); - let state = validate_luks2_headers(fs::File::open(&header).unwrap()).unwrap(); - assert_eq!(state.root_epoch, Some(2)); - - let restore = cryptsetup( - &["luksHeaderRestore", "--batch-mode", "--header-backup-file"], - &header, - ) - .arg(&image) - .status() - .unwrap(); - assert!(restore.success()); - assert!(unlocks("new")); - assert!(!unlocks("old")); - - // A later rotation replaces the recorded epoch. - write_luks_root_epoch(&image, 3).unwrap(); - fs::remove_file(&header).unwrap(); - let backup = cryptsetup(&["luksHeaderBackup", "--header-backup-file"], &header) - .arg(&image) - .status() - .unwrap(); - assert!(backup.success()); - let state = validate_luks2_headers(fs::File::open(&header).unwrap()).unwrap(); - assert_eq!(state.root_epoch, Some(3)); -} - -#[test] -fn test_validate_luks2_header_root_epoch_token() { - let with_token = |token: &str| { - let mut header = include_bytes!("../tests/fixtures/luks_header_good").to_vec(); - let replacement = format!(r#""tokens":{{"0":{token}}}"#); - let patched = patch_luks_json(&mut header, br#""tokens":{}"#, replacement.as_bytes()); - assert_eq!(patched, 2, "expected to patch both header copies"); - validate_luks2_headers(&mut &header[..]) - }; - let state = with_token(r#"{"type":"dstack-kms","keyslots":["0"],"root_epoch":2}"#).unwrap(); - assert_eq!(state.root_epoch, Some(2)); - for token in [ - r#"{"type":"luks2-keyring","keyslots":["0"],"root_epoch":2}"#, - r#"{"type":"dstack-kms","keyslots":["1"],"root_epoch":2}"#, - r#"{"type":"dstack-kms","keyslots":["0"]}"#, - ] { - assert!(with_token(token).is_err(), "{token}"); - } -} - -#[test] -fn test_validate_luks2_header_rejects_keyslot_area_that_overflows() { - let mut header = include_bytes!("../tests/fixtures/luks_header_good").to_vec(); - // 2**64 - 32768 + 1000: added to the accepted offset 32768 it wraps to - // 1000, which is below PAYLOAD_OFFSET. - let patched = patch_luks_json( - &mut header, - br#""size":"258048""#, - br#""size":"18446744073709519848""#, - ); - assert_eq!(patched, 2, "expected to patch both header copies"); - let error = validate_luks2_headers(&mut &header[..]).unwrap_err(); - assert!( - error.to_string().contains("Invalid LUKS keyslot area"), - "{error:#}" - ); -} - #[cfg(test)] fn test_app_compose(manifest_version: serde_json::Value, platforms: Option<&[&str]>) -> AppCompose { let mut value = serde_json::json!({ diff --git a/dstack/dstack-util/src/system_setup/luks.rs b/dstack/dstack-util/src/system_setup/luks.rs index aaa184a1c..5b5e2fef1 100644 --- a/dstack/dstack-util/src/system_setup/luks.rs +++ b/dstack/dstack-util/src/system_setup/luks.rs @@ -2,20 +2,28 @@ // // SPDX-License-Identifier: Apache-2.0 -//! The subset of the LUKS2 on-disk format that `validate_luks2_headers` -//! checks. Fields it does not inspect are not modeled; serde ignores them. +//! LUKS2 header validation and key epoch migration. +//! Fields validation does not inspect are not modeled; serde ignores them. //! See . -use std::collections::HashMap; +use std::{ + collections::HashMap, + io::Write, + path::Path, + process::{Command, Stdio}, +}; -use anyhow::{Context, Result}; +use anyhow::{bail, Context, Result}; +use fs_err as fs; +use safe_write::safe_write_with_mode; +use scopeguard::defer; use serde::{de::IgnoredAny, Deserialize}; use serde_with::{serde_as, DisplayFromStr}; /// Size of the binary header that precedes each JSON area. -pub(super) const BINARY_HEADER_SIZE: u64 = 4096; +const BINARY_HEADER_SIZE: u64 = 4096; -pub(super) struct LuksHeader { +struct LuksHeader { pub magic: [u8; 6], pub version: u16, pub hdr_size: u64, @@ -47,7 +55,7 @@ impl LuksHeader { } #[derive(Deserialize)] -pub(super) struct LuksJson { +struct LuksJson { pub keyslots: HashMap, pub tokens: HashMap, pub segments: HashMap, @@ -56,10 +64,10 @@ pub(super) struct LuksJson { } /// Type of the only token dstack writes: the KMS key epoch of keyslot 0. -pub(super) const ROOT_EPOCH_TOKEN_TYPE: &str = "dstack-kms"; +const ROOT_EPOCH_TOKEN_TYPE: &str = "dstack-kms"; #[derive(Deserialize)] -pub(super) struct LuksToken { +struct LuksToken { #[serde(rename = "type")] pub kind: String, pub keyslots: Vec, @@ -69,7 +77,7 @@ pub(super) struct LuksToken { #[derive(Deserialize)] #[serde(tag = "type", rename_all = "lowercase")] -pub(super) enum LuksKeyslot { +enum LuksKeyslot { Luks2 { key_size: u16, area: LuksArea, @@ -83,7 +91,7 @@ pub(super) enum LuksKeyslot { #[serde_as] #[derive(Deserialize)] #[serde(tag = "type", rename_all = "lowercase")] -pub(super) enum LuksArea { +enum LuksArea { Raw { encryption: String, key_size: u32, @@ -99,7 +107,7 @@ pub(super) enum LuksArea { /// it (see security report #552). #[derive(Deserialize)] #[serde(tag = "type", rename_all = "lowercase")] -pub(super) enum LuksKdf { +enum LuksKdf { Pbkdf2 { hash: String, }, @@ -109,14 +117,14 @@ pub(super) enum LuksKdf { #[derive(Deserialize)] #[serde(tag = "type", rename_all = "lowercase")] -pub(super) enum LuksAf { +enum LuksAf { Luks1 { hash: String, stripes: u16 }, } #[serde_as] #[derive(Deserialize)] #[serde(tag = "type", rename_all = "lowercase")] -pub(super) enum LuksSegment { +enum LuksSegment { Crypt { #[serde_as(as = "DisplayFromStr")] offset: u64, @@ -135,7 +143,7 @@ pub(super) enum LuksSegment { #[serde_as] #[derive(Deserialize)] #[serde(tag = "type", rename_all = "lowercase")] -pub(super) enum LuksDigest { +enum LuksDigest { Pbkdf2 { #[serde_as(as = "Vec")] keyslots: Vec, @@ -147,7 +155,7 @@ pub(super) enum LuksDigest { #[serde_as] #[derive(Deserialize)] -pub(super) struct LuksConfig { +struct LuksConfig { #[serde_as(as = "DisplayFromStr")] pub keyslots_size: u64, #[serde(default)] @@ -155,3 +163,504 @@ pub(super) struct LuksConfig { #[serde(default)] pub requirements: Option, } + +macro_rules! const_pad { + ($s:expr, $len:expr) => { + const { + assert!($s.len() <= $len, "The s is too long"); + let mut padded: [u8; $len] = [0; $len]; + let mut i = 0; + while i < $s.len() { + padded[i] = $s[i]; + i += 1; + } + padded + } + }; +} + +pub(super) const PAYLOAD_OFFSET: u64 = 16777216; + +fn cryptsetup(args: &[&str], device: &Path) -> Command { + let mut command = Command::new("cryptsetup"); + command.args(args).arg(device); + command +} + +/// Run `command` with `input` on stdin and return its exit code. +fn run_with_input(mut command: Command, input: &[u8]) -> Result> { + let mut child = command + .stdin(Stdio::piped()) + .spawn() + .context("failed to start cryptsetup")?; + child + .stdin + .take() + .context("cryptsetup stdin is unavailable")? + .write_all(input) + .context("failed to write to cryptsetup")?; + Ok(child + .wait() + .context("failed to wait for cryptsetup")? + .code()) +} + +pub(super) fn write_luks_root_epoch(device: &Path, epoch: u32) -> Result<()> { + let token = serde_json::json!({ + "type": ROOT_EPOCH_TOKEN_TYPE, + "keyslots": ["0"], + "root_epoch": epoch, + }); + let args = ["token", "import", "--token-id", "0", "--token-replace"]; + let code = run_with_input(cryptsetup(&args, device), token.to_string().as_bytes())?; + if code != Some(0) { + bail!("failed to record root epoch {epoch} in the LUKS header: {code:?}"); + } + Ok(()) +} + +/// 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( + header: &Path, + old_key: &str, + new_key: &str, + epoch: u32, +) -> Result<()> { + use std::os::unix::fs::PermissionsExt; + + // luksHeaderBackup leaves the copy read-only. + fs::set_permissions(header, std::fs::Permissions::from_mode(0o600))?; + let old_file = header.with_file_name("old-key"); + defer! { + fs::remove_file(&old_file).ok(); + } + safe_write_with_mode(&old_file, old_key, 0o600)?; + let mut command = cryptsetup( + &[ + "luksChangeKey", + "--batch-mode", + "--pbkdf", + "pbkdf2", + "--key-slot", + "0", + ], + header, + ); + command.arg("--key-file").arg(&old_file); + if run_with_input(command, new_key.as_bytes())? != Some(0) { + bail!("cryptsetup luksChangeKey failed"); + } + write_luks_root_epoch(header, epoch) +} + +/// Set by luksFormat and cleared once the filesystem exists, so a disk still +/// carrying it was interrupted mid-initialization and holds no data. +pub(super) const LUKS_INITIALIZING_LABEL: &str = "dstack-initializing"; + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(super) struct LuksHeaderState { + /// The header carries [`LUKS_INITIALIZING_LABEL`]. + pub(super) initializing: bool, + /// Root epoch recorded in the dstack token; `None` means epoch 0. + pub(super) root_epoch: Option, +} + +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 { + bail!("LUKS header copies disagree on the key epoch"); + } + Ok(primary) +} + +fn validate_single_luks2_header( + mut reader: impl std::io::Read, + hdr_ind: u64, +) -> Result { + let mut hdr_data = vec![0u8; BINARY_HEADER_SIZE as usize]; + reader + .read_exact(&mut hdr_data) + .context("Failed to read LUKS header")?; + let LuksHeader { + magic, + version, + hdr_size, + label, + csum_alg, + subsystem, + hdr_offset, + } = LuksHeader::parse(&hdr_data).context("Failed to decode LUKS header")?; + + let expected_magic = match hdr_ind { + 0 => [76, 85, 75, 83, 186, 190], + 1 => [83, 75, 85, 76, 186, 190], + _ => bail!("Invalid LUKS header index: {hdr_ind}"), + }; + if magic != expected_magic { + bail!("Invalid LUKS magic: {magic:?}"); + } + if version != 2 { + bail!("Invalid LUKS version: {version}"); + } + let initializing = label == const_pad!(LUKS_INITIALIZING_LABEL.as_bytes(), 48); + if !initializing && label != [0; 48] { + bail!("Invalid LUKS label: {:?}", label); + } + if csum_alg != const_pad!(b"sha256", 32) { + bail!("Invalid LUKS checksum algorithm"); + } + if subsystem != [0; 48] { + bail!("Invalid LUKS subsystem"); + } + if hdr_offset != hdr_ind * hdr_size { + bail!("Invalid LUKS header offset: {hdr_offset}"); + } + if !(4096..=1024 * 1024 * 16).contains(&hdr_size) { + bail!("Invalid LUKS header size: {hdr_size}"); + } + + // Check JSON + let json_size = hdr_size - BINARY_HEADER_SIZE; + let mut jsn_data = vec![0u8; json_size as usize]; + reader + .read_exact(&mut jsn_data) + .context("Failed to read LUKS JSON")?; + let json_end = jsn_data + .iter() + .position(|&b| b == 0) + .unwrap_or(jsn_data.len()); + jsn_data.truncate(json_end); + + let LuksJson { + keyslots, + tokens, + segments, + digests, + config: + LuksConfig { + keyslots_size, + flags, + requirements, + }, + } = serde_json::from_slice(&jsn_data).context("Failed to decode LUKS JSON")?; + + if keyslots.len() != 1 { + bail!("Invalid LUKS keyslots"); + } + let root_epoch = match tokens.len() { + 0 => None, + 1 => { + let token = tokens.get(&0).context("Invalid LUKS token id")?; + if token.kind != ROOT_EPOCH_TOKEN_TYPE || token.keyslots != ["0"] { + bail!("Invalid LUKS token"); + } + Some(token.root_epoch.context("LUKS token has no root epoch")?) + } + _ => bail!("Invalid LUKS tokens"), + }; + if segments.len() != 1 { + bail!("Invalid LUKS segments"); + } + if digests.len() != 1 { + bail!("Invalid LUKS digests"); + } + if keyslots_size % 4096 != 0 { + bail!("Invalid LUKS keyslots size: {keyslots_size}"); + } + if flags.is_some() { + bail!("Invalid LUKS flags"); + } + if requirements.is_some() { + bail!("Invalid LUKS requirements"); + } + + { + let LuksKeyslot::Luks2 { + key_size, + area: + LuksArea::Raw { + encryption, + key_size: area_key_size, + offset, + size, + }, + kdf, + af: + LuksAf::Luks1 { + hash: af_hash, + stripes, + }, + priority, + } = keyslots.get(&0).context("no LUKS keyslot")?; + if encryption != "aes-xts-plain64" { + bail!("Invalid LUKS keyslot encryption: {encryption}"); + } + // Pin where the encrypted key material is read from. The binary area + // must sit between the two header copies and the encrypted payload; + // otherwise a host with raw disk access could redirect it elsewhere. + let area_end = offset.checked_add(*size); + if *offset < 2 * hdr_size || area_end.is_none_or(|end| end > PAYLOAD_OFFSET) { + bail!("Invalid LUKS keyslot area: offset={offset} size={size}"); + } + if *key_size != 64 { + bail!("Invalid LUKS keyslot key size: {key_size}"); + } + if *area_key_size != 64 { + bail!("Invalid LUKS keyslot key size: {area_key_size}"); + } + let LuksKdf::Pbkdf2 { hash: kdf_hash } = kdf else { + bail!("Invalid LUKS keyslot KDF"); + }; + if kdf_hash != "sha256" { + bail!("Invalid LUKS keyslot hash: {kdf_hash}"); + } + if af_hash != "sha256" { + bail!("Invalid LUKS keyslot hash: {af_hash}"); + } + if *stripes != 4000 { + bail!("Invalid LUKS keyslot stripes: {stripes}"); + } + if priority.is_some() { + bail!("Invalid LUKS keyslot priority"); + } + } + + { + let LuksSegment::Crypt { + offset, + size, + iv_tweak, + encryption, + sector_size, + integrity, + flags, + } = segments.get(&0).context("no LUKS segment")?; + if *offset != PAYLOAD_OFFSET { + bail!("Invalid LUKS segment offset"); + } + if size != "dynamic" { + bail!("Invalid LUKS segment size"); + } + if *iv_tweak != 0 { + bail!("Invalid LUKS segment IV tweak"); + } + if encryption != "aes-xts-plain64" { + bail!("Invalid LUKS segment encryption"); + } + if *sector_size != 512 { + bail!("Invalid LUKS segment sector size"); + } + if integrity.is_some() { + bail!("Invalid LUKS segment integrity"); + } + if flags.is_some() { + bail!("Invalid LUKS segment flags"); + } + } + { + let LuksDigest::Pbkdf2 { + keyslots, + segments, + hash, + } = digests.get(&0).context("no LUKS digest")?; + if hash != "sha256" { + bail!("Invalid LUKS digest hash: {hash}"); + } + if keyslots != &[0] { + bail!("Invalid LUKS digest keyslots: {keyslots:?}"); + } + if segments != &[0] { + bail!("Invalid LUKS digest segments: {segments:?}"); + } + } + Ok(LuksHeaderState { + initializing, + root_epoch, + }) +} + +#[test] +fn test_validate_luks2_header() { + let mut header_data = include_bytes!("../../tests/fixtures/luks_header_good").to_vec(); + let state = + validate_luks2_headers(&mut &header_data[..]).expect("Failed to validate LUKS2 header"); + assert_eq!( + state, + LuksHeaderState { + initializing: false, + root_epoch: None + } + ); + header_data[24..24 + LUKS_INITIALIZING_LABEL.len()] + .copy_from_slice(LUKS_INITIALIZING_LABEL.as_bytes()); + assert!( + validate_luks2_headers(&mut &header_data[..]) + .expect("Failed to validate LUKS2 header") + .initializing + ); + let header_data = include_bytes!("../../tests/fixtures/luks_header_cipher_null").to_vec(); + let error = validate_luks2_headers(&mut &header_data[..]).unwrap_err(); + assert!(error + .to_string() + .contains("Invalid LUKS keyslot encryption")); +} + +#[test] +fn test_validate_luks2_header_rejects_out_of_range_keyslot_area() { + // Redirect the keyslot binary area below the header region. Same length + // so the surrounding header stays intact; "00768" parses to 768, which is + // inside the header copies (< 2 * hdr_size) rather than the metadata gap. + let mut header = include_bytes!("../../tests/fixtures/luks_header_good").to_vec(); + let patched = patch_luks_json(&mut header, br#""offset":"32768""#, br#""offset":"00768""#); + assert_eq!(patched, 2, "expected to patch both header copies"); + let error = validate_luks2_headers(&mut &header[..]).unwrap_err(); + assert!(error.to_string().contains("Invalid LUKS keyslot area")); +} + +/// Replace `needle` with `replacement` inside every LUKS JSON region, keeping +/// each region the same length by giving back trailing NUL padding. +#[cfg(test)] +fn patch_luks_json(header: &mut [u8], needle: &[u8], replacement: &[u8]) -> usize { + assert!(replacement.len() >= needle.len()); + let grow = replacement.len() - needle.len(); + let mut patched = 0; + let mut i = 0; + while i + needle.len() <= header.len() { + if &header[i..i + needle.len()] != needle { + i += 1; + continue; + } + let tail = &mut header[i..]; + let end = tail.iter().position(|b| *b == 0).expect("NUL padding"); + assert!(end + grow < tail.len(), "not enough NUL padding"); + tail.copy_within(needle.len()..end, replacement.len()); + tail[..replacement.len()].copy_from_slice(replacement); + patched += 1; + i += replacement.len(); + } + patched +} + +#[test] +fn test_disk_key_migration_with_cryptsetup() { + if Command::new("cryptsetup") + .arg("--version") + .output() + .is_err() + { + eprintln!("cryptsetup is unavailable; skipping"); + return; + } + let dir = tempfile::tempdir().unwrap(); + let image = dir.path().join("disk.img"); + fs::File::create(&image) + .unwrap() + .set_len(2 * PAYLOAD_OFFSET) + .unwrap(); + let offset = (PAYLOAD_OFFSET / 512).to_string(); + let format = cryptsetup( + &[ + "luksFormat", + "--batch-mode", + "--type", + "luks2", + "--offset", + &offset, + "--cipher", + "aes-xts-plain64", + "--pbkdf", + "pbkdf2", + "--sector-size", + "512", + "-d-", + ], + &image, + ); + assert_eq!(run_with_input(format, b"old").unwrap(), Some(0)); + let unlocks = |key: &str| { + let test = cryptsetup( + &["open", "--test-passphrase", "--key-slot", "0", "-d-"], + &image, + ); + run_with_input(test, key.as_bytes()).unwrap() == Some(0) + }; + + // Re-key a detached copy of the header, as the boot does, and restore it. + let header = dir.path().join("header"); + let backup = cryptsetup(&["luksHeaderBackup", "--header-backup-file"], &header) + .arg(&image) + .status() + .unwrap(); + assert!(backup.success()); + rekey_luks_header(&header, "old", "new", 2).unwrap(); + assert!(!dir.path().join("old-key").exists()); + assert!( + unlocks("old"), + "the disk is unchanged until the header is restored" + ); + let state = validate_luks2_headers(fs::File::open(&header).unwrap()).unwrap(); + assert_eq!(state.root_epoch, Some(2)); + + let restore = cryptsetup( + &["luksHeaderRestore", "--batch-mode", "--header-backup-file"], + &header, + ) + .arg(&image) + .status() + .unwrap(); + assert!(restore.success()); + assert!(unlocks("new")); + assert!(!unlocks("old")); + + // A later rotation replaces the recorded epoch. + write_luks_root_epoch(&image, 3).unwrap(); + fs::remove_file(&header).unwrap(); + let backup = cryptsetup(&["luksHeaderBackup", "--header-backup-file"], &header) + .arg(&image) + .status() + .unwrap(); + assert!(backup.success()); + let state = validate_luks2_headers(fs::File::open(&header).unwrap()).unwrap(); + assert_eq!(state.root_epoch, Some(3)); +} + +#[test] +fn test_validate_luks2_header_root_epoch_token() { + let with_token = |token: &str| { + let mut header = include_bytes!("../../tests/fixtures/luks_header_good").to_vec(); + let replacement = format!(r#""tokens":{{"0":{token}}}"#); + let patched = patch_luks_json(&mut header, br#""tokens":{}"#, replacement.as_bytes()); + assert_eq!(patched, 2, "expected to patch both header copies"); + validate_luks2_headers(&mut &header[..]) + }; + let state = with_token(r#"{"type":"dstack-kms","keyslots":["0"],"root_epoch":2}"#).unwrap(); + assert_eq!(state.root_epoch, Some(2)); + for token in [ + r#"{"type":"luks2-keyring","keyslots":["0"],"root_epoch":2}"#, + r#"{"type":"dstack-kms","keyslots":["1"],"root_epoch":2}"#, + r#"{"type":"dstack-kms","keyslots":["0"]}"#, + ] { + assert!(with_token(token).is_err(), "{token}"); + } +} + +#[test] +fn test_validate_luks2_header_rejects_keyslot_area_that_overflows() { + let mut header = include_bytes!("../../tests/fixtures/luks_header_good").to_vec(); + // 2**64 - 32768 + 1000: added to the accepted offset 32768 it wraps to + // 1000, which is below PAYLOAD_OFFSET. + let patched = patch_luks_json( + &mut header, + br#""size":"258048""#, + br#""size":"18446744073709519848""#, + ); + assert_eq!(patched, 2, "expected to patch both header copies"); + let error = validate_luks2_headers(&mut &header[..]).unwrap_err(); + assert!( + error.to_string().contains("Invalid LUKS keyslot area"), + "{error:#}" + ); +} From 5c42c4781c9a28d5a0d12dc72c61014eaf856ac2 Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Fri, 6 Nov 2026 04:46:14 -0800 Subject: [PATCH 043/166] guest-agent: derive GetKey keys under a selected KMS root epoch GetKey takes an optional root_epoch (unset = current) and reports the epoch it used, so apps can recover keys derived before a root rotation. The agent fetches an older epoch's app root from the KMS by number on first use and keeps it in memory. The Rust, Go, Python, and JS SDKs expose the selector and the returned epoch. Signed-off-by: Kevin Wang --- dstack/cert-client/src/lib.rs | 22 +++++ .../guest-agent/rpc/proto/agent_rpc_v1.proto | 8 ++ dstack/guest-agent/src/rpc_service.rs | 81 +++++++++++++++++++ dstack/guest-agent/src/rpc_service_v1.rs | 61 +++++++++++--- sdk/go/dstack/client_v1.go | 25 +++++- sdk/js/src/client-v1.ts | 14 +++- sdk/python/src/dstack_sdk/dstack_client_v1.py | 19 ++++- sdk/rust/src/dstack_client_v1.rs | 27 ++++++- sdk/rust/types/src/dstack_v1.rs | 3 + 9 files changed, 243 insertions(+), 17 deletions(-) diff --git a/dstack/cert-client/src/lib.rs b/dstack/cert-client/src/lib.rs index 65498f275..9e222bb04 100644 --- a/dstack/cert-client/src/lib.rs +++ b/dstack/cert-client/src/lib.rs @@ -168,6 +168,28 @@ impl CertRequestClient { } } + /// The app keys of an older KMS root epoch, from the KMS that served the + /// current ones. + pub async fn get_app_key(&self, root_epoch: u32) -> Result { + let CertRequestClient::Kms { client, vm_config } = self else { + bail!("only KMS-provided keys have older root epochs"); + }; + let response = client + .get_app_key(GetAppKeyRequest { + api_version: 1, + vm_config: vm_config.clone(), + root_epoch: Some(root_epoch), + }) + .await?; + if response.root_epoch != root_epoch { + bail!( + "KMS returned keys of root epoch {} instead of {root_epoch}", + response.root_epoch + ); + } + Ok(response) + } + pub async fn get_root_ca(&self) -> Result { match self { CertRequestClient::Local { ca } => Ok(ca.pem_cert.clone()), diff --git a/dstack/guest-agent/rpc/proto/agent_rpc_v1.proto b/dstack/guest-agent/rpc/proto/agent_rpc_v1.proto index 3df408ad8..bebc85f99 100644 --- a/dstack/guest-agent/rpc/proto/agent_rpc_v1.proto +++ b/dstack/guest-agent/rpc/proto/agent_rpc_v1.proto @@ -189,6 +189,11 @@ message GetKeyRequest { // signing method: v0's `secp256k1_prehashed` was a mode wearing an // algorithm's name, and it does not appear here. string algorithm = 2; + // KMS root key epoch to derive from. Unset means the current epoch, which is + // what new keys should use; epoch 0 is the original root keys. Name an older + // 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; } message GetKeyResponse { @@ -209,6 +214,9 @@ message GetKeyResponse { // unversioned surface returns. Link [0] is not: its claim encoding is new, // and it is deliberately not one the v0 claim format can produce. repeated bytes signature_chain = 3; + // KMS root key epoch the key was derived from. Link [1] is signed by that + // epoch's KMS root key. + uint32 root_epoch = 4; } // --------------------------------------------------------------------------- diff --git a/dstack/guest-agent/src/rpc_service.rs b/dstack/guest-agent/src/rpc_service.rs index a6dcd8e5f..068d8b53d 100644 --- a/dstack/guest-agent/src/rpc_service.rs +++ b/dstack/guest-agent/src/rpc_service.rs @@ -2,6 +2,7 @@ // // SPDX-License-Identifier: Apache-2.0 +use std::collections::HashMap; use std::sync::{Arc, Mutex}; use std::time::{Duration, Instant}; @@ -92,11 +93,35 @@ pub struct AppState { inner: Arc, } +/// The app root key of one KMS root epoch, with its KMS signature. +#[derive(Clone)] +pub(crate) struct AppRoot { + pub 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 { + let signing_key = SigningKey::from_slice(&k256_key) + .context("the app root k256 key is not a valid secp256k1 scalar")?; + Ok(Self { + epoch, + k256_key, + signing_key, + kms_signature, + }) + } +} + struct AppStateInner { config: Config, keys: AppKeys, vm_config: String, cert_client: CertRequestClient, + /// Held across a fetch so concurrent requests for one epoch fetch it once. + 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>, @@ -300,6 +325,7 @@ impl AppState { config, keys, cert_client, + historical_roots: Default::default(), demo_cert: Default::default(), demo_cert_last_failure: Default::default(), vm_config, @@ -377,6 +403,58 @@ impl AppState { Ok(key) } + /// KMS root epoch of the app keys this agent serves. + pub(crate) fn root_epoch(&self) -> u32 { + self.inner.keys.root_epoch() + } + + /// 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 { + let current = self.root_epoch(); + let epoch = root_epoch.unwrap_or(current); + if epoch == current { + return Ok(AppRoot { + 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(), + }); + } + if epoch > current { + anyhow::bail!("root epoch {epoch} is newer than the current one, {current}"); + } + self.historical_root(epoch).await + } + + pub(crate) async fn historical_root(&self, epoch: u32) -> Result { + let mut roots = self.inner.historical_roots.lock().await; + if let Some(root) = roots.get(&epoch) { + return Ok(root.clone()); + } + let keys = self + .inner + .cert_client + .get_app_key(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)?; + roots.insert(epoch, root.clone()); + Ok(root) + } + + #[cfg(test)] + pub(crate) async fn insert_historical_root( + &self, + epoch: u32, + 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); + } + /// The KMS root key's signature over the app root public key: the second /// link of every signature chain, produced outside this agent and passed /// through byte-for-byte on both API surfaces. @@ -1152,6 +1230,8 @@ pub(crate) mod tests { .expect("certificate carries an attestation") } + pub(crate) const HISTORICAL_K256_KEY: [u8; 32] = [0x11; 32]; + pub(crate) async fn setup_test_state() -> (AppState, tempfile::NamedTempFile) { setup_test_state_with_platform(None).await } @@ -1442,6 +1522,7 @@ pNs85uhOZE8z2jr8Pg== keys: dummy_keys, vm_config: String::new(), cert_client: dummy_cert_client, + historical_roots: Default::default(), demo_cert: Default::default(), demo_cert_last_failure: Default::default(), platform: Arc::new(TestSimulatorPlatform { diff --git a/dstack/guest-agent/src/rpc_service_v1.rs b/dstack/guest-agent/src/rpc_service_v1.rs index 354e94618..8f178ccbf 100644 --- a/dstack/guest-agent/src/rpc_service_v1.rs +++ b/dstack/guest-agent/src/rpc_service_v1.rs @@ -159,16 +159,19 @@ impl V1RpcHandler { Self { state } } - /// Derive the key named by `(domain, algorithm)` together with its - /// signature chain. - fn derive(&self, domain: &str, algorithm: &str) -> Result<(AppKey, Vec>)> { + /// Derive the key named by `(domain, algorithm)` under the app root of + /// `root_epoch`, together with its signature chain and that epoch. + async fn derive( + &self, + domain: &str, + algorithm: &str, + root_epoch: Option, + ) -> Result<(AppKey, Vec>, u32)> { let algorithm = Algorithm::parse(algorithm)?; - let key = AppKey::derive(self.state.app_root_k256_key(), domain, algorithm)?; - let chain = vec![ - key.claim_signature(self.state.app_root_signing_key()?)?, - self.state.kms_k256_signature().to_vec(), - ]; - Ok((key, chain)) + let root = self.state.app_root(root_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)) } } @@ -196,11 +199,14 @@ impl DstackGuestRpc for V1RpcHandler { } async fn get_key(self, request: GetKeyRequest) -> Result { - let (key, signature_chain) = self.derive(&request.domain, &request.algorithm)?; + let (key, signature_chain, root_epoch) = self + .derive(&request.domain, &request.algorithm, request.root_epoch) + .await?; Ok(GetKeyResponse { key: key.secret(), public_key: key.public_key(), signature_chain, + root_epoch, }) } @@ -342,9 +348,44 @@ mod tests { GetKeyRequest { domain: domain.to_string(), algorithm: algorithm.to_string(), + root_epoch: None, } } + #[tokio::test] + async fn get_key_selects_the_root_epoch() { + let (state, _guard) = state().await; + let handler = || V1RpcHandler::new(state.clone()); + let request = |root_epoch| GetKeyRequest { + root_epoch, + ..get_key_request("storage-encryption", "secp256k1") + }; + + let current = handler().get_key(request(None)).await.unwrap(); + assert_eq!(current.root_epoch, 0); + assert_eq!( + handler().get_key(request(Some(0))).await.unwrap().key, + current.key + ); + assert!(handler().get_key(request(Some(1))).await.is_err()); + } + + /// Older epochs come from the KMS once and are then served from memory; + /// the fixture has no KMS, so the cache is filled directly. + #[tokio::test] + async fn an_older_root_is_fetched_once_and_cached() { + let (state, _guard) = state().await; + let key = crate::rpc_service::tests::HISTORICAL_K256_KEY.to_vec(); + state + .insert_historical_root(1, key.clone(), vec![0xab; 65]) + .await; + let root = state.historical_root(1).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()); + } + /// The fixture's app root key is what `keys::tests` committed its vectors /// against. If the two drift, the vectors stop describing what this /// handler serves and every assertion built on them goes quiet. diff --git a/sdk/go/dstack/client_v1.go b/sdk/go/dstack/client_v1.go index fd0f98ace..5b27c29da 100644 --- a/sdk/go/dstack/client_v1.go +++ b/sdk/go/dstack/client_v1.go @@ -42,6 +42,8 @@ type GetKeyV1Response struct { // key over the app root public key. See `docs/guest-api-v1.md` for the claim // encoding and the verification steps. SignatureChain [][]byte + // KMS root key epoch the key was derived from. + RootEpoch uint32 } // Response from a v1 attestation request. @@ -442,6 +444,18 @@ 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) +} + +// GetKeyAtRootEpoch is GetKey under the KMS root key epoch rootEpoch instead of +// the current one; epoch 0 is the original root keys. Name an older epoch only +// 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) +} + +func (c *DstackClientV1) getKey(ctx context.Context, domain string, algorithm string, rootEpoch *uint32) (*GetKeyV1Response, error) { if algorithm == "" { return nil, fmt.Errorf("algorithm is required, use `secp256k1` or `ed25519`") } @@ -450,6 +464,9 @@ func (c *DstackClientV1) GetKey(ctx context.Context, domain string, algorithm st "domain": domain, "algorithm": algorithm, } + if rootEpoch != nil { + payload["root_epoch"] = *rootEpoch + } data, err := c.sendRPCRequest(ctx, "/v1/GetKey", payload) if err != nil { @@ -464,6 +481,7 @@ func (c *DstackClientV1) GetKey(ctx context.Context, domain string, algorithm st Key *string `json:"key"` PublicKey *string `json:"public_key"` SignatureChain *[]*string `json:"signature_chain"` + RootEpoch uint32 `json:"root_epoch"` } if err := json.Unmarshal(data, &response); err != nil { return nil, err @@ -489,7 +507,12 @@ func (c *DstackClientV1) GetKey(ctx context.Context, domain string, algorithm st } } - return &GetKeyV1Response{Key: key, PublicKey: publicKey, SignatureChain: chain}, nil + return &GetKeyV1Response{ + Key: key, + PublicKey: publicKey, + SignatureChain: chain, + RootEpoch: response.RootEpoch, + }, nil } // Attest produces a versioned attestation over reportData, which must be diff --git a/sdk/js/src/client-v1.ts b/sdk/js/src/client-v1.ts index d35de8590..fbf3156f5 100644 --- a/sdk/js/src/client-v1.ts +++ b/sdk/js/src/client-v1.ts @@ -68,6 +68,8 @@ export interface GetKeyResponseV1 { public_key: Uint8Array /** Two links: the app root key over the v1 key claim, then the KMS root key. */ signature_chain: Uint8Array[] + /** KMS root key epoch the key was derived from. */ + root_epoch: number } export interface AttestResponseV1 { @@ -333,14 +335,19 @@ 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): Promise { + async getKey(domain: string, algorithm: string, rootEpoch?: 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) { throw new Error('algorithm is required, use "secp256k1" or "ed25519"') } - const payload = JSON.stringify({ domain, algorithm }) - const result = await send_rpc_request<{ key: string, public_key: string, signature_chain: string[] }>( + // Unset is the current KMS root key epoch; name an older one only to migrate off it. + if (rootEpoch !== undefined && (!Number.isInteger(rootEpoch) || rootEpoch < 0)) { + throw new Error('rootEpoch must be a non-negative epoch number') + } + 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 }>( this.endpoint, '/v1/GetKey', payload) throwOnRpcError(result) return Object.freeze({ @@ -348,6 +355,7 @@ export class DstackClientV1 { 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, __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 e57c8bba5..77c1438d0 100644 --- a/sdk/python/src/dstack_sdk/dstack_client_v1.py +++ b/sdk/python/src/dstack_sdk/dstack_client_v1.py @@ -107,6 +107,8 @@ class GetKeyResponseV1(BaseModel): # Two links: the app root key's signature over the v1 key claim, then the # KMS root key's signature over the app root public key. signature_chain: List[HexBytes] + # KMS root key epoch the key was derived from. + root_epoch: int = 0 class GpuEvidenceBundleV1(BaseModel): @@ -237,7 +239,9 @@ async def issue_cert( result = await self._send_rpc_request("IssueCert", data) return IssueCertResponseV1(**result) - async def get_key(self, domain: str, algorithm: str) -> GetKeyResponseV1: + async def get_key( + self, domain: str, algorithm: str, root_epoch: Optional[int] = None + ) -> GetKeyResponseV1: """Derive an application key from ``(domain, algorithm)``. Both arguments are required. Derivation is flat: two domains yield @@ -245,10 +249,19 @@ async def get_key(self, domain: str, algorithm: str) -> GetKeyResponseV1: exactly ``secp256k1`` or ``ed25519`` -- there is no default and no ``k256`` alias, because in v0 a typo silently produced a key of the wrong type under a name the caller thought meant something else. + + ``root_epoch`` selects the KMS root key epoch; the default ``None`` is + the current epoch, and epoch 0 is the original root keys. Name an older + epoch only to reach a key from before a KMS root key rotation, for + example to migrate off it. """ if not algorithm: raise ValueError("algorithm is required, use `secp256k1` or `ed25519`") data: Dict[str, Any] = {"domain": domain, "algorithm": algorithm} + if root_epoch is not None: + if root_epoch < 0: + raise ValueError("root_epoch must be a non-negative epoch number") + data["root_epoch"] = root_epoch result = await self._send_rpc_request("GetKey", data) return GetKeyResponseV1(**result) @@ -350,7 +363,9 @@ def issue_cert( raise NotImplementedError @call_async - def get_key(self, domain: str, algorithm: str) -> GetKeyResponseV1: + def get_key( + self, domain: str, algorithm: str, root_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 efcc78236..71f9c5f33 100644 --- a/sdk/rust/src/dstack_client_v1.rs +++ b/sdk/rust/src/dstack_client_v1.rs @@ -95,10 +95,35 @@ 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::get_key`] under the KMS root key epoch `root_epoch` instead of + /// the current one; epoch 0 is the original root keys. Name an older epoch + /// only to reach a key from before a KMS root key rotation, for example to + /// migrate off it. + pub async fn get_key_at_root_epoch( + &self, + domain: &str, + algorithm: &str, + root_epoch: u32, + ) -> Result { + self.request_key(domain, algorithm, Some(root_epoch)).await + } + + async fn request_key( + &self, + domain: &str, + algorithm: &str, + root_epoch: Option, + ) -> Result { if algorithm.is_empty() { anyhow::bail!("algorithm is required, use `secp256k1` or `ed25519`") } - let data = json!({ "domain": domain, "algorithm": algorithm }); + let mut data = json!({ "domain": domain, "algorithm": algorithm }); + if let Some(root_epoch) = root_epoch { + data["root_epoch"] = root_epoch.into(); + } let response = self.send_rpc_request("GetKey", &data).await?; Ok(serde_json::from_value::(response)?) } diff --git a/sdk/rust/types/src/dstack_v1.rs b/sdk/rust/types/src/dstack_v1.rs index 917db972e..3812b2e4f 100644 --- a/sdk/rust/types/src/dstack_v1.rs +++ b/sdk/rust/types/src/dstack_v1.rs @@ -125,6 +125,9 @@ pub struct GetKeyResponse { /// steps. Verifying is the relying party's job; this SDK does not do it. #[serde(with = "hex_vec")] pub signature_chain: Vec>, + /// KMS root key epoch the key was derived from. + #[serde(default)] + pub root_epoch: u32, } /// The value an attestation binds to: bytes, and only bytes. From 839785d7d0cbfa623d60f71cf2d72d09ce4beec6 Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Sun, 8 Nov 2026 01:08:07 -0800 Subject: [PATCH 044/166] fix(guest): release epoch cache lock during key fetch Signed-off-by: Kevin Wang --- dstack/guest-agent/src/rpc_service.rs | 16 +++++++++++----- 1 file changed, 11 insertions(+), 5 deletions(-) diff --git a/dstack/guest-agent/src/rpc_service.rs b/dstack/guest-agent/src/rpc_service.rs index 068d8b53d..d79a9ed22 100644 --- a/dstack/guest-agent/src/rpc_service.rs +++ b/dstack/guest-agent/src/rpc_service.rs @@ -429,9 +429,15 @@ impl AppState { } pub(crate) async fn historical_root(&self, epoch: u32) -> Result { - let mut roots = self.inner.historical_roots.lock().await; - if let Some(root) = roots.get(&epoch) { - return Ok(root.clone()); + if let Some(root) = self + .inner + .historical_roots + .lock() + .await + .get(&epoch) + .cloned() + { + return Ok(root); } let keys = self .inner @@ -440,8 +446,8 @@ impl AppState { .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)?; - roots.insert(epoch, root.clone()); - Ok(root) + let mut roots = self.inner.historical_roots.lock().await; + Ok(roots.entry(epoch).or_insert_with(|| root.clone()).clone()) } #[cfg(test)] From 9b236bf1b27d813a481f19a11ef9baf3ef191961 Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Sun, 8 Nov 2026 01:08:07 -0800 Subject: [PATCH 045/166] fix(sdk): validate root epochs as uint32 Signed-off-by: Kevin Wang --- sdk/js/src/client-v1.ts | 4 ++-- sdk/python/src/dstack_sdk/dstack_client_v1.py | 4 ++-- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/sdk/js/src/client-v1.ts b/sdk/js/src/client-v1.ts index fbf3156f5..9d0f94b15 100644 --- a/sdk/js/src/client-v1.ts +++ b/sdk/js/src/client-v1.ts @@ -342,8 +342,8 @@ export class DstackClientV1 { throw new Error('algorithm is required, use "secp256k1" or "ed25519"') } // Unset is the current KMS root key epoch; name an older one only to migrate off it. - if (rootEpoch !== undefined && (!Number.isInteger(rootEpoch) || rootEpoch < 0)) { - throw new Error('rootEpoch must be a non-negative epoch number') + 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 }) diff --git a/sdk/python/src/dstack_sdk/dstack_client_v1.py b/sdk/python/src/dstack_sdk/dstack_client_v1.py index 77c1438d0..750b0b942 100644 --- a/sdk/python/src/dstack_sdk/dstack_client_v1.py +++ b/sdk/python/src/dstack_sdk/dstack_client_v1.py @@ -259,8 +259,8 @@ async def get_key( raise ValueError("algorithm is required, use `secp256k1` or `ed25519`") data: Dict[str, Any] = {"domain": domain, "algorithm": algorithm} if root_epoch is not None: - if root_epoch < 0: - raise ValueError("root_epoch must be a non-negative epoch number") + 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 result = await self._send_rpc_request("GetKey", data) return GetKeyResponseV1(**result) From 067897d2839ff0b3500be1d178b94c32b765a7ec Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Sun, 8 Nov 2026 06:09:57 -0800 Subject: [PATCH 046/166] test(sdk): initialize root epoch in key response fixture Signed-off-by: Kevin Wang --- sdk/rust/types/src/dstack_v1.rs | 1 + 1 file changed, 1 insertion(+) diff --git a/sdk/rust/types/src/dstack_v1.rs b/sdk/rust/types/src/dstack_v1.rs index 3812b2e4f..1ef1d7fd1 100644 --- a/sdk/rust/types/src/dstack_v1.rs +++ b/sdk/rust/types/src/dstack_v1.rs @@ -351,6 +351,7 @@ mod tests { key: vec![0x01; 32], public_key: vec![0x02; 33], signature_chain: vec![vec![0xaa, 0xbb], vec![0xcc]], + root_epoch: 0, }; let json = serde_json::to_string(&response).expect("serializes"); assert!(json.contains(r#""signature_chain":["aabb","cc"]"#)); From 0e6ff360c04f9ecacbcdaf78b7906a41196a900d Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Sat, 7 Nov 2026 00:05:46 -0800 Subject: [PATCH 047/166] feat: seal environment variables with epoch-aware HPKE Signed-off-by: Kevin Wang --- docs/deployment.md | 5 +- docs/sealed-env.md | 98 ++++++ dstack/Cargo.lock | 202 ++++++++++-- dstack/Cargo.toml | 1 + dstack/dstack-types/src/lib.rs | 2 + dstack/dstack-types/src/shared_filenames.rs | 1 + dstack/dstack-util/Cargo.toml | 1 + dstack/dstack-util/src/crypto.rs | 130 ++++++++ dstack/dstack-util/src/host_api.rs | 18 ++ dstack/dstack-util/src/main.rs | 3 + dstack/dstack-util/src/system_setup.rs | 306 ++++++++++++++++-- dstack/guest-agent/src/rpc_service.rs | 1 + dstack/host-api/proto/host_api.proto | 9 + dstack/kms/Cargo.toml | 1 + dstack/kms/rpc/proto/kms_rpc.proto | 11 + dstack/kms/src/main_service.rs | 72 ++++- dstack/vmm/requirements.txt | 3 +- dstack/vmm/rpc/proto/vmm_rpc.proto | 13 + dstack/vmm/src/app.rs | 173 +++++++++- dstack/vmm/src/app/vm_info.rs | 1 + dstack/vmm/src/app/workdir.rs | 69 +++- dstack/vmm/src/config.rs | 13 + dstack/vmm/src/host_api_service.rs | 6 + dstack/vmm/src/main_service.rs | 31 +- dstack/vmm/src/one_shot.rs | 5 +- dstack/vmm/src/tests/test_vmm_cli.py | 264 +++++++++++++++ dstack/vmm/src/vmm-cli.py | 138 ++++++-- dstack/vmm/ui/build.mjs | 4 + dstack/vmm/ui/package-lock.json | 45 +++ dstack/vmm/ui/package.json | 7 +- dstack/vmm/ui/src/composables/useVmManager.ts | 19 +- dstack/vmm/ui/src/lib/browserCrypto.ts | 6 + dstack/vmm/ui/src/lib/envEncryption.ts | 68 ++++ dstack/vmm/ui/src/lib/vmmRpcClient.ts | 9 +- dstack/vmm/ui/tests/envEncryption.test.cjs | 215 ++++++++++++ dstack/vmm/vmm.toml | 4 + sdk/fixtures/sealed-env/javascript.json | 17 + sdk/fixtures/sealed-env/python.json | 17 + sdk/go/README.md | 42 +++ sdk/go/dstack/seal_env.go | 64 ++++ sdk/go/dstack/seal_env_test.go | 115 +++++++ sdk/go/dstack/verify_sealed_env_test.go | 123 +++++++ sdk/go/dstack/verify_signature.go | 77 ++++- sdk/go/go.mod | 1 + sdk/go/go.sum | 2 + sdk/js/README.md | 22 ++ sdk/js/bun.lock | 8 + sdk/js/package.json | 10 + sdk/js/src/__tests__/sealed-env.test.ts | 51 +++ sdk/js/src/index.ts | 4 +- sdk/js/src/seal-env.ts | 46 +++ sdk/js/src/verify-env-encrypt-public-key.ts | 32 ++ sdk/js/tsup.config.ts | 1 + sdk/python/README.md | 11 + sdk/python/pdm.lock | 18 +- sdk/python/pyproject.toml | 1 + sdk/python/src/dstack_sdk/__init__.py | 4 + sdk/python/src/dstack_sdk/seal_env.py | 34 ++ .../verify_env_encrypt_public_key.py | 31 ++ sdk/python/tests/test_seal_env.py | 63 ++++ sdk/rust/Cargo.lock | 152 ++++++++- sdk/rust/Cargo.toml | 4 + sdk/rust/README.md | 23 ++ sdk/rust/src/lib.rs | 5 + sdk/rust/src/sealed_env.rs | 240 ++++++++++++++ 65 files changed, 3051 insertions(+), 121 deletions(-) create mode 100644 docs/sealed-env.md create mode 100644 dstack/vmm/ui/src/lib/browserCrypto.ts create mode 100644 dstack/vmm/ui/src/lib/envEncryption.ts create mode 100644 dstack/vmm/ui/tests/envEncryption.test.cjs create mode 100644 sdk/fixtures/sealed-env/javascript.json create mode 100644 sdk/fixtures/sealed-env/python.json create mode 100644 sdk/go/dstack/seal_env.go create mode 100644 sdk/go/dstack/seal_env_test.go create mode 100644 sdk/go/dstack/verify_sealed_env_test.go create mode 100644 sdk/js/src/__tests__/sealed-env.test.ts create mode 100644 sdk/js/src/seal-env.ts create mode 100644 sdk/python/src/dstack_sdk/seal_env.py create mode 100644 sdk/python/tests/test_seal_env.py create mode 100644 sdk/rust/src/sealed_env.rs diff --git a/docs/deployment.md b/docs/deployment.md index d68fdea44..7917caac8 100644 --- a/docs/deployment.md +++ b/docs/deployment.md @@ -552,7 +552,10 @@ for the full procedure; the KMS steps are: `PopRootEpoch` (or remove the `rootEpochs` entry) and rotate again; nodes that imported the dropped N must be onboarded again. 6. Boot CVMs on an OS that supports root epochs. At boot they re-key their data - disk from the epoch recorded in its LUKS header to the current one. + disk from the epoch recorded in its LUKS header to the current one. If an + encrypted environment still uses an older epoch, the CVM re-encrypts it during + boot and the VMM replaces the ciphertext only if it was not changed concurrently. + See [SealedEnv](./sealed-env.md) for the new format, file precedence, and rollback requirements. Each `authorizeKms` or `authorizeApp` result carries the current epoch and its chain hash, including when authorization is denied. KMS verifies these against its keyring; diff --git a/docs/sealed-env.md b/docs/sealed-env.md new file mode 100644 index 000000000..0f02709bb --- /dev/null +++ b/docs/sealed-env.md @@ -0,0 +1,98 @@ +# Sealed environment variables + +SealedEnv protects environment variables with HPKE and identifies the KMS root +key epoch needed to decrypt them. Plaintext remains inside the client and CVM; +the VMM stores and forwards ciphertext. + +## Wire format + +The `.sealed-env` file contains UTF-8 JSON: + +```json +{"version":1,"root_epoch":3,"enc":"","ciphertext":""} +``` + +Version 1 fixes RFC 9180 Base mode with DHKEM(X25519, HKDF-SHA256), +HKDF-SHA256, and AES-256-GCM. `enc` is a 32-byte encapsulated key; `ciphertext` +includes the 16-byte authentication tag. Binary fields use standard padded +Base64. HPKE `info` is UTF-8 `dstack-sealed-env-v1:N`, where `N` is the canonical +decimal `root_epoch` uint32, with no leading zeros. AAD is empty. Unknown envelope +versions are rejected; extra JSON fields may be ignored. This provides +confidentiality, not authentication of the person uploading the environment. + +Plaintext uses the existing JSON representation: + +```json +{"env":[{"key":"MESSAGE","value":"Hello"}]} +``` + +The guest still enforces the application's environment-name allowlist. + +## Keys and signatures + +`GetSealedEnvPubKey` is available on KMS and proxied by VMM. Its request contains +`app_id` (20 bytes) and optional `root_epoch`; omission selects the current +active epoch, while future epochs are rejected. The response contains +`public_key`, `root_epoch`, `timestamp`, and `signature`. + +For each app, the KMS derives the 32-byte receiver private key directly from +that epoch's P-256 CA private scalar using HKDF-SHA256: + +- salt: ASCII `dstack-sealed-env-v1`; +- IKM: the 32-byte P-256 private scalar; +- info: the raw 20-byte app ID. + +`GetAppKey` returns it as `sealed_env_key` alongside the legacy key. The legacy +key derivation and `GetAppEnvEncryptPubKey` remain unchanged. + +The public-key response signs Keccak256 of this unambiguous byte sequence: + +```text +"dstack-sealed-env-pubkey-v1:" || app_id[20] || epoch[4, big-endian] + || public_key[32] || timestamp[8, big-endian] +``` + +The signing key is the selected epoch's secp256k1 root. The signature is +`r[32] || s[32] || recovery_id[1]`. Clients must compare the recovered signer +against an independently trusted epoch root key or pin. Recovering a signer +alone is not a trust decision. Clients should reject stale responses (default +maximum age 300 seconds, future clock tolerance 60 seconds). + +## File precedence and migration + +A new guest chooses `.sealed-env` whenever that file exists. An empty, invalid, +or unreadable sealed file is an error; it must not fall back to a legacy file. +If no sealed file exists, the guest accepts `.encrypted-env` for compatibility. +The copy limits are 384 KiB for sealed JSON and 256 KiB for legacy ciphertext. + +At boot, sealed ciphertext explicitly selects its historical root epoch. The +guest obtains that epoch's app key over the already trusted KMS connection, +decrypts, and seals to the current epoch. Legacy ciphertext uses the legacy +key path, trying earlier epochs when necessary, and is converted to SealedEnv. +A future epoch or failed authentication stops environment provisioning. + +The guest sends `UpdateSealedEnv` through the host channel before application +code runs. The VMM binds the update to the VM's CID and a snapshot of both env +files, including presence. It accepts at most one update in the boot window; +an empty update closes the window without replacing files. Owner updates and +guest replacement share the VM launch lock so an owner change cannot be lost +between the snapshot comparison and write. + +VM creation and update accept either `sealed_env` or legacy `encrypted_env`. +Supplying both nonempty fields is rejected. Writing one format removes the other: +an external legacy update after a guest sealed update therefore becomes the +sole environment for the next boot. Omitting both preserves existing files on +update. A replacement is safely written before the other file is removed; +filesystem errors are reported rather than ignored. + +Old OS images understand only `.encrypted-env`. After conversion removes that +file, reverting to an old OS requires the owner to upload legacy ciphertext +again. Rotation does not revoke credentials exposed before migration, and +cannot erase old host snapshots. + +## Scope of verification + +Local tests exercise HPKE implementations with synthetic keys, signature +contexts, file precedence and update races. They do not establish production +TEE isolation or real-chain availability. The shared fixture under +`sdk/fixtures/sealed-env/` is public test data, never deployment key material. diff --git a/dstack/Cargo.lock b/dstack/Cargo.lock index 7f86192db..cbdbc35f2 100644 --- a/dstack/Cargo.lock +++ b/dstack/Cargo.lock @@ -27,6 +27,16 @@ dependencies = [ "generic-array", ] +[[package]] +name = "aead" +version = "0.6.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1973cfbc1a2daf9cf550e74e1f088c28e7f7d8c1e1418fb6c9dc5184b7e84c99" +dependencies = [ + "crypto-common 0.2.2", + "inout 0.2.2", +] + [[package]] name = "aes" version = "0.8.4" @@ -34,24 +44,49 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b169f7a6d4742236a0a00c541b845991d0ac43e546831af1249753ab4c3aa3a0" dependencies = [ "cfg-if", - "cipher", + "cipher 0.4.4", "cpufeatures 0.2.17", ] +[[package]] +name = "aes" +version = "0.9.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "35f0f96ce78e38c3dc6d8948aa8163d06385be74000f3c7a95bf1eef35d3ea32" +dependencies = [ + "cipher 0.5.2", + "cpubits", + "cpufeatures 0.3.0", +] + [[package]] name = "aes-gcm" version = "0.10.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "831010a0f742e1209b3bcea8fab6a8e149051ba6099432c8cb2cc117dec3ead1" dependencies = [ - "aead", - "aes", - "cipher", - "ctr", - "ghash", + "aead 0.5.2", + "aes 0.8.4", + "cipher 0.4.4", + "ctr 0.9.2", + "ghash 0.5.1", "subtle", ] +[[package]] +name = "aes-gcm" +version = "0.11.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7f2b8006a0c83f52b62ba44a97b58bf76fe2f70a329e588f67f89691d93d498f" +dependencies = [ + "aead 0.6.1", + "aes 0.9.3", + "cipher 0.5.2", + "ctr 0.10.1", + "ctutils", + "ghash 0.6.0", +] + [[package]] name = "ahash" version = "0.8.12" @@ -625,7 +660,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e412e2cd0f2b2d93e02543ceae7917b3c70331573df19ee046bcbc35e45e87d7" dependencies = [ "byteorder", - "cipher", + "cipher 0.4.4", ] [[package]] @@ -969,10 +1004,21 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "773f3b9af64447d2ce9850330c473515014aa235e6a783b02db81ff39e4a3dad" dependencies = [ "crypto-common 0.1.7", - "inout", + "inout 0.1.4", "zeroize", ] +[[package]] +name = "cipher" +version = "0.5.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e8cf2a2c93cd704877c0858356ed03480ff301ee950b43f1cbe4573b088bfa6c" +dependencies = [ + "block-buffer 0.12.1", + "crypto-common 0.2.2", + "inout 0.2.2", +] + [[package]] name = "clap" version = "4.6.1" @@ -1348,7 +1394,16 @@ version = "0.9.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "0369ee1ad671834580515889b80f2ea915f23b8be8d0daa4bbaf2ac5c7590835" dependencies = [ - "cipher", + "cipher 0.4.4", +] + +[[package]] +name = "ctr" +version = "0.10.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "baaca1c4b237092596f64d571e9db6ce4109c4ef9742e27590f1709594461f21" +dependencies = [ + "cipher 0.5.2", ] [[package]] @@ -1377,6 +1432,21 @@ dependencies = [ "zeroize", ] +[[package]] +name = "curve25519-dalek" +version = "5.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b5eed333089e2e1c1ac8c6c0398e5e2497b4c9926ca6d0365ed1e099afa5bc23" +dependencies = [ + "cfg-if", + "cpufeatures 0.3.0", + "curve25519-dalek-derive", + "fiat-crypto 0.3.0", + "rustc_version", + "subtle", + "zeroize", +] + [[package]] name = "curve25519-dalek-derive" version = "0.1.1" @@ -2145,6 +2215,7 @@ dependencies = [ "k256", "load_config", "moka", + "p256 0.13.2", "parity-scale-codec", "ra-rpc", "ra-tls", @@ -2163,7 +2234,7 @@ dependencies = [ "tokio", "tracing", "tracing-subscriber", - "x25519-dalek", + "x25519-dalek 2.0.1", "x509-parser 0.16.0", "yasna 0.5.2", ] @@ -2321,7 +2392,7 @@ dependencies = [ name = "dstack-util" version = "0.6.3" dependencies = [ - "aes-gcm", + "aes-gcm 0.10.3", "anyhow", "base64 0.22.1", "binrw", @@ -2330,7 +2401,7 @@ dependencies = [ "cert-client", "clap", "cmd_lib", - "curve25519-dalek", + "curve25519-dalek 4.1.3", "dcap-qvl", "dstack-api-auth", "dstack-attest", @@ -2345,6 +2416,7 @@ dependencies = [ "hex", "hex_fmt", "host-api", + "hpke", "k256", "libc", "lspci", @@ -2377,7 +2449,7 @@ dependencies = [ "tracing", "tracing-subscriber", "url", - "x25519-dalek", + "x25519-dalek 2.0.1", "x509-parser 0.16.0", "yaml-rust2", ] @@ -2589,7 +2661,7 @@ version = "2.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "70e796c081cee67dc755e1a36a0a172b897fab85fc3f6bc48307991f64e4eca9" dependencies = [ - "curve25519-dalek", + "curve25519-dalek 4.1.3", "ed25519", "rand_core 0.6.4", "serde", @@ -2616,7 +2688,7 @@ dependencies = [ "ff 0.13.1", "generic-array", "group 0.13.0", - "hkdf", + "hkdf 0.12.4", "pem-rfc7468", "pkcs8", "rand_core 0.6.4", @@ -3177,7 +3249,16 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f0d8a4362ccb29cb0b265253fb0a2728f592895ee6854fd9bc13f2ffda266ff1" dependencies = [ "opaque-debug", - "polyval", + "polyval 0.6.2", +] + +[[package]] +name = "ghash" +version = "0.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2eecf2d5dc9b66b732b97707a0210906b1d30523eb773193ab777c0c84b3e8d5" +dependencies = [ + "polyval 0.7.3", ] [[package]] @@ -3485,6 +3566,15 @@ dependencies = [ "hmac 0.12.1", ] +[[package]] +name = "hkdf" +version = "0.13.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4aaa26c720c68b866f2c96ef5c1264b3e6f473fe5d4ce61cd44bbe913e553018" +dependencies = [ + "hmac 0.13.0", +] + [[package]] name = "hmac" version = "0.12.1" @@ -3525,6 +3615,25 @@ dependencies = [ "serde_json", ] +[[package]] +name = "hpke" +version = "0.15.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8f512e325a1a71bf3273fb9b7fc9e5d96dac136fb4a95d0e6637001b12ea3e09" +dependencies = [ + "aead 0.6.1", + "aes-gcm 0.11.1", + "curve25519-dalek 5.0.0", + "getrandom 0.4.2", + "hkdf 0.13.0", + "hybrid-array", + "rand_core 0.10.1", + "sha2 0.11.1", + "subtle", + "x25519-dalek 3.0.0", + "zeroize", +] + [[package]] name = "http" version = "1.4.1" @@ -3916,6 +4025,15 @@ dependencies = [ "generic-array", ] +[[package]] +name = "inout" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4250ce6452e92010fdf7268ccc5d14faa80bb12fc741938534c58f16804e03c7" +dependencies = [ + "hybrid-array", +] + [[package]] name = "insta" version = "1.47.2" @@ -4376,7 +4494,7 @@ dependencies = [ "tokio", "tracing", "tracing-subscriber", - "x25519-dalek", + "x25519-dalek 2.0.1", "xsalsa20poly1305", ] @@ -5474,7 +5592,7 @@ checksum = "8159bd90725d2df49889a078b54f4f79e87f1f8a8444194cdca81d38f5393abf" dependencies = [ "cpufeatures 0.2.17", "opaque-debug", - "universal-hash", + "universal-hash 0.5.1", ] [[package]] @@ -5486,7 +5604,18 @@ dependencies = [ "cfg-if", "cpufeatures 0.2.17", "opaque-debug", - "universal-hash", + "universal-hash 0.5.1", +] + +[[package]] +name = "polyval" +version = "0.7.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f0fa31d631f2b2cb2a544d0aa321ce847a94764d701ca2becc411138b93d49cd" +dependencies = [ + "cpubits", + "cpufeatures 0.3.0", + "universal-hash 0.6.1", ] [[package]] @@ -5946,7 +6075,7 @@ dependencies = [ "fs-err", "hex", "hex_fmt", - "hkdf", + "hkdf 0.12.4", "k256", "or-panic", "p256 0.13.2", @@ -6811,7 +6940,7 @@ version = "0.10.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "97a22f5af31f73a954c10289c93e8a50cc23d971e80ee446f1f6f7137a088213" dependencies = [ - "cipher", + "cipher 0.4.4", ] [[package]] @@ -6887,10 +7016,10 @@ version = "0.11.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6e9fcb6c2e176e86ec703e22560d99d65a5ee9056ae45a08e13e84ebf796296f" dependencies = [ - "aead", + "aead 0.5.2", "arrayref", "arrayvec", - "curve25519-dalek", + "curve25519-dalek 4.1.3", "getrandom_or_panic", "merlin", "rand_core 0.6.4", @@ -7477,7 +7606,7 @@ dependencies = [ "or-panic", "rand_core 0.6.4", "salsa20", - "x25519-dalek", + "x25519-dalek 2.0.1", "xsalsa20poly1305", ] @@ -8353,6 +8482,16 @@ dependencies = [ "subtle", ] +[[package]] +name = "universal-hash" +version = "0.6.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f4987bdc12753382e0bec4a65c50738ffaabc998b9cdd1f952fb5f39b0048a96" +dependencies = [ + "crypto-common 0.2.2", + "ctutils", +] + [[package]] name = "unsafe-libyaml" version = "0.2.11" @@ -9336,12 +9475,23 @@ version = "2.0.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c7e468321c81fb07fa7f4c636c3972b9100f0346e5b6a9f2bd0603a52f7ed277" dependencies = [ - "curve25519-dalek", + "curve25519-dalek 4.1.3", "rand_core 0.6.4", "serde", "zeroize", ] +[[package]] +name = "x25519-dalek" +version = "3.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e7e8131a03190127fb2263afc72b322ecadae46b6ff8c6f399ff5d02f5559af6" +dependencies = [ + "curve25519-dalek 5.0.0", + "rand_core 0.10.1", + "zeroize", +] + [[package]] name = "x509-cert" version = "0.2.5" @@ -9417,7 +9567,7 @@ version = "0.9.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "02a6dad357567f81cd78ee75f7c61f1b30bb2fe4390be8fb7c69e2ac8dffb6c7" dependencies = [ - "aead", + "aead 0.5.2", "poly1305", "salsa20", "subtle", diff --git a/dstack/Cargo.toml b/dstack/Cargo.toml index 2e8753148..7a98c2471 100644 --- a/dstack/Cargo.toml +++ b/dstack/Cargo.toml @@ -170,6 +170,7 @@ base64 = "0.22.1" bcrypt = "0.17.1" hex = { version = "0.4.3", default-features = false } hex_fmt = "0.3.0" +hpke = { version = "0.15.1", default-features = false, features = ["alloc", "getrandom", "x25519", "aes"] } hex-literal = "1.0.0" prost = "0.13.5" prost-types = "0.13.5" diff --git a/dstack/dstack-types/src/lib.rs b/dstack/dstack-types/src/lib.rs index d4a51710f..c7419103a 100644 --- a/dstack/dstack-types/src/lib.rs +++ b/dstack/dstack-types/src/lib.rs @@ -2400,6 +2400,8 @@ pub struct AppKeys { pub disk_crypt_key: Vec, #[serde(with = "hex_bytes", default)] pub env_crypt_key: Vec, + #[serde(with = "hex_bytes", default)] + pub sealed_env_key: Vec, #[serde(with = "hex_bytes")] pub k256_key: Vec, #[serde(with = "hex_bytes")] diff --git a/dstack/dstack-types/src/shared_filenames.rs b/dstack/dstack-types/src/shared_filenames.rs index 79944cff6..77c9d0cf6 100644 --- a/dstack/dstack-types/src/shared_filenames.rs +++ b/dstack/dstack-types/src/shared_filenames.rs @@ -8,6 +8,7 @@ pub const SYS_CONFIG: &str = ".sys-config.json"; pub const TEE_SIMULATOR_CONFIG: &str = ".tee-simulator.json"; pub const USER_CONFIG: &str = ".user-config"; pub const ENCRYPTED_ENV: &str = ".encrypted-env"; +pub const SEALED_ENV: &str = ".sealed-env"; pub const DECRYPTED_ENV: &str = ".decrypted-env"; pub const DECRYPTED_ENV_JSON: &str = ".decrypted-env.json"; pub const INSTANCE_INFO: &str = ".instance_info"; diff --git a/dstack/dstack-util/Cargo.toml b/dstack/dstack-util/Cargo.toml index b7fff59a5..f2401bc40 100644 --- a/dstack/dstack-util/Cargo.toml +++ b/dstack/dstack-util/Cargo.toml @@ -19,6 +19,7 @@ fs-err.workspace = true getrandom = { workspace = true, features = ["std"] } hex.workspace = true hex_fmt.workspace = true +hpke.workspace = true regex.workspace = true scale = { workspace = true, features = ["derive"] } schnorrkel.workspace = true diff --git a/dstack/dstack-util/src/crypto.rs b/dstack/dstack-util/src/crypto.rs index 1f8227f5b..4c080609c 100644 --- a/dstack/dstack-util/src/crypto.rs +++ b/dstack/dstack-util/src/crypto.rs @@ -68,6 +68,35 @@ pub fn dh_decrypt(secret: [u8; 32], ciphertext: &[u8]) -> Result> { .map_err(|e| anyhow!("Decryption failed: {}", e)) } +/// Encrypt to `remote_public_key` in the format [`dh_decrypt`] reads: +/// ephemeral public key, IV, then the AES-GCM ciphertext. +#[cfg(test)] +pub fn dh_encrypt(remote_public_key: [u8; 32], plaintext: &[u8]) -> Result> { + let mut ephemeral_secret = [0u8; 32]; + getrandom::fill(&mut ephemeral_secret).context("failed to generate ephemeral secret")?; + let ephemeral_secret = StaticSecret::from(ephemeral_secret); + let shared_secret = ephemeral_secret + .diffie_hellman(&PublicKey::from(remote_public_key)) + .to_bytes(); + ensure!( + !shared_secret.iter().all(|byte| *byte == 0), + "invalid X25519 shared secret" + ); + let cipher = Aes256Gcm::new_from_slice(&shared_secret) + .map_err(|e| anyhow!("failed to create cipher: {e}"))?; + let mut iv = [0u8; 12]; + getrandom::fill(&mut iv).context("failed to generate IV")?; + let ciphertext = cipher + .encrypt(Nonce::::from_slice(&iv), plaintext) + .map_err(|e| anyhow!("encryption failed: {e}"))?; + Ok([ + PublicKey::from(&ephemeral_secret).as_bytes(), + &iv[..], + &ciphertext, + ] + .concat()) +} + fn stream_nonce(prefix: &[u8; 8], index: u32) -> [u8; 12] { let mut nonce = [0u8; 12]; nonce[..8].copy_from_slice(prefix); @@ -266,6 +295,15 @@ pub fn dh_decrypt_stream( mod tests { use super::*; + #[test] + fn dh_encrypt_round_trips_through_dh_decrypt() { + let secret = [7u8; 32]; + let public = PublicKey::from(&StaticSecret::from(secret)).to_bytes(); + let ciphertext = dh_encrypt(public, b"env").unwrap(); + assert_eq!(dh_decrypt(secret, &ciphertext).unwrap(), b"env"); + assert!(dh_decrypt([8u8; 32], &ciphertext).is_err()); + } + #[test] fn test_dh_agree() { use rand::Rng; @@ -361,3 +399,95 @@ mod tests { .is_err()); } } + +/// Encrypt a sealed environment with HPKE DHKEM(X25519, HKDF-SHA256) and AES-256-GCM. +pub fn sealed_env_encrypt(public_key: &[u8], info: &[u8], plaintext: &[u8]) -> Result> { + use hpke::{ + aead::AesGcm256, kdf::HkdfSha256, kem::X25519HkdfSha256, Deserializable, OpModeS, + Serializable, + }; + let public_key = ::PublicKey::from_bytes(public_key) + .map_err(|_| anyhow!("invalid sealed env public key"))?; + let (enc, ciphertext) = hpke::single_shot_seal::( + &OpModeS::Base, + &public_key, + info, + plaintext, + &[], + ) + .map_err(|e| anyhow!("sealed env encryption failed: {e}"))?; + Ok([enc.to_bytes().as_ref(), ciphertext.as_slice()].concat()) +} + +/// Decrypt a sealed environment HPKE payload. +pub fn sealed_env_decrypt(secret: &[u8], info: &[u8], payload: &[u8]) -> Result> { + use hpke::{aead::AesGcm256, kdf::HkdfSha256, kem::X25519HkdfSha256, Deserializable, OpModeR}; + let enc = payload + .get(..32) + .ok_or_else(|| anyhow!("sealed env payload is missing encapsulated key"))?; + let ciphertext = payload + .get(32..) + .ok_or_else(|| anyhow!("sealed env payload is missing ciphertext"))?; + let secret = ::PrivateKey::from_bytes(secret) + .map_err(|_| anyhow!("invalid sealed env private key"))?; + let enc = ::EncappedKey::from_bytes(enc) + .map_err(|_| anyhow!("invalid sealed env encapsulated key"))?; + hpke::single_shot_open::( + &OpModeR::Base, + &secret, + &enc, + info, + ciphertext, + &[], + ) + .map_err(|e| anyhow!("sealed env decryption failed: {e}")) +} + +#[cfg(test)] +mod sealed_env_tests { + use super::*; + + #[test] + fn opens_sdk_envelopes() { + use base64::Engine as _; + for source in [ + include_str!("../../../sdk/fixtures/sealed-env/python.json"), + include_str!("../../../sdk/fixtures/sealed-env/javascript.json"), + ] { + let fixture: serde_json::Value = serde_json::from_str(source).unwrap(); + let envelope = &fixture["sealed_env"]; + let decode = |name: &str| { + base64::engine::general_purpose::STANDARD + .decode(envelope[name].as_str().unwrap()) + .unwrap() + }; + let payload = [decode("enc"), decode("ciphertext")].concat(); + let plaintext = sealed_env_decrypt( + &hex::decode(fixture["private_key"].as_str().unwrap()).unwrap(), + b"dstack-sealed-env-v1:3", + &payload, + ) + .unwrap(); + let decoded: serde_json::Value = serde_json::from_slice(&plaintext).unwrap(); + assert_eq!(decoded["env"], fixture["env"]); + } + } + + #[test] + fn sealed_env_round_trips_and_binds_info() { + let secret = StaticSecret::random_from_rng(rand::thread_rng()); + let public = PublicKey::from(&secret).to_bytes(); + let info = b"dstack-sealed-env-v1:3"; + let payload = sealed_env_encrypt(&public, info, b"secret env").unwrap(); + assert_eq!( + sealed_env_decrypt(&secret.to_bytes(), info, &payload).unwrap(), + b"secret env" + ); + assert!( + sealed_env_decrypt(&secret.to_bytes(), b"dstack-sealed-env-v1:4", &payload).is_err() + ); + let mut tampered = payload; + *tampered.last_mut().unwrap() ^= 1; + assert!(sealed_env_decrypt(&secret.to_bytes(), info, &tampered).is_err()); + } +} diff --git a/dstack/dstack-util/src/host_api.rs b/dstack/dstack-util/src/host_api.rs index e038ea3b2..743becd42 100644 --- a/dstack/dstack-util/src/host_api.rs +++ b/dstack/dstack-util/src/host_api.rs @@ -80,6 +80,24 @@ impl HostApi { Ok(()) } + /// Hand the host the env re-encrypted to the current root epoch, or an + /// empty one when nothing changed. + pub async fn update_sealed_env(&self, sealed_env: Vec) -> Result<()> { + if !matches!(Platform::detect_or_dstack(), Platform::Dstack) { + return Ok(()); + } + let Some(client) = &self.client else { + return Ok(()); + }; + tokio::time::timeout( + HOST_API_TIMEOUT, + client.update_sealed_env(host_api::SealedEnv { sealed_env }), + ) + .await + .context("Timed out updating the sealed env on the host")??; + Ok(()) + } + pub async fn notify_q(&self, event: &str, payload: &str) { if let Err(err) = self.notify(event, payload).await { warn!("Failed to notify event {event} to host: {:?}", err); diff --git a/dstack/dstack-util/src/main.rs b/dstack/dstack-util/src/main.rs index ab1f6837b..5b0eebf9c 100644 --- a/dstack/dstack-util/src/main.rs +++ b/dstack/dstack-util/src/main.rs @@ -677,6 +677,7 @@ async fn cmd_get_keys(args: GetKeysArgs) -> Result<()> { ca_cert: root_ca, disk_crypt_key: response.disk_crypt_key, env_crypt_key: response.env_crypt_key, + sealed_env_key: response.sealed_env_key, k256_key: response.k256_key, k256_signature: response.k256_signature, gateway_app_id: response.gateway_app_id, @@ -793,6 +794,7 @@ async fn cmd_encrypt(args: EncryptArgs) -> Result<()> { let response = KmsClient::new(client) .get_app_env_encrypt_pub_key(dstack_kms_rpc::AppId { app_id: app_id.to_vec(), + root_epoch: None, }) .await .context("failed to get app environment encryption public key")?; @@ -1191,6 +1193,7 @@ fn make_app_keys( Ok(AppKeys { disk_crypt_key: sha256(&disk_key.serialize_der()).to_vec(), env_crypt_key: vec![], + sealed_env_key: vec![], k256_key: k256_key.to_bytes().to_vec(), k256_signature: vec![], gateway_app_id: "".to_string(), diff --git a/dstack/dstack-util/src/system_setup.rs b/dstack/dstack-util/src/system_setup.rs index 0699befb9..54a607363 100644 --- a/dstack/dstack-util/src/system_setup.rs +++ b/dstack/dstack-util/src/system_setup.rs @@ -16,13 +16,14 @@ use std::{ }; use anyhow::{anyhow, bail, Context, Result}; +use base64::Engine as _; use dstack_attest::{default_verifier, emit_runtime_event, set_runtime_event_version}; use dstack_kms_rpc as rpc; use dstack_types::{ gpu_policy_hash, shared_filenames::{ APP_COMPOSE, APP_KEYS, DECRYPTED_ENV, DECRYPTED_ENV_JSON, ENCRYPTED_ENV, - HOST_SHARED_DIR_NAME, INSTANCE_INFO, SYS_CONFIG, USER_CONFIG, + HOST_SHARED_DIR_NAME, INSTANCE_INFO, SEALED_ENV, SYS_CONFIG, USER_CONFIG, }, GpuPolicy, KeyProvider, KeyProviderInfo, GPU_ATTESTATION_OUTPUT, }; @@ -229,6 +230,10 @@ impl HostShareDir { self.base_dir.join(ENCRYPTED_ENV) } + fn sealed_env_file(&self) -> PathBuf { + self.base_dir.join(SEALED_ENV) + } + fn sys_config_file(&self) -> PathBuf { self.base_dir.join(SYS_CONFIG) } @@ -242,11 +247,34 @@ impl HostShareDir { } } +fn env_file_exists(path: &Path) -> Result { + match fs::symlink_metadata(path) { + Ok(metadata) if metadata.is_file() => Ok(true), + Ok(_) => bail!("env {} is not a regular file", path.display()), + Err(error) if error.kind() == std::io::ErrorKind::NotFound => Ok(false), + Err(error) => { + Err(error).with_context(|| format!("failed to inspect env {}", path.display())) + } + } +} + +fn read_optional_env(path: &Path) -> Result>> { + if !env_file_exists(path)? { + return Ok(None); + } + match fs::read(path) { + Ok(bytes) => Ok(Some(bytes)), + Err(error) if error.kind() == std::io::ErrorKind::NotFound => Ok(None), + Err(error) => Err(error).with_context(|| format!("failed to read env {}", path.display())), + } +} + struct HostShared { dir: HostShareDir, sys_config: SysConfig, app_compose: AppCompose, encrypted_env: Vec, + sealed_env: Option>, instance_info: InstanceInfo, } @@ -261,12 +289,18 @@ impl HostShared { } else { InstanceInfo::default() }; - let encrypted_env = fs::read(host_shared_dir.encrypted_env_file()).unwrap_or_default(); + let sealed_env = read_optional_env(&host_shared_dir.sealed_env_file())?; + let encrypted_env = if sealed_env.is_some() { + Vec::new() + } else { + read_optional_env(&host_shared_dir.encrypted_env_file())?.unwrap_or_default() + }; Ok(Self { dir: host_shared_dir.clone(), sys_config, app_compose, encrypted_env, + sealed_env, instance_info, }) } @@ -315,7 +349,16 @@ impl HostShared { copy(APP_COMPOSE, SZ_1MB * 50, false)?; copy(SYS_CONFIG, SZ_1KB * 32, false)?; copy(INSTANCE_INFO, SZ_1KB * 10, true)?; - copy(ENCRYPTED_ENV, SZ_1KB * 256, true)?; + if env_file_exists(&host_shared_dir.join(SEALED_ENV))? { + copy(SEALED_ENV, SZ_1KB * 384, false)?; + } else { + match fs::remove_file(host_shared_copy_dir.join(SEALED_ENV)) { + Ok(()) => {} + Err(error) if error.kind() == std::io::ErrorKind::NotFound => {} + Err(error) => return Err(error).context("failed to remove stale sealed env copy"), + } + copy(ENCRYPTED_ENV, SZ_1KB * 256, true)?; + } copy(USER_CONFIG, SZ_1MB * 50, true)?; info!("Unmounting host-shared"); unmount_host_shared(host_shared_dir)?; @@ -2389,6 +2432,7 @@ impl HostShared { ca_cert: root_ca, disk_crypt_key: response.disk_crypt_key, env_crypt_key: response.env_crypt_key, + sealed_env_key: response.sealed_env_key, k256_key: response.k256_key, k256_signature: response.k256_signature, gateway_app_id: response.gateway_app_id, @@ -3277,25 +3321,74 @@ impl<'a> Stage0<'a> { } impl Stage1<'_> { - fn decrypt_env_vars( - &self, - key: &[u8], - ciphertext: &[u8], - allowed: &BTreeSet, - ) -> Result> { - let vars = if !ciphertext.is_empty() { - info!("Processing encrypted env"); - let env_crypt_key: [u8; 32] = key.try_into().ok().context( - "encrypted env is present but the key provider gave no valid env crypt key", + /// Decrypt the host-supplied env, preferring the sealed format. + async fn decrypt_env(&self) -> Result<(Vec, Vec)> { + let current_secret = self.keys.sealed_env_key.as_slice(); + 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 { + bail!("unsupported sealed env version: {}", envelope.version); + } + let enc = base64::engine::general_purpose::STANDARD.decode(envelope.enc)?; + if enc.len() != 32 { + bail!("sealed env encapsulated key must be 32 bytes"); + } + let ciphertext = + base64::engine::general_purpose::STANDARD.decode(envelope.ciphertext)?; + let payload = [enc, ciphertext].concat(); + let current = self.keys.root_epoch(); + if envelope.root_epoch > current { + bail!( + "sealed env epoch {} exceeds current epoch {current}", + envelope.root_epoch + ); + } + let historical; + let secret = if envelope.root_epoch == current { + current_secret + } else { + historical = self + .shared + .request_historical_keys(&self.keys, envelope.root_epoch) + .await?; + historical.sealed_env_key.as_slice() + }; + let plaintext = crate::crypto::sealed_env_decrypt( + secret, + &sealed_env_info(&envelope.root_epoch), + &payload, )?; - let decrypted_json = - dh_decrypt(env_crypt_key, ciphertext).context("Failed to decrypt env file")?; - crate::parse_env_file::parse_env(&decrypted_json, allowed)? - } else { - info!("No encrypted env, using default"); - Default::default() + if envelope.root_epoch == self.keys.root_epoch() { + return Ok((plaintext, Vec::new())); + } + let sealed = seal_env(current_secret, self.keys.root_epoch(), &plaintext)?; + return Ok((plaintext, sealed)); + } + let ciphertext = &self.shared.encrypted_env; + let current_key = env_crypt_key(&self.keys.env_crypt_key)?; + let plaintext = match dh_decrypt(current_key, ciphertext) { + Ok(plaintext) => plaintext, + Err(error) => { + let current = self.keys.root_epoch(); + let mut recovered = None; + for epoch in (0..current).rev() { + let old = self + .shared + .request_historical_keys(&self.keys, epoch) + .await?; + if let Ok(plaintext) = + dh_decrypt(env_crypt_key(&old.env_crypt_key)?, ciphertext) + { + recovered = Some(plaintext); + break; + } + } + recovered.ok_or(error)? + } }; - Ok(vars) + let sealed = seal_env(current_secret, self.keys.root_epoch(), &plaintext)?; + Ok((plaintext, sealed)) } fn write_env_file(&self, env_vars: &BTreeMap) -> Result<()> { @@ -3312,7 +3405,7 @@ impl Stage1<'_> { Ok(()) } - fn unseal_env_vars(&self) -> Result> { + async fn unseal_env_vars(&self) -> Result> { let allowed_envs: BTreeSet = self .shared .app_compose @@ -3320,18 +3413,29 @@ impl Stage1<'_> { .iter() .cloned() .collect(); - // Decrypt env file - let decrypted_env = self.decrypt_env_vars( - &self.keys.env_crypt_key, - &self.shared.encrypted_env, - &allowed_envs, - )?; + let (decrypted_env, sealed) = + if self.shared.sealed_env.is_none() && self.shared.encrypted_env.is_empty() { + info!("No encrypted env, using default"); + (Default::default(), Vec::new()) + } else { + info!("Processing encrypted env"); + let (plaintext, sealed) = self.decrypt_env().await?; + ( + crate::parse_env_file::parse_env(&plaintext, &allowed_envs)?, + sealed, + ) + }; self.write_env_file(&decrypted_env)?; + // Sent on every boot, empty when nothing changed: the host accepts one + // update per boot, and this one comes before any application code. + if let Err(err) = self.vmm.update_sealed_env(sealed).await { + warn!("failed to store the re-encrypted env on the host: {err:#}"); + } Ok(decrypted_env) } async fn setup(&self) -> Result<()> { - let _envs = self.unseal_env_vars()?; + let _envs = self.unseal_env_vars().await?; self.link_files()?; self.setup_socket_dir()?; self.setup_guest_agent_config()?; @@ -3428,6 +3532,37 @@ impl Stage1<'_> { } } +#[derive(serde::Deserialize)] +struct SealedEnv { + version: u32, + root_epoch: u32, + enc: String, + ciphertext: String, +} + +fn sealed_env_info(epoch: &u32) -> Vec { + format!("dstack-sealed-env-v1:{epoch}").into_bytes() +} + +fn seal_env(secret: &[u8], 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 (enc, ciphertext) = payload.split_at(32); + Ok(serde_json::to_vec(&serde_json::json!({ + "version": 1, + "root_epoch": epoch, + "enc": base64::engine::general_purpose::STANDARD.encode(enc), + "ciphertext": base64::engine::general_purpose::STANDARD.encode(ciphertext), + }))?) +} + +fn env_crypt_key(key: &[u8]) -> Result<[u8; 32]> { + key.try_into() + .ok() + .context("encrypted env is present but the key provider gave no valid env crypt key") +} fn kms_root_epoch(keys: &AppKeys) -> Option { match keys.key_provider { KeyProvider::Kms { root_epoch, .. } => Some(root_epoch), @@ -4227,3 +4362,118 @@ Endpoint = [2001:db8::1]:51822 assert!(parse_wg_ip("10.0.0.2/32\n[Interface]").is_err()); } } + +#[cfg(test)] +mod sealed_env_file_tests { + use super::*; + + fn stage(args: &SetupArgs, sealed_env: Option>, encrypted_env: Vec) -> Stage1<'_> { + Stage1 { + args, + vmm: HostApi::default(), + shared: HostShared { + dir: HostShareDir::new(&args.work_dir), + sys_config: serde_json::from_value(serde_json::json!({"vm_config": ""})).unwrap(), + app_compose: serde_json::from_value(serde_json::json!({ + "manifest_version": "2", "name": "sealed-test", "runner": "docker-compose" + })) + .unwrap(), + instance_info: InstanceInfo::default(), + sealed_env, + encrypted_env, + }, + keys: AppKeys { + disk_crypt_key: vec![], + env_crypt_key: vec![9; 32], + sealed_env_key: vec![7; 32], + k256_key: vec![], + k256_signature: vec![], + gateway_app_id: String::new(), + ca_cert: String::new(), + key_provider: KeyProvider::Kms { + url: "http://127.0.0.1:1".into(), + pubkey: vec![], + root_epoch: 3, + }, + }, + } + } + + #[tokio::test] + async fn guest_prefers_sealed_and_converts_legacy() { + 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":"MESSAGE","value":"test"}]}"#; + let sealed = seal_env(&[7; 32], 3, plaintext).unwrap(); + let current = stage(&args, Some(sealed), b"invalid legacy".to_vec()); + let (opened, update) = current.decrypt_env().await.unwrap(); + assert_eq!(opened, plaintext); + assert!(update.is_empty()); + + let public = x25519_dalek::PublicKey::from(&x25519_dalek::StaticSecret::from([9; 32])); + let legacy = crate::crypto::dh_encrypt(public.to_bytes(), plaintext).unwrap(); + let old = stage(&args, None, legacy.clone()); + let (opened, update) = old.decrypt_env().await.unwrap(); + assert_eq!(opened, plaintext); + assert!(!update.is_empty()); + let migrated = stage(&args, Some(update), vec![]); + assert_eq!(migrated.decrypt_env().await.unwrap().0, plaintext); + + for invalid in [ + Vec::new(), + b"invalid JSON".to_vec(), + br#"{"version":2,"root_epoch":3,"enc":"","ciphertext":""}"#.to_vec(), + ] { + assert!(stage(&args, Some(invalid), legacy.clone()) + .decrypt_env() + .await + .is_err()); + } + let future = seal_env(&[7; 32], 4, plaintext).unwrap(); + assert!(stage(&args, Some(future), legacy) + .decrypt_env() + .await + .unwrap_err() + .to_string() + .contains("exceeds current")); + } + + #[test] + fn sealed_file_selection_ignores_legacy_but_rejects_links() { + let dir = tempfile::tempdir().unwrap(); + fs::write(dir.path().join(SYS_CONFIG), br#"{"vm_config":""}"#).unwrap(); + fs::write( + dir.path().join(APP_COMPOSE), + br#"{"manifest_version":"2","name":"test","runner":"docker-compose"}"#, + ) + .unwrap(); + let sealed_path = dir.path().join(SEALED_ENV); + fs::write(&sealed_path, b"sealed bytes").unwrap(); + fs::create_dir(dir.path().join(ENCRYPTED_ENV)).unwrap(); + let shared = HostShared::load(dir.path()).unwrap(); + assert_eq!(shared.sealed_env.unwrap(), b"sealed bytes"); + assert!(shared.encrypted_env.is_empty()); + fs::remove_file(&sealed_path).unwrap(); + std::os::unix::fs::symlink("missing-target", &sealed_path).unwrap(); + assert!(env_file_exists(&sealed_path).is_err()); + assert!(HostShared::load(dir.path()).is_err()); + } + + #[test] + fn missing_empty_and_unreadable_env_are_distinct() { + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join(SEALED_ENV); + assert_eq!(read_optional_env(&path).unwrap(), None); + fs::write(&path, []).unwrap(); + assert_eq!(read_optional_env(&path).unwrap(), Some(Vec::new())); + assert!( + serde_json::from_slice::(&read_optional_env(&path).unwrap().unwrap()) + .is_err() + ); + assert!(read_optional_env(dir.path()).is_err()); + } +} diff --git a/dstack/guest-agent/src/rpc_service.rs b/dstack/guest-agent/src/rpc_service.rs index d79a9ed22..58ff5a865 100644 --- a/dstack/guest-agent/src/rpc_service.rs +++ b/dstack/guest-agent/src/rpc_service.rs @@ -1437,6 +1437,7 @@ pNs85uhOZE8z2jr8Pg== let dummy_keys = AppKeys { disk_crypt_key: Vec::new(), env_crypt_key: Vec::new(), + sealed_env_key: Vec::new(), k256_key: DUMMY_K256_KEY.to_vec(), k256_signature: Vec::new(), gateway_app_id: String::new(), diff --git a/dstack/host-api/proto/host_api.proto b/dstack/host-api/proto/host_api.proto index f8864f516..c4880ba74 100644 --- a/dstack/host-api/proto/host_api.proto +++ b/dstack/host-api/proto/host_api.proto @@ -27,8 +27,17 @@ message GetSealingKeyResponse { bytes provider_quote = 2; } +// The env the guest re-encrypted to the current KMS root epoch at boot; empty +// when it needed no change. +message SealedEnv { + bytes sealed_env = 1; +} + service HostApi { rpc Info(google.protobuf.Empty) returns (HostInfo); rpc Notify(Notification) returns (google.protobuf.Empty); rpc GetSealingKey(GetSealingKeyRequest) returns (GetSealingKeyResponse); + // Replace this VM's stored encrypted env. Accepted once per boot, before + // application code runs, and only if the stored env is the one it booted with. + rpc UpdateSealedEnv(SealedEnv) returns (google.protobuf.Empty); } diff --git a/dstack/kms/Cargo.toml b/dstack/kms/Cargo.toml index 5c616d7a4..2979b4230 100644 --- a/dstack/kms/Cargo.toml +++ b/dstack/kms/Cargo.toml @@ -32,6 +32,7 @@ reqwest = { workspace = true, features = ["json"] } sha2.workspace = true sha3.workspace = true k256.workspace = true +p256.workspace = true rand.workspace = true dstack-guest-agent-rpc.workspace = true http-client = { workspace = true, features = ["prpc"] } diff --git a/dstack/kms/rpc/proto/kms_rpc.proto b/dstack/kms/rpc/proto/kms_rpc.proto index 13083647b..7cfd3d8ac 100644 --- a/dstack/kms/rpc/proto/kms_rpc.proto +++ b/dstack/kms/rpc/proto/kms_rpc.proto @@ -18,6 +18,14 @@ message GetAppKeyRequest { message AppId { bytes app_id = 1; + optional uint32 root_epoch = 2; +} + +message SealedEnvPubKeyResponse { + bytes public_key = 1; + uint32 root_epoch = 2; + uint64 timestamp = 3; + bytes signature = 4; } message PublicKeyResponse { @@ -51,6 +59,8 @@ message AppKeyResponse { bytes os_image_hash = 8; // Root key epoch all epoch-dependent fields above were derived from. uint32 root_epoch = 9; + // X25519 private key for the SealedEnv HPKE receiver, derived for this epoch. + bytes sealed_env_key = 10; } message GetMetaResponse { @@ -109,6 +119,7 @@ service KMS { rpc GetKmsKey(GetKmsKeyRequest) returns (KmsKeyResponse); // Request the app environment encryption public key given the app id rpc GetAppEnvEncryptPubKey(AppId) returns (PublicKeyResponse); + rpc GetSealedEnvPubKey(AppId) returns (SealedEnvPubKeyResponse); // Request the KMS instance metadata for use as a probe and health check. rpc GetMeta(google.protobuf.Empty) returns (GetMetaResponse); // Compatibility for existing OS images and tokenless legacy onboarding. diff --git a/dstack/kms/src/main_service.rs b/dstack/kms/src/main_service.rs index 605092339..20a63afb2 100644 --- a/dstack/kms/src/main_service.rs +++ b/dstack/kms/src/main_service.rs @@ -14,13 +14,14 @@ use anyhow::{anyhow, bail, Context, Result}; use dstack_kms_rpc::{ kms_server::{KmsRpc, KmsServer}, AppId, AppKeyResponse, GetAppKeyRequest, GetKmsKeyRequest, GetMetaResponse, - GetTempCaCertResponse, KmsKeyResponse, KmsKeys, PublicKeyResponse, SignCertRequest, - SignCertResponse, + GetTempCaCertResponse, KmsKeyResponse, KmsKeys, PublicKeyResponse, SealedEnvPubKeyResponse, + SignCertRequest, SignCertResponse, }; use dstack_verifier::{CvmVerifier, VerificationDetails}; use fs_err as fs; use k256::ecdsa::SigningKey; use moka::future::Cache; +use p256::pkcs8::DecodePrivateKey; use ra_rpc::{CallContext, RpcCall}; use ra_tls::{ attestation::{resolve_vm_config, AttestationVerifier, TeeVariant, VerifiedAttestation}, @@ -366,6 +367,21 @@ struct AppKeys { env_crypt_key: Vec, k256_key: Vec, k256_signature: Vec, + sealed_env_key: Vec, +} + +fn derive_sealed_env_key(root: &RootKeys, app_id: &[u8]) -> 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], + 32, + ) + .map_err(|_| anyhow!("failed to derive sealed env key"))? + .try_into() + .map_err(|_| anyhow!("invalid sealed env key length")) } fn derive_app_keys(root: &RootKeys, app_id: &[u8], instance_id: &[u8]) -> Result { @@ -377,11 +393,13 @@ fn derive_app_keys(root: &RootKeys, app_id: &[u8], instance_id: &[u8]) -> Result 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(); Ok(AppKeys { disk_crypt_key: disk_crypt_key.to_vec(), env_crypt_key: env_crypt_key.to_vec(), k256_key: k256_key.to_bytes().to_vec(), k256_signature, + sealed_env_key, }) } @@ -660,6 +678,7 @@ impl KmsRpc for RpcHandler { gateway_app_id, os_image_hash: boot_info.os_image_hash, root_epoch, + sealed_env_key: keys.sealed_env_key, }) } @@ -711,6 +730,40 @@ impl KmsRpc for RpcHandler { }) } + async fn get_sealed_env_pub_key(self, request: AppId) -> Result { + let requested = request.root_epoch; + let current = self.ensure_self_allowed(Operation::Serve).await?; + let epoch = requested.unwrap_or(current); + if epoch > current { + bail!("sealed env root epoch {epoch} is not active; current epoch is {current}"); + } + 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 public_key = + x25519_dalek::PublicKey::from(&x25519_dalek::StaticSecret::from(secret)).to_bytes(); + let timestamp = std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .context("system time before UNIX epoch")? + .as_secs(); + let mut message = Vec::with_capacity(4 + 32 + 8); + message.extend_from_slice(&epoch.to_be_bytes()); + message.extend_from_slice(&public_key); + message.extend_from_slice(×tamp.to_be_bytes()); + let signature = sign_message( + &root.k256, + b"dstack-sealed-env-pubkey-v1", + &request.app_id, + &message, + )?; + Ok(SealedEnvPubKeyResponse { + public_key: public_key.to_vec(), + root_epoch: epoch, + timestamp, + signature, + }) + } + async fn get_meta(self) -> Result { let bootstrap_info = fs::read_to_string(self.state.config.bootstrap_info()) .ok() @@ -1532,6 +1585,21 @@ mod tests { KmsState::new(config, keyring, ca_cert.pem()).unwrap() } + #[test] + fn sealed_env_derivation_is_separate_and_shared_with_app_keys() { + 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 sealed = derive_sealed_env_key(root, &app).unwrap(); + let app_keys = derive_app_keys(root, &app, &[2; 20]).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()); + } + #[rocket::async_test] async fn authorization_selects_epochs_without_a_serving_epoch_snapshot() { let dir = tempfile::tempdir().unwrap(); diff --git a/dstack/vmm/requirements.txt b/dstack/vmm/requirements.txt index 5a4636513..9ec6ca2f0 100644 --- a/dstack/vmm/requirements.txt +++ b/dstack/vmm/requirements.txt @@ -1,5 +1,5 @@ annotated-types==0.7.0 -cffi==1.17.1 +cffi==2.0.0 cryptography==50.0.0 cytoolz==1.0.1 eth-hash==0.7.1 @@ -7,6 +7,7 @@ eth-keys==0.7.0 eth-typing==5.2.1 eth-utils==5.3.0 pycparser==2.22 +pyhpke==0.6.5 pydantic==2.11.3 pydantic_core==2.33.1 safe-pysha3==1.0.4 diff --git a/dstack/vmm/rpc/proto/vmm_rpc.proto b/dstack/vmm/rpc/proto/vmm_rpc.proto index 84b66b2d6..6fb3a61fd 100644 --- a/dstack/vmm/rpc/proto/vmm_rpc.proto +++ b/dstack/vmm/rpc/proto/vmm_rpc.proto @@ -109,6 +109,8 @@ message VmConfiguration { repeated PortMapping ports = 7; // Encrypted environment variables bytes encrypted_env = 8; + // Sealed environment variables. Takes precedence over encrypted_env. + bytes sealed_env = 24; // App ID. // If not provided, it assumes the app is newly created. // If provided, and KMS is enabled, it assumes the app is upgraded from given app_id. @@ -214,6 +216,8 @@ message UpdateVmRequest { string compose_file = 2; // Optional update encrypted environment variables. Leave empty to not update. bytes encrypted_env = 3; + // Optional sealed environment update. Leave empty to not update. + bytes sealed_env = 24; // Optional update user config. Leave empty to not update. string user_config = 4; // Optional update port mapping. @@ -287,9 +291,17 @@ message ImageInfo { // App identifier used for encrypting env vars and routing telemetry. message AppId { bytes app_id = 1; + optional uint32 root_epoch = 2; } // Response with the KMS-backed public key used for env encryption plus audit signature. +message SealedEnvPubKeyResponse { + bytes public_key = 1; + uint32 root_epoch = 2; + uint64 timestamp = 3; + bytes signature = 4; +} + message PublicKeyResponse { bytes public_key = 1; // Legacy signature without timestamp (for backward compatibility). @@ -423,6 +435,7 @@ service Vmm { // Get Env encrypt public key rpc GetAppEnvEncryptPubKey(AppId) returns (PublicKeyResponse); + rpc GetSealedEnvPubKey(AppId) returns (SealedEnvPubKeyResponse); // Get VM info by ID rpc GetInfo(Id) returns (GetInfoResponse); diff --git a/dstack/vmm/src/app.rs b/dstack/vmm/src/app.rs index 379bea90a..f060c1ece 100644 --- a/dstack/vmm/src/app.rs +++ b/dstack/vmm/src/app.rs @@ -16,7 +16,8 @@ use bon::Builder; use dstack_kms_rpc::kms_client::KmsClient; use dstack_types::mr_config::MrConfigV3; use dstack_types::shared_filenames::{ - APP_COMPOSE, ENCRYPTED_ENV, INSTANCE_INFO, SYS_CONFIG, TEE_SIMULATOR_CONFIG, USER_CONFIG, + APP_COMPOSE, ENCRYPTED_ENV, INSTANCE_INFO, SEALED_ENV, SYS_CONFIG, TEE_SIMULATOR_CONFIG, + USER_CONFIG, }; use dstack_types::version::Version; use dstack_vmm_rpc::{ @@ -530,9 +531,16 @@ impl App { } self.sync_dynamic_config(id)?; self.set_started(id, true)?; + // A guest the VMM did not launch may already run application code. + let env_update = if is_running { + None + } else { + Some(self.sha256_of_env(id)?) + }; let vm_config = { let mut state = self.lock(); let vm_state = state.get_mut(id).context("VM not found")?; + vm_state.state.env_update = env_update; // Older images does not support for progress reporting if vm_state.config.image.info.shared_ro { vm_state.state.start(is_running); @@ -1561,7 +1569,22 @@ impl App { vm.state.events.pop_front(); } match event { - "boot.progress" => vm.state.boot_progress = body, + "boot.progress" => { + // Application code (init and pre-launch scripts, containers) + // runs from here on, so the guest OS can no longer vouch for + // an env update. + if matches!( + body.as_str(), + "init-script" + | "pre-launch" + | "starting containers" + | "running main script" + | "done" + ) { + vm.state.env_update = None; + } + vm.state.boot_progress = body; + } "boot.error" => vm.state.boot_error = body, "shutdown.progress" => { let powering_off = body == "powering off"; @@ -1582,12 +1605,63 @@ impl App { Ok(()) } - pub(crate) fn compose_file_path(&self, id: &str) -> Result { - Ok(self.shared_dir(id)?.join(APP_COMPOSE)) + /// Store the encrypted env a guest re-encrypted at boot. Accepted once per + /// boot, before application code runs, and only over the env the VM booted + /// with, so an update made since then is never overwritten. + pub(crate) async fn update_sealed_env(&self, cid: u32, env: Vec) -> Result<()> { + let id = self + .lock() + .vms + .values() + .find(|vm| vm.config.cid == cid) + .context("VM not found")? + .config + .manifest + .id + .clone(); + let _launch = self.launch_lock(&id).await; + let boot_env = { + let mut state = self.lock(); + let vm = state.vms.get_mut(&id).context("VM not found")?; + if vm.config.cid != cid { + bail!("VM CID changed"); + } + vm.state.env_update.take() + }; + let Some(boot_env) = boot_env else { + bail!("the env can only be updated once per boot, before applications start"); + }; + if env.is_empty() { + return Ok(()); + } + if self.sha256_of_env(&id)? != boot_env { + bail!("the env changed since the VM booted; not replacing it"); + } + self.work_dir(&id)?.write_env(&[], &env)?; + info!(id, "stored sealed env from the guest"); + Ok(()) } - pub(crate) fn encrypted_env_path(&self, id: &str) -> Result { - Ok(self.shared_dir(id)?.join(ENCRYPTED_ENV)) + fn sha256_of_env(&self, id: &str) -> Result<[u8; 32]> { + let shared = self.shared_dir(id)?; + let mut hasher = Sha256::new(); + for name in [ENCRYPTED_ENV, SEALED_ENV] { + hasher.update(name.as_bytes()); + match fs::read(shared.join(name)) { + Ok(data) => { + hasher.update([1]); + hasher.update((data.len() as u64).to_le_bytes()); + hasher.update(data); + } + Err(err) if err.kind() == std::io::ErrorKind::NotFound => hasher.update([0]), + Err(err) => return Err(err).context("failed to read env"), + } + } + Ok(hasher.finalize().into()) + } + + pub(crate) fn compose_file_path(&self, id: &str) -> Result { + Ok(self.shared_dir(id)?.join(APP_COMPOSE)) } pub(crate) fn user_config_path(&self, id: &str) -> Result { @@ -1605,15 +1679,13 @@ impl App { req: &VmConfiguration, app_id: &str, ) -> Result { + VmWorkDir::validate_env_update(&req.encrypted_env, &req.sealed_env)?; let work_dir = self.work_dir(id)?; let shared_dir = work_dir.join("shared"); fs::create_dir_all(&shared_dir).context("Failed to create shared directory")?; fs::write(shared_dir.join(APP_COMPOSE), &req.compose_file) .context("Failed to write compose file")?; - if !req.encrypted_env.is_empty() { - fs::write(shared_dir.join(ENCRYPTED_ENV), &req.encrypted_env) - .context("Failed to write encrypted env")?; - } + work_dir.write_env(&req.encrypted_env, &req.sealed_env)?; if !req.user_config.is_empty() { fs::write(shared_dir.join(USER_CONFIG), &req.user_config) .context("Failed to write user config")?; @@ -2508,6 +2580,83 @@ mod tests { (vm, networks) } + #[test] + fn prepare_work_dir_stores_sealed_and_rejects_both_formats() { + let dir = tempfile::tempdir().unwrap(); + let app = App::new( + test_config(&dir.path().join("absent-netd.sock"), dir.path()), + SupervisorClient::new("http://127.0.0.1:0"), + ); + let mut request = VmConfiguration { + compose_file: "{}".into(), + sealed_env: b"sealed".to_vec(), + ..Default::default() + }; + let workdir = app.prepare_work_dir("vm-1", &request, "").unwrap(); + assert_eq!(fs::read(workdir.sealed_env_path()).unwrap(), b"sealed"); + assert!(!workdir.encrypted_env_path().exists()); + request.encrypted_env = b"legacy".to_vec(); + assert!(app.prepare_work_dir("vm-2", &request, "").is_err()); + assert!(!app.work_dir("vm-2").unwrap().path().exists()); + } + + #[tokio::test] + async fn sealed_env_cas_binds_both_files_and_serializes_owner_updates() { + let dir = tempfile::tempdir().unwrap(); + let app = App::new( + test_config(&dir.path().join("absent-netd.sock"), dir.path()), + SupervisorClient::new("http://127.0.0.1:0"), + ); + let (vm, _) = bridge_vm(&app, "vm-1"); + let cid = vm.cid; + app.lock().vms.insert("vm-1".into(), VmState::new(vm)); + let workdir = app.work_dir("vm-1").unwrap(); + fs::create_dir_all(workdir.shared_dir()).unwrap(); + let boot = || { + let hash = app.sha256_of_env("vm-1").unwrap(); + app.lock().vms.get_mut("vm-1").unwrap().state.env_update = Some(hash); + }; + assert!(app.update_sealed_env(cid, b"new".to_vec()).await.is_err()); + workdir.write_env(b"legacy", &[]).unwrap(); + boot(); + app.update_sealed_env(cid, b"sealed".to_vec()) + .await + .unwrap(); + assert!(!workdir.encrypted_env_path().exists()); + assert_eq!(fs::read(workdir.sealed_env_path()).unwrap(), b"sealed"); + assert!(app.update_sealed_env(cid, b"again".to_vec()).await.is_err()); + + for name in [ENCRYPTED_ENV, SEALED_ENV] { + let path = workdir.shared_dir().join(name); + let _ = fs::remove_file(&path); + boot(); + fs::write(&path, []).unwrap(); + assert!(app.update_sealed_env(cid, b"stale".to_vec()).await.is_err()); + assert_eq!(fs::read(&path).unwrap(), b""); + } + + boot(); + let held = app.launch_lock("vm-1").await; + let guest = app.update_sealed_env(cid, b"stale".to_vec()); + tokio::pin!(guest); + assert!(tokio::time::timeout(Duration::from_millis(20), &mut guest) + .await + .is_err()); + workdir.write_env(b"owner", &[]).unwrap(); + drop(held); + assert!(guest.await.is_err()); + assert_eq!(fs::read(workdir.encrypted_env_path()).unwrap(), b"owner"); + assert!(!workdir.sealed_env_path().exists()); + + boot(); + app.update_sealed_env(cid, Vec::new()).await.unwrap(); + assert!(app.update_sealed_env(cid, b"again".to_vec()).await.is_err()); + boot(); + app.vm_event_report(cid, "boot.progress", "init-script".into()) + .unwrap(); + assert!(app.update_sealed_env(cid, b"late".to_vec()).await.is_err()); + } + /// A launch that never reached a netd built nothing, so it must not leave /// a marker that makes every later stop and removal wait for one. What /// an earlier boot may still hold keeps its marker. @@ -3950,6 +4099,10 @@ struct VmStateMut { auto_restart: AutoRestartState, /// True when the VM is being removed (cleanup in progress). removing: bool, + /// SHA-256 of both env files, including presence, this boot started with, while the guest + /// may still replace it; taken by the first update and cleared once the + /// boot reaches application code. + env_update: Option<[u8; 32]>, } impl VmStateMut { diff --git a/dstack/vmm/src/app/vm_info.rs b/dstack/vmm/src/app/vm_info.rs index 26e4e1f84..557a089e9 100644 --- a/dstack/vmm/src/app/vm_info.rs +++ b/dstack/vmm/src/app/vm_info.rs @@ -203,6 +203,7 @@ impl VmInfo { compose_file: fs::read_to_string(workdir.app_compose_path()) .unwrap_or_default(), encrypted_env: fs::read(workdir.encrypted_env_path()).unwrap_or_default(), + sealed_env: fs::read(workdir.sealed_env_path()).unwrap_or_default(), user_config: fs::read_to_string(workdir.user_config_path()).unwrap_or_default(), vcpu: self.manifest.vcpu, memory: self.manifest.memory, diff --git a/dstack/vmm/src/app/workdir.rs b/dstack/vmm/src/app/workdir.rs index fd67072f3..3fce85cfb 100644 --- a/dstack/vmm/src/app/workdir.rs +++ b/dstack/vmm/src/app/workdir.rs @@ -9,7 +9,9 @@ use std::path::{Path, PathBuf}; use anyhow::{Context, Result}; use dstack_types::{ - shared_filenames::{APP_COMPOSE, ENCRYPTED_ENV, INSTANCE_INFO, SYS_CONFIG, USER_CONFIG}, + shared_filenames::{ + APP_COMPOSE, ENCRYPTED_ENV, INSTANCE_INFO, SEALED_ENV, SYS_CONFIG, USER_CONFIG, + }, AppCompose, SysConfig, }; use fs_err as fs; @@ -131,10 +133,39 @@ impl VmWorkDir { self.shared_dir().join(USER_CONFIG) } + pub(crate) fn validate_env_update(encrypted: &[u8], sealed: &[u8]) -> Result<()> { + if !encrypted.is_empty() && !sealed.is_empty() { + anyhow::bail!("encrypted_env and sealed_env are mutually exclusive"); + } + Ok(()) + } + + /// Caller serializes updates to an existing VM with its launch lock. + pub(crate) fn write_env(&self, encrypted: &[u8], sealed: &[u8]) -> Result<()> { + Self::validate_env_update(encrypted, sealed)?; + let (path, other, data) = if !sealed.is_empty() { + (self.sealed_env_path(), self.encrypted_env_path(), sealed) + } else if !encrypted.is_empty() { + (self.encrypted_env_path(), self.sealed_env_path(), encrypted) + } else { + return Ok(()); + }; + safe_write::safe_write(path, data).context("failed to write env")?; + match fs::remove_file(other) { + Ok(()) => Ok(()), + Err(err) if err.kind() == std::io::ErrorKind::NotFound => Ok(()), + Err(err) => Err(err).context("failed to remove superseded env"), + } + } + pub fn encrypted_env_path(&self) -> PathBuf { self.shared_dir().join(ENCRYPTED_ENV) } + pub fn sealed_env_path(&self) -> PathBuf { + self.shared_dir().join(SEALED_ENV) + } + pub fn instance_info_path(&self) -> PathBuf { self.shared_dir().join(INSTANCE_INFO) } @@ -389,3 +420,39 @@ mod tests { workdir.remove_all().unwrap(); } } + +#[cfg(test)] +mod env_tests { + use super::*; + + #[test] + fn env_updates_replace_the_other_format_and_reject_ambiguity() { + let dir = tempfile::tempdir().unwrap(); + let workdir = VmWorkDir::new(dir.path()); + fs::create_dir_all(workdir.shared_dir()).unwrap(); + workdir.write_env(b"legacy", &[]).unwrap(); + workdir.write_env(&[], b"sealed").unwrap(); + assert!(!workdir.encrypted_env_path().exists()); + assert_eq!(fs::read(workdir.sealed_env_path()).unwrap(), b"sealed"); + assert!(workdir.write_env(b"ambiguous", b"ambiguous").is_err()); + assert_eq!(fs::read(workdir.sealed_env_path()).unwrap(), b"sealed"); + workdir.write_env(&[], &[]).unwrap(); + assert_eq!(fs::read(workdir.sealed_env_path()).unwrap(), b"sealed"); + workdir.write_env(b"owner", &[]).unwrap(); + assert!(!workdir.sealed_env_path().exists()); + assert_eq!(fs::read(workdir.encrypted_env_path()).unwrap(), b"owner"); + } + + #[test] + fn env_write_failure_preserves_original_and_cleanup_errors_propagate() { + let dir = tempfile::tempdir().unwrap(); + let workdir = VmWorkDir::new(dir.path()); + fs::create_dir_all(workdir.shared_dir()).unwrap(); + workdir.write_env(b"original", &[]).unwrap(); + fs::create_dir(workdir.sealed_env_path()).unwrap(); + assert!(workdir.write_env(&[], b"replacement").is_err()); + assert_eq!(fs::read(workdir.encrypted_env_path()).unwrap(), b"original"); + assert!(workdir.write_env(b"updated", &[]).is_err()); + assert_eq!(fs::read(workdir.encrypted_env_path()).unwrap(), b"updated"); + } +} diff --git a/dstack/vmm/src/config.rs b/dstack/vmm/src/config.rs index c923db519..4f76e742a 100644 --- a/dstack/vmm/src/config.rs +++ b/dstack/vmm/src/config.rs @@ -1432,6 +1432,19 @@ impl Config { mod tests { use super::*; + #[test] + fn host_api_limits_sealed_env_updates() { + use rocket::data::ToByteUnit; + use rocket::figment::providers::Serialized; + // Built the way `run_host_api` builds the host API server's config. + let figment = load_config_figment(None); + let figment = figment.clone().merge(Serialized::defaults( + figment.find_value("host_api").unwrap(), + )); + let config: rocket::Config = figment.extract().unwrap(); + assert_eq!(config.limits.get("UpdateSealedEnv"), Some(1.mebibytes())); + } + #[test] fn image_dir_stays_a_direct_child_of_the_image_root() { let cfg = ImageConfig { diff --git a/dstack/vmm/src/host_api_service.rs b/dstack/vmm/src/host_api_service.rs index fc4dcc704..0ad3c53eb 100644 --- a/dstack/vmm/src/host_api_service.rs +++ b/dstack/vmm/src/host_api_service.rs @@ -46,6 +46,12 @@ impl HostApiRpc for HostApiHandler { .vm_event_report(self.endpoint.cid, &request.event, request.payload) } + async fn update_sealed_env(self, request: host_api::SealedEnv) -> Result<()> { + self.app + .update_sealed_env(self.endpoint.cid, request.sealed_env) + .await + } + async fn get_sealing_key(self, request: GetSealingKeyRequest) -> Result { let key_provider = &self.app.config.key_provider; if !key_provider.enabled { diff --git a/dstack/vmm/src/main_service.rs b/dstack/vmm/src/main_service.rs index ffd06ad61..2cd917547 100644 --- a/dstack/vmm/src/main_service.rs +++ b/dstack/vmm/src/main_service.rs @@ -14,8 +14,9 @@ use dstack_vmm_rpc::vmm_server::{VmmRpc, VmmServer}; use dstack_vmm_rpc::{ AppId, ComposeHash as RpcComposeHash, GatewaySettings, GetInfoResponse, GetMetaResponse, Id, ImageInfo as RpcImageInfo, ImageListResponse, KmsSettings, ListGpusResponse, PublicKeyResponse, - ReloadVmsResponse, ResizeVmRequest, ResourcesSettings, StatusRequest, StatusResponse, - SvListResponse, SvProcessInfo, UpdateVmRequest, VersionResponse, VmConfiguration, + ReloadVmsResponse, ResizeVmRequest, ResourcesSettings, SealedEnvPubKeyResponse, StatusRequest, + StatusResponse, SvListResponse, SvProcessInfo, UpdateVmRequest, VersionResponse, + VmConfiguration, }; use fs_err as fs; use path_absolutize::Absolutize; @@ -311,6 +312,7 @@ pub fn create_manifest_from_vm_config( request: VmConfiguration, cvm_config: &crate::config::CvmConfig, ) -> Result { + VmWorkDir::validate_env_update(&request.encrypted_env, &request.sealed_env)?; validate_label(&request.name)?; // Same checks as `validate_resize_request`. if request.vcpu == 0 { @@ -1092,6 +1094,7 @@ impl VmmRpc for RpcHandler { // deleted that directory would recreate it holding nothing but a // manifest: invisible to `list_vms`, unloadable at every start, and // claiming the VM's netd interfaces against collection forever. + VmWorkDir::validate_env_update(&request.encrypted_env, &request.sealed_env)?; let _launch = self.app.launch_lock(&request.id).await; // Again under the lock: removal can have claimed the VM while this // waited for it. @@ -1207,11 +1210,7 @@ impl VmmRpc for RpcHandler { fs::write(self.compose_file_path(&request.id)?, &request.compose_file) .context("Failed to write compose file")?; } - if !request.encrypted_env.is_empty() { - let encrypted_env_path = self.encrypted_env_path(&request.id)?; - fs::write(encrypted_env_path, &request.encrypted_env) - .context("Failed to write encrypted env")?; - } + vm_work_dir.write_env(&request.encrypted_env, &request.sealed_env)?; if !request.user_config.is_empty() { let user_config_path = self.user_config_path(&request.id)?; fs::write(user_config_path, &request.user_config) @@ -1233,6 +1232,7 @@ impl VmmRpc for RpcHandler { let response = kms .get_app_env_encrypt_pub_key(dstack_kms_rpc::AppId { app_id: request.app_id, + root_epoch: None, }) .await?; Ok(PublicKeyResponse { @@ -1243,6 +1243,22 @@ impl VmmRpc for RpcHandler { }) } + async fn get_sealed_env_pub_key(self, request: AppId) -> Result { + let kms = self.kms_client()?; + let response = kms + .get_sealed_env_pub_key(dstack_kms_rpc::AppId { + app_id: request.app_id, + root_epoch: request.root_epoch, + }) + .await?; + Ok(SealedEnvPubKeyResponse { + public_key: response.public_key, + root_epoch: response.root_epoch, + timestamp: response.timestamp, + signature: response.signature, + }) + } + async fn get_info(self, request: Id) -> Result { info!(vm_id = %request.id, "get_info RPC called"); if let Some(vm) = self.app.vm_info(&request.id).await? { @@ -1497,6 +1513,7 @@ mod tests { disk_size: 10, ports: vec![], encrypted_env: vec![], + sealed_env: vec![], app_id: None, user_config: String::new(), hugepages: false, diff --git a/dstack/vmm/src/one_shot.rs b/dstack/vmm/src/one_shot.rs index 4aa36741b..9989472ea 100644 --- a/dstack/vmm/src/one_shot.rs +++ b/dstack/vmm/src/one_shot.rs @@ -216,10 +216,7 @@ Compose file content (first 200 chars): .context("Failed to write app compose file")?; // Write other files if present - if !vm_config.encrypted_env.is_empty() { - fs_err::write(vm_work_dir.encrypted_env_path(), &vm_config.encrypted_env) - .context("Failed to write encrypted env")?; - } + vm_work_dir.write_env(&vm_config.encrypted_env, &vm_config.sealed_env)?; if !vm_config.user_config.is_empty() { fs_err::write(vm_work_dir.user_config_path(), &vm_config.user_config) diff --git a/dstack/vmm/src/tests/test_vmm_cli.py b/dstack/vmm/src/tests/test_vmm_cli.py index f3c33c371..23564bad2 100644 --- a/dstack/vmm/src/tests/test_vmm_cli.py +++ b/dstack/vmm/src/tests/test_vmm_cli.py @@ -5,12 +5,20 @@ from __future__ import annotations +import argparse +import base64 import contextlib import importlib.util import io +import json import sys import unittest from pathlib import Path +from unittest.mock import Mock, patch + +from eth_keys import keys +from eth_utils import keccak +from pyhpke import AEADId, CipherSuite, KDFId, KEMId def load_vmm_cli(): @@ -55,5 +63,261 @@ def test_kms_urls_are_sent_to_upgrade_app(self) -> None: ) +class SealedEnvTests(unittest.TestCase): + def setUp(self): + self.app_id = "11" * 20 + self.epoch = 3 + self.timestamp = 1700000000 + self.clock = patch.object(VMM_CLI.time, "time", return_value=self.timestamp) + self.clock.start() + self.addCleanup(self.clock.stop) + self.signer = keys.PrivateKey(b"\x22" * 32) + self.pin = self.signer.public_key.to_compressed_bytes().hex() + self.suite = CipherSuite.new( + KEMId.DHKEM_X25519_HKDF_SHA256, KDFId.HKDF_SHA256, AEADId.AES256_GCM + ) + self.recipient = self.suite.kem.derive_key_pair(b"test recipient key material") + self.response = { + "public_key": self.recipient.public_key.to_public_bytes().hex(), + "root_epoch": self.epoch, + "timestamp": self.timestamp, + } + digest = keccak( + b"dstack-sealed-env-pubkey-v1:" + + bytes.fromhex(self.app_id) + + self.epoch.to_bytes(4, "big") + + bytes.fromhex(self.response["public_key"]) + + self.timestamp.to_bytes(8, "big") + ) + self.response["signature"] = self.signer.sign_msg_hash(digest).to_bytes().hex() + self.cli = VMM_CLI.VmmCLI("https://vmm.example") + self.cli.client.request = Mock(return_value=(200, self.response)) + self.cli.get_app_env_encrypt_pub_key = Mock( + return_value=self.response["public_key"] + ) + self.pins = patch.object(VMM_CLI, "load_whitelist", return_value=[self.pin]) + self.pins.start() + self.addCleanup(self.pins.stop) + self.envs = {"SECRET": "synthetic test value", "UNICODE": "测试"} + + def decrypt(self, payload): + self.assertNotIn("encrypted_env", payload) + envelope = json.loads(bytes.fromhex(payload["sealed_env"])) + self.assertEqual(set(envelope), {"version", "root_epoch", "enc", "ciphertext"}) + self.assertEqual(envelope["version"], 1) + self.assertEqual(envelope["root_epoch"], self.epoch) + enc = base64.b64decode(envelope["enc"], validate=True) + ciphertext = base64.b64decode(envelope["ciphertext"], validate=True) + receiver = self.suite.create_recipient_context( + enc, + self.recipient.private_key, + info=f"dstack-sealed-env-v1:{self.epoch}".encode(), + ) + plaintext = json.loads(receiver.open(ciphertext, aad=b"")) + self.assertEqual( + plaintext, {"env": [{"key": k, "value": v} for k, v in self.envs.items()]} + ) + wrong_epoch = self.suite.create_recipient_context( + enc, self.recipient.private_key, info=b"dstack-sealed-env-v1:0" + ) + with self.assertRaises(Exception): + wrong_epoch.open(ciphertext, aad=b"") + + def test_sealed_payload_round_trip_and_normalized_pins(self): + for pin in [self.pin, "0x" + self.pin.upper()]: + with ( + self.subTest(pin=pin), + patch.object(VMM_CLI, "load_whitelist", return_value=[pin]), + ): + self.decrypt(self.cli.encrypt_env_payload(self.envs, self.app_id, None)) + self.cli.get_app_env_encrypt_pub_key.assert_not_called() + + def test_unpinned_signer_never_becomes_trusted(self): + for pins in [ + [], + [keys.PrivateKey(b"\x33" * 32).public_key.to_compressed_bytes().hex()], + ]: + with ( + self.subTest(pins=pins), + patch.object(VMM_CLI, "load_whitelist", return_value=pins), + patch.object(self.cli, "confirm_untrusted_signer") as confirm, + ): + with self.assertRaisesRegex(ValueError, "not pinned"): + self.cli.encrypt_env_payload(self.envs, self.app_id, None) + confirm.assert_not_called() + self.cli.get_app_env_encrypt_pub_key.assert_not_called() + + def test_signature_binds_every_field_and_app(self): + cases = [ + ("public_key", "44" * 32), + ("root_epoch", 4), + ("timestamp", self.timestamp + 1), + ("signature", "00" * 65), + ("signature", "00" * 64), + ("public_key", "00" * 31), + ("root_epoch", -1), + ("root_epoch", 2**32), + ("root_epoch", True), + ("timestamp", -1), + ("timestamp", 2**64), + ("timestamp", True), + ] + for field, value in cases: + with self.subTest(field=field, value=value): + self.cli.client.request.return_value = ( + 200, + {**self.response, field: value}, + ) + with self.assertRaises(Exception): + self.cli.encrypt_env_payload(self.envs, self.app_id, None) + self.cli.client.request.return_value = (200, self.response) + for app_id in ["55" * 20, "11" * 19]: + with self.assertRaises(ValueError): + self.cli.encrypt_env_payload(self.envs, app_id, None) + self.cli.get_app_env_encrypt_pub_key.assert_not_called() + + def test_pinned_signed_key_timestamp_window(self): + for age in [-60, 0, 300]: + with self.subTest(age=age), patch.object( + VMM_CLI.time, "time", return_value=self.timestamp + age + ): + self.decrypt(self.cli.encrypt_env_payload(self.envs, self.app_id, None)) + for age, reason in [(-61, "future"), (301, "stale")]: + with self.subTest(age=age), patch.object( + VMM_CLI.time, "time", return_value=self.timestamp + age + ): + with self.assertRaisesRegex(ValueError, reason): + self.cli.encrypt_env_payload(self.envs, self.app_id, None) + self.cli.get_app_env_encrypt_pub_key.assert_not_called() + + def test_requested_epoch_is_sent_and_checked(self): + self.cli.get_sealed_env_pub_key(self.app_id, root_epoch=3) + self.assertEqual( + self.cli.client.request.call_args.kwargs["body"], + {"app_id": self.app_id, "root_epoch": 3}, + ) + with self.assertRaisesRegex(ValueError, "does not match"): + self.cli.get_sealed_env_pub_key(self.app_id, root_epoch=2) + + def test_only_explicit_unsupported_method_falls_back(self): + for status in [400, 404]: + self.cli.client.request.return_value = ( + status, + {"error": "Service not found: GetSealedEnvPubKey"}, + ) + payload = self.cli.encrypt_env_payload(self.envs, self.app_id, None) + self.assertEqual(set(payload), {"encrypted_env"}) + self.cli.get_app_env_encrypt_pub_key.assert_called_with(self.app_id, None) + self.cli.get_app_env_encrypt_pub_key.side_effect = ValueError( + "legacy trust rejected" + ) + with self.assertRaisesRegex(ValueError, "legacy trust rejected"): + self.cli.encrypt_env_payload(self.envs, self.app_id, None) + + def test_errors_and_bad_responses_do_not_downgrade(self): + for status, body in [ + (401, {"error": "Service not found: GetSealedEnvPubKey"}), + (403, {"error": "forbidden"}), + (404, "Not Found"), + (400, {"error": "invalid app"}), + (500, {"error": "unavailable"}), + (200, {}), + (200, "invalid JSON"), + ]: + with self.subTest(status=status, body=body): + self.cli.client.request.return_value = (status, body) + with self.assertRaises(Exception): + self.cli.encrypt_env_payload(self.envs, self.app_id, None) + self.cli.client.request.side_effect = TimeoutError("test timeout") + with self.assertRaises(TimeoutError): + self.cli.encrypt_env_payload(self.envs, self.app_id, None) + self.cli.get_app_env_encrypt_pub_key.assert_not_called() + + def test_direct_kms_verifies_pin_and_does_not_forward_vmm_credentials(self): + with patch.object(VMM_CLI, "VmmClient") as client: + client.return_value.request.return_value = (200, self.response) + self.decrypt( + self.cli.encrypt_env_payload( + self.envs, self.app_id, "https://kms.example" + ) + ) + client.assert_called_once_with("https://kms.example", verify_tls=False) + self.cli.get_app_env_encrypt_pub_key.assert_not_called() + + def test_create_update_and_update_env_send_only_the_negotiated_field(self): + compose = json.dumps({"kms_enabled": True}) + args = argparse.Namespace( + compose="compose.json", + env_file="env", + name="vm", + image="image", + vcpu=1, + memory=1024, + disk=10, + app_id=self.app_id, + user_config=None, + port=None, + hugepages=False, + pin_numa=False, + stopped=False, + no_tee=False, + annotation=None, + disk_prealloc=None, + simulated_tee=None, + swap=None, + ppcie=False, + gpu=None, + kms_url=["https://kms.example"], + gateway_url=None, + net_queues=None, + net=None, + net_vhost=None, + kms_encrypt_url=None, + ) + + def rpc(method, params=None): + if method == "GetInfo": + return { + "found": True, + "info": { + "app_id": self.app_id, + "configuration": {"compose_file": compose}, + }, + } + return {"id": "vm"} + + self.cli.rpc_call = Mock(side_effect=rpc) + with ( + patch.object(VMM_CLI.os.path, "exists", return_value=True), + patch.object(VMM_CLI, "read_utf8", return_value=compose), + patch.object(VMM_CLI, "parse_env_file", return_value=self.envs), + patch.object(VMM_CLI, "VmmClient") as client, + contextlib.redirect_stdout(io.StringIO()), + ): + for supported in [True, False]: + client.return_value.request.return_value = ( + (200, self.response) + if supported + else (404, {"error": "Service not found: GetSealedEnvPubKey"}) + ) + for action in [ + lambda: self.cli.create_vm(args), + lambda: self.cli.update_vm( + "vm", env_file="env", kms_urls=args.kms_url + ), + lambda: self.cli.update_vm_env("vm", self.envs, args.kms_url), + ]: + with self.subTest(supported=supported, action=action): + self.cli.rpc_call.reset_mock() + action() + method, payload = self.cli.rpc_call.call_args.args + self.assertIn(method, ["CreateVm", "UpgradeApp"]) + if supported: + self.decrypt(payload) + else: + self.assertIn("encrypted_env", payload) + self.assertNotIn("sealed_env", payload) + + if __name__ == "__main__": unittest.main() diff --git a/dstack/vmm/src/vmm-cli.py b/dstack/vmm/src/vmm-cli.py index ad469f981..0001c582f 100755 --- a/dstack/vmm/src/vmm-cli.py +++ b/dstack/vmm/src/vmm-cli.py @@ -16,6 +16,7 @@ import ssl import sys import tempfile +import time import urllib.parse from typing import Any, BinaryIO, Dict, List, Optional, Tuple, Union @@ -310,6 +311,77 @@ def cmd_switch_vmm(args): print(f"Switched to VMM {selected['id'][:8]} ({url})") +def seal_env(envs, hex_public_key: str, root_epoch: int) -> str: + """Return hex-encoded SealedEnv JSON using RFC 9180 base mode.""" + from pyhpke import AEADId, CipherSuite, KDFId, KEMId + + validate_uint(root_epoch, 32, "root_epoch") + suite = CipherSuite.new( + KEMId.DHKEM_X25519_HKDF_SHA256, KDFId.HKDF_SHA256, AEADId.AES256_GCM + ) + public_key = suite.kem.deserialize_public_key( + bytes.fromhex(hex_public_key.removeprefix("0x")) + ) + enc, context = suite.create_sender_context( + public_key, info=f"dstack-sealed-env-v1:{root_epoch}".encode("utf-8") + ) + plaintext = json.dumps({"env": envs}).encode("utf-8") + envelope = { + "version": 1, + "root_epoch": root_epoch, + "enc": base64.b64encode(enc).decode("ascii"), + "ciphertext": base64.b64encode(context.seal(plaintext, aad=b"")).decode( + "ascii" + ), + } + return json.dumps(envelope, separators=(",", ":")).encode("utf-8").hex() + + +def validate_uint(value: int, bits: int, name: str) -> None: + if type(value) is not int or not 0 <= value < 1 << bits: + raise ValueError(f"{name} must be a uint{bits}") + + +def verify_sealed_env_pubkey( + response: Dict, app_id: str, root_epoch: Optional[int] = None +) -> None: + """Authenticate a SealedEnv key against the locally pinned epoch signers.""" + if not CRYPTO_AVAILABLE: + raise ImportError("Install the VMM CLI requirements to verify SealedEnv keys") + app_id_bytes = bytes.fromhex(app_id.removeprefix("0x")) + public_key = bytes.fromhex(response["public_key"].removeprefix("0x")) + signature = bytes.fromhex(response["signature"].removeprefix("0x")) + epoch, timestamp = response["root_epoch"], response["timestamp"] + validate_uint(epoch, 32, "root_epoch") + validate_uint(timestamp, 64, "timestamp") + now = int(time.time()) + if now - timestamp > 300: + raise ValueError("SealedEnv public key timestamp is stale") + if timestamp - now > 60: + raise ValueError("SealedEnv public key timestamp is in the future") + if len(app_id_bytes) != 20 or len(public_key) != 32 or len(signature) != 65: + raise ValueError("Invalid SealedEnv app ID, public key or signature length") + if root_epoch is not None: + validate_uint(root_epoch, 32, "root_epoch") + if epoch != root_epoch: + raise ValueError("SealedEnv root epoch does not match the requested epoch") + digest = keccak( + b"dstack-sealed-env-pubkey-v1:" + + app_id_bytes + + epoch.to_bytes(4, "big") + + public_key + + timestamp.to_bytes(8, "big") + ) + signer = keys.Signature(signature_bytes=signature).recover_public_key_from_msg_hash( + digest + ) + pins = {bytes.fromhex(pin.removeprefix("0x")) for pin in load_whitelist()} + if signer.to_compressed_bytes() not in pins: + raise ValueError( + "SealedEnv signer is not pinned; configure a trusted epoch signer with kms add" + ) + + def encrypt_env(envs, hex_public_key: str) -> str: """Encrypts environment variables using a one-time X25519 key exchange and AES-GCM. @@ -776,6 +848,49 @@ def list_images(self, json_output: bool = False) -> None: rows = [[img["name"], img.get("version", "-")] for img in images] print(format_table(rows, headers)) + def get_sealed_env_pub_key( + self, + app_id: str, + kms_url: Optional[str] = None, + root_epoch: Optional[int] = None, + ) -> Optional[Dict]: + """Return an authenticated key, or None only for an unsupported RPC.""" + params = {"app_id": app_id} + if root_epoch is not None: + validate_uint(root_epoch, 32, "root_epoch") + params["root_epoch"] = root_epoch + # KMS RA-TLS is authenticated by the pinned signer below. The VMM + # path retains its configured TLS verification and authentication. + client = VmmClient(kms_url, verify_tls=False) if kms_url else self.client + status, response = client.request( + "POST", + "/prpc/GetSealedEnvPubKey?json", + headers={"Content-Type": "application/json"}, + body=params, + ) + error = response.get("error") if isinstance(response, dict) else response + if status in (400, 404) and error == "Service not found: GetSealedEnvPubKey": + return None + if status != 200: + raise RuntimeError(f"GetSealedEnvPubKey failed with HTTP {status}") + if not isinstance(response, dict): + raise ValueError("Invalid GetSealedEnvPubKey response") + verify_sealed_env_pubkey(response, app_id, root_epoch) + return response + + def encrypt_env_payload( + self, envs: Dict[str, str], app_id: str, kms_url: Optional[str] + ) -> Dict: + """Select the wire field by RPC support, never by environment file contents.""" + envs_list = [{"key": k, "value": v} for k, v in envs.items()] + key = self.get_sealed_env_pub_key(app_id, kms_url) + if key is not None: + return { + "sealed_env": seal_env(envs_list, key["public_key"], key["root_epoch"]) + } + public_key = self.get_app_env_encrypt_pub_key(app_id, kms_url) + return {"encrypted_env": encrypt_env(envs_list, public_key)} + def get_app_env_encrypt_pub_key( self, app_id: str, kms_url: Optional[str] = None ) -> Dict: @@ -1039,10 +1154,7 @@ def create_vm(self, args) -> None: encrypt_url = args.kms_encrypt_url or ( args.kms_url[0] if args.kms_url else None ) - encrypt_pubkey = self.get_app_env_encrypt_pub_key(app_id, encrypt_url) - print(f"Encrypting environment variables with key: {encrypt_pubkey}") - envs_list = [{"key": k, "value": v} for k, v in envs.items()] - params["encrypted_env"] = encrypt_env(envs_list, encrypt_pubkey) + params.update(self.encrypt_env_payload(envs, app_id, encrypt_url)) response = self.rpc_call("CreateVm", params) print(f"Created VM with ID: {response.get('id')}") return response.get("id") @@ -1067,16 +1179,8 @@ def update_vm_env( vm_configuration = vm_info_response["info"].get("configuration") or {} compose_file = vm_configuration.get("compose_file") - # Now get the encryption key for the app - encrypt_pubkey = self.get_app_env_encrypt_pub_key( - app_id, kms_urls[0] if kms_urls else None - ) - print(f"Encrypting environment variables with key: {encrypt_pubkey}") - envs_list = [{"key": k, "value": v} for k, v in envs.items()] - encrypted_env = encrypt_env(envs_list, encrypt_pubkey) - - # Use UpdateApp with the VM ID - payload = {"id": vm_id, "encrypted_env": encrypted_env} + payload = {"id": vm_id} + payload.update(self.encrypt_env_payload(envs, app_id, kms_urls[0])) if compose_file: try: @@ -1238,11 +1342,7 @@ def update_vm( vm_configuration = vm_info_response["info"].get("configuration") or {} compose_file_content = vm_configuration.get("compose_file") - encrypt_pubkey = self.get_app_env_encrypt_pub_key( - app_id, kms_urls[0] if kms_urls else None - ) - envs_list = [{"key": k, "value": v} for k, v in envs.items()] - upgrade_params["encrypted_env"] = encrypt_env(envs_list, encrypt_pubkey) + upgrade_params.update(self.encrypt_env_payload(envs, app_id, kms_urls[0])) updates.append("environment variables") # update allowed_envs in compose file if needed diff --git a/dstack/vmm/ui/build.mjs b/dstack/vmm/ui/build.mjs index 00c06fc5a..85a814731 100644 --- a/dstack/vmm/ui/build.mjs +++ b/dstack/vmm/ui/build.mjs @@ -24,6 +24,9 @@ function canonicalId(absPath) { } function resolveModule(parentId, request) { + if (request === 'crypto') { + return canonicalId(path.join(TS_OUT_DIR, 'lib/browserCrypto.js')); + } const base = parentId ? path.dirname(path.resolve(ROOT, parentId)) : MODULE_DIR; const absPath = nodeRequire.resolve(request, { paths: [base] }); return canonicalId(absPath); @@ -162,6 +165,7 @@ async function compileTypeScript() { await fs.rm(TS_OUT_DIR, { recursive: true, force: true }); await run(TSC, ['--project', path.join(ROOT, 'tsconfig.json')]); await copyDir(path.join(SOURCE_DIR, 'templates'), path.join(TS_OUT_DIR, 'templates')); + await fs.writeFile(path.join(TS_OUT_DIR, 'package.json'), JSON.stringify({ type: 'commonjs' })); } async function build({ watch = false } = {}) { diff --git a/dstack/vmm/ui/package-lock.json b/dstack/vmm/ui/package-lock.json index b04b7ed8d..1d9101ebf 100644 --- a/dstack/vmm/ui/package-lock.json +++ b/dstack/vmm/ui/package-lock.json @@ -7,6 +7,10 @@ "": { "name": "vmm-ui", "version": "0.1.0", + "dependencies": { + "@hpke/core": "1.10.0", + "@hpke/dhkem-x25519": "1.8.0" + }, "devDependencies": { "@types/node": "^20.11.30", "protobufjs": "^7.5.8", @@ -64,6 +68,47 @@ "node": ">=6.9.0" } }, + "node_modules/@hpke/common": { + "version": "1.12.0", + "resolved": "https://registry.npmjs.org/@hpke/common/-/common-1.12.0.tgz", + "integrity": "sha512-os6y44MDJmrx7MVAG8rb6LyV3BEXSFtX/g4KMhKxVtpXvCWwL4fbBbbtG45fkx25JmxpCJdGzla9asLNMI99/A==", + "license": "MIT", + "engines": { + "node": ">=16.0.0" + }, + "funding": { + "type": "github", + "url": "https://github.com/sponsors/dajiaji" + } + }, + "node_modules/@hpke/core": { + "version": "1.10.0", + "resolved": "https://registry.npmjs.org/@hpke/core/-/core-1.10.0.tgz", + "integrity": "sha512-toCyUYXvU5KdDjqtWeMRRhR09R2us4OeHGo5YkdDOpo8+QJ6sIpuuOGLDu0VbNBskf5acHgsqgF53w2mU1NNYQ==", + "license": "MIT", + "dependencies": { + "@hpke/common": "^1.11.0" + }, + "engines": { + "node": ">=16.0.0" + }, + "funding": { + "type": "github", + "url": "https://github.com/sponsors/dajiaji" + } + }, + "node_modules/@hpke/dhkem-x25519": { + "version": "1.8.0", + "resolved": "https://registry.npmjs.org/@hpke/dhkem-x25519/-/dhkem-x25519-1.8.0.tgz", + "integrity": "sha512-S1MWWkAfu+TFxySgv5+2P3O4Mx/jk7BsoplzQaA1s3sfUJVJ2UsZsSzSsMc+FXJumLXncoJFlO6mK6mDGspfmA==", + "license": "MIT", + "dependencies": { + "@hpke/common": "^1.10.0" + }, + "engines": { + "node": ">=16.0.0" + } + }, "node_modules/@jsdoc/salty": { "version": "0.2.9", "resolved": "https://registry.npmjs.org/@jsdoc/salty/-/salty-0.2.9.tgz", diff --git a/dstack/vmm/ui/package.json b/dstack/vmm/ui/package.json index b88e34045..985719998 100644 --- a/dstack/vmm/ui/package.json +++ b/dstack/vmm/ui/package.json @@ -6,12 +6,17 @@ "scripts": { "generate:proto": "bash scripts/build_proto.sh", "build": "node build.mjs", - "watch": "node build.mjs --watch" + "watch": "node build.mjs --watch", + "test": "DSTACK_UI_OUT=dist/index.html npm run build && node --test tests/*.test.cjs" }, "devDependencies": { "@types/node": "^20.11.30", "protobufjs": "^7.5.8", "protobufjs-cli": "^1.3.3", "typescript": "^5.4.5" + }, + "dependencies": { + "@hpke/core": "1.10.0", + "@hpke/dhkem-x25519": "1.8.0" } } diff --git a/dstack/vmm/ui/src/composables/useVmManager.ts b/dstack/vmm/ui/src/composables/useVmManager.ts index f81aede22..5b9fbab32 100644 --- a/dstack/vmm/ui/src/composables/useVmManager.ts +++ b/dstack/vmm/ui/src/composables/useVmManager.ts @@ -4,6 +4,7 @@ declare const Vue: any; const { ref, computed, watch, onMounted } = Vue; import type { vmm as VmmTypes } from '../proto/vmm_rpc'; +import { encryptEnvPayload } from '../lib/envEncryption'; // Types based on Rust definitions type VmConfiguration = VmmTypes.IVmConfiguration; @@ -643,6 +644,7 @@ type CreateVmPayloadSource = { disk_size: number; ports: PortFormEntry[]; encrypted_env?: Uint8Array; + sealed_env?: Uint8Array; app_id?: string | null; user_config?: string; hugepages?: boolean; @@ -670,6 +672,7 @@ type CreateVmPayloadSource = { disk_size: Math.max(0, Number(source.disk_size) || 0), ports: normalizedPorts, encrypted_env: source.encrypted_env, + sealed_env: source.sealed_env, app_id: source.app_id || undefined, user_config: source.user_config || '', hugepages: !!source.hugepages, @@ -1222,7 +1225,7 @@ type CreateVmPayloadSource = { async function encryptEnv(envs: EncryptedEnvEntry[], kmsEnabled: boolean, appId: string | null) { if (!kmsEnabled || envs.length === 0) { - return undefined; + return {}; } let appIdToUse = appId; if (!appIdToUse) { @@ -1230,8 +1233,7 @@ type CreateVmPayloadSource = { appIdToUse = await calcAppId(appCompose); } const keyBytes = hexToBytes(appIdToUse); - const response = await vmmRpc.getAppEnvEncryptPubKey({ app_id: keyBytes }); - return encryptEnvWithKey(envs, response.public_key); + return encryptEnvPayload(vmmRpc, envs, keyBytes, encryptEnvWithKey); } async function encryptEnvWithKey(envs: EncryptedEnvEntry[], publicKeyBytes: Uint8Array) { @@ -1281,7 +1283,7 @@ type CreateVmPayloadSource = { memory: vmForm.value.memory, disk_size: vmForm.value.disk_size, ports: vmForm.value.ports, - encrypted_env: encryptedEnv || undefined, + ...encryptedEnv, app_id: vmForm.value.app_id || undefined, user_config: vmForm.value.user_config, hugepages: vmForm.value.hugepages, @@ -1362,14 +1364,15 @@ type CreateVmPayloadSource = { const composeWasExplicitlyUpdated = updateDialog.value.updateCompose; let composeNeedsUpdate = composeWasExplicitlyUpdated; - let encryptedEnvPayload; + let encryptedEnvPayload = {}; if (updateDialog.value.resetSecrets) { - const keyResponse = await vmmRpc.getAppEnvEncryptPubKey({ app_id: hexToBytes(vm.app_id || '') }); - encryptedEnvPayload = await encryptEnvWithKey(updateDialog.value.encryptedEnvs, keyResponse.public_key); + encryptedEnvPayload = await encryptEnvPayload( + vmmRpc, updateDialog.value.encryptedEnvs, hexToBytes(vm.app_id || ''), encryptEnvWithKey, + ); composeNeedsUpdate = true; } body.compose_file = composeNeedsUpdate ? await makeUpdateComposeFile() : undefined; - body.encrypted_env = encryptedEnvPayload; + Object.assign(body, encryptedEnvPayload); body.user_config = updated.user_config; body.update_ports = true; body.ports = normalizePorts(updated.ports); diff --git a/dstack/vmm/ui/src/lib/browserCrypto.ts b/dstack/vmm/ui/src/lib/browserCrypto.ts new file mode 100644 index 000000000..37da2d201 --- /dev/null +++ b/dstack/vmm/ui/src/lib/browserCrypto.ts @@ -0,0 +1,6 @@ +// SPDX-FileCopyrightText: © 2026 Phala Network +// +// SPDX-License-Identifier: Apache-2.0 + +// HPKE's Node fallback uses the browser's native Web Crypto in this bundle. +export const webcrypto = globalThis.crypto; diff --git a/dstack/vmm/ui/src/lib/envEncryption.ts b/dstack/vmm/ui/src/lib/envEncryption.ts new file mode 100644 index 000000000..4a29d3f8c --- /dev/null +++ b/dstack/vmm/ui/src/lib/envEncryption.ts @@ -0,0 +1,68 @@ +// SPDX-FileCopyrightText: © 2026 Phala Network +// +// SPDX-License-Identifier: Apache-2.0 + +import { Aes256Gcm, CipherSuite, HkdfSha256 } from '@hpke/core'; +import { DhkemX25519HkdfSha256 } from '@hpke/dhkem-x25519'; +import type { vmm as VmmTypes } from '../proto/vmm_rpc'; +import { RpcError } from './vmmRpcClient'; + +export type EnvEntry = { key: string; value: string }; +type EnvPayload = { sealed_env?: Uint8Array; encrypted_env?: Uint8Array }; +type KeyClient = { + getSealedEnvPubKey(request: VmmTypes.IAppId): Promise; + getAppEnvEncryptPubKey(request: VmmTypes.IAppId): Promise; +}; + +const missingService = 'Service not found: GetSealedEnvPubKey'; + +function isUnsupported(error: unknown): boolean { + if (!(error instanceof RpcError) || ![400, 404].includes(error.status)) return false; + if (error.message === missingService) return true; + // The VMM forwards an older KMS's JSON RPC error through its binary RPC. + const forwarded = /^Request failed with status=(?:400 Bad Request|404 Not Found), error=(.+)$/s.exec(error.message); + if (!forwarded) return false; + try { + return JSON.parse(forwarded[1]).error === missingService; + } catch { + return false; + } +} + +/** Keep key retrieval on the authenticated, same-origin VMM transport. */ +export async function encryptEnvPayload( + client: KeyClient, + envs: EnvEntry[], + appId: Uint8Array, + encryptLegacy: (envs: EnvEntry[], key: Uint8Array) => Promise, +): Promise { + let response: VmmTypes.ISealedEnvPubKeyResponse; + try { + response = await client.getSealedEnvPubKey({ app_id: appId }); + } catch (error) { + if (!isUnsupported(error)) throw error; + const legacy = await client.getAppEnvEncryptPubKey({ app_id: appId }); + return { encrypted_env: await encryptLegacy(envs, legacy.public_key) }; + } + const epoch = response.root_epoch; + if (!Number.isInteger(epoch) || epoch < 0 || epoch > 0xffffffff || response.public_key?.length !== 32) { + throw new Error('invalid sealed env public key response'); + } + const suite = new CipherSuite({ + kem: new DhkemX25519HkdfSha256(), + kdf: new HkdfSha256(), + aead: new Aes256Gcm(), + }); + const recipientPublicKey = await suite.kem.deserializePublicKey(new Uint8Array(response.public_key)); + const { enc, ct } = await suite.seal({ + recipientPublicKey, + info: new TextEncoder().encode(`dstack-sealed-env-v1:${epoch}`), + }, new TextEncoder().encode(JSON.stringify({ env: envs }))); + const base64 = (bytes: ArrayBuffer) => btoa(Array.from(new Uint8Array(bytes), (byte) => String.fromCharCode(byte)).join('')); + return { sealed_env: new TextEncoder().encode(JSON.stringify({ + version: 1, + root_epoch: epoch, + enc: base64(enc), + ciphertext: base64(ct), + })) }; +} diff --git a/dstack/vmm/ui/src/lib/vmmRpcClient.ts b/dstack/vmm/ui/src/lib/vmmRpcClient.ts index 683f499fd..e1ae87505 100644 --- a/dstack/vmm/ui/src/lib/vmmRpcClient.ts +++ b/dstack/vmm/ui/src/lib/vmmRpcClient.ts @@ -8,6 +8,13 @@ const textDecoder = new TextDecoder(); const EMPTY_BODY = new Uint8Array(); let cachedClient: any; +export class RpcError extends Error { + constructor(message: string, readonly status: number) { + super(message); + this.name = 'RpcError'; + } +} + function decodePrpcError(buffer: Uint8Array) { try { if (buffer && buffer.length > 0) { @@ -76,7 +83,7 @@ export function getVmmRpcClient(basePath = '/prpc') { .then(async (response) => { const buffer = new Uint8Array(await response.arrayBuffer()); if (!response.ok) { - callback(new Error(decodePrpcError(buffer))); + callback(new RpcError(decodePrpcError(buffer), response.status)); return; } callback(null, buffer); diff --git a/dstack/vmm/ui/tests/envEncryption.test.cjs b/dstack/vmm/ui/tests/envEncryption.test.cjs new file mode 100644 index 000000000..3ee0e4d12 --- /dev/null +++ b/dstack/vmm/ui/tests/envEncryption.test.cjs @@ -0,0 +1,215 @@ +// SPDX-FileCopyrightText: © 2026 Phala Network +// +// SPDX-License-Identifier: Apache-2.0 + +const { test } = require('node:test'); +const assert = require('node:assert/strict'); +const { readFileSync } = require('node:fs'); +const { runInNewContext } = require('node:vm'); +const { webcrypto } = require('node:crypto'); +const { CipherSuite, Aes256Gcm, HkdfSha256 } = require('@hpke/core'); +const { DhkemX25519HkdfSha256 } = require('@hpke/dhkem-x25519'); +const { encryptEnvPayload } = require('../build/ts/lib/envEncryption.js'); +const { RpcError } = require('../build/ts/lib/vmmRpcClient.js'); +const { vmm } = require('../build/ts/proto/vmm_rpc.js'); +const { prpc } = require('../build/ts/proto/prpc.js'); + +const suite = new CipherSuite({ kem: new DhkemX25519HkdfSha256(), kdf: new HkdfSha256(), aead: new Aes256Gcm() }); +const appId = new Uint8Array(20).fill(0x11); +const envs = [{ key: 'GREETING', value: 'Hello 世界' }]; +const missing = 'Service not found: GetSealedEnvPubKey'; + +async function recipient() { + const pair = await suite.kem.generateKeyPair(); + const public_key = new Uint8Array(await suite.kem.serializePublicKey(pair.publicKey)); + return { pair, response: { public_key, root_epoch: 7, timestamp: 123, signature: new Uint8Array(65) } }; +} + +async function open(pair, bytes, epoch) { + const envelope = JSON.parse(new TextDecoder().decode(bytes)); + assert.equal(envelope.version, 1); + const pt = await suite.open({ + recipientKey: pair.privateKey, + enc: Buffer.from(envelope.enc, 'base64'), + info: new TextEncoder().encode(`dstack-sealed-env-v1:${epoch ?? envelope.root_epoch}`), + }, Buffer.from(envelope.ciphertext, 'base64')); + return JSON.parse(new TextDecoder().decode(pt)); +} + +const noLegacy = async () => { throw new Error('unexpected legacy fallback'); }; + +test('seals JSON bytes with HPKE, binds epoch, and rejects tampering', async () => { + const { pair, response } = await recipient(); + for (const epoch of [0, 7, 0xffffffff]) { + const payload = await encryptEnvPayload({ + getSealedEnvPubKey: async (request) => { + assert.deepEqual(request, { app_id: appId }); + return { ...response, root_epoch: epoch }; + }, + getAppEnvEncryptPubKey: noLegacy, + }, envs, appId, noLegacy); + assert.deepEqual(Object.keys(payload), ['sealed_env']); + assert.deepEqual(await open(pair, payload.sealed_env), { env: envs }); + await assert.rejects(open(pair, payload.sealed_env, epoch === 0 ? 1 : 0)); + const modified = JSON.parse(new TextDecoder().decode(payload.sealed_env)); + const ct = Buffer.from(modified.ciphertext, 'base64'); + ct[0] ^= 1; + modified.ciphertext = ct.toString('base64'); + await assert.rejects(open(pair, new TextEncoder().encode(JSON.stringify(modified)))); + } +}); + +test('falls back only on explicit missing RPC, including a forwarded KMS error', async () => { + for (const error of [ + new RpcError(missing, 404), new RpcError(missing, 400), + new RpcError(`Request failed with status=404 Not Found, error=${JSON.stringify({ error: missing })}`, 400), + ]) { + let legacyCalls = 0; + const payload = await encryptEnvPayload({ + getSealedEnvPubKey: async () => { throw error; }, + getAppEnvEncryptPubKey: async (request) => { + legacyCalls++; + assert.deepEqual(request, { app_id: appId }); + return { public_key: new Uint8Array(32) }; + }, + }, envs, appId, async (entries, key) => { + assert.deepEqual(entries, envs); + assert.equal(key.length, 32); + return new Uint8Array([1, 2, 3]); + }); + assert.equal(legacyCalls, 1); + assert.deepEqual(Object.keys(payload), ['encrypted_env']); + } +}); + +test('network, HTTP, authentication and unrelated service errors never downgrade', async () => { + for (const error of [ + new Error('network failed'), new Error(missing), new RpcError('Not Found', 404), + new RpcError(missing, 401), new RpcError(missing, 403), new RpcError(missing, 500), + new RpcError('Service not found: DifferentMethod', 404), + new RpcError('failed to validate the server attestation', 400), + new RpcError(`Request failed with status=500 Internal Server Error, error={"error":"${missing}"}`, 400), + ]) { + await assert.rejects(encryptEnvPayload({ + getSealedEnvPubKey: async () => { throw error; }, + getAppEnvEncryptPubKey: noLegacy, + }, envs, appId, noLegacy), (actual) => actual === error); + } +}); + +test('invalid key or epoch never downgrades after a successful key RPC', async () => { + const { response } = await recipient(); + for (const invalid of [ + { public_key: new Uint8Array(31) }, { public_key: new Uint8Array(33) }, + { public_key: new Uint8Array(32) }, + { root_epoch: -1 }, { root_epoch: 1.5 }, { root_epoch: 0x100000000 }, + ]) { + let legacyCalls = 0; + await assert.rejects(encryptEnvPayload({ + getSealedEnvPubKey: async () => ({ ...response, ...invalid }), + getAppEnvEncryptPubKey: async () => { legacyCalls++; return response; }, + }, envs, appId, noLegacy)); + assert.equal(legacyCalls, 0); + } +}); + +// Execute the shipped bundle with browser Web Crypto and protobuf transport. +// Vue lifecycle hooks are inert so tests do not start background polling. +function consoleApp(response, unsupported = false) { + const calls = []; + const alerts = []; + let manager; + const html = readFileSync(new URL('../dist/index.html', `file://${__filename}`), 'utf8'); + const scripts = [...html.matchAll(/