From 4b05277ce42db66c23ae70243018ea8b3d63db42 Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Thu, 5 Nov 2026 21:43:56 -0800 Subject: [PATCH 1/9] 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 2/9] 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 3/9] 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 4/9] 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 5/9] 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 6/9] 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 7/9] 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 8/9] 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 9/9] 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