diff --git a/dstack/Cargo.lock b/dstack/Cargo.lock index a0167f769..6671ed8e6 100644 --- a/dstack/Cargo.lock +++ b/dstack/Cargo.lock @@ -376,6 +376,12 @@ dependencies = [ "tracing", ] +[[package]] +name = "az" +version = "1.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "be5eb007b7cacc6c660343e96f650fedf4b5a77512399eb952ca6642cf8d13f7" + [[package]] name = "base16ct" version = "0.2.0" @@ -628,7 +634,7 @@ checksum = "3f179cfbddb6e77a5472703d4b30436bff32929c0aa8a9008ecf23d1d3cdd0da" dependencies = [ "serde", "serde_repr", - "serde_with", + "serde_with 3.21.0", ] [[package]] @@ -729,6 +735,20 @@ name = "bytemuck" version = "1.25.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.11.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f65693059b6b9c588b9f62fed1cedbf0a8b805631457ea162d68f0de186f3de5" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.117", +] [[package]] name = "byteorder" @@ -849,6 +869,52 @@ version = "0.2.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "613afe47fcd5fac7ccf1db93babcb082c5994d996f20b8b159f2ad1658eb5724" +[[package]] +name = "cggmp21" +version = "0.6.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3309aeafc6311b8823faa1ee73e2e776b7fd5a0b3531ee72ec7f1acf40136ffc" +dependencies = [ + "cggmp21-keygen", + "digest 0.10.7", + "futures", + "generic-ec", + "generic-ec-zkp", + "hd-wallet", + "hex", + "key-share", + "paillier-zk", + "rand_core 0.6.4", + "rand_hash", + "round-based", + "serde", + "serde_with 2.3.3", + "sha2 0.10.9", + "thiserror 1.0.69", + "udigest", +] + +[[package]] +name = "cggmp21-keygen" +version = "0.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "aaa8c850290c494f951abe0350e56c31e4f5664863490197490ff48cb825447d" +dependencies = [ + "digest 0.10.7", + "displaydoc", + "generic-ec", + "generic-ec-zkp", + "hex", + "key-share", + "rand_core 0.6.4", + "round-based", + "serde", + "serde_with 2.3.3", + "sha2 0.10.9", + "thiserror 1.0.69", + "udigest", +] + [[package]] name = "chacha20" version = "0.10.0" @@ -2023,7 +2089,9 @@ name = "dstack-kms" version = "0.6.0" dependencies = [ "anyhow", + "async-trait", "cc-eventlog", + "cggmp21", "chrono", "clap", "dstack-api-auth", @@ -2035,11 +2103,14 @@ dependencies = [ "dstack-types", "dstack-verifier", "fs-err", + "futures", "hex", "hex_fmt", "http-client", "k256", + "key-share", "load_config", + "p256", "parity-scale-codec", "ra-rpc", "ra-tls", @@ -2047,10 +2118,12 @@ dependencies = [ "reqwest", "ring", "rocket", + "round-based", "safe-write", "serde", "serde-duration", "serde-human-bytes", + "serde_jcs", "serde_json", "sha2 0.10.9", "sha3", @@ -2445,6 +2518,18 @@ dependencies = [ "zeroize", ] +[[package]] +name = "educe" +version = "0.4.23" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0f0042ff8246a363dbe77d2ceedb073339e85a804b9a47636c6e016a9a32c05f" +dependencies = [ + "enum-ordinalize", + "proc-macro2", + "quote", + "syn 1.0.109", +] + [[package]] name = "either" version = "1.16.0" @@ -2496,6 +2581,19 @@ dependencies = [ "cfg-if", ] +[[package]] +name = "enum-ordinalize" +version = "3.1.15" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1bf1fa3f06bbff1ea5b1a9c7b14aa992a39657db60a2759457328d7e058f49ee" +dependencies = [ + "num-bigint", + "num-traits", + "proc-macro2", + "quote", + "syn 2.0.117", +] + [[package]] name = "enum_dispatch" version = "0.3.13" @@ -2602,6 +2700,19 @@ dependencies = [ "regex-syntax", ] +[[package]] +name = "fast-paillier" +version = "0.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d1108d991b54d8e3aa3eb155c07863306cbceafb713ab1ebcef085e19f3cb84c" +dependencies = [ + "bytemuck", + "rand_core 0.6.4", + "rug", + "serde", + "thiserror 1.0.69", +] + [[package]] name = "fastrand" version = "2.4.1" @@ -2904,11 +3015,77 @@ version = "0.14.7" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "85649ca51fd72272d7821adaf274ad91c288277713d9c18820d8499a7ff69e9a" dependencies = [ + "serde", "typenum", "version_check", "zeroize", ] +[[package]] +name = "generic-ec" +version = "0.4.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8de1099ac0b4d87261d67ff5d4ed400af617a1da40b58908d759b9cf5fd8ed27" +dependencies = [ + "curve25519-dalek", + "digest 0.10.7", + "generic-ec-core", + "generic-ec-curves", + "hex", + "phantom-type 0.4.2", + "rand_core 0.6.4", + "rand_hash", + "serde", + "serde_with 2.3.3", + "subtle", + "udigest", + "zeroize", +] + +[[package]] +name = "generic-ec-core" +version = "0.2.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dcba5fdf70cc3ce5805c487f8523b4ceeb32e8ec5237c71ffd93c1ca47a97fee" +dependencies = [ + "generic-array", + "rand_core 0.6.4", + "serde", + "subtle", + "zeroize", +] + +[[package]] +name = "generic-ec-curves" +version = "0.2.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4c7c6d23001a5eb60eec2b785a63d2ca965fdfbaf3314b3b46df047398369e28" +dependencies = [ + "elliptic-curve", + "generic-ec-core", + "k256", + "p256", + "rand_core 0.6.4", + "sha2 0.10.9", + "subtle", + "zeroize", +] + +[[package]] +name = "generic-ec-zkp" +version = "0.4.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bd3945c585fdddba3f86bda4e4cfba22d5e255001b3e145c9db305ad096c6d88" +dependencies = [ + "digest 0.10.7", + "generic-array", + "generic-ec", + "rand_core 0.6.4", + "serde", + "subtle", + "udigest", +] + [[package]] name = "getrandom" version = "0.1.16" @@ -3021,6 +3198,16 @@ dependencies = [ "regex-syntax", ] +[[package]] +name = "gmp-mpfr-sys" +version = "1.7.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7db155b537cb791b133341f99f68371d86ee7fa4c79aacfbc376d72d23c70531" +dependencies = [ + "libc", + "windows-sys 0.61.2", +] + [[package]] name = "gpt" version = "4.1.0" @@ -3165,6 +3352,19 @@ dependencies = [ "hashbrown 0.15.5", ] +[[package]] +name = "hd-wallet" +version = "0.6.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6522551bb35937363845f39a6d4c49e60bdb35a8f7154ebdd078cab50be97992" +dependencies = [ + "generic-array", + "generic-ec", + "hmac 0.12.1", + "sha2 0.10.9", + "subtle", +] + [[package]] name = "heck" version = "0.3.3" @@ -3191,6 +3391,9 @@ name = "hex" version = "0.4.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7f24254aa9a54b5c858eaee2f5bccdb46aaf0e486a595ed5fd8f86ba55232a70" +dependencies = [ + "serde", +] [[package]] name = "hex-literal" @@ -4005,6 +4208,22 @@ dependencies = [ "tokio", ] +[[package]] +name = "key-share" +version = "0.6.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "206b4474f861dedc6fc38e06f7c52c52f1e01180d5284aa62b58844a044fad7d" +dependencies = [ + "displaydoc", + "generic-ec", + "generic-ec-zkp", + "hex", + "rand_core 0.6.4", + "serde", + "serde_with 2.3.3", + "thiserror 1.0.69", +] + [[package]] name = "konst" version = "0.2.20" @@ -4962,6 +5181,24 @@ dependencies = [ "winapi", ] +[[package]] +name = "paillier-zk" +version = "0.4.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d9963224009a2fd339cffd8f6c5c35fc9a91732b89acd4b0ed34c30afe70a193" +dependencies = [ + "digest 0.10.7", + "fast-paillier", + "generic-ec", + "rand_core 0.6.4", + "rand_hash", + "rug", + "serde", + "serde_with 3.21.0", + "thiserror 1.0.69", + "udigest", +] + [[package]] name = "parcelona" version = "0.4.3" @@ -5161,6 +5398,24 @@ dependencies = [ "indexmap 2.14.0", ] +[[package]] +name = "phantom-type" +version = "0.3.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f710afd11c9711b04f97ab61bb9747d5a04562fdf0f9f44abc3de92490084982" +dependencies = [ + "educe", +] + +[[package]] +name = "phantom-type" +version = "0.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e68f5dc797c2a743e024e1c53215474598faf0408826a90249569ad7f47adeaa" +dependencies = [ + "educe", +] + [[package]] name = "phf" version = "0.12.1" @@ -5684,6 +5939,7 @@ name = "ra-tls" version = "0.6.0" dependencies = [ "anyhow", + "base64 0.22.1", "bon", "cc-eventlog", "dcap-qvl", @@ -5833,6 +6089,17 @@ version = "0.10.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "63b8176103e19a2643978565ca18b50549f6101881c443590420e4dc998a3c69" +[[package]] +name = "rand_hash" +version = "0.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "16bc1dd921383c6564eb0b8252f5b3f6622b84d40c6e35f5e6790e1fd7abb7a9" +dependencies = [ + "digest 0.10.7", + "rand_core 0.6.4", + "udigest", +] + [[package]] name = "rand_hc" version = "0.2.0" @@ -6249,6 +6516,30 @@ dependencies = [ "uncased", ] +[[package]] +name = "round-based" +version = "0.4.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "da76edf50de0a9d6911fc79261bb04cc9f3f3a375e0201799f5edf58499af341" +dependencies = [ + "futures-util", + "phantom-type 0.3.1", + "round-based-derive", + "thiserror 2.0.18", + "tracing", +] + +[[package]] +name = "round-based-derive" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4afa4d5b318bcafae8a7ebc57c1cb7d4b2db7358293e34d71bfd605fd327cc13" +dependencies = [ + "proc-macro2", + "quote", + "syn 1.0.109", +] + [[package]] name = "rsa" version = "0.9.10" @@ -6270,6 +6561,19 @@ dependencies = [ "zeroize", ] +[[package]] +name = "rug" +version = "1.30.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "07a8857882aec59d27254b02481c709327c13de6fad1da60bfc4f9783eaaa61e" +dependencies = [ + "az", + "gmp-mpfr-sys", + "libc", + "libm", + "serde", +] + [[package]] name = "rust-argon2" version = "0.8.3" @@ -6929,6 +7233,21 @@ dependencies = [ "serde", ] +[[package]] +name = "serde_with" +version = "2.3.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "07ff71d2c147a7b57362cead5e22f772cd52f6ab31cfcd9edcd7f6aeb2a0afbe" +dependencies = [ + "base64 0.13.1", + "chrono", + "hex", + "serde", + "serde_json", + "serde_with_macros 2.3.3", + "time", +] + [[package]] name = "serde_with" version = "3.21.0" @@ -6945,9 +7264,34 @@ dependencies = [ "schemars 1.2.1", "serde_core", "serde_json", + "serde_with_macros 3.21.0", "time", ] +[[package]] +name = "serde_with_macros" +version = "2.3.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "881b6f881b17d13214e5d494c939ebab463d01264ce1811e9d4ac3a882e7695f" +dependencies = [ + "darling 0.20.11", + "proc-macro2", + "quote", + "syn 2.0.117", +] + +[[package]] +name = "serde_with_macros" +version = "3.21.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "84d57bc0c8b9a17920c178daa6bb924850d54a9c97ab45194bb8c17ad66bb660" +dependencies = [ + "darling 0.23.0", + "proc-macro2", + "quote", + "syn 2.0.117", +] + [[package]] name = "serde_yaml" version = "0.9.34+deprecated" @@ -8067,6 +8411,27 @@ dependencies = [ "serde", ] +[[package]] +name = "udigest" +version = "0.2.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "19ff079a60bd5dc98b364ce7b5a633a8937bf558f5d19c9a390f5ae1973cf07e" +dependencies = [ + "digest 0.10.7", + "udigest-derive", +] + +[[package]] +name = "udigest-derive" +version = "0.3.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "25fd5248861b973cd5d1da5604b0ce22a35fa77f015d9f7ed9ab57078205bb86" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.117", +] + [[package]] name = "uncased" version = "0.9.10" diff --git a/dstack/Cargo.toml b/dstack/Cargo.toml index 848809b03..c32ce2564 100644 --- a/dstack/Cargo.toml +++ b/dstack/Cargo.toml @@ -123,6 +123,7 @@ wavekv = "1.0.0" # Core dependencies anyhow = { version = "1.0.97", default-features = false } +async-trait = "0.1.89" binrw = { version = "0.15.1", default-features = false, features = ["std"] } arc-swap = "1" errify = { version = "0.3.0", features = ["anyhow"] } @@ -185,6 +186,8 @@ object = "0.36.4" fatfs = "0.3.6" fscommon = "0.1.1" ciborium = "0.2" +cggmp21 = { version = "0.6.3", features = ["curve-secp256k1", "curve-secp256r1", "state-machine"] } +round-based = { version = "0.4.1", features = ["state-machine"] } # Networking/HTTP bollard = "0.18.1" diff --git a/dstack/dstack-util/src/system_setup.rs b/dstack/dstack-util/src/system_setup.rs index 6f258d34b..9ec9089c5 100644 --- a/dstack/dstack-util/src/system_setup.rs +++ b/dstack/dstack-util/src/system_setup.rs @@ -2090,7 +2090,11 @@ impl<'a> Stage0<'a> { gateway_app_id: response.gateway_app_id, key_provider: KeyProvider::Kms { url: kms_url, - pubkey: root_pubkey, + pubkey: if response.key_provider_id.is_empty() { + root_pubkey + } else { + response.key_provider_id + }, tmp_ca_key: tmp_ca.temp_ca_key, tmp_ca_cert: tmp_ca.temp_ca_cert, }, diff --git a/dstack/kms/Cargo.toml b/dstack/kms/Cargo.toml index 08a04bf4d..9a680c781 100644 --- a/dstack/kms/Cargo.toml +++ b/dstack/kms/Cargo.toml @@ -11,9 +11,11 @@ license.workspace = true [dependencies] anyhow.workspace = true +async-trait.workspace = true chrono.workspace = true clap.workspace = true fs-err.workspace = true +futures.workspace = true dstack-build-info.workspace = true hex.workspace = true hex_fmt.workspace = true @@ -33,6 +35,7 @@ reqwest = { workspace = true, features = ["json"] } sha2.workspace = true sha3.workspace = true k256.workspace = true +p256 = { workspace = true, features = ["hash2curve"] } rand.workspace = true dstack-guest-agent-rpc.workspace = true http-client = { workspace = true, features = ["prpc"] } @@ -41,6 +44,7 @@ x509-parser = { workspace = true, features = ["verify"] } ring.workspace = true safe-write.workspace = true serde_json.workspace = true +serde_jcs.workspace = true dstack-types.workspace = true tokio = { workspace = true, features = ["process"] } tempfile.workspace = true @@ -49,9 +53,14 @@ dstack-verifier = { workspace = true, default-features = false } dstack-mr.workspace = true dstack-attest.workspace = true dstack-api-auth.workspace = true +cggmp21.workspace = true +round-based.workspace = true [dev-dependencies] cc-eventlog.workspace = true +cggmp21 = { workspace = true, features = ["spof"] } +round-based = { workspace = true, features = ["sim"] } +key-share = { version = "0.6.1", features = ["spof"] } [features] default = [] diff --git a/dstack/kms/README.md b/dstack/kms/README.md index 4ea96fbda..74f3e01bb 100644 --- a/dstack/kms/README.md +++ b/dstack/kms/README.md @@ -232,3 +232,8 @@ The `SignCert` RPC is used by the dstack app to sign a TLS certificate. In this - Verify the CSR signature - Query the smart contract to check if the app is authorized - If authorized, sign the CSR with the CA root key and return the certificate chain to the app + +## Multi-instance MPC mode + +See [docs/mpc.md](docs/mpc.md) for the identity chain, genesis, membership +lifecycle, crash recovery, deployment procedure, and migration constraints. diff --git a/dstack/kms/docs/mpc.md b/dstack/kms/docs/mpc.md new file mode 100644 index 000000000..198001d92 --- /dev/null +++ b/dstack/kms/docs/mpc.md @@ -0,0 +1,183 @@ +# Multi-instance MPC KMS + +## Security identity + +MPC mode never loads or reconstructs a complete root private key. Its stable +`key_provider_id` is the SHA-256 domain-separated commitment to: + +- the threshold P-256 root CA public key; +- the threshold K-256 signing public key; +- the threshold P-256 derivation public key; and +- the MPC protocol version and cluster ID. + +Membership is deliberately excluded from this identity. A membership epoch is +a canonical manifest signed by the threshold K-256 key. Each member entry binds +the node ID, endpoint, quote-bound RA-TLS SubjectPublicKeyInfo, and a commitment +to that node's three public shares. + +The verification chain is: + +1. an application pins `key_provider_id` in its measured compose configuration; +2. `GetMeta` publishes the cluster identity, signed active manifest, and local + node evidence; +3. the manifest signature verifies under the pinned K-256 group key; +4. the manifest binds the serving node's RA-TLS key and public-share + commitment; and +5. the TEE quote binds the live TLS key to the measured KMS instance. + +The threshold P-256 key signs the root certificate and all member RPC +certificates. A member certificate is issued only after every signer verifies +that its embedded attestation is valid and bound to the authorized TLS key. + +## Genesis + +Genesis is an all-member maintenance ceremony. Before starting it, provision +only transport material (`rpc.crt`, `rpc.key`, the temporary CA files), an +operator-reviewed genesis plan, and a CA used by the provisional transport. +Do not provision a root private key. + +Example MPC configuration: + +```toml +[core.mpc] +enabled = true +protocol_version = 1 +cluster_id = "production-kms" +node_id = "kms-1" +identity_file = "/var/lib/dstack-kms/mpc-identity.json" +genesis_plan_file = "/var/lib/dstack-kms/genesis-plan.json" +genesis_tls_ca_cert = "/var/lib/dstack-kms/genesis-transport-ca.crt" +join_authorization_file = "/var/lib/dstack-kms/join-authorization.json" +client_cert_file = "/var/lib/dstack-kms/mpc-client.crt" +client_key_file = "/var/lib/dstack-kms/mpc-client.key" +manifest_file = "/var/lib/dstack-kms/epoch-manifest.json" +checkpoint_file = "/var/lib/dstack-kms/epoch-checkpoint.json" +p256_share_file = "/var/lib/dstack-kms/p256.share" +k256_share_file = "/var/lib/dstack-kms/k256.share" +derivation_share_file = "/var/lib/dstack-kms/derivation.share" +max_sessions = 128 +session_ttl = "5m" +``` + +The canonical JSON genesis plan contains sorted members: + +```json +{ + "protocol_version": 1, + "cluster_id": "production-kms", + "threshold": 3, + "coordinator": "kms-1", + "members": [ + { + "node_id": "kms-1", + "endpoint": "https://kms-1.example/prpc", + "attestation_pubkey": "" + } + ] +} +``` + +All members must receive the identical plan. The coordinator waits for every +attested peer and then runs distributed auxiliary generation and DKG for all +three keys. It threshold-signs the root CA, reissues each attested RPC +certificate under that root, signs epoch 1, and distributes the final bundle. +A durable genesis journal makes every node finish a crash-interrupted commit. +Processes shut down after commit and must be restarted by their supervisor. + +`mpc-client.crt` must be accepted by `[rpc.tls.mutual].ca_certs`, contain a +valid RA-TLS attestation, and use the same key as `rpc.crt`. The server +certificate is reissued by the threshold root; the client certificate remains +under the mutual-TLS CA so Rocket can authenticate peer connections before the +RPC quote and manifest checks run. + +Record and independently compare `key_provider_id` from multiple members before +placing it in application compose files. + +## Normal operation + +Signing uses an exact live threshold selected in canonical manifest order. +Two offline nodes are tolerated by a 3-of-5 cluster. Every protocol envelope is +bound to the epoch, protocol, session, request hash, sender, recipient, sequence +number, and expiry. Raw arbitrary-digest signing is not exposed. + +Derivation is a verifiable threshold PRF. Each response carries a DLEQ proof +against the member's public derivation share; the coordinator combines only +verified partials. Disk, environment, application K-256, and application CA +keys use separate semantic domains. + +## Revocation and subset resharing + +For a target consisting only of current members: + +1. call `Admin.PrepareReshare` with a canonical next-epoch `ResharePlan`; +2. construct an `EpochManifest` using the returned commitment for every target + member; +3. call `Admin.SignEpochManifest`; +4. send the signed result to `Admin.ActivateEpoch` on every target and removed + member; and +5. let the supervisor restart each process. + +Exactly an old threshold acts as dealers. Their Lagrange-weighted random +polynomials preserve every group public key. Private evaluations are sent +point-to-point over RA-TLS and are verified against coefficient commitments. + +Activation validates the signed chain, local pending share topology, stable +group keys, and local share commitment. It writes an activation journal before +changing active files. Startup completes an interrupted transaction before it +loads a manifest or share. Removed nodes install the signed epoch and then fail +closed because they are no longer members. + +## Adding members + +Adding a member is an explicitly authorized maintenance operation: + +1. Create a sorted `ResharePlan` containing the target membership and exactly + the old threshold in `dealers`. Existing node endpoints and attestation keys + must remain unchanged. New members use freshly attested TLS keys. +2. Call `Admin.AuthorizeReshare`. The result is signed by the active threshold + K-256 key and binds the predecessor, target epoch, threshold, dealers, + endpoints, and quote-bound keys. +3. Provision the signed authorization, active signed manifest, cluster + identity, root certificate, and provisional attested transport key on every + target node. Place the authorization at `join_authorization_file`. +4. Restart every target. They enter maintenance join mode. The first authorized + dealer coordinates fresh all-target auxiliary generation, point-to-point + verifiable resharing, target manifest signing, and threshold issuance of + each target RPC certificate. +5. Each target atomically activates the signed epoch, removes the authorization, + and shuts down for a normal-mode restart. + +The coordinator never receives all private recipient evaluations. A joining +node receives only the old public topology and its own evaluation from each +old dealer. The old root private key is never reconstructed. + +## Migration and compatibility + +The public KMS RPCs and guest key-delivery formats remain unchanged. MPC mode +adds metadata fields and admin methods without removing legacy fields. + +Moving an existing single-root deployment to a newly generated MPC cluster is +a provider rotation: the `key_provider_id` and deterministically derived +application keys change. Existing encrypted disks therefore require an +application-level migration ceremony. Do not silently switch an existing +compose pin. Run old and new providers in parallel, migrate protected data, +then update and re-measure the compose pin. + +MPC members reject legacy root-key export and handover. There is intentionally +no API that reconstructs threshold roots. + +## Operational requirements + +- Keep share, checkpoint, journal, identity, authorization, and manifest files + on durable storage; secret shares and journals must be mode `0600`. +- Use authenticated admin listeners. Never enable `insecure_no_auth` in a + network-reachable deployment. +- Require at least threshold 2; production deployments should use 3-of-5 or + stronger. +- Monitor quorum availability and restart completion. Do not activate a target + manifest until every retained member reports its expected commitment. +- Back up shares independently. A backup set smaller than the threshold cannot + recover the root; losing threshold shares permanently loses the provider. +- Treat a same-epoch fork, skipped epoch, checkpoint permission failure, quote + mismatch, share commitment mismatch, or group-key change as a security + incident. The implementation fails closed for all of them. diff --git a/dstack/kms/kms.toml b/dstack/kms/kms.toml index 08e954662..36cb79c0a 100644 --- a/dstack/kms/kms.toml +++ b/dstack/kms/kms.toml @@ -97,6 +97,32 @@ insecure_no_auth = false # or another access-control layer. enabled = true +# MPC cluster identity preview. When enabled, applications pin a stable hash of +# the group public keys rather than one KMS node. The local keys must match the +# configured signing group keys until the threshold key backend is enabled. +[core.mpc] +enabled = false +protocol_version = 1 +cluster_id = "" +identity_file = "" +genesis_plan_file = "" +genesis_tls_ca_cert = "" +join_authorization_file = "" +client_cert_file = "" +client_key_file = "" +p256_group_pubkey = "" +k256_group_pubkey = "" +derivation_group_pubkey = "" +node_id = "" +manifest_file = "" +checkpoint_file = "" +allow_unsigned_manifest = false +max_sessions = 128 +session_ttl = "5m" +p256_share_file = "" +k256_share_file = "" +derivation_share_file = "" + [core.auth_api] type = "webhook" diff --git a/dstack/kms/rpc/proto/kms_rpc.proto b/dstack/kms/rpc/proto/kms_rpc.proto index de44e28a0..40171d97f 100644 --- a/dstack/kms/rpc/proto/kms_rpc.proto +++ b/dstack/kms/rpc/proto/kms_rpc.proto @@ -46,6 +46,9 @@ message AppKeyResponse { string gateway_app_id = 7; // OS Image hash bytes os_image_hash = 8; + // Stable identity of the key provider. Empty means legacy behavior, where + // clients use the root CA public key as the provider identity. + bytes key_provider_id = 9; } message GetMetaResponse { @@ -58,6 +61,13 @@ message GetMetaResponse { optional string kms_contract_address = 7; optional uint64 chain_id = 8; optional string app_auth_implementation = 9; + bytes key_provider_id = 10; + optional string mpc_cluster_identity = 11; + // Canonical threshold-signed active epoch manifest JSON. + optional string mpc_epoch_manifest = 12; + // Local node evidence. Its attestation key is simultaneously bound by the + // RA-TLS quote and by the threshold-signed epoch manifest. + optional string mpc_node_evidence = 13; } message GetKmsKeyRequest { @@ -107,6 +117,63 @@ service KMS { rpc SignCert(SignCertRequest) returns (SignCertResponse); } +message MpcPushRequest { + bytes envelope_json = 1; +} + +message MpcPollRequest { + bytes session_id = 1; +} + +message MpcPollResponse { + repeated bytes envelope_json = 1; +} + +message MpcStartRequest { + bytes operation_json = 1; +} + +message MpcStartResponse { + bytes result = 1; +} + +// Internal interactive-protocol transport. Every method requires an +// attested RA-TLS client certificate belonging to the active epoch manifest. +service MpcTransport { + rpc Ping(google.protobuf.Empty) returns (google.protobuf.Empty); + rpc Start(MpcStartRequest) returns (MpcStartResponse); + rpc Push(MpcPushRequest) returns (google.protobuf.Empty); +} + +message MpcGenesisStartRequest { + bytes operation_json = 1; +} + +message MpcGenesisStartResponse { + bytes result_json = 1; +} + +message MpcGenesisFinalizeRequest { + bytes bundle_json = 1; +} + +// One-time DKG service used before root material exists. Peers are authorized +// by quote-bound certificate keys in the operator-provisioned genesis plan. +service MpcGenesis { + rpc Ping(google.protobuf.Empty) returns (google.protobuf.Empty); + rpc Start(MpcGenesisStartRequest) returns (MpcGenesisStartResponse); + rpc Push(MpcPushRequest) returns (google.protobuf.Empty); + rpc Finalize(MpcGenesisFinalizeRequest) returns (google.protobuf.Empty); +} + +// Maintenance-mode protocol for threshold-authorized membership transitions. +service MpcJoin { + rpc Ping(google.protobuf.Empty) returns (google.protobuf.Empty); + rpc Start(MpcGenesisStartRequest) returns (MpcGenesisStartResponse); + rpc Push(MpcPushRequest) returns (google.protobuf.Empty); + rpc Finalize(MpcGenesisFinalizeRequest) returns (google.protobuf.Empty); +} + // The KMS admin RPC service. Served on a separate admin listener // (`[core.admin]`) behind the shared HTTP authenticator, so the credential // travels in the `Authorization`/`X-Admin-Token` header rather than the @@ -114,6 +181,12 @@ service KMS { service Admin { // Clear the image cache rpc ClearImageCache(ClearImageCacheRequest) returns (google.protobuf.Empty); + // Ask the active threshold committee to authorize the next epoch manifest. + rpc SignEpochManifest(SignEpochManifestRequest) returns (SignEpochManifestResponse); + // Prepare threshold shares for a revoke/subset epoch transition. + rpc PrepareReshare(PrepareReshareRequest) returns (PrepareReshareResponse); + rpc ActivateEpoch(ActivateEpochRequest) returns (ActivateEpochResponse); + rpc AuthorizeReshare(AuthorizeReshareRequest) returns (AuthorizeReshareResponse); } message ClearImageCacheRequest { @@ -121,6 +194,28 @@ message ClearImageCacheRequest { string config_hash = 2; } +message SignEpochManifestRequest { + bytes manifest_json = 1; +} + +message SignEpochManifestResponse { + bytes signed_manifest_json = 1; +} + +message PrepareReshareRequest { + bytes plan_json = 1; +} + +message PrepareReshareResponse { + // Canonical JSON map from node ID to its new 32-byte share commitment. + bytes commitments_json = 1; +} + +message ActivateEpochRequest { bytes signed_manifest_json = 1; } +message ActivateEpochResponse { bool retained = 1; } +message AuthorizeReshareRequest { bytes plan_json = 1; } +message AuthorizeReshareResponse { bytes signed_plan_json = 1; } + message BootstrapRequest { string domain = 1; } diff --git a/dstack/kms/src/admin_service.rs b/dstack/kms/src/admin_service.rs index 1b8bec309..898330626 100644 --- a/dstack/kms/src/admin_service.rs +++ b/dstack/kms/src/admin_service.rs @@ -7,10 +7,12 @@ //! htpasswd via `Authorization`/`X-Admin-Token`, so the handlers do no token //! checks of their own. -use anyhow::Result; +use anyhow::{Context, Result}; use dstack_kms_rpc::{ admin_server::{AdminRpc, AdminServer}, - ClearImageCacheRequest, + ActivateEpochRequest, ActivateEpochResponse, AuthorizeReshareRequest, AuthorizeReshareResponse, + ClearImageCacheRequest, PrepareReshareRequest, PrepareReshareResponse, + SignEpochManifestRequest, SignEpochManifestResponse, }; use ra_rpc::{CallContext, RpcCall}; @@ -25,6 +27,57 @@ impl AdminRpc for AdminRpcHandler { self.state .clear_image_cache(&request.image_hash, &request.config_hash) } + + async fn sign_epoch_manifest( + self, + request: SignEpochManifestRequest, + ) -> Result { + let manifest = serde_json::from_slice(&request.manifest_json) + .context("invalid epoch manifest JSON")?; + let signed = self + .state + .key_backend() + .sign_epoch_manifest(manifest) + .await?; + Ok(SignEpochManifestResponse { + signed_manifest_json: serde_jcs::to_vec(&signed) + .context("failed to encode signed epoch manifest")?, + }) + } + + async fn prepare_reshare( + self, + request: PrepareReshareRequest, + ) -> Result { + let plan = + serde_json::from_slice(&request.plan_json).context("invalid reshare plan JSON")?; + let commitments = self.state.key_backend().prepare_reshare(plan).await?; + Ok(PrepareReshareResponse { + commitments_json: serde_jcs::to_vec(&commitments) + .context("failed to encode reshare commitments")?, + }) + } + + async fn activate_epoch(self, request: ActivateEpochRequest) -> Result { + let signed = serde_json::from_slice(&request.signed_manifest_json) + .context("invalid signed epoch manifest JSON")?; + let retained = self.state.key_backend().activate_epoch(signed).await?; + self.state.request_restart(); + Ok(ActivateEpochResponse { retained }) + } + + async fn authorize_reshare( + self, + request: AuthorizeReshareRequest, + ) -> Result { + let plan = + serde_json::from_slice(&request.plan_json).context("invalid reshare plan JSON")?; + let signed = self.state.key_backend().authorize_reshare(plan).await?; + Ok(AuthorizeReshareResponse { + signed_plan_json: serde_jcs::to_vec(&signed) + .context("failed to encode signed reshare authorization")?, + }) + } } impl RpcCall for AdminRpcHandler { diff --git a/dstack/kms/src/cggmp_engine.rs b/dstack/kms/src/cggmp_engine.rs new file mode 100644 index 000000000..10984873f --- /dev/null +++ b/dstack/kms/src/cggmp_engine.rs @@ -0,0 +1,317 @@ +// SPDX-FileCopyrightText: © 2024-2026 Phala Network +// +// SPDX-License-Identifier: Apache-2.0 + +//! CGGMP21 key-share persistence and protocol context binding. + +use std::{fs::Permissions, os::unix::fs::PermissionsExt, path::Path}; + +use anyhow::{ensure, Context, Result}; +use cggmp21::{ + generic_ec::{Curve, Scalar}, + supported_curves::{Secp256k1, Secp256r1}, + IncompleteKeyShare, KeyShare, +}; +use fs_err as fs; +use serde::{de::DeserializeOwned, Deserialize, Serialize}; +use sha2::{Digest, Sha256}; + +pub(crate) type K256IncompleteShare = IncompleteKeyShare; +pub(crate) type P256IncompleteShare = IncompleteKeyShare; +pub(crate) type K256KeyShare = KeyShare; +pub(crate) type P256KeyShare = KeyShare; + +const SHARE_FORMAT_VERSION: u16 = 1; +const EXECUTION_ID_DOMAIN: &[u8] = b"dstack-cggmp21-execution-id-v1"; +const SHARE_COMMITMENT_DOMAIN: &[u8] = b"dstack-mpc-share-commitment-v1"; + +#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub(crate) enum CggmpCurve { + P256, + K256, +} + +#[derive(Debug, Serialize, Deserialize)] +struct StoredShare { + version: u16, + cluster_id: String, + epoch: u64, + node_id: String, + curve: CggmpCurve, + payload: serde_json::Value, + checksum: String, +} + +#[derive(Serialize)] +struct ShareChecksum<'a> { + version: u16, + cluster_id: &'a str, + epoch: u64, + node_id: &'a str, + curve: CggmpCurve, + payload: &'a serde_json::Value, +} + +pub(crate) fn execution_id( + cluster_id: &str, + epoch: u64, + curve: CggmpCurve, + operation_nonce: &[u8; 32], +) -> [u8; 32] { + let mut hash = Sha256::new(); + for field in [ + EXECUTION_ID_DOMAIN, + cluster_id.as_bytes(), + &epoch.to_be_bytes(), + match curve { + CggmpCurve::P256 => b"p256" as &[u8], + CggmpCurve::K256 => b"k256" as &[u8], + }, + operation_nonce, + ] { + hash.update((field.len() as u32).to_be_bytes()); + hash.update(field); + } + hash.finalize().into() +} + +pub(crate) fn share_commitment( + p256_public_share: &[u8], + k256_public_share: &[u8], + derivation_public_share: &[u8], +) -> [u8; 32] { + let mut hash = Sha256::new(); + hash.update((SHARE_COMMITMENT_DOMAIN.len() as u32).to_be_bytes()); + hash.update(SHARE_COMMITMENT_DOMAIN); + for value in [ + p256_public_share, + k256_public_share, + derivation_public_share, + ] { + hash.update((value.len() as u32).to_be_bytes()); + hash.update(value); + } + hash.finalize().into() +} + +pub(crate) fn validate_share_topology( + share: &KeyShare, + expected_index: usize, + expected_members: usize, + expected_threshold: u16, +) -> Result<()> { + ensure!( + usize::from(share.core.i) == expected_index, + "CGGMP share index does not match manifest ordering" + ); + ensure!( + share.core.key_info.public_shares.len() == expected_members, + "CGGMP public-share count does not match manifest" + ); + let setup = share + .core + .key_info + .vss_setup + .as_ref() + .context("CGGMP share is not a threshold Shamir share")?; + ensure!( + setup.min_signers == expected_threshold, + "CGGMP threshold does not match manifest" + ); + ensure!( + setup.I.len() == expected_members, + "CGGMP VSS index count mismatch" + ); + for (index, evaluation_point) in setup.I.iter().enumerate() { + ensure!( + Scalar::from(evaluation_point.clone()) == Scalar::from(index as u64 + 1), + "CGGMP VSS indexes must use canonical points 1..=n" + ); + } + Ok(()) +} + +pub(crate) fn store_share( + path: &Path, + cluster_id: &str, + epoch: u64, + node_id: &str, + curve: CggmpCurve, + share: &T, +) -> Result<()> { + ensure!(!cluster_id.is_empty(), "cluster_id must not be empty"); + ensure!(epoch > 0, "epoch must be greater than zero"); + ensure!(!node_id.is_empty(), "node_id must not be empty"); + let payload = serde_json::to_value(share).context("failed to encode CGGMP share")?; + let checksum = checksum(cluster_id, epoch, node_id, curve, &payload)?; + let stored = StoredShare { + version: SHARE_FORMAT_VERSION, + cluster_id: cluster_id.into(), + epoch, + node_id: node_id.into(), + curve, + payload, + checksum: hex::encode(checksum), + }; + let encoded = serde_jcs::to_vec(&stored).context("failed to serialize stored CGGMP share")?; + safe_write::safe_write(path, &encoded).context("failed to persist CGGMP share")?; + fs::set_permissions(path, Permissions::from_mode(0o600)) + .context("failed to restrict CGGMP share permissions")?; + Ok(()) +} + +pub(crate) fn load_share( + path: &Path, + expected_cluster_id: &str, + expected_epoch: u64, + expected_node_id: &str, + expected_curve: CggmpCurve, +) -> Result { + let metadata = fs::metadata(path).context("failed to stat CGGMP share")?; + ensure!( + metadata.permissions().mode() & 0o077 == 0, + "CGGMP share file permissions are too broad" + ); + let stored: StoredShare = + serde_json::from_slice(&fs::read(path).context("failed to read CGGMP share")?) + .context("failed to parse CGGMP share")?; + ensure!( + stored.version == SHARE_FORMAT_VERSION, + "unsupported CGGMP share version" + ); + ensure!( + stored.cluster_id == expected_cluster_id, + "CGGMP share cluster mismatch" + ); + ensure!( + stored.epoch == expected_epoch, + "CGGMP share epoch mismatch or rollback" + ); + ensure!( + stored.node_id == expected_node_id, + "CGGMP share node mismatch" + ); + ensure!(stored.curve == expected_curve, "CGGMP share curve mismatch"); + let expected_checksum = checksum( + &stored.cluster_id, + stored.epoch, + &stored.node_id, + stored.curve, + &stored.payload, + )?; + ensure!( + stored.checksum == hex::encode(expected_checksum), + "CGGMP share checksum mismatch" + ); + serde_json::from_value(stored.payload).context("CGGMP share validation failed") +} + +fn checksum( + cluster_id: &str, + epoch: u64, + node_id: &str, + curve: CggmpCurve, + payload: &serde_json::Value, +) -> Result<[u8; 32]> { + let canonical = serde_jcs::to_vec(&ShareChecksum { + version: SHARE_FORMAT_VERSION, + cluster_id, + epoch, + node_id, + curve, + payload, + })?; + Ok(Sha256::digest(canonical).into()) +} + +#[cfg(test)] +mod tests { + use super::*; + use cggmp21::ExecutionId; + use rand::rngs::OsRng; + + #[test] + fn cggmp_threshold_dkg_produces_consistent_validated_shares() { + let n = 3; + let threshold = 2; + let eid = execution_id("test-cluster", 1, CggmpCurve::K256, &[9; 32]); + let outputs = round_based::sim::run(n, |i, party| async move { + cggmp21::keygen::(ExecutionId::new(&eid), i, n) + .set_threshold(threshold) + .start(&mut OsRng, party) + .await + }) + .unwrap() + .expect_ok() + .into_vec(); + let public_key = serde_json::to_vec(&outputs[0].shared_public_key()).unwrap(); + assert!(outputs.iter().all(|share| { + serde_json::to_vec(&share.shared_public_key()).unwrap() == public_key + })); + + let directory = tempfile::tempdir().unwrap(); + let path = directory.path().join("k256-share.json"); + store_share( + &path, + "test-cluster", + 1, + "kms-1", + CggmpCurve::K256, + &outputs[0], + ) + .unwrap(); + let loaded: K256IncompleteShare = + load_share(&path, "test-cluster", 1, "kms-1", CggmpCurve::K256).unwrap(); + assert_eq!( + serde_json::to_vec(&loaded.shared_public_key()).unwrap(), + public_key + ); + assert!(load_share::( + &path, + "test-cluster", + 2, + "kms-1", + CggmpCurve::K256 + ) + .is_err()); + } + + #[test] + fn cggmp_p256_threshold_dkg_roundtrips_validated_share() { + let n = 3; + let threshold = 2; + let eid = execution_id("test-cluster", 1, CggmpCurve::P256, &[8; 32]); + let outputs = round_based::sim::run(n, |i, party| async move { + cggmp21::keygen::(ExecutionId::new(&eid), i, n) + .set_threshold(threshold) + .start(&mut OsRng, party) + .await + }) + .unwrap() + .expect_ok() + .into_vec(); + let public_key = serde_json::to_vec(&outputs[0].shared_public_key()).unwrap(); + assert!(outputs.iter().all(|share| { + serde_json::to_vec(&share.shared_public_key()).unwrap() == public_key + })); + + let directory = tempfile::tempdir().unwrap(); + let path = directory.path().join("p256-share.json"); + store_share( + &path, + "test-cluster", + 1, + "kms-1", + CggmpCurve::P256, + &outputs[0], + ) + .unwrap(); + let loaded: P256IncompleteShare = + load_share(&path, "test-cluster", 1, "kms-1", CggmpCurve::P256).unwrap(); + assert_eq!( + serde_json::to_vec(&loaded.shared_public_key()).unwrap(), + public_key + ); + } +} diff --git a/dstack/kms/src/config.rs b/dstack/kms/src/config.rs index 4d95b0e17..5202ff5ef 100644 --- a/dstack/kms/src/config.rs +++ b/dstack/kms/src/config.rs @@ -74,12 +74,76 @@ pub(crate) struct KmsConfig { #[serde(default = "default_true")] pub enforce_self_authorization: bool, pub metrics: MetricsConfig, + /// Optional stable identity for an MPC-backed KMS. Cryptographic MPC + /// operations are provided by the key backend; this section defines the + /// identity that applications pin independently of cluster membership. + #[serde(default)] + pub mpc: MpcConfig, /// Admin API listener + authentication. The admin RPCs (e.g. /// `ClearImageCache`) are served here, behind the shared HTTP authenticator. #[serde(default)] pub admin: AdminConfig, } +#[derive(Debug, Clone, Default, Deserialize)] +pub(crate) struct MpcConfig { + #[serde(default)] + pub enabled: bool, + #[serde(default)] + pub protocol_version: u16, + #[serde(default)] + pub cluster_id: String, + #[serde(default)] + pub identity_file: PathBuf, + #[serde(default)] + pub genesis_plan_file: PathBuf, + #[serde(default)] + pub genesis_tls_ca_cert: PathBuf, + /// Threshold-signed transition authorization. When present, the process + /// enters maintenance join mode until the target epoch is activated. + #[serde(default)] + pub join_authorization_file: PathBuf, + /// RA-TLS client identity signed by the Rocket mutual-TLS CA. It should + /// use the same quote-bound key as the server RPC certificate. + #[serde(default)] + pub client_cert_file: PathBuf, + #[serde(default)] + pub client_key_file: PathBuf, + #[serde(default)] + pub p256_group_pubkey: String, + #[serde(default)] + pub k256_group_pubkey: String, + #[serde(default)] + pub derivation_group_pubkey: String, + #[serde(default)] + pub node_id: String, + #[serde(default)] + pub manifest_file: PathBuf, + #[serde(default)] + pub checkpoint_file: PathBuf, + /// Development-only compatibility escape hatch for unsigned manifests. + #[serde(default)] + pub allow_unsigned_manifest: bool, + #[serde(default = "default_mpc_max_sessions")] + pub max_sessions: usize, + #[serde(default = "default_mpc_session_ttl", with = "serde_duration")] + pub session_ttl: Duration, + #[serde(default)] + pub p256_share_file: PathBuf, + #[serde(default)] + pub k256_share_file: PathBuf, + #[serde(default)] + pub derivation_share_file: PathBuf, +} + +fn default_mpc_max_sessions() -> usize { + 128 +} + +fn default_mpc_session_ttl() -> Duration { + Duration::from_secs(300) +} + #[derive(Debug, Clone, Deserialize)] pub(crate) struct MetricsConfig { /// Whether to expose the unauthenticated Prometheus `/metrics` endpoint. @@ -114,7 +178,34 @@ fn default_true() -> bool { } impl KmsConfig { + /// Transport identity required to run attested MPC genesis. Unlike + /// `keys_exists`, this deliberately does not require threshold shares or + /// the final root certificate. + pub fn genesis_transport_exists(&self) -> bool { + self.tmp_ca_cert().exists() + && self.tmp_ca_key().exists() + && self.rpc_cert().exists() + && self.rpc_key().exists() + && self.mpc.client_cert_file.exists() + && self.mpc.client_key_file.exists() + } + pub fn keys_exists(&self) -> bool { + if self.mpc.enabled { + return self.tmp_ca_cert().exists() + && self.tmp_ca_key().exists() + && self.root_ca_cert().exists() + && self.rpc_cert().exists() + && self.rpc_key().exists() + && self.mpc.client_cert_file.exists() + && self.mpc.client_key_file.exists() + && self.mpc.p256_share_file.exists() + && self.mpc.k256_share_file.exists() + && self.mpc.derivation_share_file.exists() + && self.mpc.identity_file.exists() + && self.mpc.manifest_file.exists() + && self.mpc.checkpoint_file.exists(); + } self.tmp_ca_cert().exists() && self.tmp_ca_key().exists() && self.root_ca_cert().exists() diff --git a/dstack/kms/src/key_backend.rs b/dstack/kms/src/key_backend.rs new file mode 100644 index 000000000..27739125d --- /dev/null +++ b/dstack/kms/src/key_backend.rs @@ -0,0 +1,1569 @@ +// SPDX-FileCopyrightText: © 2024-2026 Phala Network +// +// SPDX-License-Identifier: Apache-2.0 + +//! Root-key operation boundary. +//! +//! The public KMS service must only depend on this interface. The local +//! implementation preserves the legacy behavior; the MPC implementation can +//! replace it without making root key material available to request handlers. + +use std::{ + collections::BTreeMap, + sync::{Arc, Mutex}, + time::Duration, +}; + +use anyhow::{bail, ensure, Context, Result}; +use async_trait::async_trait; +use cggmp21::key_share::{AnyKeyShare as _, Valid}; +use cggmp21::{ + generic_ec::Scalar, + supported_curves::{Secp256k1, Secp256r1}, + DataToSign, ExecutionId, +}; +use futures::future::join_all; +use k256::ecdsa::{RecoveryId, Signature as K256Signature, SigningKey, VerifyingKey}; +use ra_tls::{ + cert::{prepare_external_certificate, CaCert, CertRequest}, + kdf, +}; +use rand::{rngs::OsRng, RngCore}; + +use crate::{ + cggmp_engine::{ + execution_id, load_share, share_commitment, store_share, validate_share_topology, + CggmpCurve, K256KeyShare, P256KeyShare, + }, + crypto::{derive_k256_key, sign_message, sign_message_with_timestamp}, + mpc_driver::{ + drive_state_machine_blocking, BlockingTransport, DriverContext, EnvelopeTransport, + MpcHttpTransport, + }, + mpc_identity::{ClusterIdentity, EpochManifest, SignedEpochManifest}, + mpc_lifecycle::{ + activate_pending_epoch, record_epoch_authorization, EpochPaths, ResharePlan, + SignedResharePlan, + }, + mpc_operation::{ + DerivePayload, DerivePurpose, K256SignPayload, ManifestSignaturePayload, MpcOperation, + MpcOperationPayload, P256CertificatePayload, + }, + mpc_reshare::{self, PrivateContribution, PublicContribution}, + mpc_session::{MpcEnvelope, MpcProtocol, SessionRouter}, + threshold_prf::{self, PrfPartial}, +}; + +pub(crate) struct DerivedAppKeys { + pub disk_key: [u8; 32], + pub env_key: [u8; 32], + pub k256_key: Vec, + pub k256_signature: Vec, +} + +#[async_trait] +pub(crate) trait KeyBackend: Send + Sync { + async fn derive_app_keys(&self, app_id: &[u8], instance_id: &[u8]) -> Result; + async fn derive_env_key(&self, app_id: &[u8]) -> Result<[u8; 32]>; + async fn sign_k256(&self, prefix: &[u8], app_id: &[u8], message: &[u8]) -> Result>; + async fn sign_k256_timestamped( + &self, + prefix: &[u8], + app_id: &[u8], + timestamp: u64, + message: &[u8], + ) -> Result>; + fn k256_public_key(&self) -> Vec; + fn p256_public_key(&self) -> Vec; + fn derivation_public_key(&self) -> Vec; + fn root_ca_cert(&self) -> &str; + async fn derive_app_ca(&self, app_id: &[u8]) -> Result; + async fn sign_epoch_manifest(&self, manifest: EpochManifest) -> Result; + async fn prepare_reshare(&self, plan: ResharePlan) -> Result>>; + async fn activate_epoch(&self, signed: SignedEpochManifest) -> Result; + async fn authorize_reshare(&self, plan: ResharePlan) -> Result; + async fn run_mpc_operation( + &self, + operation: crate::mpc_operation::MpcOperation, + initiator: &str, + ) -> Result>; + /// Legacy KMS-to-KMS migration. MPC backends must reject this operation. + async fn export_root_keys(&self) -> Result<(String, Vec)>; +} + +pub(crate) struct LocalKeyBackend { + root_ca: CaCert, + k256_key: SigningKey, +} + +impl LocalKeyBackend { + pub(crate) fn from_pem_and_bytes( + root_ca_cert: String, + root_ca_key: String, + k256_key: &[u8], + ) -> Result { + Ok(Self { + root_ca: CaCert::new(root_ca_cert, root_ca_key).context("invalid root CA")?, + k256_key: SigningKey::from_slice(k256_key).context("invalid K-256 root key")?, + }) + } +} + +#[async_trait] +impl KeyBackend for LocalKeyBackend { + async fn derive_app_keys(&self, app_id: &[u8], instance_id: &[u8]) -> Result { + let disk_key = kdf::derive_dh_secret( + &self.root_ca.key, + &[app_id, instance_id, b"app-disk-crypt-key"], + )?; + let env_key = self.derive_env_key(app_id).await?; + let (k256_key, k256_signature) = derive_k256_key(&self.k256_key, app_id)?; + Ok(DerivedAppKeys { + disk_key, + env_key, + k256_key: k256_key.to_bytes().to_vec(), + k256_signature, + }) + } + + async fn derive_env_key(&self, app_id: &[u8]) -> Result<[u8; 32]> { + kdf::derive_dh_secret(&self.root_ca.key, &[app_id, b"env-encrypt-key"]) + } + + async fn sign_k256(&self, prefix: &[u8], app_id: &[u8], message: &[u8]) -> Result> { + sign_message(&self.k256_key, prefix, app_id, message) + } + + async fn sign_k256_timestamped( + &self, + prefix: &[u8], + app_id: &[u8], + timestamp: u64, + message: &[u8], + ) -> Result> { + sign_message_with_timestamp(&self.k256_key, prefix, app_id, timestamp, message) + } + + fn k256_public_key(&self) -> Vec { + self.k256_key.verifying_key().to_sec1_bytes().to_vec() + } + + fn p256_public_key(&self) -> Vec { + self.root_ca.key.public_key_raw().to_vec() + } + + fn derivation_public_key(&self) -> Vec { + self.p256_public_key() + } + + fn root_ca_cert(&self) -> &str { + &self.root_ca.pem_cert + } + + async fn derive_app_ca(&self, app_id: &[u8]) -> Result { + let app_key = kdf::derive_p256_key_pair(&self.root_ca.key, &[app_id, b"app-ca"])?; + let request = CertRequest::builder() + .key(&app_key) + .org_name("Dstack") + .subject("Dstack App CA") + .ca_level(0) + .app_id(app_id) + .special_usage("app:ca") + .build(); + let certificate = self + .root_ca + .sign(request) + .context("failed to sign app CA")?; + Ok(CaCert::from_parts(app_key, certificate)) + } + + async fn sign_epoch_manifest(&self, manifest: EpochManifest) -> Result { + use k256::ecdsa::signature::hazmat::PrehashSigner as _; + let hash = manifest.manifest_hash()?; + let signature: K256Signature = self + .k256_key + .sign_prehash(&hash) + .context("failed to sign epoch manifest")?; + Ok(SignedEpochManifest { + manifest, + signature: signature.to_bytes().to_vec(), + }) + } + + async fn prepare_reshare(&self, _plan: ResharePlan) -> Result>> { + bail!("resharing is unavailable on the local key backend") + } + + async fn activate_epoch(&self, _signed: SignedEpochManifest) -> Result { + bail!("epoch activation is unavailable on the local key backend") + } + + async fn authorize_reshare(&self, _plan: ResharePlan) -> Result { + bail!("reshare authorization is unavailable on the local key backend") + } + + async fn export_root_keys(&self) -> Result<(String, Vec)> { + Ok(( + self.root_ca.key.serialize_pem(), + self.k256_key.to_bytes().to_vec(), + )) + } + + async fn run_mpc_operation( + &self, + _operation: crate::mpc_operation::MpcOperation, + _initiator: &str, + ) -> Result> { + bail!("MPC operations are unavailable on the local key backend") + } +} + +/// MPC backend containing only validated threshold shares. No complete root +/// private key is loaded or reconstructed by this backend. +pub(crate) struct MpcKeyBackend { + root_ca_cert: String, + p256_share: P256KeyShare, + k256_share: K256KeyShare, + derivation_share: P256KeyShare, + transport: Arc, + manifest: EpochManifest, + cluster_id: String, + node_id: String, + p256_share_file: std::path::PathBuf, + k256_share_file: std::path::PathBuf, + derivation_share_file: std::path::PathBuf, + manifest_file: std::path::PathBuf, + checkpoint_file: std::path::PathBuf, + identity: ClusterIdentity, + operations: Mutex, Arc>>, +} + +struct OperationRecord { + request_hash: Vec, + initiator: String, + expires_at: u64, + result: Mutex, String>>>, + finished: tokio::sync::Notify, +} + +#[derive(serde::Serialize, serde::Deserialize)] +struct ReshareWire { + p256_public: PublicContribution, + p256_private: PrivateContribution, + k256_public: PublicContribution, + k256_private: PrivateContribution, + derivation_public: PublicContribution, + derivation_private: PrivateContribution, +} + +const MAX_ACTIVE_OPERATIONS: usize = 128; + +impl MpcKeyBackend { + pub(crate) fn load( + root_ca_cert: String, + identity: &ClusterIdentity, + epoch: u64, + node_id: &str, + p256_share_file: &std::path::Path, + k256_share_file: &std::path::Path, + derivation_share_file: &std::path::Path, + manifest_file: &std::path::Path, + checkpoint_file: &std::path::Path, + manifest: &EpochManifest, + router: std::sync::Arc, + rpc_cert: String, + rpc_key: String, + attestation_verifier: std::sync::Arc, + ) -> Result { + let (_, root_pem) = x509_parser::pem::parse_x509_pem(root_ca_cert.as_bytes()) + .context("invalid MPC root CA PEM")?; + let root_certificate = root_pem + .parse_x509() + .context("invalid MPC root CA certificate")?; + ensure!(root_certificate.is_ca(), "MPC root certificate is not a CA"); + let transport = MpcHttpTransport::new( + manifest, + node_id, + router, + rpc_cert, + rpc_key, + root_ca_cert.clone(), + attestation_verifier, + )?; + let backend = Self { + root_ca_cert, + p256_share: load_share( + p256_share_file, + &identity.cluster_id, + epoch, + node_id, + CggmpCurve::P256, + )?, + k256_share: load_share( + k256_share_file, + &identity.cluster_id, + epoch, + node_id, + CggmpCurve::K256, + )?, + derivation_share: load_share( + derivation_share_file, + &identity.cluster_id, + epoch, + node_id, + CggmpCurve::P256, + )?, + transport: Arc::new(transport), + manifest: manifest.clone(), + cluster_id: identity.cluster_id.clone(), + node_id: node_id.into(), + p256_share_file: p256_share_file.into(), + k256_share_file: k256_share_file.into(), + derivation_share_file: derivation_share_file.into(), + manifest_file: manifest_file.into(), + checkpoint_file: checkpoint_file.into(), + identity: identity.clone(), + operations: Mutex::new(BTreeMap::new()), + }; + ensure!( + root_certificate + .public_key() + .subject_public_key + .data + .as_ref() + == backend.p256_public_key(), + "MPC root certificate public key does not match the P-256 group key" + ); + let member_index = manifest + .members + .iter() + .position(|member| member.node_id == node_id) + .context("local MPC member is absent from manifest")?; + validate_share_topology( + &backend.p256_share, + member_index, + manifest.members.len(), + manifest.threshold, + )?; + validate_share_topology( + &backend.k256_share, + member_index, + manifest.members.len(), + manifest.threshold, + )?; + validate_share_topology( + &backend.derivation_share, + member_index, + manifest.members.len(), + manifest.threshold, + )?; + ensure!( + backend.p256_share.core.key_info.public_shares.len() == manifest.members.len() + && backend.k256_share.core.key_info.public_shares.len() == manifest.members.len() + && backend.derivation_share.core.key_info.public_shares.len() + == manifest.members.len(), + "threshold share participant count does not match manifest" + ); + let p256 = backend + .p256_share + .core + .key_info + .public_shares + .get(member_index) + .context("manifest has no matching P-256 public share")? + .to_bytes(false); + let k256 = backend + .k256_share + .core + .key_info + .public_shares + .get(member_index) + .context("manifest has no matching K-256 public share")? + .to_bytes(true); + let derivation = backend + .derivation_share + .core + .key_info + .public_shares + .get(member_index) + .context("manifest has no matching derivation public share")? + .to_bytes(false); + ensure!( + manifest.members[member_index].share_commitment + == share_commitment(p256.as_bytes(), k256.as_bytes(), derivation.as_bytes()), + "local threshold shares do not match the epoch manifest commitment" + ); + Ok(backend) + } + + async fn participants(&self) -> Result> { + let threshold = usize::from(self.manifest.threshold); + let reachable = self.transport.reachable_nodes().await; + let mut selected: Vec<_> = self + .manifest + .members + .iter() + .filter(|member| member.node_id == self.node_id) + .map(|member| member.node_id.clone()) + .collect(); + selected.extend( + self.manifest + .members + .iter() + .filter(|member| { + member.node_id != self.node_id && reachable.contains(&member.node_id) + }) + .take(threshold.saturating_sub(1)) + .map(|member| member.node_id.clone()), + ); + selected.sort_by_key(|node| { + self.manifest + .members + .iter() + .position(|member| &member.node_id == node) + .expect("selected manifest member") + }); + ensure!(selected.len() == threshold, "not enough MPC members"); + Ok(selected) + } + + async fn coordinate_k256(&self, payload: K256SignPayload) -> Result> { + let mut session_id = [0u8; 32]; + OsRng.fill_bytes(&mut session_id); + let operation = MpcOperation::new_k256( + session_id, + self.manifest.epoch, + self.participants().await?, + unix_time()? + .checked_add(60) + .context("MPC expiry overflow")?, + payload, + )?; + operation.validate(&self.manifest, &self.node_id)?; + + let remote = operation + .participants + .iter() + .filter(|node| *node != &self.node_id) + .map(|node| self.transport.start_operation(node, &operation)); + let (local, remote) = tokio::join!( + self.run_mpc_operation(operation.clone(), &self.node_id), + join_all(remote) + ); + let local = local?; + for result in remote { + ensure!( + result? == local, + "MPC participants returned different signatures" + ); + } + self.add_recovery_id(&operation, local) + } + + async fn derive(&self, payload: DerivePayload) -> Result<[u8; 32]> { + let mut session_id = [0u8; 32]; + OsRng.fill_bytes(&mut session_id); + let operation = MpcOperation::new_derivation( + session_id, + self.manifest.epoch, + self.participants().await?, + unix_time()? + .checked_add(60) + .context("MPC expiry overflow")?, + payload.clone(), + )?; + operation.validate(&self.manifest, &self.node_id)?; + let remote = operation + .participants + .iter() + .filter(|node| *node != &self.node_id) + .map(|node| async { + ( + node.clone(), + self.transport.start_operation(node, &operation).await, + ) + }); + let (local, remote) = tokio::join!( + self.run_mpc_operation(operation.clone(), &self.node_id), + join_all(remote) + ); + let mut results = vec![(self.node_id.clone(), local?)]; + for (node, result) in remote { + results.push((node, result?)); + } + let input = payload.input()?; + let request_hash: [u8; 32] = operation.request_hash.as_slice().try_into()?; + let mut evaluations = Vec::with_capacity(results.len()); + for (node, encoded) in results { + let expected_index = self + .manifest + .members + .iter() + .position(|member| member.node_id == node) + .context("derivation responder is not in manifest")?; + let partial: PrfPartial = + serde_json::from_slice(&encoded).context("invalid threshold derivation partial")?; + ensure!( + usize::from(partial.keygen_index) == expected_index, + "threshold derivation partial has wrong signer index" + ); + let public_share = self + .derivation_share + .core + .key_info + .public_shares + .get(expected_index) + .context("missing derivation public share")? + .to_bytes(false); + evaluations.push(( + partial.keygen_index, + threshold_prf::verify(&partial, public_share.as_bytes(), &input, &request_hash)?, + )); + } + threshold_prf::combine(evaluations, &input) + } + + async fn sign_p256_certificate(&self, payload: P256CertificatePayload) -> Result> { + self.validate_p256_certificate(&payload)?; + let mut session_id = [0u8; 32]; + OsRng.fill_bytes(&mut session_id); + let operation = MpcOperation::new_p256_certificate( + session_id, + self.manifest.epoch, + self.participants().await?, + unix_time()? + .checked_add(60) + .context("MPC expiry overflow")?, + payload, + )?; + operation.validate(&self.manifest, &self.node_id)?; + let remote = operation + .participants + .iter() + .filter(|node| *node != &self.node_id) + .map(|node| self.transport.start_operation(node, &operation)); + let (local, remote) = tokio::join!( + self.run_mpc_operation(operation.clone(), &self.node_id), + join_all(remote) + ); + let local = local?; + ensure!( + local.len() == 64, + "invalid threshold P-256 signature length" + ); + for result in remote { + ensure!( + result? == local, + "MPC participants returned different P-256 signatures" + ); + } + Ok(local) + } + + async fn coordinate_manifest_signature( + &self, + manifest: EpochManifest, + ) -> Result { + let mut session_id = [0u8; 32]; + OsRng.fill_bytes(&mut session_id); + let operation = MpcOperation::new_manifest_signature( + session_id, + self.manifest.epoch, + self.participants().await?, + unix_time()? + .checked_add(60) + .context("MPC expiry overflow")?, + ManifestSignaturePayload { + manifest: manifest.clone(), + }, + )?; + operation.validate(&self.manifest, &self.node_id)?; + let remote = operation + .participants + .iter() + .filter(|node| *node != &self.node_id) + .map(|node| self.transport.start_operation(node, &operation)); + let (local, remote) = tokio::join!( + self.run_mpc_operation(operation.clone(), &self.node_id), + join_all(remote) + ); + let local = local?; + ensure!(local.len() == 64, "invalid manifest signature length"); + for result in remote { + ensure!( + result? == local, + "MPC participants returned different manifest signatures" + ); + } + Ok(SignedEpochManifest { + manifest, + signature: local, + }) + } + + async fn coordinate_reshare_authorization( + &self, + plan: ResharePlan, + ) -> Result { + plan.validate_transition(&self.manifest)?; + let mut session_id = [0u8; 32]; + OsRng.fill_bytes(&mut session_id); + let operation = MpcOperation::new_reshare_authorization( + session_id, + self.manifest.epoch, + self.participants().await?, + unix_time()? + .checked_add(60) + .context("MPC expiry overflow")?, + plan.clone(), + )?; + operation.validate(&self.manifest, &self.node_id)?; + let remote = operation + .participants + .iter() + .filter(|node| *node != &self.node_id) + .map(|node| self.transport.start_operation(node, &operation)); + let (local, remote) = tokio::join!( + self.run_mpc_operation(operation.clone(), &self.node_id), + join_all(remote), + ); + let local = local?; + ensure!( + local.len() == 64, + "invalid reshare authorization signature length" + ); + for result in remote { + ensure!( + result? == local, + "MPC participants returned different reshare authorization signatures" + ); + } + let signed = SignedResharePlan { + plan, + signature: local, + }; + signed.verify(&self.identity, &self.manifest)?; + Ok(signed) + } + + async fn coordinate_reshare(&self, plan: ResharePlan) -> Result>> { + plan.validate_subset(&self.manifest)?; + let dealers = if plan.dealers.is_empty() { + self.participants().await? + } else { + let reachable = self.transport.reachable_nodes().await; + ensure!( + plan.dealers + .iter() + .all(|node| node == &self.node_id || reachable.contains(node)), + "an authorized reshare dealer is unavailable" + ); + plan.dealers.clone() + }; + let mut session_id = [0u8; 32]; + OsRng.fill_bytes(&mut session_id); + let operation = MpcOperation::new_reshare( + session_id, + self.manifest.epoch, + dealers, + unix_time()? + .checked_add(300) + .context("MPC expiry overflow")?, + plan.clone(), + )?; + operation.validate(&self.manifest, &self.node_id)?; + let mut executors = operation.participants.clone(); + executors.extend(plan.members.iter().map(|member| member.node_id.clone())); + executors.sort(); + executors.dedup(); + let remote = executors + .iter() + .filter(|node| *node != &self.node_id) + .map(|node| async { + ( + node.clone(), + self.transport.start_operation(node, &operation).await, + ) + }); + let (local, remote) = tokio::join!( + self.run_mpc_operation(operation.clone(), &self.node_id), + join_all(remote) + ); + let mut results = BTreeMap::from([(self.node_id.clone(), local?)]); + for (node, result) in remote { + results.insert(node, result?); + } + let mut commitments = BTreeMap::new(); + for member in &plan.members { + let commitment = results + .remove(&member.node_id) + .context("target reshare member did not return a commitment")?; + ensure!(commitment.len() == 32, "invalid reshared member commitment"); + commitments.insert(member.node_id.clone(), commitment); + } + Ok(commitments) + } + + fn validate_p256_certificate(&self, payload: &P256CertificatePayload) -> Result<()> { + use ra_tls::oids::{PHALA_RATLS_APP_ID, PHALA_RATLS_CERT_USAGE}; + use x509_parser::{ + certificate::TbsCertificate, + der_parser::oid::Oid, + prelude::{FromDer as _, ParsedExtension}, + }; + let (_, root_pem) = x509_parser::pem::parse_x509_pem(self.root_ca_cert.as_bytes()) + .context("invalid root CA PEM")?; + let root = root_pem + .parse_x509() + .context("invalid root CA certificate")?; + let (remaining, tbs) = TbsCertificate::from_der(&payload.tbs_der) + .map_err(|error| anyhow::anyhow!("invalid certificate TBS: {error}"))?; + ensure!(remaining.is_empty(), "certificate TBS has trailing bytes"); + ensure!( + tbs.issuer() == root.subject(), + "certificate issuer is not the root CA" + ); + ensure!( + tbs.signature.algorithm.to_id_string() == "1.2.840.10045.4.3.2", + "certificate does not use ECDSA-with-SHA256" + ); + let common_names = tbs + .subject() + .iter_common_name() + .map(|value| value.as_str().context("invalid certificate common name")) + .collect::>>()?; + ensure!( + common_names == ["Dstack App CA"], + "app CA certificate has an unexpected subject" + ); + let validity_seconds = tbs + .validity() + .not_after + .timestamp() + .checked_sub(tbs.validity().not_before.timestamp()) + .context("app CA certificate validity is inverted")?; + ensure!( + validity_seconds <= 10 * 366 * 24 * 60 * 60, + "app CA certificate validity exceeds ten years" + ); + let constraints = tbs + .basic_constraints() + .context("invalid certificate basic constraints")? + .context("app CA certificate lacks basic constraints")?; + ensure!( + constraints.value.ca && constraints.value.path_len_constraint == Some(0), + "app CA certificate has invalid CA constraints" + ); + let app_oid = Oid::from(PHALA_RATLS_APP_ID) + .map_err(|error| anyhow::anyhow!("invalid app ID OID: {error:?}"))?; + let app_extension = tbs + .get_extension_unique(&app_oid) + .context("duplicate app ID extension")? + .context("app CA certificate lacks app ID")?; + let app_id = yasna::parse_der(app_extension.value, |reader| reader.read_bytes()) + .context("invalid app ID extension")?; + ensure!(app_id == payload.app_id, "certificate app ID mismatch"); + let usage_oid = Oid::from(PHALA_RATLS_CERT_USAGE) + .map_err(|error| anyhow::anyhow!("invalid usage OID: {error:?}"))?; + let usage_extension = tbs + .get_extension_unique(&usage_oid) + .context("duplicate certificate usage extension")? + .context("app CA certificate lacks usage extension")?; + let usage = yasna::parse_der(usage_extension.value, |reader| reader.read_bytes()) + .context("invalid certificate usage extension")?; + ensure!(usage == b"app:ca", "certificate usage is not app:ca"); + ensure!( + matches!( + tbs.public_key().parsed(), + Ok(x509_parser::public_key::PublicKey::EC(_)) + ), + "app CA public key is not an EC key" + ); + // Force duplicate/invalid key-usage extensions to fail parsing as well. + let usage = tbs.key_usage().context("invalid key usage extension")?; + ensure!( + usage.is_some_and(|usage| usage.value.key_cert_sign() && usage.value.crl_sign()), + "app CA certificate lacks certificate-signing key usage" + ); + for extension in tbs.extensions() { + if matches!( + extension.parsed_extension(), + ParsedExtension::ParseError { .. } + ) { + bail!("app CA certificate contains a malformed extension") + } + } + Ok(()) + } + + fn add_recovery_id(&self, operation: &MpcOperation, signature: Vec) -> Result> { + ensure!(signature.len() == 64, "invalid MPC K-256 signature length"); + let signature = K256Signature::from_slice(&signature).context("invalid MPC signature")?; + let public_key = VerifyingKey::from_sec1_bytes(&self.k256_public_key()) + .context("invalid MPC group public key")?; + let digest = match &operation.payload { + MpcOperationPayload::SignK256(payload) => payload.digest(), + MpcOperationPayload::SignManifest(_) | MpcOperationPayload::AuthorizeReshare(_) => { + bail!("policy signatures do not have recovery IDs") + } + MpcOperationPayload::SignP256Certificate(_) => { + bail!("P-256 certificate result does not have a recovery ID") + } + MpcOperationPayload::Derive(_) => bail!("derivation result is not a signature"), + MpcOperationPayload::Reshare(_) => bail!("resharing result is not a signature"), + }; + for recovery_id in 0u8..4 { + let recovery_id = RecoveryId::try_from(recovery_id).context("invalid recovery ID")?; + if VerifyingKey::recover_from_prehash(&digest, &signature, recovery_id) + .is_ok_and(|recovered| recovered == public_key) + { + let mut encoded = signature.to_bytes().to_vec(); + encoded.push(recovery_id.to_byte()); + return Ok(encoded); + } + } + bail!("MPC signature cannot be recovered to the group public key") + } + + async fn execute_once(&self, operation: MpcOperation, initiator: &str) -> Result> { + let (record, owner) = { + let now = unix_time()?; + let mut operations = self.operations.lock().expect("operation mutex poisoned"); + operations.retain(|_, record| record.expires_at >= now); + if let Some(record) = operations.get(&operation.session_id) { + ensure!( + record.request_hash == operation.request_hash && record.initiator == initiator, + "MPC session ID was reused with different operation context" + ); + (record.clone(), false) + } else { + ensure!( + operations + .values() + .filter(|record| record + .result + .lock() + .expect("operation result mutex poisoned") + .is_none()) + .count() + < MAX_ACTIVE_OPERATIONS, + "too many active MPC operations" + ); + ensure!(operations.len() < 4096, "too many cached MPC operations"); + let record = Arc::new(OperationRecord { + request_hash: operation.request_hash.clone(), + initiator: initiator.into(), + expires_at: operation.expires_at, + result: Mutex::new(None), + finished: tokio::sync::Notify::new(), + }); + operations.insert(operation.session_id.clone(), record.clone()); + (record, true) + } + }; + if owner { + let result = self.execute_operation(operation).await; + *record + .result + .lock() + .expect("operation result mutex poisoned") = Some( + result + .as_ref() + .map(|value| value.clone()) + .map_err(|error| format!("{error:#}")), + ); + record.finished.notify_waiters(); + result + } else { + loop { + let notified = record.finished.notified(); + if let Some(result) = record + .result + .lock() + .expect("operation result mutex poisoned") + .clone() + { + return result.map_err(anyhow::Error::msg); + } + notified.await; + } + } + } + + async fn execute_operation(&self, operation: MpcOperation) -> Result> { + match &operation.payload { + MpcOperationPayload::SignK256(_) + | MpcOperationPayload::SignManifest(_) + | MpcOperationPayload::AuthorizeReshare(_) => self.execute_k256(operation).await, + MpcOperationPayload::SignP256Certificate(payload) => { + self.validate_p256_certificate(payload)?; + self.execute_p256(operation).await + } + MpcOperationPayload::Derive(payload) => { + let input = payload.input()?; + let own_index = self + .manifest + .members + .iter() + .position(|member| member.node_id == self.node_id) + .context("local node is not in manifest")?; + ensure!( + usize::from(self.derivation_share.core.i) == own_index, + "derivation share index does not match manifest ordering" + ); + let secret_scalar: &Scalar = self.derivation_share.core.x.as_ref(); + let secret: [u8; 32] = secret_scalar + .to_be_bytes() + .as_bytes() + .try_into() + .context("invalid derivation secret share length")?; + let public = self + .derivation_share + .core + .key_info + .public_shares + .get(own_index) + .context("missing local derivation public share")? + .to_bytes(false); + let request_hash: [u8; 32] = operation.request_hash.as_slice().try_into()?; + let partial = threshold_prf::evaluate( + own_index + .try_into() + .context("derivation share index overflow")?, + &secret, + public.as_bytes(), + &input, + &request_hash, + &mut OsRng, + )?; + serde_json::to_vec(&partial).context("failed to encode derivation partial") + } + MpcOperationPayload::Reshare(plan) => self.execute_reshare(&operation, plan).await, + } + } + + async fn execute_reshare( + &self, + operation: &MpcOperation, + plan: &ResharePlan, + ) -> Result> { + self.authorize_policy_once("reshare", plan.epoch, &plan.authorization_hash()?)?; + let old_dealers = operation + .participants + .iter() + .map(|node| { + self.manifest + .members + .iter() + .position(|member| &member.node_id == node) + .context("reshare dealer is not in active manifest")? + .try_into() + .context("reshare dealer index overflow") + }) + .collect::>>()?; + let local_old_index: u16 = self + .manifest + .members + .iter() + .position(|member| member.node_id == self.node_id) + .context("local node is absent from active manifest")? + .try_into()?; + let local_new_index = plan + .members + .iter() + .position(|member| member.node_id == self.node_id); + let p256_core = Valid::validate(self.p256_share.core.clone()) + .map_err(|error| anyhow::anyhow!("invalid active P-256 core share: {error}"))?; + let k256_core = Valid::validate(self.k256_share.core.clone()) + .map_err(|error| anyhow::anyhow!("invalid active K-256 core share: {error}"))?; + let derivation_core = Valid::validate(self.derivation_share.core.clone()) + .map_err(|error| anyhow::anyhow!("invalid active derivation core share: {error}"))?; + let mut received = BTreeMap::::new(); + if old_dealers.contains(&local_old_index) { + let (p256_public, p256_private) = mpc_reshare::create_contribution( + &p256_core, + &old_dealers, + plan.members.len().try_into()?, + plan.threshold, + &mut OsRng, + )?; + let (k256_public, k256_private) = mpc_reshare::create_contribution( + &k256_core, + &old_dealers, + plan.members.len().try_into()?, + plan.threshold, + &mut OsRng, + )?; + let (derivation_public, derivation_private) = mpc_reshare::create_contribution( + &derivation_core, + &old_dealers, + plan.members.len().try_into()?, + plan.threshold, + &mut OsRng, + )?; + for (new_index, member) in plan.members.iter().enumerate() { + let wire = ReshareWire { + p256_public: p256_public.clone(), + p256_private: p256_private[new_index].clone(), + k256_public: k256_public.clone(), + k256_private: k256_private[new_index].clone(), + derivation_public: derivation_public.clone(), + derivation_private: derivation_private[new_index].clone(), + }; + if member.node_id == self.node_id { + received.insert(local_old_index, wire); + } else { + self.transport + .send(MpcEnvelope { + session_id: operation.session_id.clone(), + epoch: operation.epoch, + protocol: MpcProtocol::Reshare, + request_hash: operation.request_hash.clone(), + sender: self.node_id.clone(), + recipient: member.node_id.clone(), + sequence: u64::from(local_old_index) + 1, + expires_at: operation.expires_at, + payload: serde_json::to_vec(&wire) + .context("failed to encode resharing contribution")?, + broadcast: false, + }) + .await?; + } + } + } + let Some(local_new_index) = local_new_index else { + return Ok(vec![]); + }; + while received.len() < old_dealers.len() { + ensure!( + unix_time()? <= operation.expires_at, + "resharing session expired" + ); + for envelope in self + .transport + .receive( + operation + .session_id + .as_slice() + .try_into() + .context("invalid session ID")?, + ) + .await? + { + let sender_index: u16 = self + .manifest + .members + .iter() + .position(|member| member.node_id == envelope.sender) + .context("resharing message sender is not active")? + .try_into()?; + ensure!( + old_dealers.contains(&sender_index), + "message sender is not a dealer" + ); + ensure!( + !received.contains_key(&sender_index), + "duplicate dealer contribution" + ); + let wire: ReshareWire = serde_json::from_slice(&envelope.payload) + .context("invalid resharing contribution")?; + received.insert(sender_index, wire); + } + if received.len() < old_dealers.len() { + tokio::time::sleep(Duration::from_millis(10)).await; + } + } + let wires = old_dealers + .iter() + .map(|dealer| received.get(dealer).expect("all dealer messages received")) + .collect::>(); + let p256 = mpc_reshare::verify_and_combine( + &p256_core, + &old_dealers, + plan.members.len().try_into()?, + plan.threshold, + local_new_index.try_into()?, + &wires + .iter() + .map(|wire| wire.p256_public.clone()) + .collect::>(), + &wires + .iter() + .map(|wire| wire.p256_private.clone()) + .collect::>(), + )? + .into_incomplete_share(local_new_index.try_into()?, plan.threshold, &p256_core)?; + let k256 = mpc_reshare::verify_and_combine( + &k256_core, + &old_dealers, + plan.members.len().try_into()?, + plan.threshold, + local_new_index.try_into()?, + &wires + .iter() + .map(|wire| wire.k256_public.clone()) + .collect::>(), + &wires + .iter() + .map(|wire| wire.k256_private.clone()) + .collect::>(), + )? + .into_incomplete_share(local_new_index.try_into()?, plan.threshold, &k256_core)?; + let derivation = mpc_reshare::verify_and_combine( + &derivation_core, + &old_dealers, + plan.members.len().try_into()?, + plan.threshold, + local_new_index.try_into()?, + &wires + .iter() + .map(|wire| wire.derivation_public.clone()) + .collect::>(), + &wires + .iter() + .map(|wire| wire.derivation_private.clone()) + .collect::>(), + )? + .into_incomplete_share( + local_new_index.try_into()?, + plan.threshold, + &derivation_core, + )?; + + let target_old_indexes = plan + .members + .iter() + .map(|member| { + self.manifest + .members + .iter() + .position(|active| active.node_id == member.node_id) + .context("reshare target is not active") + }) + .collect::>>()?; + let filtered_aux = |share: &P256KeyShare| -> Result<_> { + let mut aux = share.aux.clone(); + aux.parties = target_old_indexes + .iter() + .map(|index| aux.parties[*index].clone()) + .collect(); + Valid::validate(aux) + .map_err(|error| anyhow::anyhow!("invalid filtered auxiliary share: {error}")) + }; + let p256_complete = cggmp21::KeyShare::from_parts((p256, filtered_aux(&self.p256_share)?))?; + let k256_complete = cggmp21::KeyShare::from_parts((k256, { + let mut aux = self.k256_share.aux.clone(); + aux.parties = target_old_indexes + .iter() + .map(|index| aux.parties[*index].clone()) + .collect(); + Valid::validate(aux).map_err(|error| anyhow::anyhow!("invalid K-256 aux: {error}"))? + }))?; + let derivation_complete = + cggmp21::KeyShare::from_parts((derivation, filtered_aux(&self.derivation_share)?))?; + let pending = + |path: &std::path::Path| path.with_extension(format!("epoch-{}.pending", plan.epoch)); + store_share( + &pending(&self.p256_share_file), + &self.cluster_id, + plan.epoch, + &self.node_id, + CggmpCurve::P256, + &p256_complete, + )?; + store_share( + &pending(&self.k256_share_file), + &self.cluster_id, + plan.epoch, + &self.node_id, + CggmpCurve::K256, + &k256_complete, + )?; + store_share( + &pending(&self.derivation_share_file), + &self.cluster_id, + plan.epoch, + &self.node_id, + CggmpCurve::P256, + &derivation_complete, + )?; + let p256_public = + p256_complete.core.key_info.public_shares[local_new_index].to_bytes(false); + let k256_public = k256_complete.core.key_info.public_shares[local_new_index].to_bytes(true); + let derivation_public = + derivation_complete.core.key_info.public_shares[local_new_index].to_bytes(false); + Ok(share_commitment( + p256_public.as_bytes(), + k256_public.as_bytes(), + derivation_public.as_bytes(), + ) + .to_vec()) + } + + async fn execute_k256(&self, operation: MpcOperation) -> Result> { + let digest = match &operation.payload { + MpcOperationPayload::SignK256(payload) => payload.digest(), + MpcOperationPayload::SignManifest(payload) => { + let digest = payload.digest()?; + self.authorize_policy_once("manifest", payload.manifest.epoch, &digest)?; + digest + } + MpcOperationPayload::AuthorizeReshare(plan) => { + let digest = plan.authorization_hash()?; + self.authorize_policy_once("reshare", plan.epoch, &digest)?; + digest + } + _ => bail!("K-256 executor received a different MPC operation"), + }; + let local_protocol_index: u16 = operation + .participants + .iter() + .position(|node| node == &self.node_id) + .context("local node is not an MPC participant")? + .try_into() + .context("MPC participant index overflow")?; + let keygen_indexes = operation + .participants + .iter() + .map(|node| { + self.manifest + .members + .iter() + .position(|member| &member.node_id == node) + .context("MPC participant is not in manifest")? + .try_into() + .context("MPC keygen index overflow") + }) + .collect::>>()?; + let own_keygen_index = self + .manifest + .members + .iter() + .position(|member| member.node_id == self.node_id) + .context("local node is not in manifest")?; + ensure!( + usize::from(self.k256_share.core.i) == own_keygen_index, + "K-256 share index does not match manifest ordering" + ); + + let eid = execution_id( + &self.cluster_id, + operation.epoch, + CggmpCurve::K256, + operation + .session_id + .as_slice() + .try_into() + .context("invalid MPC session ID")?, + ); + let context = DriverContext { + session_id: operation.session_id.as_slice().try_into()?, + epoch: operation.epoch, + protocol: MpcProtocol::SignK256, + request_hash: operation.request_hash.as_slice().try_into()?, + local_node_id: self.node_id.clone(), + participants: operation.participants.clone(), + expires_at: operation.expires_at, + poll_interval: Duration::from_millis(10), + }; + let share = self.k256_share.clone(); + let transport = self.transport.clone(); + let runtime = tokio::runtime::Handle::current(); + tokio::task::spawn_blocking(move || { + let mut rng = OsRng; + let data = + DataToSign::from_scalar(Scalar::::from_be_bytes_mod_order(digest)); + let state = cggmp21::signing( + ExecutionId::new(&eid), + local_protocol_index, + &keygen_indexes, + &share, + ) + .sign_sync(&mut rng, data); + let signature = drive_state_machine_blocking( + state, + &BlockingTransport::new(transport, runtime), + context, + )??; + let mut encoded = vec![0u8; cggmp21::Signature::::serialized_len()]; + signature.write_to_slice(&mut encoded); + Ok(encoded) + }) + .await + .context("MPC signing worker panicked")? + } + + fn authorize_policy_once(&self, label: &str, epoch: u64, digest: &[u8; 32]) -> Result<()> { + record_epoch_authorization(&self.manifest_file, label, epoch, digest) + } + + async fn execute_p256(&self, operation: MpcOperation) -> Result> { + let MpcOperationPayload::SignP256Certificate(payload) = &operation.payload else { + bail!("P-256 executor received a different MPC operation") + }; + let digest = payload.digest(); + let local_protocol_index: u16 = operation + .participants + .iter() + .position(|node| node == &self.node_id) + .context("local node is not an MPC participant")? + .try_into() + .context("MPC participant index overflow")?; + let keygen_indexes = operation + .participants + .iter() + .map(|node| { + self.manifest + .members + .iter() + .position(|member| &member.node_id == node) + .context("MPC participant is not in manifest")? + .try_into() + .context("MPC keygen index overflow") + }) + .collect::>>()?; + let own_keygen_index = self + .manifest + .members + .iter() + .position(|member| member.node_id == self.node_id) + .context("local node is not in manifest")?; + ensure!( + usize::from(self.p256_share.core.i) == own_keygen_index, + "P-256 share index does not match manifest ordering" + ); + let eid = execution_id( + &self.cluster_id, + operation.epoch, + CggmpCurve::P256, + operation.session_id.as_slice().try_into()?, + ); + let context = DriverContext { + session_id: operation.session_id.as_slice().try_into()?, + epoch: operation.epoch, + protocol: MpcProtocol::SignP256, + request_hash: operation.request_hash.as_slice().try_into()?, + local_node_id: self.node_id.clone(), + participants: operation.participants.clone(), + expires_at: operation.expires_at, + poll_interval: Duration::from_millis(10), + }; + let share = self.p256_share.clone(); + let transport = self.transport.clone(); + let runtime = tokio::runtime::Handle::current(); + tokio::task::spawn_blocking(move || { + let mut rng = OsRng; + let data = + DataToSign::from_scalar(Scalar::::from_be_bytes_mod_order(digest)); + let state = cggmp21::signing( + ExecutionId::new(&eid), + local_protocol_index, + &keygen_indexes, + &share, + ) + .sign_sync(&mut rng, data); + let signature = drive_state_machine_blocking( + state, + &BlockingTransport::new(transport, runtime), + context, + )??; + let mut encoded = vec![0u8; cggmp21::Signature::::serialized_len()]; + signature.write_to_slice(&mut encoded); + Ok(encoded) + }) + .await + .context("MPC P-256 signing worker panicked")? + } +} + +fn unix_time() -> Result { + Ok(std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .context("system time before UNIX epoch")? + .as_secs()) +} + +#[async_trait] +impl KeyBackend for MpcKeyBackend { + async fn derive_app_keys(&self, app_id: &[u8], instance_id: &[u8]) -> Result { + let disk_key = self + .derive(DerivePayload { + purpose: DerivePurpose::DiskKey, + app_id: app_id.to_vec(), + instance_id: Some(instance_id.to_vec()), + counter: 0, + }) + .await?; + let env_key = self.derive_env_key(app_id).await?; + let mut k256_key = None; + for counter in 0..=16 { + let app_seed = self + .derive(DerivePayload { + purpose: DerivePurpose::AppK256, + app_id: app_id.to_vec(), + instance_id: None, + counter, + }) + .await?; + if let Ok(key) = SigningKey::from_slice(&app_seed) { + k256_key = Some(key); + break; + } + } + let k256_key = k256_key.context("threshold derivation could not produce a K-256 key")?; + let public_key = k256_key.verifying_key().to_sec1_bytes(); + let k256_signature = self + .sign_k256(b"dstack-kms-issued", app_id, &public_key) + .await?; + Ok(DerivedAppKeys { + disk_key, + env_key, + k256_key: k256_key.to_bytes().to_vec(), + k256_signature, + }) + } + + async fn derive_env_key(&self, app_id: &[u8]) -> Result<[u8; 32]> { + self.derive(DerivePayload { + purpose: DerivePurpose::EnvKey, + app_id: app_id.to_vec(), + instance_id: None, + counter: 0, + }) + .await + } + + async fn sign_k256(&self, prefix: &[u8], app_id: &[u8], message: &[u8]) -> Result> { + self.coordinate_k256(K256SignPayload { + prefix: prefix.to_vec(), + app_id: app_id.to_vec(), + timestamp: None, + message: message.to_vec(), + }) + .await + } + + async fn sign_k256_timestamped( + &self, + prefix: &[u8], + app_id: &[u8], + timestamp: u64, + message: &[u8], + ) -> Result> { + self.coordinate_k256(K256SignPayload { + prefix: prefix.to_vec(), + app_id: app_id.to_vec(), + timestamp: Some(timestamp), + message: message.to_vec(), + }) + .await + } + + fn k256_public_key(&self) -> Vec { + self.k256_share + .shared_public_key() + .to_bytes(true) + .as_bytes() + .to_vec() + } + + fn p256_public_key(&self) -> Vec { + self.p256_share + .shared_public_key() + .to_bytes(false) + .as_bytes() + .to_vec() + } + + fn derivation_public_key(&self) -> Vec { + self.derivation_share + .shared_public_key() + .to_bytes(false) + .as_bytes() + .to_vec() + } + + fn root_ca_cert(&self) -> &str { + &self.root_ca_cert + } + + async fn derive_app_ca(&self, app_id: &[u8]) -> Result { + let mut app_key = None; + for counter in 0..=16 { + let seed = self + .derive(DerivePayload { + purpose: DerivePurpose::AppCa, + app_id: app_id.to_vec(), + instance_id: None, + counter, + }) + .await?; + if let Ok(key) = kdf::derive_p256_key_pair_from_bytes(&seed, &[app_id, b"app-ca"]) { + app_key = Some(key); + break; + } + } + let app_key = + app_key.context("threshold derivation could not produce a P-256 app CA key")?; + let now = std::time::SystemTime::now(); + let request = CertRequest::builder() + .key(&app_key) + .org_name("Dstack") + .subject("Dstack App CA") + .ca_level(0) + .app_id(app_id) + .special_usage("app:ca") + .not_before(now - Duration::from_secs(5 * 60)) + .not_after(now + Duration::from_secs(10 * 365 * 24 * 60 * 60)) + .build(); + let external = prepare_external_certificate(request, &self.root_ca_cert) + .context("failed to prepare threshold-signed app CA")?; + let signature = self + .sign_p256_certificate(P256CertificatePayload { + app_id: app_id.to_vec(), + tbs_der: external.tbs_der().to_vec(), + }) + .await?; + let certificate = external.finish(&signature)?; + CaCert::new(certificate, app_key.serialize_pem()) + .context("failed to construct threshold-signed app CA") + } + + async fn sign_epoch_manifest(&self, manifest: EpochManifest) -> Result { + self.coordinate_manifest_signature(manifest).await + } + + async fn prepare_reshare(&self, plan: ResharePlan) -> Result>> { + self.coordinate_reshare(plan).await + } + + async fn activate_epoch(&self, signed: SignedEpochManifest) -> Result { + activate_pending_epoch( + &EpochPaths { + manifest: &self.manifest_file, + checkpoint: &self.checkpoint_file, + p256_share: &self.p256_share_file, + k256_share: &self.k256_share_file, + derivation_share: &self.derivation_share_file, + }, + &self.identity, + &self.manifest, + &signed, + &self.cluster_id, + &self.node_id, + ) + } + + async fn authorize_reshare(&self, plan: ResharePlan) -> Result { + self.coordinate_reshare_authorization(plan).await + } + + async fn run_mpc_operation(&self, operation: MpcOperation, initiator: &str) -> Result> { + operation.validate(&self.manifest, initiator)?; + self.execute_once(operation, initiator).await + } + + async fn export_root_keys(&self) -> Result<(String, Vec)> { + bail!("MPC root shares cannot be exported as root keys") + } +} diff --git a/dstack/kms/src/main.rs b/dstack/kms/src/main.rs index 478c5c52a..765c6bd26 100644 --- a/dstack/kms/src/main.rs +++ b/dstack/kms/src/main.rs @@ -18,11 +18,23 @@ use tracing::{info, warn}; mod admin_auth; mod admin_service; +mod cggmp_engine; mod config; // mod ct_log; mod crypto; +mod key_backend; mod main_service; +mod mpc_driver; +mod mpc_genesis; +mod mpc_identity; +mod mpc_join; +mod mpc_lifecycle; +mod mpc_operation; +mod mpc_reshare; +mod mpc_service; +mod mpc_session; mod onboard_service; +mod threshold_prf; fn app_version() -> String { dstack_build_info::app_version!() @@ -69,13 +81,100 @@ async fn run_onboard_service(kms_config: KmsConfig, figment: Figment) -> Result< "/prpc", ra_rpc::prpc_routes!(OnboardState, OnboardHandler, trim: "Onboard."), ) - .manage(state) + .manage(state.clone()) .launch() .await .map_err(|err| anyhow!(err.to_string()))?; Ok(()) } +async fn run_mpc_genesis_service(kms_config: KmsConfig, figment: Figment) -> Result<()> { + use mpc_genesis::{GenesisHandler, GenesisState}; + + let state = GenesisState::new(kms_config)?; + let quote_verifier = QuoteVerifier::new(state.verifier()); + info!("Starting attested MPC genesis service"); + for method in mpc_genesis::rpc_methods() { + info!(" /prpc/{method}"); + } + let figment = figment + .clone() + .merge(Serialized::defaults(figment.find_value("rpc")?)); + let rocket = rocket::custom(figment) + .mount( + "/prpc", + ra_rpc::prpc_routes!(GenesisState, GenesisHandler, trim: "MpcGenesis."), + ) + .manage(state.clone()) + .manage(quote_verifier); + let rocket = rocket + .ignite() + .await + .map_err(|error| anyhow!(error.to_string()))?; + let shutdown = rocket.shutdown(); + let finalized = state.clone(); + tokio::spawn(async move { + finalized.wait_finalized().await; + info!("MPC genesis finalized; shutting down for normal-mode restart"); + shutdown.notify(); + }); + let coordinator = state.clone(); + if coordinator.is_coordinator() { + tokio::spawn(async move { + if let Err(error) = coordinator.coordinate_dkg().await { + tracing::error!("MPC genesis failed: {error:#}"); + } + }); + } + rocket + .launch() + .await + .map_err(|error| anyhow!(error.to_string()))?; + Ok(()) +} + +async fn run_mpc_join_service(kms_config: KmsConfig, figment: Figment) -> Result<()> { + use mpc_join::{JoinHandler, JoinState}; + let state = JoinState::new(kms_config)?; + let quote_verifier = QuoteVerifier::new(state.verifier()); + info!("Starting threshold-authorized MPC join service"); + for method in mpc_join::rpc_methods() { + info!(" /prpc/{method}"); + } + let figment = figment + .clone() + .merge(Serialized::defaults(figment.find_value("rpc")?)); + let rocket = rocket::custom(figment) + .mount( + "/prpc", + ra_rpc::prpc_routes!(JoinState, JoinHandler, trim: "MpcJoin."), + ) + .manage(state.clone()) + .manage(quote_verifier) + .ignite() + .await + .map_err(|error| anyhow!(error.to_string()))?; + let shutdown = rocket.shutdown(); + let completed = state.clone(); + tokio::spawn(async move { + completed.wait_finalized().await; + shutdown.notify(); + }); + let coordinator = state.clone(); + if coordinator.is_coordinator() { + tokio::spawn(async move { + if let Err(error) = coordinator.coordinate().await { + tracing::error!("MPC join failed: {error:#}"); + } + }); + } + rocket + .launch() + .await + .map_err(|error| anyhow!(error.to_string()))?; + Ok(()) +} + #[rocket::get("/metrics")] fn metrics(state: &State) -> RawText { RawText(state.metrics().render_prometheus()) @@ -124,7 +223,10 @@ async fn main() -> Result<()> { ); } - if config.onboard.enabled && !config.keys_exists() { + if config.mpc.enabled && !config.genesis_transport_exists() { + bail!("MPC requires operator-provisioned attested transport material; legacy root-key onboarding is disabled"); + } + if config.onboard.enabled && !config.mpc.enabled && !config.keys_exists() { info!("Onboarding"); run_onboard_service(config.clone(), figment.clone()).await?; if !config.keys_exists() { @@ -132,22 +234,50 @@ async fn main() -> Result<()> { } } - info!("Updating certs"); - if let Err(err) = onboard_service::update_certs(&config).await { - if config.attest_rpc_cert { - return Err(err).context( - "Failed to reissue the attested KMS RPC certificate; refusing to start with a \ - potentially unattested certificate", - ); + if config.mpc.enabled && mpc_genesis::recover_if_needed(&config)? { + info!("Recovered interrupted MPC genesis commit"); + } + if config.mpc.enabled && mpc_join::recover_if_needed(&config)? { + info!("Recovered interrupted MPC membership activation"); + } + + if config.mpc.enabled + && !config.mpc.join_authorization_file.as_os_str().is_empty() + && config.mpc.join_authorization_file.exists() + { + return run_mpc_join_service(config, figment).await; + } + + if config.mpc.enabled && !config.keys_exists() { + if !config.genesis_transport_exists() { + bail!("MPC genesis requires onboarded RPC transport certificates"); } - warn!("Failed to update certs: {err}"); - }; + return run_mpc_genesis_service(config, figment).await; + } + + if config.mpc.enabled { + info!("Using threshold-issued MPC RPC certificate"); + } else { + info!("Updating certs"); + if let Err(err) = onboard_service::update_certs(&config).await { + if config.attest_rpc_cert { + return Err(err).context( + "Failed to reissue the attested KMS RPC certificate; refusing to start with a \ + potentially unattested certificate", + ); + } + warn!("Failed to update certs: {err}"); + }; + } info!("Starting KMS"); info!("Supported methods:"); for method in main_service::rpc_methods() { info!(" /prpc/{method}"); } + for method in mpc_service::rpc_methods() { + info!(" /prpc/{method}"); + } let metrics_enabled = config.metrics.enabled; let admin_config = config.admin.clone(); @@ -180,6 +310,10 @@ async fn main() -> Result<()> { "/prpc", ra_rpc::prpc_routes!(KmsState, RpcHandler, trim: "KMS."), ) + .mount( + "/prpc", + ra_rpc::prpc_routes!(KmsState, mpc_service::MpcHandler, trim: "MpcTransport."), + ) .manage(state.clone()); if metrics_enabled { @@ -194,6 +328,11 @@ async fn main() -> Result<()> { rocket = rocket.manage(quote_verifier); + let rocket = rocket + .ignite() + .await + .map_err(|error| anyhow!(error.to_string()))?; + let main_shutdown = rocket.shutdown(); let main_srv = rocket.launch(); match admin_setup { Some((admin_figment, admin_fairing)) => { @@ -204,21 +343,37 @@ async fn main() -> Result<()> { } else { info!("admin API authentication enabled"); } - let admin_srv = rocket::custom(admin_figment) + let admin_rocket = rocket::custom(admin_figment) .attach(admin_fairing) .mount("/", admin_auth::routes()) .mount( "/prpc", ra_rpc::prpc_routes!(KmsState, admin_service::AdminRpcHandler, trim: "Admin."), ) - .manage(state) - .launch(); + .manage(state.clone()) + .ignite() + .await + .map_err(|error| anyhow!(error.to_string()))?; + let admin_shutdown = admin_rocket.shutdown(); + let restart_state = state.clone(); + tokio::spawn(async move { + restart_state.wait_restart_requested().await; + info!("MPC epoch activated; shutting down for epoch restart"); + main_shutdown.notify(); + admin_shutdown.notify(); + }); + let admin_srv = admin_rocket.launch(); tokio::try_join!( async { main_srv.await.map_err(|err| anyhow!(err.to_string())) }, async { admin_srv.await.map_err(|err| anyhow!(err.to_string())) }, )?; } None => { + let restart_state = state.clone(); + tokio::spawn(async move { + restart_state.wait_restart_requested().await; + main_shutdown.notify(); + }); main_srv.await.map_err(|err| anyhow!(err.to_string()))?; } } diff --git a/dstack/kms/src/main_service.rs b/dstack/kms/src/main_service.rs index 339fbd3bc..51ef95d29 100644 --- a/dstack/kms/src/main_service.rs +++ b/dstack/kms/src/main_service.rs @@ -5,7 +5,7 @@ use std::{ path::Path, sync::{ - atomic::{AtomicU64, Ordering}, + atomic::{AtomicBool, AtomicU64, Ordering}, Arc, }, }; @@ -19,12 +19,10 @@ use dstack_kms_rpc::{ }; use dstack_verifier::{CvmVerifier, VerificationDetails}; use fs_err as fs; -use k256::ecdsa::SigningKey; use ra_rpc::{CallContext, RpcCall}; use ra_tls::{ attestation::{AttestationVerifier, TeeVariant, VerifiedAttestation}, - cert::{CaCert, CertRequest, CertSigningRequestV1, CertSigningRequestV2, Csr}, - kdf, + cert::{CertSigningRequestV1, CertSigningRequestV2, Csr}, }; use scale::Decode; use tokio::sync::OnceCell; @@ -33,7 +31,10 @@ use upgrade_authority::{build_boot_info, ensure_app_id_len, local_kms_boot_info, use crate::{ config::KmsConfig, - crypto::{derive_k256_key, sign_message, sign_message_with_timestamp}, + key_backend::{KeyBackend, LocalKeyBackend, MpcKeyBackend}, + mpc_identity::{decode_hex, ClusterIdentity, EpochManifest, NodeEvidence, SignedEpochManifest}, + mpc_lifecycle, + mpc_session::SessionRouter, }; pub(crate) mod amd_attest; @@ -54,14 +55,18 @@ impl std::ops::Deref for KmsState { pub struct KmsStateInner { config: KmsConfig, - root_ca: CaCert, - k256_key: SigningKey, temp_ca_cert: String, temp_ca_key: String, verifier: CvmVerifier, attestation_verifier: Arc, self_boot_info: OnceCell, metrics: KmsMetrics, + mpc_identity: Option, + signed_manifest: Option, + key_backend: Arc, + mpc_router: Option>, + restart_requested: AtomicBool, + restart_notify: tokio::sync::Notify, } #[derive(Default)] @@ -134,19 +139,183 @@ impl KmsState { } pub fn new(config: KmsConfig) -> Result { - let root_ca = CaCert::load(config.root_ca_cert(), config.root_ca_key()) + let root_ca_cert = fs::read_to_string(config.root_ca_cert()) .context("Failed to load root CA certificate")?; - let key_bytes = fs::read(config.k256_key()).context("Failed to read ECDSA root key")?; - let k256_key = - SigningKey::from_slice(&key_bytes).context("Failed to load ECDSA root key")?; - let temp_ca_key = - fs::read_to_string(config.tmp_ca_key()).context("Faeild to read temp ca key")?; - let temp_ca_cert = - fs::read_to_string(config.tmp_ca_cert()).context("Faeild to read temp ca cert")?; let attestation_verifier = Arc::new( AttestationVerifier::load(&config.attestation) .context("failed to load attestation verifier")?, ); + let configured_identity = if config.mpc.enabled { + Some(if config.mpc.identity_file.as_os_str().is_empty() { + ClusterIdentity::new( + config.mpc.protocol_version, + config.mpc.cluster_id.clone(), + decode_hex("mpc.p256_group_pubkey", &config.mpc.p256_group_pubkey)?, + decode_hex("mpc.k256_group_pubkey", &config.mpc.k256_group_pubkey)?, + decode_hex( + "mpc.derivation_group_pubkey", + &config.mpc.derivation_group_pubkey, + )?, + )? + } else { + serde_json::from_slice( + &fs::read(&config.mpc.identity_file) + .context("failed to read MPC identity file")?, + ) + .context("failed to parse MPC identity file")? + }) + } else { + None + }; + if let Some(identity) = &configured_identity { + identity.validate()?; + crate::mpc_lifecycle::recover_pending_activation( + &crate::mpc_lifecycle::EpochPaths { + manifest: &config.mpc.manifest_file, + checkpoint: &config.mpc.checkpoint_file, + p256_share: &config.mpc.p256_share_file, + k256_share: &config.mpc.k256_share_file, + derivation_share: &config.mpc.derivation_share_file, + }, + identity, + &identity.cluster_id, + &config.mpc.node_id, + )?; + } + let mut signed_manifest: Option = None; + let manifest: Option = if config.mpc.enabled { + anyhow::ensure!( + !config.mpc.manifest_file.as_os_str().is_empty(), + "MPC manifest_file must not be empty" + ); + let encoded = + fs::read(&config.mpc.manifest_file).context("failed to read MPC manifest")?; + match serde_json::from_slice::(&encoded) { + Ok(signed) => { + signed.verify( + configured_identity + .as_ref() + .context("MPC identity is missing")?, + )?; + let manifest = signed.manifest.clone(); + signed_manifest = Some(signed); + Some(manifest) + } + Err(signed_error) => { + anyhow::ensure!( + config.mpc.allow_unsigned_manifest, + "MPC manifest must carry a valid threshold signature: {signed_error}" + ); + Some( + serde_json::from_slice(&encoded) + .context("failed to parse unsigned MPC manifest")?, + ) + } + } + } else { + None + }; + let mpc_router = manifest + .as_ref() + .map(|manifest| { + SessionRouter::new( + manifest.clone(), + config.mpc.max_sessions, + config.mpc.session_ttl, + ) + .map(Arc::new) + }) + .transpose()?; + if let Some(manifest) = &manifest { + let local = manifest + .members + .iter() + .find(|member| member.node_id == config.mpc.node_id) + .context("local MPC member is absent from manifest")?; + for path in [config.rpc_cert(), config.mpc.client_cert_file.clone()] { + let (_, pem) = x509_parser::pem::parse_x509_pem(&fs::read(path)?)?; + anyhow::ensure!( + pem.parse_x509()?.public_key().raw == local.attestation_pubkey, + "local MPC server/client certificate key differs from manifest" + ); + } + } + let key_backend: Arc = if config.mpc.enabled { + let manifest = manifest.as_ref().context("MPC manifest is missing")?; + Arc::new(MpcKeyBackend::load( + root_ca_cert, + configured_identity + .as_ref() + .context("MPC identity is missing")?, + manifest.epoch, + &config.mpc.node_id, + &config.mpc.p256_share_file, + &config.mpc.k256_share_file, + &config.mpc.derivation_share_file, + &config.mpc.manifest_file, + &config.mpc.checkpoint_file, + manifest, + mpc_router.clone().context("MPC router is missing")?, + fs::read_to_string(&config.mpc.client_cert_file) + .context("failed to read MPC client cert")?, + fs::read_to_string(&config.mpc.client_key_file) + .context("failed to read MPC client key")?, + attestation_verifier.clone(), + )?) + } else { + let key_bytes = fs::read(config.k256_key()).context("Failed to read ECDSA root key")?; + Arc::new(LocalKeyBackend::from_pem_and_bytes( + root_ca_cert, + fs::read_to_string(config.root_ca_key())?, + &key_bytes, + )?) + }; + let mpc_identity = if config.mpc.enabled { + let identity = configured_identity.context("MPC identity is missing")?; + anyhow::ensure!( + identity.p256_group_pubkey == key_backend.p256_public_key(), + "MPC P-256 group public key does not match the active root CA" + ); + anyhow::ensure!( + identity.k256_group_pubkey == key_backend.k256_public_key(), + "MPC K-256 group public key does not match the active signing key" + ); + anyhow::ensure!( + identity.derivation_group_pubkey == key_backend.derivation_public_key(), + "MPC derivation group public key does not match the active derivation share" + ); + anyhow::ensure!( + !config.mpc.node_id.is_empty(), + "MPC node_id must not be empty" + ); + let manifest = manifest.as_ref().context("MPC manifest is missing")?; + anyhow::ensure!( + manifest.provider_id == identity.provider_id(), + "MPC manifest provider ID does not match cluster identity" + ); + anyhow::ensure!( + manifest.contains_member(&config.mpc.node_id), + "local node is not a member of the MPC epoch" + ); + if let Some(signed) = &signed_manifest { + anyhow::ensure!( + !config.mpc.checkpoint_file.as_os_str().is_empty(), + "MPC checkpoint_file must not be empty for signed manifests" + ); + mpc_lifecycle::validate_and_checkpoint( + &config.mpc.checkpoint_file, + signed, + &identity, + )?; + } + Some(identity) + } else { + None + }; + let temp_ca_key = + fs::read_to_string(config.tmp_ca_key()).context("Faeild to read temp ca key")?; + let temp_ca_cert = + fs::read_to_string(config.tmp_ca_cert()).context("Faeild to read temp ca cert")?; let verifier = CvmVerifier::new( config.image.cache_dir.display().to_string(), config.image.download_url.clone(), @@ -161,25 +330,59 @@ impl KmsState { Ok(Self { inner: Arc::new(KmsStateInner { config, - root_ca, - k256_key, temp_ca_cert, temp_ca_key, verifier, attestation_verifier, self_boot_info: OnceCell::new(), metrics: KmsMetrics::default(), + mpc_identity, + signed_manifest, + key_backend, + mpc_router, + restart_requested: AtomicBool::new(false), + restart_notify: tokio::sync::Notify::new(), }), }) } + fn key_provider_id(&self) -> Vec { + self.mpc_identity + .as_ref() + .map(|identity| identity.provider_id().to_vec()) + .unwrap_or_default() + } + pub(crate) fn metrics(&self) -> &KmsMetrics { &self.inner.metrics } + pub(crate) fn request_restart(&self) { + self.restart_requested.store(true, Ordering::Release); + self.restart_notify.notify_waiters(); + } + + pub(crate) async fn wait_restart_requested(&self) { + loop { + let notified = self.restart_notify.notified(); + if self.restart_requested.load(Ordering::Acquire) { + return; + } + notified.await; + } + } + pub(crate) fn attestation_verifier(&self) -> Arc { self.inner.attestation_verifier.clone() } + + pub(crate) fn mpc_router(&self) -> Option<&SessionRouter> { + self.inner.mpc_router.as_deref() + } + + pub(crate) fn key_backend(&self) -> &dyn KeyBackend { + self.inner.key_backend.as_ref() + } } pub struct RpcHandler { @@ -330,26 +533,6 @@ impl RpcHandler { gateway_app_id: response.gateway_app_id, }) } - - fn derive_app_ca(&self, app_id: &[u8]) -> Result { - let context_data = vec![app_id, b"app-ca"]; - let app_key = kdf::derive_p256_key_pair(&self.state.root_ca.key, &context_data) - .context("Failed to derive app disk key")?; - let req = CertRequest::builder() - .key(&app_key) - .org_name("Dstack") - .subject("Dstack App CA") - .ca_level(0) - .app_id(app_id) - .special_usage("app:ca") - .build(); - let app_ca = self - .state - .root_ca - .sign(req) - .context("Failed to sign App CA")?; - Ok(CaCert::from_parts(app_key, app_ca)) - } } impl KmsRpc for RpcHandler { @@ -376,32 +559,23 @@ impl KmsRpc for RpcHandler { let instance_id = boot_info.instance_id; let os_image_hash = boot_info.os_image_hash; - let context_data = vec![&app_id[..], &instance_id[..], b"app-disk-crypt-key"]; - let app_disk_key = kdf::derive_dh_secret(&self.state.root_ca.key, &context_data) - .context("Failed to derive app disk key")?; - let env_crypt_key = { - let secret = - kdf::derive_dh_secret(&self.state.root_ca.key, &[&app_id[..], b"env-encrypt-key"]) - .context("Failed to derive env encrypt key")?; - let secret = x25519_dalek::StaticSecret::from(secret); - secret.to_bytes() - }; - - let (k256_key, k256_signature) = { - let (k256_app_key, signature) = derive_k256_key(&self.state.k256_key, &app_id) - .context("Failed to derive app ecdsa key")?; - (k256_app_key.to_bytes().to_vec(), signature) - }; + let derived = self + .state + .key_backend + .derive_app_keys(&app_id, &instance_id) + .await + .context("Failed to derive app keys")?; Ok(AppKeyResponse { - ca_cert: self.state.root_ca.pem_cert.clone(), - disk_crypt_key: app_disk_key.to_vec(), - env_crypt_key: env_crypt_key.to_vec(), - k256_key, - k256_signature, + ca_cert: self.state.key_backend.root_ca_cert().to_string(), + disk_crypt_key: derived.disk_key.to_vec(), + env_crypt_key: derived.env_key.to_vec(), + k256_key: derived.k256_key, + k256_signature: derived.k256_signature, tproxy_app_id: gateway_app_id.clone(), gateway_app_id, os_image_hash, + key_provider_id: self.state.key_provider_id(), }) } @@ -410,11 +584,11 @@ impl KmsRpc for RpcHandler { .await .context("KMS self authorization failed")?; ensure_app_id_len(&request.app_id)?; - let secret = kdf::derive_dh_secret( - &self.state.root_ca.key, - &[&request.app_id[..], "env-encrypt-key".as_bytes()], - ) - .context("Failed to derive env encrypt key")?; + let secret = self + .state + .key_backend + .derive_env_key(&request.app_id) + .await?; let secret = x25519_dalek::StaticSecret::from(secret); let pubkey = x25519_dalek::PublicKey::from(&secret); @@ -425,23 +599,25 @@ impl KmsRpc for RpcHandler { .as_secs(); // Legacy signature (without timestamp) for backward compatibility - let signature = sign_message( - &self.state.k256_key, - b"dstack-env-encrypt-pubkey", - &request.app_id, - &public_key, - ) - .context("Failed to sign the public key")?; + let signature = self + .state + .key_backend + .sign_k256(b"dstack-env-encrypt-pubkey", &request.app_id, &public_key) + .await + .context("Failed to sign the public key")?; // New signature with timestamp to prevent replay attacks - let signature_v1 = sign_message_with_timestamp( - &self.state.k256_key, - b"dstack-env-encrypt-pubkey", - &request.app_id, - timestamp, - &public_key, - ) - .context("Failed to sign the public key with timestamp")?; + let signature_v1 = self + .state + .key_backend + .sign_k256_timestamped( + b"dstack-env-encrypt-pubkey", + &request.app_id, + timestamp, + &public_key, + ) + .await + .context("Failed to sign the public key with timestamp")?; Ok(PublicKeyResponse { public_key, @@ -457,21 +633,54 @@ impl KmsRpc for RpcHandler { .and_then(|s| serde_json::from_str(&s).ok()); let info = self.state.config.auth_api.get_info().await?; Ok(GetMetaResponse { - ca_cert: self.state.inner.root_ca.pem_cert.clone(), + ca_cert: self.state.key_backend.root_ca_cert().to_string(), allow_any_upgrade: self.state.inner.config.auth_api.is_dev(), - k256_pubkey: self - .state - .inner - .k256_key - .verifying_key() - .to_sec1_bytes() - .to_vec(), + k256_pubkey: self.state.key_backend.k256_public_key(), bootstrap_info, is_dev: self.state.config.auth_api.is_dev(), kms_contract_address: info.kms_contract_address, chain_id: info.chain_id, gateway_app_id: info.gateway_app_id, app_auth_implementation: info.app_implementation, + key_provider_id: self.state.key_provider_id(), + mpc_cluster_identity: self + .state + .mpc_identity + .as_ref() + .map(serde_json::to_string) + .transpose() + .context("failed to encode MPC cluster identity")?, + mpc_epoch_manifest: self + .state + .signed_manifest + .as_ref() + .map(|manifest| { + String::from_utf8(serde_jcs::to_vec(manifest)?) + .context("signed manifest JSON is not UTF-8") + }) + .transpose() + .context("failed to encode signed MPC epoch manifest")?, + mpc_node_evidence: match (&self.state.mpc_identity, &self.state.signed_manifest) { + (Some(identity), Some(signed)) => { + let member = signed + .manifest + .members + .iter() + .find(|member| member.node_id == self.state.config.mpc.node_id) + .context("local node is absent from active MPC manifest")?; + let evidence = NodeEvidence { + provider_id: identity.provider_id().to_vec(), + epoch: signed.manifest.epoch, + manifest_hash: signed.manifest.manifest_hash()?.to_vec(), + node_id: member.node_id.clone(), + attestation_pubkey: member.attestation_pubkey.clone(), + share_commitment: member.share_commitment.clone(), + }; + evidence.report_data_hash()?; + Some(String::from_utf8(serde_jcs::to_vec(&evidence)?)?) + } + _ => None, + }, }) } @@ -485,12 +694,15 @@ impl KmsRpc for RpcHandler { self.state.config.sev_snp_key_release, self.state.config.aws_nitro_tpm_key_release, )?; + let (ca_key, k256_key) = self + .state + .key_backend + .export_root_keys() + .await + .context("key backend does not permit root key export")?; Ok(KmsKeyResponse { temp_ca_key: self.state.inner.temp_ca_key.clone(), - keys: vec![KmsKeys { - ca_key: self.state.inner.root_ca.key.serialize_pem(), - k256_key: self.state.inner.k256_key.to_bytes().to_vec(), - }], + keys: vec![KmsKeys { ca_key, k256_key }], }) } @@ -507,7 +719,7 @@ impl KmsRpc for RpcHandler { Ok(GetTempCaCertResponse { temp_ca_cert: self.state.inner.temp_ca_cert.clone(), temp_ca_key: self.state.inner.temp_ca_key.clone(), - ca_cert: self.state.inner.root_ca.pem_cert.clone(), + ca_cert: self.state.key_backend.root_ca_cert().to_string(), }) } @@ -547,7 +759,11 @@ impl KmsRpc for RpcHandler { self.state.config.sev_snp_key_release, self.state.config.aws_nitro_tpm_key_release, )?; - let app_ca = self.derive_app_ca(&app_info.boot_info.app_id)?; + let app_ca = self + .state + .key_backend + .derive_app_ca(&app_info.boot_info.app_id) + .await?; let cert = app_ca .sign_csr(&csr, Some(&app_info.boot_info.app_id), "app:custom") .context("Failed to sign certificate")?; @@ -555,7 +771,7 @@ impl KmsRpc for RpcHandler { certificate_chain: vec![ cert.pem(), app_ca.pem_cert.clone(), - self.state.root_ca.pem_cert.clone(), + self.state.key_backend.root_ca_cert().to_string(), ], }) } diff --git a/dstack/kms/src/mpc_driver.rs b/dstack/kms/src/mpc_driver.rs new file mode 100644 index 000000000..3558f65f4 --- /dev/null +++ b/dstack/kms/src/mpc_driver.rs @@ -0,0 +1,709 @@ +// SPDX-FileCopyrightText: © 2024-2026 Phala Network +// +// SPDX-License-Identifier: Apache-2.0 + +//! Drives `round-based` state machines over the attested MPC envelope layer. + +use std::{collections::VecDeque, time::Duration}; + +use anyhow::{bail, ensure, Context, Result}; +use async_trait::async_trait; +use round_based::{ + state_machine::{ProceedResult, StateMachine}, + Incoming, MessageDestination, MessageType, +}; +use serde::{de::DeserializeOwned, Serialize}; + +use dstack_kms_rpc::{mpc_transport_client::MpcTransportClient, MpcPushRequest, MpcStartRequest}; +use ra_rpc::client::{RaClient, RaClientConfig}; +use ra_tls::attestation::AttestationVerifier; + +use crate::{ + mpc_identity::EpochManifest, + mpc_session::{MpcEnvelope, MpcProtocol, SessionRouter}, +}; + +#[derive(Clone)] +pub(crate) struct DriverContext { + pub session_id: [u8; 32], + pub epoch: u64, + pub protocol: MpcProtocol, + pub request_hash: [u8; 32], + pub local_node_id: String, + /// Protocol index -> epoch node ID. Ordering is consensus-critical. + pub participants: Vec, + pub expires_at: u64, + pub poll_interval: Duration, +} + +impl DriverContext { + fn local_index(&self) -> Result { + self.participants + .iter() + .position(|node| node == &self.local_node_id) + .context("local node is not a protocol participant")? + .try_into() + .context("too many protocol participants") + } +} + +#[async_trait] +pub(crate) trait EnvelopeTransport: Send + Sync { + async fn send(&self, envelope: MpcEnvelope) -> Result<()>; + async fn receive(&self, session_id: &[u8; 32]) -> Result>; +} + +pub(crate) struct BlockingTransport { + inner: std::sync::Arc, + runtime: tokio::runtime::Handle, +} + +impl BlockingTransport { + pub(crate) fn new(inner: std::sync::Arc, runtime: tokio::runtime::Handle) -> Self { + Self { inner, runtime } + } + + fn send(&self, envelope: MpcEnvelope) -> Result<()> { + self.runtime.block_on(self.inner.send(envelope)) + } + + fn receive(&self, session_id: &[u8; 32]) -> Result> { + self.runtime.block_on(self.inner.receive(session_id)) + } +} + +pub(crate) struct MpcHttpTransport { + local_node_id: String, + local_router: std::sync::Arc, + clients: std::collections::BTreeMap>, +} + +impl MpcHttpTransport { + #[allow(clippy::too_many_arguments)] + pub(crate) fn new( + manifest: &EpochManifest, + local_node_id: &str, + local_router: std::sync::Arc, + client_cert: String, + client_key: String, + ca_cert: String, + attestation_verifier: std::sync::Arc, + ) -> Result { + ensure!( + manifest.contains_member(local_node_id), + "local MPC node is not in the epoch manifest" + ); + let mut clients = std::collections::BTreeMap::new(); + for member in &manifest.members { + if member.node_id == local_node_id { + continue; + } + let expected_public_key = member.attestation_pubkey.clone(); + let expected_node_id = member.node_id.clone(); + let client = RaClientConfig::builder() + .remote_uri(member.endpoint.clone()) + .tls_client_cert(client_cert.clone()) + .tls_client_key(client_key.clone()) + .tls_ca_cert(ca_cert.clone()) + .tls_built_in_root_certs(false) + .attestation_verifier(attestation_verifier.clone()) + .cert_validator(Box::new(move |info| { + let info = info.context("MPC peer did not present a certificate")?; + ensure!( + info.attestation.is_some(), + "MPC peer certificate is not attested" + ); + let (_, certificate) = x509_parser::parse_x509_certificate(&info.cert_der) + .context("invalid MPC peer certificate")?; + ensure!( + certificate.public_key().raw == expected_public_key, + "MPC peer certificate key does not match manifest member {expected_node_id}" + ); + Ok(()) + })) + .build() + .into_client() + .with_context(|| format!("failed to build MPC client for {}", member.node_id))?; + clients.insert(member.node_id.clone(), MpcTransportClient::new(client)); + } + Ok(Self { + local_node_id: local_node_id.into(), + local_router, + clients, + }) + } + + pub(crate) async fn start_operation( + &self, + node_id: &str, + operation: &crate::mpc_operation::MpcOperation, + ) -> Result> { + self.clients + .get(node_id) + .context("MPC operation participant has no attested client")? + .start(MpcStartRequest { + operation_json: serde_json::to_vec(operation) + .context("failed to encode MPC operation")?, + }) + .await + .with_context(|| format!("MPC operation failed on participant {node_id}")) + .map(|response| response.result) + } + + pub(crate) async fn reachable_nodes(&self) -> Vec { + futures::future::join_all(self.clients.iter().map(|(node_id, client)| async move { + tokio::time::timeout(Duration::from_secs(3), client.ping()) + .await + .ok() + .and_then(Result::ok) + .map(|_| node_id.clone()) + })) + .await + .into_iter() + .flatten() + .collect() + } +} + +#[async_trait] +impl EnvelopeTransport for MpcHttpTransport { + async fn send(&self, envelope: MpcEnvelope) -> Result<()> { + ensure!( + envelope.sender == self.local_node_id, + "outgoing MPC sender does not match local node" + ); + self.clients + .get(&envelope.recipient) + .context("MPC recipient has no attested client")? + .push(MpcPushRequest { + envelope_json: serde_json::to_vec(&envelope) + .context("failed to encode MPC envelope")?, + }) + .await + .context("failed to deliver MPC envelope") + } + + async fn receive(&self, session_id: &[u8; 32]) -> Result> { + self.local_router.drain(&self.local_node_id, session_id) + } +} + +pub(crate) async fn drive_state_machine( + mut state: S, + transport: &T, + context: DriverContext, +) -> Result +where + S: StateMachine, + S::Msg: Serialize + DeserializeOwned, + T: EnvelopeTransport, +{ + ensure!( + context.participants.len() >= 2, + "MPC needs at least two participants" + ); + let local_index = context.local_index()?; + let mut sequence = 0u64; + let mut pending = VecDeque::new(); + loop { + match state.proceed() { + ProceedResult::SendMsg(outgoing) => { + sequence = sequence.checked_add(1).context("MPC sequence overflow")?; + let payload = serde_json::to_vec(&outgoing.msg) + .context("failed to encode MPC protocol message")?; + let (recipients, broadcast): (Vec, bool) = match outgoing.recipient { + MessageDestination::AllParties => ( + (0..context.participants.len()) + .filter(|index| *index != usize::from(local_index)) + .map(|index| u16::try_from(index).expect("participant index overflow")) + .collect(), + true, + ), + MessageDestination::OneParty(index) => (vec![index], false), + }; + for recipient in recipients { + let recipient = context + .participants + .get(usize::from(recipient)) + .context("protocol emitted an invalid recipient index")?; + transport + .send(MpcEnvelope { + session_id: context.session_id.to_vec(), + epoch: context.epoch, + protocol: context.protocol, + request_hash: context.request_hash.to_vec(), + sender: context.local_node_id.clone(), + recipient: recipient.clone(), + sequence, + expires_at: context.expires_at, + payload: payload.clone(), + broadcast, + }) + .await?; + } + } + ProceedResult::NeedsOneMoreMessage => { + while pending.is_empty() { + ensure_not_expired(context.expires_at)?; + pending.extend(transport.receive(&context.session_id).await?); + if pending.is_empty() { + tokio::time::sleep(context.poll_interval).await; + } + } + let envelope = pending.pop_front().expect("pending queue is not empty"); + ensure!( + envelope.recipient == context.local_node_id, + "misrouted MPC message" + ); + let sender: u16 = context + .participants + .iter() + .position(|node| node == &envelope.sender) + .context("message sender is not a protocol participant")? + .try_into() + .context("sender index overflow")?; + let incoming = Incoming { + id: envelope.sequence, + sender, + msg_type: if envelope.broadcast { + MessageType::Broadcast + } else { + MessageType::P2P + }, + msg: serde_json::from_slice(&envelope.payload) + .context("failed to decode MPC protocol message")?, + }; + state + .received_msg(incoming) + .map_err(|_| anyhow::anyhow!("MPC state machine rejected incoming message"))?; + } + ProceedResult::Yielded => tokio::task::yield_now().await, + ProceedResult::Output(output) => return Ok(output), + ProceedResult::Error(error) => bail!("MPC state machine failed: {error}"), + } + } +} + +/// Executes a non-`Send` round-based state machine on a dedicated blocking +/// thread while network I/O is serviced by the Tokio runtime. CGGMP's sync +/// adapter internally uses `Rc`, so it must never be moved between threads. +pub(crate) fn drive_state_machine_blocking( + mut state: S, + transport: &BlockingTransport, + context: DriverContext, +) -> Result +where + S: StateMachine, + S::Msg: Serialize + DeserializeOwned, + T: EnvelopeTransport, +{ + ensure!( + context.participants.len() >= 2, + "MPC needs at least two participants" + ); + let local_index = context.local_index()?; + let mut sequence = 0u64; + let mut pending = VecDeque::new(); + loop { + match state.proceed() { + ProceedResult::SendMsg(outgoing) => { + sequence = sequence.checked_add(1).context("MPC sequence overflow")?; + let payload = serde_json::to_vec(&outgoing.msg) + .context("failed to encode MPC protocol message")?; + let (recipients, broadcast): (Vec, bool) = match outgoing.recipient { + MessageDestination::AllParties => ( + (0..context.participants.len()) + .filter(|index| *index != usize::from(local_index)) + .map(|index| u16::try_from(index).expect("participant index overflow")) + .collect(), + true, + ), + MessageDestination::OneParty(index) => (vec![index], false), + }; + for recipient in recipients { + let recipient = context + .participants + .get(usize::from(recipient)) + .context("protocol emitted an invalid recipient index")?; + transport.send(MpcEnvelope { + session_id: context.session_id.to_vec(), + epoch: context.epoch, + protocol: context.protocol, + request_hash: context.request_hash.to_vec(), + sender: context.local_node_id.clone(), + recipient: recipient.clone(), + sequence, + expires_at: context.expires_at, + payload: payload.clone(), + broadcast, + })?; + } + } + ProceedResult::NeedsOneMoreMessage => { + while pending.is_empty() { + ensure_not_expired(context.expires_at)?; + pending.extend(transport.receive(&context.session_id)?); + if pending.is_empty() { + std::thread::sleep(context.poll_interval); + } + } + let envelope = pending.pop_front().expect("pending queue is not empty"); + ensure!( + envelope.recipient == context.local_node_id, + "misrouted MPC message" + ); + let sender: u16 = context + .participants + .iter() + .position(|node| node == &envelope.sender) + .context("message sender is not a protocol participant")? + .try_into() + .context("sender index overflow")?; + state + .received_msg(Incoming { + id: envelope.sequence, + sender, + msg_type: if envelope.broadcast { + MessageType::Broadcast + } else { + MessageType::P2P + }, + msg: serde_json::from_slice(&envelope.payload) + .context("failed to decode MPC protocol message")?, + }) + .map_err(|_| anyhow::anyhow!("MPC state machine rejected incoming message"))?; + } + ProceedResult::Yielded => std::thread::yield_now(), + ProceedResult::Output(output) => return Ok(output), + ProceedResult::Error(error) => bail!("MPC state machine failed: {error}"), + } + } +} + +fn ensure_not_expired(expires_at: u64) -> Result<()> { + let now = std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .context("system time before UNIX epoch")? + .as_secs(); + ensure!(now <= expires_at, "MPC session expired"); + Ok(()) +} + +#[cfg(test)] +mod tests { + use std::{collections::BTreeMap, sync::Arc}; + + use cggmp21::{key_share::AnyKeyShare as _, supported_curves::Secp256k1, ExecutionId}; + use rand::rngs::OsRng; + + use super::*; + use crate::{ + cggmp_engine::{execution_id, CggmpCurve}, + mpc_identity::{EpochManifest, EpochMember}, + mpc_session::SessionRouter, + }; + + #[derive(Clone)] + struct MemoryTransport { + local: String, + routers: Arc>>, + } + + #[async_trait] + impl EnvelopeTransport for MemoryTransport { + async fn send(&self, envelope: MpcEnvelope) -> Result<()> { + self.routers + .get(&envelope.recipient) + .context("missing recipient router")? + .push(&self.local, envelope) + } + + async fn receive(&self, session_id: &[u8; 32]) -> Result> { + self.routers + .get(&self.local) + .context("missing local router")? + .drain(&self.local, session_id) + } + } + + fn manifest() -> EpochManifest { + EpochManifest { + provider_id: vec![3; 32], + epoch: 1, + threshold: 2, + previous_manifest_hash: vec![], + members: ["kms-1", "kms-2", "kms-3"] + .into_iter() + .map(|node_id| EpochMember { + node_id: node_id.into(), + endpoint: format!("https://{node_id}:8443/prpc"), + attestation_pubkey: vec![1; 32], + share_commitment: vec![2; 33], + }) + .collect(), + } + } + + fn five_member_manifest() -> EpochManifest { + EpochManifest { + provider_id: vec![3; 32], + epoch: 1, + threshold: 3, + previous_manifest_hash: vec![], + members: (1..=5) + .map(|index| EpochMember { + node_id: format!("kms-{index}"), + endpoint: format!("https://kms-{index}:8443/prpc"), + attestation_pubkey: vec![index as u8; 32], + share_commitment: vec![2; 33], + }) + .collect(), + } + } + + async fn dkg_party( + index: u16, + transport: MemoryTransport, + context: DriverContext, + eid: [u8; 32], + ) -> Result { + let mut rng = OsRng; + let state = cggmp21::keygen::(ExecutionId::new(&eid), index, 3) + .set_threshold(2) + .into_state_machine(&mut rng); + drive_state_machine(state, &transport, context) + .await? + .map_err(Into::into) + } + + #[tokio::test] + async fn drives_real_cggmp_dkg_over_authenticated_envelopes() { + let manifest = manifest(); + let routers: Arc>> = Arc::new( + manifest + .members + .iter() + .map(|member| { + ( + member.node_id.clone(), + Arc::new( + SessionRouter::new(manifest.clone(), 8, Duration::from_secs(30)) + .unwrap(), + ), + ) + }) + .collect(), + ); + let participants: Vec<_> = manifest + .members + .iter() + .map(|member| member.node_id.clone()) + .collect(); + let session_id = [7; 32]; + let request_hash = [8; 32]; + let expires_at = std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .unwrap() + .as_secs() + + 20; + let eid = execution_id("test-cluster", 1, CggmpCurve::K256, &session_id); + let make = |index: u16| { + let local = participants[usize::from(index)].clone(); + dkg_party( + index, + MemoryTransport { + local: local.clone(), + routers: routers.clone(), + }, + DriverContext { + session_id, + epoch: 1, + protocol: MpcProtocol::DkgK256, + request_hash, + local_node_id: local, + participants: participants.clone(), + expires_at, + poll_interval: Duration::from_millis(1), + }, + eid, + ) + }; + let (a, b, c) = tokio::join!(make(0), make(1), make(2)); + let shares = [a.unwrap(), b.unwrap(), c.unwrap()]; + let public_key = serde_json::to_vec(&shares[0].shared_public_key()).unwrap(); + assert!(shares.iter().all(|share| { + serde_json::to_vec(&share.shared_public_key()).unwrap() == public_key + })); + } + + #[tokio::test] + async fn drives_real_threshold_signature_over_authenticated_envelopes() { + use cggmp21::{generic_ec::Scalar, security_level::SecurityLevel128, DataToSign}; + use k256::ecdsa::{signature::hazmat::PrehashVerifier, Signature, VerifyingKey}; + use sha2::Digest as _; + + let shares = cggmp21::trusted_dealer::builder::(5) + .set_threshold(Some(3)) + .generate_shares(&mut OsRng) + .unwrap(); + let manifest = five_member_manifest(); + let routers: Arc>> = Arc::new( + manifest + .members + .iter() + .map(|member| { + ( + member.node_id.clone(), + Arc::new( + SessionRouter::new(manifest.clone(), 8, Duration::from_secs(30)) + .unwrap(), + ), + ) + }) + .collect(), + ); + // kms-2 and kms-4 are deliberately offline. + let participants = vec![ + "kms-1".to_string(), + "kms-3".to_string(), + "kms-5".to_string(), + ]; + let session_id = [31; 32]; + let request_hash = [32; 32]; + let expires_at = std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .unwrap() + .as_secs() + + 30; + let digest = [33u8; 32]; + let eid = execution_id("test-cluster", 1, CggmpCurve::K256, &session_id); + let sign = |protocol_index: u16, keygen_index: usize| { + let local = participants[usize::from(protocol_index)].clone(); + let transport = MemoryTransport { + local: local.clone(), + routers: routers.clone(), + }; + let context = DriverContext { + session_id, + epoch: 1, + protocol: MpcProtocol::SignK256, + request_hash, + local_node_id: local, + participants: participants.clone(), + expires_at, + poll_interval: Duration::from_millis(1), + }; + let share = shares[keygen_index].clone(); + async move { + let mut rng = OsRng; + let data = + DataToSign::from_scalar(Scalar::::from_be_bytes_mod_order(digest)); + let keygen_indexes = [0, 2, 4]; + let state = cggmp21::signing( + ExecutionId::new(&eid), + protocol_index, + &keygen_indexes, + &share, + ) + .sign_sync(&mut rng, data); + drive_state_machine(state, &transport, context) + .await + .unwrap() + .unwrap() + } + }; + let (first, second, third) = tokio::join!(sign(0, 0), sign(1, 2), sign(2, 4)); + assert_eq!(first, second); + assert_eq!(first, third); + let mut encoded = [0u8; 64]; + first.write_to_slice(&mut encoded); + let signature = Signature::from_slice(&encoded).unwrap(); + let public_key = + VerifyingKey::from_sec1_bytes(shares[0].shared_public_key().to_bytes(true).as_bytes()) + .unwrap(); + public_key.verify_prehash(&digest, &signature).unwrap(); + + // Reuse the curve-independent auxiliary information to exercise the + // P-256 threshold certificate path without another safe-prime setup. + use cggmp21::supported_curves::Secp256r1; + let p256_cores = cggmp21::trusted_dealer::builder::(5) + .set_threshold(Some(3)) + .generate_core_shares(&mut OsRng) + .unwrap(); + let p256_shares = p256_cores + .into_iter() + .zip(shares.iter()) + .map(|(core, share)| { + let aux = cggmp21::key_share::Valid::validate(share.aux.clone()).unwrap(); + cggmp21::KeyShare::from_parts((core, aux)).unwrap() + }) + .collect::>(); + let raw_public = p256_shares[0].shared_public_key().to_bytes(false); + const P256_SPKI_PREFIX: &[u8] = &[ + 0x30, 0x59, 0x30, 0x13, 0x06, 0x07, 0x2a, 0x86, 0x48, 0xce, 0x3d, 0x02, 0x01, 0x06, + 0x08, 0x2a, 0x86, 0x48, 0xce, 0x3d, 0x03, 0x01, 0x07, 0x03, 0x42, 0x00, + ]; + let mut spki = P256_SPKI_PREFIX.to_vec(); + spki.extend_from_slice(raw_public.as_bytes()); + let public = ra_tls::rcgen::SubjectPublicKeyInfo::from_der(&spki).unwrap(); + let external = ra_tls::cert::prepare_external_self_signed( + ra_tls::cert::CertRequest::builder() + .key(&public) + .subject("Threshold Test Root") + .ca_level(1) + .build(), + ) + .unwrap(); + let p256_digest: [u8; 32] = sha2::Sha256::digest(external.tbs_der()).into(); + let p256_session = [41; 32]; + let p256_expires_at = std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .unwrap() + .as_secs() + + 30; + let p256_eid = execution_id("test-cluster", 1, CggmpCurve::P256, &p256_session); + let p256_sign = |protocol_index: u16, keygen_index: usize| { + let local = participants[usize::from(protocol_index)].clone(); + let context = DriverContext { + session_id: p256_session, + epoch: 1, + protocol: MpcProtocol::SignP256, + request_hash: [42; 32], + local_node_id: local.clone(), + participants: participants.clone(), + expires_at: p256_expires_at, + poll_interval: Duration::from_millis(1), + }; + let transport = MemoryTransport { + local, + routers: routers.clone(), + }; + let share = p256_shares[keygen_index].clone(); + async move { + let mut rng = OsRng; + let data = DataToSign::from_scalar(Scalar::::from_be_bytes_mod_order( + p256_digest, + )); + let state = cggmp21::signing( + ExecutionId::new(&p256_eid), + protocol_index, + &[0, 2, 4], + &share, + ) + .sign_sync(&mut rng, data); + drive_state_machine(state, &transport, context) + .await + .unwrap() + .unwrap() + } + }; + let (p1, p2, p3) = tokio::join!(p256_sign(0, 0), p256_sign(1, 2), p256_sign(2, 4)); + assert_eq!(p1, p2); + assert_eq!(p1, p3); + let mut encoded = [0u8; 64]; + p1.write_to_slice(&mut encoded); + let pem = external.finish(&encoded).unwrap(); + let (_, pem) = x509_parser::pem::parse_x509_pem(pem.as_bytes()).unwrap(); + pem.parse_x509().unwrap().verify_signature(None).unwrap(); + } +} diff --git a/dstack/kms/src/mpc_genesis.rs b/dstack/kms/src/mpc_genesis.rs new file mode 100644 index 000000000..806c6be7b --- /dev/null +++ b/dstack/kms/src/mpc_genesis.rs @@ -0,0 +1,1547 @@ +// SPDX-FileCopyrightText: © 2024-2026 Phala Network +// +// SPDX-License-Identifier: Apache-2.0 + +//! One-time, attested distributed genesis DKG. + +use std::{ + collections::BTreeMap, + os::unix::fs::PermissionsExt, + sync::{ + atomic::{AtomicBool, Ordering}, + Arc, Mutex, + }, + time::Duration, +}; + +use anyhow::{ensure, Context, Result}; +use async_trait::async_trait; +use cggmp21::{ + generic_ec::Scalar, + key_share::AnyKeyShare as _, + supported_curves::{Secp256k1, Secp256r1}, + DataToSign, ExecutionId, KeyShare, +}; +use dstack_kms_rpc::{ + mpc_genesis_client::MpcGenesisClient, + mpc_genesis_server::{MpcGenesisRpc, MpcGenesisServer}, + MpcGenesisFinalizeRequest, MpcGenesisStartRequest, MpcGenesisStartResponse, MpcPushRequest, +}; +use ra_rpc::{ + client::{RaClient, RaClientConfig}, + CallContext, RpcCall, +}; +use ra_tls::attestation::AttestationVerifier; +use ra_tls::{ + cert::{ + prepare_external_certificate, prepare_external_self_signed, CertRequest, + ExternalCertificate, + }, + rcgen::SubjectPublicKeyInfo, +}; +use rand::{rngs::OsRng, RngCore}; +use serde::{Deserialize, Serialize}; +use sha2::{Digest, Sha256}; + +use crate::{ + cggmp_engine::{ + execution_id, load_share, share_commitment, store_share, validate_share_topology, + CggmpCurve, K256KeyShare, P256KeyShare, + }, + config::KmsConfig, + mpc_driver::{ + drive_state_machine_blocking, BlockingTransport, DriverContext, EnvelopeTransport, + }, + mpc_identity::{ClusterIdentity, EpochManifest, EpochMember, SignedEpochManifest}, + mpc_lifecycle::validate_and_checkpoint, + mpc_session::{MpcEnvelope, MpcProtocol, SessionRouter}, +}; + +const GENESIS_DOMAIN: &[u8] = b"dstack-mpc-genesis-plan-v1"; +const GENESIS_TTL: Duration = Duration::from_secs(30 * 60); + +#[derive(Clone, Debug, Serialize, Deserialize)] +pub(crate) struct GenesisMember { + pub node_id: String, + pub endpoint: String, + #[serde(with = "hex_bytes")] + pub attestation_pubkey: Vec, +} + +#[derive(Clone, Debug, Serialize, Deserialize)] +pub(crate) struct GenesisPlan { + pub protocol_version: u16, + pub cluster_id: String, + pub threshold: u16, + pub coordinator: String, + pub members: Vec, +} + +impl GenesisPlan { + fn validate(&self) -> Result<()> { + ensure!( + self.protocol_version == 1, + "unsupported genesis protocol version" + ); + ensure!(!self.cluster_id.is_empty(), "genesis cluster ID is empty"); + ensure!( + self.threshold >= 2 && usize::from(self.threshold) <= self.members.len(), + "invalid genesis threshold" + ); + ensure!( + self.members + .iter() + .any(|member| member.node_id == self.coordinator), + "genesis coordinator is not a member" + ); + let mut previous: Option<&str> = None; + for member in &self.members { + ensure!( + previous.is_none_or(|value| value < member.node_id.as_str()), + "genesis members must be unique and ordered" + ); + ensure!( + member.endpoint.starts_with("https://"), + "genesis endpoint must use HTTPS" + ); + ensure!( + !member.attestation_pubkey.is_empty(), + "missing genesis attestation key" + ); + previous = Some(&member.node_id); + } + Ok(()) + } + + fn hash(&self) -> Result<[u8; 32]> { + self.validate()?; + let encoded = serde_jcs::to_vec(self).context("failed to encode genesis plan")?; + let mut hash = Sha256::new(); + hash.update((GENESIS_DOMAIN.len() as u32).to_be_bytes()); + hash.update(GENESIS_DOMAIN); + hash.update((encoded.len() as u32).to_be_bytes()); + hash.update(encoded); + Ok(hash.finalize().into()) + } + + fn provisional_manifest(&self) -> Result { + Ok(EpochManifest { + provider_id: self.hash()?.to_vec(), + epoch: 1, + threshold: self.threshold, + previous_manifest_hash: vec![], + members: self + .members + .iter() + .map(|member| EpochMember { + node_id: member.node_id.clone(), + endpoint: member.endpoint.clone(), + attestation_pubkey: member.attestation_pubkey.clone(), + share_commitment: vec![1], + }) + .collect(), + }) + } +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +enum GenesisKind { + AuxiliaryInfo, + DkgP256, + DkgK256, + DkgDerivation, + Commitments, + SignRoot, + SignManifest, + PrepareRpcCertificate, + SignRpcCertificate, +} + +#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] +struct GenesisOperation { + #[serde(with = "hex_bytes")] + session_id: Vec, + kind: GenesisKind, + #[serde(with = "hex_bytes")] + payload: Vec, + expires_at: u64, + #[serde(with = "hex_bytes")] + request_hash: Vec, +} + +#[derive(Serialize)] +struct GenesisOperationPreimage<'a> { + #[serde(with = "hex_bytes")] + session_id: &'a [u8], + kind: GenesisKind, + #[serde(with = "hex_bytes")] + payload: &'a [u8], + expires_at: u64, +} + +impl GenesisOperation { + fn new(kind: GenesisKind, payload: Vec) -> Result { + let mut session = [0u8; 32]; + OsRng.fill_bytes(&mut session); + let mut operation = Self { + session_id: session.to_vec(), + kind, + payload, + expires_at: unix_time()? + .checked_add(GENESIS_TTL.as_secs()) + .context("expiry overflow")?, + request_hash: vec![], + }; + operation.request_hash = operation.hash()?.to_vec(); + Ok(operation) + } + + fn hash(&self) -> Result<[u8; 32]> { + let encoded = serde_jcs::to_vec(&GenesisOperationPreimage { + session_id: &self.session_id, + kind: self.kind, + payload: &self.payload, + expires_at: self.expires_at, + })?; + Ok(Sha256::digest(encoded).into()) + } + + fn validate(&self) -> Result<()> { + ensure!(self.session_id.len() == 32, "invalid genesis session ID"); + ensure!( + self.request_hash == self.hash()?, + "genesis request hash mismatch" + ); + let now = unix_time()?; + ensure!(self.expires_at >= now, "genesis operation expired"); + ensure!( + self.expires_at - now <= GENESIS_TTL.as_secs(), + "genesis operation TTL too large" + ); + Ok(()) + } +} + +#[derive(Default)] +struct GenesisArtifacts { + aux: Option, + p256: Option, + k256: Option, + derivation: Option, +} + +enum GenesisSigningShare { + P256(P256KeyShare), + K256(K256KeyShare), +} + +#[derive(Clone, Debug, Serialize, Deserialize)] +struct GenesisBundle { + identity: ClusterIdentity, + root_ca_pem: String, + signed_manifest: SignedEpochManifest, + rpc_certificates: BTreeMap, +} + +#[derive(Clone, Debug, Serialize, Deserialize)] +struct RpcCertificateSigningPayload { + node_id: String, + root_ca_pem: String, + #[serde(with = "hex_bytes")] + tbs_der: Vec, +} + +#[derive(Serialize, Deserialize)] +struct GenesisJournal { + version: u16, + bundle: GenesisBundle, +} + +struct GenesisTransport { + local_node_id: String, + local_router: Arc, + clients: BTreeMap>, +} + +impl GenesisTransport { + fn new( + plan: &GenesisPlan, + local_node_id: &str, + router: Arc, + cert: String, + key: String, + ca: String, + verifier: Arc, + ) -> Result { + let mut clients = BTreeMap::new(); + for member in &plan.members { + if member.node_id == local_node_id { + continue; + } + let expected_key = member.attestation_pubkey.clone(); + let node_id = member.node_id.clone(); + let client = RaClientConfig::builder() + .remote_uri(member.endpoint.clone()) + .tls_client_cert(cert.clone()) + .tls_client_key(key.clone()) + .tls_ca_cert(ca.clone()) + .tls_built_in_root_certs(false) + .attestation_verifier(verifier.clone()) + .cert_validator(Box::new(move |info| { + let info = info.context("genesis peer has no certificate")?; + ensure!(info.attestation.is_some(), "genesis peer is not attested"); + let (_, certificate) = x509_parser::parse_x509_certificate(&info.cert_der) + .context("invalid genesis peer certificate")?; + ensure!( + certificate.public_key().raw == expected_key, + "wrong genesis peer key for {node_id}" + ); + Ok(()) + })) + .build() + .into_client() + .with_context(|| { + format!("failed to build genesis client for {}", member.node_id) + })?; + clients.insert(member.node_id.clone(), MpcGenesisClient::new(client)); + } + Ok(Self { + local_node_id: local_node_id.into(), + local_router: router, + clients, + }) + } + + async fn start(&self, node: &str, operation: &GenesisOperation) -> Result> { + Ok(self + .clients + .get(node) + .context("missing genesis client")? + .start(MpcGenesisStartRequest { + operation_json: serde_json::to_vec(operation)?, + }) + .await? + .result_json) + } + + async fn wait_until_ready(&self, timeout: Duration) -> Result<()> { + let deadline = tokio::time::Instant::now() + timeout; + loop { + let probes = self.clients.iter().map(|(node, client)| async move { + client + .ping() + .await + .with_context(|| format!("genesis peer {node} is unavailable")) + }); + let results = futures::future::join_all(probes).await; + if results.iter().all(Result::is_ok) { + return Ok(()); + } + ensure!( + tokio::time::Instant::now() < deadline, + "timed out waiting for genesis peers: {}", + results + .into_iter() + .filter_map(Result::err) + .map(|error| error.to_string()) + .collect::>() + .join("; ") + ); + tokio::time::sleep(Duration::from_secs(1)).await; + } + } +} + +#[async_trait] +impl EnvelopeTransport for GenesisTransport { + async fn send(&self, envelope: MpcEnvelope) -> Result<()> { + ensure!( + envelope.sender == self.local_node_id, + "wrong genesis sender" + ); + self.clients + .get(&envelope.recipient) + .context("missing genesis recipient")? + .push(MpcPushRequest { + envelope_json: serde_json::to_vec(&envelope)?, + }) + .await?; + Ok(()) + } + + async fn receive(&self, session_id: &[u8; 32]) -> Result> { + self.local_router.drain(&self.local_node_id, session_id) + } +} + +struct GenesisInner { + config: KmsConfig, + plan: GenesisPlan, + router: Arc, + transport: Arc, + artifacts: Mutex, + verifier: Arc, + finalized: AtomicBool, + finalized_notify: tokio::sync::Notify, + operations: Mutex, Arc>>, +} + +struct GenesisOperationRecord { + operation: GenesisOperation, + result: Mutex, String>>>, + finished: tokio::sync::Notify, +} + +#[derive(Clone)] +pub(crate) struct GenesisState(Arc); + +impl GenesisState { + pub(crate) fn new(config: KmsConfig) -> Result { + ensure!( + !config.mpc.genesis_plan_file.as_os_str().is_empty(), + "missing genesis plan file" + ); + let plan: GenesisPlan = + serde_json::from_slice(&fs_err::read(&config.mpc.genesis_plan_file)?)?; + plan.validate()?; + ensure!( + plan.cluster_id == config.mpc.cluster_id, + "genesis cluster ID mismatch" + ); + ensure!( + plan.members + .iter() + .any(|member| member.node_id == config.mpc.node_id), + "local node absent from genesis plan" + ); + let local = plan + .members + .iter() + .find(|member| member.node_id == config.mpc.node_id) + .unwrap(); + for path in [config.rpc_cert(), config.mpc.client_cert_file.clone()] { + let (_, pem) = x509_parser::pem::parse_x509_pem(&fs_err::read(path)?)?; + ensure!( + pem.parse_x509()?.public_key().raw == local.attestation_pubkey, + "local genesis server/client certificate key differs from plan" + ); + } + let manifest = plan.provisional_manifest()?; + let router = Arc::new(SessionRouter::new(manifest, 32, GENESIS_TTL)?); + let verifier = Arc::new(AttestationVerifier::load(&config.attestation)?); + let ca = fs_err::read_to_string(&config.mpc.genesis_tls_ca_cert) + .context("failed to read genesis TLS CA")?; + let transport = Arc::new(GenesisTransport::new( + &plan, + &config.mpc.node_id, + router.clone(), + fs_err::read_to_string(&config.mpc.client_cert_file)?, + fs_err::read_to_string(&config.mpc.client_key_file)?, + ca, + verifier.clone(), + )?); + Ok(Self(Arc::new(GenesisInner { + config, + plan, + router, + transport, + artifacts: Mutex::new(GenesisArtifacts::default()), + verifier, + finalized: AtomicBool::new(false), + finalized_notify: tokio::sync::Notify::new(), + operations: Mutex::new(BTreeMap::new()), + }))) + } + + pub(crate) fn verifier(&self) -> Arc { + self.0.verifier.clone() + } + + pub(crate) fn is_coordinator(&self) -> bool { + self.0.config.mpc.node_id == self.0.plan.coordinator + } + + pub(crate) async fn wait_finalized(&self) { + loop { + let notified = self.0.finalized_notify.notified(); + if self.0.finalized.load(Ordering::Acquire) { + return; + } + notified.await; + } + } + + fn authenticated_node(&self, key: &[u8]) -> Result<&str> { + self.0 + .plan + .members + .iter() + .find(|member| member.attestation_pubkey == key) + .map(|member| member.node_id.as_str()) + .context("attested peer is not in genesis plan") + } + + fn identity_from_artifacts(&self) -> Result { + let artifacts = self.0.artifacts.lock().expect("genesis mutex poisoned"); + ClusterIdentity::new( + self.0.plan.protocol_version, + self.0.plan.cluster_id.clone(), + artifacts + .p256 + .as_ref() + .context("missing P-256 genesis share")? + .shared_public_key() + .to_bytes(false) + .as_bytes() + .to_vec(), + artifacts + .k256 + .as_ref() + .context("missing K-256 genesis share")? + .shared_public_key() + .to_bytes(true) + .as_bytes() + .to_vec(), + artifacts + .derivation + .as_ref() + .context("missing derivation genesis share")? + .shared_public_key() + .to_bytes(false) + .as_bytes() + .to_vec(), + ) + } + + fn validate_root_tbs(&self, tbs_der: &[u8]) -> Result<()> { + use x509_parser::{certificate::TbsCertificate, prelude::FromDer as _}; + let (remaining, tbs) = TbsCertificate::from_der(tbs_der) + .map_err(|error| anyhow::anyhow!("invalid genesis root TBS: {error}"))?; + ensure!(remaining.is_empty(), "genesis root TBS has trailing bytes"); + ensure!( + tbs.subject() == tbs.issuer(), + "genesis root is not self-issued" + ); + ensure!(tbs.is_ca(), "genesis root is not a CA"); + ensure!( + tbs.signature.algorithm.to_id_string() == "1.2.840.10045.4.3.2", + "genesis root does not use ECDSA/SHA-256" + ); + ensure!( + tbs.public_key().subject_public_key.data.as_ref() + == self.identity_from_artifacts()?.p256_group_pubkey, + "genesis root public key does not match DKG" + ); + Ok(()) + } + + fn validate_genesis_manifest(&self, manifest: &EpochManifest) -> Result<()> { + let identity = self.identity_from_artifacts()?; + ensure!( + manifest.provider_id == identity.provider_id(), + "wrong genesis provider ID" + ); + ensure!( + manifest.epoch == 1 && manifest.previous_manifest_hash.is_empty(), + "invalid genesis epoch chain" + ); + ensure!( + manifest.threshold == self.0.plan.threshold, + "wrong genesis threshold" + ); + ensure!( + manifest.members.len() == self.0.plan.members.len(), + "wrong genesis member count" + ); + for (member, planned) in manifest.members.iter().zip(&self.0.plan.members) { + ensure!( + member.node_id == planned.node_id + && member.endpoint == planned.endpoint + && member.attestation_pubkey == planned.attestation_pubkey, + "genesis member differs from operator plan" + ); + } + let local_index = self + .0 + .plan + .members + .iter() + .position(|member| member.node_id == self.0.config.mpc.node_id) + .context("local genesis member missing")?; + let artifacts = self.0.artifacts.lock().expect("genesis mutex poisoned"); + let expected = local_share_commitment( + artifacts + .p256 + .as_ref() + .context("missing P-256 genesis share")?, + artifacts + .k256 + .as_ref() + .context("missing K-256 genesis share")?, + artifacts + .derivation + .as_ref() + .context("missing derivation genesis share")?, + local_index, + )?; + ensure!( + manifest.members[local_index].share_commitment == expected, + "wrong local genesis share commitment" + ); + Ok(()) + } + + fn prepare_rpc_certificate(&self, root_ca_pem: &str) -> Result> { + let key = + ra_tls::rcgen::KeyPair::from_pem(&fs_err::read_to_string(self.0.config.rpc_key())?) + .context("invalid genesis RPC key")?; + let (_, current_pem) = + x509_parser::pem::parse_x509_pem(&fs_err::read(self.0.config.rpc_cert())?)?; + let current = current_pem.parse_x509()?; + let attestation = ra_tls::attestation::from_cert(¤t)? + .context("genesis RPC certificate lacks attestation")?; + let domain = fs_err::read_to_string(self.0.config.rpc_domain())?; + let names = [domain.clone()]; + let external = prepare_external_certificate( + CertRequest::builder() + .key(&key) + .subject(&domain) + .alt_names(&names) + .special_usage("kms:rpc") + .attestation(&attestation) + .usage_server_auth(true) + .usage_client_auth(true) + .build(), + root_ca_pem, + )?; + Ok(external.tbs_der().to_vec()) + } + + async fn validate_rpc_certificate_tbs( + &self, + payload: &RpcCertificateSigningPayload, + ) -> Result<()> { + use ra_tls::oids::{PHALA_RATLS_ATTESTATION, PHALA_RATLS_CERT_USAGE}; + use x509_parser::{ + certificate::TbsCertificate, der_parser::oid::Oid, prelude::FromDer as _, + }; + let member = self + .0 + .plan + .members + .iter() + .find(|member| member.node_id == payload.node_id) + .context("RPC certificate target is not in genesis plan")?; + let (_, root_pem) = x509_parser::pem::parse_x509_pem(payload.root_ca_pem.as_bytes())?; + let root = root_pem.parse_x509()?; + root.verify_signature(None) + .context("invalid threshold root certificate")?; + ensure!( + root.public_key().subject_public_key.data.as_ref() + == self.identity_from_artifacts()?.p256_group_pubkey, + "RPC certificate root differs from DKG" + ); + let (remaining, tbs) = TbsCertificate::from_der(&payload.tbs_der) + .map_err(|error| anyhow::anyhow!("invalid RPC certificate TBS: {error}"))?; + ensure!( + remaining.is_empty(), + "RPC certificate TBS has trailing bytes" + ); + ensure!( + tbs.issuer() == root.subject(), + "RPC certificate has wrong issuer" + ); + ensure!( + tbs.public_key().raw == member.attestation_pubkey, + "RPC certificate key differs from quote-bound genesis key" + ); + ensure!( + tbs.signature.algorithm.to_id_string() == "1.2.840.10045.4.3.2", + "RPC certificate does not use ECDSA/SHA-256" + ); + let usage_oid = Oid::from(PHALA_RATLS_CERT_USAGE) + .map_err(|error| anyhow::anyhow!("invalid usage OID: {error:?}"))?; + let usage = tbs + .get_extension_unique(&usage_oid)? + .context("RPC certificate lacks usage")?; + ensure!( + yasna::parse_der(usage.value, |reader| reader.read_bytes())? == b"kms:rpc", + "wrong RPC certificate usage" + ); + let attestation_oid = Oid::from(PHALA_RATLS_ATTESTATION) + .map_err(|error| anyhow::anyhow!("invalid attestation OID: {error:?}"))?; + let extension = tbs + .get_extension_unique(&attestation_oid)? + .context("RPC certificate lacks attestation")?; + let encoded = yasna::parse_der(extension.value, |reader| reader.read_bytes())?; + let attestation = ra_tls::attestation::VersionedAttestation::from_bytes(&encoded) + .context("invalid RPC certificate attestation")?; + let public_key = tbs.public_key().raw.to_vec(); + attestation + .into_v1() + .verify_with_ra_pubkey(&public_key, &self.0.verifier) + .await + .context("RPC certificate attestation is not bound to its key")?; + Ok(()) + } + + async fn execute(&self, operation: GenesisOperation, initiator: &str) -> Result> { + operation.validate()?; + ensure!( + initiator == self.0.plan.coordinator, + "only genesis coordinator may start DKG" + ); + let (record, owner) = { + let mut operations = self.0.operations.lock().expect("genesis mutex poisoned"); + if let Some(record) = operations.get(&operation.request_hash) { + ensure!( + record.operation == operation, + "conflicting genesis operation hash" + ); + (record.clone(), false) + } else { + ensure!(operations.len() < 256, "too many genesis operations"); + let record = Arc::new(GenesisOperationRecord { + operation: operation.clone(), + result: Mutex::new(None), + finished: tokio::sync::Notify::new(), + }); + operations.insert(operation.request_hash.clone(), record.clone()); + (record, true) + } + }; + if owner { + let result = self.execute_owned(operation).await; + *record.result.lock().expect("genesis mutex poisoned") = Some( + result + .as_ref() + .cloned() + .map_err(|error| format!("{error:#}")), + ); + record.finished.notify_waiters(); + result + } else { + loop { + let notified = record.finished.notified(); + if let Some(result) = record + .result + .lock() + .expect("genesis mutex poisoned") + .clone() + { + return result.map_err(anyhow::Error::msg); + } + notified.await; + } + } + } + + async fn execute_owned(&self, operation: GenesisOperation) -> Result> { + let local_index: u16 = self + .0 + .plan + .members + .iter() + .position(|member| member.node_id == self.0.config.mpc.node_id) + .context("local genesis index missing")? + .try_into()?; + let participants = self + .0 + .plan + .members + .iter() + .map(|member| member.node_id.clone()) + .collect::>(); + let protocol = match operation.kind { + GenesisKind::AuxiliaryInfo => MpcProtocol::AuxiliaryInfo, + GenesisKind::DkgP256 => MpcProtocol::DkgP256, + GenesisKind::DkgK256 => MpcProtocol::DkgK256, + GenesisKind::DkgDerivation => MpcProtocol::DkgDerivation, + GenesisKind::Commitments => MpcProtocol::SignK256, + GenesisKind::SignRoot => MpcProtocol::SignP256, + GenesisKind::SignManifest => MpcProtocol::SignK256, + GenesisKind::PrepareRpcCertificate | GenesisKind::SignRpcCertificate => { + MpcProtocol::SignP256 + } + }; + let context = DriverContext { + session_id: operation.session_id.as_slice().try_into()?, + epoch: 1, + protocol, + request_hash: operation.request_hash.as_slice().try_into()?, + local_node_id: self.0.config.mpc.node_id.clone(), + participants, + expires_at: operation.expires_at, + poll_interval: Duration::from_millis(20), + }; + let transport = self.0.transport.clone(); + let runtime = tokio::runtime::Handle::current(); + let n: u16 = self.0.plan.members.len().try_into()?; + let threshold = self.0.plan.threshold; + let cluster = self.0.plan.cluster_id.clone(); + let nonce: [u8; 32] = operation.session_id.as_slice().try_into()?; + match operation.kind { + GenesisKind::AuxiliaryInfo => { + let aux = tokio::task::spawn_blocking(move || { + let mut rng = OsRng; + let primes = cggmp21::PregeneratedPrimes::generate(&mut rng); + let eid = execution_id(&cluster, 1, CggmpCurve::K256, &nonce); + let state = + cggmp21::aux_info_gen(ExecutionId::new(&eid), local_index, n, primes) + .into_state_machine(&mut rng); + drive_state_machine_blocking( + state, + &BlockingTransport::new(transport, runtime), + context, + )? + .map_err(|error| anyhow::anyhow!("genesis auxiliary protocol failed: {error}")) + }) + .await + .context("auxiliary worker panicked")??; + self.0.artifacts.lock().expect("genesis mutex poisoned").aux = Some(aux); + Ok(vec![]) + } + GenesisKind::DkgP256 | GenesisKind::DkgDerivation => { + let aux = self + .0 + .artifacts + .lock() + .expect("genesis mutex poisoned") + .aux + .clone() + .context("auxiliary info has not completed")?; + let curve_label = if operation.kind == GenesisKind::DkgP256 { + CggmpCurve::P256 + } else { + CggmpCurve::K256 + }; + let share = tokio::task::spawn_blocking(move || { + let mut rng = OsRng; + let eid = execution_id(&cluster, 1, curve_label, &nonce); + let state = + cggmp21::keygen::(ExecutionId::new(&eid), local_index, n) + .set_threshold(threshold) + .into_state_machine(&mut rng); + let core = drive_state_machine_blocking( + state, + &BlockingTransport::new(transport, runtime), + context, + )??; + KeyShare::from_parts((core, aux)).map_err(|error| { + anyhow::anyhow!("failed to complete P-256 genesis share: {error}") + }) + }) + .await + .context("P-256 DKG worker panicked")??; + let public = share + .shared_public_key() + .to_bytes(false) + .as_bytes() + .to_vec(); + let mut artifacts = self.0.artifacts.lock().expect("genesis mutex poisoned"); + if operation.kind == GenesisKind::DkgP256 { + artifacts.p256 = Some(share); + } else { + artifacts.derivation = Some(share); + } + Ok(public) + } + GenesisKind::DkgK256 => { + let aux = self + .0 + .artifacts + .lock() + .expect("genesis mutex poisoned") + .aux + .clone() + .context("auxiliary info has not completed")?; + let share = tokio::task::spawn_blocking(move || { + let mut rng = OsRng; + let eid = execution_id(&cluster, 1, CggmpCurve::K256, &nonce); + let state = + cggmp21::keygen::(ExecutionId::new(&eid), local_index, n) + .set_threshold(threshold) + .into_state_machine(&mut rng); + let core = drive_state_machine_blocking( + state, + &BlockingTransport::new(transport, runtime), + context, + )??; + KeyShare::from_parts((core, aux)).map_err(|error| { + anyhow::anyhow!("failed to complete K-256 genesis share: {error}") + }) + }) + .await + .context("K-256 DKG worker panicked")??; + let public = share.shared_public_key().to_bytes(true).as_bytes().to_vec(); + self.0 + .artifacts + .lock() + .expect("genesis mutex poisoned") + .k256 = Some(share); + Ok(public) + } + GenesisKind::Commitments => { + let artifacts = self.0.artifacts.lock().expect("genesis mutex poisoned"); + let p256 = artifacts + .p256 + .as_ref() + .context("missing P-256 genesis share")?; + let k256 = artifacts + .k256 + .as_ref() + .context("missing K-256 genesis share")?; + let derivation = artifacts + .derivation + .as_ref() + .context("missing derivation genesis share")?; + Ok( + local_share_commitment(p256, k256, derivation, usize::from(local_index))? + .to_vec(), + ) + } + GenesisKind::PrepareRpcCertificate => { + let root_ca_pem = std::str::from_utf8(&operation.payload) + .context("invalid genesis root PEM payload")?; + self.prepare_rpc_certificate(root_ca_pem) + } + GenesisKind::SignRoot | GenesisKind::SignManifest | GenesisKind::SignRpcCertificate => { + ensure!( + operation.payload.len() <= 64 * 1024, + "genesis signing payload is too large" + ); + let digest: [u8; 32] = match operation.kind { + GenesisKind::SignRoot => { + self.validate_root_tbs(&operation.payload)?; + Sha256::digest(&operation.payload).into() + } + GenesisKind::SignManifest => { + let manifest: EpochManifest = serde_json::from_slice(&operation.payload) + .context("invalid genesis manifest")?; + self.validate_genesis_manifest(&manifest)?; + manifest.manifest_hash()? + } + GenesisKind::SignRpcCertificate => { + let payload: RpcCertificateSigningPayload = + serde_json::from_slice(&operation.payload) + .context("invalid RPC certificate signing payload")?; + self.validate_rpc_certificate_tbs(&payload).await?; + Sha256::digest(&payload.tbs_der).into() + } + _ => unreachable!(), + }; + let keygen_indexes = (0..n).collect::>(); + let local_protocol_index = local_index; + let (curve, share) = { + let artifacts = self.0.artifacts.lock().expect("genesis mutex poisoned"); + match operation.kind { + GenesisKind::SignRoot | GenesisKind::SignRpcCertificate => ( + CggmpCurve::P256, + GenesisSigningShare::P256( + artifacts + .p256 + .clone() + .context("missing P-256 genesis share")?, + ), + ), + GenesisKind::SignManifest => ( + CggmpCurve::K256, + GenesisSigningShare::K256( + artifacts + .k256 + .clone() + .context("missing K-256 genesis share")?, + ), + ), + _ => unreachable!(), + } + }; + tokio::task::spawn_blocking(move || { + let mut rng = OsRng; + let eid = execution_id(&cluster, 1, curve, &nonce); + match share { + GenesisSigningShare::P256(share) => { + let data = DataToSign::from_scalar( + Scalar::::from_be_bytes_mod_order(digest), + ); + let state = cggmp21::signing( + ExecutionId::new(&eid), + local_protocol_index, + &keygen_indexes, + &share, + ) + .sign_sync(&mut rng, data); + let signature = drive_state_machine_blocking( + state, + &BlockingTransport::new(transport, runtime), + context, + )??; + let mut encoded = + vec![0; cggmp21::Signature::::serialized_len()]; + signature.write_to_slice(&mut encoded); + Ok(encoded) + } + GenesisSigningShare::K256(share) => { + let data = DataToSign::from_scalar( + Scalar::::from_be_bytes_mod_order(digest), + ); + let state = cggmp21::signing( + ExecutionId::new(&eid), + local_protocol_index, + &keygen_indexes, + &share, + ) + .sign_sync(&mut rng, data); + let signature = drive_state_machine_blocking( + state, + &BlockingTransport::new(transport, runtime), + context, + )??; + let mut encoded = + vec![0; cggmp21::Signature::::serialized_len()]; + signature.write_to_slice(&mut encoded); + Ok(encoded) + } + } + }) + .await + .context("genesis signing worker panicked")? + } + } + } + + pub(crate) async fn coordinate_dkg(&self) -> Result<()> { + ensure!( + self.0.config.mpc.node_id == self.0.plan.coordinator, + "not genesis coordinator" + ); + self.0.transport.wait_until_ready(GENESIS_TTL).await?; + for kind in [ + GenesisKind::AuxiliaryInfo, + GenesisKind::DkgP256, + GenesisKind::DkgK256, + GenesisKind::DkgDerivation, + ] { + let results = self + .coordinate_operation(GenesisOperation::new(kind, vec![])?) + .await?; + let first = results + .values() + .next() + .context("genesis has no participants")?; + ensure!( + results.values().all(|value| value == first), + "genesis group public keys differ across nodes" + ); + } + let commitments = self + .coordinate_operation(GenesisOperation::new(GenesisKind::Commitments, vec![])?) + .await?; + let identity = self.identity_from_artifacts()?; + let manifest = EpochManifest { + provider_id: identity.provider_id().to_vec(), + epoch: 1, + threshold: self.0.plan.threshold, + previous_manifest_hash: vec![], + members: self + .0 + .plan + .members + .iter() + .map(|member| EpochMember { + node_id: member.node_id.clone(), + endpoint: member.endpoint.clone(), + attestation_pubkey: member.attestation_pubkey.clone(), + share_commitment: commitments + .get(&member.node_id) + .cloned() + .unwrap_or_default(), + }) + .collect(), + }; + manifest.manifest_hash()?; + + let public_key = p256_subject_public_key(&identity.p256_group_pubkey)?; + let external_root = prepare_external_self_signed( + CertRequest::builder() + .key(&public_key) + .org_name("Dstack") + .subject("Dstack Threshold Root CA") + .ca_level(1) + .build(), + )?; + let root_signatures = self + .coordinate_operation(GenesisOperation::new( + GenesisKind::SignRoot, + external_root.tbs_der().to_vec(), + )?) + .await?; + let root_signature = identical_result(&root_signatures, "root signatures")?; + let root_ca_pem = external_root.finish(root_signature)?; + + let prepared_rpc_certificates = self + .coordinate_operation(GenesisOperation::new( + GenesisKind::PrepareRpcCertificate, + root_ca_pem.as_bytes().to_vec(), + )?) + .await?; + let mut rpc_certificates = BTreeMap::new(); + for (node_id, tbs_der) in prepared_rpc_certificates { + let payload = RpcCertificateSigningPayload { + node_id: node_id.clone(), + root_ca_pem: root_ca_pem.clone(), + tbs_der: tbs_der.clone(), + }; + let signatures = self + .coordinate_operation(GenesisOperation::new( + GenesisKind::SignRpcCertificate, + serde_jcs::to_vec(&payload)?, + )?) + .await?; + let signature = identical_result(&signatures, "RPC certificate signatures")?; + let certificate = ExternalCertificate::from_tbs_der(tbs_der)?.finish(signature)?; + rpc_certificates.insert(node_id, certificate); + } + + let manifest_json = serde_jcs::to_vec(&manifest)?; + let manifest_signatures = self + .coordinate_operation(GenesisOperation::new( + GenesisKind::SignManifest, + manifest_json, + )?) + .await?; + let signature = identical_result(&manifest_signatures, "manifest signatures")?.to_vec(); + let bundle = GenesisBundle { + identity, + root_ca_pem, + signed_manifest: SignedEpochManifest { + manifest, + signature, + }, + rpc_certificates, + }; + let bundle_json = serde_jcs::to_vec(&bundle)?; + let finalizations = self + .0 + .plan + .members + .iter() + .filter(|member| member.node_id != self.0.config.mpc.node_id) + .map(|member| async { + self.0 + .transport + .clients + .get(&member.node_id) + .context("missing genesis client")? + .finalize(MpcGenesisFinalizeRequest { + bundle_json: bundle_json.clone(), + }) + .await + .with_context(|| format!("failed to finalize genesis peer {}", member.node_id)) + }); + for result in futures::future::join_all(finalizations).await { + result?; + } + self.finalize_bundle(bundle) + } + + async fn coordinate_operation( + &self, + operation: GenesisOperation, + ) -> Result>> { + let remote = self + .0 + .plan + .members + .iter() + .filter(|member| member.node_id != self.0.config.mpc.node_id) + .map(|member| async { + Ok::<_, anyhow::Error>(( + member.node_id.clone(), + self.0.transport.start(&member.node_id, &operation).await?, + )) + }); + let local_operation = operation.clone(); + let (local, remote) = tokio::join!( + self.execute(local_operation, &self.0.plan.coordinator), + futures::future::join_all(remote), + ); + let mut results = BTreeMap::new(); + results.insert(self.0.config.mpc.node_id.clone(), local?); + for result in remote { + let (node, value) = result?; + results.insert(node, value); + } + Ok(results) + } + + fn finalize_bundle(&self, bundle: GenesisBundle) -> Result<()> { + ensure!( + bundle.identity == self.identity_from_artifacts()?, + "genesis identity differs from local DKG" + ); + bundle.signed_manifest.verify(&bundle.identity)?; + self.validate_genesis_manifest(&bundle.signed_manifest.manifest)?; + let (_, pem) = x509_parser::pem::parse_x509_pem(bundle.root_ca_pem.as_bytes()) + .context("invalid genesis root PEM")?; + let certificate = pem + .parse_x509() + .context("invalid genesis root certificate")?; + ensure!( + certificate.subject() == certificate.issuer(), + "genesis root is not self-signed" + ); + ensure!( + certificate.public_key().subject_public_key.data.as_ref() + == bundle.identity.p256_group_pubkey, + "wrong genesis root public key" + ); + certificate + .verify_signature(None) + .context("invalid genesis root threshold signature")?; + let journal_path = genesis_journal_path(&self.0.config); + safe_write::safe_write( + &journal_path, + &serde_jcs::to_vec(&GenesisJournal { + version: 1, + bundle: bundle.clone(), + })?, + )?; + fs_err::set_permissions(&journal_path, std::fs::Permissions::from_mode(0o600))?; + self.persist_shares()?; + persist_bundle_files(&self.0.config, &bundle)?; + fs_err::remove_file(journal_path)?; + self.0.finalized.store(true, Ordering::Release); + self.0.finalized_notify.notify_waiters(); + Ok(()) + } + + fn persist_shares(&self) -> Result<()> { + let artifacts = self.0.artifacts.lock().expect("genesis mutex poisoned"); + store_share( + &self.0.config.mpc.p256_share_file, + &self.0.plan.cluster_id, + 1, + &self.0.config.mpc.node_id, + CggmpCurve::P256, + artifacts + .p256 + .as_ref() + .context("missing P-256 genesis share")?, + )?; + store_share( + &self.0.config.mpc.k256_share_file, + &self.0.plan.cluster_id, + 1, + &self.0.config.mpc.node_id, + CggmpCurve::K256, + artifacts + .k256 + .as_ref() + .context("missing K-256 genesis share")?, + )?; + store_share( + &self.0.config.mpc.derivation_share_file, + &self.0.plan.cluster_id, + 1, + &self.0.config.mpc.node_id, + CggmpCurve::P256, + artifacts + .derivation + .as_ref() + .context("missing derivation genesis share")?, + )?; + Ok(()) + } +} + +pub(crate) struct GenesisHandler { + state: GenesisState, + peer_key: Vec, +} + +impl RpcCall for GenesisHandler { + type PrpcService = MpcGenesisServer; + fn construct(context: CallContext<'_, GenesisState>) -> Result { + context + .attestation + .context("genesis peer must be attested")?; + Ok(Self { + state: context.state.clone(), + peer_key: context + .remote_public_key + .context("genesis peer has no key")?, + }) + } +} + +impl MpcGenesisRpc for GenesisHandler { + async fn ping(self) -> Result<()> { + self.state.authenticated_node(&self.peer_key)?; + Ok(()) + } + async fn start(self, request: MpcGenesisStartRequest) -> Result { + let initiator = self.state.authenticated_node(&self.peer_key)?.to_string(); + let operation = serde_json::from_slice(&request.operation_json)?; + Ok(MpcGenesisStartResponse { + result_json: self.state.execute(operation, &initiator).await?, + }) + } + async fn push(self, request: MpcPushRequest) -> Result<()> { + let sender = self.state.authenticated_node(&self.peer_key)?.to_string(); + let envelope = serde_json::from_slice(&request.envelope_json)?; + self.state.0.router.push(&sender, envelope) + } + async fn finalize(self, request: MpcGenesisFinalizeRequest) -> Result<()> { + let initiator = self.state.authenticated_node(&self.peer_key)?; + ensure!( + initiator == self.state.0.plan.coordinator, + "only genesis coordinator may finalize" + ); + let bundle = + serde_json::from_slice(&request.bundle_json).context("invalid genesis bundle")?; + self.state.finalize_bundle(bundle) + } +} + +pub(crate) fn rpc_methods() -> &'static [&'static str] { + >::supported_methods() +} + +fn genesis_journal_path(config: &KmsConfig) -> std::path::PathBuf { + config.mpc.identity_file.with_extension("genesis-journal") +} + +fn persist_bundle_files(config: &KmsConfig, bundle: &GenesisBundle) -> Result<()> { + use ra_tls::traits::CertExt as _; + ensure!( + !config.mpc.identity_file.as_os_str().is_empty(), + "missing MPC identity file" + ); + let rpc_pem = bundle + .rpc_certificates + .get(&config.mpc.node_id) + .context("genesis bundle lacks local RPC certificate")?; + let (_, rpc_der) = x509_parser::pem::parse_x509_pem(rpc_pem.as_bytes())?; + let rpc = rpc_der.parse_x509()?; + let member = bundle + .signed_manifest + .manifest + .members + .iter() + .find(|member| member.node_id == config.mpc.node_id) + .context("local node absent from genesis manifest")?; + ensure!( + rpc.public_key().raw == member.attestation_pubkey, + "genesis RPC certificate key mismatch" + ); + ensure!( + rpc.get_special_usage()?.as_deref() == Some("kms:rpc"), + "genesis RPC certificate usage mismatch" + ); + ensure!( + ra_tls::attestation::from_cert(&rpc)?.is_some(), + "genesis RPC certificate lacks attestation" + ); + let (_, root_der) = x509_parser::pem::parse_x509_pem(bundle.root_ca_pem.as_bytes())?; + let root = root_der.parse_x509()?; + rpc.verify_signature(Some(root.public_key())) + .context("invalid threshold RPC certificate signature")?; + ensure!( + !config.mpc.manifest_file.as_os_str().is_empty(), + "missing MPC manifest file" + ); + safe_write::safe_write( + &config.mpc.identity_file, + &serde_jcs::to_vec(&bundle.identity)?, + )?; + safe_write::safe_write(config.root_ca_cert(), bundle.root_ca_pem.as_bytes())?; + safe_write::safe_write(config.rpc_cert(), rpc_pem.as_bytes())?; + // A directory converted from legacy mode must not retain exportable root + // keys after MPC genesis has established the threshold identity. + for legacy_secret in [config.root_ca_key(), config.k256_key()] { + if legacy_secret.exists() { + fs_err::remove_file(legacy_secret)?; + } + } + safe_write::safe_write( + &config.mpc.manifest_file, + &serde_jcs::to_vec(&bundle.signed_manifest)?, + )?; + validate_and_checkpoint( + &config.mpc.checkpoint_file, + &bundle.signed_manifest, + &bundle.identity, + ) +} + +/// Finish a crash-interrupted genesis commit before startup decides whether +/// to enter genesis or normal mode. +pub(crate) fn recover_if_needed(config: &KmsConfig) -> Result { + if config.mpc.identity_file.as_os_str().is_empty() { + return Ok(false); + } + let journal_path = genesis_journal_path(config); + if !journal_path.exists() { + return Ok(false); + } + let metadata = fs_err::symlink_metadata(&journal_path)?; + ensure!( + metadata.file_type().is_file(), + "MPC genesis journal is not a file" + ); + ensure!( + metadata.permissions().mode() & 0o077 == 0, + "MPC genesis journal permissions are too broad" + ); + let journal: GenesisJournal = serde_json::from_slice(&fs_err::read(&journal_path)?) + .context("failed to parse MPC genesis journal")?; + ensure!( + journal.version == 1, + "unsupported MPC genesis journal version" + ); + let bundle = journal.bundle; + bundle.signed_manifest.verify(&bundle.identity)?; + ensure!( + bundle.signed_manifest.manifest.epoch == 1, + "genesis journal contains a non-genesis epoch" + ); + let (_, pem) = x509_parser::pem::parse_x509_pem(bundle.root_ca_pem.as_bytes())?; + let certificate = pem.parse_x509()?; + ensure!( + certificate.public_key().subject_public_key.data.as_ref() + == bundle.identity.p256_group_pubkey, + "genesis journal root key mismatch" + ); + certificate + .verify_signature(None) + .context("invalid journaled genesis root signature")?; + let manifest = &bundle.signed_manifest.manifest; + let index = manifest + .members + .iter() + .position(|member| member.node_id == config.mpc.node_id) + .context("local member absent from journaled genesis")?; + let p: P256KeyShare = load_share( + &config.mpc.p256_share_file, + &bundle.identity.cluster_id, + 1, + &config.mpc.node_id, + CggmpCurve::P256, + )?; + let k: K256KeyShare = load_share( + &config.mpc.k256_share_file, + &bundle.identity.cluster_id, + 1, + &config.mpc.node_id, + CggmpCurve::K256, + )?; + let d: P256KeyShare = load_share( + &config.mpc.derivation_share_file, + &bundle.identity.cluster_id, + 1, + &config.mpc.node_id, + CggmpCurve::P256, + )?; + validate_share_topology(&p, index, manifest.members.len(), manifest.threshold)?; + validate_share_topology(&k, index, manifest.members.len(), manifest.threshold)?; + validate_share_topology(&d, index, manifest.members.len(), manifest.threshold)?; + ensure!( + p.shared_public_key().to_bytes(false).as_bytes() == bundle.identity.p256_group_pubkey + && k.shared_public_key().to_bytes(true).as_bytes() == bundle.identity.k256_group_pubkey + && d.shared_public_key().to_bytes(false).as_bytes() + == bundle.identity.derivation_group_pubkey, + "journaled genesis shares changed stable identity" + ); + let pp = p.core.key_info.public_shares[index].to_bytes(false); + let kp = k.core.key_info.public_shares[index].to_bytes(true); + let dp = d.core.key_info.public_shares[index].to_bytes(false); + ensure!( + manifest.members[index].share_commitment + == share_commitment(pp.as_bytes(), kp.as_bytes(), dp.as_bytes()), + "journaled genesis shares do not match manifest" + ); + persist_bundle_files(config, &bundle)?; + fs_err::remove_file(journal_path)?; + Ok(true) +} + +fn unix_time() -> Result { + Ok(std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH)? + .as_secs()) +} + +fn local_share_commitment( + p256: &P256KeyShare, + k256: &K256KeyShare, + derivation: &P256KeyShare, + index: usize, +) -> Result<[u8; 32]> { + let p256 = p256 + .core + .key_info + .public_shares + .get(index) + .context("missing P-256 public share")? + .to_bytes(false); + let k256 = k256 + .core + .key_info + .public_shares + .get(index) + .context("missing K-256 public share")? + .to_bytes(true); + let derivation = derivation + .core + .key_info + .public_shares + .get(index) + .context("missing derivation public share")? + .to_bytes(false); + Ok(share_commitment( + p256.as_bytes(), + k256.as_bytes(), + derivation.as_bytes(), + )) +} + +fn identical_result<'a>(results: &'a BTreeMap>, label: &str) -> Result<&'a [u8]> { + let first = results + .values() + .next() + .context("genesis operation has no results")?; + ensure!( + results.values().all(|value| value == first), + "genesis {label} differ across nodes" + ); + Ok(first) +} + +fn p256_subject_public_key(raw: &[u8]) -> Result { + ensure!( + raw.len() == 65 && raw[0] == 4, + "invalid P-256 genesis public key" + ); + // SubjectPublicKeyInfo for id-ecPublicKey / prime256v1 followed by the + // uncompressed SEC1 point. + const PREFIX: &[u8] = &[ + 0x30, 0x59, 0x30, 0x13, 0x06, 0x07, 0x2a, 0x86, 0x48, 0xce, 0x3d, 0x02, 0x01, 0x06, 0x08, + 0x2a, 0x86, 0x48, 0xce, 0x3d, 0x03, 0x01, 0x07, 0x03, 0x42, 0x00, + ]; + let mut der = PREFIX.to_vec(); + der.extend_from_slice(raw); + SubjectPublicKeyInfo::from_der(&der).context("failed to construct genesis root public key") +} + +mod hex_bytes { + use serde::{Deserialize, Deserializer, Serializer}; + pub fn serialize(value: &[u8], serializer: S) -> Result { + serializer.serialize_str(&hex::encode(value)) + } + pub fn deserialize<'de, D: Deserializer<'de>>(deserializer: D) -> Result, D::Error> { + let value = String::deserialize(deserializer)?; + hex::decode(value.strip_prefix("0x").unwrap_or(&value)).map_err(serde::de::Error::custom) + } +} diff --git a/dstack/kms/src/mpc_identity.rs b/dstack/kms/src/mpc_identity.rs new file mode 100644 index 000000000..a071998f8 --- /dev/null +++ b/dstack/kms/src/mpc_identity.rs @@ -0,0 +1,343 @@ +// SPDX-FileCopyrightText: © 2024-2026 Phala Network +// +// SPDX-License-Identifier: Apache-2.0 + +//! Stable identities and attestation bindings for an MPC KMS cluster. +//! +//! Membership and epochs deliberately do not enter [`ClusterIdentity::provider_id`]. +//! This lets a cluster reshare its keys without changing the identity pinned by apps. + +use anyhow::{bail, ensure, Context, Result}; +use k256::ecdsa::{signature::hazmat::PrehashVerifier, Signature, VerifyingKey}; +use serde::{Deserialize, Serialize}; +use sha2::{Digest, Sha256}; + +const ID_DOMAIN: &[u8] = b"dstack-mpc-kms-provider-id-v1"; +const NODE_EVIDENCE_DOMAIN: &[u8] = b"dstack-mpc-kms-node-evidence-v1"; +const MANIFEST_DOMAIN: &[u8] = b"dstack-mpc-kms-epoch-manifest-v1"; + +#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] +pub(crate) struct ClusterIdentity { + pub protocol_version: u16, + pub cluster_id: String, + #[serde(with = "hex_bytes")] + pub p256_group_pubkey: Vec, + #[serde(with = "hex_bytes")] + pub k256_group_pubkey: Vec, + #[serde(with = "hex_bytes")] + pub derivation_group_pubkey: Vec, +} + +impl ClusterIdentity { + pub(crate) fn new( + protocol_version: u16, + cluster_id: String, + p256_group_pubkey: Vec, + k256_group_pubkey: Vec, + derivation_group_pubkey: Vec, + ) -> Result { + let identity = Self { + protocol_version, + cluster_id, + p256_group_pubkey, + k256_group_pubkey, + derivation_group_pubkey, + }; + identity.validate()?; + Ok(identity) + } + + pub(crate) fn validate(&self) -> Result<()> { + ensure!( + self.protocol_version == 1, + "unsupported MPC protocol version" + ); + ensure!( + !self.cluster_id.is_empty(), + "MPC cluster_id must not be empty" + ); + ensure!(self.cluster_id.len() <= 128, "MPC cluster_id is too long"); + ensure!( + !self.p256_group_pubkey.is_empty(), + "missing P-256 group public key" + ); + ensure!( + !self.k256_group_pubkey.is_empty(), + "missing K-256 group public key" + ); + ensure!( + !self.derivation_group_pubkey.is_empty(), + "missing derivation group public key" + ); + Ok(()) + } + + /// Stable key-provider ID pinned in app compose and launch measurements. + pub(crate) fn provider_id(&self) -> [u8; 32] { + hash_fields( + ID_DOMAIN, + &[ + &self.protocol_version.to_be_bytes(), + self.cluster_id.as_bytes(), + &self.p256_group_pubkey, + &self.k256_group_pubkey, + &self.derivation_group_pubkey, + ], + ) + } +} + +#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] +pub(crate) struct EpochMember { + pub node_id: String, + pub endpoint: String, + #[serde(with = "hex_bytes")] + pub attestation_pubkey: Vec, + #[serde(with = "hex_bytes")] + pub share_commitment: Vec, +} + +#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] +pub(crate) struct EpochManifest { + #[serde(with = "hex_bytes")] + pub provider_id: Vec, + pub epoch: u64, + pub threshold: u16, + #[serde(with = "hex_bytes")] + pub previous_manifest_hash: Vec, + pub members: Vec, +} + +#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] +pub(crate) struct SignedEpochManifest { + pub manifest: EpochManifest, + #[serde(with = "hex_bytes")] + pub signature: Vec, +} + +impl SignedEpochManifest { + pub(crate) fn verify(&self, identity: &ClusterIdentity) -> Result<[u8; 32]> { + let hash = self.manifest.manifest_hash()?; + ensure!( + self.manifest.provider_id == identity.provider_id(), + "manifest provider ID does not match cluster identity" + ); + let key = VerifyingKey::from_sec1_bytes(&identity.k256_group_pubkey) + .context("invalid manifest verification key")?; + let signature = Signature::from_slice(&self.signature) + .context("invalid epoch manifest signature encoding")?; + key.verify_prehash(&hash, &signature) + .context("invalid epoch manifest threshold signature")?; + Ok(hash) + } +} + +impl EpochManifest { + pub(crate) fn manifest_hash(&self) -> Result<[u8; 32]> { + ensure!(self.provider_id.len() == 32, "provider_id must be 32 bytes"); + ensure!(self.epoch > 0, "epoch must be greater than zero"); + ensure!( + self.previous_manifest_hash.is_empty() || self.previous_manifest_hash.len() == 32, + "previous_manifest_hash must be empty or 32 bytes" + ); + ensure!( + self.threshold > 0 && usize::from(self.threshold) <= self.members.len(), + "threshold must be between one and the member count" + ); + let mut fields: Vec> = vec![ + self.provider_id.clone(), + self.epoch.to_be_bytes().to_vec(), + self.threshold.to_be_bytes().to_vec(), + self.previous_manifest_hash.clone(), + ]; + let mut previous_node_id: Option<&str> = None; + for member in &self.members { + ensure!( + !member.node_id.is_empty(), + "member node_id must not be empty" + ); + ensure!( + member.endpoint.starts_with("https://"), + "member endpoint must use HTTPS" + ); + ensure!( + previous_node_id.is_none_or(|previous| previous < member.node_id.as_str()), + "members must be sorted by unique node_id" + ); + ensure!( + !member.attestation_pubkey.is_empty(), + "missing member attestation public key" + ); + ensure!( + !member.share_commitment.is_empty(), + "missing member share commitment" + ); + fields.push(member.node_id.as_bytes().to_vec()); + fields.push(member.endpoint.as_bytes().to_vec()); + fields.push(member.attestation_pubkey.clone()); + fields.push(member.share_commitment.clone()); + previous_node_id = Some(&member.node_id); + } + let references: Vec<&[u8]> = fields.iter().map(Vec::as_slice).collect(); + Ok(hash_fields(MANIFEST_DOMAIN, &references)) + } + + pub(crate) fn contains_member(&self, node_id: &str) -> bool { + self.members + .binary_search_by_key(&node_id, |member| member.node_id.as_str()) + .is_ok() + } + + pub(crate) fn member_for_attestation_key(&self, public_key: &[u8]) -> Option<&EpochMember> { + self.members + .iter() + .find(|member| member.attestation_pubkey == public_key) + } +} + +#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] +pub(crate) struct NodeEvidence { + #[serde(with = "hex_bytes")] + pub provider_id: Vec, + pub epoch: u64, + #[serde(with = "hex_bytes")] + pub manifest_hash: Vec, + pub node_id: String, + #[serde(with = "hex_bytes")] + pub attestation_pubkey: Vec, + #[serde(with = "hex_bytes")] + pub share_commitment: Vec, +} + +impl NodeEvidence { + /// Value to place in the first 32 bytes of TDX `report_data`. + pub(crate) fn report_data_hash(&self) -> Result<[u8; 32]> { + ensure!(self.provider_id.len() == 32, "provider_id must be 32 bytes"); + ensure!( + self.manifest_hash.len() == 32, + "manifest_hash must be 32 bytes" + ); + ensure!(!self.node_id.is_empty(), "node_id must not be empty"); + ensure!( + !self.attestation_pubkey.is_empty(), + "missing attestation public key" + ); + ensure!( + !self.share_commitment.is_empty(), + "missing share commitment" + ); + Ok(hash_fields( + NODE_EVIDENCE_DOMAIN, + &[ + &self.provider_id, + &self.epoch.to_be_bytes(), + &self.manifest_hash, + self.node_id.as_bytes(), + &self.attestation_pubkey, + &self.share_commitment, + ], + )) + } +} + +pub(crate) fn decode_hex(name: &str, value: &str) -> Result> { + let value = value.strip_prefix("0x").unwrap_or(value); + if value.is_empty() { + bail!("{name} must not be empty"); + } + hex::decode(value).with_context(|| format!("invalid hex in {name}")) +} + +fn hash_fields(domain: &[u8], fields: &[&[u8]]) -> [u8; 32] { + let mut hash = Sha256::new(); + hash.update((domain.len() as u32).to_be_bytes()); + hash.update(domain); + for field in fields { + hash.update((field.len() as u32).to_be_bytes()); + hash.update(field); + } + hash.finalize().into() +} + +mod hex_bytes { + use serde::{Deserialize, Deserializer, Serializer}; + + pub fn serialize(value: &[u8], serializer: S) -> Result + where + S: Serializer, + { + serializer.serialize_str(&hex::encode(value)) + } + + pub fn deserialize<'de, D>(deserializer: D) -> Result, D::Error> + where + D: Deserializer<'de>, + { + let value = String::deserialize(deserializer)?; + hex::decode(value.strip_prefix("0x").unwrap_or(&value)).map_err(serde::de::Error::custom) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn identity() -> ClusterIdentity { + ClusterIdentity::new( + 1, + "production".into(), + vec![2; 33], + vec![3; 33], + vec![4; 32], + ) + .unwrap() + } + + #[test] + fn provider_id_is_deterministic_and_domain_separated() { + let id = identity(); + assert_eq!(id.provider_id(), id.clone().provider_id()); + let mut other = id; + other.cluster_id = "staging".into(); + assert_ne!(other.provider_id(), identity().provider_id()); + } + + #[test] + fn node_evidence_binds_epoch_and_share_commitment() { + let evidence = NodeEvidence { + provider_id: identity().provider_id().to_vec(), + epoch: 7, + manifest_hash: vec![5; 32], + node_id: "kms-1".into(), + attestation_pubkey: vec![6; 32], + share_commitment: vec![7; 33], + }; + let hash = evidence.report_data_hash().unwrap(); + let mut next = evidence; + next.epoch += 1; + assert_ne!(hash, next.report_data_hash().unwrap()); + } + + #[test] + fn manifest_hash_binds_threshold_and_requires_canonical_members() { + let member = |node_id: &str| EpochMember { + node_id: node_id.into(), + endpoint: format!("https://{node_id}:8443/prpc"), + attestation_pubkey: vec![6; 32], + share_commitment: vec![7; 33], + }; + let manifest = EpochManifest { + provider_id: identity().provider_id().to_vec(), + epoch: 1, + threshold: 2, + previous_manifest_hash: vec![], + members: vec![member("kms-1"), member("kms-2"), member("kms-3")], + }; + let hash = manifest.manifest_hash().unwrap(); + let mut other = manifest.clone(); + other.threshold = 3; + assert_ne!(hash, other.manifest_hash().unwrap()); + other.members.swap(0, 1); + assert!(other.manifest_hash().is_err()); + } +} diff --git a/dstack/kms/src/mpc_join.rs b/dstack/kms/src/mpc_join.rs new file mode 100644 index 000000000..9d0cecc80 --- /dev/null +++ b/dstack/kms/src/mpc_join.rs @@ -0,0 +1,1441 @@ +// SPDX-FileCopyrightText: © 2026 Phala Network +// SPDX-License-Identifier: Apache-2.0 + +//! Attested maintenance protocol for adding members without reconstructing a root key. + +use std::{ + collections::BTreeMap, + sync::{ + atomic::{AtomicBool, Ordering}, + Arc, Mutex, + }, + time::Duration, +}; + +use anyhow::{ensure, Context, Result}; +use async_trait::async_trait; +use cggmp21::{ + generic_ec::Scalar, + key_share::{AnyKeyShare as _, DirtyKeyInfo, Valid}, + supported_curves::{Secp256k1, Secp256r1}, + DataToSign, ExecutionId, KeyShare, +}; +use dstack_kms_rpc::{ + mpc_join_client::MpcJoinClient, + mpc_join_server::{MpcJoinRpc, MpcJoinServer}, + MpcGenesisFinalizeRequest, MpcGenesisStartRequest, MpcGenesisStartResponse, MpcPushRequest, +}; +use ra_rpc::{ + client::{RaClient, RaClientConfig}, + CallContext, RpcCall, +}; +use ra_tls::{ + attestation::AttestationVerifier, + cert::{prepare_external_certificate, CertRequest, ExternalCertificate}, +}; +use rand::{rngs::OsRng, RngCore}; +use serde::{Deserialize, Serialize}; +use sha2::{Digest, Sha256}; + +use crate::{ + cggmp_engine::{ + execution_id, load_share, share_commitment, store_share, CggmpCurve, K256KeyShare, + P256KeyShare, + }, + config::KmsConfig, + mpc_driver::{ + drive_state_machine_blocking, BlockingTransport, DriverContext, EnvelopeTransport, + }, + mpc_identity::{ClusterIdentity, EpochManifest, EpochMember, SignedEpochManifest}, + mpc_lifecycle::{ + activate_pending_epoch, initialize_join_checkpoint, pending_share_path, + record_epoch_authorization, EpochPaths, SignedResharePlan, + }, + mpc_reshare::{self, PrivateContribution, PublicContribution}, + mpc_session::{MpcEnvelope, MpcProtocol, SessionRouter}, +}; + +const JOIN_DOMAIN: &[u8] = b"dstack-mpc-join-operation-v1"; +const JOIN_TTL: Duration = Duration::from_secs(30 * 60); + +#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +enum JoinKind { + AuxiliaryInfo, + Reshare, + SignManifest, + PrepareRpcCertificate, + SignRpcCertificate, +} + +#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] +struct JoinOperation { + #[serde(with = "hex_bytes")] + session_id: Vec, + kind: JoinKind, + #[serde(with = "hex_bytes")] + payload: Vec, + expires_at: u64, + #[serde(with = "hex_bytes")] + request_hash: Vec, +} + +impl JoinOperation { + fn new(kind: JoinKind, payload: Vec) -> Result { + let mut session = [0u8; 32]; + OsRng.fill_bytes(&mut session); + let mut value = Self { + session_id: session.to_vec(), + kind, + payload, + expires_at: unix_time()? + JOIN_TTL.as_secs(), + request_hash: vec![], + }; + value.request_hash = value.hash()?.to_vec(); + Ok(value) + } + fn hash(&self) -> Result<[u8; 32]> { + #[derive(Serialize)] + struct Preimage<'a> { + session_id: &'a [u8], + kind: JoinKind, + payload: &'a [u8], + expires_at: u64, + } + let encoded = serde_jcs::to_vec(&Preimage { + session_id: &self.session_id, + kind: self.kind, + payload: &self.payload, + expires_at: self.expires_at, + })?; + let mut hash = Sha256::new(); + hash.update(JOIN_DOMAIN); + hash.update(encoded); + Ok(hash.finalize().into()) + } + fn validate(&self) -> Result<()> { + ensure!( + self.session_id.len() == 32 && self.request_hash == self.hash()?, + "invalid join operation binding" + ); + let now = unix_time()?; + ensure!( + self.expires_at >= now && self.expires_at - now <= JOIN_TTL.as_secs(), + "invalid join operation expiry" + ); + Ok(()) + } +} + +#[derive(Clone, Serialize, Deserialize)] +struct JoinWire { + p256_reference: DirtyKeyInfo, + p256_public: PublicContribution, + p256_private: PrivateContribution, + k256_reference: DirtyKeyInfo, + k256_public: PublicContribution, + k256_private: PrivateContribution, + derivation_reference: DirtyKeyInfo, + derivation_public: PublicContribution, + derivation_private: PrivateContribution, +} + +#[derive(Clone, Serialize, Deserialize)] +struct JoinRpcCertificatePayload { + node_id: String, + #[serde(with = "hex_bytes")] + tbs_der: Vec, +} + +#[derive(Clone, Serialize, Deserialize)] +struct JoinFinalBundle { + signed_manifest: SignedEpochManifest, + rpc_certificates: BTreeMap, +} + +#[derive(Default)] +struct JoinArtifacts { + aux: Option, + pending: Option<(P256KeyShare, K256KeyShare, P256KeyShare)>, +} + +struct JoinTransport { + local_node_id: String, + router: Arc, + clients: BTreeMap>, +} + +impl JoinTransport { + fn new( + plan: &SignedResharePlan, + local: &str, + router: Arc, + cert: String, + key: String, + ca: String, + verifier: Arc, + ) -> Result { + let mut clients = BTreeMap::new(); + for member in &plan.plan.members { + if member.node_id == local { + continue; + } + let expected = member.attestation_pubkey.clone(); + let node = member.node_id.clone(); + let client = RaClientConfig::builder() + .remote_uri(member.endpoint.clone()) + .tls_client_cert(cert.clone()) + .tls_client_key(key.clone()) + .tls_ca_cert(ca.clone()) + .tls_built_in_root_certs(false) + .tls_no_check(true) + .tls_no_check_hostname(true) + .attestation_verifier(verifier.clone()) + .cert_validator(Box::new(move |info| { + let info = info.context("join peer has no certificate")?; + ensure!(info.attestation.is_some(), "join peer is not attested"); + let (_, certificate) = x509_parser::parse_x509_certificate(&info.cert_der) + .context("invalid join peer certificate")?; + ensure!( + certificate.public_key().raw == expected, + "wrong quote-bound key for {node}" + ); + Ok(()) + })) + .build() + .into_client() + .context("failed to build join client")?; + clients.insert(member.node_id.clone(), MpcJoinClient::new(client)); + } + Ok(Self { + local_node_id: local.into(), + router, + clients, + }) + } + async fn start(&self, node: &str, operation: &JoinOperation) -> Result> { + Ok(self + .clients + .get(node) + .context("missing join client")? + .start(MpcGenesisStartRequest { + operation_json: serde_json::to_vec(operation)?, + }) + .await? + .result_json) + } + async fn ready(&self) -> Result<()> { + let deadline = tokio::time::Instant::now() + JOIN_TTL; + loop { + let results = + futures::future::join_all(self.clients.values().map(|client| client.ping())).await; + if results.iter().all(Result::is_ok) { + return Ok(()); + } + ensure!( + tokio::time::Instant::now() < deadline, + "timed out waiting for join peers" + ); + tokio::time::sleep(Duration::from_secs(1)).await; + } + } +} + +#[async_trait] +impl EnvelopeTransport for JoinTransport { + async fn send(&self, envelope: MpcEnvelope) -> Result<()> { + ensure!(envelope.sender == self.local_node_id, "wrong join sender"); + self.clients + .get(&envelope.recipient) + .context("missing join recipient")? + .push(MpcPushRequest { + envelope_json: serde_json::to_vec(&envelope)?, + }) + .await?; + Ok(()) + } + async fn receive(&self, session_id: &[u8; 32]) -> Result> { + self.router.drain(&self.local_node_id, session_id) + } +} + +struct JoinInner { + config: KmsConfig, + identity: ClusterIdentity, + active: SignedEpochManifest, + authorization: SignedResharePlan, + router: Arc, + transport: Arc, + verifier: Arc, + artifacts: Mutex, + old_p256: Option, + old_k256: Option, + old_derivation: Option, + finalized: tokio::sync::Notify, + is_finalized: AtomicBool, + operations: Mutex, Arc>>, +} + +struct JoinOperationRecord { + operation: JoinOperation, + result: Mutex, String>>>, + finished: tokio::sync::Notify, +} + +#[derive(Clone)] +pub(crate) struct JoinState(Arc); + +impl JoinState { + pub(crate) fn new(config: KmsConfig) -> Result { + let identity: ClusterIdentity = + serde_json::from_slice(&fs_err::read(&config.mpc.identity_file)?)?; + let active: SignedEpochManifest = + serde_json::from_slice(&fs_err::read(&config.mpc.manifest_file)?)?; + active.verify(&identity)?; + let authorization: SignedResharePlan = + serde_json::from_slice(&fs_err::read(&config.mpc.join_authorization_file)?)?; + authorization.verify(&identity, &active.manifest)?; + ensure!( + authorization + .plan + .members + .iter() + .any(|member| member.node_id == config.mpc.node_id), + "local node is absent from join target" + ); + let local = authorization + .plan + .members + .iter() + .find(|member| member.node_id == config.mpc.node_id) + .unwrap(); + for path in [config.rpc_cert(), config.mpc.client_cert_file.clone()] { + let (_, pem) = x509_parser::pem::parse_x509_pem(&fs_err::read(path)?)?; + ensure!( + pem.parse_x509()?.public_key().raw == local.attestation_pubkey, + "local join server/client certificate key differs from authorization" + ); + } + initialize_join_checkpoint( + &config.mpc.checkpoint_file, + &active, + &authorization, + &identity, + )?; + let provisional = EpochManifest { + provider_id: identity.provider_id().to_vec(), + epoch: authorization.plan.epoch, + threshold: authorization.plan.threshold, + previous_manifest_hash: authorization.plan.previous_manifest_hash.clone(), + members: authorization + .plan + .members + .iter() + .map(|member| EpochMember { + node_id: member.node_id.clone(), + endpoint: member.endpoint.clone(), + attestation_pubkey: member.attestation_pubkey.clone(), + share_commitment: vec![1], + }) + .collect(), + }; + let router = Arc::new(SessionRouter::new(provisional, 32, JOIN_TTL)?); + let verifier = Arc::new(AttestationVerifier::load(&config.attestation)?); + let ca = fs_err::read_to_string(config.root_ca_cert())?; + let transport = Arc::new(JoinTransport::new( + &authorization, + &config.mpc.node_id, + router.clone(), + fs_err::read_to_string(&config.mpc.client_cert_file)?, + fs_err::read_to_string(&config.mpc.client_key_file)?, + ca, + verifier.clone(), + )?); + let old_index = active + .manifest + .members + .iter() + .position(|member| member.node_id == config.mpc.node_id); + let load_old = |curve, path: &std::path::Path| -> Result { + match curve { + CggmpCurve::P256 => Ok(serde_json::to_value(load_share::( + path, + &identity.cluster_id, + active.manifest.epoch, + &config.mpc.node_id, + curve, + )?)?), + CggmpCurve::K256 => Ok(serde_json::to_value(load_share::( + path, + &identity.cluster_id, + active.manifest.epoch, + &config.mpc.node_id, + curve, + )?)?), + } + }; + let (old_p256, old_k256, old_derivation) = if old_index.is_some() { + ( + Some(serde_json::from_value(load_old( + CggmpCurve::P256, + &config.mpc.p256_share_file, + )?)?), + Some(serde_json::from_value(load_old( + CggmpCurve::K256, + &config.mpc.k256_share_file, + )?)?), + Some(serde_json::from_value(load_old( + CggmpCurve::P256, + &config.mpc.derivation_share_file, + )?)?), + ) + } else { + (None, None, None) + }; + Ok(Self(Arc::new(JoinInner { + config, + identity, + active, + authorization, + router, + transport, + verifier, + artifacts: Mutex::new(JoinArtifacts::default()), + old_p256, + old_k256, + old_derivation, + finalized: tokio::sync::Notify::new(), + is_finalized: AtomicBool::new(false), + operations: Mutex::new(BTreeMap::new()), + }))) + } + pub(crate) fn verifier(&self) -> Arc { + self.0.verifier.clone() + } + pub(crate) fn is_coordinator(&self) -> bool { + self.0 + .authorization + .plan + .dealers + .first() + .is_some_and(|node| node == &self.0.config.mpc.node_id) + } + pub(crate) async fn wait_finalized(&self) { + loop { + let notified = self.0.finalized.notified(); + if self.0.is_finalized.load(Ordering::Acquire) { + return; + } + notified.await; + } + } + fn authenticated_node(&self, key: &[u8]) -> Result<&str> { + self.0 + .authorization + .plan + .members + .iter() + .find(|member| member.attestation_pubkey == key) + .map(|member| member.node_id.as_str()) + .context("peer is absent from authorized join") + } + + async fn execute(&self, operation: JoinOperation, initiator: &str) -> Result> { + operation.validate()?; + ensure!( + self.0 + .authorization + .plan + .dealers + .first() + .is_some_and(|node| node == initiator), + "only authorized join coordinator may start" + ); + let (record, owner) = { + let mut records = self.0.operations.lock().expect("join mutex poisoned"); + if let Some(record) = records.get(&operation.request_hash) { + ensure!( + record.operation == operation, + "conflicting join operation hash" + ); + (record.clone(), false) + } else { + ensure!(records.len() < 256, "too many join operations"); + let record = Arc::new(JoinOperationRecord { + operation: operation.clone(), + result: Mutex::new(None), + finished: tokio::sync::Notify::new(), + }); + records.insert(operation.request_hash.clone(), record.clone()); + (record, true) + } + }; + if owner { + let result = self.execute_owned(operation).await; + *record.result.lock().expect("join mutex poisoned") = Some( + result + .as_ref() + .cloned() + .map_err(|error| format!("{error:#}")), + ); + record.finished.notify_waiters(); + result + } else { + loop { + let notified = record.finished.notified(); + if let Some(result) = record.result.lock().expect("join mutex poisoned").clone() { + return result.map_err(anyhow::Error::msg); + } + notified.await; + } + } + } + + async fn execute_owned(&self, operation: JoinOperation) -> Result> { + match operation.kind { + JoinKind::AuxiliaryInfo => self.execute_aux(operation).await, + JoinKind::Reshare => self.execute_reshare(operation).await, + JoinKind::SignManifest => self.execute_sign_manifest(operation).await, + JoinKind::PrepareRpcCertificate => self.prepare_rpc_certificate(operation), + JoinKind::SignRpcCertificate => self.execute_sign_rpc_certificate(operation).await, + } + } + + async fn execute_aux(&self, operation: JoinOperation) -> Result> { + ensure!( + operation.payload.is_empty(), + "auxiliary operation payload must be empty" + ); + let members = &self.0.authorization.plan.members; + let local_index: u16 = members + .iter() + .position(|member| member.node_id == self.0.config.mpc.node_id) + .context("local target index missing")? + .try_into()?; + let n: u16 = members.len().try_into()?; + let context = self.context( + &operation, + MpcProtocol::AuxiliaryInfo, + members + .iter() + .map(|member| member.node_id.clone()) + .collect(), + )?; + let transport = self.0.transport.clone(); + let runtime = tokio::runtime::Handle::current(); + let eid = execution_id( + &self.0.identity.cluster_id, + self.0.authorization.plan.epoch, + CggmpCurve::K256, + operation.session_id.as_slice().try_into()?, + ); + let aux = tokio::task::spawn_blocking(move || { + let mut rng = OsRng; + let primes = cggmp21::PregeneratedPrimes::generate(&mut rng); + let state = cggmp21::aux_info_gen(ExecutionId::new(&eid), local_index, n, primes) + .into_state_machine(&mut rng); + drive_state_machine_blocking( + state, + &BlockingTransport::new(transport, runtime), + context, + )? + .map_err(|error| anyhow::anyhow!("join auxiliary protocol failed: {error}")) + }) + .await + .context("join auxiliary worker panicked")??; + self.0.artifacts.lock().expect("join mutex poisoned").aux = Some(aux); + Ok(vec![]) + } + + async fn execute_reshare(&self, operation: JoinOperation) -> Result> { + ensure!( + operation.payload == serde_jcs::to_vec(&self.0.authorization)?, + "wrong join authorization payload" + ); + let plan = &self.0.authorization.plan; + let recipient: u16 = plan + .members + .iter() + .position(|member| member.node_id == self.0.config.mpc.node_id) + .context("local target index missing")? + .try_into()?; + let dealer_indexes = plan + .dealers + .iter() + .map(|node| { + self.0 + .active + .manifest + .members + .iter() + .position(|member| &member.node_id == node) + .context("authorized dealer is not active")? + .try_into() + .context("dealer index overflow") + }) + .collect::>>()?; + let mut received = BTreeMap::::new(); + if let Some(position) = plan + .dealers + .iter() + .position(|node| node == &self.0.config.mpc.node_id) + { + let old_index = dealer_indexes[position]; + let p = self + .0 + .old_p256 + .as_ref() + .context("dealer lacks old P-256 share")?; + let k = self + .0 + .old_k256 + .as_ref() + .context("dealer lacks old K-256 share")?; + let d = self + .0 + .old_derivation + .as_ref() + .context("dealer lacks old derivation share")?; + let p_core = Valid::validate(p.core.clone()) + .map_err(|error| anyhow::anyhow!("invalid old P-256 share: {error}"))?; + let k_core = Valid::validate(k.core.clone()) + .map_err(|error| anyhow::anyhow!("invalid old K-256 share: {error}"))?; + let d_core = Valid::validate(d.core.clone()) + .map_err(|error| anyhow::anyhow!("invalid old derivation share: {error}"))?; + let (pp, ps) = mpc_reshare::create_contribution( + &p_core, + &dealer_indexes, + plan.members.len().try_into()?, + plan.threshold, + &mut OsRng, + )?; + let (kp, ks) = mpc_reshare::create_contribution( + &k_core, + &dealer_indexes, + plan.members.len().try_into()?, + plan.threshold, + &mut OsRng, + )?; + let (dp, ds) = mpc_reshare::create_contribution( + &d_core, + &dealer_indexes, + plan.members.len().try_into()?, + plan.threshold, + &mut OsRng, + )?; + for (index, member) in plan.members.iter().enumerate() { + let wire = JoinWire { + p256_reference: p_core.key_info.clone(), + p256_public: pp.clone(), + p256_private: ps[index].clone(), + k256_reference: k_core.key_info.clone(), + k256_public: kp.clone(), + k256_private: ks[index].clone(), + derivation_reference: d_core.key_info.clone(), + derivation_public: dp.clone(), + derivation_private: ds[index].clone(), + }; + if index == usize::from(recipient) { + received.insert(old_index, wire); + } else { + self.0 + .transport + .send(MpcEnvelope { + session_id: operation.session_id.clone(), + epoch: plan.epoch, + protocol: MpcProtocol::Reshare, + request_hash: operation.request_hash.clone(), + sender: self.0.config.mpc.node_id.clone(), + recipient: member.node_id.clone(), + sequence: u64::from(old_index) + 1, + expires_at: operation.expires_at, + payload: serde_json::to_vec(&wire)?, + broadcast: false, + }) + .await?; + } + } + } + while received.len() < dealer_indexes.len() { + ensure!( + unix_time()? <= operation.expires_at, + "join resharing expired" + ); + for envelope in self + .0 + .transport + .receive(operation.session_id.as_slice().try_into()?) + .await? + { + let index: u16 = self + .0 + .active + .manifest + .members + .iter() + .position(|member| member.node_id == envelope.sender) + .context("join sender is not active")? + .try_into()?; + ensure!( + dealer_indexes.contains(&index) && !received.contains_key(&index), + "invalid duplicate join dealer" + ); + received.insert( + index, + serde_json::from_slice(&envelope.payload) + .context("invalid join contribution")?, + ); + } + if received.len() < dealer_indexes.len() { + tokio::time::sleep(Duration::from_millis(10)).await; + } + } + let wires = dealer_indexes + .iter() + .map(|index| received.get(index).expect("all join dealers received")) + .collect::>(); + for wire in &wires[1..] { + ensure!( + serde_jcs::to_vec(&wire.p256_reference)? + == serde_jcs::to_vec(&wires[0].p256_reference)? + && serde_jcs::to_vec(&wire.k256_reference)? + == serde_jcs::to_vec(&wires[0].k256_reference)? + && serde_jcs::to_vec(&wire.derivation_reference)? + == serde_jcs::to_vec(&wires[0].derivation_reference)?, + "dealers disagree on old public topology" + ); + } + let combine_p = mpc_reshare::verify_and_combine_with_key_info( + &wires[0].p256_reference, + &dealer_indexes, + plan.members.len().try_into()?, + plan.threshold, + recipient, + &wires + .iter() + .map(|w| w.p256_public.clone()) + .collect::>(), + &wires + .iter() + .map(|w| w.p256_private.clone()) + .collect::>(), + )?; + let combine_k = mpc_reshare::verify_and_combine_with_key_info( + &wires[0].k256_reference, + &dealer_indexes, + plan.members.len().try_into()?, + plan.threshold, + recipient, + &wires + .iter() + .map(|w| w.k256_public.clone()) + .collect::>(), + &wires + .iter() + .map(|w| w.k256_private.clone()) + .collect::>(), + )?; + let combine_d = mpc_reshare::verify_and_combine_with_key_info( + &wires[0].derivation_reference, + &dealer_indexes, + plan.members.len().try_into()?, + plan.threshold, + recipient, + &wires + .iter() + .map(|w| w.derivation_public.clone()) + .collect::>(), + &wires + .iter() + .map(|w| w.derivation_private.clone()) + .collect::>(), + )?; + let p_core = combine_p.into_incomplete_share_with_key_info( + recipient, + plan.threshold, + &wires[0].p256_reference, + )?; + let k_core = combine_k.into_incomplete_share_with_key_info( + recipient, + plan.threshold, + &wires[0].k256_reference, + )?; + let d_core = combine_d.into_incomplete_share_with_key_info( + recipient, + plan.threshold, + &wires[0].derivation_reference, + )?; + let aux = self + .0 + .artifacts + .lock() + .expect("join mutex poisoned") + .aux + .clone() + .context("join auxiliary info missing")?; + let p = KeyShare::from_parts((p_core, aux.clone())) + .map_err(|error| anyhow::anyhow!("invalid joined P-256 share: {error}"))?; + let k = KeyShare::from_parts((k_core, aux.clone())) + .map_err(|error| anyhow::anyhow!("invalid joined K-256 share: {error}"))?; + let d = KeyShare::from_parts((d_core, aux)) + .map_err(|error| anyhow::anyhow!("invalid joined derivation share: {error}"))?; + ensure!( + p.shared_public_key().to_bytes(false).as_bytes() == self.0.identity.p256_group_pubkey + && k.shared_public_key().to_bytes(true).as_bytes() + == self.0.identity.k256_group_pubkey + && d.shared_public_key().to_bytes(false).as_bytes() + == self.0.identity.derivation_group_pubkey, + "join changed stable group identity" + ); + let p_pub = p.core.key_info.public_shares[usize::from(recipient)].to_bytes(false); + let k_pub = k.core.key_info.public_shares[usize::from(recipient)].to_bytes(true); + let d_pub = d.core.key_info.public_shares[usize::from(recipient)].to_bytes(false); + let commitment = share_commitment(p_pub.as_bytes(), k_pub.as_bytes(), d_pub.as_bytes()); + let epoch = plan.epoch; + let node = &self.0.config.mpc.node_id; + let cluster = &self.0.identity.cluster_id; + store_share( + &pending_share_path(&self.0.config.mpc.p256_share_file, epoch), + cluster, + epoch, + node, + CggmpCurve::P256, + &p, + )?; + store_share( + &pending_share_path(&self.0.config.mpc.k256_share_file, epoch), + cluster, + epoch, + node, + CggmpCurve::K256, + &k, + )?; + store_share( + &pending_share_path(&self.0.config.mpc.derivation_share_file, epoch), + cluster, + epoch, + node, + CggmpCurve::P256, + &d, + )?; + self.0 + .artifacts + .lock() + .expect("join mutex poisoned") + .pending = Some((p, k, d)); + Ok(commitment.to_vec()) + } + + async fn execute_sign_manifest(&self, operation: JoinOperation) -> Result> { + ensure!( + self.0 + .authorization + .plan + .dealers + .iter() + .any(|node| node == &self.0.config.mpc.node_id), + "only authorized old dealers sign target manifest" + ); + let manifest: EpochManifest = + serde_json::from_slice(&operation.payload).context("invalid target manifest")?; + self.validate_target_manifest(&manifest)?; + let digest = manifest.manifest_hash()?; + record_epoch_authorization( + &self.0.config.mpc.manifest_file, + "manifest", + manifest.epoch, + &digest, + )?; + let dealers = self.0.authorization.plan.dealers.clone(); + let local: u16 = dealers + .iter() + .position(|node| node == &self.0.config.mpc.node_id) + .unwrap() + .try_into()?; + let keygen = dealers + .iter() + .map(|node| { + self.0 + .active + .manifest + .members + .iter() + .position(|member| &member.node_id == node) + .context("dealer is not active")? + .try_into() + .context("dealer index overflow") + }) + .collect::>>()?; + let share = self + .0 + .old_k256 + .clone() + .context("manifest signer lacks old K-256 share")?; + let context = self.context(&operation, MpcProtocol::SignK256, dealers)?; + let transport = self.0.transport.clone(); + let runtime = tokio::runtime::Handle::current(); + let eid = execution_id( + &self.0.identity.cluster_id, + self.0.active.manifest.epoch, + CggmpCurve::K256, + operation.session_id.as_slice().try_into()?, + ); + tokio::task::spawn_blocking(move || { + let mut rng = OsRng; + let data = + DataToSign::from_scalar(Scalar::::from_be_bytes_mod_order(digest)); + let state = cggmp21::signing(ExecutionId::new(&eid), local, &keygen, &share) + .sign_sync(&mut rng, data); + let signature = drive_state_machine_blocking( + state, + &BlockingTransport::new(transport, runtime), + context, + )??; + let mut encoded = vec![0; cggmp21::Signature::::serialized_len()]; + signature.write_to_slice(&mut encoded); + Ok(encoded) + }) + .await + .context("join manifest signing worker panicked")? + } + + fn prepare_rpc_certificate(&self, operation: JoinOperation) -> Result> { + ensure!( + operation.payload.is_empty(), + "RPC certificate preparation payload must be empty" + ); + let key = + ra_tls::rcgen::KeyPair::from_pem(&fs_err::read_to_string(self.0.config.rpc_key())?)?; + let (_, current_pem) = + x509_parser::pem::parse_x509_pem(&fs_err::read(self.0.config.rpc_cert())?)?; + let current = current_pem.parse_x509()?; + let attestation = ra_tls::attestation::from_cert(¤t)? + .context("join RPC certificate lacks attestation")?; + let domain = fs_err::read_to_string(self.0.config.rpc_domain())?; + let names = [domain.clone()]; + Ok(prepare_external_certificate( + CertRequest::builder() + .key(&key) + .subject(&domain) + .alt_names(&names) + .special_usage("kms:rpc") + .attestation(&attestation) + .usage_server_auth(true) + .usage_client_auth(true) + .build(), + &fs_err::read_to_string(self.0.config.root_ca_cert())?, + )? + .tbs_der() + .to_vec()) + } + + async fn validate_rpc_certificate_tbs( + &self, + payload: &JoinRpcCertificatePayload, + ) -> Result<()> { + use ra_tls::oids::{PHALA_RATLS_ATTESTATION, PHALA_RATLS_CERT_USAGE}; + use x509_parser::{ + certificate::TbsCertificate, der_parser::oid::Oid, prelude::FromDer as _, + }; + let member = self + .0 + .authorization + .plan + .members + .iter() + .find(|member| member.node_id == payload.node_id) + .context("RPC certificate target is absent from join authorization")?; + let (_, root_pem) = + x509_parser::pem::parse_x509_pem(&fs_err::read(self.0.config.root_ca_cert())?)?; + let root = root_pem.parse_x509()?; + ensure!( + root.public_key().subject_public_key.data.as_ref() == self.0.identity.p256_group_pubkey, + "join RPC issuer key mismatch" + ); + let (remaining, tbs) = TbsCertificate::from_der(&payload.tbs_der) + .map_err(|error| anyhow::anyhow!("invalid join RPC TBS: {error}"))?; + ensure!( + remaining.is_empty() && tbs.issuer() == root.subject(), + "join RPC certificate issuer mismatch" + ); + ensure!( + tbs.public_key().raw == member.attestation_pubkey, + "join RPC certificate key is not quote-bound authorization key" + ); + let usage_oid = Oid::from(PHALA_RATLS_CERT_USAGE) + .map_err(|error| anyhow::anyhow!("invalid usage OID: {error:?}"))?; + let usage = tbs + .get_extension_unique(&usage_oid)? + .context("join RPC certificate lacks usage")?; + ensure!( + yasna::parse_der(usage.value, |reader| reader.read_bytes())? == b"kms:rpc", + "join RPC certificate usage mismatch" + ); + let attestation_oid = Oid::from(PHALA_RATLS_ATTESTATION) + .map_err(|error| anyhow::anyhow!("invalid attestation OID: {error:?}"))?; + let extension = tbs + .get_extension_unique(&attestation_oid)? + .context("join RPC certificate lacks attestation")?; + let encoded = yasna::parse_der(extension.value, |reader| reader.read_bytes())?; + let attestation = ra_tls::attestation::VersionedAttestation::from_bytes(&encoded) + .context("invalid joined RPC certificate attestation")?; + let public_key = tbs.public_key().raw.to_vec(); + attestation + .into_v1() + .verify_with_ra_pubkey(&public_key, &self.0.verifier) + .await + .context("joined RPC certificate attestation is not bound to its key")?; + Ok(()) + } + + async fn execute_sign_rpc_certificate(&self, operation: JoinOperation) -> Result> { + ensure!( + self.0 + .authorization + .plan + .dealers + .iter() + .any(|node| node == &self.0.config.mpc.node_id), + "only old dealers sign joined RPC certificates" + ); + let payload: JoinRpcCertificatePayload = serde_json::from_slice(&operation.payload)?; + self.validate_rpc_certificate_tbs(&payload).await?; + let digest: [u8; 32] = Sha256::digest(&payload.tbs_der).into(); + let dealers = self.0.authorization.plan.dealers.clone(); + let local: u16 = dealers + .iter() + .position(|node| node == &self.0.config.mpc.node_id) + .unwrap() + .try_into()?; + let keygen = dealers + .iter() + .map(|node| { + self.0 + .active + .manifest + .members + .iter() + .position(|member| &member.node_id == node) + .context("dealer is not active")? + .try_into() + .context("dealer index overflow") + }) + .collect::>>()?; + let share = self + .0 + .old_p256 + .clone() + .context("RPC signer lacks old P-256 share")?; + let context = self.context(&operation, MpcProtocol::SignP256, dealers)?; + let transport = self.0.transport.clone(); + let runtime = tokio::runtime::Handle::current(); + let eid = execution_id( + &self.0.identity.cluster_id, + self.0.active.manifest.epoch, + CggmpCurve::P256, + operation.session_id.as_slice().try_into()?, + ); + tokio::task::spawn_blocking(move || { + let mut rng = OsRng; + let data = + DataToSign::from_scalar(Scalar::::from_be_bytes_mod_order(digest)); + let state = cggmp21::signing(ExecutionId::new(&eid), local, &keygen, &share) + .sign_sync(&mut rng, data); + let signature = drive_state_machine_blocking( + state, + &BlockingTransport::new(transport, runtime), + context, + )??; + let mut encoded = vec![0; cggmp21::Signature::::serialized_len()]; + signature.write_to_slice(&mut encoded); + Ok(encoded) + }) + .await + .context("join RPC signing worker panicked")? + } + + fn validate_target_manifest(&self, manifest: &EpochManifest) -> Result<()> { + let plan = &self.0.authorization.plan; + ensure!( + manifest.provider_id == self.0.identity.provider_id() + && manifest.epoch == plan.epoch + && manifest.threshold == plan.threshold + && manifest.previous_manifest_hash == plan.previous_manifest_hash, + "target manifest differs from join authorization" + ); + ensure!( + manifest.members.len() == plan.members.len(), + "target manifest member count mismatch" + ); + for (actual, expected) in manifest.members.iter().zip(&plan.members) { + ensure!( + actual.node_id == expected.node_id + && actual.endpoint == expected.endpoint + && actual.attestation_pubkey == expected.attestation_pubkey + && actual.share_commitment.len() == 32, + "target manifest member mismatch" + ); + } + let local_index = manifest + .members + .iter() + .position(|member| member.node_id == self.0.config.mpc.node_id) + .context("local member missing from target manifest")?; + if let Some((p, k, d)) = &self + .0 + .artifacts + .lock() + .expect("join mutex poisoned") + .pending + { + let p = p.core.key_info.public_shares[local_index].to_bytes(false); + let k = k.core.key_info.public_shares[local_index].to_bytes(true); + let d = d.core.key_info.public_shares[local_index].to_bytes(false); + ensure!( + manifest.members[local_index].share_commitment + == share_commitment(p.as_bytes(), k.as_bytes(), d.as_bytes()), + "target manifest does not bind local joined shares" + ); + } + Ok(()) + } + + fn context( + &self, + operation: &JoinOperation, + protocol: MpcProtocol, + participants: Vec, + ) -> Result { + Ok(DriverContext { + session_id: operation.session_id.as_slice().try_into()?, + epoch: self.0.authorization.plan.epoch, + protocol, + request_hash: operation.request_hash.as_slice().try_into()?, + local_node_id: self.0.config.mpc.node_id.clone(), + participants, + expires_at: operation.expires_at, + poll_interval: Duration::from_millis(20), + }) + } + + pub(crate) async fn coordinate(&self) -> Result<()> { + ensure!(self.is_coordinator(), "not authorized join coordinator"); + self.0.transport.ready().await?; + let aux = JoinOperation::new(JoinKind::AuxiliaryInfo, vec![])?; + self.run_targets(aux, true).await?; + let reshare = + JoinOperation::new(JoinKind::Reshare, serde_jcs::to_vec(&self.0.authorization)?)?; + let commitments = self.run_targets(reshare, false).await?; + let plan = &self.0.authorization.plan; + let manifest = EpochManifest { + provider_id: self.0.identity.provider_id().to_vec(), + epoch: plan.epoch, + threshold: plan.threshold, + previous_manifest_hash: plan.previous_manifest_hash.clone(), + members: plan + .members + .iter() + .map(|member| EpochMember { + node_id: member.node_id.clone(), + endpoint: member.endpoint.clone(), + attestation_pubkey: member.attestation_pubkey.clone(), + share_commitment: commitments + .get(&member.node_id) + .cloned() + .unwrap_or_default(), + }) + .collect(), + }; + manifest.manifest_hash()?; + let sign = JoinOperation::new(JoinKind::SignManifest, serde_jcs::to_vec(&manifest)?)?; + let signatures = self.run_dealers(sign).await?; + let first = identical(&signatures)?; + let signed = SignedEpochManifest { + manifest, + signature: first.to_vec(), + }; + signed.verify(&self.0.identity)?; + let prepared = self + .run_targets( + JoinOperation::new(JoinKind::PrepareRpcCertificate, vec![])?, + false, + ) + .await?; + let mut rpc_certificates = BTreeMap::new(); + for (node_id, tbs_der) in prepared { + let payload = JoinRpcCertificatePayload { + node_id: node_id.clone(), + tbs_der: tbs_der.clone(), + }; + let signatures = self + .run_dealers(JoinOperation::new( + JoinKind::SignRpcCertificate, + serde_jcs::to_vec(&payload)?, + )?) + .await?; + let signature = identical(&signatures)?; + rpc_certificates.insert( + node_id, + ExternalCertificate::from_tbs_der(tbs_der)?.finish(signature)?, + ); + } + let bundle = JoinFinalBundle { + signed_manifest: signed, + rpc_certificates, + }; + let encoded = serde_jcs::to_vec(&bundle)?; + let calls = plan + .members + .iter() + .filter(|member| member.node_id != self.0.config.mpc.node_id) + .map(|member| async { + self.0 + .transport + .clients + .get(&member.node_id) + .context("missing join client")? + .finalize(MpcGenesisFinalizeRequest { + bundle_json: encoded.clone(), + }) + .await?; + Ok::<_, anyhow::Error>(()) + }); + for result in futures::future::join_all(calls).await { + result?; + } + self.finalize(bundle) + } + + async fn run_targets( + &self, + operation: JoinOperation, + equal: bool, + ) -> Result>> { + let calls = self + .0 + .authorization + .plan + .members + .iter() + .filter(|member| member.node_id != self.0.config.mpc.node_id) + .map(|member| async { + Ok::<_, anyhow::Error>(( + member.node_id.clone(), + self.0.transport.start(&member.node_id, &operation).await?, + )) + }); + let local_operation = operation.clone(); + let (local, remote) = tokio::join!( + self.execute( + local_operation, + self.0.authorization.plan.dealers.first().unwrap() + ), + futures::future::join_all(calls) + ); + let mut out = BTreeMap::from([(self.0.config.mpc.node_id.clone(), local?)]); + for result in remote { + let (node, value) = result?; + out.insert(node, value); + } + if equal { + identical(&out)?; + } + Ok(out) + } + async fn run_dealers(&self, operation: JoinOperation) -> Result>> { + let calls = self + .0 + .authorization + .plan + .dealers + .iter() + .filter(|node| *node != &self.0.config.mpc.node_id) + .map(|node| async { + Ok::<_, anyhow::Error>(( + node.clone(), + self.0.transport.start(node, &operation).await?, + )) + }); + let local_operation = operation.clone(); + let (local, remote) = tokio::join!( + self.execute( + local_operation, + self.0.authorization.plan.dealers.first().unwrap() + ), + futures::future::join_all(calls) + ); + let mut out = BTreeMap::from([(self.0.config.mpc.node_id.clone(), local?)]); + for result in remote { + let (node, value) = result?; + out.insert(node, value); + } + Ok(out) + } + fn finalize(&self, bundle: JoinFinalBundle) -> Result<()> { + use ra_tls::traits::CertExt as _; + let signed = &bundle.signed_manifest; + self.validate_target_manifest(&signed.manifest)?; + signed.verify(&self.0.identity)?; + let rpc_pem = bundle + .rpc_certificates + .get(&self.0.config.mpc.node_id) + .context("join bundle lacks local RPC certificate")?; + let (_, rpc_der) = x509_parser::pem::parse_x509_pem(rpc_pem.as_bytes())?; + let rpc = rpc_der.parse_x509()?; + let member = signed + .manifest + .members + .iter() + .find(|member| member.node_id == self.0.config.mpc.node_id) + .unwrap(); + ensure!( + rpc.public_key().raw == member.attestation_pubkey + && rpc.get_special_usage()?.as_deref() == Some("kms:rpc") + && ra_tls::attestation::from_cert(&rpc)?.is_some(), + "invalid joined RPC certificate identity" + ); + let (_, root_der) = + x509_parser::pem::parse_x509_pem(&fs_err::read(self.0.config.root_ca_cert())?)?; + rpc.verify_signature(Some(root_der.parse_x509()?.public_key())) + .context("invalid joined RPC certificate signature")?; + safe_write::safe_write(self.0.config.rpc_cert(), rpc_pem.as_bytes())?; + activate_pending_epoch( + &EpochPaths { + manifest: &self.0.config.mpc.manifest_file, + checkpoint: &self.0.config.mpc.checkpoint_file, + p256_share: &self.0.config.mpc.p256_share_file, + k256_share: &self.0.config.mpc.k256_share_file, + derivation_share: &self.0.config.mpc.derivation_share_file, + }, + &self.0.identity, + &self.0.active.manifest, + signed, + &self.0.identity.cluster_id, + &self.0.config.mpc.node_id, + )?; + if self.0.config.mpc.join_authorization_file.exists() { + fs_err::remove_file(&self.0.config.mpc.join_authorization_file)?; + } + self.0.is_finalized.store(true, Ordering::Release); + self.0.finalized.notify_waiters(); + Ok(()) + } +} + +pub(crate) struct JoinHandler { + state: JoinState, + peer_key: Vec, +} +impl RpcCall for JoinHandler { + type PrpcService = MpcJoinServer; + fn construct(context: CallContext<'_, JoinState>) -> Result { + context.attestation.context("join peer must be attested")?; + Ok(Self { + state: context.state.clone(), + peer_key: context.remote_public_key.context("join peer has no key")?, + }) + } +} +impl MpcJoinRpc for JoinHandler { + async fn ping(self) -> Result<()> { + self.state.authenticated_node(&self.peer_key)?; + Ok(()) + } + async fn start(self, request: MpcGenesisStartRequest) -> Result { + let initiator = self.state.authenticated_node(&self.peer_key)?.to_string(); + let operation = serde_json::from_slice(&request.operation_json)?; + Ok(MpcGenesisStartResponse { + result_json: self.state.execute(operation, &initiator).await?, + }) + } + async fn push(self, request: MpcPushRequest) -> Result<()> { + let sender = self.state.authenticated_node(&self.peer_key)?.to_string(); + self.state + .0 + .router + .push(&sender, serde_json::from_slice(&request.envelope_json)?) + } + async fn finalize(self, request: MpcGenesisFinalizeRequest) -> Result<()> { + let initiator = self.state.authenticated_node(&self.peer_key)?; + ensure!( + self.state + .0 + .authorization + .plan + .dealers + .first() + .is_some_and(|node| node == initiator), + "only join coordinator may finalize" + ); + self.state + .finalize(serde_json::from_slice(&request.bundle_json)?) + } +} +pub(crate) fn rpc_methods() -> &'static [&'static str] { + >::supported_methods() +} + +/// Finish an interrupted activation before startup chooses maintenance mode. +pub(crate) fn recover_if_needed(config: &KmsConfig) -> Result { + if config.mpc.join_authorization_file.as_os_str().is_empty() + || !config.mpc.join_authorization_file.exists() + || !config.mpc.identity_file.exists() + { + return Ok(false); + } + let identity: ClusterIdentity = + serde_json::from_slice(&fs_err::read(&config.mpc.identity_file)?)?; + let recovered = crate::mpc_lifecycle::recover_pending_activation( + &EpochPaths { + manifest: &config.mpc.manifest_file, + checkpoint: &config.mpc.checkpoint_file, + p256_share: &config.mpc.p256_share_file, + k256_share: &config.mpc.k256_share_file, + derivation_share: &config.mpc.derivation_share_file, + }, + &identity, + &identity.cluster_id, + &config.mpc.node_id, + )?; + let authorization: SignedResharePlan = + serde_json::from_slice(&fs_err::read(&config.mpc.join_authorization_file)?)?; + if config.mpc.manifest_file.exists() { + let installed: SignedEpochManifest = + serde_json::from_slice(&fs_err::read(&config.mpc.manifest_file)?)?; + if installed.manifest.epoch == authorization.plan.epoch + && installed.manifest.previous_manifest_hash + == authorization.plan.previous_manifest_hash + { + installed.verify(&identity)?; + fs_err::remove_file(&config.mpc.join_authorization_file)?; + return Ok(true); + } + } + Ok(recovered) +} +fn identical(results: &BTreeMap>) -> Result<&[u8]> { + let first = results + .values() + .next() + .context("join operation has no results")?; + ensure!( + results.values().all(|value| value == first), + "join participants returned different results" + ); + Ok(first) +} +fn unix_time() -> Result { + Ok(std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH)? + .as_secs()) +} +mod hex_bytes { + use serde::{Deserialize, Deserializer, Serializer}; + pub fn serialize(value: &[u8], serializer: S) -> Result { + serializer.serialize_str(&hex::encode(value)) + } + pub fn deserialize<'de, D: Deserializer<'de>>(deserializer: D) -> Result, D::Error> { + let value = String::deserialize(deserializer)?; + hex::decode(value.strip_prefix("0x").unwrap_or(&value)).map_err(serde::de::Error::custom) + } +} diff --git a/dstack/kms/src/mpc_lifecycle.rs b/dstack/kms/src/mpc_lifecycle.rs new file mode 100644 index 000000000..018ca14c3 --- /dev/null +++ b/dstack/kms/src/mpc_lifecycle.rs @@ -0,0 +1,821 @@ +// SPDX-FileCopyrightText: © 2024-2026 Phala Network +// +// SPDX-License-Identifier: Apache-2.0 + +//! Durable epoch-chain validation for MPC membership lifecycle changes. + +use std::{ + fs::Permissions, + io::Write, + os::unix::fs::PermissionsExt, + path::{Path, PathBuf}, +}; + +use anyhow::{ensure, Context, Result}; +use fs_err as fs; +use k256::ecdsa::{signature::hazmat::PrehashVerifier, Signature, VerifyingKey}; +use serde::{Deserialize, Serialize}; +use sha2::{Digest, Sha256}; + +use crate::{ + cggmp_engine::{ + load_share, share_commitment, validate_share_topology, CggmpCurve, K256KeyShare, + P256KeyShare, + }, + mpc_identity::{ClusterIdentity, EpochManifest, SignedEpochManifest}, +}; + +#[derive(Clone, Debug)] +pub(crate) struct EpochPaths<'a> { + pub manifest: &'a Path, + pub checkpoint: &'a Path, + pub p256_share: &'a Path, + pub k256_share: &'a Path, + pub derivation_share: &'a Path, +} + +#[derive(Serialize, Deserialize)] +struct ActivationJournal { + version: u16, + signed_manifest: SignedEpochManifest, +} + +pub(crate) fn pending_share_path(path: &Path, epoch: u64) -> PathBuf { + path.with_extension(format!("epoch-{epoch}.pending")) +} + +fn activation_journal_path(manifest: &Path) -> PathBuf { + manifest.with_extension("activation-journal") +} + +/// Persist a single consensus proposal per policy and epoch using an atomic +/// create, rejecting concurrent or post-restart forks. +pub(crate) fn record_epoch_authorization( + manifest_path: &Path, + label: &str, + epoch: u64, + digest: &[u8; 32], +) -> Result<()> { + let path = manifest_path.with_extension(format!("epoch-{epoch}.{label}.authorized")); + let validate_existing = || -> Result<()> { + let metadata = fs::symlink_metadata(&path)?; + ensure!( + metadata.file_type().is_file(), + "MPC policy authorization record is not a file" + ); + ensure!( + metadata.permissions().mode() & 0o077 == 0, + "MPC policy authorization permissions are too broad" + ); + ensure!( + fs::read(&path)? == digest, + "conflicting threshold authorization for epoch {epoch}" + ); + Ok(()) + }; + if path.exists() { + return validate_existing(); + } + if let Some(parent) = path.parent() { + fs::create_dir_all(parent)?; + } + match std::fs::OpenOptions::new() + .write(true) + .create_new(true) + .open(&path) + { + Ok(mut file) => { + file.set_permissions(Permissions::from_mode(0o600))?; + file.write_all(digest)?; + file.sync_all()?; + Ok(()) + } + Err(error) if error.kind() == std::io::ErrorKind::AlreadyExists => validate_existing(), + Err(error) => Err(error.into()), + } +} + +#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] +pub(crate) struct ReshareMember { + pub node_id: String, + pub endpoint: String, + #[serde(with = "hex_bytes")] + pub attestation_pubkey: Vec, +} + +#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] +pub(crate) struct ResharePlan { + pub epoch: u64, + #[serde(with = "hex_bytes")] + pub previous_manifest_hash: Vec, + pub threshold: u16, + /// Exact old-epoch threshold that deals the transition. Required when the + /// target contains joining members; subset-only transitions may leave it + /// empty and select a live quorum at execution time. + #[serde(default)] + pub dealers: Vec, + pub members: Vec, +} + +const RESHARE_AUTH_DOMAIN: &[u8] = b"dstack-mpc-reshare-authorization-v1"; + +#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] +pub(crate) struct SignedResharePlan { + pub plan: ResharePlan, + #[serde(with = "hex_bytes")] + pub signature: Vec, +} + +impl ResharePlan { + pub(crate) fn validate_transition( + &self, + active: &crate::mpc_identity::EpochManifest, + ) -> Result<()> { + ensure!( + self.epoch == active.epoch + 1, + "reshare epoch must advance by one" + ); + ensure!( + self.previous_manifest_hash == active.manifest_hash()?, + "reshare plan does not extend the active manifest" + ); + ensure!( + self.threshold >= 2 && usize::from(self.threshold) <= self.members.len(), + "invalid reshare threshold" + ); + let mut previous: Option<&str> = None; + for member in &self.members { + ensure!( + previous.is_none_or(|value| value < member.node_id.as_str()), + "reshare members must be unique and ordered" + ); + if let Some(current) = active + .members + .iter() + .find(|current| current.node_id == member.node_id) + { + ensure!( + current.endpoint == member.endpoint + && current.attestation_pubkey == member.attestation_pubkey, + "existing member identity changed during reshare" + ); + } + ensure!( + !member.node_id.is_empty(), + "reshare member node ID is empty" + ); + ensure!( + member.endpoint.starts_with("https://"), + "reshare member endpoint must use HTTPS" + ); + ensure!( + !member.attestation_pubkey.is_empty(), + "reshare member attestation key is empty" + ); + previous = Some(&member.node_id); + } + if !self.dealers.is_empty() { + ensure!( + self.dealers.len() == usize::from(active.threshold), + "reshare authorization must bind exactly the old threshold of dealers" + ); + let indexes = self + .dealers + .iter() + .map(|dealer| { + active + .members + .binary_search_by_key(&dealer.as_str(), |member| member.node_id.as_str()) + .map_err(|_| anyhow::anyhow!("reshare dealer is not active")) + }) + .collect::>>()?; + ensure!( + indexes.windows(2).all(|pair| pair[0] < pair[1]), + "reshare dealers must be unique and in active manifest order" + ); + } + let has_joiner = self + .members + .iter() + .any(|member| !active.contains_member(&member.node_id)); + ensure!( + !has_joiner || !self.dealers.is_empty(), + "join transition must bind its old-epoch dealers" + ); + Ok(()) + } + + pub(crate) fn authorization_hash(&self) -> Result<[u8; 32]> { + self.validate_encoding()?; + let encoded = serde_jcs::to_vec(self)?; + let mut hash = Sha256::new(); + hash.update((RESHARE_AUTH_DOMAIN.len() as u32).to_be_bytes()); + hash.update(RESHARE_AUTH_DOMAIN); + hash.update((encoded.len() as u32).to_be_bytes()); + hash.update(encoded); + Ok(hash.finalize().into()) + } + + fn validate_encoding(&self) -> Result<()> { + ensure!(self.epoch > 1, "reshare epoch is invalid"); + ensure!( + self.previous_manifest_hash.len() == 32, + "reshare predecessor hash is invalid" + ); + ensure!( + self.threshold >= 2 && usize::from(self.threshold) <= self.members.len(), + "invalid reshare threshold" + ); + Ok(()) + } + + pub(crate) fn validate_subset( + &self, + active: &crate::mpc_identity::EpochManifest, + ) -> Result<()> { + self.validate_transition(active)?; + let mut previous: Option<&str> = None; + for member in &self.members { + ensure!( + previous.is_none_or(|value| value < member.node_id.as_str()), + "reshare members must be unique and ordered" + ); + let active_member = active + .members + .iter() + .find(|active| active.node_id == member.node_id) + .context("joining new members requires the join protocol")?; + ensure!( + active_member.endpoint == member.endpoint + && active_member.attestation_pubkey == member.attestation_pubkey, + "reshare member identity differs from active manifest" + ); + previous = Some(&member.node_id); + } + Ok(()) + } +} + +impl SignedResharePlan { + pub(crate) fn verify(&self, identity: &ClusterIdentity, active: &EpochManifest) -> Result<()> { + self.plan.validate_transition(active)?; + let key = VerifyingKey::from_sec1_bytes(&identity.k256_group_pubkey) + .context("invalid reshare authorization key")?; + let signature = Signature::from_slice(&self.signature) + .context("invalid reshare authorization signature")?; + key.verify_prehash(&self.plan.authorization_hash()?, &signature) + .context("invalid threshold reshare authorization") + } +} + +#[derive(Serialize, Deserialize)] +struct EpochCheckpoint { + version: u16, + #[serde(with = "hex_bytes")] + provider_id: Vec, + epoch: u64, + #[serde(with = "hex_bytes")] + manifest_hash: Vec, +} + +/// Verify a threshold-signed epoch and atomically advance the local checkpoint. +/// Epochs cannot be skipped, forked, or rolled back relative to the checkpoint. +pub(crate) fn validate_and_checkpoint( + path: &Path, + signed: &SignedEpochManifest, + identity: &ClusterIdentity, +) -> Result<()> { + ensure!(!path.as_os_str().is_empty(), "MPC checkpoint path is empty"); + let hash = signed.verify(identity)?; + let provider_id = identity.provider_id(); + let previous = if path.exists() { + let metadata = fs::symlink_metadata(path).context("failed to stat MPC checkpoint")?; + ensure!( + metadata.file_type().is_file(), + "MPC checkpoint is not a file" + ); + ensure!( + metadata.permissions().mode() & 0o077 == 0, + "MPC checkpoint permissions are too broad" + ); + Some( + serde_json::from_slice::( + &fs::read(path).context("failed to read MPC checkpoint")?, + ) + .context("failed to parse MPC checkpoint")?, + ) + } else { + None + }; + match previous { + None => { + ensure!( + signed.manifest.epoch == 1, + "first MPC epoch must be genesis epoch 1" + ); + ensure!( + signed.manifest.previous_manifest_hash.is_empty(), + "genesis manifest must not have a predecessor" + ); + } + Some(previous) => { + ensure!(previous.version == 1, "unsupported MPC checkpoint version"); + ensure!( + previous.provider_id == provider_id, + "MPC checkpoint provider mismatch" + ); + if signed.manifest.epoch == previous.epoch { + ensure!( + previous.manifest_hash == hash, + "conflicting manifest for active epoch" + ); + return Ok(()); + } + ensure!( + signed.manifest.epoch == previous.epoch + 1, + "MPC epoch rollback or skipped transition" + ); + ensure!( + signed.manifest.previous_manifest_hash == previous.manifest_hash, + "MPC manifest does not extend the checkpointed epoch" + ); + } + } + let checkpoint = EpochCheckpoint { + version: 1, + provider_id: provider_id.to_vec(), + epoch: signed.manifest.epoch, + manifest_hash: hash.to_vec(), + }; + let encoded = serde_jcs::to_vec(&checkpoint).context("failed to encode MPC checkpoint")?; + if let Some(parent) = path.parent() { + fs::create_dir_all(parent).context("failed to create MPC checkpoint directory")?; + } + safe_write::safe_write(path, &encoded).context("failed to persist MPC checkpoint")?; + fs::set_permissions(path, Permissions::from_mode(0o600)) + .context("failed to restrict MPC checkpoint permissions")?; + Ok(()) +} + +/// Seed a newly authorized joiner with the currently signed checkpoint. This +/// is only valid when the signed transition explicitly names that exact +/// manifest hash as its predecessor. +pub(crate) fn initialize_join_checkpoint( + path: &Path, + active: &SignedEpochManifest, + authorization: &SignedResharePlan, + identity: &ClusterIdentity, +) -> Result<()> { + let hash = active.verify(identity)?; + authorization.verify(identity, &active.manifest)?; + ensure!( + authorization.plan.previous_manifest_hash == hash, + "join authorization predecessor mismatch" + ); + if path.exists() { + return validate_and_checkpoint(path, active, identity); + } + let checkpoint = EpochCheckpoint { + version: 1, + provider_id: identity.provider_id().to_vec(), + epoch: active.manifest.epoch, + manifest_hash: hash.to_vec(), + }; + if let Some(parent) = path.parent() { + fs::create_dir_all(parent)?; + } + safe_write::safe_write(path, &serde_jcs::to_vec(&checkpoint)?)?; + fs::set_permissions(path, Permissions::from_mode(0o600))?; + Ok(()) +} + +/// Durably install a threshold-authorized next epoch. A journal is written +/// before any active file changes, so startup can finish an interrupted +/// activation instead of observing a mixed manifest/share generation. +pub(crate) fn activate_pending_epoch( + paths: &EpochPaths<'_>, + identity: &ClusterIdentity, + active: &EpochManifest, + signed: &SignedEpochManifest, + cluster_id: &str, + node_id: &str, +) -> Result { + signed.verify(identity)?; + ensure!( + signed.manifest.epoch == active.epoch + 1, + "activation must advance exactly one epoch" + ); + ensure!( + signed.manifest.previous_manifest_hash == active.manifest_hash()?, + "activation does not extend active epoch" + ); + let journal = ActivationJournal { + version: 1, + signed_manifest: signed.clone(), + }; + let journal_path = activation_journal_path(paths.manifest); + safe_write::safe_write(&journal_path, &serde_jcs::to_vec(&journal)?) + .context("failed to persist MPC activation journal")?; + fs::set_permissions(&journal_path, Permissions::from_mode(0o600))?; + finish_activation(paths, identity, signed, cluster_id, node_id)?; + fs::remove_file(&journal_path).context("failed to clear MPC activation journal")?; + for path in [paths.p256_share, paths.k256_share, paths.derivation_share] { + let pending = pending_share_path(path, signed.manifest.epoch); + if pending.exists() { + fs::remove_file(pending)?; + } + } + Ok(signed.manifest.contains_member(node_id)) +} + +/// Complete a journaled activation before normal KMS state reads either the +/// manifest or shares. +pub(crate) fn recover_pending_activation( + paths: &EpochPaths<'_>, + identity: &ClusterIdentity, + cluster_id: &str, + node_id: &str, +) -> Result { + let journal_path = activation_journal_path(paths.manifest); + if !journal_path.exists() { + return Ok(false); + } + let metadata = fs::symlink_metadata(&journal_path)?; + ensure!( + metadata.file_type().is_file(), + "MPC activation journal is not a file" + ); + ensure!( + metadata.permissions().mode() & 0o077 == 0, + "MPC activation journal permissions are too broad" + ); + let journal: ActivationJournal = serde_json::from_slice(&fs::read(&journal_path)?) + .context("failed to parse MPC activation journal")?; + ensure!( + journal.version == 1, + "unsupported MPC activation journal version" + ); + finish_activation( + paths, + identity, + &journal.signed_manifest, + cluster_id, + node_id, + )?; + fs::remove_file(&journal_path)?; + for path in [paths.p256_share, paths.k256_share, paths.derivation_share] { + let pending = pending_share_path(path, journal.signed_manifest.manifest.epoch); + if pending.exists() { + fs::remove_file(pending)?; + } + } + Ok(true) +} + +fn finish_activation( + paths: &EpochPaths<'_>, + identity: &ClusterIdentity, + signed: &SignedEpochManifest, + cluster_id: &str, + node_id: &str, +) -> Result<()> { + signed.verify(identity)?; + if let Some(index) = signed + .manifest + .members + .iter() + .position(|member| member.node_id == node_id) + { + let epoch = signed.manifest.epoch; + let p256_pending = pending_share_path(paths.p256_share, epoch); + let k256_pending = pending_share_path(paths.k256_share, epoch); + let derivation_pending = pending_share_path(paths.derivation_share, epoch); + let p256: P256KeyShare = + load_share(&p256_pending, cluster_id, epoch, node_id, CggmpCurve::P256)?; + let k256: K256KeyShare = + load_share(&k256_pending, cluster_id, epoch, node_id, CggmpCurve::K256)?; + let derivation: P256KeyShare = load_share( + &derivation_pending, + cluster_id, + epoch, + node_id, + CggmpCurve::P256, + )?; + for result in [ + validate_share_topology( + &p256, + index, + signed.manifest.members.len(), + signed.manifest.threshold, + ), + validate_share_topology( + &k256, + index, + signed.manifest.members.len(), + signed.manifest.threshold, + ), + validate_share_topology( + &derivation, + index, + signed.manifest.members.len(), + signed.manifest.threshold, + ), + ] { + result?; + } + use cggmp21::key_share::AnyKeyShare as _; + ensure!( + p256.shared_public_key().to_bytes(false).as_bytes() == identity.p256_group_pubkey, + "activated P-256 group key changed" + ); + ensure!( + k256.shared_public_key().to_bytes(true).as_bytes() == identity.k256_group_pubkey, + "activated K-256 group key changed" + ); + ensure!( + derivation.shared_public_key().to_bytes(false).as_bytes() + == identity.derivation_group_pubkey, + "activated derivation group key changed" + ); + let p = p256.core.key_info.public_shares[index].to_bytes(false); + let k = k256.core.key_info.public_shares[index].to_bytes(true); + let d = derivation.core.key_info.public_shares[index].to_bytes(false); + ensure!( + signed.manifest.members[index].share_commitment + == share_commitment(p.as_bytes(), k.as_bytes(), d.as_bytes()), + "pending shares do not match signed manifest" + ); + for (pending, active) in [ + (&p256_pending, paths.p256_share), + (&k256_pending, paths.k256_share), + (&derivation_pending, paths.derivation_share), + ] { + safe_write::safe_write(active, &fs::read(pending)?)?; + fs::set_permissions(active, Permissions::from_mode(0o600))?; + } + } + safe_write::safe_write(paths.manifest, &serde_jcs::to_vec(signed)?)?; + fs::set_permissions(paths.manifest, Permissions::from_mode(0o600))?; + validate_and_checkpoint(paths.checkpoint, signed, identity) +} + +mod hex_bytes { + use serde::{Deserialize, Deserializer, Serializer}; + + pub fn serialize(value: &[u8], serializer: S) -> Result + where + S: Serializer, + { + serializer.serialize_str(&hex::encode(value)) + } + + pub fn deserialize<'de, D>(deserializer: D) -> Result, D::Error> + where + D: Deserializer<'de>, + { + let value = String::deserialize(deserializer)?; + hex::decode(value.strip_prefix("0x").unwrap_or(&value)).map_err(serde::de::Error::custom) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::mpc_identity::{EpochManifest, EpochMember}; + use k256::ecdsa::{signature::hazmat::PrehashSigner, Signature, SigningKey}; + use std::sync::{Arc, Barrier}; + + fn identity(key: &SigningKey) -> ClusterIdentity { + ClusterIdentity::new( + 1, + "cluster".into(), + vec![2; 65], + key.verifying_key().to_sec1_bytes().to_vec(), + vec![3; 65], + ) + .unwrap() + } + + fn signed( + key: &SigningKey, + identity: &ClusterIdentity, + epoch: u64, + previous: Vec, + ) -> SignedEpochManifest { + let manifest = EpochManifest { + provider_id: identity.provider_id().to_vec(), + epoch, + threshold: 2, + previous_manifest_hash: previous, + members: ["kms-1", "kms-2", "kms-3"] + .into_iter() + .map(|node_id| EpochMember { + node_id: node_id.into(), + endpoint: format!("https://{node_id}/prpc"), + attestation_pubkey: vec![4; 32], + share_commitment: vec![5; 33], + }) + .collect(), + }; + let signature: Signature = key + .sign_prehash(&manifest.manifest_hash().unwrap()) + .unwrap(); + SignedEpochManifest { + manifest, + signature: signature.to_bytes().to_vec(), + } + } + + #[test] + fn epoch_authorization_is_idempotent_and_private() { + let directory = tempfile::tempdir().unwrap(); + let manifest = directory.path().join("manifest.json"); + let digest = [7; 32]; + + record_epoch_authorization(&manifest, "manifest", 2, &digest).unwrap(); + record_epoch_authorization(&manifest, "manifest", 2, &digest).unwrap(); + + let record = manifest.with_extension("epoch-2.manifest.authorized"); + assert_eq!(fs::read(&record).unwrap(), digest); + assert_eq!( + fs::metadata(record).unwrap().permissions().mode() & 0o777, + 0o600 + ); + } + + #[test] + fn epoch_authorization_rejects_a_different_digest() { + let directory = tempfile::tempdir().unwrap(); + let manifest = directory.path().join("manifest.json"); + record_epoch_authorization(&manifest, "manifest", 2, &[7; 32]).unwrap(); + + let error = record_epoch_authorization(&manifest, "manifest", 2, &[8; 32]) + .unwrap_err() + .to_string(); + assert!(error.contains("conflicting threshold authorization")); + } + + #[test] + fn concurrent_epoch_authorization_allows_only_one_fork() { + let directory = tempfile::tempdir().unwrap(); + let manifest = Arc::new(directory.path().join("manifest.json")); + let barrier = Arc::new(Barrier::new(2)); + let handles = [7_u8, 8] + .into_iter() + .map(|byte| { + let manifest = manifest.clone(); + let barrier = barrier.clone(); + std::thread::spawn(move || { + barrier.wait(); + record_epoch_authorization(&manifest, "manifest", 2, &[byte; 32]) + }) + }) + .collect::>(); + let results = handles + .into_iter() + .map(|handle| handle.join().unwrap()) + .collect::>(); + + assert_eq!(results.iter().filter(|result| result.is_ok()).count(), 1); + assert_eq!(results.iter().filter(|result| result.is_err()).count(), 1); + } + + #[test] + fn checkpoint_rejects_rollback_skip_and_fork() { + let key = SigningKey::from_slice(&[7; 32]).unwrap(); + let identity = identity(&key); + let directory = tempfile::tempdir().unwrap(); + let checkpoint = directory.path().join("epoch.json"); + let first = signed(&key, &identity, 1, vec![]); + validate_and_checkpoint(&checkpoint, &first, &identity).unwrap(); + validate_and_checkpoint(&checkpoint, &first, &identity).unwrap(); + let first_hash = first.manifest.manifest_hash().unwrap().to_vec(); + let second = signed(&key, &identity, 2, first_hash); + validate_and_checkpoint(&checkpoint, &second, &identity).unwrap(); + assert!(validate_and_checkpoint(&checkpoint, &first, &identity).is_err()); + let skipped = signed( + &key, + &identity, + 4, + second.manifest.manifest_hash().unwrap().to_vec(), + ); + assert!(validate_and_checkpoint(&checkpoint, &skipped, &identity).is_err()); + let fork = signed(&key, &identity, 3, vec![9; 32]); + assert!(validate_and_checkpoint(&checkpoint, &fork, &identity).is_err()); + } + + #[test] + fn subset_reshare_plan_rejects_unapproved_joiners() { + let key = SigningKey::from_slice(&[7; 32]).unwrap(); + let identity = identity(&key); + let active = signed(&key, &identity, 1, vec![]).manifest; + let plan = ResharePlan { + epoch: 2, + previous_manifest_hash: active.manifest_hash().unwrap().to_vec(), + threshold: 2, + dealers: vec![], + members: active.members[..2] + .iter() + .map(|member| ReshareMember { + node_id: member.node_id.clone(), + endpoint: member.endpoint.clone(), + attestation_pubkey: member.attestation_pubkey.clone(), + }) + .collect(), + }; + plan.validate_subset(&active).unwrap(); + let mut join = plan; + join.members[1].node_id = "kms-new".into(); + assert!(join.validate_subset(&active).is_err()); + } + + #[test] + fn signed_transition_authorizes_quote_bound_joiner() { + let key = SigningKey::from_slice(&[9; 32]).unwrap(); + let identity = identity(&key); + let active = signed(&key, &identity, 1, vec![]).manifest; + let mut members = active + .members + .iter() + .map(|member| ReshareMember { + node_id: member.node_id.clone(), + endpoint: member.endpoint.clone(), + attestation_pubkey: member.attestation_pubkey.clone(), + }) + .collect::>(); + members.push(ReshareMember { + node_id: "kms-4".into(), + endpoint: "https://kms-4/prpc".into(), + attestation_pubkey: vec![8; 32], + }); + let plan = ResharePlan { + epoch: 2, + previous_manifest_hash: active.manifest_hash().unwrap().to_vec(), + threshold: 3, + dealers: active.members[..2] + .iter() + .map(|member| member.node_id.clone()) + .collect(), + members, + }; + let signature: Signature = key + .sign_prehash(&plan.authorization_hash().unwrap()) + .unwrap(); + let signed = SignedResharePlan { + plan, + signature: signature.to_bytes().to_vec(), + }; + signed.verify(&identity, &active).unwrap(); + let mut forged = signed; + forged.plan.members.last_mut().unwrap().attestation_pubkey[0] ^= 1; + assert!(forged.verify(&identity, &active).is_err()); + } + + #[test] + fn startup_recovers_journaled_revocation() { + let key = SigningKey::from_slice(&[8; 32]).unwrap(); + let identity = identity(&key); + let directory = tempfile::tempdir().unwrap(); + let manifest_path = directory.path().join("manifest.json"); + let checkpoint_path = directory.path().join("checkpoint.json"); + let p256 = directory.path().join("p256.share"); + let k256 = directory.path().join("k256.share"); + let derivation = directory.path().join("derivation.share"); + let paths = EpochPaths { + manifest: &manifest_path, + checkpoint: &checkpoint_path, + p256_share: &p256, + k256_share: &k256, + derivation_share: &derivation, + }; + let active = signed(&key, &identity, 1, vec![]); + validate_and_checkpoint(paths.checkpoint, &active, &identity).unwrap(); + let mut next = signed( + &key, + &identity, + 2, + active.manifest.manifest_hash().unwrap().to_vec(), + ); + next.manifest.members.remove(0); + let signature: Signature = key + .sign_prehash(&next.manifest.manifest_hash().unwrap()) + .unwrap(); + next.signature = signature.to_bytes().to_vec(); + let journal_path = activation_journal_path(paths.manifest); + safe_write::safe_write( + &journal_path, + &serde_jcs::to_vec(&ActivationJournal { + version: 1, + signed_manifest: next.clone(), + }) + .unwrap(), + ) + .unwrap(); + fs::set_permissions(&journal_path, Permissions::from_mode(0o600)).unwrap(); + assert!(recover_pending_activation(&paths, &identity, "cluster", "kms-1").unwrap()); + let installed: SignedEpochManifest = + serde_json::from_slice(&fs::read(paths.manifest).unwrap()).unwrap(); + assert_eq!(installed, next); + assert!(!journal_path.exists()); + } +} diff --git a/dstack/kms/src/mpc_operation.rs b/dstack/kms/src/mpc_operation.rs new file mode 100644 index 000000000..d3921d15f --- /dev/null +++ b/dstack/kms/src/mpc_operation.rs @@ -0,0 +1,489 @@ +// SPDX-FileCopyrightText: © 2024-2026 Phala Network +// +// SPDX-License-Identifier: Apache-2.0 + +use std::time::{SystemTime, UNIX_EPOCH}; + +use anyhow::{ensure, Context, Result}; +use serde::{Deserialize, Serialize}; +use sha2::{Digest, Sha256}; +use sha3::Keccak256; + +use crate::mpc_identity::EpochManifest; +use crate::mpc_lifecycle::ResharePlan; + +const REQUEST_DOMAIN: &[u8] = b"dstack-mpc-operation-v1"; +const MAX_SIGN_MESSAGE_BYTES: usize = 1024 * 1024; +const MAX_OPERATION_TTL_SECS: u64 = 300; + +#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] +pub(crate) struct K256SignPayload { + #[serde(with = "hex_bytes")] + pub prefix: Vec, + #[serde(with = "hex_bytes")] + pub app_id: Vec, + pub timestamp: Option, + #[serde(with = "hex_bytes")] + pub message: Vec, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub(crate) enum DerivePurpose { + DiskKey, + EnvKey, + AppK256, + AppCa, +} + +#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] +pub(crate) struct DerivePayload { + pub purpose: DerivePurpose, + #[serde(with = "hex_bytes")] + pub app_id: Vec, + #[serde(default, with = "option_hex_bytes")] + pub instance_id: Option>, + #[serde(default)] + pub counter: u8, +} + +#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] +pub(crate) struct P256CertificatePayload { + #[serde(with = "hex_bytes")] + pub app_id: Vec, + #[serde(with = "hex_bytes")] + pub tbs_der: Vec, +} + +#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] +pub(crate) struct ManifestSignaturePayload { + pub manifest: EpochManifest, +} + +impl ManifestSignaturePayload { + fn validate(&self) -> Result<()> { + self.manifest.manifest_hash().map(|_| ()) + } + + pub(crate) fn digest(&self) -> Result<[u8; 32]> { + self.manifest.manifest_hash() + } +} + +impl P256CertificatePayload { + fn validate(&self) -> Result<()> { + use x509_parser::{certificate::TbsCertificate, prelude::FromDer as _}; + ensure!( + self.app_id.len() == 20, + "certificate app_id must be 20 bytes" + ); + ensure!( + !self.tbs_der.is_empty() && self.tbs_der.len() <= 64 * 1024, + "certificate TBS length is invalid" + ); + let (remaining, _) = TbsCertificate::from_der(&self.tbs_der) + .map_err(|error| anyhow::anyhow!("invalid certificate TBS: {error}"))?; + ensure!(remaining.is_empty(), "certificate TBS has trailing bytes"); + Ok(()) + } + + pub(crate) fn digest(&self) -> [u8; 32] { + Sha256::digest(&self.tbs_der).into() + } +} + +impl DerivePayload { + fn validate(&self) -> Result<()> { + ensure!( + self.app_id.len() == 20, + "MPC derivation app_id must be 20 bytes" + ); + ensure!(self.counter <= 16, "MPC derivation counter is too large"); + match self.purpose { + DerivePurpose::DiskKey => ensure!( + self.instance_id + .as_ref() + .is_some_and(|value| value.len() == 20), + "disk-key derivation requires a 20-byte instance_id" + ), + _ => ensure!( + self.instance_id.is_none(), + "instance_id is only valid for disk-key derivation" + ), + } + Ok(()) + } + + pub(crate) fn input(&self) -> Result> { + let encoded = serde_jcs::to_vec(self).context("failed to encode derivation input")?; + let mut input = b"dstack-mpc-derivation-v1".to_vec(); + input.extend_from_slice(&(encoded.len() as u32).to_be_bytes()); + input.extend_from_slice(&encoded); + Ok(input) + } +} + +impl K256SignPayload { + pub(crate) fn digest(&self) -> [u8; 32] { + let mut digest = Keccak256::new(); + digest.update(&self.prefix); + digest.update(b":"); + digest.update(&self.app_id); + if let Some(timestamp) = self.timestamp { + digest.update(timestamp.to_be_bytes()); + } + digest.update(&self.message); + digest.finalize().into() + } + + fn validate(&self) -> Result<()> { + ensure!( + matches!( + self.prefix.as_slice(), + b"dstack-env-encrypt-pubkey" | b"dstack-kms-issued" + ), + "MPC K-256 signing domain is not allowed" + ); + ensure!( + self.app_id.len() == 20, + "MPC signing app_id must be 20 bytes" + ); + ensure!( + !self.message.is_empty() && self.message.len() <= MAX_SIGN_MESSAGE_BYTES, + "MPC signing message length is invalid" + ); + Ok(()) + } +} + +#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] +#[serde(tag = "kind", rename_all = "snake_case")] +pub(crate) enum MpcOperationPayload { + SignK256(K256SignPayload), + SignManifest(ManifestSignaturePayload), + SignP256Certificate(P256CertificatePayload), + Derive(DerivePayload), + Reshare(ResharePlan), + AuthorizeReshare(ResharePlan), +} + +#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] +pub(crate) struct MpcOperation { + #[serde(with = "hex_bytes")] + pub session_id: Vec, + pub epoch: u64, + pub participants: Vec, + pub expires_at: u64, + pub payload: MpcOperationPayload, + #[serde(with = "hex_bytes")] + pub request_hash: Vec, +} + +#[derive(Serialize)] +struct RequestPreimage<'a> { + #[serde(with = "hex_bytes")] + session_id: &'a [u8], + epoch: u64, + participants: &'a [String], + expires_at: u64, + payload: &'a MpcOperationPayload, +} + +impl MpcOperation { + pub(crate) fn new_k256( + session_id: [u8; 32], + epoch: u64, + participants: Vec, + expires_at: u64, + payload: K256SignPayload, + ) -> Result { + let mut operation = Self { + session_id: session_id.to_vec(), + epoch, + participants, + expires_at, + payload: MpcOperationPayload::SignK256(payload), + request_hash: vec![], + }; + operation.request_hash = operation.compute_request_hash()?.to_vec(); + Ok(operation) + } + + pub(crate) fn new_derivation( + session_id: [u8; 32], + epoch: u64, + participants: Vec, + expires_at: u64, + payload: DerivePayload, + ) -> Result { + let mut operation = Self { + session_id: session_id.to_vec(), + epoch, + participants, + expires_at, + payload: MpcOperationPayload::Derive(payload), + request_hash: vec![], + }; + operation.request_hash = operation.compute_request_hash()?.to_vec(); + Ok(operation) + } + + pub(crate) fn new_p256_certificate( + session_id: [u8; 32], + epoch: u64, + participants: Vec, + expires_at: u64, + payload: P256CertificatePayload, + ) -> Result { + let mut operation = Self { + session_id: session_id.to_vec(), + epoch, + participants, + expires_at, + payload: MpcOperationPayload::SignP256Certificate(payload), + request_hash: vec![], + }; + operation.request_hash = operation.compute_request_hash()?.to_vec(); + Ok(operation) + } + + pub(crate) fn new_manifest_signature( + session_id: [u8; 32], + epoch: u64, + participants: Vec, + expires_at: u64, + payload: ManifestSignaturePayload, + ) -> Result { + let mut operation = Self { + session_id: session_id.to_vec(), + epoch, + participants, + expires_at, + payload: MpcOperationPayload::SignManifest(payload), + request_hash: vec![], + }; + operation.request_hash = operation.compute_request_hash()?.to_vec(); + Ok(operation) + } + + pub(crate) fn new_reshare( + session_id: [u8; 32], + epoch: u64, + participants: Vec, + expires_at: u64, + plan: ResharePlan, + ) -> Result { + let mut operation = Self { + session_id: session_id.to_vec(), + epoch, + participants, + expires_at, + payload: MpcOperationPayload::Reshare(plan), + request_hash: vec![], + }; + operation.request_hash = operation.compute_request_hash()?.to_vec(); + Ok(operation) + } + + pub(crate) fn new_reshare_authorization( + session_id: [u8; 32], + epoch: u64, + participants: Vec, + expires_at: u64, + plan: ResharePlan, + ) -> Result { + let mut operation = Self { + session_id: session_id.to_vec(), + epoch, + participants, + expires_at, + payload: MpcOperationPayload::AuthorizeReshare(plan), + request_hash: vec![], + }; + operation.request_hash = operation.compute_request_hash()?.to_vec(); + Ok(operation) + } + + pub(crate) fn validate(&self, manifest: &EpochManifest, initiator: &str) -> Result<()> { + ensure!( + self.session_id.len() == 32, + "MPC operation session ID must be 32 bytes" + ); + ensure!(self.epoch == manifest.epoch, "MPC operation epoch mismatch"); + ensure!( + self.participants.len() == usize::from(manifest.threshold), + "MPC operation must use exactly threshold participants" + ); + ensure!( + self.participants.iter().any(|node| node == initiator), + "MPC operation initiator is not a participant" + ); + let indexes = self + .participants + .iter() + .map(|node| { + manifest + .members + .binary_search_by_key(&node.as_str(), |member| member.node_id.as_str()) + .map_err(|_| anyhow::anyhow!("MPC operation contains a non-member")) + }) + .collect::>>()?; + ensure!( + indexes.windows(2).all(|pair| pair[0] < pair[1]), + "MPC participants must be unique and in manifest order" + ); + let now = unix_time()?; + ensure!(self.expires_at >= now, "MPC operation expired"); + ensure!( + self.expires_at - now <= MAX_OPERATION_TTL_SECS, + "MPC operation expiration exceeds maximum TTL" + ); + match &self.payload { + MpcOperationPayload::SignK256(payload) => payload.validate()?, + MpcOperationPayload::SignManifest(payload) => { + payload.validate()?; + ensure!( + payload.manifest.provider_id == manifest.provider_id, + "proposed manifest changes provider ID" + ); + let genesis_self_authorization = manifest.epoch == 1 + && manifest.previous_manifest_hash.is_empty() + && payload.manifest == *manifest; + if !genesis_self_authorization { + ensure!( + payload.manifest.epoch == manifest.epoch + 1, + "proposed manifest must advance exactly one epoch" + ); + ensure!( + payload.manifest.previous_manifest_hash == manifest.manifest_hash()?, + "proposed manifest does not extend active manifest" + ); + } + } + MpcOperationPayload::SignP256Certificate(payload) => payload.validate()?, + MpcOperationPayload::Derive(payload) => payload.validate()?, + MpcOperationPayload::Reshare(plan) => plan.validate_subset(manifest)?, + MpcOperationPayload::AuthorizeReshare(plan) => plan.validate_transition(manifest)?, + } + ensure!( + self.request_hash == self.compute_request_hash()?, + "MPC operation request hash mismatch" + ); + Ok(()) + } + + pub(crate) fn compute_request_hash(&self) -> Result<[u8; 32]> { + let encoded = serde_jcs::to_vec(&RequestPreimage { + session_id: &self.session_id, + epoch: self.epoch, + participants: &self.participants, + expires_at: self.expires_at, + payload: &self.payload, + }) + .context("failed to encode MPC operation")?; + let mut digest = Sha256::new(); + digest.update((REQUEST_DOMAIN.len() as u32).to_be_bytes()); + digest.update(REQUEST_DOMAIN); + digest.update((encoded.len() as u32).to_be_bytes()); + digest.update(encoded); + Ok(digest.finalize().into()) + } +} + +fn unix_time() -> Result { + Ok(SystemTime::now() + .duration_since(UNIX_EPOCH) + .context("system time before UNIX epoch")? + .as_secs()) +} + +mod hex_bytes { + use serde::{Deserialize, Deserializer, Serializer}; + + pub fn serialize(value: &[u8], serializer: S) -> Result + where + S: Serializer, + { + serializer.serialize_str(&hex::encode(value)) + } + + pub fn deserialize<'de, D>(deserializer: D) -> Result, D::Error> + where + D: Deserializer<'de>, + { + let value = String::deserialize(deserializer)?; + hex::decode(value.strip_prefix("0x").unwrap_or(&value)).map_err(serde::de::Error::custom) + } +} + +mod option_hex_bytes { + use serde::{Deserialize, Deserializer, Serialize, Serializer}; + + pub fn serialize(value: &Option>, serializer: S) -> Result + where + S: Serializer, + { + value.as_ref().map(hex::encode).serialize(serializer) + } + + pub fn deserialize<'de, D>(deserializer: D) -> Result>, D::Error> + where + D: Deserializer<'de>, + { + Option::::deserialize(deserializer)? + .map(|value| hex::decode(value.strip_prefix("0x").unwrap_or(&value))) + .transpose() + .map_err(serde::de::Error::custom) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::mpc_identity::{EpochManifest, EpochMember}; + + fn manifest() -> EpochManifest { + EpochManifest { + provider_id: vec![1; 32], + epoch: 3, + threshold: 2, + previous_manifest_hash: vec![], + members: ["kms-1", "kms-2", "kms-3"] + .into_iter() + .map(|node_id| EpochMember { + node_id: node_id.into(), + endpoint: format!("https://{node_id}/prpc"), + attestation_pubkey: vec![2; 32], + share_commitment: vec![3; 33], + }) + .collect(), + } + } + + #[test] + fn operation_hash_binds_semantics_and_canonical_quorum() { + let operation = MpcOperation::new_k256( + [4; 32], + 3, + vec!["kms-1".into(), "kms-2".into()], + unix_time().unwrap() + 30, + K256SignPayload { + prefix: b"dstack-env-encrypt-pubkey".to_vec(), + app_id: vec![5; 20], + timestamp: Some(7), + message: vec![6; 32], + }, + ) + .unwrap(); + operation.validate(&manifest(), "kms-1").unwrap(); + let mut tampered = operation.clone(); + if let MpcOperationPayload::SignK256(payload) = &mut tampered.payload { + payload.message[0] ^= 1; + } + assert!(tampered.validate(&manifest(), "kms-1").is_err()); + let mut reordered = operation; + reordered.participants.swap(0, 1); + reordered.request_hash = reordered.compute_request_hash().unwrap().to_vec(); + assert!(reordered.validate(&manifest(), "kms-1").is_err()); + } +} diff --git a/dstack/kms/src/mpc_reshare.rs b/dstack/kms/src/mpc_reshare.rs new file mode 100644 index 000000000..7c7cc88da --- /dev/null +++ b/dstack/kms/src/mpc_reshare.rs @@ -0,0 +1,442 @@ +// SPDX-FileCopyrightText: © 2024-2026 Phala Network +// +// SPDX-License-Identifier: Apache-2.0 + +//! Verifiable Shamir resharing while preserving the group public key. +//! +//! Exactly an old threshold of dealers contributes a polynomial whose constant +//! is its Lagrange-weighted old share. New recipients verify private +//! evaluations against public coefficient commitments before summing them. + +use anyhow::{ensure, Context, Result}; +use cggmp21::{ + generic_ec::{Curve, NonZero, Point, Scalar, SecretScalar}, + key_share::{DirtyIncompleteKeyShare, DirtyKeyInfo, Valid, VssSetup}, + IncompleteKeyShare, +}; +use rand::{CryptoRng, RngCore}; +use serde::{Deserialize, Serialize}; + +#[derive(Clone, Debug, Serialize, Deserialize)] +pub(crate) struct PublicContribution { + pub dealer_index: u16, + pub commitments: Vec>, +} + +#[derive(Clone, Debug, Serialize, Deserialize)] +pub(crate) struct PrivateContribution { + pub dealer_index: u16, + pub recipient_index: u16, + pub evaluation: Vec, +} + +pub(crate) struct ResharedMaterial { + pub secret_share: Scalar, + pub public_shares: Vec>, + pub shared_public_key: Point, +} + +impl ResharedMaterial { + pub(crate) fn into_incomplete_share( + self, + recipient_index: u16, + threshold: u16, + template: &IncompleteKeyShare, + ) -> Result> { + ensure!( + usize::from(recipient_index) < self.public_shares.len(), + "reshared recipient index out of bounds" + ); + self.into_incomplete_share_with_key_info(recipient_index, threshold, &template.key_info) + } + + pub(crate) fn into_incomplete_share_with_key_info( + self, + recipient_index: u16, + threshold: u16, + template: &DirtyKeyInfo, + ) -> Result> { + let mut key_info = template.clone(); + key_info.shared_public_key = NonZero::from_point(self.shared_public_key) + .context("reshared group public key is identity")?; + key_info.public_shares = self + .public_shares + .into_iter() + .map(|point| NonZero::from_point(point).context("reshared public share is identity")) + .collect::>>()?; + key_info.vss_setup = Some(VssSetup { + min_signers: threshold, + I: (1..=key_info.public_shares.len()) + .map(|index| { + NonZero::from_scalar(Scalar::from(index as u64)) + .expect("positive share index is non-zero") + }) + .collect(), + }); + let mut secret = self.secret_share; + let secret = NonZero::from_secret_scalar(SecretScalar::new(&mut secret)) + .context("reshared secret share is zero")?; + Valid::validate(DirtyIncompleteKeyShare { + i: recipient_index, + key_info, + x: secret, + }) + .map_err(|error| anyhow::anyhow!("invalid reshared CGGMP core share: {error}")) + } +} + +pub(crate) fn create_contribution( + old_share: &IncompleteKeyShare, + old_dealers: &[u16], + new_member_count: u16, + new_threshold: u16, + rng: &mut R, +) -> Result<(PublicContribution, Vec)> { + ensure!( + old_share.vss_setup.is_some(), + "resharing requires an old Shamir share" + ); + ensure!( + old_dealers.len() == usize::from(old_share.vss_setup.as_ref().unwrap().min_signers), + "resharing must use exactly the old threshold of dealers" + ); + ensure!( + old_dealers.windows(2).all(|pair| pair[0] < pair[1]), + "old resharing dealers must be unique and ordered" + ); + ensure!( + old_dealers.contains(&old_share.i), + "local old share is not a resharing dealer" + ); + ensure!( + new_threshold >= 2 && new_threshold <= new_member_count, + "invalid new resharing threshold" + ); + let setup = old_share.vss_setup.as_ref().unwrap(); + let dealer_position = old_dealers + .iter() + .position(|index| *index == old_share.i) + .expect("dealer membership checked"); + let dealer_points = old_dealers + .iter() + .map(|index| { + setup + .I + .get(usize::from(*index)) + .map(|point| Scalar::from(point.clone())) + .context("old dealer index is out of bounds") + }) + .collect::>>()?; + let lambda = lagrange_at_zero(dealer_position, &dealer_points)?; + let mut coefficients = Vec::with_capacity(usize::from(new_threshold)); + let old_secret: &Scalar = old_share.x.as_ref(); + coefficients.push(lambda * old_secret); + coefficients.extend((1..new_threshold).map(|_| Scalar::::random(rng))); + let commitments = coefficients + .iter() + .map(|coefficient| { + (Point::generator() * coefficient) + .to_bytes(true) + .as_bytes() + .to_vec() + }) + .collect(); + let private = (0..new_member_count) + .map(|recipient_index| { + let x = Scalar::from(u64::from(recipient_index) + 1); + PrivateContribution { + dealer_index: old_share.i, + recipient_index, + evaluation: evaluate_polynomial(&coefficients, x) + .to_be_bytes() + .as_bytes() + .to_vec(), + } + }) + .collect(); + Ok(( + PublicContribution { + dealer_index: old_share.i, + commitments, + }, + private, + )) +} + +#[allow(clippy::too_many_arguments)] +pub(crate) fn verify_and_combine( + old_reference: &IncompleteKeyShare, + old_dealers: &[u16], + new_member_count: u16, + new_threshold: u16, + recipient_index: u16, + public: &[PublicContribution], + private: &[PrivateContribution], +) -> Result> { + verify_and_combine_with_key_info( + &old_reference.key_info, + old_dealers, + new_member_count, + new_threshold, + recipient_index, + public, + private, + ) +} + +#[allow(clippy::too_many_arguments)] +pub(crate) fn verify_and_combine_with_key_info( + old_key_info: &DirtyKeyInfo, + old_dealers: &[u16], + new_member_count: u16, + new_threshold: u16, + recipient_index: u16, + public: &[PublicContribution], + private: &[PrivateContribution], +) -> Result> { + ensure!( + recipient_index < new_member_count, + "new recipient index out of bounds" + ); + ensure!( + public.len() == old_dealers.len(), + "missing public resharing contribution" + ); + ensure!( + private.len() == old_dealers.len(), + "missing private resharing contribution" + ); + let setup = old_key_info + .vss_setup + .as_ref() + .context("resharing requires an old Shamir setup")?; + let dealer_points = old_dealers + .iter() + .map(|index| { + setup + .I + .get(usize::from(*index)) + .map(|point| Scalar::from(point.clone())) + .context("old dealer index is out of bounds") + }) + .collect::>>()?; + let recipient_x = Scalar::from(u64::from(recipient_index) + 1); + let mut aggregate_commitments = vec![Point::::zero(); usize::from(new_threshold)]; + let mut secret_share = Scalar::::zero(); + for (position, dealer_index) in old_dealers.iter().enumerate() { + let public = public + .iter() + .find(|item| item.dealer_index == *dealer_index) + .context("missing dealer public contribution")?; + ensure!( + public.commitments.len() == usize::from(new_threshold), + "dealer commitment degree does not match new threshold" + ); + let commitments = public + .commitments + .iter() + .map(|encoded| Point::::from_bytes(encoded).context("invalid resharing commitment")) + .collect::>>()?; + let lambda = lagrange_at_zero(position, &dealer_points)?; + let expected_constant = old_key_info + .public_shares + .get(usize::from(*dealer_index)) + .context("missing old dealer public share")? + * lambda; + ensure!( + commitments[0] == expected_constant, + "dealer changed its weighted old share" + ); + let private = private + .iter() + .find(|item| { + item.dealer_index == *dealer_index && item.recipient_index == recipient_index + }) + .context("missing recipient resharing contribution")?; + let evaluation = Scalar::::from_be_bytes(&private.evaluation) + .context("invalid private resharing scalar")?; + ensure!( + Point::generator() * evaluation == evaluate_commitments(&commitments, recipient_x), + "private resharing contribution does not match commitments" + ); + secret_share += evaluation; + for (aggregate, commitment) in aggregate_commitments.iter_mut().zip(commitments) { + *aggregate += commitment; + } + } + ensure!( + aggregate_commitments[0] == *old_key_info.shared_public_key, + "resharing changed the group public key" + ); + let public_shares = (0..new_member_count) + .map(|index| { + evaluate_commitments(&aggregate_commitments, Scalar::from(u64::from(index) + 1)) + }) + .collect::>(); + ensure!( + Point::generator() * secret_share == public_shares[usize::from(recipient_index)], + "combined reshared secret does not match its public share" + ); + Ok(ResharedMaterial { + secret_share, + public_shares, + shared_public_key: aggregate_commitments[0], + }) +} + +fn lagrange_at_zero(position: usize, indexes: &[Scalar]) -> Result> { + let x_i = indexes[position]; + let mut numerator = Scalar::one(); + let mut denominator = Scalar::one(); + for (other, x_j) in indexes.iter().enumerate() { + if other != position { + numerator *= -*x_j; + denominator *= x_i - x_j; + } + } + Ok(numerator + * denominator + .invert() + .context("duplicate resharing evaluation point")?) +} + +fn evaluate_polynomial(coefficients: &[Scalar], x: Scalar) -> Scalar { + coefficients + .iter() + .rev() + .fold(Scalar::zero(), |value, coefficient| value * x + coefficient) +} + +fn evaluate_commitments(coefficients: &[Point], x: Scalar) -> Point { + coefficients + .iter() + .rev() + .fold(Point::zero(), |value, coefficient| value * x + coefficient) +} + +#[cfg(test)] +mod tests { + use super::*; + use cggmp21::{generic_ec::NonZero, supported_curves::Secp256k1}; + use key_share::trusted_dealer; + use rand::rngs::OsRng; + + #[test] + fn reshares_two_of_three_into_three_of_five_without_changing_key() { + let old = trusted_dealer::builder::(3) + .set_threshold(Some(2)) + .generate_shares(&mut OsRng) + .unwrap(); + let dealers = [0u16, 2]; + let mut publics = vec![]; + let mut private_by_recipient = vec![vec![]; 5]; + for dealer in dealers { + let (public, private) = + create_contribution(&old[usize::from(dealer)], &dealers, 5, 3, &mut OsRng).unwrap(); + publics.push(public); + for contribution in private { + private_by_recipient[usize::from(contribution.recipient_index)].push(contribution); + } + } + let materials = private_by_recipient + .iter() + .enumerate() + .map(|(recipient, private)| { + verify_and_combine(&old[0], &dealers, 5, 3, recipient as u16, &publics, private) + .unwrap() + }) + .collect::>(); + assert!(materials.iter().all(|material| { + material.shared_public_key == *old[0].shared_public_key() + && material.public_shares == materials[0].public_shares + })); + let indexes = [Scalar::from(1u16), Scalar::from(2u16), Scalar::from(3u16)]; + let reconstructed = (0..3).fold(Scalar::zero(), |secret, position| { + secret + + lagrange_at_zero(position, &indexes).unwrap() * materials[position].secret_share + }); + assert_eq!( + Point::generator() * reconstructed, + *old[0].shared_public_key() + ); + assert!(NonZero::from_point(materials[0].shared_public_key).is_some()); + for (index, material) in materials.into_iter().enumerate() { + let share = material + .into_incomplete_share(index as u16, 3, &old[0]) + .unwrap(); + assert_eq!(share.i, index as u16); + assert_eq!(*share.shared_public_key(), *old[0].shared_public_key()); + } + } + + #[test] + fn joining_recipient_materializes_from_public_old_topology_only() { + let old = trusted_dealer::builder::(3) + .set_threshold(Some(2)) + .generate_shares(&mut OsRng) + .unwrap(); + let dealers = [0u16, 1u16]; + let contributions = dealers + .iter() + .map(|index| { + create_contribution(&old[usize::from(*index)], &dealers, 5, 3, &mut OsRng).unwrap() + }) + .collect::>(); + let public = contributions + .iter() + .map(|item| item.0.clone()) + .collect::>(); + let private = contributions + .iter() + .map(|item| item.1[4].clone()) + .collect::>(); + // The joining party receives only public old-key topology plus its own + // verifiable evaluations; it never receives any old secret share. + let material = verify_and_combine_with_key_info( + &old[0].key_info, + &dealers, + 5, + 3, + 4, + &public, + &private, + ) + .unwrap(); + let joined = material + .into_incomplete_share_with_key_info(4, 3, &old[0].key_info) + .unwrap(); + assert_eq!(joined.i, 4); + assert_eq!( + joined.key_info.shared_public_key, + old[0].key_info.shared_public_key + ); + assert_eq!(joined.key_info.public_shares.len(), 5); + } + + #[test] + fn rejects_a_dealer_that_changes_the_group_key() { + let old = trusted_dealer::builder::(3) + .set_threshold(Some(2)) + .generate_shares(&mut OsRng) + .unwrap(); + let dealers = [0u16, 1]; + let (mut first_public, first_private) = + create_contribution(&old[0], &dealers, 3, 2, &mut OsRng).unwrap(); + let (second_public, second_private) = + create_contribution(&old[1], &dealers, 3, 2, &mut OsRng).unwrap(); + first_public.commitments[0] = (Point::::generator() * Scalar::one()) + .to_bytes(true) + .as_bytes() + .to_vec(); + assert!(verify_and_combine( + &old[0], + &dealers, + 3, + 2, + 0, + &[first_public, second_public], + &[first_private[0].clone(), second_private[0].clone()], + ) + .is_err()); + } +} diff --git a/dstack/kms/src/mpc_service.rs b/dstack/kms/src/mpc_service.rs new file mode 100644 index 000000000..bd98dd45b --- /dev/null +++ b/dstack/kms/src/mpc_service.rs @@ -0,0 +1,73 @@ +// SPDX-FileCopyrightText: © 2024-2026 Phala Network +// +// SPDX-License-Identifier: Apache-2.0 + +use anyhow::{Context, Result}; +use dstack_kms_rpc::{ + mpc_transport_server::{MpcTransportRpc, MpcTransportServer}, + MpcPushRequest, MpcStartRequest, MpcStartResponse, +}; +use ra_rpc::{CallContext, RpcCall}; + +use crate::{main_service::KmsState, mpc_session::MpcEnvelope}; + +pub(crate) struct MpcHandler { + state: KmsState, + peer_public_key: Vec, +} + +impl RpcCall for MpcHandler { + type PrpcService = MpcTransportServer; + + fn construct(context: CallContext<'_, KmsState>) -> Result { + context + .attestation + .context("MPC transport requires an attested RA-TLS peer")?; + let peer_public_key = context + .remote_public_key + .context("MPC transport peer certificate is missing")?; + Ok(Self { + state: context.state.clone(), + peer_public_key, + }) + } +} + +impl MpcHandler { + fn peer_and_router(&self) -> Result<(&str, &crate::mpc_session::SessionRouter)> { + let router = self.state.mpc_router().context("MPC mode is not enabled")?; + let node_id = router.authenticated_node_id(&self.peer_public_key)?; + Ok((node_id, router)) + } +} + +impl MpcTransportRpc for MpcHandler { + async fn ping(self) -> Result<()> { + self.peer_and_router()?; + Ok(()) + } + + async fn start(self, request: MpcStartRequest) -> Result { + let operation = serde_json::from_slice(&request.operation_json) + .context("invalid MPC operation encoding")?; + let (node_id, _) = self.peer_and_router()?; + let node_id = node_id.to_owned(); + let result = self + .state + .key_backend() + .run_mpc_operation(operation, &node_id) + .await?; + Ok(MpcStartResponse { result }) + } + + async fn push(self, request: MpcPushRequest) -> Result<()> { + let envelope: MpcEnvelope = serde_json::from_slice(&request.envelope_json) + .context("invalid MPC envelope encoding")?; + let (node_id, router) = self.peer_and_router()?; + router.push(node_id, envelope) + } +} + +pub(crate) fn rpc_methods() -> &'static [&'static str] { + >::supported_methods() +} diff --git a/dstack/kms/src/mpc_session.rs b/dstack/kms/src/mpc_session.rs new file mode 100644 index 000000000..0ae5a7a5a --- /dev/null +++ b/dstack/kms/src/mpc_session.rs @@ -0,0 +1,298 @@ +// SPDX-FileCopyrightText: © 2024-2026 Phala Network +// +// SPDX-License-Identifier: Apache-2.0 + +//! Authenticated message router for interactive MPC protocols. + +use std::{ + collections::{BTreeMap, VecDeque}, + sync::Mutex, + time::{Duration, SystemTime, UNIX_EPOCH}, +}; + +use anyhow::{bail, ensure, Context, Result}; +use serde::{Deserialize, Serialize}; + +use crate::mpc_identity::EpochManifest; + +const MAX_MESSAGE_BYTES: usize = 1024 * 1024; + +#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub(crate) enum MpcProtocol { + AuxiliaryInfo, + DkgP256, + DkgK256, + DkgDerivation, + SignP256, + SignK256, + Derive, + Reshare, +} + +#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] +pub(crate) struct MpcEnvelope { + #[serde(with = "hex_bytes")] + pub session_id: Vec, + pub epoch: u64, + pub protocol: MpcProtocol, + #[serde(with = "hex_bytes")] + pub request_hash: Vec, + pub sender: String, + pub recipient: String, + pub sequence: u64, + pub expires_at: u64, + #[serde(with = "hex_bytes")] + pub payload: Vec, + pub broadcast: bool, +} + +struct Session { + protocol: MpcProtocol, + request_hash: Vec, + expires_at: u64, + last_sequence: BTreeMap, + inboxes: BTreeMap>, +} + +pub(crate) struct SessionRouter { + manifest: EpochManifest, + max_sessions: usize, + max_ttl: Duration, + sessions: Mutex, Session>>, +} + +impl SessionRouter { + pub(crate) fn new( + manifest: EpochManifest, + max_sessions: usize, + max_ttl: Duration, + ) -> Result { + manifest.manifest_hash().context("invalid epoch manifest")?; + ensure!(max_sessions > 0, "max_sessions must be greater than zero"); + ensure!(!max_ttl.is_zero(), "max_ttl must be greater than zero"); + Ok(Self { + manifest, + max_sessions, + max_ttl, + sessions: Mutex::new(BTreeMap::new()), + }) + } + + /// Enqueue a protocol message after binding its claimed sender to the + /// independently authenticated RA-TLS peer identity. + pub(crate) fn push(&self, authenticated_node_id: &str, message: MpcEnvelope) -> Result<()> { + self.validate_envelope(authenticated_node_id, &message)?; + let now = unix_time()?; + let mut sessions = self.sessions.lock().expect("session mutex poisoned"); + sessions.retain(|_, session| session.expires_at >= now); + if !sessions.contains_key(&message.session_id) && sessions.len() >= self.max_sessions { + bail!("too many active MPC sessions"); + } + let session = sessions + .entry(message.session_id.clone()) + .or_insert_with(|| Session { + protocol: message.protocol, + request_hash: message.request_hash.clone(), + expires_at: message.expires_at, + last_sequence: BTreeMap::new(), + inboxes: BTreeMap::new(), + }); + ensure!( + session.protocol == message.protocol, + "session protocol changed" + ); + ensure!( + session.request_hash == message.request_hash, + "session request hash changed" + ); + ensure!( + session.expires_at == message.expires_at, + "session expiration changed" + ); + let last = session + .last_sequence + .entry(message.sender.clone()) + .or_default(); + ensure!( + message.sequence > *last, + "replayed or reordered MPC message" + ); + *last = message.sequence; + session + .inboxes + .entry(message.recipient.clone()) + .or_default() + .push_back(message); + Ok(()) + } + + pub(crate) fn authenticated_node_id(&self, attested_public_key: &[u8]) -> Result<&str> { + self.manifest + .member_for_attestation_key(attested_public_key) + .map(|member| member.node_id.as_str()) + .context("attested RA-TLS key is not an epoch member") + } + + pub(crate) fn drain( + &self, + authenticated_node_id: &str, + session_id: &[u8], + ) -> Result> { + ensure!(session_id.len() == 32, "session_id must be 32 bytes"); + ensure!( + self.manifest.contains_member(authenticated_node_id), + "authenticated peer is not an epoch member" + ); + let mut sessions = self.sessions.lock().expect("session mutex poisoned"); + let Some(session) = sessions.get_mut(session_id) else { + return Ok(vec![]); + }; + Ok(session + .inboxes + .remove(authenticated_node_id) + .unwrap_or_default() + .into()) + } + + fn validate_envelope(&self, authenticated_node_id: &str, message: &MpcEnvelope) -> Result<()> { + ensure!( + message.session_id.len() == 32, + "session_id must be 32 bytes" + ); + ensure!( + message.request_hash.len() == 32, + "request_hash must be 32 bytes" + ); + ensure!(message.epoch == self.manifest.epoch, "MPC epoch mismatch"); + ensure!( + message.sender == authenticated_node_id, + "MPC sender identity mismatch" + ); + ensure!( + self.manifest.contains_member(&message.sender), + "sender is not an epoch member" + ); + ensure!( + self.manifest.contains_member(&message.recipient), + "recipient is not an epoch member" + ); + ensure!( + message.sender != message.recipient, + "sender and recipient must differ" + ); + ensure!(!message.payload.is_empty(), "MPC payload must not be empty"); + ensure!( + message.payload.len() <= MAX_MESSAGE_BYTES, + "MPC payload is too large" + ); + let now = unix_time()?; + ensure!(message.expires_at >= now, "MPC message expired"); + ensure!( + message.expires_at - now <= self.max_ttl.as_secs(), + "MPC message expiration exceeds maximum TTL" + ); + Ok(()) + } +} + +fn unix_time() -> Result { + Ok(SystemTime::now() + .duration_since(UNIX_EPOCH) + .context("system time before UNIX epoch")? + .as_secs()) +} + +mod hex_bytes { + use serde::{Deserialize, Deserializer, Serializer}; + + pub fn serialize(value: &[u8], serializer: S) -> Result + where + S: Serializer, + { + serializer.serialize_str(&hex::encode(value)) + } + + pub fn deserialize<'de, D>(deserializer: D) -> Result, D::Error> + where + D: Deserializer<'de>, + { + let value = String::deserialize(deserializer)?; + hex::decode(value.strip_prefix("0x").unwrap_or(&value)).map_err(serde::de::Error::custom) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::mpc_identity::EpochMember; + + fn router() -> SessionRouter { + let members = ["kms-1", "kms-2", "kms-3"] + .into_iter() + .map(|node_id| EpochMember { + node_id: node_id.into(), + endpoint: format!("https://{node_id}:8443/prpc"), + attestation_pubkey: vec![1; 32], + share_commitment: vec![2; 33], + }) + .collect(); + SessionRouter::new( + EpochManifest { + provider_id: vec![3; 32], + epoch: 4, + threshold: 2, + previous_manifest_hash: vec![], + members, + }, + 8, + Duration::from_secs(60), + ) + .unwrap() + } + + fn message() -> MpcEnvelope { + MpcEnvelope { + session_id: vec![4; 32], + epoch: 4, + protocol: MpcProtocol::SignK256, + request_hash: vec![5; 32], + sender: "kms-1".into(), + recipient: "kms-2".into(), + sequence: 1, + expires_at: unix_time().unwrap() + 30, + payload: vec![6], + broadcast: false, + } + } + + #[test] + fn routes_only_authenticated_epoch_members() { + let router = router(); + router.push("kms-1", message()).unwrap(); + assert_eq!(router.drain("kms-2", &[4; 32]).unwrap(), vec![message()]); + assert!(router.push("kms-3", message()).is_err()); + } + + #[test] + fn rejects_replay_and_session_context_changes() { + let router = router(); + let original = message(); + router.push("kms-1", original.clone()).unwrap(); + assert!(router.push("kms-1", original.clone()).is_err()); + let mut changed = original; + changed.sequence = 2; + changed.request_hash[0] ^= 1; + assert!(router.push("kms-1", changed).is_err()); + + let mut wrong_epoch = message(); + wrong_epoch.epoch = 3; + assert!(router.push("kms-1", wrong_epoch).is_err()); + let mut forged_sender = message(); + forged_sender.sender = "kms-3".into(); + assert!(router.push("kms-1", forged_sender).is_err()); + let mut unknown_recipient = message(); + unknown_recipient.recipient = "kms-9".into(); + assert!(router.push("kms-1", unknown_recipient).is_err()); + } +} diff --git a/dstack/kms/src/threshold_prf.rs b/dstack/kms/src/threshold_prf.rs new file mode 100644 index 000000000..4000e345e --- /dev/null +++ b/dstack/kms/src/threshold_prf.rs @@ -0,0 +1,286 @@ +// SPDX-FileCopyrightText: © 2024-2026 Phala Network +// +// SPDX-License-Identifier: Apache-2.0 + +//! Verifiable threshold PRF over P-256. +//! +//! Each node evaluates a hash-to-curve point with its Shamir share and proves +//! equality of discrete logarithms against the public share commitment. A +//! quorum combines verified evaluations at zero without reconstructing the +//! cluster secret scalar. + +use anyhow::{ensure, Context, Result}; +use p256::{ + elliptic_curve::{ + bigint::U256, + ff::Field, + group::Group, + hash2curve::{ExpandMsgXmd, GroupDigest}, + ops::Reduce, + sec1::{FromEncodedPoint, ToEncodedPoint}, + PrimeField, + }, + AffinePoint, EncodedPoint, NistP256, ProjectivePoint, Scalar, +}; +use rand::{CryptoRng, RngCore}; +use serde::{Deserialize, Serialize}; +use sha2::{Digest, Sha256}; + +const HASH_TO_CURVE_DST: &[u8] = b"DSTACK-MPC-THRESHOLD-PRF-P256_XMD:SHA-256_SSWU_RO_V1"; +const CHALLENGE_DOMAIN: &[u8] = b"dstack-mpc-threshold-prf-dleq-v1"; +const OUTPUT_DOMAIN: &[u8] = b"dstack-mpc-threshold-prf-output-v1"; + +#[derive(Clone, Debug, Serialize, Deserialize)] +pub(crate) struct PrfPartial { + pub keygen_index: u16, + #[serde(with = "hex_bytes")] + pub evaluation: Vec, + #[serde(with = "hex_bytes")] + pub commitment_g: Vec, + #[serde(with = "hex_bytes")] + pub commitment_h: Vec, + #[serde(with = "hex_bytes")] + pub response: Vec, +} + +pub(crate) fn evaluate( + keygen_index: u16, + secret_share: &[u8; 32], + public_share: &[u8], + input: &[u8], + request_hash: &[u8; 32], + rng: &mut R, +) -> Result { + let secret = scalar(secret_share).context("invalid threshold PRF secret share")?; + ensure!(!bool::from(secret.is_zero()), "threshold PRF share is zero"); + let public = point(public_share).context("invalid threshold PRF public share")?; + ensure!( + ProjectivePoint::GENERATOR * secret == public, + "threshold PRF secret does not match public share" + ); + let h = hash_to_curve(input)?; + let evaluation = h * secret; + let nonce = Scalar::random(rng); + let commitment_g = ProjectivePoint::GENERATOR * nonce; + let commitment_h = h * nonce; + let challenge = challenge( + keygen_index, + input, + request_hash, + public, + h, + evaluation, + commitment_g, + commitment_h, + ); + let response = nonce + challenge * secret; + Ok(PrfPartial { + keygen_index, + evaluation: encode_point(evaluation), + commitment_g: encode_point(commitment_g), + commitment_h: encode_point(commitment_h), + response: response.to_bytes().to_vec(), + }) +} + +pub(crate) fn verify( + partial: &PrfPartial, + public_share: &[u8], + input: &[u8], + request_hash: &[u8; 32], +) -> Result { + let public = point(public_share).context("invalid threshold PRF public share")?; + let h = hash_to_curve(input)?; + let evaluation = point(&partial.evaluation).context("invalid PRF evaluation")?; + let commitment_g = point(&partial.commitment_g).context("invalid DLEQ commitment")?; + let commitment_h = point(&partial.commitment_h).context("invalid DLEQ commitment")?; + let response_bytes: [u8; 32] = partial + .response + .as_slice() + .try_into() + .context("invalid DLEQ response length")?; + let response = scalar(&response_bytes).context("invalid DLEQ response")?; + let challenge = challenge( + partial.keygen_index, + input, + request_hash, + public, + h, + evaluation, + commitment_g, + commitment_h, + ); + ensure!( + ProjectivePoint::GENERATOR * response == commitment_g + public * challenge, + "invalid threshold PRF proof for generator relation" + ); + ensure!( + h * response == commitment_h + evaluation * challenge, + "invalid threshold PRF proof for evaluation relation" + ); + Ok(evaluation) +} + +pub(crate) fn combine( + mut evaluations: Vec<(u16, ProjectivePoint)>, + context: &[u8], +) -> Result<[u8; 32]> { + ensure!(!evaluations.is_empty(), "no threshold PRF evaluations"); + evaluations.sort_by_key(|(index, _)| *index); + ensure!( + evaluations.windows(2).all(|pair| pair[0].0 != pair[1].0), + "duplicate threshold PRF share index" + ); + let indexes: Vec<_> = evaluations + .iter() + .map(|(index, _)| Scalar::from(u64::from(*index) + 1)) + .collect(); + let mut combined = ProjectivePoint::IDENTITY; + for (position, (_, evaluation)) in evaluations.iter().enumerate() { + let x_i = indexes[position]; + let mut numerator = Scalar::ONE; + let mut denominator = Scalar::ONE; + for (other, x_j) in indexes.iter().enumerate() { + if other != position { + numerator *= -*x_j; + denominator *= x_i - x_j; + } + } + let inverse = Option::::from(denominator.invert()) + .context("invalid threshold PRF interpolation set")?; + combined += *evaluation * (numerator * inverse); + } + ensure!( + !bool::from(combined.is_identity()), + "threshold PRF output is identity" + ); + let mut hash = Sha256::new(); + hash.update((OUTPUT_DOMAIN.len() as u32).to_be_bytes()); + hash.update(OUTPUT_DOMAIN); + hash.update((context.len() as u32).to_be_bytes()); + hash.update(context); + hash.update(encode_point(combined)); + Ok(hash.finalize().into()) +} + +fn hash_to_curve(input: &[u8]) -> Result { + NistP256::hash_from_bytes::>(&[input], &[HASH_TO_CURVE_DST]) + .map_err(|_| anyhow::anyhow!("failed to hash threshold PRF input to P-256")) +} + +#[allow(clippy::too_many_arguments)] +fn challenge( + keygen_index: u16, + input: &[u8], + request_hash: &[u8; 32], + public: ProjectivePoint, + h: ProjectivePoint, + evaluation: ProjectivePoint, + commitment_g: ProjectivePoint, + commitment_h: ProjectivePoint, +) -> Scalar { + let mut hash = Sha256::new(); + hash.update((CHALLENGE_DOMAIN.len() as u32).to_be_bytes()); + hash.update(CHALLENGE_DOMAIN); + hash.update(keygen_index.to_be_bytes()); + hash.update((input.len() as u32).to_be_bytes()); + hash.update(input); + hash.update(request_hash); + for value in [ + ProjectivePoint::GENERATOR, + public, + h, + evaluation, + commitment_g, + commitment_h, + ] { + hash.update(encode_point(value)); + } + >::reduce_bytes(&hash.finalize()) +} + +fn scalar(bytes: &[u8; 32]) -> Option { + Option::from(Scalar::from_repr((*bytes).into())) +} + +fn point(bytes: &[u8]) -> Result { + let encoded = EncodedPoint::from_bytes(bytes).context("invalid P-256 point encoding")?; + Option::::from(AffinePoint::from_encoded_point(&encoded)) + .map(ProjectivePoint::from) + .context("P-256 point is not on curve") +} + +fn encode_point(point: ProjectivePoint) -> Vec { + point.to_affine().to_encoded_point(true).as_bytes().to_vec() +} + +mod hex_bytes { + use serde::{Deserialize, Deserializer, Serializer}; + + pub fn serialize(value: &[u8], serializer: S) -> Result + where + S: Serializer, + { + serializer.serialize_str(&hex::encode(value)) + } + + pub fn deserialize<'de, D>(deserializer: D) -> Result, D::Error> + where + D: Deserializer<'de>, + { + let value = String::deserialize(deserializer)?; + hex::decode(value.strip_prefix("0x").unwrap_or(&value)).map_err(serde::de::Error::custom) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use p256::elliptic_curve::sec1::ToEncodedPoint; + use rand::rngs::OsRng; + + #[test] + fn verifies_and_combines_two_of_three_without_reconstructing_secret() { + let secret = Scalar::random(&mut OsRng); + let slope = Scalar::random(&mut OsRng); + let shares: Vec<_> = (1u64..=3) + .map(|index| secret + slope * Scalar::from(index)) + .collect(); + let input = b"app-id\0instance-id\0disk-key"; + let request_hash = [9u8; 32]; + let mut evaluations = vec![]; + for index in [0usize, 2] { + let share_bytes: [u8; 32] = shares[index].to_bytes().into(); + let public = (ProjectivePoint::GENERATOR * shares[index]) + .to_affine() + .to_encoded_point(true); + let partial = evaluate( + index as u16, + &share_bytes, + public.as_bytes(), + input, + &request_hash, + &mut OsRng, + ) + .unwrap(); + evaluations.push(( + index as u16, + verify(&partial, public.as_bytes(), input, &request_hash).unwrap(), + )); + } + let output = combine(evaluations, input).unwrap(); + let expected = combine(vec![(0, hash_to_curve(input).unwrap() * secret)], input).unwrap(); + assert_eq!(output, expected); + } + + #[test] + fn rejects_tampered_partial() { + let share = Scalar::random(&mut OsRng); + let share_bytes: [u8; 32] = share.to_bytes().into(); + let public = encode_point(ProjectivePoint::GENERATOR * share); + let mut partial = + evaluate(0, &share_bytes, &public, b"input", &[1; 32], &mut OsRng).unwrap(); + partial.response[0] ^= 1; + assert!(verify(&partial, &public, b"input", &[1; 32]).is_err()); + } +} diff --git a/dstack/ra-rpc/src/lib.rs b/dstack/ra-rpc/src/lib.rs index bf1e22080..1ac1419cf 100644 --- a/dstack/ra-rpc/src/lib.rs +++ b/dstack/ra-rpc/src/lib.rs @@ -57,6 +57,10 @@ pub struct CallContext<'a, State> { pub remote_endpoint: Option, pub remote_app_id: Option>, pub remote_app_info: Option, + /// SubjectPublicKeyInfo bytes from the authenticated peer certificate. + /// When attestation is present, quote verification has bound this key to + /// the report data. + pub remote_public_key: Option>, } pub trait RpcCall: Sized { diff --git a/dstack/ra-rpc/src/rocket_helper.rs b/dstack/ra-rpc/src/rocket_helper.rs index 2cc170630..d803a99d5 100644 --- a/dstack/ra-rpc/src/rocket_helper.rs +++ b/dstack/ra-rpc/src/rocket_helper.rs @@ -532,6 +532,10 @@ pub async fn handle_prpc_impl>( .map(|cert| ra_tls::attestation::from_der(cert.as_bytes())) .transpose()? .flatten(); + let remote_public_key = request + .certificate + .as_ref() + .map(|certificate| certificate.public_key().raw.to_vec()); let attestation = match (request.quote_verifier, attestation) { (Some(quote_verifier), Some(attestation)) => { let pubkey = request @@ -568,6 +572,7 @@ pub async fn handle_prpc_impl>( remote_endpoint: request.remote_addr.cloned().map(RemoteEndpoint::from), remote_app_id, remote_app_info, + remote_public_key, }; let call = Call::construct(context).context("failed to construct call")?; let (status_code, output) = call diff --git a/dstack/ra-tls/Cargo.toml b/dstack/ra-tls/Cargo.toml index 3725f3330..90de13b1a 100644 --- a/dstack/ra-tls/Cargo.toml +++ b/dstack/ra-tls/Cargo.toml @@ -11,6 +11,7 @@ license.workspace = true [dependencies] anyhow.workspace = true +base64.workspace = true bon.workspace = true dcap-qvl.workspace = true elliptic-curve.workspace = true diff --git a/dstack/ra-tls/src/cert.rs b/dstack/ra-tls/src/cert.rs index df8f4aa05..e1c47950a 100644 --- a/dstack/ra-tls/src/cert.rs +++ b/dstack/ra-tls/src/cert.rs @@ -8,16 +8,18 @@ use std::time::{SystemTime, UNIX_EPOCH}; use std::{path::Path, time::Duration}; use anyhow::{anyhow, bail, Context, Result}; +use base64::{engine::general_purpose::STANDARD as BASE64, Engine as _}; use fs_err as fs; use rcgen::{ BasicConstraints, Certificate, CertificateParams, CustomExtension, DistinguishedName, DnType, - ExtendedKeyUsagePurpose, IsCa, KeyPair, KeyUsagePurpose, PublicKeyData, + ExtendedKeyUsagePurpose, IsCa, KeyPair, KeyUsagePurpose, PublicKeyData, PKCS_ECDSA_P256_SHA256, }; use ring::rand::SystemRandom; use ring::signature::{ EcdsaKeyPair, UnparsedPublicKey, ECDSA_P256_SHA256_ASN1, ECDSA_P256_SHA256_ASN1_SIGNING, }; use scale::{Decode, Encode}; +use sha2::{Digest as _, Sha256}; use x509_parser::der_parser::Oid; use x509_parser::prelude::{FromDer as _, X509Certificate}; use x509_parser::public_key::PublicKey; @@ -437,6 +439,136 @@ impl CertRequest<'_, Key> { } } +/// A certificate body prepared for an external P-256 signer, such as a +/// threshold signing quorum. The exact TBS bytes are consensus-critical and +/// must be signed with ECDSA/SHA-256. +pub struct ExternalCertificate { + tbs_der: Vec, +} + +impl ExternalCertificate { + /// Recreate an externally signable certificate from previously validated + /// TBS bytes. This is used when a coordinator transports a member-created + /// TBS body to a threshold signing quorum. + pub fn from_tbs_der(tbs_der: Vec) -> Result { + let (remaining, _) = x509_parser::certificate::TbsCertificate::from_der(&tbs_der) + .map_err(|error| anyhow::anyhow!("invalid external certificate TBS: {error}"))?; + anyhow::ensure!( + remaining.is_empty(), + "external certificate TBS has trailing bytes" + ); + Ok(Self { tbs_der }) + } + + /// DER-encoded `TBSCertificate` bytes to hash with SHA-256 and sign. + pub fn tbs_der(&self) -> &[u8] { + &self.tbs_der + } + + /// Assemble the final PEM certificate from a fixed-width `(r || s)` P-256 + /// signature returned by a threshold signer. + pub fn finish(self, signature: &[u8]) -> Result { + let signature = p256::ecdsa::Signature::from_slice(signature) + .context("invalid external P-256 signature")?; + let signature_der = signature.to_der(); + // AlgorithmIdentifier { id-ecdsa-with-SHA256, absent parameters } + const ECDSA_SHA256_ALGORITHM: &[u8] = &[ + 0x30, 0x0a, 0x06, 0x08, 0x2a, 0x86, 0x48, 0xce, 0x3d, 0x04, 0x03, 0x02, + ]; + let der = yasna::construct_der(|writer| { + writer.write_sequence(|writer| { + writer.next().write_der(&self.tbs_der); + writer.next().write_der(ECDSA_SHA256_ALGORITHM); + writer + .next() + .write_bitvec_bytes(signature_der.as_bytes(), signature_der.len() * 8); + }); + }); + let encoded = BASE64.encode(der); + let mut pem = String::from("-----BEGIN CERTIFICATE-----\n"); + for line in encoded.as_bytes().chunks(64) { + pem.push_str(std::str::from_utf8(line).expect("base64 is UTF-8")); + pem.push('\n'); + } + pem.push_str("-----END CERTIFICATE-----\n"); + Ok(pem) + } +} + +/// Prepare a certificate request using the issuer identity and extensions from +/// an existing root certificate, without requiring its private key. A throwaway +/// key is used only to make rcgen serialize the deterministic TBS structure; +/// its signature is discarded. +pub fn prepare_external_certificate( + request: CertRequest<'_, Key>, + issuer_pem: &str, +) -> Result { + let issuer_key = KeyPair::generate_for(&PKCS_ECDSA_P256_SHA256) + .context("failed to create external-certificate serialization key")?; + let issuer_params = CertificateParams::from_ca_cert_pem(issuer_pem) + .context("failed to parse external certificate issuer")?; + let issuer = issuer_params + .self_signed(&issuer_key) + .context("failed to create external certificate issuer template")?; + let public_key = request.key; + let placeholder = request + .into_cert_params()? + .signed_by(public_key, &issuer, &issuer_key) + .context("failed to serialize external certificate template")?; + let (_, parsed) = x509_parser::parse_x509_certificate(placeholder.der()) + .context("failed to parse external certificate template")?; + Ok(ExternalCertificate { + tbs_der: parsed.tbs_certificate.as_ref().to_vec(), + }) +} + +/// Prepare a self-signed certificate body whose public key has no locally +/// available private key. This is used to create a threshold root CA after DKG. +pub fn prepare_external_self_signed( + request: CertRequest<'_, Key>, +) -> Result { + let public_key = request.key; + let params = request.into_cert_params()?; + let subject_public_key_info = yasna::construct_der(|writer| { + writer.write_sequence(|writer| { + writer.next().write_sequence(|writer| { + writer + .next() + .write_oid(&yasna::models::ObjectIdentifier::from_slice(&[ + 1, 2, 840, 10045, 2, 1, + ])); + writer + .next() + .write_oid(&yasna::models::ObjectIdentifier::from_slice(&[ + 1, 2, 840, 10045, 3, 1, 7, + ])); + }); + writer + .next() + .write_bitvec_bytes(public_key.der_bytes(), public_key.der_bytes().len() * 8); + }); + }); + let key_id = Sha256::digest(subject_public_key_info); + let mut issuer_params = CertificateParams::default(); + issuer_params.distinguished_name = params.distinguished_name.clone(); + issuer_params.key_usages = params.key_usages.clone(); + issuer_params.key_identifier_method = rcgen::KeyIdMethod::PreSpecified(key_id[..20].to_vec()); + issuer_params.is_ca = params.is_ca.clone(); + let placeholder_key = KeyPair::generate_for(&PKCS_ECDSA_P256_SHA256) + .context("failed to create self-signed certificate serialization key")?; + let issuer = issuer_params + .self_signed(&placeholder_key) + .context("failed to create self-signed issuer template")?; + let placeholder = params + .signed_by(public_key, &issuer, &placeholder_key) + .context("failed to serialize self-signed external certificate template")?; + let (_, parsed) = x509_parser::parse_x509_certificate(placeholder.der()) + .context("failed to parse self-signed external certificate template")?; + Ok(ExternalCertificate { + tbs_der: parsed.tbs_certificate.as_ref().to_vec(), + }) +} + impl CertExt for Certificate { fn get_extension_der(&self, oid: &[u64]) -> Result>> { let found = self @@ -575,6 +707,101 @@ mod tests { use scale::Encode; use std::net::IpAddr; + #[test] + fn external_certificate_roundtrip_uses_the_real_issuer_key() { + let root_key = KeyPair::generate_for(&PKCS_ECDSA_P256_SHA256).unwrap(); + let root = CertRequest::builder() + .key(&root_key) + .subject("Threshold Root") + .ca_level(1) + .build() + .self_signed() + .unwrap(); + let app_key = KeyPair::generate_for(&PKCS_ECDSA_P256_SHA256).unwrap(); + let request = CertRequest::builder() + .key(&app_key) + .subject("Dstack App CA") + .ca_level(0) + .app_id(&[7u8; 20]) + .special_usage("app:ca") + .build(); + let external = prepare_external_certificate(request, &root.pem()).unwrap(); + let signer = EcdsaKeyPair::from_pkcs8( + &ECDSA_P256_SHA256_ASN1_SIGNING, + &root_key.serialize_der(), + &SystemRandom::new(), + ) + .unwrap(); + let der_signature = signer + .sign(&SystemRandom::new(), external.tbs_der()) + .unwrap(); + let signature = p256::ecdsa::Signature::from_der(der_signature.as_ref()).unwrap(); + let pem = external.finish(signature.to_bytes().as_slice()).unwrap(); + + let (_, root_pem) = x509_parser::pem::parse_x509_pem(root.pem().as_bytes()).unwrap(); + let root_x509 = root_pem.parse_x509().unwrap(); + let (_, app_pem) = x509_parser::pem::parse_x509_pem(pem.as_bytes()).unwrap(); + let app_x509 = app_pem.parse_x509().unwrap(); + app_x509 + .verify_signature(Some(root_x509.public_key())) + .unwrap(); + } + + #[test] + fn external_self_signed_root_roundtrip() { + let root_key = KeyPair::generate_for(&PKCS_ECDSA_P256_SHA256).unwrap(); + let request = CertRequest::builder() + .key(&root_key) + .org_name("Dstack") + .subject("Dstack Threshold Root CA") + .ca_level(1) + .build(); + let external = prepare_external_self_signed(request).unwrap(); + let signer = EcdsaKeyPair::from_pkcs8( + &ECDSA_P256_SHA256_ASN1_SIGNING, + &root_key.serialize_der(), + &SystemRandom::new(), + ) + .unwrap(); + let der_signature = signer + .sign(&SystemRandom::new(), external.tbs_der()) + .unwrap(); + let signature = p256::ecdsa::Signature::from_der(der_signature.as_ref()).unwrap(); + let pem = external.finish(signature.to_bytes().as_slice()).unwrap(); + let (_, pem) = x509_parser::pem::parse_x509_pem(pem.as_bytes()).unwrap(); + let certificate = pem.parse_x509().unwrap(); + assert_eq!(certificate.subject(), certificate.issuer()); + certificate.verify_signature(None).unwrap(); + } + + #[test] + fn transported_external_tbs_roundtrip() { + let root_key = KeyPair::generate_for(&PKCS_ECDSA_P256_SHA256).unwrap(); + let external = prepare_external_self_signed( + CertRequest::builder() + .key(&root_key) + .subject("Transported Root") + .ca_level(1) + .build(), + ) + .unwrap(); + let transported = ExternalCertificate::from_tbs_der(external.tbs_der().to_vec()).unwrap(); + let signer = EcdsaKeyPair::from_pkcs8( + &ECDSA_P256_SHA256_ASN1_SIGNING, + &root_key.serialize_der(), + &SystemRandom::new(), + ) + .unwrap(); + let der = signer + .sign(&SystemRandom::new(), transported.tbs_der()) + .unwrap(); + let signature = p256::ecdsa::Signature::from_der(der.as_ref()).unwrap(); + let pem = transported.finish(signature.to_bytes().as_slice()).unwrap(); + let (_, pem) = x509_parser::pem::parse_x509_pem(pem.as_bytes()).unwrap(); + pem.parse_x509().unwrap().verify_signature(None).unwrap(); + assert!(ExternalCertificate::from_tbs_der(vec![1, 2, 3]).is_err()); + } + #[test] fn test_csr_signing_and_verification() { let key_pair = KeyPair::generate_for(&PKCS_ECDSA_P256_SHA256).unwrap();