From 0e6ff360c04f9ecacbcdaf78b7906a41196a900d Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Sat, 7 Nov 2026 00:05:46 -0800 Subject: [PATCH 01/17] feat: seal environment variables with epoch-aware HPKE Signed-off-by: Kevin Wang --- docs/deployment.md | 5 +- docs/sealed-env.md | 98 ++++++ dstack/Cargo.lock | 202 ++++++++++-- dstack/Cargo.toml | 1 + dstack/dstack-types/src/lib.rs | 2 + dstack/dstack-types/src/shared_filenames.rs | 1 + dstack/dstack-util/Cargo.toml | 1 + dstack/dstack-util/src/crypto.rs | 130 ++++++++ dstack/dstack-util/src/host_api.rs | 18 ++ dstack/dstack-util/src/main.rs | 3 + dstack/dstack-util/src/system_setup.rs | 306 ++++++++++++++++-- dstack/guest-agent/src/rpc_service.rs | 1 + dstack/host-api/proto/host_api.proto | 9 + dstack/kms/Cargo.toml | 1 + dstack/kms/rpc/proto/kms_rpc.proto | 11 + dstack/kms/src/main_service.rs | 72 ++++- dstack/vmm/requirements.txt | 3 +- dstack/vmm/rpc/proto/vmm_rpc.proto | 13 + dstack/vmm/src/app.rs | 173 +++++++++- dstack/vmm/src/app/vm_info.rs | 1 + dstack/vmm/src/app/workdir.rs | 69 +++- dstack/vmm/src/config.rs | 13 + dstack/vmm/src/host_api_service.rs | 6 + dstack/vmm/src/main_service.rs | 31 +- dstack/vmm/src/one_shot.rs | 5 +- dstack/vmm/src/tests/test_vmm_cli.py | 264 +++++++++++++++ dstack/vmm/src/vmm-cli.py | 138 ++++++-- dstack/vmm/ui/build.mjs | 4 + dstack/vmm/ui/package-lock.json | 45 +++ dstack/vmm/ui/package.json | 7 +- dstack/vmm/ui/src/composables/useVmManager.ts | 19 +- dstack/vmm/ui/src/lib/browserCrypto.ts | 6 + dstack/vmm/ui/src/lib/envEncryption.ts | 68 ++++ dstack/vmm/ui/src/lib/vmmRpcClient.ts | 9 +- dstack/vmm/ui/tests/envEncryption.test.cjs | 215 ++++++++++++ dstack/vmm/vmm.toml | 4 + sdk/fixtures/sealed-env/javascript.json | 17 + sdk/fixtures/sealed-env/python.json | 17 + sdk/go/README.md | 42 +++ sdk/go/dstack/seal_env.go | 64 ++++ sdk/go/dstack/seal_env_test.go | 115 +++++++ sdk/go/dstack/verify_sealed_env_test.go | 123 +++++++ sdk/go/dstack/verify_signature.go | 77 ++++- sdk/go/go.mod | 1 + sdk/go/go.sum | 2 + sdk/js/README.md | 22 ++ sdk/js/bun.lock | 8 + sdk/js/package.json | 10 + sdk/js/src/__tests__/sealed-env.test.ts | 51 +++ sdk/js/src/index.ts | 4 +- sdk/js/src/seal-env.ts | 46 +++ sdk/js/src/verify-env-encrypt-public-key.ts | 32 ++ sdk/js/tsup.config.ts | 1 + sdk/python/README.md | 11 + sdk/python/pdm.lock | 18 +- sdk/python/pyproject.toml | 1 + sdk/python/src/dstack_sdk/__init__.py | 4 + sdk/python/src/dstack_sdk/seal_env.py | 34 ++ .../verify_env_encrypt_public_key.py | 31 ++ sdk/python/tests/test_seal_env.py | 63 ++++ sdk/rust/Cargo.lock | 152 ++++++++- sdk/rust/Cargo.toml | 4 + sdk/rust/README.md | 23 ++ sdk/rust/src/lib.rs | 5 + sdk/rust/src/sealed_env.rs | 240 ++++++++++++++ 65 files changed, 3051 insertions(+), 121 deletions(-) create mode 100644 docs/sealed-env.md create mode 100644 dstack/vmm/ui/src/lib/browserCrypto.ts create mode 100644 dstack/vmm/ui/src/lib/envEncryption.ts create mode 100644 dstack/vmm/ui/tests/envEncryption.test.cjs create mode 100644 sdk/fixtures/sealed-env/javascript.json create mode 100644 sdk/fixtures/sealed-env/python.json create mode 100644 sdk/go/dstack/seal_env.go create mode 100644 sdk/go/dstack/seal_env_test.go create mode 100644 sdk/go/dstack/verify_sealed_env_test.go create mode 100644 sdk/js/src/__tests__/sealed-env.test.ts create mode 100644 sdk/js/src/seal-env.ts create mode 100644 sdk/python/src/dstack_sdk/seal_env.py create mode 100644 sdk/python/tests/test_seal_env.py create mode 100644 sdk/rust/src/sealed_env.rs diff --git a/docs/deployment.md b/docs/deployment.md index d68fdea44..7917caac8 100644 --- a/docs/deployment.md +++ b/docs/deployment.md @@ -552,7 +552,10 @@ for the full procedure; the KMS steps are: `PopRootEpoch` (or remove the `rootEpochs` entry) and rotate again; nodes that imported the dropped N must be onboarded again. 6. Boot CVMs on an OS that supports root epochs. At boot they re-key their data - disk from the epoch recorded in its LUKS header to the current one. + disk from the epoch recorded in its LUKS header to the current one. If an + encrypted environment still uses an older epoch, the CVM re-encrypts it during + boot and the VMM replaces the ciphertext only if it was not changed concurrently. + See [SealedEnv](./sealed-env.md) for the new format, file precedence, and rollback requirements. Each `authorizeKms` or `authorizeApp` result carries the current epoch and its chain hash, including when authorization is denied. KMS verifies these against its keyring; diff --git a/docs/sealed-env.md b/docs/sealed-env.md new file mode 100644 index 000000000..0f02709bb --- /dev/null +++ b/docs/sealed-env.md @@ -0,0 +1,98 @@ +# Sealed environment variables + +SealedEnv protects environment variables with HPKE and identifies the KMS root +key epoch needed to decrypt them. Plaintext remains inside the client and CVM; +the VMM stores and forwards ciphertext. + +## Wire format + +The `.sealed-env` file contains UTF-8 JSON: + +```json +{"version":1,"root_epoch":3,"enc":"","ciphertext":""} +``` + +Version 1 fixes RFC 9180 Base mode with DHKEM(X25519, HKDF-SHA256), +HKDF-SHA256, and AES-256-GCM. `enc` is a 32-byte encapsulated key; `ciphertext` +includes the 16-byte authentication tag. Binary fields use standard padded +Base64. HPKE `info` is UTF-8 `dstack-sealed-env-v1:N`, where `N` is the canonical +decimal `root_epoch` uint32, with no leading zeros. AAD is empty. Unknown envelope +versions are rejected; extra JSON fields may be ignored. This provides +confidentiality, not authentication of the person uploading the environment. + +Plaintext uses the existing JSON representation: + +```json +{"env":[{"key":"MESSAGE","value":"Hello"}]} +``` + +The guest still enforces the application's environment-name allowlist. + +## Keys and signatures + +`GetSealedEnvPubKey` is available on KMS and proxied by VMM. Its request contains +`app_id` (20 bytes) and optional `root_epoch`; omission selects the current +active epoch, while future epochs are rejected. The response contains +`public_key`, `root_epoch`, `timestamp`, and `signature`. + +For each app, the KMS derives the 32-byte receiver private key directly from +that epoch's P-256 CA private scalar using HKDF-SHA256: + +- salt: ASCII `dstack-sealed-env-v1`; +- IKM: the 32-byte P-256 private scalar; +- info: the raw 20-byte app ID. + +`GetAppKey` returns it as `sealed_env_key` alongside the legacy key. The legacy +key derivation and `GetAppEnvEncryptPubKey` remain unchanged. + +The public-key response signs Keccak256 of this unambiguous byte sequence: + +```text +"dstack-sealed-env-pubkey-v1:" || app_id[20] || epoch[4, big-endian] + || public_key[32] || timestamp[8, big-endian] +``` + +The signing key is the selected epoch's secp256k1 root. The signature is +`r[32] || s[32] || recovery_id[1]`. Clients must compare the recovered signer +against an independently trusted epoch root key or pin. Recovering a signer +alone is not a trust decision. Clients should reject stale responses (default +maximum age 300 seconds, future clock tolerance 60 seconds). + +## File precedence and migration + +A new guest chooses `.sealed-env` whenever that file exists. An empty, invalid, +or unreadable sealed file is an error; it must not fall back to a legacy file. +If no sealed file exists, the guest accepts `.encrypted-env` for compatibility. +The copy limits are 384 KiB for sealed JSON and 256 KiB for legacy ciphertext. + +At boot, sealed ciphertext explicitly selects its historical root epoch. The +guest obtains that epoch's app key over the already trusted KMS connection, +decrypts, and seals to the current epoch. Legacy ciphertext uses the legacy +key path, trying earlier epochs when necessary, and is converted to SealedEnv. +A future epoch or failed authentication stops environment provisioning. + +The guest sends `UpdateSealedEnv` through the host channel before application +code runs. The VMM binds the update to the VM's CID and a snapshot of both env +files, including presence. It accepts at most one update in the boot window; +an empty update closes the window without replacing files. Owner updates and +guest replacement share the VM launch lock so an owner change cannot be lost +between the snapshot comparison and write. + +VM creation and update accept either `sealed_env` or legacy `encrypted_env`. +Supplying both nonempty fields is rejected. Writing one format removes the other: +an external legacy update after a guest sealed update therefore becomes the +sole environment for the next boot. Omitting both preserves existing files on +update. A replacement is safely written before the other file is removed; +filesystem errors are reported rather than ignored. + +Old OS images understand only `.encrypted-env`. After conversion removes that +file, reverting to an old OS requires the owner to upload legacy ciphertext +again. Rotation does not revoke credentials exposed before migration, and +cannot erase old host snapshots. + +## Scope of verification + +Local tests exercise HPKE implementations with synthetic keys, signature +contexts, file precedence and update races. They do not establish production +TEE isolation or real-chain availability. The shared fixture under +`sdk/fixtures/sealed-env/` is public test data, never deployment key material. diff --git a/dstack/Cargo.lock b/dstack/Cargo.lock index 7f86192db..cbdbc35f2 100644 --- a/dstack/Cargo.lock +++ b/dstack/Cargo.lock @@ -27,6 +27,16 @@ dependencies = [ "generic-array", ] +[[package]] +name = "aead" +version = "0.6.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1973cfbc1a2daf9cf550e74e1f088c28e7f7d8c1e1418fb6c9dc5184b7e84c99" +dependencies = [ + "crypto-common 0.2.2", + "inout 0.2.2", +] + [[package]] name = "aes" version = "0.8.4" @@ -34,24 +44,49 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b169f7a6d4742236a0a00c541b845991d0ac43e546831af1249753ab4c3aa3a0" dependencies = [ "cfg-if", - "cipher", + "cipher 0.4.4", "cpufeatures 0.2.17", ] +[[package]] +name = "aes" +version = "0.9.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "35f0f96ce78e38c3dc6d8948aa8163d06385be74000f3c7a95bf1eef35d3ea32" +dependencies = [ + "cipher 0.5.2", + "cpubits", + "cpufeatures 0.3.0", +] + [[package]] name = "aes-gcm" version = "0.10.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "831010a0f742e1209b3bcea8fab6a8e149051ba6099432c8cb2cc117dec3ead1" dependencies = [ - "aead", - "aes", - "cipher", - "ctr", - "ghash", + "aead 0.5.2", + "aes 0.8.4", + "cipher 0.4.4", + "ctr 0.9.2", + "ghash 0.5.1", "subtle", ] +[[package]] +name = "aes-gcm" +version = "0.11.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7f2b8006a0c83f52b62ba44a97b58bf76fe2f70a329e588f67f89691d93d498f" +dependencies = [ + "aead 0.6.1", + "aes 0.9.3", + "cipher 0.5.2", + "ctr 0.10.1", + "ctutils", + "ghash 0.6.0", +] + [[package]] name = "ahash" version = "0.8.12" @@ -625,7 +660,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e412e2cd0f2b2d93e02543ceae7917b3c70331573df19ee046bcbc35e45e87d7" dependencies = [ "byteorder", - "cipher", + "cipher 0.4.4", ] [[package]] @@ -969,10 +1004,21 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "773f3b9af64447d2ce9850330c473515014aa235e6a783b02db81ff39e4a3dad" dependencies = [ "crypto-common 0.1.7", - "inout", + "inout 0.1.4", "zeroize", ] +[[package]] +name = "cipher" +version = "0.5.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e8cf2a2c93cd704877c0858356ed03480ff301ee950b43f1cbe4573b088bfa6c" +dependencies = [ + "block-buffer 0.12.1", + "crypto-common 0.2.2", + "inout 0.2.2", +] + [[package]] name = "clap" version = "4.6.1" @@ -1348,7 +1394,16 @@ version = "0.9.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "0369ee1ad671834580515889b80f2ea915f23b8be8d0daa4bbaf2ac5c7590835" dependencies = [ - "cipher", + "cipher 0.4.4", +] + +[[package]] +name = "ctr" +version = "0.10.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "baaca1c4b237092596f64d571e9db6ce4109c4ef9742e27590f1709594461f21" +dependencies = [ + "cipher 0.5.2", ] [[package]] @@ -1377,6 +1432,21 @@ dependencies = [ "zeroize", ] +[[package]] +name = "curve25519-dalek" +version = "5.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b5eed333089e2e1c1ac8c6c0398e5e2497b4c9926ca6d0365ed1e099afa5bc23" +dependencies = [ + "cfg-if", + "cpufeatures 0.3.0", + "curve25519-dalek-derive", + "fiat-crypto 0.3.0", + "rustc_version", + "subtle", + "zeroize", +] + [[package]] name = "curve25519-dalek-derive" version = "0.1.1" @@ -2145,6 +2215,7 @@ dependencies = [ "k256", "load_config", "moka", + "p256 0.13.2", "parity-scale-codec", "ra-rpc", "ra-tls", @@ -2163,7 +2234,7 @@ dependencies = [ "tokio", "tracing", "tracing-subscriber", - "x25519-dalek", + "x25519-dalek 2.0.1", "x509-parser 0.16.0", "yasna 0.5.2", ] @@ -2321,7 +2392,7 @@ dependencies = [ name = "dstack-util" version = "0.6.3" dependencies = [ - "aes-gcm", + "aes-gcm 0.10.3", "anyhow", "base64 0.22.1", "binrw", @@ -2330,7 +2401,7 @@ dependencies = [ "cert-client", "clap", "cmd_lib", - "curve25519-dalek", + "curve25519-dalek 4.1.3", "dcap-qvl", "dstack-api-auth", "dstack-attest", @@ -2345,6 +2416,7 @@ dependencies = [ "hex", "hex_fmt", "host-api", + "hpke", "k256", "libc", "lspci", @@ -2377,7 +2449,7 @@ dependencies = [ "tracing", "tracing-subscriber", "url", - "x25519-dalek", + "x25519-dalek 2.0.1", "x509-parser 0.16.0", "yaml-rust2", ] @@ -2589,7 +2661,7 @@ version = "2.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "70e796c081cee67dc755e1a36a0a172b897fab85fc3f6bc48307991f64e4eca9" dependencies = [ - "curve25519-dalek", + "curve25519-dalek 4.1.3", "ed25519", "rand_core 0.6.4", "serde", @@ -2616,7 +2688,7 @@ dependencies = [ "ff 0.13.1", "generic-array", "group 0.13.0", - "hkdf", + "hkdf 0.12.4", "pem-rfc7468", "pkcs8", "rand_core 0.6.4", @@ -3177,7 +3249,16 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f0d8a4362ccb29cb0b265253fb0a2728f592895ee6854fd9bc13f2ffda266ff1" dependencies = [ "opaque-debug", - "polyval", + "polyval 0.6.2", +] + +[[package]] +name = "ghash" +version = "0.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2eecf2d5dc9b66b732b97707a0210906b1d30523eb773193ab777c0c84b3e8d5" +dependencies = [ + "polyval 0.7.3", ] [[package]] @@ -3485,6 +3566,15 @@ dependencies = [ "hmac 0.12.1", ] +[[package]] +name = "hkdf" +version = "0.13.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4aaa26c720c68b866f2c96ef5c1264b3e6f473fe5d4ce61cd44bbe913e553018" +dependencies = [ + "hmac 0.13.0", +] + [[package]] name = "hmac" version = "0.12.1" @@ -3525,6 +3615,25 @@ dependencies = [ "serde_json", ] +[[package]] +name = "hpke" +version = "0.15.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8f512e325a1a71bf3273fb9b7fc9e5d96dac136fb4a95d0e6637001b12ea3e09" +dependencies = [ + "aead 0.6.1", + "aes-gcm 0.11.1", + "curve25519-dalek 5.0.0", + "getrandom 0.4.2", + "hkdf 0.13.0", + "hybrid-array", + "rand_core 0.10.1", + "sha2 0.11.1", + "subtle", + "x25519-dalek 3.0.0", + "zeroize", +] + [[package]] name = "http" version = "1.4.1" @@ -3916,6 +4025,15 @@ dependencies = [ "generic-array", ] +[[package]] +name = "inout" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4250ce6452e92010fdf7268ccc5d14faa80bb12fc741938534c58f16804e03c7" +dependencies = [ + "hybrid-array", +] + [[package]] name = "insta" version = "1.47.2" @@ -4376,7 +4494,7 @@ dependencies = [ "tokio", "tracing", "tracing-subscriber", - "x25519-dalek", + "x25519-dalek 2.0.1", "xsalsa20poly1305", ] @@ -5474,7 +5592,7 @@ checksum = "8159bd90725d2df49889a078b54f4f79e87f1f8a8444194cdca81d38f5393abf" dependencies = [ "cpufeatures 0.2.17", "opaque-debug", - "universal-hash", + "universal-hash 0.5.1", ] [[package]] @@ -5486,7 +5604,18 @@ dependencies = [ "cfg-if", "cpufeatures 0.2.17", "opaque-debug", - "universal-hash", + "universal-hash 0.5.1", +] + +[[package]] +name = "polyval" +version = "0.7.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f0fa31d631f2b2cb2a544d0aa321ce847a94764d701ca2becc411138b93d49cd" +dependencies = [ + "cpubits", + "cpufeatures 0.3.0", + "universal-hash 0.6.1", ] [[package]] @@ -5946,7 +6075,7 @@ dependencies = [ "fs-err", "hex", "hex_fmt", - "hkdf", + "hkdf 0.12.4", "k256", "or-panic", "p256 0.13.2", @@ -6811,7 +6940,7 @@ version = "0.10.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "97a22f5af31f73a954c10289c93e8a50cc23d971e80ee446f1f6f7137a088213" dependencies = [ - "cipher", + "cipher 0.4.4", ] [[package]] @@ -6887,10 +7016,10 @@ version = "0.11.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6e9fcb6c2e176e86ec703e22560d99d65a5ee9056ae45a08e13e84ebf796296f" dependencies = [ - "aead", + "aead 0.5.2", "arrayref", "arrayvec", - "curve25519-dalek", + "curve25519-dalek 4.1.3", "getrandom_or_panic", "merlin", "rand_core 0.6.4", @@ -7477,7 +7606,7 @@ dependencies = [ "or-panic", "rand_core 0.6.4", "salsa20", - "x25519-dalek", + "x25519-dalek 2.0.1", "xsalsa20poly1305", ] @@ -8353,6 +8482,16 @@ dependencies = [ "subtle", ] +[[package]] +name = "universal-hash" +version = "0.6.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f4987bdc12753382e0bec4a65c50738ffaabc998b9cdd1f952fb5f39b0048a96" +dependencies = [ + "crypto-common 0.2.2", + "ctutils", +] + [[package]] name = "unsafe-libyaml" version = "0.2.11" @@ -9336,12 +9475,23 @@ version = "2.0.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c7e468321c81fb07fa7f4c636c3972b9100f0346e5b6a9f2bd0603a52f7ed277" dependencies = [ - "curve25519-dalek", + "curve25519-dalek 4.1.3", "rand_core 0.6.4", "serde", "zeroize", ] +[[package]] +name = "x25519-dalek" +version = "3.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e7e8131a03190127fb2263afc72b322ecadae46b6ff8c6f399ff5d02f5559af6" +dependencies = [ + "curve25519-dalek 5.0.0", + "rand_core 0.10.1", + "zeroize", +] + [[package]] name = "x509-cert" version = "0.2.5" @@ -9417,7 +9567,7 @@ version = "0.9.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "02a6dad357567f81cd78ee75f7c61f1b30bb2fe4390be8fb7c69e2ac8dffb6c7" dependencies = [ - "aead", + "aead 0.5.2", "poly1305", "salsa20", "subtle", diff --git a/dstack/Cargo.toml b/dstack/Cargo.toml index 2e8753148..7a98c2471 100644 --- a/dstack/Cargo.toml +++ b/dstack/Cargo.toml @@ -170,6 +170,7 @@ base64 = "0.22.1" bcrypt = "0.17.1" hex = { version = "0.4.3", default-features = false } hex_fmt = "0.3.0" +hpke = { version = "0.15.1", default-features = false, features = ["alloc", "getrandom", "x25519", "aes"] } hex-literal = "1.0.0" prost = "0.13.5" prost-types = "0.13.5" diff --git a/dstack/dstack-types/src/lib.rs b/dstack/dstack-types/src/lib.rs index d4a51710f..c7419103a 100644 --- a/dstack/dstack-types/src/lib.rs +++ b/dstack/dstack-types/src/lib.rs @@ -2400,6 +2400,8 @@ pub struct AppKeys { pub disk_crypt_key: Vec, #[serde(with = "hex_bytes", default)] pub env_crypt_key: Vec, + #[serde(with = "hex_bytes", default)] + pub sealed_env_key: Vec, #[serde(with = "hex_bytes")] pub k256_key: Vec, #[serde(with = "hex_bytes")] diff --git a/dstack/dstack-types/src/shared_filenames.rs b/dstack/dstack-types/src/shared_filenames.rs index 79944cff6..77c9d0cf6 100644 --- a/dstack/dstack-types/src/shared_filenames.rs +++ b/dstack/dstack-types/src/shared_filenames.rs @@ -8,6 +8,7 @@ pub const SYS_CONFIG: &str = ".sys-config.json"; pub const TEE_SIMULATOR_CONFIG: &str = ".tee-simulator.json"; pub const USER_CONFIG: &str = ".user-config"; pub const ENCRYPTED_ENV: &str = ".encrypted-env"; +pub const SEALED_ENV: &str = ".sealed-env"; pub const DECRYPTED_ENV: &str = ".decrypted-env"; pub const DECRYPTED_ENV_JSON: &str = ".decrypted-env.json"; pub const INSTANCE_INFO: &str = ".instance_info"; diff --git a/dstack/dstack-util/Cargo.toml b/dstack/dstack-util/Cargo.toml index b7fff59a5..f2401bc40 100644 --- a/dstack/dstack-util/Cargo.toml +++ b/dstack/dstack-util/Cargo.toml @@ -19,6 +19,7 @@ fs-err.workspace = true getrandom = { workspace = true, features = ["std"] } hex.workspace = true hex_fmt.workspace = true +hpke.workspace = true regex.workspace = true scale = { workspace = true, features = ["derive"] } schnorrkel.workspace = true diff --git a/dstack/dstack-util/src/crypto.rs b/dstack/dstack-util/src/crypto.rs index 1f8227f5b..4c080609c 100644 --- a/dstack/dstack-util/src/crypto.rs +++ b/dstack/dstack-util/src/crypto.rs @@ -68,6 +68,35 @@ pub fn dh_decrypt(secret: [u8; 32], ciphertext: &[u8]) -> Result> { .map_err(|e| anyhow!("Decryption failed: {}", e)) } +/// Encrypt to `remote_public_key` in the format [`dh_decrypt`] reads: +/// ephemeral public key, IV, then the AES-GCM ciphertext. +#[cfg(test)] +pub fn dh_encrypt(remote_public_key: [u8; 32], plaintext: &[u8]) -> Result> { + let mut ephemeral_secret = [0u8; 32]; + getrandom::fill(&mut ephemeral_secret).context("failed to generate ephemeral secret")?; + let ephemeral_secret = StaticSecret::from(ephemeral_secret); + let shared_secret = ephemeral_secret + .diffie_hellman(&PublicKey::from(remote_public_key)) + .to_bytes(); + ensure!( + !shared_secret.iter().all(|byte| *byte == 0), + "invalid X25519 shared secret" + ); + let cipher = Aes256Gcm::new_from_slice(&shared_secret) + .map_err(|e| anyhow!("failed to create cipher: {e}"))?; + let mut iv = [0u8; 12]; + getrandom::fill(&mut iv).context("failed to generate IV")?; + let ciphertext = cipher + .encrypt(Nonce::::from_slice(&iv), plaintext) + .map_err(|e| anyhow!("encryption failed: {e}"))?; + Ok([ + PublicKey::from(&ephemeral_secret).as_bytes(), + &iv[..], + &ciphertext, + ] + .concat()) +} + fn stream_nonce(prefix: &[u8; 8], index: u32) -> [u8; 12] { let mut nonce = [0u8; 12]; nonce[..8].copy_from_slice(prefix); @@ -266,6 +295,15 @@ pub fn dh_decrypt_stream( mod tests { use super::*; + #[test] + fn dh_encrypt_round_trips_through_dh_decrypt() { + let secret = [7u8; 32]; + let public = PublicKey::from(&StaticSecret::from(secret)).to_bytes(); + let ciphertext = dh_encrypt(public, b"env").unwrap(); + assert_eq!(dh_decrypt(secret, &ciphertext).unwrap(), b"env"); + assert!(dh_decrypt([8u8; 32], &ciphertext).is_err()); + } + #[test] fn test_dh_agree() { use rand::Rng; @@ -361,3 +399,95 @@ mod tests { .is_err()); } } + +/// Encrypt a sealed environment with HPKE DHKEM(X25519, HKDF-SHA256) and AES-256-GCM. +pub fn sealed_env_encrypt(public_key: &[u8], info: &[u8], plaintext: &[u8]) -> Result> { + use hpke::{ + aead::AesGcm256, kdf::HkdfSha256, kem::X25519HkdfSha256, Deserializable, OpModeS, + Serializable, + }; + let public_key = ::PublicKey::from_bytes(public_key) + .map_err(|_| anyhow!("invalid sealed env public key"))?; + let (enc, ciphertext) = hpke::single_shot_seal::( + &OpModeS::Base, + &public_key, + info, + plaintext, + &[], + ) + .map_err(|e| anyhow!("sealed env encryption failed: {e}"))?; + Ok([enc.to_bytes().as_ref(), ciphertext.as_slice()].concat()) +} + +/// Decrypt a sealed environment HPKE payload. +pub fn sealed_env_decrypt(secret: &[u8], info: &[u8], payload: &[u8]) -> Result> { + use hpke::{aead::AesGcm256, kdf::HkdfSha256, kem::X25519HkdfSha256, Deserializable, OpModeR}; + let enc = payload + .get(..32) + .ok_or_else(|| anyhow!("sealed env payload is missing encapsulated key"))?; + let ciphertext = payload + .get(32..) + .ok_or_else(|| anyhow!("sealed env payload is missing ciphertext"))?; + let secret = ::PrivateKey::from_bytes(secret) + .map_err(|_| anyhow!("invalid sealed env private key"))?; + let enc = ::EncappedKey::from_bytes(enc) + .map_err(|_| anyhow!("invalid sealed env encapsulated key"))?; + hpke::single_shot_open::( + &OpModeR::Base, + &secret, + &enc, + info, + ciphertext, + &[], + ) + .map_err(|e| anyhow!("sealed env decryption failed: {e}")) +} + +#[cfg(test)] +mod sealed_env_tests { + use super::*; + + #[test] + fn opens_sdk_envelopes() { + use base64::Engine as _; + for source in [ + include_str!("../../../sdk/fixtures/sealed-env/python.json"), + include_str!("../../../sdk/fixtures/sealed-env/javascript.json"), + ] { + let fixture: serde_json::Value = serde_json::from_str(source).unwrap(); + let envelope = &fixture["sealed_env"]; + let decode = |name: &str| { + base64::engine::general_purpose::STANDARD + .decode(envelope[name].as_str().unwrap()) + .unwrap() + }; + let payload = [decode("enc"), decode("ciphertext")].concat(); + let plaintext = sealed_env_decrypt( + &hex::decode(fixture["private_key"].as_str().unwrap()).unwrap(), + b"dstack-sealed-env-v1:3", + &payload, + ) + .unwrap(); + let decoded: serde_json::Value = serde_json::from_slice(&plaintext).unwrap(); + assert_eq!(decoded["env"], fixture["env"]); + } + } + + #[test] + fn sealed_env_round_trips_and_binds_info() { + let secret = StaticSecret::random_from_rng(rand::thread_rng()); + let public = PublicKey::from(&secret).to_bytes(); + let info = b"dstack-sealed-env-v1:3"; + let payload = sealed_env_encrypt(&public, info, b"secret env").unwrap(); + assert_eq!( + sealed_env_decrypt(&secret.to_bytes(), info, &payload).unwrap(), + b"secret env" + ); + assert!( + sealed_env_decrypt(&secret.to_bytes(), b"dstack-sealed-env-v1:4", &payload).is_err() + ); + let mut tampered = payload; + *tampered.last_mut().unwrap() ^= 1; + assert!(sealed_env_decrypt(&secret.to_bytes(), info, &tampered).is_err()); + } +} diff --git a/dstack/dstack-util/src/host_api.rs b/dstack/dstack-util/src/host_api.rs index e038ea3b2..743becd42 100644 --- a/dstack/dstack-util/src/host_api.rs +++ b/dstack/dstack-util/src/host_api.rs @@ -80,6 +80,24 @@ impl HostApi { Ok(()) } + /// Hand the host the env re-encrypted to the current root epoch, or an + /// empty one when nothing changed. + pub async fn update_sealed_env(&self, sealed_env: Vec) -> Result<()> { + if !matches!(Platform::detect_or_dstack(), Platform::Dstack) { + return Ok(()); + } + let Some(client) = &self.client else { + return Ok(()); + }; + tokio::time::timeout( + HOST_API_TIMEOUT, + client.update_sealed_env(host_api::SealedEnv { sealed_env }), + ) + .await + .context("Timed out updating the sealed env on the host")??; + Ok(()) + } + pub async fn notify_q(&self, event: &str, payload: &str) { if let Err(err) = self.notify(event, payload).await { warn!("Failed to notify event {event} to host: {:?}", err); diff --git a/dstack/dstack-util/src/main.rs b/dstack/dstack-util/src/main.rs index ab1f6837b..5b0eebf9c 100644 --- a/dstack/dstack-util/src/main.rs +++ b/dstack/dstack-util/src/main.rs @@ -677,6 +677,7 @@ async fn cmd_get_keys(args: GetKeysArgs) -> Result<()> { ca_cert: root_ca, disk_crypt_key: response.disk_crypt_key, env_crypt_key: response.env_crypt_key, + sealed_env_key: response.sealed_env_key, k256_key: response.k256_key, k256_signature: response.k256_signature, gateway_app_id: response.gateway_app_id, @@ -793,6 +794,7 @@ async fn cmd_encrypt(args: EncryptArgs) -> Result<()> { let response = KmsClient::new(client) .get_app_env_encrypt_pub_key(dstack_kms_rpc::AppId { app_id: app_id.to_vec(), + root_epoch: None, }) .await .context("failed to get app environment encryption public key")?; @@ -1191,6 +1193,7 @@ fn make_app_keys( Ok(AppKeys { disk_crypt_key: sha256(&disk_key.serialize_der()).to_vec(), env_crypt_key: vec![], + sealed_env_key: vec![], k256_key: k256_key.to_bytes().to_vec(), k256_signature: vec![], gateway_app_id: "".to_string(), diff --git a/dstack/dstack-util/src/system_setup.rs b/dstack/dstack-util/src/system_setup.rs index 0699befb9..54a607363 100644 --- a/dstack/dstack-util/src/system_setup.rs +++ b/dstack/dstack-util/src/system_setup.rs @@ -16,13 +16,14 @@ use std::{ }; use anyhow::{anyhow, bail, Context, Result}; +use base64::Engine as _; use dstack_attest::{default_verifier, emit_runtime_event, set_runtime_event_version}; use dstack_kms_rpc as rpc; use dstack_types::{ gpu_policy_hash, shared_filenames::{ APP_COMPOSE, APP_KEYS, DECRYPTED_ENV, DECRYPTED_ENV_JSON, ENCRYPTED_ENV, - HOST_SHARED_DIR_NAME, INSTANCE_INFO, SYS_CONFIG, USER_CONFIG, + HOST_SHARED_DIR_NAME, INSTANCE_INFO, SEALED_ENV, SYS_CONFIG, USER_CONFIG, }, GpuPolicy, KeyProvider, KeyProviderInfo, GPU_ATTESTATION_OUTPUT, }; @@ -229,6 +230,10 @@ impl HostShareDir { self.base_dir.join(ENCRYPTED_ENV) } + fn sealed_env_file(&self) -> PathBuf { + self.base_dir.join(SEALED_ENV) + } + fn sys_config_file(&self) -> PathBuf { self.base_dir.join(SYS_CONFIG) } @@ -242,11 +247,34 @@ impl HostShareDir { } } +fn env_file_exists(path: &Path) -> Result { + match fs::symlink_metadata(path) { + Ok(metadata) if metadata.is_file() => Ok(true), + Ok(_) => bail!("env {} is not a regular file", path.display()), + Err(error) if error.kind() == std::io::ErrorKind::NotFound => Ok(false), + Err(error) => { + Err(error).with_context(|| format!("failed to inspect env {}", path.display())) + } + } +} + +fn read_optional_env(path: &Path) -> Result>> { + if !env_file_exists(path)? { + return Ok(None); + } + match fs::read(path) { + Ok(bytes) => Ok(Some(bytes)), + Err(error) if error.kind() == std::io::ErrorKind::NotFound => Ok(None), + Err(error) => Err(error).with_context(|| format!("failed to read env {}", path.display())), + } +} + struct HostShared { dir: HostShareDir, sys_config: SysConfig, app_compose: AppCompose, encrypted_env: Vec, + sealed_env: Option>, instance_info: InstanceInfo, } @@ -261,12 +289,18 @@ impl HostShared { } else { InstanceInfo::default() }; - let encrypted_env = fs::read(host_shared_dir.encrypted_env_file()).unwrap_or_default(); + let sealed_env = read_optional_env(&host_shared_dir.sealed_env_file())?; + let encrypted_env = if sealed_env.is_some() { + Vec::new() + } else { + read_optional_env(&host_shared_dir.encrypted_env_file())?.unwrap_or_default() + }; Ok(Self { dir: host_shared_dir.clone(), sys_config, app_compose, encrypted_env, + sealed_env, instance_info, }) } @@ -315,7 +349,16 @@ impl HostShared { copy(APP_COMPOSE, SZ_1MB * 50, false)?; copy(SYS_CONFIG, SZ_1KB * 32, false)?; copy(INSTANCE_INFO, SZ_1KB * 10, true)?; - copy(ENCRYPTED_ENV, SZ_1KB * 256, true)?; + if env_file_exists(&host_shared_dir.join(SEALED_ENV))? { + copy(SEALED_ENV, SZ_1KB * 384, false)?; + } else { + match fs::remove_file(host_shared_copy_dir.join(SEALED_ENV)) { + Ok(()) => {} + Err(error) if error.kind() == std::io::ErrorKind::NotFound => {} + Err(error) => return Err(error).context("failed to remove stale sealed env copy"), + } + copy(ENCRYPTED_ENV, SZ_1KB * 256, true)?; + } copy(USER_CONFIG, SZ_1MB * 50, true)?; info!("Unmounting host-shared"); unmount_host_shared(host_shared_dir)?; @@ -2389,6 +2432,7 @@ impl HostShared { ca_cert: root_ca, disk_crypt_key: response.disk_crypt_key, env_crypt_key: response.env_crypt_key, + sealed_env_key: response.sealed_env_key, k256_key: response.k256_key, k256_signature: response.k256_signature, gateway_app_id: response.gateway_app_id, @@ -3277,25 +3321,74 @@ impl<'a> Stage0<'a> { } impl Stage1<'_> { - fn decrypt_env_vars( - &self, - key: &[u8], - ciphertext: &[u8], - allowed: &BTreeSet, - ) -> Result> { - let vars = if !ciphertext.is_empty() { - info!("Processing encrypted env"); - let env_crypt_key: [u8; 32] = key.try_into().ok().context( - "encrypted env is present but the key provider gave no valid env crypt key", + /// Decrypt the host-supplied env, preferring the sealed format. + async fn decrypt_env(&self) -> Result<(Vec, Vec)> { + let current_secret = self.keys.sealed_env_key.as_slice(); + if let Some(sealed_env) = &self.shared.sealed_env { + let envelope: SealedEnv = + serde_json::from_slice(sealed_env).context("failed to parse sealed env")?; + if envelope.version != 1 { + bail!("unsupported sealed env version: {}", envelope.version); + } + let enc = base64::engine::general_purpose::STANDARD.decode(envelope.enc)?; + if enc.len() != 32 { + bail!("sealed env encapsulated key must be 32 bytes"); + } + let ciphertext = + base64::engine::general_purpose::STANDARD.decode(envelope.ciphertext)?; + let payload = [enc, ciphertext].concat(); + let current = self.keys.root_epoch(); + if envelope.root_epoch > current { + bail!( + "sealed env epoch {} exceeds current epoch {current}", + envelope.root_epoch + ); + } + let historical; + let secret = if envelope.root_epoch == current { + current_secret + } else { + historical = self + .shared + .request_historical_keys(&self.keys, envelope.root_epoch) + .await?; + historical.sealed_env_key.as_slice() + }; + let plaintext = crate::crypto::sealed_env_decrypt( + secret, + &sealed_env_info(&envelope.root_epoch), + &payload, )?; - let decrypted_json = - dh_decrypt(env_crypt_key, ciphertext).context("Failed to decrypt env file")?; - crate::parse_env_file::parse_env(&decrypted_json, allowed)? - } else { - info!("No encrypted env, using default"); - Default::default() + if envelope.root_epoch == self.keys.root_epoch() { + return Ok((plaintext, Vec::new())); + } + let sealed = seal_env(current_secret, self.keys.root_epoch(), &plaintext)?; + return Ok((plaintext, sealed)); + } + let ciphertext = &self.shared.encrypted_env; + let current_key = env_crypt_key(&self.keys.env_crypt_key)?; + let plaintext = match dh_decrypt(current_key, ciphertext) { + Ok(plaintext) => plaintext, + Err(error) => { + let current = self.keys.root_epoch(); + let mut recovered = None; + for epoch in (0..current).rev() { + let old = self + .shared + .request_historical_keys(&self.keys, epoch) + .await?; + if let Ok(plaintext) = + dh_decrypt(env_crypt_key(&old.env_crypt_key)?, ciphertext) + { + recovered = Some(plaintext); + break; + } + } + recovered.ok_or(error)? + } }; - Ok(vars) + let sealed = seal_env(current_secret, self.keys.root_epoch(), &plaintext)?; + Ok((plaintext, sealed)) } fn write_env_file(&self, env_vars: &BTreeMap) -> Result<()> { @@ -3312,7 +3405,7 @@ impl Stage1<'_> { Ok(()) } - fn unseal_env_vars(&self) -> Result> { + async fn unseal_env_vars(&self) -> Result> { let allowed_envs: BTreeSet = self .shared .app_compose @@ -3320,18 +3413,29 @@ impl Stage1<'_> { .iter() .cloned() .collect(); - // Decrypt env file - let decrypted_env = self.decrypt_env_vars( - &self.keys.env_crypt_key, - &self.shared.encrypted_env, - &allowed_envs, - )?; + let (decrypted_env, sealed) = + if self.shared.sealed_env.is_none() && self.shared.encrypted_env.is_empty() { + info!("No encrypted env, using default"); + (Default::default(), Vec::new()) + } else { + info!("Processing encrypted env"); + let (plaintext, sealed) = self.decrypt_env().await?; + ( + crate::parse_env_file::parse_env(&plaintext, &allowed_envs)?, + sealed, + ) + }; self.write_env_file(&decrypted_env)?; + // Sent on every boot, empty when nothing changed: the host accepts one + // update per boot, and this one comes before any application code. + if let Err(err) = self.vmm.update_sealed_env(sealed).await { + warn!("failed to store the re-encrypted env on the host: {err:#}"); + } Ok(decrypted_env) } async fn setup(&self) -> Result<()> { - let _envs = self.unseal_env_vars()?; + let _envs = self.unseal_env_vars().await?; self.link_files()?; self.setup_socket_dir()?; self.setup_guest_agent_config()?; @@ -3428,6 +3532,37 @@ impl Stage1<'_> { } } +#[derive(serde::Deserialize)] +struct SealedEnv { + version: u32, + root_epoch: u32, + enc: String, + ciphertext: String, +} + +fn sealed_env_info(epoch: &u32) -> Vec { + format!("dstack-sealed-env-v1:{epoch}").into_bytes() +} + +fn seal_env(secret: &[u8], epoch: u32, plaintext: &[u8]) -> Result> { + let secret: [u8; 32] = secret.try_into().context("invalid sealed env key")?; + let public = x25519_dalek::PublicKey::from(&x25519_dalek::StaticSecret::from(secret)); + let payload = + crate::crypto::sealed_env_encrypt(&public.to_bytes(), &sealed_env_info(&epoch), plaintext)?; + let (enc, ciphertext) = payload.split_at(32); + Ok(serde_json::to_vec(&serde_json::json!({ + "version": 1, + "root_epoch": epoch, + "enc": base64::engine::general_purpose::STANDARD.encode(enc), + "ciphertext": base64::engine::general_purpose::STANDARD.encode(ciphertext), + }))?) +} + +fn env_crypt_key(key: &[u8]) -> Result<[u8; 32]> { + key.try_into() + .ok() + .context("encrypted env is present but the key provider gave no valid env crypt key") +} fn kms_root_epoch(keys: &AppKeys) -> Option { match keys.key_provider { KeyProvider::Kms { root_epoch, .. } => Some(root_epoch), @@ -4227,3 +4362,118 @@ Endpoint = [2001:db8::1]:51822 assert!(parse_wg_ip("10.0.0.2/32\n[Interface]").is_err()); } } + +#[cfg(test)] +mod sealed_env_file_tests { + use super::*; + + fn stage(args: &SetupArgs, sealed_env: Option>, encrypted_env: Vec) -> Stage1<'_> { + Stage1 { + args, + vmm: HostApi::default(), + shared: HostShared { + dir: HostShareDir::new(&args.work_dir), + sys_config: serde_json::from_value(serde_json::json!({"vm_config": ""})).unwrap(), + app_compose: serde_json::from_value(serde_json::json!({ + "manifest_version": "2", "name": "sealed-test", "runner": "docker-compose" + })) + .unwrap(), + instance_info: InstanceInfo::default(), + sealed_env, + encrypted_env, + }, + keys: AppKeys { + disk_crypt_key: vec![], + env_crypt_key: vec![9; 32], + sealed_env_key: vec![7; 32], + k256_key: vec![], + k256_signature: vec![], + gateway_app_id: String::new(), + ca_cert: String::new(), + key_provider: KeyProvider::Kms { + url: "http://127.0.0.1:1".into(), + pubkey: vec![], + root_epoch: 3, + }, + }, + } + } + + #[tokio::test] + async fn guest_prefers_sealed_and_converts_legacy() { + let dir = tempfile::tempdir().unwrap(); + let args = SetupArgs { + work_dir: dir.path().into(), + device: "unused".into(), + mount_point: "unused".into(), + }; + let plaintext = br#"{"env":[{"key":"MESSAGE","value":"test"}]}"#; + let sealed = seal_env(&[7; 32], 3, plaintext).unwrap(); + let current = stage(&args, Some(sealed), b"invalid legacy".to_vec()); + let (opened, update) = current.decrypt_env().await.unwrap(); + assert_eq!(opened, plaintext); + assert!(update.is_empty()); + + let public = x25519_dalek::PublicKey::from(&x25519_dalek::StaticSecret::from([9; 32])); + let legacy = crate::crypto::dh_encrypt(public.to_bytes(), plaintext).unwrap(); + let old = stage(&args, None, legacy.clone()); + let (opened, update) = old.decrypt_env().await.unwrap(); + assert_eq!(opened, plaintext); + assert!(!update.is_empty()); + let migrated = stage(&args, Some(update), vec![]); + assert_eq!(migrated.decrypt_env().await.unwrap().0, plaintext); + + for invalid in [ + Vec::new(), + b"invalid JSON".to_vec(), + br#"{"version":2,"root_epoch":3,"enc":"","ciphertext":""}"#.to_vec(), + ] { + assert!(stage(&args, Some(invalid), legacy.clone()) + .decrypt_env() + .await + .is_err()); + } + let future = seal_env(&[7; 32], 4, plaintext).unwrap(); + assert!(stage(&args, Some(future), legacy) + .decrypt_env() + .await + .unwrap_err() + .to_string() + .contains("exceeds current")); + } + + #[test] + fn sealed_file_selection_ignores_legacy_but_rejects_links() { + let dir = tempfile::tempdir().unwrap(); + fs::write(dir.path().join(SYS_CONFIG), br#"{"vm_config":""}"#).unwrap(); + fs::write( + dir.path().join(APP_COMPOSE), + br#"{"manifest_version":"2","name":"test","runner":"docker-compose"}"#, + ) + .unwrap(); + let sealed_path = dir.path().join(SEALED_ENV); + fs::write(&sealed_path, b"sealed bytes").unwrap(); + fs::create_dir(dir.path().join(ENCRYPTED_ENV)).unwrap(); + let shared = HostShared::load(dir.path()).unwrap(); + assert_eq!(shared.sealed_env.unwrap(), b"sealed bytes"); + assert!(shared.encrypted_env.is_empty()); + fs::remove_file(&sealed_path).unwrap(); + std::os::unix::fs::symlink("missing-target", &sealed_path).unwrap(); + assert!(env_file_exists(&sealed_path).is_err()); + assert!(HostShared::load(dir.path()).is_err()); + } + + #[test] + fn missing_empty_and_unreadable_env_are_distinct() { + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join(SEALED_ENV); + assert_eq!(read_optional_env(&path).unwrap(), None); + fs::write(&path, []).unwrap(); + assert_eq!(read_optional_env(&path).unwrap(), Some(Vec::new())); + assert!( + serde_json::from_slice::(&read_optional_env(&path).unwrap().unwrap()) + .is_err() + ); + assert!(read_optional_env(dir.path()).is_err()); + } +} diff --git a/dstack/guest-agent/src/rpc_service.rs b/dstack/guest-agent/src/rpc_service.rs index d79a9ed22..58ff5a865 100644 --- a/dstack/guest-agent/src/rpc_service.rs +++ b/dstack/guest-agent/src/rpc_service.rs @@ -1437,6 +1437,7 @@ pNs85uhOZE8z2jr8Pg== let dummy_keys = AppKeys { disk_crypt_key: Vec::new(), env_crypt_key: Vec::new(), + sealed_env_key: Vec::new(), k256_key: DUMMY_K256_KEY.to_vec(), k256_signature: Vec::new(), gateway_app_id: String::new(), diff --git a/dstack/host-api/proto/host_api.proto b/dstack/host-api/proto/host_api.proto index f8864f516..c4880ba74 100644 --- a/dstack/host-api/proto/host_api.proto +++ b/dstack/host-api/proto/host_api.proto @@ -27,8 +27,17 @@ message GetSealingKeyResponse { bytes provider_quote = 2; } +// The env the guest re-encrypted to the current KMS root epoch at boot; empty +// when it needed no change. +message SealedEnv { + bytes sealed_env = 1; +} + service HostApi { rpc Info(google.protobuf.Empty) returns (HostInfo); rpc Notify(Notification) returns (google.protobuf.Empty); rpc GetSealingKey(GetSealingKeyRequest) returns (GetSealingKeyResponse); + // Replace this VM's stored encrypted env. Accepted once per boot, before + // application code runs, and only if the stored env is the one it booted with. + rpc UpdateSealedEnv(SealedEnv) returns (google.protobuf.Empty); } diff --git a/dstack/kms/Cargo.toml b/dstack/kms/Cargo.toml index 5c616d7a4..2979b4230 100644 --- a/dstack/kms/Cargo.toml +++ b/dstack/kms/Cargo.toml @@ -32,6 +32,7 @@ reqwest = { workspace = true, features = ["json"] } sha2.workspace = true sha3.workspace = true k256.workspace = true +p256.workspace = true rand.workspace = true dstack-guest-agent-rpc.workspace = true http-client = { workspace = true, features = ["prpc"] } diff --git a/dstack/kms/rpc/proto/kms_rpc.proto b/dstack/kms/rpc/proto/kms_rpc.proto index 13083647b..7cfd3d8ac 100644 --- a/dstack/kms/rpc/proto/kms_rpc.proto +++ b/dstack/kms/rpc/proto/kms_rpc.proto @@ -18,6 +18,14 @@ message GetAppKeyRequest { message AppId { bytes app_id = 1; + optional uint32 root_epoch = 2; +} + +message SealedEnvPubKeyResponse { + bytes public_key = 1; + uint32 root_epoch = 2; + uint64 timestamp = 3; + bytes signature = 4; } message PublicKeyResponse { @@ -51,6 +59,8 @@ message AppKeyResponse { bytes os_image_hash = 8; // Root key epoch all epoch-dependent fields above were derived from. uint32 root_epoch = 9; + // X25519 private key for the SealedEnv HPKE receiver, derived for this epoch. + bytes sealed_env_key = 10; } message GetMetaResponse { @@ -109,6 +119,7 @@ service KMS { rpc GetKmsKey(GetKmsKeyRequest) returns (KmsKeyResponse); // Request the app environment encryption public key given the app id rpc GetAppEnvEncryptPubKey(AppId) returns (PublicKeyResponse); + rpc GetSealedEnvPubKey(AppId) returns (SealedEnvPubKeyResponse); // Request the KMS instance metadata for use as a probe and health check. rpc GetMeta(google.protobuf.Empty) returns (GetMetaResponse); // Compatibility for existing OS images and tokenless legacy onboarding. diff --git a/dstack/kms/src/main_service.rs b/dstack/kms/src/main_service.rs index 605092339..20a63afb2 100644 --- a/dstack/kms/src/main_service.rs +++ b/dstack/kms/src/main_service.rs @@ -14,13 +14,14 @@ use anyhow::{anyhow, bail, Context, Result}; use dstack_kms_rpc::{ kms_server::{KmsRpc, KmsServer}, AppId, AppKeyResponse, GetAppKeyRequest, GetKmsKeyRequest, GetMetaResponse, - GetTempCaCertResponse, KmsKeyResponse, KmsKeys, PublicKeyResponse, SignCertRequest, - SignCertResponse, + GetTempCaCertResponse, KmsKeyResponse, KmsKeys, PublicKeyResponse, SealedEnvPubKeyResponse, + SignCertRequest, SignCertResponse, }; use dstack_verifier::{CvmVerifier, VerificationDetails}; use fs_err as fs; use k256::ecdsa::SigningKey; use moka::future::Cache; +use p256::pkcs8::DecodePrivateKey; use ra_rpc::{CallContext, RpcCall}; use ra_tls::{ attestation::{resolve_vm_config, AttestationVerifier, TeeVariant, VerifiedAttestation}, @@ -366,6 +367,21 @@ struct AppKeys { env_crypt_key: Vec, k256_key: Vec, k256_signature: Vec, + sealed_env_key: Vec, +} + +fn derive_sealed_env_key(root: &RootKeys, app_id: &[u8]) -> Result<[u8; 32]> { + ensure_app_id_len(app_id)?; + let ca_secret = p256::SecretKey::from_pkcs8_der(root.ca.key.serialized_der())?; + kdf::derive_key_with_salt( + b"dstack-sealed-env-v1", + &ca_secret.to_bytes(), + &[app_id], + 32, + ) + .map_err(|_| anyhow!("failed to derive sealed env key"))? + .try_into() + .map_err(|_| anyhow!("invalid sealed env key length")) } fn derive_app_keys(root: &RootKeys, app_id: &[u8], instance_id: &[u8]) -> Result { @@ -377,11 +393,13 @@ fn derive_app_keys(root: &RootKeys, app_id: &[u8], instance_id: &[u8]) -> Result let env_crypt_key = x25519_dalek::StaticSecret::from(env_crypt_key).to_bytes(); let (k256_key, k256_signature) = derive_k256_key(&root.k256, app_id).context("Failed to derive app ecdsa key")?; + let sealed_env_key = derive_sealed_env_key(root, app_id)?.to_vec(); Ok(AppKeys { disk_crypt_key: disk_crypt_key.to_vec(), env_crypt_key: env_crypt_key.to_vec(), k256_key: k256_key.to_bytes().to_vec(), k256_signature, + sealed_env_key, }) } @@ -660,6 +678,7 @@ impl KmsRpc for RpcHandler { gateway_app_id, os_image_hash: boot_info.os_image_hash, root_epoch, + sealed_env_key: keys.sealed_env_key, }) } @@ -711,6 +730,40 @@ impl KmsRpc for RpcHandler { }) } + async fn get_sealed_env_pub_key(self, request: AppId) -> Result { + let requested = request.root_epoch; + let current = self.ensure_self_allowed(Operation::Serve).await?; + let epoch = requested.unwrap_or(current); + if epoch > current { + bail!("sealed env root epoch {epoch} is not active; current epoch is {current}"); + } + let root = self.state.epoch_root(epoch)?; + ensure_app_id_len(&request.app_id)?; + let secret = derive_sealed_env_key(root, &request.app_id)?; + let public_key = + x25519_dalek::PublicKey::from(&x25519_dalek::StaticSecret::from(secret)).to_bytes(); + let timestamp = std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .context("system time before UNIX epoch")? + .as_secs(); + let mut message = Vec::with_capacity(4 + 32 + 8); + message.extend_from_slice(&epoch.to_be_bytes()); + message.extend_from_slice(&public_key); + message.extend_from_slice(×tamp.to_be_bytes()); + let signature = sign_message( + &root.k256, + b"dstack-sealed-env-pubkey-v1", + &request.app_id, + &message, + )?; + Ok(SealedEnvPubKeyResponse { + public_key: public_key.to_vec(), + root_epoch: epoch, + timestamp, + signature, + }) + } + async fn get_meta(self) -> Result { let bootstrap_info = fs::read_to_string(self.state.config.bootstrap_info()) .ok() @@ -1532,6 +1585,21 @@ mod tests { KmsState::new(config, keyring, ca_cert.pem()).unwrap() } + #[test] + fn sealed_env_derivation_is_separate_and_shared_with_app_keys() { + let dir = tempfile::tempdir().unwrap(); + let state = kms_state(dir.path(), "http://127.0.0.1:1"); + let root = state.epoch_root(0).unwrap(); + let app = [1; 20]; + let sealed = derive_sealed_env_key(root, &app).unwrap(); + let app_keys = derive_app_keys(root, &app, &[2; 20]).unwrap(); + assert_eq!(sealed.as_slice(), app_keys.sealed_env_key); + assert_ne!(sealed.as_slice(), app_keys.env_crypt_key); + assert_ne!(sealed.as_slice(), app_keys.disk_crypt_key); + assert_ne!(sealed, derive_sealed_env_key(root, &[3; 20]).unwrap()); + assert!(derive_sealed_env_key(root, &[1; 19]).is_err()); + } + #[rocket::async_test] async fn authorization_selects_epochs_without_a_serving_epoch_snapshot() { let dir = tempfile::tempdir().unwrap(); diff --git a/dstack/vmm/requirements.txt b/dstack/vmm/requirements.txt index 5a4636513..9ec6ca2f0 100644 --- a/dstack/vmm/requirements.txt +++ b/dstack/vmm/requirements.txt @@ -1,5 +1,5 @@ annotated-types==0.7.0 -cffi==1.17.1 +cffi==2.0.0 cryptography==50.0.0 cytoolz==1.0.1 eth-hash==0.7.1 @@ -7,6 +7,7 @@ eth-keys==0.7.0 eth-typing==5.2.1 eth-utils==5.3.0 pycparser==2.22 +pyhpke==0.6.5 pydantic==2.11.3 pydantic_core==2.33.1 safe-pysha3==1.0.4 diff --git a/dstack/vmm/rpc/proto/vmm_rpc.proto b/dstack/vmm/rpc/proto/vmm_rpc.proto index 84b66b2d6..6fb3a61fd 100644 --- a/dstack/vmm/rpc/proto/vmm_rpc.proto +++ b/dstack/vmm/rpc/proto/vmm_rpc.proto @@ -109,6 +109,8 @@ message VmConfiguration { repeated PortMapping ports = 7; // Encrypted environment variables bytes encrypted_env = 8; + // Sealed environment variables. Takes precedence over encrypted_env. + bytes sealed_env = 24; // App ID. // If not provided, it assumes the app is newly created. // If provided, and KMS is enabled, it assumes the app is upgraded from given app_id. @@ -214,6 +216,8 @@ message UpdateVmRequest { string compose_file = 2; // Optional update encrypted environment variables. Leave empty to not update. bytes encrypted_env = 3; + // Optional sealed environment update. Leave empty to not update. + bytes sealed_env = 24; // Optional update user config. Leave empty to not update. string user_config = 4; // Optional update port mapping. @@ -287,9 +291,17 @@ message ImageInfo { // App identifier used for encrypting env vars and routing telemetry. message AppId { bytes app_id = 1; + optional uint32 root_epoch = 2; } // Response with the KMS-backed public key used for env encryption plus audit signature. +message SealedEnvPubKeyResponse { + bytes public_key = 1; + uint32 root_epoch = 2; + uint64 timestamp = 3; + bytes signature = 4; +} + message PublicKeyResponse { bytes public_key = 1; // Legacy signature without timestamp (for backward compatibility). @@ -423,6 +435,7 @@ service Vmm { // Get Env encrypt public key rpc GetAppEnvEncryptPubKey(AppId) returns (PublicKeyResponse); + rpc GetSealedEnvPubKey(AppId) returns (SealedEnvPubKeyResponse); // Get VM info by ID rpc GetInfo(Id) returns (GetInfoResponse); diff --git a/dstack/vmm/src/app.rs b/dstack/vmm/src/app.rs index 379bea90a..f060c1ece 100644 --- a/dstack/vmm/src/app.rs +++ b/dstack/vmm/src/app.rs @@ -16,7 +16,8 @@ use bon::Builder; use dstack_kms_rpc::kms_client::KmsClient; use dstack_types::mr_config::MrConfigV3; use dstack_types::shared_filenames::{ - APP_COMPOSE, ENCRYPTED_ENV, INSTANCE_INFO, SYS_CONFIG, TEE_SIMULATOR_CONFIG, USER_CONFIG, + APP_COMPOSE, ENCRYPTED_ENV, INSTANCE_INFO, SEALED_ENV, SYS_CONFIG, TEE_SIMULATOR_CONFIG, + USER_CONFIG, }; use dstack_types::version::Version; use dstack_vmm_rpc::{ @@ -530,9 +531,16 @@ impl App { } self.sync_dynamic_config(id)?; self.set_started(id, true)?; + // A guest the VMM did not launch may already run application code. + let env_update = if is_running { + None + } else { + Some(self.sha256_of_env(id)?) + }; let vm_config = { let mut state = self.lock(); let vm_state = state.get_mut(id).context("VM not found")?; + vm_state.state.env_update = env_update; // Older images does not support for progress reporting if vm_state.config.image.info.shared_ro { vm_state.state.start(is_running); @@ -1561,7 +1569,22 @@ impl App { vm.state.events.pop_front(); } match event { - "boot.progress" => vm.state.boot_progress = body, + "boot.progress" => { + // Application code (init and pre-launch scripts, containers) + // runs from here on, so the guest OS can no longer vouch for + // an env update. + if matches!( + body.as_str(), + "init-script" + | "pre-launch" + | "starting containers" + | "running main script" + | "done" + ) { + vm.state.env_update = None; + } + vm.state.boot_progress = body; + } "boot.error" => vm.state.boot_error = body, "shutdown.progress" => { let powering_off = body == "powering off"; @@ -1582,12 +1605,63 @@ impl App { Ok(()) } - pub(crate) fn compose_file_path(&self, id: &str) -> Result { - Ok(self.shared_dir(id)?.join(APP_COMPOSE)) + /// Store the encrypted env a guest re-encrypted at boot. Accepted once per + /// boot, before application code runs, and only over the env the VM booted + /// with, so an update made since then is never overwritten. + pub(crate) async fn update_sealed_env(&self, cid: u32, env: Vec) -> Result<()> { + let id = self + .lock() + .vms + .values() + .find(|vm| vm.config.cid == cid) + .context("VM not found")? + .config + .manifest + .id + .clone(); + let _launch = self.launch_lock(&id).await; + let boot_env = { + let mut state = self.lock(); + let vm = state.vms.get_mut(&id).context("VM not found")?; + if vm.config.cid != cid { + bail!("VM CID changed"); + } + vm.state.env_update.take() + }; + let Some(boot_env) = boot_env else { + bail!("the env can only be updated once per boot, before applications start"); + }; + if env.is_empty() { + return Ok(()); + } + if self.sha256_of_env(&id)? != boot_env { + bail!("the env changed since the VM booted; not replacing it"); + } + self.work_dir(&id)?.write_env(&[], &env)?; + info!(id, "stored sealed env from the guest"); + Ok(()) } - pub(crate) fn encrypted_env_path(&self, id: &str) -> Result { - Ok(self.shared_dir(id)?.join(ENCRYPTED_ENV)) + fn sha256_of_env(&self, id: &str) -> Result<[u8; 32]> { + let shared = self.shared_dir(id)?; + let mut hasher = Sha256::new(); + for name in [ENCRYPTED_ENV, SEALED_ENV] { + hasher.update(name.as_bytes()); + match fs::read(shared.join(name)) { + Ok(data) => { + hasher.update([1]); + hasher.update((data.len() as u64).to_le_bytes()); + hasher.update(data); + } + Err(err) if err.kind() == std::io::ErrorKind::NotFound => hasher.update([0]), + Err(err) => return Err(err).context("failed to read env"), + } + } + Ok(hasher.finalize().into()) + } + + pub(crate) fn compose_file_path(&self, id: &str) -> Result { + Ok(self.shared_dir(id)?.join(APP_COMPOSE)) } pub(crate) fn user_config_path(&self, id: &str) -> Result { @@ -1605,15 +1679,13 @@ impl App { req: &VmConfiguration, app_id: &str, ) -> Result { + VmWorkDir::validate_env_update(&req.encrypted_env, &req.sealed_env)?; let work_dir = self.work_dir(id)?; let shared_dir = work_dir.join("shared"); fs::create_dir_all(&shared_dir).context("Failed to create shared directory")?; fs::write(shared_dir.join(APP_COMPOSE), &req.compose_file) .context("Failed to write compose file")?; - if !req.encrypted_env.is_empty() { - fs::write(shared_dir.join(ENCRYPTED_ENV), &req.encrypted_env) - .context("Failed to write encrypted env")?; - } + work_dir.write_env(&req.encrypted_env, &req.sealed_env)?; if !req.user_config.is_empty() { fs::write(shared_dir.join(USER_CONFIG), &req.user_config) .context("Failed to write user config")?; @@ -2508,6 +2580,83 @@ mod tests { (vm, networks) } + #[test] + fn prepare_work_dir_stores_sealed_and_rejects_both_formats() { + let dir = tempfile::tempdir().unwrap(); + let app = App::new( + test_config(&dir.path().join("absent-netd.sock"), dir.path()), + SupervisorClient::new("http://127.0.0.1:0"), + ); + let mut request = VmConfiguration { + compose_file: "{}".into(), + sealed_env: b"sealed".to_vec(), + ..Default::default() + }; + let workdir = app.prepare_work_dir("vm-1", &request, "").unwrap(); + assert_eq!(fs::read(workdir.sealed_env_path()).unwrap(), b"sealed"); + assert!(!workdir.encrypted_env_path().exists()); + request.encrypted_env = b"legacy".to_vec(); + assert!(app.prepare_work_dir("vm-2", &request, "").is_err()); + assert!(!app.work_dir("vm-2").unwrap().path().exists()); + } + + #[tokio::test] + async fn sealed_env_cas_binds_both_files_and_serializes_owner_updates() { + let dir = tempfile::tempdir().unwrap(); + let app = App::new( + test_config(&dir.path().join("absent-netd.sock"), dir.path()), + SupervisorClient::new("http://127.0.0.1:0"), + ); + let (vm, _) = bridge_vm(&app, "vm-1"); + let cid = vm.cid; + app.lock().vms.insert("vm-1".into(), VmState::new(vm)); + let workdir = app.work_dir("vm-1").unwrap(); + fs::create_dir_all(workdir.shared_dir()).unwrap(); + let boot = || { + let hash = app.sha256_of_env("vm-1").unwrap(); + app.lock().vms.get_mut("vm-1").unwrap().state.env_update = Some(hash); + }; + assert!(app.update_sealed_env(cid, b"new".to_vec()).await.is_err()); + workdir.write_env(b"legacy", &[]).unwrap(); + boot(); + app.update_sealed_env(cid, b"sealed".to_vec()) + .await + .unwrap(); + assert!(!workdir.encrypted_env_path().exists()); + assert_eq!(fs::read(workdir.sealed_env_path()).unwrap(), b"sealed"); + assert!(app.update_sealed_env(cid, b"again".to_vec()).await.is_err()); + + for name in [ENCRYPTED_ENV, SEALED_ENV] { + let path = workdir.shared_dir().join(name); + let _ = fs::remove_file(&path); + boot(); + fs::write(&path, []).unwrap(); + assert!(app.update_sealed_env(cid, b"stale".to_vec()).await.is_err()); + assert_eq!(fs::read(&path).unwrap(), b""); + } + + boot(); + let held = app.launch_lock("vm-1").await; + let guest = app.update_sealed_env(cid, b"stale".to_vec()); + tokio::pin!(guest); + assert!(tokio::time::timeout(Duration::from_millis(20), &mut guest) + .await + .is_err()); + workdir.write_env(b"owner", &[]).unwrap(); + drop(held); + assert!(guest.await.is_err()); + assert_eq!(fs::read(workdir.encrypted_env_path()).unwrap(), b"owner"); + assert!(!workdir.sealed_env_path().exists()); + + boot(); + app.update_sealed_env(cid, Vec::new()).await.unwrap(); + assert!(app.update_sealed_env(cid, b"again".to_vec()).await.is_err()); + boot(); + app.vm_event_report(cid, "boot.progress", "init-script".into()) + .unwrap(); + assert!(app.update_sealed_env(cid, b"late".to_vec()).await.is_err()); + } + /// A launch that never reached a netd built nothing, so it must not leave /// a marker that makes every later stop and removal wait for one. What /// an earlier boot may still hold keeps its marker. @@ -3950,6 +4099,10 @@ struct VmStateMut { auto_restart: AutoRestartState, /// True when the VM is being removed (cleanup in progress). removing: bool, + /// SHA-256 of both env files, including presence, this boot started with, while the guest + /// may still replace it; taken by the first update and cleared once the + /// boot reaches application code. + env_update: Option<[u8; 32]>, } impl VmStateMut { diff --git a/dstack/vmm/src/app/vm_info.rs b/dstack/vmm/src/app/vm_info.rs index 26e4e1f84..557a089e9 100644 --- a/dstack/vmm/src/app/vm_info.rs +++ b/dstack/vmm/src/app/vm_info.rs @@ -203,6 +203,7 @@ impl VmInfo { compose_file: fs::read_to_string(workdir.app_compose_path()) .unwrap_or_default(), encrypted_env: fs::read(workdir.encrypted_env_path()).unwrap_or_default(), + sealed_env: fs::read(workdir.sealed_env_path()).unwrap_or_default(), user_config: fs::read_to_string(workdir.user_config_path()).unwrap_or_default(), vcpu: self.manifest.vcpu, memory: self.manifest.memory, diff --git a/dstack/vmm/src/app/workdir.rs b/dstack/vmm/src/app/workdir.rs index fd67072f3..3fce85cfb 100644 --- a/dstack/vmm/src/app/workdir.rs +++ b/dstack/vmm/src/app/workdir.rs @@ -9,7 +9,9 @@ use std::path::{Path, PathBuf}; use anyhow::{Context, Result}; use dstack_types::{ - shared_filenames::{APP_COMPOSE, ENCRYPTED_ENV, INSTANCE_INFO, SYS_CONFIG, USER_CONFIG}, + shared_filenames::{ + APP_COMPOSE, ENCRYPTED_ENV, INSTANCE_INFO, SEALED_ENV, SYS_CONFIG, USER_CONFIG, + }, AppCompose, SysConfig, }; use fs_err as fs; @@ -131,10 +133,39 @@ impl VmWorkDir { self.shared_dir().join(USER_CONFIG) } + pub(crate) fn validate_env_update(encrypted: &[u8], sealed: &[u8]) -> Result<()> { + if !encrypted.is_empty() && !sealed.is_empty() { + anyhow::bail!("encrypted_env and sealed_env are mutually exclusive"); + } + Ok(()) + } + + /// Caller serializes updates to an existing VM with its launch lock. + pub(crate) fn write_env(&self, encrypted: &[u8], sealed: &[u8]) -> Result<()> { + Self::validate_env_update(encrypted, sealed)?; + let (path, other, data) = if !sealed.is_empty() { + (self.sealed_env_path(), self.encrypted_env_path(), sealed) + } else if !encrypted.is_empty() { + (self.encrypted_env_path(), self.sealed_env_path(), encrypted) + } else { + return Ok(()); + }; + safe_write::safe_write(path, data).context("failed to write env")?; + match fs::remove_file(other) { + Ok(()) => Ok(()), + Err(err) if err.kind() == std::io::ErrorKind::NotFound => Ok(()), + Err(err) => Err(err).context("failed to remove superseded env"), + } + } + pub fn encrypted_env_path(&self) -> PathBuf { self.shared_dir().join(ENCRYPTED_ENV) } + pub fn sealed_env_path(&self) -> PathBuf { + self.shared_dir().join(SEALED_ENV) + } + pub fn instance_info_path(&self) -> PathBuf { self.shared_dir().join(INSTANCE_INFO) } @@ -389,3 +420,39 @@ mod tests { workdir.remove_all().unwrap(); } } + +#[cfg(test)] +mod env_tests { + use super::*; + + #[test] + fn env_updates_replace_the_other_format_and_reject_ambiguity() { + let dir = tempfile::tempdir().unwrap(); + let workdir = VmWorkDir::new(dir.path()); + fs::create_dir_all(workdir.shared_dir()).unwrap(); + workdir.write_env(b"legacy", &[]).unwrap(); + workdir.write_env(&[], b"sealed").unwrap(); + assert!(!workdir.encrypted_env_path().exists()); + assert_eq!(fs::read(workdir.sealed_env_path()).unwrap(), b"sealed"); + assert!(workdir.write_env(b"ambiguous", b"ambiguous").is_err()); + assert_eq!(fs::read(workdir.sealed_env_path()).unwrap(), b"sealed"); + workdir.write_env(&[], &[]).unwrap(); + assert_eq!(fs::read(workdir.sealed_env_path()).unwrap(), b"sealed"); + workdir.write_env(b"owner", &[]).unwrap(); + assert!(!workdir.sealed_env_path().exists()); + assert_eq!(fs::read(workdir.encrypted_env_path()).unwrap(), b"owner"); + } + + #[test] + fn env_write_failure_preserves_original_and_cleanup_errors_propagate() { + let dir = tempfile::tempdir().unwrap(); + let workdir = VmWorkDir::new(dir.path()); + fs::create_dir_all(workdir.shared_dir()).unwrap(); + workdir.write_env(b"original", &[]).unwrap(); + fs::create_dir(workdir.sealed_env_path()).unwrap(); + assert!(workdir.write_env(&[], b"replacement").is_err()); + assert_eq!(fs::read(workdir.encrypted_env_path()).unwrap(), b"original"); + assert!(workdir.write_env(b"updated", &[]).is_err()); + assert_eq!(fs::read(workdir.encrypted_env_path()).unwrap(), b"updated"); + } +} diff --git a/dstack/vmm/src/config.rs b/dstack/vmm/src/config.rs index c923db519..4f76e742a 100644 --- a/dstack/vmm/src/config.rs +++ b/dstack/vmm/src/config.rs @@ -1432,6 +1432,19 @@ impl Config { mod tests { use super::*; + #[test] + fn host_api_limits_sealed_env_updates() { + use rocket::data::ToByteUnit; + use rocket::figment::providers::Serialized; + // Built the way `run_host_api` builds the host API server's config. + let figment = load_config_figment(None); + let figment = figment.clone().merge(Serialized::defaults( + figment.find_value("host_api").unwrap(), + )); + let config: rocket::Config = figment.extract().unwrap(); + assert_eq!(config.limits.get("UpdateSealedEnv"), Some(1.mebibytes())); + } + #[test] fn image_dir_stays_a_direct_child_of_the_image_root() { let cfg = ImageConfig { diff --git a/dstack/vmm/src/host_api_service.rs b/dstack/vmm/src/host_api_service.rs index fc4dcc704..0ad3c53eb 100644 --- a/dstack/vmm/src/host_api_service.rs +++ b/dstack/vmm/src/host_api_service.rs @@ -46,6 +46,12 @@ impl HostApiRpc for HostApiHandler { .vm_event_report(self.endpoint.cid, &request.event, request.payload) } + async fn update_sealed_env(self, request: host_api::SealedEnv) -> Result<()> { + self.app + .update_sealed_env(self.endpoint.cid, request.sealed_env) + .await + } + async fn get_sealing_key(self, request: GetSealingKeyRequest) -> Result { let key_provider = &self.app.config.key_provider; if !key_provider.enabled { diff --git a/dstack/vmm/src/main_service.rs b/dstack/vmm/src/main_service.rs index ffd06ad61..2cd917547 100644 --- a/dstack/vmm/src/main_service.rs +++ b/dstack/vmm/src/main_service.rs @@ -14,8 +14,9 @@ use dstack_vmm_rpc::vmm_server::{VmmRpc, VmmServer}; use dstack_vmm_rpc::{ AppId, ComposeHash as RpcComposeHash, GatewaySettings, GetInfoResponse, GetMetaResponse, Id, ImageInfo as RpcImageInfo, ImageListResponse, KmsSettings, ListGpusResponse, PublicKeyResponse, - ReloadVmsResponse, ResizeVmRequest, ResourcesSettings, StatusRequest, StatusResponse, - SvListResponse, SvProcessInfo, UpdateVmRequest, VersionResponse, VmConfiguration, + ReloadVmsResponse, ResizeVmRequest, ResourcesSettings, SealedEnvPubKeyResponse, StatusRequest, + StatusResponse, SvListResponse, SvProcessInfo, UpdateVmRequest, VersionResponse, + VmConfiguration, }; use fs_err as fs; use path_absolutize::Absolutize; @@ -311,6 +312,7 @@ pub fn create_manifest_from_vm_config( request: VmConfiguration, cvm_config: &crate::config::CvmConfig, ) -> Result { + VmWorkDir::validate_env_update(&request.encrypted_env, &request.sealed_env)?; validate_label(&request.name)?; // Same checks as `validate_resize_request`. if request.vcpu == 0 { @@ -1092,6 +1094,7 @@ impl VmmRpc for RpcHandler { // deleted that directory would recreate it holding nothing but a // manifest: invisible to `list_vms`, unloadable at every start, and // claiming the VM's netd interfaces against collection forever. + VmWorkDir::validate_env_update(&request.encrypted_env, &request.sealed_env)?; let _launch = self.app.launch_lock(&request.id).await; // Again under the lock: removal can have claimed the VM while this // waited for it. @@ -1207,11 +1210,7 @@ impl VmmRpc for RpcHandler { fs::write(self.compose_file_path(&request.id)?, &request.compose_file) .context("Failed to write compose file")?; } - if !request.encrypted_env.is_empty() { - let encrypted_env_path = self.encrypted_env_path(&request.id)?; - fs::write(encrypted_env_path, &request.encrypted_env) - .context("Failed to write encrypted env")?; - } + vm_work_dir.write_env(&request.encrypted_env, &request.sealed_env)?; if !request.user_config.is_empty() { let user_config_path = self.user_config_path(&request.id)?; fs::write(user_config_path, &request.user_config) @@ -1233,6 +1232,7 @@ impl VmmRpc for RpcHandler { let response = kms .get_app_env_encrypt_pub_key(dstack_kms_rpc::AppId { app_id: request.app_id, + root_epoch: None, }) .await?; Ok(PublicKeyResponse { @@ -1243,6 +1243,22 @@ impl VmmRpc for RpcHandler { }) } + async fn get_sealed_env_pub_key(self, request: AppId) -> Result { + let kms = self.kms_client()?; + let response = kms + .get_sealed_env_pub_key(dstack_kms_rpc::AppId { + app_id: request.app_id, + root_epoch: request.root_epoch, + }) + .await?; + Ok(SealedEnvPubKeyResponse { + public_key: response.public_key, + root_epoch: response.root_epoch, + timestamp: response.timestamp, + signature: response.signature, + }) + } + async fn get_info(self, request: Id) -> Result { info!(vm_id = %request.id, "get_info RPC called"); if let Some(vm) = self.app.vm_info(&request.id).await? { @@ -1497,6 +1513,7 @@ mod tests { disk_size: 10, ports: vec![], encrypted_env: vec![], + sealed_env: vec![], app_id: None, user_config: String::new(), hugepages: false, diff --git a/dstack/vmm/src/one_shot.rs b/dstack/vmm/src/one_shot.rs index 4aa36741b..9989472ea 100644 --- a/dstack/vmm/src/one_shot.rs +++ b/dstack/vmm/src/one_shot.rs @@ -216,10 +216,7 @@ Compose file content (first 200 chars): .context("Failed to write app compose file")?; // Write other files if present - if !vm_config.encrypted_env.is_empty() { - fs_err::write(vm_work_dir.encrypted_env_path(), &vm_config.encrypted_env) - .context("Failed to write encrypted env")?; - } + vm_work_dir.write_env(&vm_config.encrypted_env, &vm_config.sealed_env)?; if !vm_config.user_config.is_empty() { fs_err::write(vm_work_dir.user_config_path(), &vm_config.user_config) diff --git a/dstack/vmm/src/tests/test_vmm_cli.py b/dstack/vmm/src/tests/test_vmm_cli.py index f3c33c371..23564bad2 100644 --- a/dstack/vmm/src/tests/test_vmm_cli.py +++ b/dstack/vmm/src/tests/test_vmm_cli.py @@ -5,12 +5,20 @@ from __future__ import annotations +import argparse +import base64 import contextlib import importlib.util import io +import json import sys import unittest from pathlib import Path +from unittest.mock import Mock, patch + +from eth_keys import keys +from eth_utils import keccak +from pyhpke import AEADId, CipherSuite, KDFId, KEMId def load_vmm_cli(): @@ -55,5 +63,261 @@ def test_kms_urls_are_sent_to_upgrade_app(self) -> None: ) +class SealedEnvTests(unittest.TestCase): + def setUp(self): + self.app_id = "11" * 20 + self.epoch = 3 + self.timestamp = 1700000000 + self.clock = patch.object(VMM_CLI.time, "time", return_value=self.timestamp) + self.clock.start() + self.addCleanup(self.clock.stop) + self.signer = keys.PrivateKey(b"\x22" * 32) + self.pin = self.signer.public_key.to_compressed_bytes().hex() + self.suite = CipherSuite.new( + KEMId.DHKEM_X25519_HKDF_SHA256, KDFId.HKDF_SHA256, AEADId.AES256_GCM + ) + self.recipient = self.suite.kem.derive_key_pair(b"test recipient key material") + self.response = { + "public_key": self.recipient.public_key.to_public_bytes().hex(), + "root_epoch": self.epoch, + "timestamp": self.timestamp, + } + digest = keccak( + b"dstack-sealed-env-pubkey-v1:" + + bytes.fromhex(self.app_id) + + self.epoch.to_bytes(4, "big") + + bytes.fromhex(self.response["public_key"]) + + self.timestamp.to_bytes(8, "big") + ) + self.response["signature"] = self.signer.sign_msg_hash(digest).to_bytes().hex() + self.cli = VMM_CLI.VmmCLI("https://vmm.example") + self.cli.client.request = Mock(return_value=(200, self.response)) + self.cli.get_app_env_encrypt_pub_key = Mock( + return_value=self.response["public_key"] + ) + self.pins = patch.object(VMM_CLI, "load_whitelist", return_value=[self.pin]) + self.pins.start() + self.addCleanup(self.pins.stop) + self.envs = {"SECRET": "synthetic test value", "UNICODE": "测试"} + + def decrypt(self, payload): + self.assertNotIn("encrypted_env", payload) + envelope = json.loads(bytes.fromhex(payload["sealed_env"])) + self.assertEqual(set(envelope), {"version", "root_epoch", "enc", "ciphertext"}) + self.assertEqual(envelope["version"], 1) + self.assertEqual(envelope["root_epoch"], self.epoch) + enc = base64.b64decode(envelope["enc"], validate=True) + ciphertext = base64.b64decode(envelope["ciphertext"], validate=True) + receiver = self.suite.create_recipient_context( + enc, + self.recipient.private_key, + info=f"dstack-sealed-env-v1:{self.epoch}".encode(), + ) + plaintext = json.loads(receiver.open(ciphertext, aad=b"")) + self.assertEqual( + plaintext, {"env": [{"key": k, "value": v} for k, v in self.envs.items()]} + ) + wrong_epoch = self.suite.create_recipient_context( + enc, self.recipient.private_key, info=b"dstack-sealed-env-v1:0" + ) + with self.assertRaises(Exception): + wrong_epoch.open(ciphertext, aad=b"") + + def test_sealed_payload_round_trip_and_normalized_pins(self): + for pin in [self.pin, "0x" + self.pin.upper()]: + with ( + self.subTest(pin=pin), + patch.object(VMM_CLI, "load_whitelist", return_value=[pin]), + ): + self.decrypt(self.cli.encrypt_env_payload(self.envs, self.app_id, None)) + self.cli.get_app_env_encrypt_pub_key.assert_not_called() + + def test_unpinned_signer_never_becomes_trusted(self): + for pins in [ + [], + [keys.PrivateKey(b"\x33" * 32).public_key.to_compressed_bytes().hex()], + ]: + with ( + self.subTest(pins=pins), + patch.object(VMM_CLI, "load_whitelist", return_value=pins), + patch.object(self.cli, "confirm_untrusted_signer") as confirm, + ): + with self.assertRaisesRegex(ValueError, "not pinned"): + self.cli.encrypt_env_payload(self.envs, self.app_id, None) + confirm.assert_not_called() + self.cli.get_app_env_encrypt_pub_key.assert_not_called() + + def test_signature_binds_every_field_and_app(self): + cases = [ + ("public_key", "44" * 32), + ("root_epoch", 4), + ("timestamp", self.timestamp + 1), + ("signature", "00" * 65), + ("signature", "00" * 64), + ("public_key", "00" * 31), + ("root_epoch", -1), + ("root_epoch", 2**32), + ("root_epoch", True), + ("timestamp", -1), + ("timestamp", 2**64), + ("timestamp", True), + ] + for field, value in cases: + with self.subTest(field=field, value=value): + self.cli.client.request.return_value = ( + 200, + {**self.response, field: value}, + ) + with self.assertRaises(Exception): + self.cli.encrypt_env_payload(self.envs, self.app_id, None) + self.cli.client.request.return_value = (200, self.response) + for app_id in ["55" * 20, "11" * 19]: + with self.assertRaises(ValueError): + self.cli.encrypt_env_payload(self.envs, app_id, None) + self.cli.get_app_env_encrypt_pub_key.assert_not_called() + + def test_pinned_signed_key_timestamp_window(self): + for age in [-60, 0, 300]: + with self.subTest(age=age), patch.object( + VMM_CLI.time, "time", return_value=self.timestamp + age + ): + self.decrypt(self.cli.encrypt_env_payload(self.envs, self.app_id, None)) + for age, reason in [(-61, "future"), (301, "stale")]: + with self.subTest(age=age), patch.object( + VMM_CLI.time, "time", return_value=self.timestamp + age + ): + with self.assertRaisesRegex(ValueError, reason): + self.cli.encrypt_env_payload(self.envs, self.app_id, None) + self.cli.get_app_env_encrypt_pub_key.assert_not_called() + + def test_requested_epoch_is_sent_and_checked(self): + self.cli.get_sealed_env_pub_key(self.app_id, root_epoch=3) + self.assertEqual( + self.cli.client.request.call_args.kwargs["body"], + {"app_id": self.app_id, "root_epoch": 3}, + ) + with self.assertRaisesRegex(ValueError, "does not match"): + self.cli.get_sealed_env_pub_key(self.app_id, root_epoch=2) + + def test_only_explicit_unsupported_method_falls_back(self): + for status in [400, 404]: + self.cli.client.request.return_value = ( + status, + {"error": "Service not found: GetSealedEnvPubKey"}, + ) + payload = self.cli.encrypt_env_payload(self.envs, self.app_id, None) + self.assertEqual(set(payload), {"encrypted_env"}) + self.cli.get_app_env_encrypt_pub_key.assert_called_with(self.app_id, None) + self.cli.get_app_env_encrypt_pub_key.side_effect = ValueError( + "legacy trust rejected" + ) + with self.assertRaisesRegex(ValueError, "legacy trust rejected"): + self.cli.encrypt_env_payload(self.envs, self.app_id, None) + + def test_errors_and_bad_responses_do_not_downgrade(self): + for status, body in [ + (401, {"error": "Service not found: GetSealedEnvPubKey"}), + (403, {"error": "forbidden"}), + (404, "Not Found"), + (400, {"error": "invalid app"}), + (500, {"error": "unavailable"}), + (200, {}), + (200, "invalid JSON"), + ]: + with self.subTest(status=status, body=body): + self.cli.client.request.return_value = (status, body) + with self.assertRaises(Exception): + self.cli.encrypt_env_payload(self.envs, self.app_id, None) + self.cli.client.request.side_effect = TimeoutError("test timeout") + with self.assertRaises(TimeoutError): + self.cli.encrypt_env_payload(self.envs, self.app_id, None) + self.cli.get_app_env_encrypt_pub_key.assert_not_called() + + def test_direct_kms_verifies_pin_and_does_not_forward_vmm_credentials(self): + with patch.object(VMM_CLI, "VmmClient") as client: + client.return_value.request.return_value = (200, self.response) + self.decrypt( + self.cli.encrypt_env_payload( + self.envs, self.app_id, "https://kms.example" + ) + ) + client.assert_called_once_with("https://kms.example", verify_tls=False) + self.cli.get_app_env_encrypt_pub_key.assert_not_called() + + def test_create_update_and_update_env_send_only_the_negotiated_field(self): + compose = json.dumps({"kms_enabled": True}) + args = argparse.Namespace( + compose="compose.json", + env_file="env", + name="vm", + image="image", + vcpu=1, + memory=1024, + disk=10, + app_id=self.app_id, + user_config=None, + port=None, + hugepages=False, + pin_numa=False, + stopped=False, + no_tee=False, + annotation=None, + disk_prealloc=None, + simulated_tee=None, + swap=None, + ppcie=False, + gpu=None, + kms_url=["https://kms.example"], + gateway_url=None, + net_queues=None, + net=None, + net_vhost=None, + kms_encrypt_url=None, + ) + + def rpc(method, params=None): + if method == "GetInfo": + return { + "found": True, + "info": { + "app_id": self.app_id, + "configuration": {"compose_file": compose}, + }, + } + return {"id": "vm"} + + self.cli.rpc_call = Mock(side_effect=rpc) + with ( + patch.object(VMM_CLI.os.path, "exists", return_value=True), + patch.object(VMM_CLI, "read_utf8", return_value=compose), + patch.object(VMM_CLI, "parse_env_file", return_value=self.envs), + patch.object(VMM_CLI, "VmmClient") as client, + contextlib.redirect_stdout(io.StringIO()), + ): + for supported in [True, False]: + client.return_value.request.return_value = ( + (200, self.response) + if supported + else (404, {"error": "Service not found: GetSealedEnvPubKey"}) + ) + for action in [ + lambda: self.cli.create_vm(args), + lambda: self.cli.update_vm( + "vm", env_file="env", kms_urls=args.kms_url + ), + lambda: self.cli.update_vm_env("vm", self.envs, args.kms_url), + ]: + with self.subTest(supported=supported, action=action): + self.cli.rpc_call.reset_mock() + action() + method, payload = self.cli.rpc_call.call_args.args + self.assertIn(method, ["CreateVm", "UpgradeApp"]) + if supported: + self.decrypt(payload) + else: + self.assertIn("encrypted_env", payload) + self.assertNotIn("sealed_env", payload) + + if __name__ == "__main__": unittest.main() diff --git a/dstack/vmm/src/vmm-cli.py b/dstack/vmm/src/vmm-cli.py index ad469f981..0001c582f 100755 --- a/dstack/vmm/src/vmm-cli.py +++ b/dstack/vmm/src/vmm-cli.py @@ -16,6 +16,7 @@ import ssl import sys import tempfile +import time import urllib.parse from typing import Any, BinaryIO, Dict, List, Optional, Tuple, Union @@ -310,6 +311,77 @@ def cmd_switch_vmm(args): print(f"Switched to VMM {selected['id'][:8]} ({url})") +def seal_env(envs, hex_public_key: str, root_epoch: int) -> str: + """Return hex-encoded SealedEnv JSON using RFC 9180 base mode.""" + from pyhpke import AEADId, CipherSuite, KDFId, KEMId + + validate_uint(root_epoch, 32, "root_epoch") + suite = CipherSuite.new( + KEMId.DHKEM_X25519_HKDF_SHA256, KDFId.HKDF_SHA256, AEADId.AES256_GCM + ) + public_key = suite.kem.deserialize_public_key( + bytes.fromhex(hex_public_key.removeprefix("0x")) + ) + enc, context = suite.create_sender_context( + public_key, info=f"dstack-sealed-env-v1:{root_epoch}".encode("utf-8") + ) + plaintext = json.dumps({"env": envs}).encode("utf-8") + envelope = { + "version": 1, + "root_epoch": root_epoch, + "enc": base64.b64encode(enc).decode("ascii"), + "ciphertext": base64.b64encode(context.seal(plaintext, aad=b"")).decode( + "ascii" + ), + } + return json.dumps(envelope, separators=(",", ":")).encode("utf-8").hex() + + +def validate_uint(value: int, bits: int, name: str) -> None: + if type(value) is not int or not 0 <= value < 1 << bits: + raise ValueError(f"{name} must be a uint{bits}") + + +def verify_sealed_env_pubkey( + response: Dict, app_id: str, root_epoch: Optional[int] = None +) -> None: + """Authenticate a SealedEnv key against the locally pinned epoch signers.""" + if not CRYPTO_AVAILABLE: + raise ImportError("Install the VMM CLI requirements to verify SealedEnv keys") + app_id_bytes = bytes.fromhex(app_id.removeprefix("0x")) + public_key = bytes.fromhex(response["public_key"].removeprefix("0x")) + signature = bytes.fromhex(response["signature"].removeprefix("0x")) + epoch, timestamp = response["root_epoch"], response["timestamp"] + validate_uint(epoch, 32, "root_epoch") + validate_uint(timestamp, 64, "timestamp") + now = int(time.time()) + if now - timestamp > 300: + raise ValueError("SealedEnv public key timestamp is stale") + if timestamp - now > 60: + raise ValueError("SealedEnv public key timestamp is in the future") + if len(app_id_bytes) != 20 or len(public_key) != 32 or len(signature) != 65: + raise ValueError("Invalid SealedEnv app ID, public key or signature length") + if root_epoch is not None: + validate_uint(root_epoch, 32, "root_epoch") + if epoch != root_epoch: + raise ValueError("SealedEnv root epoch does not match the requested epoch") + digest = keccak( + b"dstack-sealed-env-pubkey-v1:" + + app_id_bytes + + epoch.to_bytes(4, "big") + + public_key + + timestamp.to_bytes(8, "big") + ) + signer = keys.Signature(signature_bytes=signature).recover_public_key_from_msg_hash( + digest + ) + pins = {bytes.fromhex(pin.removeprefix("0x")) for pin in load_whitelist()} + if signer.to_compressed_bytes() not in pins: + raise ValueError( + "SealedEnv signer is not pinned; configure a trusted epoch signer with kms add" + ) + + def encrypt_env(envs, hex_public_key: str) -> str: """Encrypts environment variables using a one-time X25519 key exchange and AES-GCM. @@ -776,6 +848,49 @@ def list_images(self, json_output: bool = False) -> None: rows = [[img["name"], img.get("version", "-")] for img in images] print(format_table(rows, headers)) + def get_sealed_env_pub_key( + self, + app_id: str, + kms_url: Optional[str] = None, + root_epoch: Optional[int] = None, + ) -> Optional[Dict]: + """Return an authenticated key, or None only for an unsupported RPC.""" + params = {"app_id": app_id} + if root_epoch is not None: + validate_uint(root_epoch, 32, "root_epoch") + params["root_epoch"] = root_epoch + # KMS RA-TLS is authenticated by the pinned signer below. The VMM + # path retains its configured TLS verification and authentication. + client = VmmClient(kms_url, verify_tls=False) if kms_url else self.client + status, response = client.request( + "POST", + "/prpc/GetSealedEnvPubKey?json", + headers={"Content-Type": "application/json"}, + body=params, + ) + error = response.get("error") if isinstance(response, dict) else response + if status in (400, 404) and error == "Service not found: GetSealedEnvPubKey": + return None + if status != 200: + raise RuntimeError(f"GetSealedEnvPubKey failed with HTTP {status}") + if not isinstance(response, dict): + raise ValueError("Invalid GetSealedEnvPubKey response") + verify_sealed_env_pubkey(response, app_id, root_epoch) + return response + + def encrypt_env_payload( + self, envs: Dict[str, str], app_id: str, kms_url: Optional[str] + ) -> Dict: + """Select the wire field by RPC support, never by environment file contents.""" + envs_list = [{"key": k, "value": v} for k, v in envs.items()] + key = self.get_sealed_env_pub_key(app_id, kms_url) + if key is not None: + return { + "sealed_env": seal_env(envs_list, key["public_key"], key["root_epoch"]) + } + public_key = self.get_app_env_encrypt_pub_key(app_id, kms_url) + return {"encrypted_env": encrypt_env(envs_list, public_key)} + def get_app_env_encrypt_pub_key( self, app_id: str, kms_url: Optional[str] = None ) -> Dict: @@ -1039,10 +1154,7 @@ def create_vm(self, args) -> None: encrypt_url = args.kms_encrypt_url or ( args.kms_url[0] if args.kms_url else None ) - encrypt_pubkey = self.get_app_env_encrypt_pub_key(app_id, encrypt_url) - print(f"Encrypting environment variables with key: {encrypt_pubkey}") - envs_list = [{"key": k, "value": v} for k, v in envs.items()] - params["encrypted_env"] = encrypt_env(envs_list, encrypt_pubkey) + params.update(self.encrypt_env_payload(envs, app_id, encrypt_url)) response = self.rpc_call("CreateVm", params) print(f"Created VM with ID: {response.get('id')}") return response.get("id") @@ -1067,16 +1179,8 @@ def update_vm_env( vm_configuration = vm_info_response["info"].get("configuration") or {} compose_file = vm_configuration.get("compose_file") - # Now get the encryption key for the app - encrypt_pubkey = self.get_app_env_encrypt_pub_key( - app_id, kms_urls[0] if kms_urls else None - ) - print(f"Encrypting environment variables with key: {encrypt_pubkey}") - envs_list = [{"key": k, "value": v} for k, v in envs.items()] - encrypted_env = encrypt_env(envs_list, encrypt_pubkey) - - # Use UpdateApp with the VM ID - payload = {"id": vm_id, "encrypted_env": encrypted_env} + payload = {"id": vm_id} + payload.update(self.encrypt_env_payload(envs, app_id, kms_urls[0])) if compose_file: try: @@ -1238,11 +1342,7 @@ def update_vm( vm_configuration = vm_info_response["info"].get("configuration") or {} compose_file_content = vm_configuration.get("compose_file") - encrypt_pubkey = self.get_app_env_encrypt_pub_key( - app_id, kms_urls[0] if kms_urls else None - ) - envs_list = [{"key": k, "value": v} for k, v in envs.items()] - upgrade_params["encrypted_env"] = encrypt_env(envs_list, encrypt_pubkey) + upgrade_params.update(self.encrypt_env_payload(envs, app_id, kms_urls[0])) updates.append("environment variables") # update allowed_envs in compose file if needed diff --git a/dstack/vmm/ui/build.mjs b/dstack/vmm/ui/build.mjs index 00c06fc5a..85a814731 100644 --- a/dstack/vmm/ui/build.mjs +++ b/dstack/vmm/ui/build.mjs @@ -24,6 +24,9 @@ function canonicalId(absPath) { } function resolveModule(parentId, request) { + if (request === 'crypto') { + return canonicalId(path.join(TS_OUT_DIR, 'lib/browserCrypto.js')); + } const base = parentId ? path.dirname(path.resolve(ROOT, parentId)) : MODULE_DIR; const absPath = nodeRequire.resolve(request, { paths: [base] }); return canonicalId(absPath); @@ -162,6 +165,7 @@ async function compileTypeScript() { await fs.rm(TS_OUT_DIR, { recursive: true, force: true }); await run(TSC, ['--project', path.join(ROOT, 'tsconfig.json')]); await copyDir(path.join(SOURCE_DIR, 'templates'), path.join(TS_OUT_DIR, 'templates')); + await fs.writeFile(path.join(TS_OUT_DIR, 'package.json'), JSON.stringify({ type: 'commonjs' })); } async function build({ watch = false } = {}) { diff --git a/dstack/vmm/ui/package-lock.json b/dstack/vmm/ui/package-lock.json index b04b7ed8d..1d9101ebf 100644 --- a/dstack/vmm/ui/package-lock.json +++ b/dstack/vmm/ui/package-lock.json @@ -7,6 +7,10 @@ "": { "name": "vmm-ui", "version": "0.1.0", + "dependencies": { + "@hpke/core": "1.10.0", + "@hpke/dhkem-x25519": "1.8.0" + }, "devDependencies": { "@types/node": "^20.11.30", "protobufjs": "^7.5.8", @@ -64,6 +68,47 @@ "node": ">=6.9.0" } }, + "node_modules/@hpke/common": { + "version": "1.12.0", + "resolved": "https://registry.npmjs.org/@hpke/common/-/common-1.12.0.tgz", + "integrity": "sha512-os6y44MDJmrx7MVAG8rb6LyV3BEXSFtX/g4KMhKxVtpXvCWwL4fbBbbtG45fkx25JmxpCJdGzla9asLNMI99/A==", + "license": "MIT", + "engines": { + "node": ">=16.0.0" + }, + "funding": { + "type": "github", + "url": "https://github.com/sponsors/dajiaji" + } + }, + "node_modules/@hpke/core": { + "version": "1.10.0", + "resolved": "https://registry.npmjs.org/@hpke/core/-/core-1.10.0.tgz", + "integrity": "sha512-toCyUYXvU5KdDjqtWeMRRhR09R2us4OeHGo5YkdDOpo8+QJ6sIpuuOGLDu0VbNBskf5acHgsqgF53w2mU1NNYQ==", + "license": "MIT", + "dependencies": { + "@hpke/common": "^1.11.0" + }, + "engines": { + "node": ">=16.0.0" + }, + "funding": { + "type": "github", + "url": "https://github.com/sponsors/dajiaji" + } + }, + "node_modules/@hpke/dhkem-x25519": { + "version": "1.8.0", + "resolved": "https://registry.npmjs.org/@hpke/dhkem-x25519/-/dhkem-x25519-1.8.0.tgz", + "integrity": "sha512-S1MWWkAfu+TFxySgv5+2P3O4Mx/jk7BsoplzQaA1s3sfUJVJ2UsZsSzSsMc+FXJumLXncoJFlO6mK6mDGspfmA==", + "license": "MIT", + "dependencies": { + "@hpke/common": "^1.10.0" + }, + "engines": { + "node": ">=16.0.0" + } + }, "node_modules/@jsdoc/salty": { "version": "0.2.9", "resolved": "https://registry.npmjs.org/@jsdoc/salty/-/salty-0.2.9.tgz", diff --git a/dstack/vmm/ui/package.json b/dstack/vmm/ui/package.json index b88e34045..985719998 100644 --- a/dstack/vmm/ui/package.json +++ b/dstack/vmm/ui/package.json @@ -6,12 +6,17 @@ "scripts": { "generate:proto": "bash scripts/build_proto.sh", "build": "node build.mjs", - "watch": "node build.mjs --watch" + "watch": "node build.mjs --watch", + "test": "DSTACK_UI_OUT=dist/index.html npm run build && node --test tests/*.test.cjs" }, "devDependencies": { "@types/node": "^20.11.30", "protobufjs": "^7.5.8", "protobufjs-cli": "^1.3.3", "typescript": "^5.4.5" + }, + "dependencies": { + "@hpke/core": "1.10.0", + "@hpke/dhkem-x25519": "1.8.0" } } diff --git a/dstack/vmm/ui/src/composables/useVmManager.ts b/dstack/vmm/ui/src/composables/useVmManager.ts index f81aede22..5b9fbab32 100644 --- a/dstack/vmm/ui/src/composables/useVmManager.ts +++ b/dstack/vmm/ui/src/composables/useVmManager.ts @@ -4,6 +4,7 @@ declare const Vue: any; const { ref, computed, watch, onMounted } = Vue; import type { vmm as VmmTypes } from '../proto/vmm_rpc'; +import { encryptEnvPayload } from '../lib/envEncryption'; // Types based on Rust definitions type VmConfiguration = VmmTypes.IVmConfiguration; @@ -643,6 +644,7 @@ type CreateVmPayloadSource = { disk_size: number; ports: PortFormEntry[]; encrypted_env?: Uint8Array; + sealed_env?: Uint8Array; app_id?: string | null; user_config?: string; hugepages?: boolean; @@ -670,6 +672,7 @@ type CreateVmPayloadSource = { disk_size: Math.max(0, Number(source.disk_size) || 0), ports: normalizedPorts, encrypted_env: source.encrypted_env, + sealed_env: source.sealed_env, app_id: source.app_id || undefined, user_config: source.user_config || '', hugepages: !!source.hugepages, @@ -1222,7 +1225,7 @@ type CreateVmPayloadSource = { async function encryptEnv(envs: EncryptedEnvEntry[], kmsEnabled: boolean, appId: string | null) { if (!kmsEnabled || envs.length === 0) { - return undefined; + return {}; } let appIdToUse = appId; if (!appIdToUse) { @@ -1230,8 +1233,7 @@ type CreateVmPayloadSource = { appIdToUse = await calcAppId(appCompose); } const keyBytes = hexToBytes(appIdToUse); - const response = await vmmRpc.getAppEnvEncryptPubKey({ app_id: keyBytes }); - return encryptEnvWithKey(envs, response.public_key); + return encryptEnvPayload(vmmRpc, envs, keyBytes, encryptEnvWithKey); } async function encryptEnvWithKey(envs: EncryptedEnvEntry[], publicKeyBytes: Uint8Array) { @@ -1281,7 +1283,7 @@ type CreateVmPayloadSource = { memory: vmForm.value.memory, disk_size: vmForm.value.disk_size, ports: vmForm.value.ports, - encrypted_env: encryptedEnv || undefined, + ...encryptedEnv, app_id: vmForm.value.app_id || undefined, user_config: vmForm.value.user_config, hugepages: vmForm.value.hugepages, @@ -1362,14 +1364,15 @@ type CreateVmPayloadSource = { const composeWasExplicitlyUpdated = updateDialog.value.updateCompose; let composeNeedsUpdate = composeWasExplicitlyUpdated; - let encryptedEnvPayload; + let encryptedEnvPayload = {}; if (updateDialog.value.resetSecrets) { - const keyResponse = await vmmRpc.getAppEnvEncryptPubKey({ app_id: hexToBytes(vm.app_id || '') }); - encryptedEnvPayload = await encryptEnvWithKey(updateDialog.value.encryptedEnvs, keyResponse.public_key); + encryptedEnvPayload = await encryptEnvPayload( + vmmRpc, updateDialog.value.encryptedEnvs, hexToBytes(vm.app_id || ''), encryptEnvWithKey, + ); composeNeedsUpdate = true; } body.compose_file = composeNeedsUpdate ? await makeUpdateComposeFile() : undefined; - body.encrypted_env = encryptedEnvPayload; + Object.assign(body, encryptedEnvPayload); body.user_config = updated.user_config; body.update_ports = true; body.ports = normalizePorts(updated.ports); diff --git a/dstack/vmm/ui/src/lib/browserCrypto.ts b/dstack/vmm/ui/src/lib/browserCrypto.ts new file mode 100644 index 000000000..37da2d201 --- /dev/null +++ b/dstack/vmm/ui/src/lib/browserCrypto.ts @@ -0,0 +1,6 @@ +// SPDX-FileCopyrightText: © 2026 Phala Network +// +// SPDX-License-Identifier: Apache-2.0 + +// HPKE's Node fallback uses the browser's native Web Crypto in this bundle. +export const webcrypto = globalThis.crypto; diff --git a/dstack/vmm/ui/src/lib/envEncryption.ts b/dstack/vmm/ui/src/lib/envEncryption.ts new file mode 100644 index 000000000..4a29d3f8c --- /dev/null +++ b/dstack/vmm/ui/src/lib/envEncryption.ts @@ -0,0 +1,68 @@ +// SPDX-FileCopyrightText: © 2026 Phala Network +// +// SPDX-License-Identifier: Apache-2.0 + +import { Aes256Gcm, CipherSuite, HkdfSha256 } from '@hpke/core'; +import { DhkemX25519HkdfSha256 } from '@hpke/dhkem-x25519'; +import type { vmm as VmmTypes } from '../proto/vmm_rpc'; +import { RpcError } from './vmmRpcClient'; + +export type EnvEntry = { key: string; value: string }; +type EnvPayload = { sealed_env?: Uint8Array; encrypted_env?: Uint8Array }; +type KeyClient = { + getSealedEnvPubKey(request: VmmTypes.IAppId): Promise; + getAppEnvEncryptPubKey(request: VmmTypes.IAppId): Promise; +}; + +const missingService = 'Service not found: GetSealedEnvPubKey'; + +function isUnsupported(error: unknown): boolean { + if (!(error instanceof RpcError) || ![400, 404].includes(error.status)) return false; + if (error.message === missingService) return true; + // The VMM forwards an older KMS's JSON RPC error through its binary RPC. + const forwarded = /^Request failed with status=(?:400 Bad Request|404 Not Found), error=(.+)$/s.exec(error.message); + if (!forwarded) return false; + try { + return JSON.parse(forwarded[1]).error === missingService; + } catch { + return false; + } +} + +/** Keep key retrieval on the authenticated, same-origin VMM transport. */ +export async function encryptEnvPayload( + client: KeyClient, + envs: EnvEntry[], + appId: Uint8Array, + encryptLegacy: (envs: EnvEntry[], key: Uint8Array) => Promise, +): Promise { + let response: VmmTypes.ISealedEnvPubKeyResponse; + try { + response = await client.getSealedEnvPubKey({ app_id: appId }); + } catch (error) { + if (!isUnsupported(error)) throw error; + const legacy = await client.getAppEnvEncryptPubKey({ app_id: appId }); + return { encrypted_env: await encryptLegacy(envs, legacy.public_key) }; + } + const epoch = response.root_epoch; + if (!Number.isInteger(epoch) || epoch < 0 || epoch > 0xffffffff || response.public_key?.length !== 32) { + throw new Error('invalid sealed env public key response'); + } + const suite = new CipherSuite({ + kem: new DhkemX25519HkdfSha256(), + kdf: new HkdfSha256(), + aead: new Aes256Gcm(), + }); + const recipientPublicKey = await suite.kem.deserializePublicKey(new Uint8Array(response.public_key)); + const { enc, ct } = await suite.seal({ + recipientPublicKey, + info: new TextEncoder().encode(`dstack-sealed-env-v1:${epoch}`), + }, new TextEncoder().encode(JSON.stringify({ env: envs }))); + const base64 = (bytes: ArrayBuffer) => btoa(Array.from(new Uint8Array(bytes), (byte) => String.fromCharCode(byte)).join('')); + return { sealed_env: new TextEncoder().encode(JSON.stringify({ + version: 1, + root_epoch: epoch, + enc: base64(enc), + ciphertext: base64(ct), + })) }; +} diff --git a/dstack/vmm/ui/src/lib/vmmRpcClient.ts b/dstack/vmm/ui/src/lib/vmmRpcClient.ts index 683f499fd..e1ae87505 100644 --- a/dstack/vmm/ui/src/lib/vmmRpcClient.ts +++ b/dstack/vmm/ui/src/lib/vmmRpcClient.ts @@ -8,6 +8,13 @@ const textDecoder = new TextDecoder(); const EMPTY_BODY = new Uint8Array(); let cachedClient: any; +export class RpcError extends Error { + constructor(message: string, readonly status: number) { + super(message); + this.name = 'RpcError'; + } +} + function decodePrpcError(buffer: Uint8Array) { try { if (buffer && buffer.length > 0) { @@ -76,7 +83,7 @@ export function getVmmRpcClient(basePath = '/prpc') { .then(async (response) => { const buffer = new Uint8Array(await response.arrayBuffer()); if (!response.ok) { - callback(new Error(decodePrpcError(buffer))); + callback(new RpcError(decodePrpcError(buffer), response.status)); return; } callback(null, buffer); diff --git a/dstack/vmm/ui/tests/envEncryption.test.cjs b/dstack/vmm/ui/tests/envEncryption.test.cjs new file mode 100644 index 000000000..3ee0e4d12 --- /dev/null +++ b/dstack/vmm/ui/tests/envEncryption.test.cjs @@ -0,0 +1,215 @@ +// SPDX-FileCopyrightText: © 2026 Phala Network +// +// SPDX-License-Identifier: Apache-2.0 + +const { test } = require('node:test'); +const assert = require('node:assert/strict'); +const { readFileSync } = require('node:fs'); +const { runInNewContext } = require('node:vm'); +const { webcrypto } = require('node:crypto'); +const { CipherSuite, Aes256Gcm, HkdfSha256 } = require('@hpke/core'); +const { DhkemX25519HkdfSha256 } = require('@hpke/dhkem-x25519'); +const { encryptEnvPayload } = require('../build/ts/lib/envEncryption.js'); +const { RpcError } = require('../build/ts/lib/vmmRpcClient.js'); +const { vmm } = require('../build/ts/proto/vmm_rpc.js'); +const { prpc } = require('../build/ts/proto/prpc.js'); + +const suite = new CipherSuite({ kem: new DhkemX25519HkdfSha256(), kdf: new HkdfSha256(), aead: new Aes256Gcm() }); +const appId = new Uint8Array(20).fill(0x11); +const envs = [{ key: 'GREETING', value: 'Hello 世界' }]; +const missing = 'Service not found: GetSealedEnvPubKey'; + +async function recipient() { + const pair = await suite.kem.generateKeyPair(); + const public_key = new Uint8Array(await suite.kem.serializePublicKey(pair.publicKey)); + return { pair, response: { public_key, root_epoch: 7, timestamp: 123, signature: new Uint8Array(65) } }; +} + +async function open(pair, bytes, epoch) { + const envelope = JSON.parse(new TextDecoder().decode(bytes)); + assert.equal(envelope.version, 1); + const pt = await suite.open({ + recipientKey: pair.privateKey, + enc: Buffer.from(envelope.enc, 'base64'), + info: new TextEncoder().encode(`dstack-sealed-env-v1:${epoch ?? envelope.root_epoch}`), + }, Buffer.from(envelope.ciphertext, 'base64')); + return JSON.parse(new TextDecoder().decode(pt)); +} + +const noLegacy = async () => { throw new Error('unexpected legacy fallback'); }; + +test('seals JSON bytes with HPKE, binds epoch, and rejects tampering', async () => { + const { pair, response } = await recipient(); + for (const epoch of [0, 7, 0xffffffff]) { + const payload = await encryptEnvPayload({ + getSealedEnvPubKey: async (request) => { + assert.deepEqual(request, { app_id: appId }); + return { ...response, root_epoch: epoch }; + }, + getAppEnvEncryptPubKey: noLegacy, + }, envs, appId, noLegacy); + assert.deepEqual(Object.keys(payload), ['sealed_env']); + assert.deepEqual(await open(pair, payload.sealed_env), { env: envs }); + await assert.rejects(open(pair, payload.sealed_env, epoch === 0 ? 1 : 0)); + const modified = JSON.parse(new TextDecoder().decode(payload.sealed_env)); + const ct = Buffer.from(modified.ciphertext, 'base64'); + ct[0] ^= 1; + modified.ciphertext = ct.toString('base64'); + await assert.rejects(open(pair, new TextEncoder().encode(JSON.stringify(modified)))); + } +}); + +test('falls back only on explicit missing RPC, including a forwarded KMS error', async () => { + for (const error of [ + new RpcError(missing, 404), new RpcError(missing, 400), + new RpcError(`Request failed with status=404 Not Found, error=${JSON.stringify({ error: missing })}`, 400), + ]) { + let legacyCalls = 0; + const payload = await encryptEnvPayload({ + getSealedEnvPubKey: async () => { throw error; }, + getAppEnvEncryptPubKey: async (request) => { + legacyCalls++; + assert.deepEqual(request, { app_id: appId }); + return { public_key: new Uint8Array(32) }; + }, + }, envs, appId, async (entries, key) => { + assert.deepEqual(entries, envs); + assert.equal(key.length, 32); + return new Uint8Array([1, 2, 3]); + }); + assert.equal(legacyCalls, 1); + assert.deepEqual(Object.keys(payload), ['encrypted_env']); + } +}); + +test('network, HTTP, authentication and unrelated service errors never downgrade', async () => { + for (const error of [ + new Error('network failed'), new Error(missing), new RpcError('Not Found', 404), + new RpcError(missing, 401), new RpcError(missing, 403), new RpcError(missing, 500), + new RpcError('Service not found: DifferentMethod', 404), + new RpcError('failed to validate the server attestation', 400), + new RpcError(`Request failed with status=500 Internal Server Error, error={"error":"${missing}"}`, 400), + ]) { + await assert.rejects(encryptEnvPayload({ + getSealedEnvPubKey: async () => { throw error; }, + getAppEnvEncryptPubKey: noLegacy, + }, envs, appId, noLegacy), (actual) => actual === error); + } +}); + +test('invalid key or epoch never downgrades after a successful key RPC', async () => { + const { response } = await recipient(); + for (const invalid of [ + { public_key: new Uint8Array(31) }, { public_key: new Uint8Array(33) }, + { public_key: new Uint8Array(32) }, + { root_epoch: -1 }, { root_epoch: 1.5 }, { root_epoch: 0x100000000 }, + ]) { + let legacyCalls = 0; + await assert.rejects(encryptEnvPayload({ + getSealedEnvPubKey: async () => ({ ...response, ...invalid }), + getAppEnvEncryptPubKey: async () => { legacyCalls++; return response; }, + }, envs, appId, noLegacy)); + assert.equal(legacyCalls, 0); + } +}); + +// Execute the shipped bundle with browser Web Crypto and protobuf transport. +// Vue lifecycle hooks are inert so tests do not start background polling. +function consoleApp(response, unsupported = false) { + const calls = []; + const alerts = []; + let manager; + const html = readFileSync(new URL('../dist/index.html', `file://${__filename}`), 'utf8'); + const scripts = [...html.matchAll(/