From 1fc8c7b0e09db4c0dd031a39b4b2e9f635ee00f8 Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Thu, 30 Jul 2026 02:59:13 -0700 Subject: [PATCH 001/122] feat(kms): add stable MPC cluster identity v1 --- dstack/dstack-util/src/system_setup.rs | 6 +- dstack/kms/kms.toml | 11 ++ dstack/kms/rpc/proto/kms_rpc.proto | 5 + dstack/kms/src/config.rs | 21 +++ dstack/kms/src/main.rs | 1 + dstack/kms/src/main_service.rs | 46 +++++++ dstack/kms/src/mpc_identity.rs | 177 +++++++++++++++++++++++++ 7 files changed, 266 insertions(+), 1 deletion(-) create mode 100644 dstack/kms/src/mpc_identity.rs 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/kms.toml b/dstack/kms/kms.toml index 08e954662..9b8eff4de 100644 --- a/dstack/kms/kms.toml +++ b/dstack/kms/kms.toml @@ -97,6 +97,17 @@ 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 = "" +p256_group_pubkey = "" +k256_group_pubkey = "" +derivation_group_pubkey = "" + [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..280bff9a0 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,8 @@ 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; } message GetKmsKeyRequest { diff --git a/dstack/kms/src/config.rs b/dstack/kms/src/config.rs index 4d95b0e17..ce07e7cb2 100644 --- a/dstack/kms/src/config.rs +++ b/dstack/kms/src/config.rs @@ -74,12 +74,33 @@ 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 p256_group_pubkey: String, + #[serde(default)] + pub k256_group_pubkey: String, + #[serde(default)] + pub derivation_group_pubkey: String, +} + #[derive(Debug, Clone, Deserialize)] pub(crate) struct MetricsConfig { /// Whether to expose the unauthenticated Prometheus `/metrics` endpoint. diff --git a/dstack/kms/src/main.rs b/dstack/kms/src/main.rs index 478c5c52a..62095abc4 100644 --- a/dstack/kms/src/main.rs +++ b/dstack/kms/src/main.rs @@ -21,6 +21,7 @@ mod admin_service; mod config; // mod ct_log; mod crypto; +mod mpc_identity; mod main_service; mod onboard_service; diff --git a/dstack/kms/src/main_service.rs b/dstack/kms/src/main_service.rs index 339fbd3bc..e4382b9b1 100644 --- a/dstack/kms/src/main_service.rs +++ b/dstack/kms/src/main_service.rs @@ -25,6 +25,7 @@ use ra_tls::{ attestation::{AttestationVerifier, TeeVariant, VerifiedAttestation}, cert::{CaCert, CertRequest, CertSigningRequestV1, CertSigningRequestV2, Csr}, kdf, + rcgen::PublicKeyData as _, }; use scale::Decode; use tokio::sync::OnceCell; @@ -34,6 +35,7 @@ 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}, + mpc_identity::{decode_hex, ClusterIdentity}, }; pub(crate) mod amd_attest; @@ -62,6 +64,7 @@ pub struct KmsStateInner { attestation_verifier: Arc, self_boot_info: OnceCell, metrics: KmsMetrics, + mpc_identity: Option, } #[derive(Default)] @@ -139,6 +142,32 @@ impl KmsState { 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 mpc_identity = if config.mpc.enabled { + let identity = 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, + )?, + )?; + // Identity v1 still uses the local backend. Fail closed if config + // claims group keys different from the keys actually serving RPCs. + anyhow::ensure!( + identity.p256_group_pubkey == root_ca.key.public_key_der(), + "MPC P-256 group public key does not match the active root CA" + ); + anyhow::ensure!( + identity.k256_group_pubkey + == k256_key.verifying_key().to_sec1_bytes().as_ref(), + "MPC K-256 group public key does not match the active signing key" + ); + 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 = @@ -169,10 +198,18 @@ impl KmsState { attestation_verifier, self_boot_info: OnceCell::new(), metrics: KmsMetrics::default(), + mpc_identity, }), }) } + 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 } @@ -402,6 +439,7 @@ impl KmsRpc for RpcHandler { tproxy_app_id: gateway_app_id.clone(), gateway_app_id, os_image_hash, + key_provider_id: self.state.key_provider_id(), }) } @@ -472,6 +510,14 @@ impl KmsRpc for RpcHandler { 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")?, }) } diff --git a/dstack/kms/src/mpc_identity.rs b/dstack/kms/src/mpc_identity.rs new file mode 100644 index 000000000..974ec95e3 --- /dev/null +++ b/dstack/kms/src/mpc_identity.rs @@ -0,0 +1,177 @@ +// 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 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"; + +#[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 { + ensure!(protocol_version == 1, "unsupported MPC protocol version"); + ensure!(!cluster_id.is_empty(), "MPC cluster_id must not be empty"); + ensure!(cluster_id.len() <= 128, "MPC cluster_id is too long"); + ensure!(!p256_group_pubkey.is_empty(), "missing P-256 group public key"); + ensure!(!k256_group_pubkey.is_empty(), "missing K-256 group public key"); + ensure!( + !derivation_group_pubkey.is_empty(), + "missing derivation group public key" + ); + Ok(Self { + protocol_version, + cluster_id, + p256_group_pubkey, + k256_group_pubkey, + derivation_group_pubkey, + }) + } + + /// 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 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()); + } +} From ee6059a3f1de98c6820a652d91e8aeb3dd1c13d4 Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Thu, 30 Jul 2026 02:59:30 -0700 Subject: [PATCH 002/122] style(kms): format MPC identity implementation --- dstack/kms/src/main.rs | 2 +- dstack/kms/src/main_service.rs | 3 +-- dstack/kms/src/mpc_identity.rs | 35 +++++++++++++++++++++++++++------- 3 files changed, 30 insertions(+), 10 deletions(-) diff --git a/dstack/kms/src/main.rs b/dstack/kms/src/main.rs index 62095abc4..2839ebf01 100644 --- a/dstack/kms/src/main.rs +++ b/dstack/kms/src/main.rs @@ -21,8 +21,8 @@ mod admin_service; mod config; // mod ct_log; mod crypto; -mod mpc_identity; mod main_service; +mod mpc_identity; mod onboard_service; fn app_version() -> String { diff --git a/dstack/kms/src/main_service.rs b/dstack/kms/src/main_service.rs index e4382b9b1..052a0f144 100644 --- a/dstack/kms/src/main_service.rs +++ b/dstack/kms/src/main_service.rs @@ -160,8 +160,7 @@ impl KmsState { "MPC P-256 group public key does not match the active root CA" ); anyhow::ensure!( - identity.k256_group_pubkey - == k256_key.verifying_key().to_sec1_bytes().as_ref(), + identity.k256_group_pubkey == k256_key.verifying_key().to_sec1_bytes().as_ref(), "MPC K-256 group public key does not match the active signing key" ); Some(identity) diff --git a/dstack/kms/src/mpc_identity.rs b/dstack/kms/src/mpc_identity.rs index 974ec95e3..797c6f61e 100644 --- a/dstack/kms/src/mpc_identity.rs +++ b/dstack/kms/src/mpc_identity.rs @@ -37,8 +37,14 @@ impl ClusterIdentity { ensure!(protocol_version == 1, "unsupported MPC protocol version"); ensure!(!cluster_id.is_empty(), "MPC cluster_id must not be empty"); ensure!(cluster_id.len() <= 128, "MPC cluster_id is too long"); - ensure!(!p256_group_pubkey.is_empty(), "missing P-256 group public key"); - ensure!(!k256_group_pubkey.is_empty(), "missing K-256 group public key"); + ensure!( + !p256_group_pubkey.is_empty(), + "missing P-256 group public key" + ); + ensure!( + !k256_group_pubkey.is_empty(), + "missing K-256 group public key" + ); ensure!( !derivation_group_pubkey.is_empty(), "missing derivation group public key" @@ -85,10 +91,19 @@ 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.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"); + 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, &[ @@ -146,8 +161,14 @@ mod tests { use super::*; fn identity() -> ClusterIdentity { - ClusterIdentity::new(1, "production".into(), vec![2; 33], vec![3; 33], vec![4; 32]) - .unwrap() + ClusterIdentity::new( + 1, + "production".into(), + vec![2; 33], + vec![3; 33], + vec![4; 32], + ) + .unwrap() } #[test] From 5ea713fd0f510d7d3e2215f531b9e575f7594e26 Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Thu, 30 Jul 2026 03:00:16 -0700 Subject: [PATCH 003/122] chore(kms): remove unused MPC identity import --- dstack/kms/src/main_service.rs | 1 - 1 file changed, 1 deletion(-) diff --git a/dstack/kms/src/main_service.rs b/dstack/kms/src/main_service.rs index 052a0f144..23a776520 100644 --- a/dstack/kms/src/main_service.rs +++ b/dstack/kms/src/main_service.rs @@ -25,7 +25,6 @@ use ra_tls::{ attestation::{AttestationVerifier, TeeVariant, VerifiedAttestation}, cert::{CaCert, CertRequest, CertSigningRequestV1, CertSigningRequestV2, Csr}, kdf, - rcgen::PublicKeyData as _, }; use scale::Decode; use tokio::sync::OnceCell; From 94e217706b74f7e0dd3292db134b3bda3f1312e8 Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Thu, 30 Jul 2026 03:01:02 -0700 Subject: [PATCH 004/122] feat(kms): define canonical MPC epoch manifests --- dstack/kms/src/mpc_identity.rs | 89 ++++++++++++++++++++++++++++++++++ 1 file changed, 89 insertions(+) diff --git a/dstack/kms/src/mpc_identity.rs b/dstack/kms/src/mpc_identity.rs index 797c6f61e..7e0e9d735 100644 --- a/dstack/kms/src/mpc_identity.rs +++ b/dstack/kms/src/mpc_identity.rs @@ -13,6 +13,7 @@ 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 { @@ -73,6 +74,72 @@ impl ClusterIdentity { } } +#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] +pub(crate) struct EpochMember { + pub node_id: 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, +} + +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!( + 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.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)) + } +} + #[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] pub(crate) struct NodeEvidence { #[serde(with = "hex_bytes")] @@ -195,4 +262,26 @@ mod tests { 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(), + 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()); + } } From 1365f0bd1bd7d0c651b517cfafc09ba3986f8510 Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Thu, 30 Jul 2026 03:10:47 -0700 Subject: [PATCH 005/122] refactor(kms): introduce root key backend boundary --- dstack/kms/src/key_backend.rs | 89 ++++++++++++++++++++++++++++++++++ dstack/kms/src/main.rs | 1 + dstack/kms/src/main_service.rs | 81 +++++++++++++------------------ 3 files changed, 123 insertions(+), 48 deletions(-) create mode 100644 dstack/kms/src/key_backend.rs diff --git a/dstack/kms/src/key_backend.rs b/dstack/kms/src/key_backend.rs new file mode 100644 index 000000000..5ddffd643 --- /dev/null +++ b/dstack/kms/src/key_backend.rs @@ -0,0 +1,89 @@ +// 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 anyhow::{Context, Result}; +use k256::ecdsa::SigningKey; +use ra_tls::{kdf, rcgen::KeyPair}; + +use crate::crypto::{derive_k256_key, sign_message, sign_message_with_timestamp}; + +pub(crate) struct DerivedAppKeys { + pub disk_key: [u8; 32], + pub env_key: [u8; 32], + pub k256_key: Vec, + pub k256_signature: Vec, +} + +pub(crate) trait KeyBackend: Send + Sync { + fn derive_app_keys(&self, app_id: &[u8], instance_id: &[u8]) -> Result; + fn derive_env_key(&self, app_id: &[u8]) -> Result<[u8; 32]>; + fn sign_k256(&self, prefix: &[u8], app_id: &[u8], message: &[u8]) -> Result>; + fn sign_k256_timestamped( + &self, + prefix: &[u8], + app_id: &[u8], + timestamp: u64, + message: &[u8], + ) -> Result>; + fn k256_public_key(&self) -> Vec; +} + +pub(crate) struct LocalKeyBackend { + root_ca_key: KeyPair, + k256_key: SigningKey, +} + +impl LocalKeyBackend { + pub(crate) fn from_pem_and_bytes(root_ca_key: &str, k256_key: &[u8]) -> Result { + Ok(Self { + root_ca_key: KeyPair::from_pem(root_ca_key).context("invalid root CA key")?, + k256_key: SigningKey::from_slice(k256_key).context("invalid K-256 root key")?, + }) + } +} + +impl KeyBackend for LocalKeyBackend { + 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)?; + 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, + }) + } + + 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"]) + } + + fn sign_k256(&self, prefix: &[u8], app_id: &[u8], message: &[u8]) -> Result> { + sign_message(&self.k256_key, prefix, app_id, message) + } + + 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() + } +} diff --git a/dstack/kms/src/main.rs b/dstack/kms/src/main.rs index 2839ebf01..9001a305f 100644 --- a/dstack/kms/src/main.rs +++ b/dstack/kms/src/main.rs @@ -21,6 +21,7 @@ mod admin_service; mod config; // mod ct_log; mod crypto; +mod key_backend; mod main_service; mod mpc_identity; mod onboard_service; diff --git a/dstack/kms/src/main_service.rs b/dstack/kms/src/main_service.rs index 23a776520..2bf3bd483 100644 --- a/dstack/kms/src/main_service.rs +++ b/dstack/kms/src/main_service.rs @@ -33,7 +33,7 @@ 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}, mpc_identity::{decode_hex, ClusterIdentity}, }; @@ -64,6 +64,7 @@ pub struct KmsStateInner { self_boot_info: OnceCell, metrics: KmsMetrics, mpc_identity: Option, + key_backend: Arc, } #[derive(Default)] @@ -141,6 +142,10 @@ impl KmsState { 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 key_backend: Arc = Arc::new(LocalKeyBackend::from_pem_and_bytes( + &fs::read_to_string(config.root_ca_key())?, + &key_bytes, + )?); let mpc_identity = if config.mpc.enabled { let identity = ClusterIdentity::new( config.mpc.protocol_version, @@ -197,6 +202,7 @@ impl KmsState { self_boot_info: OnceCell::new(), metrics: KmsMetrics::default(), mpc_identity, + key_backend, }), }) } @@ -411,29 +417,18 @@ 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) + .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, + 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, @@ -446,11 +441,7 @@ 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)?; let secret = x25519_dalek::StaticSecret::from(secret); let pubkey = x25519_dalek::PublicKey::from(&secret); @@ -461,23 +452,23 @@ 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) + .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, + ) + .context("Failed to sign the public key with timestamp")?; Ok(PublicKeyResponse { public_key, @@ -495,13 +486,7 @@ impl KmsRpc for RpcHandler { Ok(GetMetaResponse { ca_cert: self.state.inner.root_ca.pem_cert.clone(), 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, From 17b8a5aed7b5d4bfb73087ca79c506f858cd72ac Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Thu, 30 Jul 2026 03:11:42 -0700 Subject: [PATCH 006/122] refactor(kms): route CA derivation and key export through backend --- dstack/kms/src/key_backend.rs | 35 ++++++++++++++++++++++++++++- dstack/kms/src/main_service.rs | 41 ++++++++++------------------------ 2 files changed, 46 insertions(+), 30 deletions(-) diff --git a/dstack/kms/src/key_backend.rs b/dstack/kms/src/key_backend.rs index 5ddffd643..7d0c06aa8 100644 --- a/dstack/kms/src/key_backend.rs +++ b/dstack/kms/src/key_backend.rs @@ -10,7 +10,11 @@ use anyhow::{Context, Result}; use k256::ecdsa::SigningKey; -use ra_tls::{kdf, rcgen::KeyPair}; +use ra_tls::{ + cert::{CaCert, CertRequest}, + kdf, + rcgen::{KeyPair, PublicKeyData as _}, +}; use crate::crypto::{derive_k256_key, sign_message, sign_message_with_timestamp}; @@ -33,6 +37,10 @@ pub(crate) trait KeyBackend: Send + Sync { message: &[u8], ) -> Result>; fn k256_public_key(&self) -> Vec; + fn p256_public_key(&self) -> Vec; + fn derive_app_ca(&self, root_ca: &CaCert, app_id: &[u8]) -> Result; + /// Legacy KMS-to-KMS migration. MPC backends must reject this operation. + fn export_root_keys(&self) -> Result<(String, Vec)>; } pub(crate) struct LocalKeyBackend { @@ -86,4 +94,29 @@ impl KeyBackend for LocalKeyBackend { 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_der() + } + + fn derive_app_ca(&self, root_ca: &CaCert, 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 = root_ca.sign(request).context("failed to sign app CA")?; + Ok(CaCert::from_parts(app_key, certificate)) + } + + fn export_root_keys(&self) -> Result<(String, Vec)> { + Ok(( + self.root_ca_key.serialize_pem(), + self.k256_key.to_bytes().to_vec(), + )) + } } diff --git a/dstack/kms/src/main_service.rs b/dstack/kms/src/main_service.rs index 2bf3bd483..9639bcbce 100644 --- a/dstack/kms/src/main_service.rs +++ b/dstack/kms/src/main_service.rs @@ -56,7 +56,6 @@ 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, @@ -160,11 +159,11 @@ impl KmsState { // Identity v1 still uses the local backend. Fail closed if config // claims group keys different from the keys actually serving RPCs. anyhow::ensure!( - identity.p256_group_pubkey == root_ca.key.public_key_der(), + 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 == k256_key.verifying_key().to_sec1_bytes().as_ref(), + identity.k256_group_pubkey == key_backend.k256_public_key(), "MPC K-256 group public key does not match the active signing key" ); Some(identity) @@ -194,7 +193,6 @@ impl KmsState { inner: Arc::new(KmsStateInner { config, root_ca, - k256_key, temp_ca_cert, temp_ca_key, verifier, @@ -371,26 +369,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 { @@ -514,12 +492,14 @@ 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() + .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 }], }) } @@ -576,7 +556,10 @@ 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(&self.state.root_ca, &app_info.boot_info.app_id)?; let cert = app_ca .sign_csr(&csr, Some(&app_info.boot_info.app_id), "app:custom") .context("Failed to sign certificate")?; From 23599acf5f5454916da8d52e6f27b6a1b5e1b006 Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Thu, 30 Jul 2026 03:12:08 -0700 Subject: [PATCH 007/122] chore(kms): remove superseded local key parsing --- dstack/kms/src/key_backend.rs | 2 +- dstack/kms/src/main_service.rs | 6 +----- 2 files changed, 2 insertions(+), 6 deletions(-) diff --git a/dstack/kms/src/key_backend.rs b/dstack/kms/src/key_backend.rs index 7d0c06aa8..de868a442 100644 --- a/dstack/kms/src/key_backend.rs +++ b/dstack/kms/src/key_backend.rs @@ -13,7 +13,7 @@ use k256::ecdsa::SigningKey; use ra_tls::{ cert::{CaCert, CertRequest}, kdf, - rcgen::{KeyPair, PublicKeyData as _}, + rcgen::KeyPair, }; use crate::crypto::{derive_k256_key, sign_message, sign_message_with_timestamp}; diff --git a/dstack/kms/src/main_service.rs b/dstack/kms/src/main_service.rs index 9639bcbce..3bf1852b5 100644 --- a/dstack/kms/src/main_service.rs +++ b/dstack/kms/src/main_service.rs @@ -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::{CaCert, CertSigningRequestV1, CertSigningRequestV2, Csr}, }; use scale::Decode; use tokio::sync::OnceCell; @@ -139,8 +137,6 @@ impl KmsState { let root_ca = CaCert::load(config.root_ca_cert(), config.root_ca_key()) .context("Failed to load root CA certificate")?; let key_bytes = fs::read(config.k256_key()).context("Failed to read ECDSA root key")?; - let k256_key = - SigningKey::from_slice(&key_bytes).context("Failed to load ECDSA root key")?; let key_backend: Arc = Arc::new(LocalKeyBackend::from_pem_and_bytes( &fs::read_to_string(config.root_ca_key())?, &key_bytes, From c3861f6cbdc48c51fad61f6b536e3d35e18ef215 Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Thu, 30 Jul 2026 03:13:07 -0700 Subject: [PATCH 008/122] feat(kms): add authenticated MPC session router --- dstack/kms/src/main.rs | 1 + dstack/kms/src/mpc_identity.rs | 6 + dstack/kms/src/mpc_session.rs | 278 +++++++++++++++++++++++++++++++++ 3 files changed, 285 insertions(+) create mode 100644 dstack/kms/src/mpc_session.rs diff --git a/dstack/kms/src/main.rs b/dstack/kms/src/main.rs index 9001a305f..557d479cf 100644 --- a/dstack/kms/src/main.rs +++ b/dstack/kms/src/main.rs @@ -24,6 +24,7 @@ mod crypto; mod key_backend; mod main_service; mod mpc_identity; +mod mpc_session; mod onboard_service; fn app_version() -> String { diff --git a/dstack/kms/src/mpc_identity.rs b/dstack/kms/src/mpc_identity.rs index 7e0e9d735..66ea481b5 100644 --- a/dstack/kms/src/mpc_identity.rs +++ b/dstack/kms/src/mpc_identity.rs @@ -138,6 +138,12 @@ impl EpochManifest { 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() + } } #[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] diff --git a/dstack/kms/src/mpc_session.rs b/dstack/kms/src/mpc_session.rs new file mode 100644 index 000000000..b4a6c88e2 --- /dev/null +++ b/dstack/kms/src/mpc_session.rs @@ -0,0 +1,278 @@ +// 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, +} + +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 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(), + 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], + } + } + + #[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 message = message(); + router.push("kms-1", message.clone()).unwrap(); + assert!(router.push("kms-1", message.clone()).is_err()); + let mut changed = message; + changed.sequence = 2; + changed.request_hash[0] ^= 1; + assert!(router.push("kms-1", changed).is_err()); + } +} From 27cd0c6b064a0c0fda67ebecb97fd659a5d41c9d Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Thu, 30 Jul 2026 03:13:53 -0700 Subject: [PATCH 009/122] feat(kms): load and validate active MPC epoch manifest --- dstack/kms/kms.toml | 4 ++++ dstack/kms/src/config.rs | 16 ++++++++++++++++ dstack/kms/src/main_service.rs | 32 +++++++++++++++++++++++++++++--- 3 files changed, 49 insertions(+), 3 deletions(-) diff --git a/dstack/kms/kms.toml b/dstack/kms/kms.toml index 9b8eff4de..759dac239 100644 --- a/dstack/kms/kms.toml +++ b/dstack/kms/kms.toml @@ -107,6 +107,10 @@ cluster_id = "" p256_group_pubkey = "" k256_group_pubkey = "" derivation_group_pubkey = "" +node_id = "" +manifest_file = "" +max_sessions = 128 +session_ttl = "5m" [core.auth_api] type = "webhook" diff --git a/dstack/kms/src/config.rs b/dstack/kms/src/config.rs index ce07e7cb2..0d975e245 100644 --- a/dstack/kms/src/config.rs +++ b/dstack/kms/src/config.rs @@ -99,6 +99,22 @@ pub(crate) struct MpcConfig { 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 = "default_mpc_max_sessions")] + pub max_sessions: usize, + #[serde(default = "default_mpc_session_ttl", with = "serde_duration")] + pub session_ttl: Duration, +} + +fn default_mpc_max_sessions() -> usize { + 128 +} + +fn default_mpc_session_ttl() -> Duration { + Duration::from_secs(300) } #[derive(Debug, Clone, Deserialize)] diff --git a/dstack/kms/src/main_service.rs b/dstack/kms/src/main_service.rs index 3bf1852b5..cfc08bc96 100644 --- a/dstack/kms/src/main_service.rs +++ b/dstack/kms/src/main_service.rs @@ -32,7 +32,9 @@ use upgrade_authority::{build_boot_info, ensure_app_id_len, local_kms_boot_info, use crate::{ config::KmsConfig, key_backend::{KeyBackend, LocalKeyBackend}, + mpc_identity::EpochManifest, mpc_identity::{decode_hex, ClusterIdentity}, + mpc_session::SessionRouter, }; pub(crate) mod amd_attest; @@ -62,6 +64,7 @@ pub struct KmsStateInner { metrics: KmsMetrics, mpc_identity: Option, key_backend: Arc, + _mpc_router: Option>, } #[derive(Default)] @@ -141,7 +144,7 @@ impl KmsState { &fs::read_to_string(config.root_ca_key())?, &key_bytes, )?); - let mpc_identity = if config.mpc.enabled { + let (mpc_identity, mpc_router) = if config.mpc.enabled { let identity = ClusterIdentity::new( config.mpc.protocol_version, config.mpc.cluster_id.clone(), @@ -162,9 +165,31 @@ impl KmsState { identity.k256_group_pubkey == key_backend.k256_public_key(), "MPC K-256 group public key does not match the active signing key" ); - Some(identity) + anyhow::ensure!( + !config.mpc.node_id.is_empty(), + "MPC node_id must not be empty" + ); + anyhow::ensure!( + !config.mpc.manifest_file.as_os_str().is_empty(), + "MPC manifest_file must not be empty" + ); + let manifest: EpochManifest = serde_json::from_slice( + &fs::read(&config.mpc.manifest_file).context("failed to read MPC manifest")?, + ) + .context("failed to parse MPC manifest")?; + 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" + ); + let router = + SessionRouter::new(manifest, config.mpc.max_sessions, config.mpc.session_ttl)?; + (Some(identity), Some(Arc::new(router))) } else { - None + (None, None) }; let temp_ca_key = fs::read_to_string(config.tmp_ca_key()).context("Faeild to read temp ca key")?; @@ -197,6 +222,7 @@ impl KmsState { metrics: KmsMetrics::default(), mpc_identity, key_backend, + _mpc_router: mpc_router, }), }) } From 8490b4a2f063389f52b891d4a8303c73045638ee Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Thu, 30 Jul 2026 03:15:07 -0700 Subject: [PATCH 010/122] feat(kms): expose attested MPC peer transport --- dstack/kms/rpc/proto/kms_rpc.proto | 19 +++++++++ dstack/kms/src/main.rs | 8 ++++ dstack/kms/src/main_service.rs | 8 +++- dstack/kms/src/mpc_identity.rs | 6 +++ dstack/kms/src/mpc_service.rs | 65 ++++++++++++++++++++++++++++++ dstack/kms/src/mpc_session.rs | 7 ++++ dstack/ra-rpc/src/lib.rs | 4 ++ dstack/ra-rpc/src/rocket_helper.rs | 5 +++ 8 files changed, 120 insertions(+), 2 deletions(-) create mode 100644 dstack/kms/src/mpc_service.rs diff --git a/dstack/kms/rpc/proto/kms_rpc.proto b/dstack/kms/rpc/proto/kms_rpc.proto index 280bff9a0..53bfea679 100644 --- a/dstack/kms/rpc/proto/kms_rpc.proto +++ b/dstack/kms/rpc/proto/kms_rpc.proto @@ -112,6 +112,25 @@ 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; +} + +// Internal interactive-protocol transport. Every method requires an +// attested RA-TLS client certificate belonging to the active epoch manifest. +service MpcTransport { + rpc Push(MpcPushRequest) returns (google.protobuf.Empty); + rpc Poll(MpcPollRequest) returns (MpcPollResponse); +} + // 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 diff --git a/dstack/kms/src/main.rs b/dstack/kms/src/main.rs index 557d479cf..b72dc02ff 100644 --- a/dstack/kms/src/main.rs +++ b/dstack/kms/src/main.rs @@ -24,6 +24,7 @@ mod crypto; mod key_backend; mod main_service; mod mpc_identity; +mod mpc_service; mod mpc_session; mod onboard_service; @@ -151,6 +152,9 @@ async fn main() -> Result<()> { 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(); @@ -183,6 +187,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 { diff --git a/dstack/kms/src/main_service.rs b/dstack/kms/src/main_service.rs index cfc08bc96..3e53892de 100644 --- a/dstack/kms/src/main_service.rs +++ b/dstack/kms/src/main_service.rs @@ -64,7 +64,7 @@ pub struct KmsStateInner { metrics: KmsMetrics, mpc_identity: Option, key_backend: Arc, - _mpc_router: Option>, + mpc_router: Option>, } #[derive(Default)] @@ -222,7 +222,7 @@ impl KmsState { metrics: KmsMetrics::default(), mpc_identity, key_backend, - _mpc_router: mpc_router, + mpc_router, }), }) } @@ -241,6 +241,10 @@ impl KmsState { 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 struct RpcHandler { diff --git a/dstack/kms/src/mpc_identity.rs b/dstack/kms/src/mpc_identity.rs index 66ea481b5..954cd28f6 100644 --- a/dstack/kms/src/mpc_identity.rs +++ b/dstack/kms/src/mpc_identity.rs @@ -144,6 +144,12 @@ impl EpochManifest { .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)] diff --git a/dstack/kms/src/mpc_service.rs b/dstack/kms/src/mpc_service.rs new file mode 100644 index 000000000..451a98ebb --- /dev/null +++ b/dstack/kms/src/mpc_service.rs @@ -0,0 +1,65 @@ +// SPDX-FileCopyrightText: © 2024-2026 Phala Network +// +// SPDX-License-Identifier: Apache-2.0 + +use anyhow::{Context, Result}; +use dstack_kms_rpc::{ + mpc_transport_server::{MpcTransport, MpcTransportServer}, + MpcPollRequest, MpcPollResponse, MpcPushRequest, +}; +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 MpcTransport for MpcHandler { + 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) + } + + async fn poll(self, request: MpcPollRequest) -> Result { + let (node_id, router) = self.peer_and_router()?; + let envelope_json = router + .drain(node_id, &request.session_id)? + .into_iter() + .map(|envelope| serde_json::to_vec(&envelope).context("failed to encode MPC envelope")) + .collect::>>()?; + Ok(MpcPollResponse { envelope_json }) + } +} + +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 index b4a6c88e2..cbfd67029 100644 --- a/dstack/kms/src/mpc_session.rs +++ b/dstack/kms/src/mpc_session.rs @@ -126,6 +126,13 @@ impl SessionRouter { 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, 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 From 25362768227f24ec3ffdda6f8d7c5b8f61ccfa3c Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Thu, 30 Jul 2026 03:15:21 -0700 Subject: [PATCH 011/122] fix(kms): implement generated MPC transport trait --- dstack/kms/src/mpc_service.rs | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/dstack/kms/src/mpc_service.rs b/dstack/kms/src/mpc_service.rs index 451a98ebb..1cbcaa1b6 100644 --- a/dstack/kms/src/mpc_service.rs +++ b/dstack/kms/src/mpc_service.rs @@ -4,7 +4,7 @@ use anyhow::{Context, Result}; use dstack_kms_rpc::{ - mpc_transport_server::{MpcTransport, MpcTransportServer}, + mpc_transport_server::{MpcTransportRpc, MpcTransportServer}, MpcPollRequest, MpcPollResponse, MpcPushRequest, }; use ra_rpc::{CallContext, RpcCall}; @@ -41,7 +41,7 @@ impl MpcHandler { } } -impl MpcTransport for MpcHandler { +impl MpcTransportRpc for MpcHandler { async fn push(self, request: MpcPushRequest) -> Result<()> { let envelope: MpcEnvelope = serde_json::from_slice(&request.envelope_json) .context("invalid MPC envelope encoding")?; From 0b92ad024953a275e377b58a97504f5d203b96de Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Thu, 30 Jul 2026 03:17:35 -0700 Subject: [PATCH 012/122] build(kms): add audited CGGMP21 threshold ECDSA engine --- dstack/Cargo.toml | 2 ++ dstack/kms/Cargo.toml | 2 ++ 2 files changed, 4 insertions(+) diff --git a/dstack/Cargo.toml b/dstack/Cargo.toml index 848809b03..d716ba205 100644 --- a/dstack/Cargo.toml +++ b/dstack/Cargo.toml @@ -185,6 +185,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/kms/Cargo.toml b/dstack/kms/Cargo.toml index 08a04bf4d..3154cdeec 100644 --- a/dstack/kms/Cargo.toml +++ b/dstack/kms/Cargo.toml @@ -49,6 +49,8 @@ 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 From fb5c4859e8d6a93c2ccd98aaa143060c744ce537 Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Thu, 30 Jul 2026 03:18:30 -0700 Subject: [PATCH 013/122] build(kms): lock CGGMP21 dependency graph --- dstack/Cargo.lock | 360 +++++++++++++++++++++++++++++++++++++++++++++- 1 file changed, 359 insertions(+), 1 deletion(-) diff --git a/dstack/Cargo.lock b/dstack/Cargo.lock index a0167f769..e23047c79 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" @@ -2024,6 +2090,7 @@ version = "0.6.0" dependencies = [ "anyhow", "cc-eventlog", + "cggmp21", "chrono", "clap", "dstack-api-auth", @@ -2047,6 +2114,7 @@ dependencies = [ "reqwest", "ring", "rocket", + "round-based", "safe-write", "serde", "serde-duration", @@ -2445,6 +2513,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 +2576,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 +2695,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 +3010,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 +3193,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 +3347,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 +3386,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 +4203,21 @@ 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", + "serde", + "serde_with 2.3.3", + "thiserror 1.0.69", +] + [[package]] name = "konst" version = "0.2.20" @@ -4962,6 +5175,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 +5392,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" @@ -5833,6 +6082,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 +6509,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 +6554,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 +7226,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 +7257,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 +8404,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" From df1240ee48bcf6132546797128ef3eab01ceb8ad Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Thu, 30 Jul 2026 03:18:48 -0700 Subject: [PATCH 014/122] test(kms): enable deterministic MPC protocol simulation --- dstack/kms/Cargo.toml | 1 + 1 file changed, 1 insertion(+) diff --git a/dstack/kms/Cargo.toml b/dstack/kms/Cargo.toml index 3154cdeec..2a58aaa49 100644 --- a/dstack/kms/Cargo.toml +++ b/dstack/kms/Cargo.toml @@ -54,6 +54,7 @@ round-based.workspace = true [dev-dependencies] cc-eventlog.workspace = true +round-based = { workspace = true, features = ["sim"] } [features] default = [] From b711f3ca43b09d2b2c8b6a45cf17ccaeea71840d Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Thu, 30 Jul 2026 03:19:43 -0700 Subject: [PATCH 015/122] feat(kms): add CGGMP21 DKG context and validated share store --- dstack/kms/Cargo.toml | 1 + dstack/kms/src/cggmp_engine.rs | 219 +++++++++++++++++++++++++++++++++ dstack/kms/src/main.rs | 1 + 3 files changed, 221 insertions(+) create mode 100644 dstack/kms/src/cggmp_engine.rs diff --git a/dstack/kms/Cargo.toml b/dstack/kms/Cargo.toml index 2a58aaa49..4143e51a8 100644 --- a/dstack/kms/Cargo.toml +++ b/dstack/kms/Cargo.toml @@ -41,6 +41,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 diff --git a/dstack/kms/src/cggmp_engine.rs b/dstack/kms/src/cggmp_engine.rs new file mode 100644 index 000000000..047d4fea1 --- /dev/null +++ b/dstack/kms/src/cggmp_engine.rs @@ -0,0 +1,219 @@ +// 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::{ + supported_curves::{Secp256k1, Secp256r1}, + IncompleteKeyShare, +}; +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; + +const SHARE_FORMAT_VERSION: u16 = 1; +const EXECUTION_ID_DOMAIN: &[u8] = b"dstack-cggmp21-execution-id-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 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()); + } +} diff --git a/dstack/kms/src/main.rs b/dstack/kms/src/main.rs index b72dc02ff..8562ce5d0 100644 --- a/dstack/kms/src/main.rs +++ b/dstack/kms/src/main.rs @@ -18,6 +18,7 @@ use tracing::{info, warn}; mod admin_auth; mod admin_service; +mod cggmp_engine; mod config; // mod ct_log; mod crypto; From 1150535eecf1b06de4a8f04893de7b4744ca4abc Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Thu, 30 Jul 2026 03:20:28 -0700 Subject: [PATCH 016/122] test(kms): verify P-256 threshold DKG share persistence --- dstack/Cargo.lock | 1 + dstack/kms/src/cggmp_engine.rs | 38 ++++++++++++++++++++++++++++++++++ 2 files changed, 39 insertions(+) diff --git a/dstack/Cargo.lock b/dstack/Cargo.lock index e23047c79..f67e8bbb4 100644 --- a/dstack/Cargo.lock +++ b/dstack/Cargo.lock @@ -2119,6 +2119,7 @@ dependencies = [ "serde", "serde-duration", "serde-human-bytes", + "serde_jcs", "serde_json", "sha2 0.10.9", "sha3", diff --git a/dstack/kms/src/cggmp_engine.rs b/dstack/kms/src/cggmp_engine.rs index 047d4fea1..ab2bf4e15 100644 --- a/dstack/kms/src/cggmp_engine.rs +++ b/dstack/kms/src/cggmp_engine.rs @@ -216,4 +216,42 @@ mod tests { ) .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 + ); + } } From 671fe8e251270d5082bb3fc5d78f429272e8334d Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Thu, 30 Jul 2026 03:22:22 -0700 Subject: [PATCH 017/122] refactor(kms): move root CA private key ownership into backend --- dstack/kms/src/key_backend.rs | 35 ++++++++++++++++++++++------------ dstack/kms/src/main_service.rs | 17 ++++++++--------- 2 files changed, 31 insertions(+), 21 deletions(-) diff --git a/dstack/kms/src/key_backend.rs b/dstack/kms/src/key_backend.rs index de868a442..61f876c5d 100644 --- a/dstack/kms/src/key_backend.rs +++ b/dstack/kms/src/key_backend.rs @@ -13,7 +13,6 @@ use k256::ecdsa::SigningKey; use ra_tls::{ cert::{CaCert, CertRequest}, kdf, - rcgen::KeyPair, }; use crate::crypto::{derive_k256_key, sign_message, sign_message_with_timestamp}; @@ -38,20 +37,25 @@ pub(crate) trait KeyBackend: Send + Sync { ) -> Result>; fn k256_public_key(&self) -> Vec; fn p256_public_key(&self) -> Vec; - fn derive_app_ca(&self, root_ca: &CaCert, app_id: &[u8]) -> Result; + fn root_ca_cert(&self) -> &str; + fn derive_app_ca(&self, app_id: &[u8]) -> Result; /// Legacy KMS-to-KMS migration. MPC backends must reject this operation. fn export_root_keys(&self) -> Result<(String, Vec)>; } pub(crate) struct LocalKeyBackend { - root_ca_key: KeyPair, + root_ca: CaCert, k256_key: SigningKey, } impl LocalKeyBackend { - pub(crate) fn from_pem_and_bytes(root_ca_key: &str, k256_key: &[u8]) -> Result { + pub(crate) fn from_pem_and_bytes( + root_ca_cert: String, + root_ca_key: String, + k256_key: &[u8], + ) -> Result { Ok(Self { - root_ca_key: KeyPair::from_pem(root_ca_key).context("invalid root CA key")?, + 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")?, }) } @@ -60,7 +64,7 @@ impl LocalKeyBackend { impl KeyBackend for LocalKeyBackend { fn derive_app_keys(&self, app_id: &[u8], instance_id: &[u8]) -> Result { let disk_key = kdf::derive_dh_secret( - &self.root_ca_key, + &self.root_ca.key, &[app_id, instance_id, b"app-disk-crypt-key"], )?; let env_key = self.derive_env_key(app_id)?; @@ -74,7 +78,7 @@ impl KeyBackend for LocalKeyBackend { } 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"]) + kdf::derive_dh_secret(&self.root_ca.key, &[app_id, b"env-encrypt-key"]) } fn sign_k256(&self, prefix: &[u8], app_id: &[u8], message: &[u8]) -> Result> { @@ -96,11 +100,15 @@ impl KeyBackend for LocalKeyBackend { } fn p256_public_key(&self) -> Vec { - self.root_ca_key.public_key_der() + self.root_ca.key.public_key_der() } - fn derive_app_ca(&self, root_ca: &CaCert, app_id: &[u8]) -> Result { - let app_key = kdf::derive_p256_key_pair(&self.root_ca_key, &[app_id, b"app-ca"])?; + fn root_ca_cert(&self) -> &str { + &self.root_ca.pem_cert + } + + 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") @@ -109,13 +117,16 @@ impl KeyBackend for LocalKeyBackend { .app_id(app_id) .special_usage("app:ca") .build(); - let certificate = root_ca.sign(request).context("failed to sign app CA")?; + let certificate = self + .root_ca + .sign(request) + .context("failed to sign app CA")?; Ok(CaCert::from_parts(app_key, certificate)) } fn export_root_keys(&self) -> Result<(String, Vec)> { Ok(( - self.root_ca_key.serialize_pem(), + self.root_ca.key.serialize_pem(), self.k256_key.to_bytes().to_vec(), )) } diff --git a/dstack/kms/src/main_service.rs b/dstack/kms/src/main_service.rs index 3e53892de..48c31c1f8 100644 --- a/dstack/kms/src/main_service.rs +++ b/dstack/kms/src/main_service.rs @@ -55,7 +55,6 @@ impl std::ops::Deref for KmsState { pub struct KmsStateInner { config: KmsConfig, - root_ca: CaCert, temp_ca_cert: String, temp_ca_key: String, verifier: CvmVerifier, @@ -137,11 +136,12 @@ 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 key_backend: Arc = Arc::new(LocalKeyBackend::from_pem_and_bytes( - &fs::read_to_string(config.root_ca_key())?, + root_ca_cert, + fs::read_to_string(config.root_ca_key())?, &key_bytes, )?); let (mpc_identity, mpc_router) = if config.mpc.enabled { @@ -213,7 +213,6 @@ impl KmsState { Ok(Self { inner: Arc::new(KmsStateInner { config, - root_ca, temp_ca_cert, temp_ca_key, verifier, @@ -428,7 +427,7 @@ impl KmsRpc for RpcHandler { .context("Failed to derive app keys")?; Ok(AppKeyResponse { - ca_cert: self.state.root_ca.pem_cert.clone(), + 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, @@ -488,7 +487,7 @@ 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.key_backend.k256_public_key(), bootstrap_info, @@ -542,7 +541,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(), }) } @@ -585,7 +584,7 @@ impl KmsRpc for RpcHandler { let app_ca = self .state .key_backend - .derive_app_ca(&self.state.root_ca, &app_info.boot_info.app_id)?; + .derive_app_ca(&app_info.boot_info.app_id)?; let cert = app_ca .sign_csr(&csr, Some(&app_info.boot_info.app_id), "app:custom") .context("Failed to sign certificate")?; @@ -593,7 +592,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(), ], }) } From b41fed1f4f2d3efe1836995c923e53ebf35a7b5c Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Thu, 30 Jul 2026 03:24:13 -0700 Subject: [PATCH 018/122] feat(kms): load only threshold shares in MPC mode --- dstack/kms/kms.toml | 2 + dstack/kms/src/cggmp_engine.rs | 4 +- dstack/kms/src/config.rs | 13 +++++ dstack/kms/src/key_backend.rs | 98 ++++++++++++++++++++++++++++++++-- dstack/kms/src/main_service.rs | 51 ++++++++++++------ 5 files changed, 148 insertions(+), 20 deletions(-) diff --git a/dstack/kms/kms.toml b/dstack/kms/kms.toml index 759dac239..13b0a6a3a 100644 --- a/dstack/kms/kms.toml +++ b/dstack/kms/kms.toml @@ -111,6 +111,8 @@ node_id = "" manifest_file = "" max_sessions = 128 session_ttl = "5m" +p256_share_file = "" +k256_share_file = "" [core.auth_api] type = "webhook" diff --git a/dstack/kms/src/cggmp_engine.rs b/dstack/kms/src/cggmp_engine.rs index ab2bf4e15..6ea0da614 100644 --- a/dstack/kms/src/cggmp_engine.rs +++ b/dstack/kms/src/cggmp_engine.rs @@ -9,7 +9,7 @@ use std::{fs::Permissions, os::unix::fs::PermissionsExt, path::Path}; use anyhow::{ensure, Context, Result}; use cggmp21::{ supported_curves::{Secp256k1, Secp256r1}, - IncompleteKeyShare, + IncompleteKeyShare, KeyShare, }; use fs_err as fs; use serde::{de::DeserializeOwned, Deserialize, Serialize}; @@ -17,6 +17,8 @@ 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"; diff --git a/dstack/kms/src/config.rs b/dstack/kms/src/config.rs index 0d975e245..6add3d060 100644 --- a/dstack/kms/src/config.rs +++ b/dstack/kms/src/config.rs @@ -107,6 +107,10 @@ pub(crate) struct MpcConfig { 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, } fn default_mpc_max_sessions() -> usize { @@ -152,6 +156,15 @@ fn default_true() -> bool { impl KmsConfig { 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.p256_share_file.exists() + && self.mpc.k256_share_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 index 61f876c5d..dfb881c31 100644 --- a/dstack/kms/src/key_backend.rs +++ b/dstack/kms/src/key_backend.rs @@ -8,14 +8,17 @@ //! implementation preserves the legacy behavior; the MPC implementation can //! replace it without making root key material available to request handlers. -use anyhow::{Context, Result}; +use anyhow::{bail, Context, Result}; use k256::ecdsa::SigningKey; use ra_tls::{ cert::{CaCert, CertRequest}, kdf, }; -use crate::crypto::{derive_k256_key, sign_message, sign_message_with_timestamp}; +use crate::{ + cggmp_engine::{load_share, CggmpCurve, K256KeyShare, P256KeyShare}, + crypto::{derive_k256_key, sign_message, sign_message_with_timestamp}, +}; pub(crate) struct DerivedAppKeys { pub disk_key: [u8; 32], @@ -100,7 +103,7 @@ impl KeyBackend for LocalKeyBackend { } fn p256_public_key(&self) -> Vec { - self.root_ca.key.public_key_der() + self.root_ca.key.public_key_raw().to_vec() } fn root_ca_cert(&self) -> &str { @@ -131,3 +134,92 @@ impl KeyBackend for LocalKeyBackend { )) } } + +/// 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, +} + +impl MpcKeyBackend { + pub(crate) fn load( + root_ca_cert: String, + cluster_id: &str, + epoch: u64, + node_id: &str, + p256_share_file: &std::path::Path, + k256_share_file: &std::path::Path, + ) -> Result { + Ok(Self { + root_ca_cert, + p256_share: load_share( + p256_share_file, + cluster_id, + epoch, + node_id, + CggmpCurve::P256, + )?, + k256_share: load_share( + k256_share_file, + cluster_id, + epoch, + node_id, + CggmpCurve::K256, + )?, + }) + } +} + +impl KeyBackend for MpcKeyBackend { + fn derive_app_keys(&self, _app_id: &[u8], _instance_id: &[u8]) -> Result { + bail!("MPC derivation protocol is unavailable") + } + + fn derive_env_key(&self, _app_id: &[u8]) -> Result<[u8; 32]> { + bail!("MPC derivation protocol is unavailable") + } + + fn sign_k256(&self, _prefix: &[u8], _app_id: &[u8], _message: &[u8]) -> Result> { + bail!("MPC signing protocol is unavailable") + } + + fn sign_k256_timestamped( + &self, + _prefix: &[u8], + _app_id: &[u8], + _timestamp: u64, + _message: &[u8], + ) -> Result> { + bail!("MPC signing protocol is unavailable") + } + + 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 root_ca_cert(&self) -> &str { + &self.root_ca_cert + } + + fn derive_app_ca(&self, _app_id: &[u8]) -> Result { + bail!("legacy app CA derivation is unavailable in MPC mode") + } + + 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_service.rs b/dstack/kms/src/main_service.rs index 48c31c1f8..b1de5464f 100644 --- a/dstack/kms/src/main_service.rs +++ b/dstack/kms/src/main_service.rs @@ -22,7 +22,7 @@ use fs_err as fs; use ra_rpc::{CallContext, RpcCall}; use ra_tls::{ attestation::{AttestationVerifier, TeeVariant, VerifiedAttestation}, - cert::{CaCert, CertSigningRequestV1, CertSigningRequestV2, Csr}, + cert::{CertSigningRequestV1, CertSigningRequestV2, Csr}, }; use scale::Decode; use tokio::sync::OnceCell; @@ -31,7 +31,7 @@ use upgrade_authority::{build_boot_info, ensure_app_id_len, local_kms_boot_info, use crate::{ config::KmsConfig, - key_backend::{KeyBackend, LocalKeyBackend}, + key_backend::{KeyBackend, LocalKeyBackend, MpcKeyBackend}, mpc_identity::EpochManifest, mpc_identity::{decode_hex, ClusterIdentity}, mpc_session::SessionRouter, @@ -138,12 +138,38 @@ impl KmsState { pub fn new(config: KmsConfig) -> Result { 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 key_backend: Arc = Arc::new(LocalKeyBackend::from_pem_and_bytes( - root_ca_cert, - fs::read_to_string(config.root_ca_key())?, - &key_bytes, - )?); + 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" + ); + Some( + serde_json::from_slice( + &fs::read(&config.mpc.manifest_file).context("failed to read MPC manifest")?, + ) + .context("failed to parse MPC manifest")?, + ) + } else { + None + }; + 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, + &config.mpc.cluster_id, + manifest.epoch, + &config.mpc.node_id, + &config.mpc.p256_share_file, + &config.mpc.k256_share_file, + )?) + } 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, mpc_router) = if config.mpc.enabled { let identity = ClusterIdentity::new( config.mpc.protocol_version, @@ -169,14 +195,7 @@ impl KmsState { !config.mpc.node_id.is_empty(), "MPC node_id must not be empty" ); - anyhow::ensure!( - !config.mpc.manifest_file.as_os_str().is_empty(), - "MPC manifest_file must not be empty" - ); - let manifest: EpochManifest = serde_json::from_slice( - &fs::read(&config.mpc.manifest_file).context("failed to read MPC manifest")?, - ) - .context("failed to parse MPC manifest")?; + let manifest = manifest.context("MPC manifest is missing")?; anyhow::ensure!( manifest.provider_id == identity.provider_id(), "MPC manifest provider ID does not match cluster identity" From e0fecd9127f8f2dc8d55ca0ef29001e8deae91cb Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Thu, 30 Jul 2026 03:24:30 -0700 Subject: [PATCH 019/122] fix(kms): expose CGGMP group public keys from shares --- dstack/kms/src/key_backend.rs | 1 + 1 file changed, 1 insertion(+) diff --git a/dstack/kms/src/key_backend.rs b/dstack/kms/src/key_backend.rs index dfb881c31..c57f3866f 100644 --- a/dstack/kms/src/key_backend.rs +++ b/dstack/kms/src/key_backend.rs @@ -9,6 +9,7 @@ //! replace it without making root key material available to request handlers. use anyhow::{bail, Context, Result}; +use cggmp21::key_share::AnyKeyShare as _; use k256::ecdsa::SigningKey; use ra_tls::{ cert::{CaCert, CertRequest}, From 20b12bdadb1ab16b3948aee98225faa5ef96fb92 Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Thu, 30 Jul 2026 03:24:48 -0700 Subject: [PATCH 020/122] build(kms): add async backend interface support --- dstack/Cargo.toml | 1 + dstack/kms/Cargo.toml | 1 + 2 files changed, 2 insertions(+) diff --git a/dstack/Cargo.toml b/dstack/Cargo.toml index d716ba205..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"] } diff --git a/dstack/kms/Cargo.toml b/dstack/kms/Cargo.toml index 4143e51a8..f233dadbe 100644 --- a/dstack/kms/Cargo.toml +++ b/dstack/kms/Cargo.toml @@ -11,6 +11,7 @@ license.workspace = true [dependencies] anyhow.workspace = true +async-trait.workspace = true chrono.workspace = true clap.workspace = true fs-err.workspace = true From 27529dcdeec631e12ac9dda8647761ed3dcd53d4 Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Thu, 30 Jul 2026 03:26:01 -0700 Subject: [PATCH 021/122] refactor(kms): make key backend operations asynchronous --- dstack/kms/src/key_backend.rs | 42 +++++++++++++++++++--------------- dstack/kms/src/main_service.rs | 13 +++++++++-- 2 files changed, 34 insertions(+), 21 deletions(-) diff --git a/dstack/kms/src/key_backend.rs b/dstack/kms/src/key_backend.rs index c57f3866f..33a30b041 100644 --- a/dstack/kms/src/key_backend.rs +++ b/dstack/kms/src/key_backend.rs @@ -9,6 +9,7 @@ //! replace it without making root key material available to request handlers. use anyhow::{bail, Context, Result}; +use async_trait::async_trait; use cggmp21::key_share::AnyKeyShare as _; use k256::ecdsa::SigningKey; use ra_tls::{ @@ -28,11 +29,12 @@ pub(crate) struct DerivedAppKeys { pub k256_signature: Vec, } +#[async_trait] pub(crate) trait KeyBackend: Send + Sync { - fn derive_app_keys(&self, app_id: &[u8], instance_id: &[u8]) -> Result; - fn derive_env_key(&self, app_id: &[u8]) -> Result<[u8; 32]>; - fn sign_k256(&self, prefix: &[u8], app_id: &[u8], message: &[u8]) -> Result>; - fn sign_k256_timestamped( + 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], @@ -42,9 +44,9 @@ pub(crate) trait KeyBackend: Send + Sync { fn k256_public_key(&self) -> Vec; fn p256_public_key(&self) -> Vec; fn root_ca_cert(&self) -> &str; - fn derive_app_ca(&self, app_id: &[u8]) -> Result; + async fn derive_app_ca(&self, app_id: &[u8]) -> Result; /// Legacy KMS-to-KMS migration. MPC backends must reject this operation. - fn export_root_keys(&self) -> Result<(String, Vec)>; + async fn export_root_keys(&self) -> Result<(String, Vec)>; } pub(crate) struct LocalKeyBackend { @@ -65,13 +67,14 @@ impl LocalKeyBackend { } } +#[async_trait] impl KeyBackend for LocalKeyBackend { - fn derive_app_keys(&self, app_id: &[u8], instance_id: &[u8]) -> Result { + 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)?; + 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, @@ -81,15 +84,15 @@ impl KeyBackend for LocalKeyBackend { }) } - fn derive_env_key(&self, app_id: &[u8]) -> Result<[u8; 32]> { + 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"]) } - fn sign_k256(&self, prefix: &[u8], app_id: &[u8], message: &[u8]) -> Result> { + async fn sign_k256(&self, prefix: &[u8], app_id: &[u8], message: &[u8]) -> Result> { sign_message(&self.k256_key, prefix, app_id, message) } - fn sign_k256_timestamped( + async fn sign_k256_timestamped( &self, prefix: &[u8], app_id: &[u8], @@ -111,7 +114,7 @@ impl KeyBackend for LocalKeyBackend { &self.root_ca.pem_cert } - fn derive_app_ca(&self, app_id: &[u8]) -> Result { + 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) @@ -128,7 +131,7 @@ impl KeyBackend for LocalKeyBackend { Ok(CaCert::from_parts(app_key, certificate)) } - fn export_root_keys(&self) -> Result<(String, Vec)> { + async fn export_root_keys(&self) -> Result<(String, Vec)> { Ok(( self.root_ca.key.serialize_pem(), self.k256_key.to_bytes().to_vec(), @@ -173,20 +176,21 @@ impl MpcKeyBackend { } } +#[async_trait] impl KeyBackend for MpcKeyBackend { - fn derive_app_keys(&self, _app_id: &[u8], _instance_id: &[u8]) -> Result { + async fn derive_app_keys(&self, _app_id: &[u8], _instance_id: &[u8]) -> Result { bail!("MPC derivation protocol is unavailable") } - fn derive_env_key(&self, _app_id: &[u8]) -> Result<[u8; 32]> { + async fn derive_env_key(&self, _app_id: &[u8]) -> Result<[u8; 32]> { bail!("MPC derivation protocol is unavailable") } - fn sign_k256(&self, _prefix: &[u8], _app_id: &[u8], _message: &[u8]) -> Result> { + async fn sign_k256(&self, _prefix: &[u8], _app_id: &[u8], _message: &[u8]) -> Result> { bail!("MPC signing protocol is unavailable") } - fn sign_k256_timestamped( + async fn sign_k256_timestamped( &self, _prefix: &[u8], _app_id: &[u8], @@ -216,11 +220,11 @@ impl KeyBackend for MpcKeyBackend { &self.root_ca_cert } - fn derive_app_ca(&self, _app_id: &[u8]) -> Result { + async fn derive_app_ca(&self, _app_id: &[u8]) -> Result { bail!("legacy app CA derivation is unavailable in MPC mode") } - fn export_root_keys(&self) -> Result<(String, Vec)> { + 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_service.rs b/dstack/kms/src/main_service.rs index b1de5464f..9dfdecd04 100644 --- a/dstack/kms/src/main_service.rs +++ b/dstack/kms/src/main_service.rs @@ -443,6 +443,7 @@ impl KmsRpc for RpcHandler { .state .key_backend .derive_app_keys(&app_id, &instance_id) + .await .context("Failed to derive app keys")?; Ok(AppKeyResponse { @@ -463,7 +464,11 @@ impl KmsRpc for RpcHandler { .await .context("KMS self authorization failed")?; ensure_app_id_len(&request.app_id)?; - let secret = self.state.key_backend.derive_env_key(&request.app_id)?; + 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); @@ -478,6 +483,7 @@ impl KmsRpc for RpcHandler { .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 @@ -490,6 +496,7 @@ impl KmsRpc for RpcHandler { timestamp, &public_key, ) + .await .context("Failed to sign the public key with timestamp")?; Ok(PublicKeyResponse { @@ -540,6 +547,7 @@ impl KmsRpc for RpcHandler { .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(), @@ -603,7 +611,8 @@ impl KmsRpc for RpcHandler { let app_ca = self .state .key_backend - .derive_app_ca(&app_info.boot_info.app_id)?; + .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")?; From db1cca5d47cb30643c3a4821ed0d9a9cb64a7f70 Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Thu, 30 Jul 2026 03:26:43 -0700 Subject: [PATCH 022/122] build(kms): lock async backend dependency --- dstack/Cargo.lock | 1 + 1 file changed, 1 insertion(+) diff --git a/dstack/Cargo.lock b/dstack/Cargo.lock index f67e8bbb4..7cc73b647 100644 --- a/dstack/Cargo.lock +++ b/dstack/Cargo.lock @@ -2089,6 +2089,7 @@ name = "dstack-kms" version = "0.6.0" dependencies = [ "anyhow", + "async-trait", "cc-eventlog", "cggmp21", "chrono", From 87e3ad241fd3dbfff72e2edab0e8d1807a83a16f Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Thu, 30 Jul 2026 03:28:20 -0700 Subject: [PATCH 023/122] feat(kms): bind MPC peer endpoints and delivery type in epoch state --- dstack/kms/src/mpc_identity.rs | 7 +++++++ dstack/kms/src/mpc_session.rs | 3 +++ 2 files changed, 10 insertions(+) diff --git a/dstack/kms/src/mpc_identity.rs b/dstack/kms/src/mpc_identity.rs index 954cd28f6..ad7960b1c 100644 --- a/dstack/kms/src/mpc_identity.rs +++ b/dstack/kms/src/mpc_identity.rs @@ -77,6 +77,7 @@ impl ClusterIdentity { #[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")] @@ -118,6 +119,10 @@ impl EpochManifest { !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" @@ -131,6 +136,7 @@ impl EpochManifest { "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); @@ -279,6 +285,7 @@ mod tests { 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], }; diff --git a/dstack/kms/src/mpc_session.rs b/dstack/kms/src/mpc_session.rs index cbfd67029..f8a63b41b 100644 --- a/dstack/kms/src/mpc_session.rs +++ b/dstack/kms/src/mpc_session.rs @@ -44,6 +44,7 @@ pub(crate) struct MpcEnvelope { pub expires_at: u64, #[serde(with = "hex_bytes")] pub payload: Vec, + pub broadcast: bool, } struct Session { @@ -231,6 +232,7 @@ mod tests { .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], }) @@ -260,6 +262,7 @@ mod tests { sequence: 1, expires_at: unix_time().unwrap() + 30, payload: vec![6], + broadcast: false, } } From 9d685eec9c9a055caa88ead57de0b8450171004c Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Thu, 30 Jul 2026 03:29:33 -0700 Subject: [PATCH 024/122] feat(kms): drive round-based MPC over attested envelopes --- dstack/kms/src/main.rs | 1 + dstack/kms/src/mpc_driver.rs | 284 +++++++++++++++++++++++++++++++++++ 2 files changed, 285 insertions(+) create mode 100644 dstack/kms/src/mpc_driver.rs diff --git a/dstack/kms/src/main.rs b/dstack/kms/src/main.rs index 8562ce5d0..d23c05fa0 100644 --- a/dstack/kms/src/main.rs +++ b/dstack/kms/src/main.rs @@ -24,6 +24,7 @@ mod config; mod crypto; mod key_backend; mod main_service; +mod mpc_driver; mod mpc_identity; mod mpc_service; mod mpc_session; diff --git a/dstack/kms/src/mpc_driver.rs b/dstack/kms/src/mpc_driver.rs new file mode 100644 index 000000000..334fc8f41 --- /dev/null +++ b/dstack/kms/src/mpc_driver.rs @@ -0,0 +1,284 @@ +// 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 crate::mpc_session::{MpcEnvelope, MpcProtocol}; + +#[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) async fn drive_state_machine( + mut state: S, + transport: &T, + context: DriverContext, +) -> Result +where + S: StateMachine + Send, + S::Msg: Serialize + DeserializeOwned + Send, + S::Output: Send, + 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}"), + } + } +} + +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::{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(), + } + } + + 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::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 + })); + } +} From 2ff0644023f09335284cc4895e00938043fe20b4 Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Thu, 30 Jul 2026 03:29:48 -0700 Subject: [PATCH 025/122] fix(kms): allow local CGGMP state-machine execution --- dstack/kms/src/mpc_driver.rs | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/dstack/kms/src/mpc_driver.rs b/dstack/kms/src/mpc_driver.rs index 334fc8f41..709303261 100644 --- a/dstack/kms/src/mpc_driver.rs +++ b/dstack/kms/src/mpc_driver.rs @@ -52,9 +52,8 @@ pub(crate) async fn drive_state_machine( context: DriverContext, ) -> Result where - S: StateMachine + Send, - S::Msg: Serialize + DeserializeOwned + Send, - S::Output: Send, + S: StateMachine, + S::Msg: Serialize + DeserializeOwned, T: EnvelopeTransport, { ensure!( @@ -225,7 +224,7 @@ mod tests { #[tokio::test] async fn drives_real_cggmp_dkg_over_authenticated_envelopes() { let manifest = manifest(); - let routers = Arc::new( + let routers: Arc>> = Arc::new( manifest .members .iter() From 0da8fb51b1989af9904a86821a6ec5f11a5ae0e8 Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Thu, 30 Jul 2026 03:31:50 -0700 Subject: [PATCH 026/122] feat(kms): add RA-TLS verified MPC HTTP delivery --- dstack/kms/src/key_backend.rs | 20 +++++++ dstack/kms/src/main_service.rs | 34 ++++++++---- dstack/kms/src/mpc_driver.rs | 95 +++++++++++++++++++++++++++++++++- 3 files changed, 138 insertions(+), 11 deletions(-) diff --git a/dstack/kms/src/key_backend.rs b/dstack/kms/src/key_backend.rs index 33a30b041..82d627646 100644 --- a/dstack/kms/src/key_backend.rs +++ b/dstack/kms/src/key_backend.rs @@ -20,6 +20,9 @@ use ra_tls::{ use crate::{ cggmp_engine::{load_share, CggmpCurve, K256KeyShare, P256KeyShare}, crypto::{derive_k256_key, sign_message, sign_message_with_timestamp}, + mpc_driver::MpcHttpTransport, + mpc_identity::EpochManifest, + mpc_session::SessionRouter, }; pub(crate) struct DerivedAppKeys { @@ -145,6 +148,8 @@ pub(crate) struct MpcKeyBackend { root_ca_cert: String, p256_share: P256KeyShare, k256_share: K256KeyShare, + #[allow(dead_code)] + transport: MpcHttpTransport, } impl MpcKeyBackend { @@ -155,7 +160,21 @@ impl MpcKeyBackend { node_id: &str, p256_share_file: &std::path::Path, k256_share_file: &std::path::Path, + manifest: &EpochManifest, + router: std::sync::Arc, + rpc_cert: String, + rpc_key: String, + attestation_verifier: std::sync::Arc, ) -> Result { + let transport = MpcHttpTransport::new( + manifest, + node_id, + router, + rpc_cert, + rpc_key, + root_ca_cert.clone(), + attestation_verifier, + )?; Ok(Self { root_ca_cert, p256_share: load_share( @@ -172,6 +191,7 @@ impl MpcKeyBackend { node_id, CggmpCurve::K256, )?, + transport, }) } } diff --git a/dstack/kms/src/main_service.rs b/dstack/kms/src/main_service.rs index 9dfdecd04..e7aeb82b0 100644 --- a/dstack/kms/src/main_service.rs +++ b/dstack/kms/src/main_service.rs @@ -138,6 +138,10 @@ impl KmsState { pub fn new(config: KmsConfig) -> Result { let root_ca_cert = fs::read_to_string(config.root_ca_cert()) .context("Failed to load root CA certificate")?; + let attestation_verifier = Arc::new( + AttestationVerifier::load(&config.attestation) + .context("failed to load attestation verifier")?, + ); let manifest: Option = if config.mpc.enabled { anyhow::ensure!( !config.mpc.manifest_file.as_os_str().is_empty(), @@ -152,6 +156,17 @@ impl KmsState { } 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()?; let key_backend: Arc = if config.mpc.enabled { let manifest = manifest.as_ref().context("MPC manifest is missing")?; Arc::new(MpcKeyBackend::load( @@ -161,6 +176,11 @@ impl KmsState { &config.mpc.node_id, &config.mpc.p256_share_file, &config.mpc.k256_share_file, + manifest, + mpc_router.clone().context("MPC router is missing")?, + fs::read_to_string(config.rpc_cert()).context("failed to read MPC RPC cert")?, + fs::read_to_string(config.rpc_key()).context("failed to read MPC RPC key")?, + attestation_verifier.clone(), )?) } else { let key_bytes = fs::read(config.k256_key()).context("Failed to read ECDSA root key")?; @@ -170,7 +190,7 @@ impl KmsState { &key_bytes, )?) }; - let (mpc_identity, mpc_router) = if config.mpc.enabled { + let mpc_identity = if config.mpc.enabled { let identity = ClusterIdentity::new( config.mpc.protocol_version, config.mpc.cluster_id.clone(), @@ -195,7 +215,7 @@ impl KmsState { !config.mpc.node_id.is_empty(), "MPC node_id must not be empty" ); - let manifest = manifest.context("MPC manifest is missing")?; + 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" @@ -204,20 +224,14 @@ impl KmsState { manifest.contains_member(&config.mpc.node_id), "local node is not a member of the MPC epoch" ); - let router = - SessionRouter::new(manifest, config.mpc.max_sessions, config.mpc.session_ttl)?; - (Some(identity), Some(Arc::new(router))) + Some(identity) } else { - (None, None) + 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 attestation_verifier = Arc::new( - AttestationVerifier::load(&config.attestation) - .context("failed to load attestation verifier")?, - ); let verifier = CvmVerifier::new( config.image.cache_dir.display().to_string(), config.image.download_url.clone(), diff --git a/dstack/kms/src/mpc_driver.rs b/dstack/kms/src/mpc_driver.rs index 709303261..a7ddc0fbb 100644 --- a/dstack/kms/src/mpc_driver.rs +++ b/dstack/kms/src/mpc_driver.rs @@ -14,7 +14,15 @@ use round_based::{ }; use serde::{de::DeserializeOwned, Serialize}; -use crate::mpc_session::{MpcEnvelope, MpcProtocol}; +use dstack_kms_rpc::{mpc_transport_client::MpcTransportClient, MpcPushRequest}; +use ra_rpc::client::{RaClient, RaClientConfig}; +use ra_tls::attestation::AttestationVerifier; +use x509_parser::prelude::FromDer as _; + +use crate::{ + mpc_identity::EpochManifest, + mpc_session::{MpcEnvelope, MpcProtocol, SessionRouter}, +}; #[derive(Clone)] pub(crate) struct DriverContext { @@ -46,6 +54,91 @@ pub(crate) trait EnvelopeTransport: Send + Sync { async fn receive(&self, session_id: &[u8; 32]) -> Result>; } +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, + }) + } +} + +#[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, From d8000f28c83dae06537c2f7981036f290af0082a Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Thu, 30 Jul 2026 03:33:46 -0700 Subject: [PATCH 027/122] chore(kms): remove unused MPC certificate import --- dstack/kms/src/mpc_driver.rs | 1 - 1 file changed, 1 deletion(-) diff --git a/dstack/kms/src/mpc_driver.rs b/dstack/kms/src/mpc_driver.rs index a7ddc0fbb..6078ec596 100644 --- a/dstack/kms/src/mpc_driver.rs +++ b/dstack/kms/src/mpc_driver.rs @@ -17,7 +17,6 @@ use serde::{de::DeserializeOwned, Serialize}; use dstack_kms_rpc::{mpc_transport_client::MpcTransportClient, MpcPushRequest}; use ra_rpc::client::{RaClient, RaClientConfig}; use ra_tls::attestation::AttestationVerifier; -use x509_parser::prelude::FromDer as _; use crate::{ mpc_identity::EpochManifest, From 73f1a6412c3901b77d500135c8501456c8c32e11 Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Thu, 30 Jul 2026 03:34:48 -0700 Subject: [PATCH 028/122] feat(kms): add quorum operation start RPC --- dstack/kms/rpc/proto/kms_rpc.proto | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/dstack/kms/rpc/proto/kms_rpc.proto b/dstack/kms/rpc/proto/kms_rpc.proto index 53bfea679..9db247c96 100644 --- a/dstack/kms/rpc/proto/kms_rpc.proto +++ b/dstack/kms/rpc/proto/kms_rpc.proto @@ -124,9 +124,18 @@ 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 Start(MpcStartRequest) returns (MpcStartResponse); rpc Push(MpcPushRequest) returns (google.protobuf.Empty); rpc Poll(MpcPollRequest) returns (MpcPollResponse); } From e85fc9351ca3c399505e6c899994eee46712b79d Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Thu, 30 Jul 2026 03:34:54 -0700 Subject: [PATCH 029/122] build(kms): enable concurrent quorum coordination --- dstack/kms/Cargo.toml | 1 + 1 file changed, 1 insertion(+) diff --git a/dstack/kms/Cargo.toml b/dstack/kms/Cargo.toml index f233dadbe..70c10c091 100644 --- a/dstack/kms/Cargo.toml +++ b/dstack/kms/Cargo.toml @@ -15,6 +15,7 @@ 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 From 055a33dde3a0d5d05c1d92de953a0aaeeb0bb045 Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Thu, 30 Jul 2026 03:35:47 -0700 Subject: [PATCH 030/122] feat(kms): define policy-scoped quorum operations --- dstack/kms/src/main.rs | 1 + dstack/kms/src/mpc_operation.rs | 241 ++++++++++++++++++++++++++++++++ 2 files changed, 242 insertions(+) create mode 100644 dstack/kms/src/mpc_operation.rs diff --git a/dstack/kms/src/main.rs b/dstack/kms/src/main.rs index d23c05fa0..5036cf4a2 100644 --- a/dstack/kms/src/main.rs +++ b/dstack/kms/src/main.rs @@ -26,6 +26,7 @@ mod key_backend; mod main_service; mod mpc_driver; mod mpc_identity; +mod mpc_operation; mod mpc_service; mod mpc_session; mod onboard_service; diff --git a/dstack/kms/src/mpc_operation.rs b/dstack/kms/src/mpc_operation.rs new file mode 100644 index 000000000..149afd011 --- /dev/null +++ b/dstack/kms/src/mpc_operation.rs @@ -0,0 +1,241 @@ +// 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; + +const REQUEST_DOMAIN: &[u8] = b"dstack-mpc-operation-v1"; +const MAX_SIGN_MESSAGE_BYTES: usize = 1024 * 1024; + +#[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, +} + +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), +} + +#[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 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()) + .context("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" + ); + ensure!(self.expires_at >= unix_time()?, "MPC operation expired"); + match &self.payload { + MpcOperationPayload::SignK256(payload) => payload.validate()?, + } + 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) + } +} + +#[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()); + } +} From bece3303f1a125bb221f9a4ef3499ab766d6935d Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Thu, 30 Jul 2026 03:37:38 -0700 Subject: [PATCH 031/122] feat(kms): drive CGGMP operations on dedicated threads --- dstack/kms/src/mpc_driver.rs | 117 +++++++++++++++++++++++++++++++++++ 1 file changed, 117 insertions(+) diff --git a/dstack/kms/src/mpc_driver.rs b/dstack/kms/src/mpc_driver.rs index 6078ec596..efcfa9498 100644 --- a/dstack/kms/src/mpc_driver.rs +++ b/dstack/kms/src/mpc_driver.rs @@ -53,6 +53,28 @@ pub(crate) trait EnvelopeTransport: Send + Sync { async fn receive(&self, session_id: &[u8; 32]) -> Result>; } +pub(crate) struct BlockingHttpTransport { + inner: std::sync::Arc, + runtime: tokio::runtime::Handle, +} + +impl BlockingHttpTransport { + 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, @@ -234,6 +256,101 @@ where } } +/// 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: &BlockingHttpTransport, + context: DriverContext, +) -> Result +where + S: StateMachine, + S::Msg: Serialize + DeserializeOwned, +{ + 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) From 13bb254d9a6c9d4fcdd29abff563fbe6c4303618 Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Thu, 30 Jul 2026 03:38:33 -0700 Subject: [PATCH 032/122] fix(kms): wire authenticated MPC operation dispatch --- dstack/kms/src/key_backend.rs | 24 ++++++++++++++++++++++++ dstack/kms/src/main_service.rs | 4 ++++ dstack/kms/src/mpc_operation.rs | 2 +- dstack/kms/src/mpc_service.rs | 15 ++++++++++++++- 4 files changed, 43 insertions(+), 2 deletions(-) diff --git a/dstack/kms/src/key_backend.rs b/dstack/kms/src/key_backend.rs index 82d627646..a6a512376 100644 --- a/dstack/kms/src/key_backend.rs +++ b/dstack/kms/src/key_backend.rs @@ -48,6 +48,11 @@ pub(crate) trait KeyBackend: Send + Sync { fn p256_public_key(&self) -> Vec; fn root_ca_cert(&self) -> &str; async fn derive_app_ca(&self, app_id: &[u8]) -> 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)>; } @@ -140,6 +145,14 @@ impl KeyBackend for LocalKeyBackend { 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 @@ -150,6 +163,7 @@ pub(crate) struct MpcKeyBackend { k256_share: K256KeyShare, #[allow(dead_code)] transport: MpcHttpTransport, + manifest: EpochManifest, } impl MpcKeyBackend { @@ -192,6 +206,7 @@ impl MpcKeyBackend { CggmpCurve::K256, )?, transport, + manifest: manifest.clone(), }) } } @@ -244,6 +259,15 @@ impl KeyBackend for MpcKeyBackend { bail!("legacy app CA derivation is unavailable in MPC mode") } + async fn run_mpc_operation( + &self, + operation: crate::mpc_operation::MpcOperation, + initiator: &str, + ) -> Result> { + operation.validate(&self.manifest, initiator)?; + bail!("MPC signing protocol is unavailable") + } + 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_service.rs b/dstack/kms/src/main_service.rs index e7aeb82b0..b58135b85 100644 --- a/dstack/kms/src/main_service.rs +++ b/dstack/kms/src/main_service.rs @@ -277,6 +277,10 @@ impl KmsState { 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 { diff --git a/dstack/kms/src/mpc_operation.rs b/dstack/kms/src/mpc_operation.rs index 149afd011..9611f06a0 100644 --- a/dstack/kms/src/mpc_operation.rs +++ b/dstack/kms/src/mpc_operation.rs @@ -127,7 +127,7 @@ impl MpcOperation { manifest .members .binary_search_by_key(&node.as_str(), |member| member.node_id.as_str()) - .context("MPC operation contains a non-member") + .map_err(|_| anyhow::anyhow!("MPC operation contains a non-member")) }) .collect::>>()?; ensure!( diff --git a/dstack/kms/src/mpc_service.rs b/dstack/kms/src/mpc_service.rs index 1cbcaa1b6..6bcbcc685 100644 --- a/dstack/kms/src/mpc_service.rs +++ b/dstack/kms/src/mpc_service.rs @@ -5,7 +5,7 @@ use anyhow::{Context, Result}; use dstack_kms_rpc::{ mpc_transport_server::{MpcTransportRpc, MpcTransportServer}, - MpcPollRequest, MpcPollResponse, MpcPushRequest, + MpcPollRequest, MpcPollResponse, MpcPushRequest, MpcStartRequest, MpcStartResponse, }; use ra_rpc::{CallContext, RpcCall}; @@ -42,6 +42,19 @@ impl MpcHandler { } impl MpcTransportRpc for MpcHandler { + 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")?; From af587c968392d953bb231387a10fe995c34debe0 Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Thu, 30 Jul 2026 03:40:51 -0700 Subject: [PATCH 033/122] feat(kms): execute threshold K-256 signing quorums --- dstack/kms/src/key_backend.rs | 226 ++++++++++++++++++++++++++++++---- dstack/kms/src/mpc_driver.rs | 19 ++- 2 files changed, 223 insertions(+), 22 deletions(-) diff --git a/dstack/kms/src/key_backend.rs b/dstack/kms/src/key_backend.rs index a6a512376..64caca04e 100644 --- a/dstack/kms/src/key_backend.rs +++ b/dstack/kms/src/key_backend.rs @@ -8,21 +8,29 @@ //! implementation preserves the legacy behavior; the MPC implementation can //! replace it without making root key material available to request handlers. -use anyhow::{bail, Context, Result}; +use std::{sync::Arc, time::Duration}; + +use anyhow::{bail, ensure, Context, Result}; use async_trait::async_trait; use cggmp21::key_share::AnyKeyShare as _; -use k256::ecdsa::SigningKey; +use cggmp21::{generic_ec::Scalar, supported_curves::Secp256k1, DataToSign, ExecutionId}; +use futures::future::join_all; +use k256::ecdsa::{RecoveryId, Signature as K256Signature, SigningKey, VerifyingKey}; use ra_tls::{ cert::{CaCert, CertRequest}, kdf, }; +use rand::{rngs::OsRng, RngCore}; use crate::{ - cggmp_engine::{load_share, CggmpCurve, K256KeyShare, P256KeyShare}, + cggmp_engine::{execution_id, load_share, CggmpCurve, K256KeyShare, P256KeyShare}, crypto::{derive_k256_key, sign_message, sign_message_with_timestamp}, - mpc_driver::MpcHttpTransport, + mpc_driver::{ + drive_state_machine_blocking, BlockingHttpTransport, DriverContext, MpcHttpTransport, + }, mpc_identity::EpochManifest, - mpc_session::SessionRouter, + mpc_operation::{K256SignPayload, MpcOperation, MpcOperationPayload}, + mpc_session::{MpcProtocol, SessionRouter}, }; pub(crate) struct DerivedAppKeys { @@ -161,9 +169,10 @@ pub(crate) struct MpcKeyBackend { root_ca_cert: String, p256_share: P256KeyShare, k256_share: K256KeyShare, - #[allow(dead_code)] - transport: MpcHttpTransport, + transport: Arc, manifest: EpochManifest, + cluster_id: String, + node_id: String, } impl MpcKeyBackend { @@ -205,10 +214,101 @@ impl MpcKeyBackend { node_id, CggmpCurve::K256, )?, - transport, + transport: Arc::new(transport), manifest: manifest.clone(), + cluster_id: cluster_id.into(), + node_id: node_id.into(), }) } + + fn participants(&self) -> Result> { + let threshold = usize::from(self.manifest.threshold); + 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) + .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()?, + 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) + } + + 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(), + }; + 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") + } +} + +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] @@ -221,18 +321,30 @@ impl KeyBackend for MpcKeyBackend { bail!("MPC derivation protocol is unavailable") } - async fn sign_k256(&self, _prefix: &[u8], _app_id: &[u8], _message: &[u8]) -> Result> { - bail!("MPC signing protocol is unavailable") + 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], + prefix: &[u8], + app_id: &[u8], + timestamp: u64, + message: &[u8], ) -> Result> { - bail!("MPC signing protocol is unavailable") + 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 { @@ -259,13 +371,85 @@ impl KeyBackend for MpcKeyBackend { bail!("legacy app CA derivation is unavailable in MPC mode") } - async fn run_mpc_operation( - &self, - operation: crate::mpc_operation::MpcOperation, - initiator: &str, - ) -> Result> { + async fn run_mpc_operation(&self, operation: MpcOperation, initiator: &str) -> Result> { operation.validate(&self.manifest, initiator)?; - bail!("MPC signing protocol is unavailable") + let MpcOperationPayload::SignK256(payload) = &operation.payload; + 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.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(); + 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, + &BlockingHttpTransport::new(transport, tokio::runtime::Handle::current()), + context, + )??; + let mut encoded = vec![0u8; signature.serialized_len()]; + signature.write_to_slice(&mut encoded); + Ok(encoded) + }) + .await + .context("MPC signing worker panicked")? } async fn export_root_keys(&self) -> Result<(String, Vec)> { diff --git a/dstack/kms/src/mpc_driver.rs b/dstack/kms/src/mpc_driver.rs index efcfa9498..45ffa19c4 100644 --- a/dstack/kms/src/mpc_driver.rs +++ b/dstack/kms/src/mpc_driver.rs @@ -14,7 +14,7 @@ use round_based::{ }; use serde::{de::DeserializeOwned, Serialize}; -use dstack_kms_rpc::{mpc_transport_client::MpcTransportClient, MpcPushRequest}; +use dstack_kms_rpc::{mpc_transport_client::MpcTransportClient, MpcPushRequest, MpcStartRequest}; use ra_rpc::client::{RaClient, RaClientConfig}; use ra_tls::attestation::AttestationVerifier; @@ -135,6 +135,23 @@ impl MpcHttpTransport { 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) + } } #[async_trait] From 7e22abd19a259273093d8f531c1799ddb7f3192c Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Thu, 30 Jul 2026 03:41:03 -0700 Subject: [PATCH 034/122] fix(kms): size threshold signatures by curve --- dstack/kms/src/key_backend.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/dstack/kms/src/key_backend.rs b/dstack/kms/src/key_backend.rs index 64caca04e..24fa7d18d 100644 --- a/dstack/kms/src/key_backend.rs +++ b/dstack/kms/src/key_backend.rs @@ -444,7 +444,7 @@ impl KeyBackend for MpcKeyBackend { &BlockingHttpTransport::new(transport, tokio::runtime::Handle::current()), context, )??; - let mut encoded = vec![0u8; signature.serialized_len()]; + let mut encoded = vec![0u8; cggmp21::Signature::::serialized_len()]; signature.write_to_slice(&mut encoded); Ok(encoded) }) From c8f9c728d717b40a165d0e46b257abba9c413d7c Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Thu, 30 Jul 2026 03:41:39 -0700 Subject: [PATCH 035/122] build(kms): lock quorum coordination dependency --- dstack/Cargo.lock | 1 + 1 file changed, 1 insertion(+) diff --git a/dstack/Cargo.lock b/dstack/Cargo.lock index 7cc73b647..759a23481 100644 --- a/dstack/Cargo.lock +++ b/dstack/Cargo.lock @@ -2103,6 +2103,7 @@ dependencies = [ "dstack-types", "dstack-verifier", "fs-err", + "futures", "hex", "hex_fmt", "http-client", From c4a1b5c27139a10725ae4d4d30c102cc527329f8 Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Thu, 30 Jul 2026 03:44:39 -0700 Subject: [PATCH 036/122] fix(kms): make quorum operation starts replay-safe --- dstack/kms/src/key_backend.rs | 211 +++++++++++++++++++++----------- dstack/kms/src/mpc_operation.rs | 8 +- 2 files changed, 148 insertions(+), 71 deletions(-) diff --git a/dstack/kms/src/key_backend.rs b/dstack/kms/src/key_backend.rs index 24fa7d18d..8bd415537 100644 --- a/dstack/kms/src/key_backend.rs +++ b/dstack/kms/src/key_backend.rs @@ -8,7 +8,11 @@ //! implementation preserves the legacy behavior; the MPC implementation can //! replace it without making root key material available to request handlers. -use std::{sync::Arc, time::Duration}; +use std::{ + collections::BTreeMap, + sync::{Arc, Mutex}, + time::Duration, +}; use anyhow::{bail, ensure, Context, Result}; use async_trait::async_trait; @@ -173,6 +177,15 @@ pub(crate) struct MpcKeyBackend { manifest: EpochManifest, cluster_id: String, node_id: String, + operations: Mutex, Arc>>, +} + +struct OperationRecord { + request_hash: Vec, + initiator: String, + expires_at: u64, + result: Mutex, String>>>, + finished: tokio::sync::Notify, } impl MpcKeyBackend { @@ -218,6 +231,7 @@ impl MpcKeyBackend { manifest: manifest.clone(), cluster_id: cluster_id.into(), node_id: node_id.into(), + operations: Mutex::new(BTreeMap::new()), }) } @@ -302,77 +316,60 @@ impl MpcKeyBackend { } bail!("MPC signature cannot be recovered to the group public key") } -} - -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 { - bail!("MPC derivation protocol is unavailable") - } - - async fn derive_env_key(&self, _app_id: &[u8]) -> Result<[u8; 32]> { - bail!("MPC derivation protocol is unavailable") - } - - 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 root_ca_cert(&self) -> &str { - &self.root_ca_cert - } - - async fn derive_app_ca(&self, _app_id: &[u8]) -> Result { - bail!("legacy app CA derivation is unavailable in MPC mode") + 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 { + 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_k256(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 run_mpc_operation(&self, operation: MpcOperation, initiator: &str) -> Result> { - operation.validate(&self.manifest, initiator)?; + async fn execute_k256(&self, operation: MpcOperation) -> Result> { let MpcOperationPayload::SignK256(payload) = &operation.payload; let digest = payload.digest(); let local_protocol_index: u16 = operation @@ -428,6 +425,7 @@ impl KeyBackend for MpcKeyBackend { }; 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 = @@ -441,7 +439,7 @@ impl KeyBackend for MpcKeyBackend { .sign_sync(&mut rng, data); let signature = drive_state_machine_blocking( state, - &BlockingHttpTransport::new(transport, tokio::runtime::Handle::current()), + &BlockingHttpTransport::new(transport, runtime), context, )??; let mut encoded = vec![0u8; cggmp21::Signature::::serialized_len()]; @@ -451,6 +449,79 @@ impl KeyBackend for MpcKeyBackend { .await .context("MPC 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 { + bail!("MPC derivation protocol is unavailable") + } + + async fn derive_env_key(&self, _app_id: &[u8]) -> Result<[u8; 32]> { + bail!("MPC derivation protocol is unavailable") + } + + 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 root_ca_cert(&self) -> &str { + &self.root_ca_cert + } + + async fn derive_app_ca(&self, _app_id: &[u8]) -> Result { + bail!("legacy app CA derivation is unavailable in MPC mode") + } + + 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/mpc_operation.rs b/dstack/kms/src/mpc_operation.rs index 9611f06a0..56c9d2da4 100644 --- a/dstack/kms/src/mpc_operation.rs +++ b/dstack/kms/src/mpc_operation.rs @@ -13,6 +13,7 @@ use crate::mpc_identity::EpochManifest; 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 { @@ -134,7 +135,12 @@ impl MpcOperation { indexes.windows(2).all(|pair| pair[0] < pair[1]), "MPC participants must be unique and in manifest order" ); - ensure!(self.expires_at >= unix_time()?, "MPC operation expired"); + 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()?, } From ff1c8d50386476e75677089992595b4c90baabe1 Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Thu, 30 Jul 2026 03:46:14 -0700 Subject: [PATCH 037/122] feat(kms): add verifiable threshold PRF primitive --- dstack/kms/Cargo.toml | 1 + dstack/kms/src/main.rs | 1 + dstack/kms/src/threshold_prf.rs | 286 ++++++++++++++++++++++++++++++++ 3 files changed, 288 insertions(+) create mode 100644 dstack/kms/src/threshold_prf.rs diff --git a/dstack/kms/Cargo.toml b/dstack/kms/Cargo.toml index 70c10c091..b6ea134cd 100644 --- a/dstack/kms/Cargo.toml +++ b/dstack/kms/Cargo.toml @@ -35,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"] } diff --git a/dstack/kms/src/main.rs b/dstack/kms/src/main.rs index 5036cf4a2..7b9120b39 100644 --- a/dstack/kms/src/main.rs +++ b/dstack/kms/src/main.rs @@ -30,6 +30,7 @@ mod mpc_operation; mod mpc_service; mod mpc_session; mod onboard_service; +mod threshold_prf; fn app_version() -> String { dstack_build_info::app_version!() diff --git a/dstack/kms/src/threshold_prf.rs b/dstack/kms/src/threshold_prf.rs new file mode 100644 index 000000000..f9a60330a --- /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()); + } +} From 91cb1aea228df8567197242405fa49e28b866edb Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Thu, 30 Jul 2026 03:46:26 -0700 Subject: [PATCH 038/122] fix(kms): interpolate owned PRF evaluations --- dstack/kms/src/threshold_prf.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/dstack/kms/src/threshold_prf.rs b/dstack/kms/src/threshold_prf.rs index f9a60330a..4000e345e 100644 --- a/dstack/kms/src/threshold_prf.rs +++ b/dstack/kms/src/threshold_prf.rs @@ -148,7 +148,7 @@ pub(crate) fn combine( } let inverse = Option::::from(denominator.invert()) .context("invalid threshold PRF interpolation set")?; - combined += evaluation * (numerator * inverse); + combined += *evaluation * (numerator * inverse); } ensure!( !bool::from(combined.is_identity()), From d06bcd9f0e7d2d5acef9dcd346d2633346e065b7 Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Thu, 30 Jul 2026 03:49:00 -0700 Subject: [PATCH 039/122] feat(kms): integrate verified threshold key derivation --- dstack/kms/kms.toml | 1 + dstack/kms/src/config.rs | 5 +- dstack/kms/src/key_backend.rs | 179 ++++++++++++++++++++++++++++++-- dstack/kms/src/main_service.rs | 5 + dstack/kms/src/mpc_operation.rs | 90 ++++++++++++++++ 5 files changed, 273 insertions(+), 7 deletions(-) diff --git a/dstack/kms/kms.toml b/dstack/kms/kms.toml index 13b0a6a3a..fea00f02c 100644 --- a/dstack/kms/kms.toml +++ b/dstack/kms/kms.toml @@ -113,6 +113,7 @@ 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/src/config.rs b/dstack/kms/src/config.rs index 6add3d060..d272d5e13 100644 --- a/dstack/kms/src/config.rs +++ b/dstack/kms/src/config.rs @@ -111,6 +111,8 @@ pub(crate) struct MpcConfig { 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 { @@ -163,7 +165,8 @@ impl KmsConfig { && self.rpc_cert().exists() && self.rpc_key().exists() && self.mpc.p256_share_file.exists() - && self.mpc.k256_share_file.exists(); + && self.mpc.k256_share_file.exists() + && self.mpc.derivation_share_file.exists(); } self.tmp_ca_cert().exists() && self.tmp_ca_key().exists() diff --git a/dstack/kms/src/key_backend.rs b/dstack/kms/src/key_backend.rs index 8bd415537..93b852750 100644 --- a/dstack/kms/src/key_backend.rs +++ b/dstack/kms/src/key_backend.rs @@ -33,8 +33,11 @@ use crate::{ drive_state_machine_blocking, BlockingHttpTransport, DriverContext, MpcHttpTransport, }, mpc_identity::EpochManifest, - mpc_operation::{K256SignPayload, MpcOperation, MpcOperationPayload}, + mpc_operation::{ + DerivePayload, DerivePurpose, K256SignPayload, MpcOperation, MpcOperationPayload, + }, mpc_session::{MpcProtocol, SessionRouter}, + threshold_prf::{self, PrfPartial}, }; pub(crate) struct DerivedAppKeys { @@ -58,6 +61,7 @@ pub(crate) trait KeyBackend: Send + Sync { ) -> 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 run_mpc_operation( @@ -130,6 +134,10 @@ impl KeyBackend for LocalKeyBackend { 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 } @@ -173,6 +181,7 @@ pub(crate) struct MpcKeyBackend { root_ca_cert: String, p256_share: P256KeyShare, k256_share: K256KeyShare, + derivation_share: P256KeyShare, transport: Arc, manifest: EpochManifest, cluster_id: String, @@ -196,6 +205,7 @@ impl MpcKeyBackend { node_id: &str, p256_share_file: &std::path::Path, k256_share_file: &std::path::Path, + derivation_share_file: &std::path::Path, manifest: &EpochManifest, router: std::sync::Arc, rpc_cert: String, @@ -227,6 +237,13 @@ impl MpcKeyBackend { node_id, CggmpCurve::K256, )?, + derivation_share: load_share( + derivation_share_file, + cluster_id, + epoch, + node_id, + CggmpCurve::P256, + )?, transport: Arc::new(transport), manifest: manifest.clone(), cluster_id: cluster_id.into(), @@ -296,6 +313,69 @@ impl MpcKeyBackend { 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()?, + 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) + } + 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")?; @@ -341,7 +421,7 @@ impl MpcKeyBackend { } }; if owner { - let result = self.execute_k256(operation).await; + let result = self.execute_operation(operation).await; *record .result .lock() @@ -369,6 +449,54 @@ impl MpcKeyBackend { } } + async fn execute_operation(&self, operation: MpcOperation) -> Result> { + match &operation.payload { + MpcOperationPayload::SignK256(_) => self.execute_k256(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: [u8; 32] = self + .derivation_share + .core + .x + .as_ref() + .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") + } + } + } + async fn execute_k256(&self, operation: MpcOperation) -> Result> { let MpcOperationPayload::SignK256(payload) = &operation.payload; let digest = payload.digest(); @@ -460,12 +588,43 @@ fn unix_time() -> Result { #[async_trait] impl KeyBackend for MpcKeyBackend { - async fn derive_app_keys(&self, _app_id: &[u8], _instance_id: &[u8]) -> Result { - bail!("MPC derivation protocol is unavailable") + 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()), + }) + .await?; + let env_key = self.derive_env_key(app_id).await?; + let app_seed = self + .derive(DerivePayload { + purpose: DerivePurpose::AppK256, + app_id: app_id.to_vec(), + instance_id: None, + }) + .await?; + let k256_key = SigningKey::from_slice(&app_seed) + .context("threshold derivation produced an invalid 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]> { - bail!("MPC derivation protocol is unavailable") + 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, + }) + .await } async fn sign_k256(&self, prefix: &[u8], app_id: &[u8], message: &[u8]) -> Result> { @@ -510,6 +669,14 @@ impl KeyBackend for MpcKeyBackend { .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 } diff --git a/dstack/kms/src/main_service.rs b/dstack/kms/src/main_service.rs index b58135b85..d17907331 100644 --- a/dstack/kms/src/main_service.rs +++ b/dstack/kms/src/main_service.rs @@ -176,6 +176,7 @@ impl KmsState { &config.mpc.node_id, &config.mpc.p256_share_file, &config.mpc.k256_share_file, + &config.mpc.derivation_share_file, manifest, mpc_router.clone().context("MPC router is missing")?, fs::read_to_string(config.rpc_cert()).context("failed to read MPC RPC cert")?, @@ -211,6 +212,10 @@ impl KmsState { 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" diff --git a/dstack/kms/src/mpc_operation.rs b/dstack/kms/src/mpc_operation.rs index 56c9d2da4..df20c1047 100644 --- a/dstack/kms/src/mpc_operation.rs +++ b/dstack/kms/src/mpc_operation.rs @@ -26,6 +26,54 @@ pub(crate) struct K256SignPayload { 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>, +} + +impl DerivePayload { + fn validate(&self) -> Result<()> { + ensure!( + self.app_id.len() == 20, + "MPC derivation app_id must be 20 bytes" + ); + 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(); @@ -63,6 +111,7 @@ impl K256SignPayload { #[serde(tag = "kind", rename_all = "snake_case")] pub(crate) enum MpcOperationPayload { SignK256(K256SignPayload), + Derive(DerivePayload), } #[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] @@ -107,6 +156,25 @@ impl MpcOperation { 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 validate(&self, manifest: &EpochManifest, initiator: &str) -> Result<()> { ensure!( self.session_id.len() == 32, @@ -143,6 +211,7 @@ impl MpcOperation { ); match &self.payload { MpcOperationPayload::SignK256(payload) => payload.validate()?, + MpcOperationPayload::Derive(payload) => payload.validate()?, } ensure!( self.request_hash == self.compute_request_hash()?, @@ -195,6 +264,27 @@ mod hex_bytes { } } +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::*; From e72f93aef8a3085528dc3d7c312f07ef7875ff67 Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Thu, 30 Jul 2026 03:49:27 -0700 Subject: [PATCH 040/122] fix(kms): dispatch derivation and signing payloads explicitly --- dstack/kms/src/key_backend.rs | 18 +++++++++++------- 1 file changed, 11 insertions(+), 7 deletions(-) diff --git a/dstack/kms/src/key_backend.rs b/dstack/kms/src/key_backend.rs index 93b852750..923163f7d 100644 --- a/dstack/kms/src/key_backend.rs +++ b/dstack/kms/src/key_backend.rs @@ -17,7 +17,11 @@ use std::{ use anyhow::{bail, ensure, Context, Result}; use async_trait::async_trait; use cggmp21::key_share::AnyKeyShare as _; -use cggmp21::{generic_ec::Scalar, supported_curves::Secp256k1, DataToSign, ExecutionId}; +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::{ @@ -383,6 +387,7 @@ impl MpcKeyBackend { .context("invalid MPC group public key")?; let digest = match &operation.payload { MpcOperationPayload::SignK256(payload) => payload.digest(), + MpcOperationPayload::Derive(_) => bail!("derivation result is not a signature"), }; for recovery_id in 0u8..4 { let recovery_id = RecoveryId::try_from(recovery_id).context("invalid recovery ID")?; @@ -464,11 +469,8 @@ impl MpcKeyBackend { usize::from(self.derivation_share.core.i) == own_index, "derivation share index does not match manifest ordering" ); - let secret: [u8; 32] = self - .derivation_share - .core - .x - .as_ref() + let secret_scalar: &Scalar = self.derivation_share.core.x.as_ref(); + let secret: [u8; 32] = secret_scalar .to_be_bytes() .as_bytes() .try_into() @@ -498,7 +500,9 @@ impl MpcKeyBackend { } async fn execute_k256(&self, operation: MpcOperation) -> Result> { - let MpcOperationPayload::SignK256(payload) = &operation.payload; + let MpcOperationPayload::SignK256(payload) = &operation.payload else { + bail!("K-256 executor received a different MPC operation") + }; let digest = payload.digest(); let local_protocol_index: u16 = operation .participants From 4238e671fc0fd529507901f3863cc55d9c24f05a Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Thu, 30 Jul 2026 03:49:54 -0700 Subject: [PATCH 041/122] build(kms): lock threshold PRF curve dependency --- dstack/Cargo.lock | 1 + 1 file changed, 1 insertion(+) diff --git a/dstack/Cargo.lock b/dstack/Cargo.lock index 759a23481..56eeb5e69 100644 --- a/dstack/Cargo.lock +++ b/dstack/Cargo.lock @@ -2109,6 +2109,7 @@ dependencies = [ "http-client", "k256", "load_config", + "p256", "parity-scale-codec", "ra-rpc", "ra-tls", From 4b2eb1bfaa8031df20ebb42cde5ceac280e4b2ce Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Thu, 30 Jul 2026 03:51:41 -0700 Subject: [PATCH 042/122] feat(ra-tls): prepare certificates for threshold signers --- dstack/ra-tls/Cargo.toml | 1 + dstack/ra-tls/src/cert.rs | 73 ++++++++++++++++++++++++++++++++++++++- 2 files changed, 73 insertions(+), 1 deletion(-) 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..2c038d61b 100644 --- a/dstack/ra-tls/src/cert.rs +++ b/dstack/ra-tls/src/cert.rs @@ -8,10 +8,11 @@ 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::{ @@ -437,6 +438,76 @@ 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 { + /// 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(), 0); + }); + }); + 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(), + }) +} + impl CertExt for Certificate { fn get_extension_der(&self, oid: &[u64]) -> Result>> { let found = self From 9d816ed1bb1fba8a11e9b51c61081acac462257c Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Thu, 30 Jul 2026 03:53:31 -0700 Subject: [PATCH 043/122] feat(kms): issue app CAs with threshold P-256 signing --- dstack/kms/src/key_backend.rs | 223 +++++++++++++++++++++++++++++++- dstack/kms/src/mpc_operation.rs | 51 ++++++++ 2 files changed, 271 insertions(+), 3 deletions(-) diff --git a/dstack/kms/src/key_backend.rs b/dstack/kms/src/key_backend.rs index 923163f7d..788cd65e2 100644 --- a/dstack/kms/src/key_backend.rs +++ b/dstack/kms/src/key_backend.rs @@ -25,7 +25,7 @@ use cggmp21::{ use futures::future::join_all; use k256::ecdsa::{RecoveryId, Signature as K256Signature, SigningKey, VerifyingKey}; use ra_tls::{ - cert::{CaCert, CertRequest}, + cert::{prepare_external_certificate, CaCert, CertRequest}, kdf, }; use rand::{rngs::OsRng, RngCore}; @@ -39,6 +39,7 @@ use crate::{ mpc_identity::EpochManifest, mpc_operation::{ DerivePayload, DerivePurpose, K256SignPayload, MpcOperation, MpcOperationPayload, + P256CertificatePayload, }, mpc_session::{MpcProtocol, SessionRouter}, threshold_prf::{self, PrfPartial}, @@ -380,6 +381,114 @@ impl MpcKeyBackend { 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()?, + 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) + } + + 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 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).context("invalid app ID OID")?; + 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).context("invalid usage OID")?; + 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")?; @@ -457,6 +566,10 @@ impl MpcKeyBackend { async fn execute_operation(&self, operation: MpcOperation) -> Result> { match &operation.payload { MpcOperationPayload::SignK256(_) => 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 @@ -581,6 +694,84 @@ impl MpcKeyBackend { .await .context("MPC signing worker panicked")? } + + 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, + &BlockingHttpTransport::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 { @@ -685,8 +876,34 @@ impl KeyBackend for MpcKeyBackend { &self.root_ca_cert } - async fn derive_app_ca(&self, _app_id: &[u8]) -> Result { - bail!("legacy app CA derivation is unavailable in MPC mode") + async fn derive_app_ca(&self, app_id: &[u8]) -> Result { + let seed = self + .derive(DerivePayload { + purpose: DerivePurpose::AppCa, + app_id: app_id.to_vec(), + instance_id: None, + }) + .await?; + let app_key = kdf::derive_p256_key_pair_from_bytes(&seed, &[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 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 run_mpc_operation(&self, operation: MpcOperation, initiator: &str) -> Result> { diff --git a/dstack/kms/src/mpc_operation.rs b/dstack/kms/src/mpc_operation.rs index df20c1047..c952357f0 100644 --- a/dstack/kms/src/mpc_operation.rs +++ b/dstack/kms/src/mpc_operation.rs @@ -44,6 +44,36 @@ pub(crate) struct DerivePayload { pub instance_id: Option>, } +#[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, +} + +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!( @@ -111,6 +141,7 @@ impl K256SignPayload { #[serde(tag = "kind", rename_all = "snake_case")] pub(crate) enum MpcOperationPayload { SignK256(K256SignPayload), + SignP256Certificate(P256CertificatePayload), Derive(DerivePayload), } @@ -175,6 +206,25 @@ impl MpcOperation { 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 validate(&self, manifest: &EpochManifest, initiator: &str) -> Result<()> { ensure!( self.session_id.len() == 32, @@ -211,6 +261,7 @@ impl MpcOperation { ); match &self.payload { MpcOperationPayload::SignK256(payload) => payload.validate()?, + MpcOperationPayload::SignP256Certificate(payload) => payload.validate()?, MpcOperationPayload::Derive(payload) => payload.validate()?, } ensure!( From b6e3a380e091c060099ccd7c3b60d2efadc3982a Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Thu, 30 Jul 2026 03:53:51 -0700 Subject: [PATCH 044/122] fix(kms): validate certificate OIDs without error coercion --- dstack/kms/src/key_backend.rs | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/dstack/kms/src/key_backend.rs b/dstack/kms/src/key_backend.rs index 788cd65e2..0cd94454e 100644 --- a/dstack/kms/src/key_backend.rs +++ b/dstack/kms/src/key_backend.rs @@ -449,7 +449,8 @@ impl MpcKeyBackend { 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).context("invalid app ID OID")?; + 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")? @@ -457,7 +458,8 @@ impl MpcKeyBackend { 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).context("invalid usage OID")?; + 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")? @@ -496,6 +498,9 @@ impl MpcKeyBackend { .context("invalid MPC group public key")?; let digest = match &operation.payload { MpcOperationPayload::SignK256(payload) => payload.digest(), + MpcOperationPayload::SignP256Certificate(_) => { + bail!("P-256 certificate result does not have a recovery ID") + } MpcOperationPayload::Derive(_) => bail!("derivation result is not a signature"), }; for recovery_id in 0u8..4 { From 2424ac4117b358c6b19217bef07afe0a30c50aa5 Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Thu, 30 Jul 2026 03:54:18 -0700 Subject: [PATCH 045/122] test(ra-tls): verify externally signed certificate assembly --- dstack/ra-tls/src/cert.rs | 40 +++++++++++++++++++++++++++++++++++++++ 1 file changed, 40 insertions(+) diff --git a/dstack/ra-tls/src/cert.rs b/dstack/ra-tls/src/cert.rs index 2c038d61b..69fb5e55e 100644 --- a/dstack/ra-tls/src/cert.rs +++ b/dstack/ra-tls/src/cert.rs @@ -646,6 +646,46 @@ 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 test_csr_signing_and_verification() { let key_pair = KeyPair::generate_for(&PKCS_ECDSA_P256_SHA256).unwrap(); From ac5dc7704b9ebbb68ab555c524c1331e12e14f62 Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Thu, 30 Jul 2026 03:55:04 -0700 Subject: [PATCH 046/122] fix(ra-tls): encode external signature bit length --- dstack/ra-tls/src/cert.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/dstack/ra-tls/src/cert.rs b/dstack/ra-tls/src/cert.rs index 69fb5e55e..3522475b0 100644 --- a/dstack/ra-tls/src/cert.rs +++ b/dstack/ra-tls/src/cert.rs @@ -467,7 +467,7 @@ impl ExternalCertificate { writer.next().write_der(ECDSA_SHA256_ALGORITHM); writer .next() - .write_bitvec_bytes(signature_der.as_bytes(), 0); + .write_bitvec_bytes(signature_der.as_bytes(), signature_der.len() * 8); }); }); let encoded = BASE64.encode(der); From 85449c8fb61b82d2c2bbff72df38de36584e8d45 Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Thu, 30 Jul 2026 03:55:23 -0700 Subject: [PATCH 047/122] build(ra-tls): lock external certificate encoding dependency --- dstack/Cargo.lock | 1 + 1 file changed, 1 insertion(+) diff --git a/dstack/Cargo.lock b/dstack/Cargo.lock index 56eeb5e69..0fd9e9c16 100644 --- a/dstack/Cargo.lock +++ b/dstack/Cargo.lock @@ -5937,6 +5937,7 @@ name = "ra-tls" version = "0.6.0" dependencies = [ "anyhow", + "base64 0.22.1", "bon", "cc-eventlog", "dcap-qvl", From 770b871677b8f5b8c1c27b817f25e5ee1f6b3f3d Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Thu, 30 Jul 2026 03:58:09 -0700 Subject: [PATCH 048/122] feat(kms): enforce threshold-signed monotonic epoch manifests --- dstack/kms/kms.toml | 2 + dstack/kms/src/config.rs | 5 + dstack/kms/src/main.rs | 1 + dstack/kms/src/main_service.rs | 73 ++++++++---- dstack/kms/src/mpc_identity.rs | 25 +++++ dstack/kms/src/mpc_lifecycle.rs | 193 ++++++++++++++++++++++++++++++++ 6 files changed, 279 insertions(+), 20 deletions(-) create mode 100644 dstack/kms/src/mpc_lifecycle.rs diff --git a/dstack/kms/kms.toml b/dstack/kms/kms.toml index fea00f02c..f193dda1b 100644 --- a/dstack/kms/kms.toml +++ b/dstack/kms/kms.toml @@ -109,6 +109,8 @@ 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 = "" diff --git a/dstack/kms/src/config.rs b/dstack/kms/src/config.rs index d272d5e13..48ed01fc7 100644 --- a/dstack/kms/src/config.rs +++ b/dstack/kms/src/config.rs @@ -103,6 +103,11 @@ pub(crate) struct MpcConfig { 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")] diff --git a/dstack/kms/src/main.rs b/dstack/kms/src/main.rs index 7b9120b39..a6de759c6 100644 --- a/dstack/kms/src/main.rs +++ b/dstack/kms/src/main.rs @@ -26,6 +26,7 @@ mod key_backend; mod main_service; mod mpc_driver; mod mpc_identity; +mod mpc_lifecycle; mod mpc_operation; mod mpc_service; mod mpc_session; diff --git a/dstack/kms/src/main_service.rs b/dstack/kms/src/main_service.rs index d17907331..fdc5efbe0 100644 --- a/dstack/kms/src/main_service.rs +++ b/dstack/kms/src/main_service.rs @@ -32,8 +32,8 @@ use upgrade_authority::{build_boot_info, ensure_app_id_len, local_kms_boot_info, use crate::{ config::KmsConfig, key_backend::{KeyBackend, LocalKeyBackend, MpcKeyBackend}, - mpc_identity::EpochManifest, - mpc_identity::{decode_hex, ClusterIdentity}, + mpc_identity::{decode_hex, ClusterIdentity, EpochManifest, SignedEpochManifest}, + mpc_lifecycle, mpc_session::SessionRouter, }; @@ -142,17 +142,50 @@ impl KmsState { AttestationVerifier::load(&config.attestation) .context("failed to load attestation verifier")?, ); + let configured_identity = if config.mpc.enabled { + Some(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 { + None + }; + 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" ); - Some( - serde_json::from_slice( - &fs::read(&config.mpc.manifest_file).context("failed to read MPC manifest")?, - ) - .context("failed to parse MPC manifest")?, - ) + 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 }; @@ -192,18 +225,7 @@ impl KmsState { )?) }; let mpc_identity = if config.mpc.enabled { - let identity = 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, - )?, - )?; - // Identity v1 still uses the local backend. Fail closed if config - // claims group keys different from the keys actually serving RPCs. + 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" @@ -229,6 +251,17 @@ impl KmsState { 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 diff --git a/dstack/kms/src/mpc_identity.rs b/dstack/kms/src/mpc_identity.rs index ad7960b1c..26f1fd6c2 100644 --- a/dstack/kms/src/mpc_identity.rs +++ b/dstack/kms/src/mpc_identity.rs @@ -8,6 +8,7 @@ //! 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}; @@ -95,6 +96,30 @@ pub(crate) struct EpochManifest { 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"); diff --git a/dstack/kms/src/mpc_lifecycle.rs b/dstack/kms/src/mpc_lifecycle.rs new file mode 100644 index 000000000..24541a332 --- /dev/null +++ b/dstack/kms/src/mpc_lifecycle.rs @@ -0,0 +1,193 @@ +// 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, os::unix::fs::PermissionsExt, path::Path}; + +use anyhow::{ensure, Context, Result}; +use fs_err as fs; +use serde::{Deserialize, Serialize}; + +use crate::mpc_identity::{ClusterIdentity, SignedEpochManifest}; + +#[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(()) +} + +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}; + + 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 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()); + } +} From 463c769e035bfa8740f9b3294167fea395202785 Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Thu, 30 Jul 2026 03:58:40 -0700 Subject: [PATCH 049/122] fix(kms): bind root certificate to threshold P-256 key --- dstack/kms/src/key_backend.rs | 20 ++++++++++++++++++-- 1 file changed, 18 insertions(+), 2 deletions(-) diff --git a/dstack/kms/src/key_backend.rs b/dstack/kms/src/key_backend.rs index 0cd94454e..e924ec7e8 100644 --- a/dstack/kms/src/key_backend.rs +++ b/dstack/kms/src/key_backend.rs @@ -217,6 +217,12 @@ impl MpcKeyBackend { 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, @@ -226,7 +232,7 @@ impl MpcKeyBackend { root_ca_cert.clone(), attestation_verifier, )?; - Ok(Self { + let backend = Self { root_ca_cert, p256_share: load_share( p256_share_file, @@ -254,7 +260,17 @@ impl MpcKeyBackend { cluster_id: cluster_id.into(), node_id: node_id.into(), 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" + ); + Ok(backend) } fn participants(&self) -> Result> { From 6231ba753983825d3aea39fa63e0ead4c64d10cf Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Thu, 30 Jul 2026 03:59:50 -0700 Subject: [PATCH 050/122] fix(kms): bind all local threshold shares to epoch membership --- dstack/kms/src/cggmp_engine.rs | 20 +++++++++++++++ dstack/kms/src/key_backend.rs | 45 +++++++++++++++++++++++++++++++++- 2 files changed, 64 insertions(+), 1 deletion(-) diff --git a/dstack/kms/src/cggmp_engine.rs b/dstack/kms/src/cggmp_engine.rs index 6ea0da614..21f258d68 100644 --- a/dstack/kms/src/cggmp_engine.rs +++ b/dstack/kms/src/cggmp_engine.rs @@ -22,6 +22,7 @@ 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")] @@ -74,6 +75,25 @@ pub(crate) fn execution_id( 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 store_share( path: &Path, cluster_id: &str, diff --git a/dstack/kms/src/key_backend.rs b/dstack/kms/src/key_backend.rs index e924ec7e8..8410dbd37 100644 --- a/dstack/kms/src/key_backend.rs +++ b/dstack/kms/src/key_backend.rs @@ -31,7 +31,9 @@ use ra_tls::{ use rand::{rngs::OsRng, RngCore}; use crate::{ - cggmp_engine::{execution_id, load_share, CggmpCurve, K256KeyShare, P256KeyShare}, + cggmp_engine::{ + execution_id, load_share, share_commitment, CggmpCurve, K256KeyShare, P256KeyShare, + }, crypto::{derive_k256_key, sign_message, sign_message_with_timestamp}, mpc_driver::{ drive_state_machine_blocking, BlockingHttpTransport, DriverContext, MpcHttpTransport, @@ -270,6 +272,47 @@ impl MpcKeyBackend { == 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")?; + 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) } From 717ad898cd74d2e497220ceb00bbdce3db8c535e Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Thu, 30 Jul 2026 04:00:58 -0700 Subject: [PATCH 051/122] feat(kms): select live members for threshold quorums --- dstack/kms/rpc/proto/kms_rpc.proto | 1 + dstack/kms/src/key_backend.rs | 13 ++++++++----- dstack/kms/src/mpc_driver.rs | 10 ++++++++++ dstack/kms/src/mpc_service.rs | 5 +++++ 4 files changed, 24 insertions(+), 5 deletions(-) diff --git a/dstack/kms/rpc/proto/kms_rpc.proto b/dstack/kms/rpc/proto/kms_rpc.proto index 9db247c96..3b36e8fd4 100644 --- a/dstack/kms/rpc/proto/kms_rpc.proto +++ b/dstack/kms/rpc/proto/kms_rpc.proto @@ -135,6 +135,7 @@ message MpcStartResponse { // 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); rpc Poll(MpcPollRequest) returns (MpcPollResponse); diff --git a/dstack/kms/src/key_backend.rs b/dstack/kms/src/key_backend.rs index 8410dbd37..6c00dd550 100644 --- a/dstack/kms/src/key_backend.rs +++ b/dstack/kms/src/key_backend.rs @@ -316,8 +316,9 @@ impl MpcKeyBackend { Ok(backend) } - fn participants(&self) -> Result> { + 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 @@ -329,7 +330,9 @@ impl MpcKeyBackend { self.manifest .members .iter() - .filter(|member| member.node_id != self.node_id) + .filter(|member| { + member.node_id != self.node_id && reachable.contains(&member.node_id) + }) .take(threshold.saturating_sub(1)) .map(|member| member.node_id.clone()), ); @@ -350,7 +353,7 @@ impl MpcKeyBackend { let operation = MpcOperation::new_k256( session_id, self.manifest.epoch, - self.participants()?, + self.participants().await?, unix_time()? .checked_add(60) .context("MPC expiry overflow")?, @@ -383,7 +386,7 @@ impl MpcKeyBackend { let operation = MpcOperation::new_derivation( session_id, self.manifest.epoch, - self.participants()?, + self.participants().await?, unix_time()? .checked_add(60) .context("MPC expiry overflow")?, @@ -447,7 +450,7 @@ impl MpcKeyBackend { let operation = MpcOperation::new_p256_certificate( session_id, self.manifest.epoch, - self.participants()?, + self.participants().await?, unix_time()? .checked_add(60) .context("MPC expiry overflow")?, diff --git a/dstack/kms/src/mpc_driver.rs b/dstack/kms/src/mpc_driver.rs index 45ffa19c4..e753dd2cf 100644 --- a/dstack/kms/src/mpc_driver.rs +++ b/dstack/kms/src/mpc_driver.rs @@ -152,6 +152,16 @@ impl MpcHttpTransport { .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 { + client.ping().await.ok().map(|_| node_id.clone()) + })) + .await + .into_iter() + .flatten() + .collect() + } } #[async_trait] diff --git a/dstack/kms/src/mpc_service.rs b/dstack/kms/src/mpc_service.rs index 6bcbcc685..d2605f337 100644 --- a/dstack/kms/src/mpc_service.rs +++ b/dstack/kms/src/mpc_service.rs @@ -42,6 +42,11 @@ impl MpcHandler { } 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")?; From 3b99e505868e1dc276c7f0b9bf139897f3fd959b Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Thu, 30 Jul 2026 04:01:20 -0700 Subject: [PATCH 052/122] fix(kms): bound MPC liveness probes --- dstack/kms/src/mpc_driver.rs | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/dstack/kms/src/mpc_driver.rs b/dstack/kms/src/mpc_driver.rs index e753dd2cf..a853d5602 100644 --- a/dstack/kms/src/mpc_driver.rs +++ b/dstack/kms/src/mpc_driver.rs @@ -155,7 +155,11 @@ impl MpcHttpTransport { pub(crate) async fn reachable_nodes(&self) -> Vec { futures::future::join_all(self.clients.iter().map(|(node_id, client)| async move { - client.ping().await.ok().map(|_| node_id.clone()) + tokio::time::timeout(Duration::from_secs(3), client.ping()) + .await + .ok() + .and_then(Result::ok) + .map(|_| node_id.clone()) })) .await .into_iter() From 3aa8a05dded7dd47243654fc429318cd33157db5 Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Thu, 30 Jul 2026 04:01:33 -0700 Subject: [PATCH 053/122] fix(kms): keep MPC inboxes local to their recipient --- dstack/kms/rpc/proto/kms_rpc.proto | 1 - dstack/kms/src/mpc_service.rs | 12 +----------- 2 files changed, 1 insertion(+), 12 deletions(-) diff --git a/dstack/kms/rpc/proto/kms_rpc.proto b/dstack/kms/rpc/proto/kms_rpc.proto index 3b36e8fd4..73888f4b7 100644 --- a/dstack/kms/rpc/proto/kms_rpc.proto +++ b/dstack/kms/rpc/proto/kms_rpc.proto @@ -138,7 +138,6 @@ service MpcTransport { rpc Ping(google.protobuf.Empty) returns (google.protobuf.Empty); rpc Start(MpcStartRequest) returns (MpcStartResponse); rpc Push(MpcPushRequest) returns (google.protobuf.Empty); - rpc Poll(MpcPollRequest) returns (MpcPollResponse); } // The KMS admin RPC service. Served on a separate admin listener diff --git a/dstack/kms/src/mpc_service.rs b/dstack/kms/src/mpc_service.rs index d2605f337..bd98dd45b 100644 --- a/dstack/kms/src/mpc_service.rs +++ b/dstack/kms/src/mpc_service.rs @@ -5,7 +5,7 @@ use anyhow::{Context, Result}; use dstack_kms_rpc::{ mpc_transport_server::{MpcTransportRpc, MpcTransportServer}, - MpcPollRequest, MpcPollResponse, MpcPushRequest, MpcStartRequest, MpcStartResponse, + MpcPushRequest, MpcStartRequest, MpcStartResponse, }; use ra_rpc::{CallContext, RpcCall}; @@ -66,16 +66,6 @@ impl MpcTransportRpc for MpcHandler { let (node_id, router) = self.peer_and_router()?; router.push(node_id, envelope) } - - async fn poll(self, request: MpcPollRequest) -> Result { - let (node_id, router) = self.peer_and_router()?; - let envelope_json = router - .drain(node_id, &request.session_id)? - .into_iter() - .map(|envelope| serde_json::to_vec(&envelope).context("failed to encode MPC envelope")) - .collect::>>()?; - Ok(MpcPollResponse { envelope_json }) - } } pub(crate) fn rpc_methods() -> &'static [&'static str] { From 420ea2f748ad90f4b3a477c857e908cf324bc61d Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Thu, 30 Jul 2026 04:04:14 -0700 Subject: [PATCH 054/122] feat(kms): add verifiable threshold-preserving resharing --- dstack/kms/src/main.rs | 1 + dstack/kms/src/mpc_reshare.rs | 321 ++++++++++++++++++++++++++++++++++ 2 files changed, 322 insertions(+) create mode 100644 dstack/kms/src/mpc_reshare.rs diff --git a/dstack/kms/src/main.rs b/dstack/kms/src/main.rs index a6de759c6..418d85506 100644 --- a/dstack/kms/src/main.rs +++ b/dstack/kms/src/main.rs @@ -28,6 +28,7 @@ mod mpc_driver; mod mpc_identity; mod mpc_lifecycle; mod mpc_operation; +mod mpc_reshare; mod mpc_service; mod mpc_session; mod onboard_service; diff --git a/dstack/kms/src/mpc_reshare.rs b/dstack/kms/src/mpc_reshare.rs new file mode 100644 index 000000000..b0bdb7210 --- /dev/null +++ b/dstack/kms/src/mpc_reshare.rs @@ -0,0 +1,321 @@ +// 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, Point, Scalar}, + 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, +} + +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)); + coefficients.push(lambda * old_share.x.as_ref()); + 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> { + 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_reference + .core + .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_reference + .core + .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_reference.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, key_share::trusted_dealer, supported_curves::Secp256k1}; + 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()); + } + + #[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() + .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()); + } +} From 986438f46c5ede65785805c8c0badd1f0c0a3c8a Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Thu, 30 Jul 2026 04:05:04 -0700 Subject: [PATCH 055/122] fix(kms): construct resharing fixtures from core shares --- dstack/kms/Cargo.toml | 1 + dstack/kms/src/mpc_reshare.rs | 10 +++++----- 2 files changed, 6 insertions(+), 5 deletions(-) diff --git a/dstack/kms/Cargo.toml b/dstack/kms/Cargo.toml index b6ea134cd..ce7eb987e 100644 --- a/dstack/kms/Cargo.toml +++ b/dstack/kms/Cargo.toml @@ -59,6 +59,7 @@ round-based.workspace = true [dev-dependencies] cc-eventlog.workspace = true round-based = { workspace = true, features = ["sim"] } +key-share = { version = "0.6.1", features = ["spof"] } [features] default = [] diff --git a/dstack/kms/src/mpc_reshare.rs b/dstack/kms/src/mpc_reshare.rs index b0bdb7210..c0ee5653c 100644 --- a/dstack/kms/src/mpc_reshare.rs +++ b/dstack/kms/src/mpc_reshare.rs @@ -79,7 +79,8 @@ pub(crate) fn create_contribution( .collect::>>()?; let lambda = lagrange_at_zero(dealer_position, &dealer_points)?; let mut coefficients = Vec::with_capacity(usize::from(new_threshold)); - coefficients.push(lambda * old_share.x.as_ref()); + 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() @@ -135,7 +136,6 @@ pub(crate) fn verify_and_combine( "missing private resharing contribution" ); let setup = old_reference - .core .vss_setup .as_ref() .context("resharing requires an old Shamir setup")?; @@ -168,7 +168,6 @@ pub(crate) fn verify_and_combine( .collect::>>()?; let lambda = lagrange_at_zero(position, &dealer_points)?; let expected_constant = old_reference - .core .key_info .public_shares .get(usize::from(*dealer_index)) @@ -248,7 +247,8 @@ fn evaluate_commitments(coefficients: &[Point], x: Scalar) -> Po #[cfg(test)] mod tests { use super::*; - use cggmp21::{generic_ec::NonZero, key_share::trusted_dealer, supported_curves::Secp256k1}; + use cggmp21::{generic_ec::NonZero, supported_curves::Secp256k1}; + use key_share::trusted_dealer; use rand::rngs::OsRng; #[test] @@ -303,7 +303,7 @@ mod tests { 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() + first_public.commitments[0] = (Point::::generator() * Scalar::one()) .to_bytes(true) .as_bytes() .to_vec(); From 9d18edc01f0feaac03aa19a7a318967a683f61cb Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Thu, 30 Jul 2026 04:05:26 -0700 Subject: [PATCH 056/122] build(kms): lock resharing test dependency --- dstack/Cargo.lock | 2 ++ 1 file changed, 2 insertions(+) diff --git a/dstack/Cargo.lock b/dstack/Cargo.lock index 0fd9e9c16..6671ed8e6 100644 --- a/dstack/Cargo.lock +++ b/dstack/Cargo.lock @@ -2108,6 +2108,7 @@ dependencies = [ "hex_fmt", "http-client", "k256", + "key-share", "load_config", "p256", "parity-scale-codec", @@ -4217,6 +4218,7 @@ dependencies = [ "generic-ec", "generic-ec-zkp", "hex", + "rand_core 0.6.4", "serde", "serde_with 2.3.3", "thiserror 1.0.69", From ab292d2569ba7c5a9791499786ac6fe2d5839af7 Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Thu, 30 Jul 2026 04:05:32 -0700 Subject: [PATCH 057/122] fix(kms): cap authenticated MPC operation concurrency --- dstack/kms/src/key_backend.rs | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/dstack/kms/src/key_backend.rs b/dstack/kms/src/key_backend.rs index 6c00dd550..9a2a29455 100644 --- a/dstack/kms/src/key_backend.rs +++ b/dstack/kms/src/key_backend.rs @@ -204,6 +204,8 @@ struct OperationRecord { finished: tokio::sync::Notify, } +const MAX_ACTIVE_OPERATIONS: usize = 128; + impl MpcKeyBackend { pub(crate) fn load( root_ca_cert: String, @@ -590,6 +592,10 @@ impl MpcKeyBackend { ); (record.clone(), false) } else { + ensure!( + operations.len() < MAX_ACTIVE_OPERATIONS, + "too many active MPC operations" + ); let record = Arc::new(OperationRecord { request_hash: operation.request_hash.clone(), initiator: initiator.into(), From 4b9560f390a1be6208e6ce96ce0df12300bad40f Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Thu, 30 Jul 2026 04:06:06 -0700 Subject: [PATCH 058/122] feat(kms): materialize reshared CGGMP core shares --- dstack/kms/src/mpc_reshare.rs | 50 ++++++++++++++++++++++++++++++++++- 1 file changed, 49 insertions(+), 1 deletion(-) diff --git a/dstack/kms/src/mpc_reshare.rs b/dstack/kms/src/mpc_reshare.rs index c0ee5653c..740aa1a2d 100644 --- a/dstack/kms/src/mpc_reshare.rs +++ b/dstack/kms/src/mpc_reshare.rs @@ -10,7 +10,8 @@ use anyhow::{ensure, Context, Result}; use cggmp21::{ - generic_ec::{Curve, Point, Scalar}, + generic_ec::{Curve, NonZero, Point, Scalar, SecretScalar}, + key_share::{DirtyIncompleteKeyShare, Valid, VssSetup}, IncompleteKeyShare, }; use rand::{CryptoRng, RngCore}; @@ -35,6 +36,46 @@ pub(crate) struct ResharedMaterial { 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" + ); + let mut key_info = template.key_info.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], @@ -290,6 +331,13 @@ mod tests { *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] From ce1bec5b0601e6319905cb207f682fbe6b07fada Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Thu, 30 Jul 2026 04:07:40 -0700 Subject: [PATCH 059/122] test(kms): exercise real CGGMP threshold signing transport --- dstack/kms/Cargo.toml | 2 +- dstack/kms/src/mpc_driver.rs | 80 ++++++++++++++++++++++++++++++++++++ 2 files changed, 81 insertions(+), 1 deletion(-) diff --git a/dstack/kms/Cargo.toml b/dstack/kms/Cargo.toml index ce7eb987e..b1ac2938c 100644 --- a/dstack/kms/Cargo.toml +++ b/dstack/kms/Cargo.toml @@ -53,7 +53,7 @@ dstack-verifier = { workspace = true, default-features = false } dstack-mr.workspace = true dstack-attest.workspace = true dstack-api-auth.workspace = true -cggmp21.workspace = true +cggmp21 = { workspace = true, features = ["spof"] } round-based.workspace = true [dev-dependencies] diff --git a/dstack/kms/src/mpc_driver.rs b/dstack/kms/src/mpc_driver.rs index a853d5602..56892695f 100644 --- a/dstack/kms/src/mpc_driver.rs +++ b/dstack/kms/src/mpc_driver.rs @@ -520,4 +520,84 @@ mod tests { 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}; + + let shares = cggmp21::trusted_dealer::builder::(3) + .set_threshold(Some(2)) + .generate_shares(&mut OsRng) + .unwrap(); + 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!["kms-1".to_string(), "kms-3".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), + }; + async { + let mut rng = OsRng; + let data = + DataToSign::from_scalar(Scalar::::from_be_bytes_mod_order(digest)); + let keygen_indexes = [0, 2]; + let state = cggmp21::signing( + ExecutionId::new(&eid), + protocol_index, + &keygen_indexes, + &shares[keygen_index], + ) + .sign_sync(&mut rng, data); + drive_state_machine(state, &transport, context) + .await + .unwrap() + .unwrap() + } + }; + let (first, second) = tokio::join!(sign(0, 0), sign(1, 2)); + assert_eq!(first, second); + 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(); + } } From d29a96b978a3989116279e5adfb5d48c6ef807c0 Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Thu, 30 Jul 2026 04:08:01 -0700 Subject: [PATCH 060/122] fix(kms): import CGGMP share metadata in signing test --- dstack/kms/src/mpc_driver.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/dstack/kms/src/mpc_driver.rs b/dstack/kms/src/mpc_driver.rs index 56892695f..00edb6cab 100644 --- a/dstack/kms/src/mpc_driver.rs +++ b/dstack/kms/src/mpc_driver.rs @@ -395,7 +395,7 @@ fn ensure_not_expired(expires_at: u64) -> Result<()> { mod tests { use std::{collections::BTreeMap, sync::Arc}; - use cggmp21::{supported_curves::Secp256k1, ExecutionId}; + use cggmp21::{key_share::AnyKeyShare as _, supported_curves::Secp256k1, ExecutionId}; use rand::rngs::OsRng; use super::*; From 1077dab8b60a2ced464627192d3bdb250c363c2f Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Thu, 30 Jul 2026 04:08:31 -0700 Subject: [PATCH 061/122] fix(kms): own signing fixtures across async test tasks --- dstack/kms/src/mpc_driver.rs | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/dstack/kms/src/mpc_driver.rs b/dstack/kms/src/mpc_driver.rs index 00edb6cab..79753e8a6 100644 --- a/dstack/kms/src/mpc_driver.rs +++ b/dstack/kms/src/mpc_driver.rs @@ -572,7 +572,8 @@ mod tests { expires_at, poll_interval: Duration::from_millis(1), }; - async { + let share = shares[keygen_index].clone(); + async move { let mut rng = OsRng; let data = DataToSign::from_scalar(Scalar::::from_be_bytes_mod_order(digest)); @@ -581,7 +582,7 @@ mod tests { ExecutionId::new(&eid), protocol_index, &keygen_indexes, - &shares[keygen_index], + &share, ) .sign_sync(&mut rng, data); drive_state_machine(state, &transport, context) From a94078a1272a62d9b8b348da99983b0dea17227c Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Thu, 30 Jul 2026 04:11:53 -0700 Subject: [PATCH 062/122] build(kms): confine trusted dealer support to tests --- dstack/kms/Cargo.toml | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/dstack/kms/Cargo.toml b/dstack/kms/Cargo.toml index b1ac2938c..9a680c781 100644 --- a/dstack/kms/Cargo.toml +++ b/dstack/kms/Cargo.toml @@ -53,11 +53,12 @@ dstack-verifier = { workspace = true, default-features = false } dstack-mr.workspace = true dstack-attest.workspace = true dstack-api-auth.workspace = true -cggmp21 = { workspace = true, features = ["spof"] } +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"] } From 3287d83798b0531621d6a93facef800a7fd17ebe Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Thu, 30 Jul 2026 04:13:10 -0700 Subject: [PATCH 063/122] feat(kms): threshold-sign epoch transitions --- dstack/kms/src/key_backend.rs | 79 ++++++++++++++++++++++++++++++--- dstack/kms/src/mpc_operation.rs | 50 +++++++++++++++++++++ 2 files changed, 122 insertions(+), 7 deletions(-) diff --git a/dstack/kms/src/key_backend.rs b/dstack/kms/src/key_backend.rs index 9a2a29455..27ce31795 100644 --- a/dstack/kms/src/key_backend.rs +++ b/dstack/kms/src/key_backend.rs @@ -38,10 +38,10 @@ use crate::{ mpc_driver::{ drive_state_machine_blocking, BlockingHttpTransport, DriverContext, MpcHttpTransport, }, - mpc_identity::EpochManifest, + mpc_identity::{EpochManifest, SignedEpochManifest}, mpc_operation::{ - DerivePayload, DerivePurpose, K256SignPayload, MpcOperation, MpcOperationPayload, - P256CertificatePayload, + DerivePayload, DerivePurpose, K256SignPayload, ManifestSignaturePayload, MpcOperation, + MpcOperationPayload, P256CertificatePayload, }, mpc_session::{MpcProtocol, SessionRouter}, threshold_prf::{self, PrfPartial}, @@ -71,6 +71,7 @@ pub(crate) trait KeyBackend: Send + Sync { 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 run_mpc_operation( &self, operation: crate::mpc_operation::MpcOperation, @@ -166,6 +167,19 @@ impl KeyBackend for LocalKeyBackend { 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 export_root_keys(&self) -> Result<(String, Vec)> { Ok(( self.root_ca.key.serialize_pem(), @@ -482,6 +496,47 @@ impl MpcKeyBackend { 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, + }) + } + fn validate_p256_certificate(&self, payload: &P256CertificatePayload) -> Result<()> { use ra_tls::oids::{PHALA_RATLS_APP_ID, PHALA_RATLS_CERT_USAGE}; use x509_parser::{ @@ -562,6 +617,9 @@ impl MpcKeyBackend { .context("invalid MPC group public key")?; let digest = match &operation.payload { MpcOperationPayload::SignK256(payload) => payload.digest(), + MpcOperationPayload::SignManifest(_) => { + bail!("manifest signatures do not have recovery IDs") + } MpcOperationPayload::SignP256Certificate(_) => { bail!("P-256 certificate result does not have a recovery ID") } @@ -638,7 +696,9 @@ impl MpcKeyBackend { async fn execute_operation(&self, operation: MpcOperation) -> Result> { match &operation.payload { - MpcOperationPayload::SignK256(_) => self.execute_k256(operation).await, + MpcOperationPayload::SignK256(_) | MpcOperationPayload::SignManifest(_) => { + self.execute_k256(operation).await + } MpcOperationPayload::SignP256Certificate(payload) => { self.validate_p256_certificate(payload)?; self.execute_p256(operation).await @@ -686,10 +746,11 @@ impl MpcKeyBackend { } async fn execute_k256(&self, operation: MpcOperation) -> Result> { - let MpcOperationPayload::SignK256(payload) = &operation.payload else { - bail!("K-256 executor received a different MPC operation") + let digest = match &operation.payload { + MpcOperationPayload::SignK256(payload) => payload.digest(), + MpcOperationPayload::SignManifest(payload) => payload.digest()?, + _ => bail!("K-256 executor received a different MPC operation"), }; - let digest = payload.digest(); let local_protocol_index: u16 = operation .participants .iter() @@ -979,6 +1040,10 @@ impl KeyBackend for MpcKeyBackend { .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 run_mpc_operation(&self, operation: MpcOperation, initiator: &str) -> Result> { operation.validate(&self.manifest, initiator)?; self.execute_once(operation, initiator).await diff --git a/dstack/kms/src/mpc_operation.rs b/dstack/kms/src/mpc_operation.rs index c952357f0..24e8ab568 100644 --- a/dstack/kms/src/mpc_operation.rs +++ b/dstack/kms/src/mpc_operation.rs @@ -52,6 +52,21 @@ pub(crate) struct P256CertificatePayload { 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 _}; @@ -141,6 +156,7 @@ impl K256SignPayload { #[serde(tag = "kind", rename_all = "snake_case")] pub(crate) enum MpcOperationPayload { SignK256(K256SignPayload), + SignManifest(ManifestSignaturePayload), SignP256Certificate(P256CertificatePayload), Derive(DerivePayload), } @@ -225,6 +241,25 @@ impl MpcOperation { 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 validate(&self, manifest: &EpochManifest, initiator: &str) -> Result<()> { ensure!( self.session_id.len() == 32, @@ -261,6 +296,21 @@ impl MpcOperation { ); 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" + ); + 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()?, } From c84c6696687c5f310cbb41084ac9ad1a246318b0 Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Thu, 30 Jul 2026 04:13:37 -0700 Subject: [PATCH 064/122] feat(kms): expose threshold epoch authorization to admin --- dstack/kms/rpc/proto/kms_rpc.proto | 10 ++++++++++ dstack/kms/src/admin_service.rs | 21 +++++++++++++++++++-- 2 files changed, 29 insertions(+), 2 deletions(-) diff --git a/dstack/kms/rpc/proto/kms_rpc.proto b/dstack/kms/rpc/proto/kms_rpc.proto index 73888f4b7..2ff216e10 100644 --- a/dstack/kms/rpc/proto/kms_rpc.proto +++ b/dstack/kms/rpc/proto/kms_rpc.proto @@ -147,6 +147,8 @@ service MpcTransport { 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); } message ClearImageCacheRequest { @@ -154,6 +156,14 @@ message ClearImageCacheRequest { string config_hash = 2; } +message SignEpochManifestRequest { + bytes manifest_json = 1; +} + +message SignEpochManifestResponse { + bytes signed_manifest_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..834d577c8 100644 --- a/dstack/kms/src/admin_service.rs +++ b/dstack/kms/src/admin_service.rs @@ -7,10 +7,10 @@ //! 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, + ClearImageCacheRequest, SignEpochManifestRequest, SignEpochManifestResponse, }; use ra_rpc::{CallContext, RpcCall}; @@ -25,6 +25,23 @@ 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")?, + }) + } } impl RpcCall for AdminRpcHandler { From 035e9cd954b6369db71395774f7f3a4bc65cc9c8 Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Thu, 30 Jul 2026 04:14:00 -0700 Subject: [PATCH 065/122] feat(kms): authorize the unsigned genesis manifest once --- dstack/kms/src/mpc_operation.rs | 21 +++++++++++++-------- 1 file changed, 13 insertions(+), 8 deletions(-) diff --git a/dstack/kms/src/mpc_operation.rs b/dstack/kms/src/mpc_operation.rs index 24e8ab568..bcd1457cf 100644 --- a/dstack/kms/src/mpc_operation.rs +++ b/dstack/kms/src/mpc_operation.rs @@ -302,14 +302,19 @@ impl MpcOperation { payload.manifest.provider_id == manifest.provider_id, "proposed manifest changes provider ID" ); - 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" - ); + 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()?, From a761fc5f2e0990005589af9cd24303b068ace364 Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Thu, 30 Jul 2026 04:14:57 -0700 Subject: [PATCH 066/122] feat(ra-tls): prepare threshold self-signed root CAs --- dstack/ra-tls/src/cert.rs | 58 +++++++++++++++++++++++++++++++++++++++ 1 file changed, 58 insertions(+) diff --git a/dstack/ra-tls/src/cert.rs b/dstack/ra-tls/src/cert.rs index 3522475b0..3d55cd770 100644 --- a/dstack/ra-tls/src/cert.rs +++ b/dstack/ra-tls/src/cert.rs @@ -19,6 +19,7 @@ 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; @@ -508,6 +509,36 @@ pub fn prepare_external_certificate( }) } +/// 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 mut key_id = Sha256::new(); + key_id.update(public_key.public_key_der()); + 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.finalize().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 @@ -686,6 +717,33 @@ mod tests { .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 test_csr_signing_and_verification() { let key_pair = KeyPair::generate_for(&PKCS_ECDSA_P256_SHA256).unwrap(); From 49674594e2c3373780da88803177751cd99d1ca8 Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Thu, 30 Jul 2026 04:15:26 -0700 Subject: [PATCH 067/122] fix(ra-tls): derive external root key identifiers from SPKI --- dstack/ra-tls/src/cert.rs | 25 +++++++++++++++++++++---- 1 file changed, 21 insertions(+), 4 deletions(-) diff --git a/dstack/ra-tls/src/cert.rs b/dstack/ra-tls/src/cert.rs index 3d55cd770..dac28ddeb 100644 --- a/dstack/ra-tls/src/cert.rs +++ b/dstack/ra-tls/src/cert.rs @@ -516,13 +516,30 @@ pub fn prepare_external_self_signed( ) -> Result { let public_key = request.key; let params = request.into_cert_params()?; - let mut key_id = Sha256::new(); - key_id.update(public_key.public_key_der()); + 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.finalize().to_vec()); + 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")?; From 95807781e86d8d980ae0db9e8c219d15166f5601 Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Thu, 30 Jul 2026 04:16:08 -0700 Subject: [PATCH 068/122] fix(kms): make threshold-derived signing keys total --- dstack/kms/src/key_backend.rs | 52 ++++++++++++++++++++++----------- dstack/kms/src/mpc_operation.rs | 3 ++ 2 files changed, 38 insertions(+), 17 deletions(-) diff --git a/dstack/kms/src/key_backend.rs b/dstack/kms/src/key_backend.rs index 27ce31795..e582758e1 100644 --- a/dstack/kms/src/key_backend.rs +++ b/dstack/kms/src/key_backend.rs @@ -923,18 +923,26 @@ impl KeyBackend for MpcKeyBackend { 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 app_seed = self - .derive(DerivePayload { - purpose: DerivePurpose::AppK256, - app_id: app_id.to_vec(), - instance_id: None, - }) - .await?; - let k256_key = SigningKey::from_slice(&app_seed) - .context("threshold derivation produced an invalid K-256 key")?; + 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) @@ -952,6 +960,7 @@ impl KeyBackend for MpcKeyBackend { purpose: DerivePurpose::EnvKey, app_id: app_id.to_vec(), instance_id: None, + counter: 0, }) .await } @@ -1011,14 +1020,23 @@ impl KeyBackend for MpcKeyBackend { } async fn derive_app_ca(&self, app_id: &[u8]) -> Result { - let seed = self - .derive(DerivePayload { - purpose: DerivePurpose::AppCa, - app_id: app_id.to_vec(), - instance_id: None, - }) - .await?; - let app_key = kdf::derive_p256_key_pair_from_bytes(&seed, &[app_id, b"app-ca"])?; + 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 request = CertRequest::builder() .key(&app_key) .org_name("Dstack") diff --git a/dstack/kms/src/mpc_operation.rs b/dstack/kms/src/mpc_operation.rs index bcd1457cf..e406e1202 100644 --- a/dstack/kms/src/mpc_operation.rs +++ b/dstack/kms/src/mpc_operation.rs @@ -42,6 +42,8 @@ pub(crate) struct DerivePayload { 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)] @@ -95,6 +97,7 @@ impl DerivePayload { 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 From 397874c92ce6460947fd3989bb4bee1234073487 Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Thu, 30 Jul 2026 04:16:35 -0700 Subject: [PATCH 069/122] fix(kms): constrain threshold-signed app CA identity --- dstack/kms/src/key_backend.rs | 19 +++++++++++++++++++ 1 file changed, 19 insertions(+) diff --git a/dstack/kms/src/key_backend.rs b/dstack/kms/src/key_backend.rs index e582758e1..951dd6d81 100644 --- a/dstack/kms/src/key_backend.rs +++ b/dstack/kms/src/key_backend.rs @@ -560,6 +560,25 @@ impl MpcKeyBackend { 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")? From b604be538f3fa856145f5188076c5eb4662e602c Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Thu, 30 Jul 2026 04:17:02 -0700 Subject: [PATCH 070/122] fix(kms): bound threshold app CA validity from issuance --- dstack/kms/src/key_backend.rs | 3 +++ 1 file changed, 3 insertions(+) diff --git a/dstack/kms/src/key_backend.rs b/dstack/kms/src/key_backend.rs index 951dd6d81..f16dd9256 100644 --- a/dstack/kms/src/key_backend.rs +++ b/dstack/kms/src/key_backend.rs @@ -1056,6 +1056,7 @@ impl KeyBackend for MpcKeyBackend { } 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") @@ -1063,6 +1064,8 @@ impl KeyBackend for MpcKeyBackend { .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")?; From 02e0dbf8c30fdf396e2e11b22ec9fd481fca873d Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Thu, 30 Jul 2026 04:20:29 -0700 Subject: [PATCH 071/122] feat(kms): prepare threshold-preserving revoke epochs --- dstack/kms/rpc/proto/kms_rpc.proto | 11 + dstack/kms/src/admin_service.rs | 16 +- dstack/kms/src/key_backend.rs | 342 ++++++++++++++++++++++++++++- dstack/kms/src/mpc_lifecycle.rs | 56 +++++ dstack/kms/src/mpc_operation.rs | 22 ++ 5 files changed, 442 insertions(+), 5 deletions(-) diff --git a/dstack/kms/rpc/proto/kms_rpc.proto b/dstack/kms/rpc/proto/kms_rpc.proto index 2ff216e10..dd7b30773 100644 --- a/dstack/kms/rpc/proto/kms_rpc.proto +++ b/dstack/kms/rpc/proto/kms_rpc.proto @@ -149,6 +149,8 @@ service Admin { 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); } message ClearImageCacheRequest { @@ -164,6 +166,15 @@ 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 BootstrapRequest { string domain = 1; } diff --git a/dstack/kms/src/admin_service.rs b/dstack/kms/src/admin_service.rs index 834d577c8..81f423be8 100644 --- a/dstack/kms/src/admin_service.rs +++ b/dstack/kms/src/admin_service.rs @@ -10,7 +10,8 @@ use anyhow::{Context, Result}; use dstack_kms_rpc::{ admin_server::{AdminRpc, AdminServer}, - ClearImageCacheRequest, SignEpochManifestRequest, SignEpochManifestResponse, + ClearImageCacheRequest, PrepareReshareRequest, PrepareReshareResponse, + SignEpochManifestRequest, SignEpochManifestResponse, }; use ra_rpc::{CallContext, RpcCall}; @@ -42,6 +43,19 @@ impl AdminRpc for AdminRpcHandler { .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")?, + }) + } } impl RpcCall for AdminRpcHandler { diff --git a/dstack/kms/src/key_backend.rs b/dstack/kms/src/key_backend.rs index f16dd9256..b6d8de852 100644 --- a/dstack/kms/src/key_backend.rs +++ b/dstack/kms/src/key_backend.rs @@ -16,7 +16,7 @@ use std::{ use anyhow::{bail, ensure, Context, Result}; use async_trait::async_trait; -use cggmp21::key_share::AnyKeyShare as _; +use cggmp21::key_share::{AnyKeyShare as _, Valid, ValidateFromParts as _}; use cggmp21::{ generic_ec::Scalar, supported_curves::{Secp256k1, Secp256r1}, @@ -32,18 +32,22 @@ use rand::{rngs::OsRng, RngCore}; use crate::{ cggmp_engine::{ - execution_id, load_share, share_commitment, CggmpCurve, K256KeyShare, P256KeyShare, + execution_id, load_share, share_commitment, store_share, CggmpCurve, K256KeyShare, + P256KeyShare, }, crypto::{derive_k256_key, sign_message, sign_message_with_timestamp}, mpc_driver::{ - drive_state_machine_blocking, BlockingHttpTransport, DriverContext, MpcHttpTransport, + drive_state_machine_blocking, BlockingHttpTransport, DriverContext, EnvelopeTransport, + MpcHttpTransport, }, mpc_identity::{EpochManifest, SignedEpochManifest}, + mpc_lifecycle::ResharePlan, mpc_operation::{ DerivePayload, DerivePurpose, K256SignPayload, ManifestSignaturePayload, MpcOperation, MpcOperationPayload, P256CertificatePayload, }, - mpc_session::{MpcProtocol, SessionRouter}, + mpc_reshare::{self, PrivateContribution, PublicContribution}, + mpc_session::{MpcEnvelope, MpcProtocol, SessionRouter}, threshold_prf::{self, PrfPartial}, }; @@ -72,6 +76,7 @@ pub(crate) trait KeyBackend: Send + Sync { 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 run_mpc_operation( &self, operation: crate::mpc_operation::MpcOperation, @@ -180,6 +185,10 @@ impl KeyBackend for LocalKeyBackend { }) } + async fn prepare_reshare(&self, _plan: ResharePlan) -> Result>> { + bail!("resharing is unavailable on the local key backend") + } + async fn export_root_keys(&self) -> Result<(String, Vec)> { Ok(( self.root_ca.key.serialize_pem(), @@ -207,6 +216,9 @@ pub(crate) struct MpcKeyBackend { 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, operations: Mutex, Arc>>, } @@ -218,6 +230,16 @@ struct OperationRecord { 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 { @@ -277,6 +299,9 @@ impl MpcKeyBackend { manifest: manifest.clone(), cluster_id: cluster_id.into(), 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(), operations: Mutex::new(BTreeMap::new()), }; ensure!( @@ -537,6 +562,53 @@ impl MpcKeyBackend { }) } + async fn coordinate_reshare(&self, plan: ResharePlan) -> Result>> { + plan.validate_subset(&self.manifest)?; + let dealers = self.participants().await?; + 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::{ @@ -761,7 +833,265 @@ impl MpcKeyBackend { )?; 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> { + 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> { @@ -1084,6 +1414,10 @@ impl KeyBackend for MpcKeyBackend { self.coordinate_manifest_signature(manifest).await } + async fn prepare_reshare(&self, plan: ResharePlan) -> Result>> { + self.coordinate_reshare(plan).await + } + async fn run_mpc_operation(&self, operation: MpcOperation, initiator: &str) -> Result> { operation.validate(&self.manifest, initiator)?; self.execute_once(operation, initiator).await diff --git a/dstack/kms/src/mpc_lifecycle.rs b/dstack/kms/src/mpc_lifecycle.rs index 24541a332..d70681479 100644 --- a/dstack/kms/src/mpc_lifecycle.rs +++ b/dstack/kms/src/mpc_lifecycle.rs @@ -12,6 +12,62 @@ use serde::{Deserialize, Serialize}; use crate::mpc_identity::{ClusterIdentity, SignedEpochManifest}; +#[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, + pub members: Vec, +} + +impl ResharePlan { + pub(crate) fn validate_subset( + &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" + ); + 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(()) + } +} + #[derive(Serialize, Deserialize)] struct EpochCheckpoint { version: u16, diff --git a/dstack/kms/src/mpc_operation.rs b/dstack/kms/src/mpc_operation.rs index e406e1202..b7c6e9acc 100644 --- a/dstack/kms/src/mpc_operation.rs +++ b/dstack/kms/src/mpc_operation.rs @@ -10,6 +10,7 @@ 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; @@ -162,6 +163,7 @@ pub(crate) enum MpcOperationPayload { SignManifest(ManifestSignaturePayload), SignP256Certificate(P256CertificatePayload), Derive(DerivePayload), + Reshare(ResharePlan), } #[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] @@ -263,6 +265,25 @@ impl MpcOperation { 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 validate(&self, manifest: &EpochManifest, initiator: &str) -> Result<()> { ensure!( self.session_id.len() == 32, @@ -321,6 +342,7 @@ impl MpcOperation { } MpcOperationPayload::SignP256Certificate(payload) => payload.validate()?, MpcOperationPayload::Derive(payload) => payload.validate()?, + MpcOperationPayload::Reshare(plan) => plan.validate_subset(manifest)?, } ensure!( self.request_hash == self.compute_request_hash()?, From 52cff72189e9d2fec656ea216c1c550dc4e54593 Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Thu, 30 Jul 2026 04:20:46 -0700 Subject: [PATCH 072/122] fix(kms): reject recovery IDs for reshare results --- dstack/kms/src/key_backend.rs | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/dstack/kms/src/key_backend.rs b/dstack/kms/src/key_backend.rs index b6d8de852..004e39a8e 100644 --- a/dstack/kms/src/key_backend.rs +++ b/dstack/kms/src/key_backend.rs @@ -16,7 +16,7 @@ use std::{ use anyhow::{bail, ensure, Context, Result}; use async_trait::async_trait; -use cggmp21::key_share::{AnyKeyShare as _, Valid, ValidateFromParts as _}; +use cggmp21::key_share::{AnyKeyShare as _, Valid}; use cggmp21::{ generic_ec::Scalar, supported_curves::{Secp256k1, Secp256r1}, @@ -715,6 +715,7 @@ impl MpcKeyBackend { 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")?; From 53d27bf2fc5e23947511b3c9d71cc67441cc5834 Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Thu, 30 Jul 2026 04:21:34 -0700 Subject: [PATCH 073/122] test(kms): separate revoke resharing from join authorization --- dstack/kms/src/mpc_lifecycle.rs | 24 ++++++++++++++++++++++++ 1 file changed, 24 insertions(+) diff --git a/dstack/kms/src/mpc_lifecycle.rs b/dstack/kms/src/mpc_lifecycle.rs index d70681479..ffd437e12 100644 --- a/dstack/kms/src/mpc_lifecycle.rs +++ b/dstack/kms/src/mpc_lifecycle.rs @@ -246,4 +246,28 @@ mod tests { 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, + 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()); + } } From a69f23dced7b0ffb798ff3e8c82bf540380b64b1 Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Thu, 30 Jul 2026 04:22:28 -0700 Subject: [PATCH 074/122] feat(kms): publish the signed active MPC epoch --- dstack/kms/rpc/proto/kms_rpc.proto | 2 ++ dstack/kms/src/main_service.rs | 12 ++++++++++++ 2 files changed, 14 insertions(+) diff --git a/dstack/kms/rpc/proto/kms_rpc.proto b/dstack/kms/rpc/proto/kms_rpc.proto index dd7b30773..0d1c1876c 100644 --- a/dstack/kms/rpc/proto/kms_rpc.proto +++ b/dstack/kms/rpc/proto/kms_rpc.proto @@ -63,6 +63,8 @@ message GetMetaResponse { 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; } message GetKmsKeyRequest { diff --git a/dstack/kms/src/main_service.rs b/dstack/kms/src/main_service.rs index fdc5efbe0..3a4ea5ee2 100644 --- a/dstack/kms/src/main_service.rs +++ b/dstack/kms/src/main_service.rs @@ -62,6 +62,7 @@ pub struct KmsStateInner { self_boot_info: OnceCell, metrics: KmsMetrics, mpc_identity: Option, + signed_manifest: Option, key_backend: Arc, mpc_router: Option>, } @@ -291,6 +292,7 @@ impl KmsState { self_boot_info: OnceCell::new(), metrics: KmsMetrics::default(), mpc_identity, + signed_manifest, key_backend, mpc_router, }), @@ -586,6 +588,16 @@ impl KmsRpc for RpcHandler { .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")?, }) } From cb519aaa3d36ad7763e989d505784bdf9f607f35 Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Thu, 30 Jul 2026 04:23:52 -0700 Subject: [PATCH 075/122] fix(kms): enforce canonical threshold-share topology --- dstack/kms/src/cggmp_engine.rs | 38 ++++++++++++++++++++++++++++++++++ dstack/kms/src/key_backend.rs | 22 ++++++++++++++++++-- 2 files changed, 58 insertions(+), 2 deletions(-) diff --git a/dstack/kms/src/cggmp_engine.rs b/dstack/kms/src/cggmp_engine.rs index 21f258d68..10984873f 100644 --- a/dstack/kms/src/cggmp_engine.rs +++ b/dstack/kms/src/cggmp_engine.rs @@ -8,6 +8,7 @@ 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, }; @@ -94,6 +95,43 @@ pub(crate) fn share_commitment( 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, diff --git a/dstack/kms/src/key_backend.rs b/dstack/kms/src/key_backend.rs index 004e39a8e..4e7686bff 100644 --- a/dstack/kms/src/key_backend.rs +++ b/dstack/kms/src/key_backend.rs @@ -32,8 +32,8 @@ use rand::{rngs::OsRng, RngCore}; use crate::{ cggmp_engine::{ - execution_id, load_share, share_commitment, store_share, CggmpCurve, K256KeyShare, - P256KeyShare, + 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::{ @@ -318,6 +318,24 @@ impl MpcKeyBackend { .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() From 97d546e7169894dc9fbe6555aadc9bf8a1e149f8 Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Thu, 30 Jul 2026 04:26:50 -0700 Subject: [PATCH 076/122] feat(kms): load generated MPC cluster identities --- dstack/kms/kms.toml | 1 + dstack/kms/src/config.rs | 2 ++ dstack/kms/src/main_service.rs | 36 ++++++++++++++++++++++---------- dstack/kms/src/mpc_identity.rs | 38 ++++++++++++++++++++++------------ 4 files changed, 53 insertions(+), 24 deletions(-) diff --git a/dstack/kms/kms.toml b/dstack/kms/kms.toml index f193dda1b..fda4ae1a0 100644 --- a/dstack/kms/kms.toml +++ b/dstack/kms/kms.toml @@ -104,6 +104,7 @@ enabled = true enabled = false protocol_version = 1 cluster_id = "" +identity_file = "" p256_group_pubkey = "" k256_group_pubkey = "" derivation_group_pubkey = "" diff --git a/dstack/kms/src/config.rs b/dstack/kms/src/config.rs index 48ed01fc7..a3b90bf23 100644 --- a/dstack/kms/src/config.rs +++ b/dstack/kms/src/config.rs @@ -94,6 +94,8 @@ pub(crate) struct MpcConfig { #[serde(default)] pub cluster_id: String, #[serde(default)] + pub identity_file: PathBuf, + #[serde(default)] pub p256_group_pubkey: String, #[serde(default)] pub k256_group_pubkey: String, diff --git a/dstack/kms/src/main_service.rs b/dstack/kms/src/main_service.rs index 3a4ea5ee2..33f38b198 100644 --- a/dstack/kms/src/main_service.rs +++ b/dstack/kms/src/main_service.rs @@ -144,19 +144,30 @@ impl KmsState { .context("failed to load attestation verifier")?, ); let configured_identity = if config.mpc.enabled { - Some(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, - )?, - )?) + 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()?; + } let mut signed_manifest: Option = None; let manifest: Option = if config.mpc.enabled { anyhow::ensure!( @@ -205,7 +216,10 @@ impl KmsState { let manifest = manifest.as_ref().context("MPC manifest is missing")?; Arc::new(MpcKeyBackend::load( root_ca_cert, - &config.mpc.cluster_id, + &configured_identity + .as_ref() + .context("MPC identity is missing")? + .cluster_id, manifest.epoch, &config.mpc.node_id, &config.mpc.p256_share_file, diff --git a/dstack/kms/src/mpc_identity.rs b/dstack/kms/src/mpc_identity.rs index 26f1fd6c2..a071998f8 100644 --- a/dstack/kms/src/mpc_identity.rs +++ b/dstack/kms/src/mpc_identity.rs @@ -36,28 +36,40 @@ impl ClusterIdentity { k256_group_pubkey: Vec, derivation_group_pubkey: Vec, ) -> Result { - ensure!(protocol_version == 1, "unsupported MPC protocol version"); - ensure!(!cluster_id.is_empty(), "MPC cluster_id must not be empty"); - ensure!(cluster_id.len() <= 128, "MPC cluster_id is too long"); + 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!( - !p256_group_pubkey.is_empty(), + 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!( - !k256_group_pubkey.is_empty(), + !self.k256_group_pubkey.is_empty(), "missing K-256 group public key" ); ensure!( - !derivation_group_pubkey.is_empty(), + !self.derivation_group_pubkey.is_empty(), "missing derivation group public key" ); - Ok(Self { - protocol_version, - cluster_id, - p256_group_pubkey, - k256_group_pubkey, - derivation_group_pubkey, - }) + Ok(()) } /// Stable key-provider ID pinned in app compose and launch measurements. From d990b36402999ddee7e022a6164d103e6fb37aec Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Thu, 30 Jul 2026 04:27:50 -0700 Subject: [PATCH 077/122] refactor(kms): generalize blocking MPC transports --- dstack/kms/src/key_backend.rs | 6 +++--- dstack/kms/src/mpc_driver.rs | 16 +++++++--------- 2 files changed, 10 insertions(+), 12 deletions(-) diff --git a/dstack/kms/src/key_backend.rs b/dstack/kms/src/key_backend.rs index 4e7686bff..f4a65ec8a 100644 --- a/dstack/kms/src/key_backend.rs +++ b/dstack/kms/src/key_backend.rs @@ -37,7 +37,7 @@ use crate::{ }, crypto::{derive_k256_key, sign_message, sign_message_with_timestamp}, mpc_driver::{ - drive_state_machine_blocking, BlockingHttpTransport, DriverContext, EnvelopeTransport, + drive_state_machine_blocking, BlockingTransport, DriverContext, EnvelopeTransport, MpcHttpTransport, }, mpc_identity::{EpochManifest, SignedEpochManifest}, @@ -1186,7 +1186,7 @@ impl MpcKeyBackend { .sign_sync(&mut rng, data); let signature = drive_state_machine_blocking( state, - &BlockingHttpTransport::new(transport, runtime), + &BlockingTransport::new(transport, runtime), context, )??; let mut encoded = vec![0u8; cggmp21::Signature::::serialized_len()]; @@ -1264,7 +1264,7 @@ impl MpcKeyBackend { .sign_sync(&mut rng, data); let signature = drive_state_machine_blocking( state, - &BlockingHttpTransport::new(transport, runtime), + &BlockingTransport::new(transport, runtime), context, )??; let mut encoded = vec![0u8; cggmp21::Signature::::serialized_len()]; diff --git a/dstack/kms/src/mpc_driver.rs b/dstack/kms/src/mpc_driver.rs index 79753e8a6..011f219d8 100644 --- a/dstack/kms/src/mpc_driver.rs +++ b/dstack/kms/src/mpc_driver.rs @@ -53,16 +53,13 @@ pub(crate) trait EnvelopeTransport: Send + Sync { async fn receive(&self, session_id: &[u8; 32]) -> Result>; } -pub(crate) struct BlockingHttpTransport { - inner: std::sync::Arc, +pub(crate) struct BlockingTransport { + inner: std::sync::Arc, runtime: tokio::runtime::Handle, } -impl BlockingHttpTransport { - pub(crate) fn new( - inner: std::sync::Arc, - runtime: tokio::runtime::Handle, - ) -> Self { +impl BlockingTransport { + pub(crate) fn new(inner: std::sync::Arc, runtime: tokio::runtime::Handle) -> Self { Self { inner, runtime } } @@ -290,14 +287,15 @@ where /// 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( +pub(crate) fn drive_state_machine_blocking( mut state: S, - transport: &BlockingHttpTransport, + transport: &BlockingTransport, context: DriverContext, ) -> Result where S: StateMachine, S::Msg: Serialize + DeserializeOwned, + T: EnvelopeTransport, { ensure!( context.participants.len() >= 2, From dfbdce8233ffbc1c65c4e28d96f2e13bab57279c Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Thu, 30 Jul 2026 04:29:58 -0700 Subject: [PATCH 078/122] feat(kms): add attested distributed genesis DKG service --- dstack/kms/kms.toml | 2 + dstack/kms/rpc/proto/kms_rpc.proto | 21 + dstack/kms/src/config.rs | 4 + dstack/kms/src/main.rs | 1 + dstack/kms/src/mpc_genesis.rs | 643 +++++++++++++++++++++++++++++ 5 files changed, 671 insertions(+) create mode 100644 dstack/kms/src/mpc_genesis.rs diff --git a/dstack/kms/kms.toml b/dstack/kms/kms.toml index fda4ae1a0..c7ab1e8cb 100644 --- a/dstack/kms/kms.toml +++ b/dstack/kms/kms.toml @@ -105,6 +105,8 @@ enabled = false protocol_version = 1 cluster_id = "" identity_file = "" +genesis_plan_file = "" +genesis_tls_ca_cert = "" p256_group_pubkey = "" k256_group_pubkey = "" derivation_group_pubkey = "" diff --git a/dstack/kms/rpc/proto/kms_rpc.proto b/dstack/kms/rpc/proto/kms_rpc.proto index 0d1c1876c..883dd12a9 100644 --- a/dstack/kms/rpc/proto/kms_rpc.proto +++ b/dstack/kms/rpc/proto/kms_rpc.proto @@ -142,6 +142,27 @@ service MpcTransport { 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); +} + // 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 diff --git a/dstack/kms/src/config.rs b/dstack/kms/src/config.rs index a3b90bf23..880ba6a83 100644 --- a/dstack/kms/src/config.rs +++ b/dstack/kms/src/config.rs @@ -96,6 +96,10 @@ pub(crate) struct MpcConfig { #[serde(default)] pub identity_file: PathBuf, #[serde(default)] + pub genesis_plan_file: PathBuf, + #[serde(default)] + pub genesis_tls_ca_cert: PathBuf, + #[serde(default)] pub p256_group_pubkey: String, #[serde(default)] pub k256_group_pubkey: String, diff --git a/dstack/kms/src/main.rs b/dstack/kms/src/main.rs index 418d85506..8f96fa082 100644 --- a/dstack/kms/src/main.rs +++ b/dstack/kms/src/main.rs @@ -25,6 +25,7 @@ mod crypto; mod key_backend; mod main_service; mod mpc_driver; +mod mpc_genesis; mod mpc_identity; mod mpc_lifecycle; mod mpc_operation; diff --git a/dstack/kms/src/mpc_genesis.rs b/dstack/kms/src/mpc_genesis.rs new file mode 100644 index 000000000..d6a1418ba --- /dev/null +++ b/dstack/kms/src/mpc_genesis.rs @@ -0,0 +1,643 @@ +// SPDX-FileCopyrightText: © 2024-2026 Phala Network +// +// SPDX-License-Identifier: Apache-2.0 + +//! One-time, attested distributed genesis DKG. + +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 _, ValidateFromParts as _}, + supported_curves::{Secp256k1, Secp256r1}, + 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 rand::{rngs::OsRng, RngCore}; +use serde::{Deserialize, Serialize}; +use sha2::{Digest, Sha256}; + +use crate::{ + cggmp_engine::{execution_id, store_share, CggmpCurve, K256KeyShare, P256KeyShare}, + config::KmsConfig, + mpc_driver::{ + drive_state_machine_blocking, BlockingTransport, DriverContext, EnvelopeTransport, + }, + mpc_identity::{EpochManifest, EpochMember}, + 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, +} + +#[derive(Clone, Debug, Serialize, Deserialize)] +struct GenesisOperation { + #[serde(with = "hex_bytes")] + session_id: Vec, + kind: GenesisKind, + 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, + expires_at: u64, +} + +impl GenesisOperation { + fn new(kind: GenesisKind) -> Result { + let mut session = [0u8; 32]; + OsRng.fill_bytes(&mut session); + let mut operation = Self { + session_id: session.to_vec(), + kind, + 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, + 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, +} + +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_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, +} + +#[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 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.rpc_cert())?, + fs_err::read_to_string(config.rpc_key())?, + ca, + verifier.clone(), + )?); + Ok(Self(Arc::new(GenesisInner { + config, + plan, + router, + transport, + artifacts: Mutex::new(GenesisArtifacts::default()), + verifier, + }))) + } + + pub(crate) fn verifier(&self) -> Arc { + self.0.verifier.clone() + } + + 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") + } + + async fn execute(&self, operation: GenesisOperation, initiator: &str) -> Result> { + operation.validate()?; + ensure!( + initiator == self.0.plan.coordinator, + "only genesis coordinator may start DKG" + ); + 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, + }; + 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(Into::into) + }) + .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(Into::into) + }) + .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(Into::into) + }) + .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) + } + } + } + + pub(crate) async fn coordinate_dkg(&self) -> Result<()> { + ensure!( + self.0.config.mpc.node_id == self.0.plan.coordinator, + "not genesis coordinator" + ); + for kind in [ + GenesisKind::AuxiliaryInfo, + GenesisKind::DkgP256, + GenesisKind::DkgK256, + GenesisKind::DkgDerivation, + ] { + let operation = GenesisOperation::new(kind)?; + let remote = self + .0 + .plan + .members + .iter() + .filter(|member| member.node_id != self.0.config.mpc.node_id) + .map(|member| self.0.transport.start(&member.node_id, &operation)); + let (local, remote) = tokio::join!( + self.execute(operation.clone(), &self.0.plan.coordinator), + futures::future::join_all(remote) + ); + let local = local?; + for result in remote { + ensure!( + result? == local, + "genesis group public keys differ across nodes" + ); + } + } + self.persist_shares() + } + + 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<()> { + bail!("genesis finalization has not been authorized") + } +} + +pub(crate) fn rpc_methods() -> &'static [&'static str] { + >::supported_methods() +} + +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) + } +} From 46a89e90a9bb0dec66d94b91e626baf4bc3ffb3f Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Thu, 30 Jul 2026 04:37:29 -0700 Subject: [PATCH 079/122] fix(kms): type genesis protocol failures explicitly --- dstack/kms/src/mpc_genesis.rs | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) diff --git a/dstack/kms/src/mpc_genesis.rs b/dstack/kms/src/mpc_genesis.rs index d6a1418ba..d62eeadd0 100644 --- a/dstack/kms/src/mpc_genesis.rs +++ b/dstack/kms/src/mpc_genesis.rs @@ -13,7 +13,6 @@ use std::{ use anyhow::{bail, ensure, Context, Result}; use async_trait::async_trait; use cggmp21::{ - key_share::{AnyKeyShare as _, ValidateFromParts as _}, supported_curves::{Secp256k1, Secp256r1}, ExecutionId, KeyShare, }; @@ -420,7 +419,7 @@ impl GenesisState { &BlockingTransport::new(transport, runtime), context, )? - .map_err(Into::into) + .map_err(|error| anyhow::anyhow!("genesis auxiliary protocol failed: {error}")) }) .await .context("auxiliary worker panicked")??; @@ -453,7 +452,9 @@ impl GenesisState { &BlockingTransport::new(transport, runtime), context, )??; - KeyShare::from_parts((core, aux)).map_err(Into::into) + 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")??; @@ -491,7 +492,9 @@ impl GenesisState { &BlockingTransport::new(transport, runtime), context, )??; - KeyShare::from_parts((core, aux)).map_err(Into::into) + 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")??; From e8cc7f2f387534185a33a497c34122b1e7f6b51e Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Thu, 30 Jul 2026 04:37:52 -0700 Subject: [PATCH 080/122] fix(kms): retain genesis key share trait import --- dstack/kms/src/mpc_genesis.rs | 1 + 1 file changed, 1 insertion(+) diff --git a/dstack/kms/src/mpc_genesis.rs b/dstack/kms/src/mpc_genesis.rs index d62eeadd0..a719693f8 100644 --- a/dstack/kms/src/mpc_genesis.rs +++ b/dstack/kms/src/mpc_genesis.rs @@ -13,6 +13,7 @@ use std::{ use anyhow::{bail, ensure, Context, Result}; use async_trait::async_trait; use cggmp21::{ + key_share::AnyKeyShare as _, supported_curves::{Secp256k1, Secp256r1}, ExecutionId, KeyShare, }; From a605c4c2c0cfbba9fa280a1bdf33a28e81c1cf04 Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Thu, 30 Jul 2026 04:39:03 -0700 Subject: [PATCH 081/122] feat(kms): launch attested MPC genesis mode --- dstack/kms/src/config.rs | 10 +++++++ dstack/kms/src/main.rs | 51 +++++++++++++++++++++++++++++++++-- dstack/kms/src/mpc_genesis.rs | 33 +++++++++++++++++++++++ 3 files changed, 92 insertions(+), 2 deletions(-) diff --git a/dstack/kms/src/config.rs b/dstack/kms/src/config.rs index 880ba6a83..ad66710d1 100644 --- a/dstack/kms/src/config.rs +++ b/dstack/kms/src/config.rs @@ -168,6 +168,16 @@ 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() + } + pub fn keys_exists(&self) -> bool { if self.mpc.enabled { return self.tmp_ca_cert().exists() diff --git a/dstack/kms/src/main.rs b/dstack/kms/src/main.rs index 8f96fa082..b63e0f171 100644 --- a/dstack/kms/src/main.rs +++ b/dstack/kms/src/main.rs @@ -87,6 +87,41 @@ async fn run_onboard_service(kms_config: KmsConfig, figment: Figment) -> Result< 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()); + let coordinator = state.clone(); + if coordinator.is_coordinator() { + tokio::spawn(async move { + if let Err(error) = coordinator.coordinate_dkg().await { + tracing::error!("MPC genesis DKG failed: {error:#}"); + } else { + info!("MPC genesis DKG completed; waiting for threshold finalization"); + } + }); + } + 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")?)); + rocket::custom(figment) + .mount( + "/prpc", + ra_rpc::prpc_routes!(GenesisState, GenesisHandler, trim: "MpcGenesis."), + ) + .manage(state) + .manage(quote_verifier) + .launch() + .await + .map_err(|error| anyhow!(error.to_string()))?; + Ok(()) +} + #[rocket::get("/metrics")] fn metrics(state: &State) -> RawText { RawText(state.metrics().render_prometheus()) @@ -135,14 +170,26 @@ async fn main() -> Result<()> { ); } - if config.onboard.enabled && !config.keys_exists() { + if config.onboard.enabled + && !(config.mpc.enabled && config.genesis_transport_exists()) + && !config.keys_exists() + { info!("Onboarding"); run_onboard_service(config.clone(), figment.clone()).await?; - if !config.keys_exists() { + if config.mpc.enabled && !config.genesis_transport_exists() + || !config.mpc.enabled && !config.keys_exists() + { bail!("Failed to onboard"); } } + if config.mpc.enabled && !config.keys_exists() { + if !config.genesis_transport_exists() { + bail!("MPC genesis requires onboarded RPC transport certificates"); + } + return run_mpc_genesis_service(config, figment).await; + } + info!("Updating certs"); if let Err(err) = onboard_service::update_certs(&config).await { if config.attest_rpc_cert { diff --git a/dstack/kms/src/mpc_genesis.rs b/dstack/kms/src/mpc_genesis.rs index a719693f8..1a8626283 100644 --- a/dstack/kms/src/mpc_genesis.rs +++ b/dstack/kms/src/mpc_genesis.rs @@ -22,6 +22,7 @@ use dstack_kms_rpc::{ mpc_genesis_server::{MpcGenesisRpc, MpcGenesisServer}, MpcGenesisFinalizeRequest, MpcGenesisStartRequest, MpcGenesisStartResponse, MpcPushRequest, }; +use prost_types::Empty; use ra_rpc::{ client::{RaClient, RaClientConfig}, CallContext, RpcCall, @@ -270,6 +271,33 @@ impl GenesisTransport { .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(Empty {}) + .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] @@ -353,6 +381,10 @@ impl GenesisState { self.0.verifier.clone() } + pub(crate) fn is_coordinator(&self) -> bool { + self.0.config.mpc.node_id == self.0.plan.coordinator + } + fn authenticated_node(&self, key: &[u8]) -> Result<&str> { self.0 .plan @@ -515,6 +547,7 @@ impl GenesisState { 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, From 1388e419ad7c4ee171593b5d0268a9d8d765bbdf Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Thu, 30 Jul 2026 04:39:21 -0700 Subject: [PATCH 082/122] fix(kms): call generated genesis ping API --- dstack/kms/src/mpc_genesis.rs | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/dstack/kms/src/mpc_genesis.rs b/dstack/kms/src/mpc_genesis.rs index 1a8626283..5aabecd61 100644 --- a/dstack/kms/src/mpc_genesis.rs +++ b/dstack/kms/src/mpc_genesis.rs @@ -22,7 +22,6 @@ use dstack_kms_rpc::{ mpc_genesis_server::{MpcGenesisRpc, MpcGenesisServer}, MpcGenesisFinalizeRequest, MpcGenesisStartRequest, MpcGenesisStartResponse, MpcPushRequest, }; -use prost_types::Empty; use ra_rpc::{ client::{RaClient, RaClientConfig}, CallContext, RpcCall, @@ -277,7 +276,7 @@ impl GenesisTransport { loop { let probes = self.clients.iter().map(|(node, client)| async move { client - .ping(Empty {}) + .ping() .await .with_context(|| format!("genesis peer {node} is unavailable")) }); From 5f0726cacafee5926f41362512da9ebd1ef2a020 Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Thu, 30 Jul 2026 04:42:52 -0700 Subject: [PATCH 083/122] feat(kms): threshold-finalize MPC genesis identity --- dstack/kms/src/mpc_genesis.rs | 524 ++++++++++++++++++++++++++++++++-- 1 file changed, 501 insertions(+), 23 deletions(-) diff --git a/dstack/kms/src/mpc_genesis.rs b/dstack/kms/src/mpc_genesis.rs index 5aabecd61..840efcc54 100644 --- a/dstack/kms/src/mpc_genesis.rs +++ b/dstack/kms/src/mpc_genesis.rs @@ -10,12 +10,13 @@ use std::{ time::Duration, }; -use anyhow::{bail, ensure, Context, Result}; +use anyhow::{ensure, Context, Result}; use async_trait::async_trait; use cggmp21::{ + generic_ec::Scalar, key_share::AnyKeyShare as _, supported_curves::{Secp256k1, Secp256r1}, - ExecutionId, KeyShare, + DataToSign, ExecutionId, KeyShare, }; use dstack_kms_rpc::{ mpc_genesis_client::MpcGenesisClient, @@ -27,17 +28,24 @@ use ra_rpc::{ CallContext, RpcCall, }; use ra_tls::attestation::AttestationVerifier; +use ra_tls::{ + cert::{prepare_external_self_signed, CertRequest}, + rcgen::SubjectPublicKeyInfo, +}; use rand::{rngs::OsRng, RngCore}; use serde::{Deserialize, Serialize}; use sha2::{Digest, Sha256}; use crate::{ - cggmp_engine::{execution_id, store_share, CggmpCurve, K256KeyShare, P256KeyShare}, + cggmp_engine::{ + execution_id, share_commitment, store_share, CggmpCurve, K256KeyShare, P256KeyShare, + }, config::KmsConfig, mpc_driver::{ drive_state_machine_blocking, BlockingTransport, DriverContext, EnvelopeTransport, }, - mpc_identity::{EpochManifest, EpochMember}, + mpc_identity::{ClusterIdentity, EpochManifest, EpochMember, SignedEpochManifest}, + mpc_lifecycle::validate_and_checkpoint, mpc_session::{MpcEnvelope, MpcProtocol, SessionRouter}, }; @@ -135,6 +143,9 @@ enum GenesisKind { DkgP256, DkgK256, DkgDerivation, + Commitments, + SignRoot, + SignManifest, } #[derive(Clone, Debug, Serialize, Deserialize)] @@ -142,6 +153,8 @@ 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, @@ -152,16 +165,19 @@ 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) -> Result { + 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")?, @@ -175,6 +191,7 @@ impl GenesisOperation { 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()) @@ -204,6 +221,18 @@ struct GenesisArtifacts { derivation: Option, } +enum GenesisSigningShare { + P256(P256KeyShare), + K256(K256KeyShare), +} + +#[derive(Clone, Debug, Serialize, Deserialize)] +struct GenesisBundle { + identity: ClusterIdentity, + root_ca_pem: String, + signed_manifest: SignedEpochManifest, +} + struct GenesisTransport { local_node_id: String, local_router: Arc, @@ -394,6 +423,116 @@ impl GenesisState { .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(()) + } + async fn execute(&self, operation: GenesisOperation, initiator: &str) -> Result> { operation.validate()?; ensure!( @@ -420,6 +559,9 @@ impl GenesisState { GenesisKind::DkgP256 => MpcProtocol::DkgP256, GenesisKind::DkgK256 => MpcProtocol::DkgK256, GenesisKind::DkgDerivation => MpcProtocol::DkgDerivation, + GenesisKind::Commitments => MpcProtocol::SignK256, + GenesisKind::SignRoot => MpcProtocol::SignP256, + GenesisKind::SignManifest => MpcProtocol::SignK256, }; let context = DriverContext { session_id: operation.session_id.as_slice().try_into()?, @@ -538,6 +680,119 @@ impl GenesisState { .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::SignRoot | GenesisKind::SignManifest => { + 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()? + } + _ => unreachable!(), + }; + let artifacts = self.0.artifacts.lock().expect("genesis mutex poisoned"); + let keygen_indexes = (0..n).collect::>(); + let local_protocol_index = local_index; + let (curve, share) = match operation.kind { + GenesisKind::SignRoot => ( + 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!(), + }; + drop(artifacts); + 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")? + } } } @@ -553,27 +808,181 @@ impl GenesisState { GenesisKind::DkgK256, GenesisKind::DkgDerivation, ] { - let operation = GenesisOperation::new(kind)?; - let remote = self + 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() - .filter(|member| member.node_id != self.0.config.mpc.node_id) - .map(|member| self.0.transport.start(&member.node_id, &operation)); - let (local, remote) = tokio::join!( - self.execute(operation.clone(), &self.0.plan.coordinator), - futures::future::join_all(remote) - ); - let local = local?; - for result in remote { - ensure!( - result? == local, - "genesis group public keys differ across nodes" - ); - } + .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 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, + }, + }; + 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.persist_shares() + 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, remote) = tokio::join!( + self.execute(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")?; + self.persist_shares()?; + ensure!( + !self.0.config.mpc.identity_file.as_os_str().is_empty(), + "missing MPC identity file" + ); + ensure!( + !self.0.config.mpc.manifest_file.as_os_str().is_empty(), + "missing MPC manifest file" + ); + safe_write::safe_write( + &self.0.config.mpc.identity_file, + serde_jcs::to_vec(&bundle.identity)?, + )?; + safe_write::safe_write(self.0.config.root_ca_cert(), bundle.root_ca_pem.as_bytes())?; + safe_write::safe_write( + &self.0.config.mpc.manifest_file, + serde_jcs::to_vec(&bundle.signed_manifest)?, + )?; + validate_and_checkpoint( + &self.0.config.mpc.checkpoint_file, + &bundle.signed_manifest, + &bundle.identity, + )?; + Ok(()) } fn persist_shares(&self) -> Result<()> { @@ -652,8 +1061,15 @@ impl MpcGenesisRpc for GenesisHandler { let envelope = serde_json::from_slice(&request.envelope_json)?; self.state.0.router.push(&sender, envelope) } - async fn finalize(self, _request: MpcGenesisFinalizeRequest) -> Result<()> { - bail!("genesis finalization has not been authorized") + 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) } } @@ -667,6 +1083,68 @@ fn unix_time() -> Result { .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 { From 42988fe77fcb0005f71251a3f61aa06af14c8854 Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Thu, 30 Jul 2026 04:43:45 -0700 Subject: [PATCH 084/122] fix(kms): make genesis finalization futures sendable --- dstack/kms/src/mpc_genesis.rs | 48 ++++++++++++++++++----------------- 1 file changed, 25 insertions(+), 23 deletions(-) diff --git a/dstack/kms/src/mpc_genesis.rs b/dstack/kms/src/mpc_genesis.rs index 840efcc54..e12c163cd 100644 --- a/dstack/kms/src/mpc_genesis.rs +++ b/dstack/kms/src/mpc_genesis.rs @@ -717,31 +717,32 @@ impl GenesisState { } _ => unreachable!(), }; - let artifacts = self.0.artifacts.lock().expect("genesis mutex poisoned"); let keygen_indexes = (0..n).collect::>(); let local_protocol_index = local_index; - let (curve, share) = match operation.kind { - GenesisKind::SignRoot => ( - CggmpCurve::P256, - GenesisSigningShare::P256( - artifacts - .p256 - .clone() - .context("missing P-256 genesis share")?, + let (curve, share) = { + let artifacts = self.0.artifacts.lock().expect("genesis mutex poisoned"); + match operation.kind { + GenesisKind::SignRoot => ( + 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")?, + GenesisKind::SignManifest => ( + CggmpCurve::K256, + GenesisSigningShare::K256( + artifacts + .k256 + .clone() + .context("missing K-256 genesis share")?, + ), ), - ), - _ => unreachable!(), + _ => unreachable!(), + } }; - drop(artifacts); tokio::task::spawn_blocking(move || { let mut rng = OsRng; let eid = execution_id(&cluster, 1, curve, &nonce); @@ -860,7 +861,7 @@ impl GenesisState { .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)?; @@ -870,7 +871,7 @@ impl GenesisState { .coordinate_operation(GenesisOperation::new( GenesisKind::SignManifest, manifest_json, - )) + )?) .await?; let signature = identical_result(&manifest_signatures, "manifest signatures")?.to_vec(); let bundle = GenesisBundle { @@ -922,8 +923,9 @@ impl GenesisState { self.0.transport.start(&member.node_id, &operation).await?, )) }); + let local_operation = operation.clone(); let (local, remote) = tokio::join!( - self.execute(operation, &self.0.plan.coordinator), + self.execute(local_operation, &self.0.plan.coordinator), futures::future::join_all(remote), ); let mut results = BTreeMap::new(); From 1dacd9cbcdbc3a58b6616caf53b905eb92bb1ebe Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Thu, 30 Jul 2026 04:44:28 -0700 Subject: [PATCH 085/122] feat(kms): restart after durable MPC genesis --- dstack/kms/src/main.rs | 32 +++++++++++++++++++------------- dstack/kms/src/mpc_genesis.rs | 21 ++++++++++++++++++++- 2 files changed, 39 insertions(+), 14 deletions(-) diff --git a/dstack/kms/src/main.rs b/dstack/kms/src/main.rs index b63e0f171..a15edaaf5 100644 --- a/dstack/kms/src/main.rs +++ b/dstack/kms/src/main.rs @@ -80,7 +80,7 @@ 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()))?; @@ -92,16 +92,6 @@ async fn run_mpc_genesis_service(kms_config: KmsConfig, figment: Figment) -> Res let state = GenesisState::new(kms_config)?; let quote_verifier = QuoteVerifier::new(state.verifier()); - let coordinator = state.clone(); - if coordinator.is_coordinator() { - tokio::spawn(async move { - if let Err(error) = coordinator.coordinate_dkg().await { - tracing::error!("MPC genesis DKG failed: {error:#}"); - } else { - info!("MPC genesis DKG completed; waiting for threshold finalization"); - } - }); - } info!("Starting attested MPC genesis service"); for method in mpc_genesis::rpc_methods() { info!(" /prpc/{method}"); @@ -109,13 +99,29 @@ async fn run_mpc_genesis_service(kms_config: KmsConfig, figment: Figment) -> Res let figment = figment .clone() .merge(Serialized::defaults(figment.find_value("rpc")?)); - rocket::custom(figment) + let rocket = rocket::custom(figment) .mount( "/prpc", ra_rpc::prpc_routes!(GenesisState, GenesisHandler, trim: "MpcGenesis."), ) .manage(state) - .manage(quote_verifier) + .manage(quote_verifier); + 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()))?; diff --git a/dstack/kms/src/mpc_genesis.rs b/dstack/kms/src/mpc_genesis.rs index e12c163cd..fce0de757 100644 --- a/dstack/kms/src/mpc_genesis.rs +++ b/dstack/kms/src/mpc_genesis.rs @@ -6,7 +6,10 @@ use std::{ collections::BTreeMap, - sync::{Arc, Mutex}, + sync::{ + atomic::{AtomicBool, Ordering}, + Arc, Mutex, + }, time::Duration, }; @@ -357,6 +360,8 @@ struct GenesisInner { transport: Arc, artifacts: Mutex, verifier: Arc, + finalized: AtomicBool, + finalized_notify: tokio::sync::Notify, } #[derive(Clone)] @@ -402,6 +407,8 @@ impl GenesisState { transport, artifacts: Mutex::new(GenesisArtifacts::default()), verifier, + finalized: AtomicBool::new(false), + finalized_notify: tokio::sync::Notify::new(), }))) } @@ -413,6 +420,16 @@ impl GenesisState { 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 @@ -984,6 +1001,8 @@ impl GenesisState { &bundle.signed_manifest, &bundle.identity, )?; + self.0.finalized.store(true, Ordering::Release); + self.0.finalized_notify.notify_waiters(); Ok(()) } From 7ac960887a741f0eea7bb52d430aa17a1dc8e8bf Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Thu, 30 Jul 2026 04:44:40 -0700 Subject: [PATCH 086/122] fix(kms): ignite genesis server before shutdown handle --- dstack/kms/src/main.rs | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/dstack/kms/src/main.rs b/dstack/kms/src/main.rs index a15edaaf5..2f9adac05 100644 --- a/dstack/kms/src/main.rs +++ b/dstack/kms/src/main.rs @@ -106,6 +106,10 @@ async fn run_mpc_genesis_service(kms_config: KmsConfig, figment: Figment) -> Res ) .manage(state) .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 { From ce51cdf90c4c3103165464dc3fb25de545c7d3bb Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Thu, 30 Jul 2026 04:44:56 -0700 Subject: [PATCH 087/122] fix(kms): retain genesis state for lifecycle watcher --- dstack/kms/src/main.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/dstack/kms/src/main.rs b/dstack/kms/src/main.rs index 2f9adac05..d9e45452e 100644 --- a/dstack/kms/src/main.rs +++ b/dstack/kms/src/main.rs @@ -104,7 +104,7 @@ async fn run_mpc_genesis_service(kms_config: KmsConfig, figment: Figment) -> Res "/prpc", ra_rpc::prpc_routes!(GenesisState, GenesisHandler, trim: "MpcGenesis."), ) - .manage(state) + .manage(state.clone()) .manage(quote_verifier); let rocket = rocket .ignite() From 905d0c50aa3f134fc51a278fa8ef4a61b916681d Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Thu, 30 Jul 2026 04:45:31 -0700 Subject: [PATCH 088/122] feat(kms): make genesis operations replay safe --- dstack/kms/src/mpc_genesis.rs | 56 ++++++++++++++++++++++++++++++++++- 1 file changed, 55 insertions(+), 1 deletion(-) diff --git a/dstack/kms/src/mpc_genesis.rs b/dstack/kms/src/mpc_genesis.rs index fce0de757..305a38874 100644 --- a/dstack/kms/src/mpc_genesis.rs +++ b/dstack/kms/src/mpc_genesis.rs @@ -151,7 +151,7 @@ enum GenesisKind { SignManifest, } -#[derive(Clone, Debug, Serialize, Deserialize)] +#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] struct GenesisOperation { #[serde(with = "hex_bytes")] session_id: Vec, @@ -362,6 +362,13 @@ struct GenesisInner { 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)] @@ -409,6 +416,7 @@ impl GenesisState { verifier, finalized: AtomicBool::new(false), finalized_notify: tokio::sync::Notify::new(), + operations: Mutex::new(BTreeMap::new()), }))) } @@ -556,6 +564,52 @@ impl GenesisState { 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() < 32, "too many active 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 From 057802d5d207aeca27855c770123dfce8a1ee431 Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Thu, 30 Jul 2026 04:48:40 -0700 Subject: [PATCH 089/122] feat(kms): atomically activate signed MPC epochs --- dstack/kms/rpc/proto/kms_rpc.proto | 4 + dstack/kms/src/admin_service.rs | 11 +- dstack/kms/src/key_backend.rs | 44 +++++- dstack/kms/src/main_service.rs | 19 ++- dstack/kms/src/mpc_lifecycle.rs | 207 ++++++++++++++++++++++++++++- 5 files changed, 271 insertions(+), 14 deletions(-) diff --git a/dstack/kms/rpc/proto/kms_rpc.proto b/dstack/kms/rpc/proto/kms_rpc.proto index 883dd12a9..6535a4096 100644 --- a/dstack/kms/rpc/proto/kms_rpc.proto +++ b/dstack/kms/rpc/proto/kms_rpc.proto @@ -174,6 +174,7 @@ service Admin { rpc SignEpochManifest(SignEpochManifestRequest) returns (SignEpochManifestResponse); // Prepare threshold shares for a revoke/subset epoch transition. rpc PrepareReshare(PrepareReshareRequest) returns (PrepareReshareResponse); + rpc ActivateEpoch(ActivateEpochRequest) returns (ActivateEpochResponse); } message ClearImageCacheRequest { @@ -198,6 +199,9 @@ message PrepareReshareResponse { bytes commitments_json = 1; } +message ActivateEpochRequest { bytes signed_manifest_json = 1; } +message ActivateEpochResponse { bool retained = 1; } + message BootstrapRequest { string domain = 1; } diff --git a/dstack/kms/src/admin_service.rs b/dstack/kms/src/admin_service.rs index 81f423be8..eefeaaad4 100644 --- a/dstack/kms/src/admin_service.rs +++ b/dstack/kms/src/admin_service.rs @@ -10,8 +10,8 @@ use anyhow::{Context, Result}; use dstack_kms_rpc::{ admin_server::{AdminRpc, AdminServer}, - ClearImageCacheRequest, PrepareReshareRequest, PrepareReshareResponse, - SignEpochManifestRequest, SignEpochManifestResponse, + ActivateEpochRequest, ActivateEpochResponse, ClearImageCacheRequest, PrepareReshareRequest, + PrepareReshareResponse, SignEpochManifestRequest, SignEpochManifestResponse, }; use ra_rpc::{CallContext, RpcCall}; @@ -56,6 +56,13 @@ impl AdminRpc for AdminRpcHandler { .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?; + Ok(ActivateEpochResponse { retained }) + } } impl RpcCall for AdminRpcHandler { diff --git a/dstack/kms/src/key_backend.rs b/dstack/kms/src/key_backend.rs index f4a65ec8a..f12a75a4a 100644 --- a/dstack/kms/src/key_backend.rs +++ b/dstack/kms/src/key_backend.rs @@ -40,8 +40,8 @@ use crate::{ drive_state_machine_blocking, BlockingTransport, DriverContext, EnvelopeTransport, MpcHttpTransport, }, - mpc_identity::{EpochManifest, SignedEpochManifest}, - mpc_lifecycle::ResharePlan, + mpc_identity::{ClusterIdentity, EpochManifest, SignedEpochManifest}, + mpc_lifecycle::{activate_pending_epoch, EpochPaths, ResharePlan}, mpc_operation::{ DerivePayload, DerivePurpose, K256SignPayload, ManifestSignaturePayload, MpcOperation, MpcOperationPayload, P256CertificatePayload, @@ -77,6 +77,7 @@ pub(crate) trait KeyBackend: Send + Sync { 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 run_mpc_operation( &self, operation: crate::mpc_operation::MpcOperation, @@ -189,6 +190,10 @@ impl KeyBackend for LocalKeyBackend { 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 export_root_keys(&self) -> Result<(String, Vec)> { Ok(( self.root_ca.key.serialize_pem(), @@ -219,6 +224,9 @@ pub(crate) struct MpcKeyBackend { 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>>, } @@ -245,12 +253,14 @@ const MAX_ACTIVE_OPERATIONS: usize = 128; impl MpcKeyBackend { pub(crate) fn load( root_ca_cert: String, - cluster_id: &str, + 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, @@ -276,32 +286,35 @@ impl MpcKeyBackend { root_ca_cert, p256_share: load_share( p256_share_file, - cluster_id, + &identity.cluster_id, epoch, node_id, CggmpCurve::P256, )?, k256_share: load_share( k256_share_file, - cluster_id, + &identity.cluster_id, epoch, node_id, CggmpCurve::K256, )?, derivation_share: load_share( derivation_share_file, - cluster_id, + &identity.cluster_id, epoch, node_id, CggmpCurve::P256, )?, transport: Arc::new(transport), manifest: manifest.clone(), - cluster_id: cluster_id.into(), + 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!( @@ -1437,6 +1450,23 @@ impl KeyBackend for MpcKeyBackend { 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 run_mpc_operation(&self, operation: MpcOperation, initiator: &str) -> Result> { operation.validate(&self.manifest, initiator)?; self.execute_once(operation, initiator).await diff --git a/dstack/kms/src/main_service.rs b/dstack/kms/src/main_service.rs index 33f38b198..02e527a9d 100644 --- a/dstack/kms/src/main_service.rs +++ b/dstack/kms/src/main_service.rs @@ -167,6 +167,18 @@ impl KmsState { }; 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 { @@ -216,15 +228,16 @@ impl KmsState { let manifest = manifest.as_ref().context("MPC manifest is missing")?; Arc::new(MpcKeyBackend::load( root_ca_cert, - &configured_identity + configured_identity .as_ref() - .context("MPC identity is missing")? - .cluster_id, + .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.rpc_cert()).context("failed to read MPC RPC cert")?, diff --git a/dstack/kms/src/mpc_lifecycle.rs b/dstack/kms/src/mpc_lifecycle.rs index ffd437e12..f50fbbae6 100644 --- a/dstack/kms/src/mpc_lifecycle.rs +++ b/dstack/kms/src/mpc_lifecycle.rs @@ -4,13 +4,46 @@ //! Durable epoch-chain validation for MPC membership lifecycle changes. -use std::{fs::Permissions, os::unix::fs::PermissionsExt, path::Path}; +use std::{ + fs::Permissions, + os::unix::fs::PermissionsExt, + path::{Path, PathBuf}, +}; use anyhow::{ensure, Context, Result}; use fs_err as fs; use serde::{Deserialize, Serialize}; -use crate::mpc_identity::{ClusterIdentity, SignedEpochManifest}; +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") +} #[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] pub(crate) struct ReshareMember { @@ -157,6 +190,176 @@ pub(crate) fn validate_and_checkpoint( 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}; From ff7449fb8fd2e0588170439d20b4a2589e68e627 Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Thu, 30 Jul 2026 04:49:32 -0700 Subject: [PATCH 090/122] feat(kms): restart after MPC epoch activation --- dstack/kms/src/admin_service.rs | 1 + dstack/kms/src/main.rs | 27 ++++++++++++++++++++++++--- dstack/kms/src/main_service.rs | 21 ++++++++++++++++++++- 3 files changed, 45 insertions(+), 4 deletions(-) diff --git a/dstack/kms/src/admin_service.rs b/dstack/kms/src/admin_service.rs index eefeaaad4..43e763fe5 100644 --- a/dstack/kms/src/admin_service.rs +++ b/dstack/kms/src/admin_service.rs @@ -61,6 +61,7 @@ impl AdminRpc for AdminRpcHandler { 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 }) } } diff --git a/dstack/kms/src/main.rs b/dstack/kms/src/main.rs index d9e45452e..1c18304dd 100644 --- a/dstack/kms/src/main.rs +++ b/dstack/kms/src/main.rs @@ -269,6 +269,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)) => { @@ -279,21 +284,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 02e527a9d..26b423420 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, }, }; @@ -65,6 +65,8 @@ pub struct KmsStateInner { signed_manifest: Option, key_backend: Arc, mpc_router: Option>, + restart_requested: AtomicBool, + restart_notify: tokio::sync::Notify, } #[derive(Default)] @@ -322,6 +324,8 @@ impl KmsState { signed_manifest, key_backend, mpc_router, + restart_requested: AtomicBool::new(false), + restart_notify: tokio::sync::Notify::new(), }), }) } @@ -337,6 +341,21 @@ impl KmsState { &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() } From 6f5023c8e31b55aa20c384508306a2a6a2827e27 Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Thu, 30 Jul 2026 04:50:36 -0700 Subject: [PATCH 091/122] test(kms): recover interrupted revoked epoch activation --- dstack/kms/src/mpc_lifecycle.rs | 48 +++++++++++++++++++++++++++++++++ 1 file changed, 48 insertions(+) diff --git a/dstack/kms/src/mpc_lifecycle.rs b/dstack/kms/src/mpc_lifecycle.rs index f50fbbae6..96775352d 100644 --- a/dstack/kms/src/mpc_lifecycle.rs +++ b/dstack/kms/src/mpc_lifecycle.rs @@ -473,4 +473,52 @@ mod tests { join.members[1].node_id = "kms-new".into(); assert!(join.validate_subset(&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()); + } } From 99aa071b37a934c3019b34d06b9317d7135b7445 Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Thu, 30 Jul 2026 04:52:39 -0700 Subject: [PATCH 092/122] feat(kms): threshold-authorize membership transitions --- dstack/kms/rpc/proto/kms_rpc.proto | 3 ++ dstack/kms/src/admin_service.rs | 18 ++++++- dstack/kms/src/key_backend.rs | 63 +++++++++++++++++++++++-- dstack/kms/src/mpc_lifecycle.rs | 76 +++++++++++++++++++++++++++++- dstack/kms/src/mpc_operation.rs | 21 +++++++++ 5 files changed, 174 insertions(+), 7 deletions(-) diff --git a/dstack/kms/rpc/proto/kms_rpc.proto b/dstack/kms/rpc/proto/kms_rpc.proto index 6535a4096..b250cfe88 100644 --- a/dstack/kms/rpc/proto/kms_rpc.proto +++ b/dstack/kms/rpc/proto/kms_rpc.proto @@ -175,6 +175,7 @@ service Admin { // 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 { @@ -201,6 +202,8 @@ message PrepareReshareResponse { 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 43e763fe5..898330626 100644 --- a/dstack/kms/src/admin_service.rs +++ b/dstack/kms/src/admin_service.rs @@ -10,8 +10,9 @@ use anyhow::{Context, Result}; use dstack_kms_rpc::{ admin_server::{AdminRpc, AdminServer}, - ActivateEpochRequest, ActivateEpochResponse, ClearImageCacheRequest, PrepareReshareRequest, - PrepareReshareResponse, SignEpochManifestRequest, SignEpochManifestResponse, + ActivateEpochRequest, ActivateEpochResponse, AuthorizeReshareRequest, AuthorizeReshareResponse, + ClearImageCacheRequest, PrepareReshareRequest, PrepareReshareResponse, + SignEpochManifestRequest, SignEpochManifestResponse, }; use ra_rpc::{CallContext, RpcCall}; @@ -64,6 +65,19 @@ impl AdminRpc for AdminRpcHandler { 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/key_backend.rs b/dstack/kms/src/key_backend.rs index f12a75a4a..4c3985e3d 100644 --- a/dstack/kms/src/key_backend.rs +++ b/dstack/kms/src/key_backend.rs @@ -41,7 +41,7 @@ use crate::{ MpcHttpTransport, }, mpc_identity::{ClusterIdentity, EpochManifest, SignedEpochManifest}, - mpc_lifecycle::{activate_pending_epoch, EpochPaths, ResharePlan}, + mpc_lifecycle::{activate_pending_epoch, EpochPaths, ResharePlan, SignedResharePlan}, mpc_operation::{ DerivePayload, DerivePurpose, K256SignPayload, ManifestSignaturePayload, MpcOperation, MpcOperationPayload, P256CertificatePayload, @@ -78,6 +78,7 @@ pub(crate) trait KeyBackend: Send + Sync { 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, @@ -194,6 +195,10 @@ impl KeyBackend for LocalKeyBackend { 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(), @@ -593,6 +598,51 @@ impl MpcKeyBackend { }) } + 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 = self.participants().await?; @@ -819,9 +869,9 @@ impl MpcKeyBackend { async fn execute_operation(&self, operation: MpcOperation) -> Result> { match &operation.payload { - MpcOperationPayload::SignK256(_) | MpcOperationPayload::SignManifest(_) => { - self.execute_k256(operation).await - } + MpcOperationPayload::SignK256(_) + | MpcOperationPayload::SignManifest(_) + | MpcOperationPayload::AuthorizeReshare(_) => self.execute_k256(operation).await, MpcOperationPayload::SignP256Certificate(payload) => { self.validate_p256_certificate(payload)?; self.execute_p256(operation).await @@ -1130,6 +1180,7 @@ impl MpcKeyBackend { let digest = match &operation.payload { MpcOperationPayload::SignK256(payload) => payload.digest(), MpcOperationPayload::SignManifest(payload) => payload.digest()?, + MpcOperationPayload::AuthorizeReshare(plan) => plan.authorization_hash()?, _ => bail!("K-256 executor received a different MPC operation"), }; let local_protocol_index: u16 = operation @@ -1467,6 +1518,10 @@ impl KeyBackend for MpcKeyBackend { ) } + 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 diff --git a/dstack/kms/src/mpc_lifecycle.rs b/dstack/kms/src/mpc_lifecycle.rs index 96775352d..668dc086b 100644 --- a/dstack/kms/src/mpc_lifecycle.rs +++ b/dstack/kms/src/mpc_lifecycle.rs @@ -12,7 +12,9 @@ use std::{ 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::{ @@ -62,8 +64,17 @@ pub(crate) struct ResharePlan { 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_subset( + pub(crate) fn validate_transition( &self, active: &crate::mpc_identity::EpochManifest, ) -> Result<()> { @@ -80,6 +91,57 @@ impl ResharePlan { "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" + ); + } + previous = Some(&member.node_id); + } + 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()), @@ -101,6 +163,18 @@ impl ResharePlan { } } +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, diff --git a/dstack/kms/src/mpc_operation.rs b/dstack/kms/src/mpc_operation.rs index b7c6e9acc..d3921d15f 100644 --- a/dstack/kms/src/mpc_operation.rs +++ b/dstack/kms/src/mpc_operation.rs @@ -164,6 +164,7 @@ pub(crate) enum MpcOperationPayload { SignP256Certificate(P256CertificatePayload), Derive(DerivePayload), Reshare(ResharePlan), + AuthorizeReshare(ResharePlan), } #[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] @@ -284,6 +285,25 @@ impl MpcOperation { 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, @@ -343,6 +363,7 @@ impl MpcOperation { 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()?, From 3cb9135957976ad33ce4c69fc2a52d967c93ab4d Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Thu, 30 Jul 2026 04:53:00 -0700 Subject: [PATCH 093/122] fix(kms): classify reshare authorization signatures --- dstack/kms/src/key_backend.rs | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/dstack/kms/src/key_backend.rs b/dstack/kms/src/key_backend.rs index 4c3985e3d..67cdf157e 100644 --- a/dstack/kms/src/key_backend.rs +++ b/dstack/kms/src/key_backend.rs @@ -789,8 +789,8 @@ impl MpcKeyBackend { .context("invalid MPC group public key")?; let digest = match &operation.payload { MpcOperationPayload::SignK256(payload) => payload.digest(), - MpcOperationPayload::SignManifest(_) => { - bail!("manifest signatures do not have recovery IDs") + 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") From d95ad0966870618938720e7ae5087c001a95d9e2 Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Thu, 30 Jul 2026 04:53:22 -0700 Subject: [PATCH 094/122] test(kms): bind joiners in threshold authorization --- dstack/kms/src/mpc_lifecycle.rs | 38 +++++++++++++++++++++++++++++++++ 1 file changed, 38 insertions(+) diff --git a/dstack/kms/src/mpc_lifecycle.rs b/dstack/kms/src/mpc_lifecycle.rs index 668dc086b..750337666 100644 --- a/dstack/kms/src/mpc_lifecycle.rs +++ b/dstack/kms/src/mpc_lifecycle.rs @@ -548,6 +548,44 @@ mod tests { 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, + 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(); From ddaafea9d0da525f67951a27a664b82325a78b82 Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Thu, 30 Jul 2026 04:56:31 -0700 Subject: [PATCH 095/122] refactor(kms): reshare from public key topology --- dstack/kms/src/mpc_reshare.rs | 41 ++++++++++++++++++++++++++++++----- 1 file changed, 35 insertions(+), 6 deletions(-) diff --git a/dstack/kms/src/mpc_reshare.rs b/dstack/kms/src/mpc_reshare.rs index 740aa1a2d..731f05437 100644 --- a/dstack/kms/src/mpc_reshare.rs +++ b/dstack/kms/src/mpc_reshare.rs @@ -11,7 +11,7 @@ use anyhow::{ensure, Context, Result}; use cggmp21::{ generic_ec::{Curve, NonZero, Point, Scalar, SecretScalar}, - key_share::{DirtyIncompleteKeyShare, Valid, VssSetup}, + key_share::{DirtyIncompleteKeyShare, DirtyKeyInfo, Valid, VssSetup}, IncompleteKeyShare, }; use rand::{CryptoRng, RngCore}; @@ -47,7 +47,16 @@ impl ResharedMaterial { usize::from(recipient_index) < self.public_shares.len(), "reshared recipient index out of bounds" ); - let mut key_info = template.key_info.clone(); + 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 @@ -163,6 +172,27 @@ pub(crate) fn verify_and_combine( 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, @@ -176,7 +206,7 @@ pub(crate) fn verify_and_combine( private.len() == old_dealers.len(), "missing private resharing contribution" ); - let setup = old_reference + let setup = old_key_info .vss_setup .as_ref() .context("resharing requires an old Shamir setup")?; @@ -208,8 +238,7 @@ pub(crate) fn verify_and_combine( .map(|encoded| Point::::from_bytes(encoded).context("invalid resharing commitment")) .collect::>>()?; let lambda = lagrange_at_zero(position, &dealer_points)?; - let expected_constant = old_reference - .key_info + let expected_constant = old_key_info .public_shares .get(usize::from(*dealer_index)) .context("missing old dealer public share")? @@ -236,7 +265,7 @@ pub(crate) fn verify_and_combine( } } ensure!( - aggregate_commitments[0] == *old_reference.shared_public_key(), + aggregate_commitments[0] == *old_key_info.shared_public_key, "resharing changed the group public key" ); let public_shares = (0..new_member_count) From adc19f0cb381e2857a17b2151c5ccacabc45c00c Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Thu, 30 Jul 2026 04:57:11 -0700 Subject: [PATCH 096/122] feat(kms): bind resharing dealers in join authorization --- dstack/kms/src/key_backend.rs | 13 ++++++++++- dstack/kms/src/mpc_lifecycle.rs | 38 +++++++++++++++++++++++++++++++++ 2 files changed, 50 insertions(+), 1 deletion(-) diff --git a/dstack/kms/src/key_backend.rs b/dstack/kms/src/key_backend.rs index 67cdf157e..443145f43 100644 --- a/dstack/kms/src/key_backend.rs +++ b/dstack/kms/src/key_backend.rs @@ -645,7 +645,18 @@ impl MpcKeyBackend { async fn coordinate_reshare(&self, plan: ResharePlan) -> Result>> { plan.validate_subset(&self.manifest)?; - let dealers = self.participants().await?; + 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( diff --git a/dstack/kms/src/mpc_lifecycle.rs b/dstack/kms/src/mpc_lifecycle.rs index 750337666..4832aa6cd 100644 --- a/dstack/kms/src/mpc_lifecycle.rs +++ b/dstack/kms/src/mpc_lifecycle.rs @@ -61,6 +61,11 @@ pub(crate) struct ResharePlan { #[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, } @@ -109,6 +114,34 @@ impl ResharePlan { } 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(()) } @@ -533,6 +566,7 @@ mod tests { epoch: 2, previous_manifest_hash: active.manifest_hash().unwrap().to_vec(), threshold: 2, + dealers: vec![], members: active.members[..2] .iter() .map(|member| ReshareMember { @@ -571,6 +605,10 @@ mod tests { 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 From 73d016172d6a5a584dfbf4d5aca31aed95b0af7f Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Thu, 30 Jul 2026 05:01:49 -0700 Subject: [PATCH 097/122] feat(kms): add attested authorized member join protocol --- dstack/kms/rpc/proto/kms_rpc.proto | 8 + dstack/kms/src/config.rs | 4 + dstack/kms/src/main.rs | 50 ++ dstack/kms/src/mpc_join.rs | 1063 ++++++++++++++++++++++++++++ dstack/kms/src/mpc_lifecycle.rs | 32 + 5 files changed, 1157 insertions(+) create mode 100644 dstack/kms/src/mpc_join.rs diff --git a/dstack/kms/rpc/proto/kms_rpc.proto b/dstack/kms/rpc/proto/kms_rpc.proto index b250cfe88..65a221913 100644 --- a/dstack/kms/rpc/proto/kms_rpc.proto +++ b/dstack/kms/rpc/proto/kms_rpc.proto @@ -163,6 +163,14 @@ service MpcGenesis { 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 diff --git a/dstack/kms/src/config.rs b/dstack/kms/src/config.rs index ad66710d1..ec6b39a6f 100644 --- a/dstack/kms/src/config.rs +++ b/dstack/kms/src/config.rs @@ -99,6 +99,10 @@ pub(crate) struct MpcConfig { 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, #[serde(default)] pub p256_group_pubkey: String, #[serde(default)] diff --git a/dstack/kms/src/main.rs b/dstack/kms/src/main.rs index 1c18304dd..eb2eab5a8 100644 --- a/dstack/kms/src/main.rs +++ b/dstack/kms/src/main.rs @@ -27,6 +27,7 @@ mod main_service; mod mpc_driver; mod mpc_genesis; mod mpc_identity; +mod mpc_join; mod mpc_lifecycle; mod mpc_operation; mod mpc_reshare; @@ -132,6 +133,48 @@ async fn run_mpc_genesis_service(kms_config: KmsConfig, figment: Figment) -> Res 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()) @@ -193,6 +236,13 @@ async fn main() -> Result<()> { } } + 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"); diff --git a/dstack/kms/src/mpc_join.rs b/dstack/kms/src/mpc_join.rs new file mode 100644 index 000000000..6406f6da4 --- /dev/null +++ b/dstack/kms/src/mpc_join.rs @@ -0,0 +1,1063 @@ +// 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::{Arc, Mutex}, + time::Duration, +}; + +use anyhow::{ensure, Context, Result}; +use async_trait::async_trait; +use cggmp21::{ + generic_ec::Scalar, + key_share::{AnyKeyShare as _, DirtyKeyInfo}, + 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; +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, 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, +} + +#[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(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) + .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, +} + +#[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" + ); + 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.rpc_cert())?, + fs_err::read_to_string(config.rpc_key())?, + 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(), + }))) + } + 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) { + self.0.finalized.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" + ); + 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, + } + } + + 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 (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()?; + 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 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" + ); + } + 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 encoded = serde_jcs::to_vec(&signed)?; + 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(signed) + } + + 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, remote) = tokio::join!( + self.execute( + 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, remote) = tokio::join!( + self.execute( + 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, signed: SignedEpochManifest) -> Result<()> { + self.validate_target_manifest(&signed.manifest)?; + signed.verify(&self.0.identity)?; + 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.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() +} +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 index 4832aa6cd..698267c57 100644 --- a/dstack/kms/src/mpc_lifecycle.rs +++ b/dstack/kms/src/mpc_lifecycle.rs @@ -297,6 +297,38 @@ pub(crate) fn validate_and_checkpoint( 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. From ee6017512e9789f87e8c06e6cf9377cb461f358a Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Thu, 30 Jul 2026 05:02:21 -0700 Subject: [PATCH 098/122] fix(kms): validate old shares during member join --- dstack/kms/src/mpc_join.rs | 26 +++++++++++++++++--------- 1 file changed, 17 insertions(+), 9 deletions(-) diff --git a/dstack/kms/src/mpc_join.rs b/dstack/kms/src/mpc_join.rs index 6406f6da4..78e37dda1 100644 --- a/dstack/kms/src/mpc_join.rs +++ b/dstack/kms/src/mpc_join.rs @@ -13,7 +13,7 @@ use anyhow::{ensure, Context, Result}; use async_trait::async_trait; use cggmp21::{ generic_ec::Scalar, - key_share::{AnyKeyShare as _, DirtyKeyInfo}, + key_share::{AnyKeyShare as _, DirtyKeyInfo, Valid}, supported_curves::{Secp256k1, Secp256r1}, DataToSign, ExecutionId, KeyShare, }; @@ -501,22 +501,28 @@ impl JoinState { .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, + &p_core, &dealer_indexes, plan.members.len().try_into()?, plan.threshold, &mut OsRng, )?; let (kp, ks) = mpc_reshare::create_contribution( - &k.core, + &k_core, &dealer_indexes, plan.members.len().try_into()?, plan.threshold, &mut OsRng, )?; let (dp, ds) = mpc_reshare::create_contribution( - &d.core, + &d_core, &dealer_indexes, plan.members.len().try_into()?, plan.threshold, @@ -524,13 +530,13 @@ impl JoinState { )?; for (index, member) in plan.members.iter().enumerate() { let wire = JoinWire { - p256_reference: p.core.key_info.clone(), + p256_reference: p_core.key_info.clone(), p256_public: pp.clone(), p256_private: ps[index].clone(), - k256_reference: k.core.key_info.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_reference: d_core.key_info.clone(), derivation_public: dp.clone(), derivation_private: ds[index].clone(), }; @@ -913,9 +919,10 @@ impl JoinState { self.0.transport.start(&member.node_id, &operation).await?, )) }); + let local_operation = operation.clone(); let (local, remote) = tokio::join!( self.execute( - operation, + local_operation, self.0.authorization.plan.dealers.first().unwrap() ), futures::future::join_all(calls) @@ -944,9 +951,10 @@ impl JoinState { self.0.transport.start(node, &operation).await?, )) }); + let local_operation = operation.clone(); let (local, remote) = tokio::join!( self.execute( - operation, + local_operation, self.0.authorization.plan.dealers.first().unwrap() ), futures::future::join_all(calls) From 95b21d62c920d74607d87b19ae2a12702779bdd5 Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Thu, 30 Jul 2026 05:03:41 -0700 Subject: [PATCH 099/122] fix(kms): harden authorized join execution --- dstack/kms/src/mpc_join.rs | 86 ++++++++++++++++++++++++++++++++- dstack/kms/src/mpc_lifecycle.rs | 12 +++++ 2 files changed, 96 insertions(+), 2 deletions(-) diff --git a/dstack/kms/src/mpc_join.rs b/dstack/kms/src/mpc_join.rs index 78e37dda1..381e3ebd4 100644 --- a/dstack/kms/src/mpc_join.rs +++ b/dstack/kms/src/mpc_join.rs @@ -5,7 +5,10 @@ use std::{ collections::BTreeMap, - sync::{Arc, Mutex}, + sync::{ + atomic::{AtomicBool, Ordering}, + Arc, Mutex, + }, time::Duration, }; @@ -249,6 +252,14 @@ struct JoinInner { 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)] @@ -361,6 +372,8 @@ impl JoinState { 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 { @@ -375,7 +388,13 @@ impl JoinState { .is_some_and(|node| node == &self.0.config.mpc.node_id) } pub(crate) async fn wait_finalized(&self) { - self.0.finalized.notified().await; + 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 @@ -399,6 +418,47 @@ impl JoinState { .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() < 16, "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, @@ -820,6 +880,27 @@ impl JoinState { "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(()) } @@ -986,6 +1067,7 @@ impl JoinState { 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(()) } diff --git a/dstack/kms/src/mpc_lifecycle.rs b/dstack/kms/src/mpc_lifecycle.rs index 698267c57..8711bb574 100644 --- a/dstack/kms/src/mpc_lifecycle.rs +++ b/dstack/kms/src/mpc_lifecycle.rs @@ -112,6 +112,18 @@ impl ResharePlan { "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() { From 18869bc73dd91379ac6ce48c413a7615addac5b5 Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Thu, 30 Jul 2026 05:04:05 -0700 Subject: [PATCH 100/122] test(kms): materialize joiner without old secret shares --- dstack/kms/src/mpc_reshare.rs | 42 +++++++++++++++++++++++++++++++++++ 1 file changed, 42 insertions(+) diff --git a/dstack/kms/src/mpc_reshare.rs b/dstack/kms/src/mpc_reshare.rs index 731f05437..9d02dd21c 100644 --- a/dstack/kms/src/mpc_reshare.rs +++ b/dstack/kms/src/mpc_reshare.rs @@ -369,6 +369,48 @@ mod tests { } } + #[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 = [0, 1]; + let contributions = dealers + .iter() + .map(|index| create_contribution(&old[*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) From e1e4f92059f2fd4dd4b4d9beccfc497ff63d729e Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Thu, 30 Jul 2026 05:04:24 -0700 Subject: [PATCH 101/122] fix(kms): type join test dealer indexes canonically --- dstack/kms/src/mpc_reshare.rs | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/dstack/kms/src/mpc_reshare.rs b/dstack/kms/src/mpc_reshare.rs index 9d02dd21c..7c7cc88da 100644 --- a/dstack/kms/src/mpc_reshare.rs +++ b/dstack/kms/src/mpc_reshare.rs @@ -375,10 +375,12 @@ mod tests { .set_threshold(Some(2)) .generate_shares(&mut OsRng) .unwrap(); - let dealers = [0, 1]; + let dealers = [0u16, 1u16]; let contributions = dealers .iter() - .map(|index| create_contribution(&old[*index], &dealers, 5, 3, &mut OsRng).unwrap()) + .map(|index| { + create_contribution(&old[usize::from(*index)], &dealers, 5, 3, &mut OsRng).unwrap() + }) .collect::>(); let public = contributions .iter() From c714586f3fcabb5aec42e0c9c38e7fdd9c64a8f3 Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Thu, 30 Jul 2026 05:06:31 -0700 Subject: [PATCH 102/122] feat(kms): recover crash-interrupted MPC genesis --- dstack/kms/kms.toml | 1 + dstack/kms/src/config.rs | 5 +- dstack/kms/src/main.rs | 4 + dstack/kms/src/mpc_genesis.rs | 161 +++++++++++++++++++++++++++++----- 4 files changed, 148 insertions(+), 23 deletions(-) diff --git a/dstack/kms/kms.toml b/dstack/kms/kms.toml index c7ab1e8cb..6a2efde69 100644 --- a/dstack/kms/kms.toml +++ b/dstack/kms/kms.toml @@ -107,6 +107,7 @@ cluster_id = "" identity_file = "" genesis_plan_file = "" genesis_tls_ca_cert = "" +join_authorization_file = "" p256_group_pubkey = "" k256_group_pubkey = "" derivation_group_pubkey = "" diff --git a/dstack/kms/src/config.rs b/dstack/kms/src/config.rs index ec6b39a6f..c9c18be0f 100644 --- a/dstack/kms/src/config.rs +++ b/dstack/kms/src/config.rs @@ -191,7 +191,10 @@ impl KmsConfig { && self.rpc_key().exists() && self.mpc.p256_share_file.exists() && self.mpc.k256_share_file.exists() - && self.mpc.derivation_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() diff --git a/dstack/kms/src/main.rs b/dstack/kms/src/main.rs index eb2eab5a8..b839aeaa1 100644 --- a/dstack/kms/src/main.rs +++ b/dstack/kms/src/main.rs @@ -236,6 +236,10 @@ async fn main() -> Result<()> { } } + if config.mpc.enabled && mpc_genesis::recover_if_needed(&config)? { + info!("Recovered interrupted MPC genesis commit"); + } + if config.mpc.enabled && !config.mpc.join_authorization_file.as_os_str().is_empty() && config.mpc.join_authorization_file.exists() diff --git a/dstack/kms/src/mpc_genesis.rs b/dstack/kms/src/mpc_genesis.rs index 305a38874..1d47b7fb6 100644 --- a/dstack/kms/src/mpc_genesis.rs +++ b/dstack/kms/src/mpc_genesis.rs @@ -6,6 +6,7 @@ use std::{ collections::BTreeMap, + os::unix::fs::PermissionsExt, sync::{ atomic::{AtomicBool, Ordering}, Arc, Mutex, @@ -41,7 +42,8 @@ use sha2::{Digest, Sha256}; use crate::{ cggmp_engine::{ - execution_id, share_commitment, store_share, CggmpCurve, K256KeyShare, P256KeyShare, + execution_id, load_share, share_commitment, store_share, validate_share_topology, + CggmpCurve, K256KeyShare, P256KeyShare, }, config::KmsConfig, mpc_driver::{ @@ -236,6 +238,12 @@ struct GenesisBundle { signed_manifest: SignedEpochManifest, } +#[derive(Serialize, Deserialize)] +struct GenesisJournal { + version: u16, + bundle: GenesisBundle, +} + struct GenesisTransport { local_node_id: String, local_router: Arc, @@ -1032,29 +1040,18 @@ impl GenesisState { certificate .verify_signature(None) .context("invalid genesis root threshold signature")?; - self.persist_shares()?; - ensure!( - !self.0.config.mpc.identity_file.as_os_str().is_empty(), - "missing MPC identity file" - ); - ensure!( - !self.0.config.mpc.manifest_file.as_os_str().is_empty(), - "missing MPC manifest file" - ); + let journal_path = genesis_journal_path(&self.0.config); safe_write::safe_write( - &self.0.config.mpc.identity_file, - serde_jcs::to_vec(&bundle.identity)?, - )?; - safe_write::safe_write(self.0.config.root_ca_cert(), bundle.root_ca_pem.as_bytes())?; - safe_write::safe_write( - &self.0.config.mpc.manifest_file, - serde_jcs::to_vec(&bundle.signed_manifest)?, - )?; - validate_and_checkpoint( - &self.0.config.mpc.checkpoint_file, - &bundle.signed_manifest, - &bundle.identity, + &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(()) @@ -1152,6 +1149,126 @@ 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<()> { + ensure!( + !config.mpc.identity_file.as_os_str().is_empty(), + "missing MPC identity file" + ); + 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.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)? From df966f266fa255a30daaf390a62e628a7b8cafde Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Thu, 30 Jul 2026 05:10:16 -0700 Subject: [PATCH 103/122] feat(kms): threshold-issue attested genesis RPC certificates --- dstack/kms/src/mpc_genesis.rs | 174 +++++++++++++++++++++++++++++++++- dstack/ra-tls/src/cert.rs | 13 +++ 2 files changed, 184 insertions(+), 3 deletions(-) diff --git a/dstack/kms/src/mpc_genesis.rs b/dstack/kms/src/mpc_genesis.rs index 1d47b7fb6..7723a4752 100644 --- a/dstack/kms/src/mpc_genesis.rs +++ b/dstack/kms/src/mpc_genesis.rs @@ -33,7 +33,10 @@ use ra_rpc::{ }; use ra_tls::attestation::AttestationVerifier; use ra_tls::{ - cert::{prepare_external_self_signed, CertRequest}, + cert::{ + prepare_external_certificate, prepare_external_self_signed, CertRequest, + ExternalCertificate, + }, rcgen::SubjectPublicKeyInfo, }; use rand::{rngs::OsRng, RngCore}; @@ -151,6 +154,8 @@ enum GenesisKind { Commitments, SignRoot, SignManifest, + PrepareRpcCertificate, + SignRpcCertificate, } #[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] @@ -236,6 +241,15 @@ 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)] @@ -566,6 +580,89 @@ impl GenesisState { 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()) + } + + 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:?}"))?; + ensure!( + tbs.get_extension_unique(&attestation_oid)?.is_some(), + "RPC certificate lacks attestation" + ); + Ok(()) + } + async fn execute(&self, operation: GenesisOperation, initiator: &str) -> Result> { operation.validate()?; ensure!( @@ -641,6 +738,9 @@ impl GenesisState { 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()?, @@ -778,7 +878,12 @@ impl GenesisState { .to_vec(), ) } - GenesisKind::SignRoot | GenesisKind::SignManifest => { + 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" @@ -794,6 +899,13 @@ impl GenesisState { 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)?; + Sha256::digest(&payload.tbs_der).into() + } _ => unreachable!(), }; let keygen_indexes = (0..n).collect::>(); @@ -801,7 +913,7 @@ impl GenesisState { let (curve, share) = { let artifacts = self.0.artifacts.lock().expect("genesis mutex poisoned"); match operation.kind { - GenesisKind::SignRoot => ( + GenesisKind::SignRoot | GenesisKind::SignRpcCertificate => ( CggmpCurve::P256, GenesisSigningShare::P256( artifacts @@ -945,6 +1057,30 @@ impl GenesisState { 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( @@ -960,6 +1096,7 @@ impl GenesisState { manifest, signature, }, + rpc_certificates, }; let bundle_json = serde_jcs::to_vec(&bundle)?; let finalizations = self @@ -1154,10 +1291,40 @@ fn genesis_journal_path(config: &KmsConfig) -> std::path::PathBuf { } 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" @@ -1167,6 +1334,7 @@ fn persist_bundle_files(config: &KmsConfig, bundle: &GenesisBundle) -> Result<() &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())?; safe_write::safe_write( &config.mpc.manifest_file, &serde_jcs::to_vec(&bundle.signed_manifest)?, diff --git a/dstack/ra-tls/src/cert.rs b/dstack/ra-tls/src/cert.rs index dac28ddeb..b40e6af3b 100644 --- a/dstack/ra-tls/src/cert.rs +++ b/dstack/ra-tls/src/cert.rs @@ -447,6 +447,19 @@ pub struct ExternalCertificate { } 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 From ac479a70d6586b54a75555e28bdb9be22d9ace30 Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Thu, 30 Jul 2026 05:12:25 -0700 Subject: [PATCH 104/122] feat(kms): threshold-issue joined member RPC certificates --- dstack/kms/src/main.rs | 3 + dstack/kms/src/mpc_join.rs | 268 ++++++++++++++++++++++++++++++++++++- 2 files changed, 266 insertions(+), 5 deletions(-) diff --git a/dstack/kms/src/main.rs b/dstack/kms/src/main.rs index b839aeaa1..8a3025782 100644 --- a/dstack/kms/src/main.rs +++ b/dstack/kms/src/main.rs @@ -239,6 +239,9 @@ async fn main() -> Result<()> { 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() diff --git a/dstack/kms/src/mpc_join.rs b/dstack/kms/src/mpc_join.rs index 381e3ebd4..5ea66c6ea 100644 --- a/dstack/kms/src/mpc_join.rs +++ b/dstack/kms/src/mpc_join.rs @@ -29,7 +29,10 @@ use ra_rpc::{ client::{RaClient, RaClientConfig}, CallContext, RpcCall, }; -use ra_tls::attestation::AttestationVerifier; +use ra_tls::{ + attestation::AttestationVerifier, + cert::{prepare_external_certificate, CertRequest, ExternalCertificate}, +}; use rand::{rngs::OsRng, RngCore}; use serde::{Deserialize, Serialize}; use sha2::{Digest, Sha256}; @@ -61,6 +64,8 @@ enum JoinKind { AuxiliaryInfo, Reshare, SignManifest, + PrepareRpcCertificate, + SignRpcCertificate, } #[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] @@ -135,6 +140,19 @@ struct JoinWire { 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, @@ -170,6 +188,8 @@ impl JoinTransport { .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")?; @@ -463,6 +483,8 @@ impl JoinState { 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, } } @@ -858,6 +880,150 @@ impl JoinState { .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()) + } + + 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:?}"))?; + ensure!( + tbs.get_extension_unique(&attestation_oid)?.is_some(), + "join RPC certificate lacks attestation" + ); + 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)?; + 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!( @@ -959,7 +1125,35 @@ impl JoinState { signature: first.to_vec(), }; signed.verify(&self.0.identity)?; - let encoded = serde_jcs::to_vec(&signed)?; + 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() @@ -979,7 +1173,7 @@ impl JoinState { for result in futures::future::join_all(calls).await { result?; } - self.finalize(signed) + self.finalize(bundle) } async fn run_targets( @@ -1047,9 +1241,34 @@ impl JoinState { } Ok(out) } - fn finalize(&self, signed: SignedEpochManifest) -> Result<()> { + 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, @@ -1060,7 +1279,7 @@ impl JoinState { }, &self.0.identity, &self.0.active.manifest, - &signed, + signed, &self.0.identity.cluster_id, &self.0.config.mpc.node_id, )?; @@ -1125,6 +1344,45 @@ impl MpcJoinRpc for JoinHandler { 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() From e0a2a5be44421c7f87e36e9f2f9132e20f11ade1 Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Thu, 30 Jul 2026 05:13:12 -0700 Subject: [PATCH 105/122] feat(kms): publish quote-bound MPC node evidence --- dstack/kms/rpc/proto/kms_rpc.proto | 3 +++ dstack/kms/src/main_service.rs | 23 ++++++++++++++++++++++- 2 files changed, 25 insertions(+), 1 deletion(-) diff --git a/dstack/kms/rpc/proto/kms_rpc.proto b/dstack/kms/rpc/proto/kms_rpc.proto index 65a221913..40171d97f 100644 --- a/dstack/kms/rpc/proto/kms_rpc.proto +++ b/dstack/kms/rpc/proto/kms_rpc.proto @@ -65,6 +65,9 @@ message GetMetaResponse { 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 { diff --git a/dstack/kms/src/main_service.rs b/dstack/kms/src/main_service.rs index 26b423420..f1cde63bb 100644 --- a/dstack/kms/src/main_service.rs +++ b/dstack/kms/src/main_service.rs @@ -32,7 +32,7 @@ use upgrade_authority::{build_boot_info, ensure_app_id_len, local_kms_boot_info, use crate::{ config::KmsConfig, key_backend::{KeyBackend, LocalKeyBackend, MpcKeyBackend}, - mpc_identity::{decode_hex, ClusterIdentity, EpochManifest, SignedEpochManifest}, + mpc_identity::{decode_hex, ClusterIdentity, EpochManifest, NodeEvidence, SignedEpochManifest}, mpc_lifecycle, mpc_session::SessionRouter, }; @@ -644,6 +644,27 @@ impl KmsRpc for RpcHandler { }) .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, + }, }) } From 6edec2673dc9cccc24f943ce2d04e06ed2a91672 Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Thu, 30 Jul 2026 05:14:38 -0700 Subject: [PATCH 106/122] fix(kms): verify attestation before threshold RPC issuance --- dstack/kms/src/mpc_genesis.rs | 23 +++++++++++++++++------ dstack/kms/src/mpc_join.rs | 23 +++++++++++++++++------ 2 files changed, 34 insertions(+), 12 deletions(-) diff --git a/dstack/kms/src/mpc_genesis.rs b/dstack/kms/src/mpc_genesis.rs index 7723a4752..fd6dcf56c 100644 --- a/dstack/kms/src/mpc_genesis.rs +++ b/dstack/kms/src/mpc_genesis.rs @@ -606,7 +606,10 @@ impl GenesisState { Ok(external.tbs_der().to_vec()) } - fn validate_rpc_certificate_tbs(&self, payload: &RpcCertificateSigningPayload) -> Result<()> { + 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 _, @@ -656,10 +659,18 @@ impl GenesisState { ); let attestation_oid = Oid::from(PHALA_RATLS_ATTESTATION) .map_err(|error| anyhow::anyhow!("invalid attestation OID: {error:?}"))?; - ensure!( - tbs.get_extension_unique(&attestation_oid)?.is_some(), - "RPC certificate lacks attestation" - ); + 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(()) } @@ -903,7 +914,7 @@ impl GenesisState { let payload: RpcCertificateSigningPayload = serde_json::from_slice(&operation.payload) .context("invalid RPC certificate signing payload")?; - self.validate_rpc_certificate_tbs(&payload)?; + self.validate_rpc_certificate_tbs(&payload).await?; Sha256::digest(&payload.tbs_der).into() } _ => unreachable!(), diff --git a/dstack/kms/src/mpc_join.rs b/dstack/kms/src/mpc_join.rs index 5ea66c6ea..c1e36d964 100644 --- a/dstack/kms/src/mpc_join.rs +++ b/dstack/kms/src/mpc_join.rs @@ -910,7 +910,10 @@ impl JoinState { .to_vec()) } - fn validate_rpc_certificate_tbs(&self, payload: &JoinRpcCertificatePayload) -> Result<()> { + 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 _, @@ -951,10 +954,18 @@ impl JoinState { ); let attestation_oid = Oid::from(PHALA_RATLS_ATTESTATION) .map_err(|error| anyhow::anyhow!("invalid attestation OID: {error:?}"))?; - ensure!( - tbs.get_extension_unique(&attestation_oid)?.is_some(), - "join RPC certificate lacks attestation" - ); + 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(()) } @@ -969,7 +980,7 @@ impl JoinState { "only old dealers sign joined RPC certificates" ); let payload: JoinRpcCertificatePayload = serde_json::from_slice(&operation.payload)?; - self.validate_rpc_certificate_tbs(&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 From 64089ad3392d6188dce761f2a1d66956e1a614c7 Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Thu, 30 Jul 2026 05:15:07 -0700 Subject: [PATCH 107/122] test(ra-tls): roundtrip transported external certificate TBS --- dstack/ra-tls/src/cert.rs | 28 ++++++++++++++++++++++++++++ 1 file changed, 28 insertions(+) diff --git a/dstack/ra-tls/src/cert.rs b/dstack/ra-tls/src/cert.rs index b40e6af3b..e1c47950a 100644 --- a/dstack/ra-tls/src/cert.rs +++ b/dstack/ra-tls/src/cert.rs @@ -774,6 +774,34 @@ mod tests { 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(); From f13f568afe6457810de7d80067a898138acdf2f5 Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Thu, 30 Jul 2026 05:15:49 -0700 Subject: [PATCH 108/122] test(kms): sign three-of-five with two MPC nodes offline --- dstack/kms/src/mpc_driver.rs | 35 +++++++++++++++++++++++++++++------ 1 file changed, 29 insertions(+), 6 deletions(-) diff --git a/dstack/kms/src/mpc_driver.rs b/dstack/kms/src/mpc_driver.rs index 011f219d8..863f89997 100644 --- a/dstack/kms/src/mpc_driver.rs +++ b/dstack/kms/src/mpc_driver.rs @@ -444,6 +444,23 @@ mod tests { } } + 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, @@ -524,11 +541,11 @@ mod tests { use cggmp21::{generic_ec::Scalar, security_level::SecurityLevel128, DataToSign}; use k256::ecdsa::{signature::hazmat::PrehashVerifier, Signature, VerifyingKey}; - let shares = cggmp21::trusted_dealer::builder::(3) - .set_threshold(Some(2)) + let shares = cggmp21::trusted_dealer::builder::(5) + .set_threshold(Some(3)) .generate_shares(&mut OsRng) .unwrap(); - let manifest = manifest(); + let manifest = five_member_manifest(); let routers: Arc>> = Arc::new( manifest .members @@ -544,7 +561,12 @@ mod tests { }) .collect(), ); - let participants = vec!["kms-1".to_string(), "kms-3".to_string()]; + // 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() @@ -575,7 +597,7 @@ mod tests { let mut rng = OsRng; let data = DataToSign::from_scalar(Scalar::::from_be_bytes_mod_order(digest)); - let keygen_indexes = [0, 2]; + let keygen_indexes = [0, 2, 4]; let state = cggmp21::signing( ExecutionId::new(&eid), protocol_index, @@ -589,8 +611,9 @@ mod tests { .unwrap() } }; - let (first, second) = tokio::join!(sign(0, 0), sign(1, 2)); + 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(); From 6b9255ade38aed1f5984d9f9a0517ded6e0b31f5 Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Thu, 30 Jul 2026 05:20:55 -0700 Subject: [PATCH 109/122] test(kms): verify three-of-five threshold root certificate --- dstack/kms/src/mpc_driver.rs | 76 ++++++++++++++++++++++++++++++++++++ 1 file changed, 76 insertions(+) diff --git a/dstack/kms/src/mpc_driver.rs b/dstack/kms/src/mpc_driver.rs index 863f89997..a555ed313 100644 --- a/dstack/kms/src/mpc_driver.rs +++ b/dstack/kms/src/mpc_driver.rs @@ -540,6 +540,7 @@ mod tests { 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)) @@ -621,5 +622,80 @@ mod tests { 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)| cggmp21::KeyShare::from_parts((core, share.aux.clone())).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_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: expires_at + 60, + 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(); } } From 67e6cc70c222372c273f4a51e4e412c6011af448 Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Thu, 30 Jul 2026 05:21:42 -0700 Subject: [PATCH 110/122] fix(kms): validate reused P-256 test auxiliary info --- dstack/kms/src/mpc_driver.rs | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/dstack/kms/src/mpc_driver.rs b/dstack/kms/src/mpc_driver.rs index a555ed313..0c4329ec5 100644 --- a/dstack/kms/src/mpc_driver.rs +++ b/dstack/kms/src/mpc_driver.rs @@ -633,7 +633,10 @@ mod tests { let p256_shares = p256_cores .into_iter() .zip(shares.iter()) - .map(|(core, share)| cggmp21::KeyShare::from_parts((core, share.aux.clone())).unwrap()) + .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] = &[ From 65c3a6e4326cd330174e6442754521d53f6d5085 Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Thu, 30 Jul 2026 05:22:09 -0700 Subject: [PATCH 111/122] test(kms): reject wrong epoch sender and recipient --- dstack/kms/src/mpc_session.rs | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/dstack/kms/src/mpc_session.rs b/dstack/kms/src/mpc_session.rs index f8a63b41b..dd621bd2f 100644 --- a/dstack/kms/src/mpc_session.rs +++ b/dstack/kms/src/mpc_session.rs @@ -284,5 +284,15 @@ mod tests { 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()); } } From faf782f62c1c92cd95e6ea61814a7677560b9205 Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Thu, 30 Jul 2026 05:22:24 -0700 Subject: [PATCH 112/122] fix(kms): avoid shadowing envelope test helper --- dstack/kms/src/mpc_session.rs | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/dstack/kms/src/mpc_session.rs b/dstack/kms/src/mpc_session.rs index dd621bd2f..0ae5a7a5a 100644 --- a/dstack/kms/src/mpc_session.rs +++ b/dstack/kms/src/mpc_session.rs @@ -277,10 +277,10 @@ mod tests { #[test] fn rejects_replay_and_session_context_changes() { let router = router(); - let message = message(); - router.push("kms-1", message.clone()).unwrap(); - assert!(router.push("kms-1", message.clone()).is_err()); - let mut changed = message; + 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()); From a889eb93fb4623618415f82fda97dcf0f1369a31 Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Thu, 30 Jul 2026 05:23:46 -0700 Subject: [PATCH 113/122] docs(kms): add MPC deployment and lifecycle runbook --- dstack/kms/README.md | 5 ++ dstack/kms/docs/mpc.md | 175 +++++++++++++++++++++++++++++++++++++++++ 2 files changed, 180 insertions(+) create mode 100644 dstack/kms/docs/mpc.md 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..1379781b5 --- /dev/null +++ b/dstack/kms/docs/mpc.md @@ -0,0 +1,175 @@ +# 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" +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. + +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. From dd3ef78efb16642890e8c5ffaf3fa2907803ec40 Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Thu, 30 Jul 2026 05:24:04 -0700 Subject: [PATCH 114/122] fix(kms): preserve threshold-issued RPC certificates --- dstack/kms/src/main.rs | 24 ++++++++++++++---------- 1 file changed, 14 insertions(+), 10 deletions(-) diff --git a/dstack/kms/src/main.rs b/dstack/kms/src/main.rs index 8a3025782..39e36ab6c 100644 --- a/dstack/kms/src/main.rs +++ b/dstack/kms/src/main.rs @@ -257,16 +257,20 @@ async fn main() -> Result<()> { return run_mpc_genesis_service(config, figment).await; } - 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}"); - }; + 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:"); From 55e98da35a4fb04cf3dda3b24ba45102839e36a8 Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Thu, 30 Jul 2026 05:24:42 -0700 Subject: [PATCH 115/122] fix(kms): forbid full-root onboarding in MPC mode --- dstack/kms/src/main.rs | 12 +++++------- dstack/kms/src/mpc_genesis.rs | 7 +++++++ 2 files changed, 12 insertions(+), 7 deletions(-) diff --git a/dstack/kms/src/main.rs b/dstack/kms/src/main.rs index 39e36ab6c..765c6bd26 100644 --- a/dstack/kms/src/main.rs +++ b/dstack/kms/src/main.rs @@ -223,15 +223,13 @@ async fn main() -> Result<()> { ); } - if config.onboard.enabled - && !(config.mpc.enabled && config.genesis_transport_exists()) - && !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.mpc.enabled && !config.genesis_transport_exists() - || !config.mpc.enabled && !config.keys_exists() - { + if !config.keys_exists() { bail!("Failed to onboard"); } } diff --git a/dstack/kms/src/mpc_genesis.rs b/dstack/kms/src/mpc_genesis.rs index fd6dcf56c..0897b33bd 100644 --- a/dstack/kms/src/mpc_genesis.rs +++ b/dstack/kms/src/mpc_genesis.rs @@ -1346,6 +1346,13 @@ fn persist_bundle_files(config: &KmsConfig, bundle: &GenesisBundle) -> Result<() )?; 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)?, From 02ec2761920fa908b49a05bfd43f7775b8b0c685 Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Thu, 30 Jul 2026 05:25:25 -0700 Subject: [PATCH 116/122] fix(kms): bound active operations without throttling replay cache --- dstack/kms/src/key_backend.rs | 11 ++++++++++- dstack/kms/src/mpc_genesis.rs | 2 +- dstack/kms/src/mpc_join.rs | 2 +- 3 files changed, 12 insertions(+), 3 deletions(-) diff --git a/dstack/kms/src/key_backend.rs b/dstack/kms/src/key_backend.rs index 443145f43..b27cea68b 100644 --- a/dstack/kms/src/key_backend.rs +++ b/dstack/kms/src/key_backend.rs @@ -835,9 +835,18 @@ impl MpcKeyBackend { (record.clone(), false) } else { ensure!( - operations.len() < MAX_ACTIVE_OPERATIONS, + 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(), diff --git a/dstack/kms/src/mpc_genesis.rs b/dstack/kms/src/mpc_genesis.rs index 0897b33bd..fe8c7cbe5 100644 --- a/dstack/kms/src/mpc_genesis.rs +++ b/dstack/kms/src/mpc_genesis.rs @@ -689,7 +689,7 @@ impl GenesisState { ); (record.clone(), false) } else { - ensure!(operations.len() < 32, "too many active genesis operations"); + ensure!(operations.len() < 256, "too many genesis operations"); let record = Arc::new(GenesisOperationRecord { operation: operation.clone(), result: Mutex::new(None), diff --git a/dstack/kms/src/mpc_join.rs b/dstack/kms/src/mpc_join.rs index c1e36d964..1778a046f 100644 --- a/dstack/kms/src/mpc_join.rs +++ b/dstack/kms/src/mpc_join.rs @@ -447,7 +447,7 @@ impl JoinState { ); (record.clone(), false) } else { - ensure!(records.len() < 16, "too many join operations"); + ensure!(records.len() < 256, "too many join operations"); let record = Arc::new(JoinOperationRecord { operation: operation.clone(), result: Mutex::new(None), From 8651673086e4a149e0bb3bd9bfe81b547e69e308 Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Thu, 30 Jul 2026 05:28:23 -0700 Subject: [PATCH 117/122] fix(kms): refresh P-256 threshold test expiry --- dstack/kms/src/mpc_driver.rs | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/dstack/kms/src/mpc_driver.rs b/dstack/kms/src/mpc_driver.rs index 0c4329ec5..3558f65f4 100644 --- a/dstack/kms/src/mpc_driver.rs +++ b/dstack/kms/src/mpc_driver.rs @@ -656,6 +656,11 @@ mod tests { .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(); @@ -666,7 +671,7 @@ mod tests { request_hash: [42; 32], local_node_id: local.clone(), participants: participants.clone(), - expires_at: expires_at + 60, + expires_at: p256_expires_at, poll_interval: Duration::from_millis(1), }; let transport = MemoryTransport { From 799e4381dd68d7bc0f50f577e35dde81ef071a94 Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Thu, 30 Jul 2026 05:34:17 -0700 Subject: [PATCH 118/122] fix(kms): separate MPC mutual TLS client identity --- dstack/kms/docs/mpc.md | 8 ++++++++ dstack/kms/kms.toml | 2 ++ dstack/kms/src/config.rs | 10 ++++++++++ dstack/kms/src/main_service.rs | 20 ++++++++++++++++++-- dstack/kms/src/mpc_genesis.rs | 16 ++++++++++++++-- dstack/kms/src/mpc_join.rs | 17 +++++++++++++++-- 6 files changed, 67 insertions(+), 6 deletions(-) diff --git a/dstack/kms/docs/mpc.md b/dstack/kms/docs/mpc.md index 1379781b5..198001d92 100644 --- a/dstack/kms/docs/mpc.md +++ b/dstack/kms/docs/mpc.md @@ -48,6 +48,8 @@ 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" @@ -82,6 +84,12 @@ 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. diff --git a/dstack/kms/kms.toml b/dstack/kms/kms.toml index 6a2efde69..36cb79c0a 100644 --- a/dstack/kms/kms.toml +++ b/dstack/kms/kms.toml @@ -108,6 +108,8 @@ 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 = "" diff --git a/dstack/kms/src/config.rs b/dstack/kms/src/config.rs index c9c18be0f..5202ff5ef 100644 --- a/dstack/kms/src/config.rs +++ b/dstack/kms/src/config.rs @@ -103,6 +103,12 @@ pub(crate) struct MpcConfig { /// 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)] @@ -180,6 +186,8 @@ impl KmsConfig { && 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 { @@ -189,6 +197,8 @@ impl KmsConfig { && 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() diff --git a/dstack/kms/src/main_service.rs b/dstack/kms/src/main_service.rs index f1cde63bb..51ef95d29 100644 --- a/dstack/kms/src/main_service.rs +++ b/dstack/kms/src/main_service.rs @@ -226,6 +226,20 @@ impl KmsState { .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( @@ -242,8 +256,10 @@ impl KmsState { &config.mpc.checkpoint_file, manifest, mpc_router.clone().context("MPC router is missing")?, - fs::read_to_string(config.rpc_cert()).context("failed to read MPC RPC cert")?, - fs::read_to_string(config.rpc_key()).context("failed to read MPC RPC key")?, + 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 { diff --git a/dstack/kms/src/mpc_genesis.rs b/dstack/kms/src/mpc_genesis.rs index fe8c7cbe5..806c6be7b 100644 --- a/dstack/kms/src/mpc_genesis.rs +++ b/dstack/kms/src/mpc_genesis.rs @@ -415,6 +415,18 @@ impl GenesisState { .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)?); @@ -424,8 +436,8 @@ impl GenesisState { &plan, &config.mpc.node_id, router.clone(), - fs_err::read_to_string(config.rpc_cert())?, - fs_err::read_to_string(config.rpc_key())?, + fs_err::read_to_string(&config.mpc.client_cert_file)?, + fs_err::read_to_string(&config.mpc.client_key_file)?, ca, verifier.clone(), )?); diff --git a/dstack/kms/src/mpc_join.rs b/dstack/kms/src/mpc_join.rs index 1778a046f..5b9e931dc 100644 --- a/dstack/kms/src/mpc_join.rs +++ b/dstack/kms/src/mpc_join.rs @@ -303,6 +303,19 @@ impl JoinState { .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, @@ -333,8 +346,8 @@ impl JoinState { &authorization, &config.mpc.node_id, router.clone(), - fs_err::read_to_string(config.rpc_cert())?, - fs_err::read_to_string(config.rpc_key())?, + fs_err::read_to_string(&config.mpc.client_cert_file)?, + fs_err::read_to_string(&config.mpc.client_key_file)?, ca, verifier.clone(), )?); From 4a62d24b5952db72e9a785404c9a6819ce56707c Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Thu, 30 Jul 2026 05:36:10 -0700 Subject: [PATCH 119/122] fix(kms): durably reject epoch authorization forks --- dstack/kms/src/key_backend.rs | 56 +++++++++++++++++++++++++++++++++-- 1 file changed, 54 insertions(+), 2 deletions(-) diff --git a/dstack/kms/src/key_backend.rs b/dstack/kms/src/key_backend.rs index b27cea68b..e21ebadb9 100644 --- a/dstack/kms/src/key_backend.rs +++ b/dstack/kms/src/key_backend.rs @@ -10,6 +10,8 @@ use std::{ collections::BTreeMap, + io::Write, + os::unix::fs::PermissionsExt, sync::{Arc, Mutex}, time::Duration, }; @@ -1199,8 +1201,16 @@ impl MpcKeyBackend { async fn execute_k256(&self, operation: MpcOperation) -> Result> { let digest = match &operation.payload { MpcOperationPayload::SignK256(payload) => payload.digest(), - MpcOperationPayload::SignManifest(payload) => payload.digest()?, - MpcOperationPayload::AuthorizeReshare(plan) => plan.authorization_hash()?, + 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 @@ -1281,6 +1291,48 @@ impl MpcKeyBackend { .context("MPC signing worker panicked")? } + fn authorize_policy_once(&self, label: &str, epoch: u64, digest: &[u8; 32]) -> Result<()> { + let path = self + .manifest_file + .with_extension(format!("epoch-{epoch}.{label}.authorized")); + let validate_existing = || -> Result<()> { + let metadata = fs_err::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_err::read(&path)? == digest, + "conflicting threshold authorization for epoch {epoch}" + ); + Ok(()) + }; + if path.exists() { + return validate_existing(); + } + if let Some(parent) = path.parent() { + fs_err::create_dir_all(parent)?; + } + match std::fs::OpenOptions::new() + .write(true) + .create_new(true) + .open(&path) + { + Ok(mut file) => { + file.set_permissions(std::fs::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()), + } + } + async fn execute_p256(&self, operation: MpcOperation) -> Result> { let MpcOperationPayload::SignP256Certificate(payload) = &operation.payload else { bail!("P-256 executor received a different MPC operation") From be34040eb0e48a734b4fd6b2e9d43a6fa9d9ee5d Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Thu, 30 Jul 2026 05:36:29 -0700 Subject: [PATCH 120/122] fix(kms): serialize pending reshare generations --- dstack/kms/src/key_backend.rs | 1 + 1 file changed, 1 insertion(+) diff --git a/dstack/kms/src/key_backend.rs b/dstack/kms/src/key_backend.rs index e21ebadb9..6b7efa4c1 100644 --- a/dstack/kms/src/key_backend.rs +++ b/dstack/kms/src/key_backend.rs @@ -946,6 +946,7 @@ impl MpcKeyBackend { operation: &MpcOperation, plan: &ResharePlan, ) -> Result> { + self.authorize_policy_once("reshare", plan.epoch, &plan.authorization_hash()?)?; let old_dealers = operation .participants .iter() From 0b56fcb3a2cd8253a599b49f23557715a95e2801 Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Thu, 30 Jul 2026 05:41:48 -0700 Subject: [PATCH 121/122] fix(kms): prevent maintenance join manifest forks --- dstack/kms/src/key_backend.rs | 47 ++++---------------------------- dstack/kms/src/mpc_join.rs | 10 +++++-- dstack/kms/src/mpc_lifecycle.rs | 48 +++++++++++++++++++++++++++++++++ 3 files changed, 61 insertions(+), 44 deletions(-) diff --git a/dstack/kms/src/key_backend.rs b/dstack/kms/src/key_backend.rs index 6b7efa4c1..27739125d 100644 --- a/dstack/kms/src/key_backend.rs +++ b/dstack/kms/src/key_backend.rs @@ -10,8 +10,6 @@ use std::{ collections::BTreeMap, - io::Write, - os::unix::fs::PermissionsExt, sync::{Arc, Mutex}, time::Duration, }; @@ -43,7 +41,10 @@ use crate::{ MpcHttpTransport, }, mpc_identity::{ClusterIdentity, EpochManifest, SignedEpochManifest}, - mpc_lifecycle::{activate_pending_epoch, EpochPaths, ResharePlan, SignedResharePlan}, + mpc_lifecycle::{ + activate_pending_epoch, record_epoch_authorization, EpochPaths, ResharePlan, + SignedResharePlan, + }, mpc_operation::{ DerivePayload, DerivePurpose, K256SignPayload, ManifestSignaturePayload, MpcOperation, MpcOperationPayload, P256CertificatePayload, @@ -1293,45 +1294,7 @@ impl MpcKeyBackend { } fn authorize_policy_once(&self, label: &str, epoch: u64, digest: &[u8; 32]) -> Result<()> { - let path = self - .manifest_file - .with_extension(format!("epoch-{epoch}.{label}.authorized")); - let validate_existing = || -> Result<()> { - let metadata = fs_err::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_err::read(&path)? == digest, - "conflicting threshold authorization for epoch {epoch}" - ); - Ok(()) - }; - if path.exists() { - return validate_existing(); - } - if let Some(parent) = path.parent() { - fs_err::create_dir_all(parent)?; - } - match std::fs::OpenOptions::new() - .write(true) - .create_new(true) - .open(&path) - { - Ok(mut file) => { - file.set_permissions(std::fs::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()), - } + record_epoch_authorization(&self.manifest_file, label, epoch, digest) } async fn execute_p256(&self, operation: MpcOperation) -> Result> { diff --git a/dstack/kms/src/mpc_join.rs b/dstack/kms/src/mpc_join.rs index 5b9e931dc..9d0cecc80 100644 --- a/dstack/kms/src/mpc_join.rs +++ b/dstack/kms/src/mpc_join.rs @@ -48,8 +48,8 @@ use crate::{ }, mpc_identity::{ClusterIdentity, EpochManifest, EpochMember, SignedEpochManifest}, mpc_lifecycle::{ - activate_pending_epoch, initialize_join_checkpoint, pending_share_path, EpochPaths, - SignedResharePlan, + activate_pending_epoch, initialize_join_checkpoint, pending_share_path, + record_epoch_authorization, EpochPaths, SignedResharePlan, }, mpc_reshare::{self, PrivateContribution, PublicContribution}, mpc_session::{MpcEnvelope, MpcProtocol, SessionRouter}, @@ -840,6 +840,12 @@ impl JoinState { 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() diff --git a/dstack/kms/src/mpc_lifecycle.rs b/dstack/kms/src/mpc_lifecycle.rs index 8711bb574..7b63210e7 100644 --- a/dstack/kms/src/mpc_lifecycle.rs +++ b/dstack/kms/src/mpc_lifecycle.rs @@ -6,6 +6,7 @@ use std::{ fs::Permissions, + io::Write, os::unix::fs::PermissionsExt, path::{Path, PathBuf}, }; @@ -47,6 +48,53 @@ 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, From fabf32dc9272b1a2b3e04f44d6b79c81a74dc6e2 Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Thu, 30 Jul 2026 05:42:30 -0700 Subject: [PATCH 122/122] test(kms): cover durable epoch authorization locks --- dstack/kms/src/mpc_lifecycle.rs | 55 +++++++++++++++++++++++++++++++++ 1 file changed, 55 insertions(+) diff --git a/dstack/kms/src/mpc_lifecycle.rs b/dstack/kms/src/mpc_lifecycle.rs index 7b63210e7..018ca14c3 100644 --- a/dstack/kms/src/mpc_lifecycle.rs +++ b/dstack/kms/src/mpc_lifecycle.rs @@ -583,6 +583,7 @@ 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( @@ -625,6 +626,60 @@ mod tests { } } + #[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();