diff --git a/dstack/cert-client/src/lib.rs b/dstack/cert-client/src/lib.rs index 65498f2751..9e222bb049 100644 --- a/dstack/cert-client/src/lib.rs +++ b/dstack/cert-client/src/lib.rs @@ -168,6 +168,28 @@ impl CertRequestClient { } } + /// The app keys of an older KMS root epoch, from the KMS that served the + /// current ones. + pub async fn get_app_key(&self, root_epoch: u32) -> Result { + let CertRequestClient::Kms { client, vm_config } = self else { + bail!("only KMS-provided keys have older root epochs"); + }; + let response = client + .get_app_key(GetAppKeyRequest { + api_version: 1, + vm_config: vm_config.clone(), + root_epoch: Some(root_epoch), + }) + .await?; + if response.root_epoch != root_epoch { + bail!( + "KMS returned keys of root epoch {} instead of {root_epoch}", + response.root_epoch + ); + } + Ok(response) + } + pub async fn get_root_ca(&self) -> Result { match self { CertRequestClient::Local { ca } => Ok(ca.pem_cert.clone()), diff --git a/dstack/guest-agent/rpc/proto/agent_rpc_v1.proto b/dstack/guest-agent/rpc/proto/agent_rpc_v1.proto index 3df408ad8e..bebc85f99b 100644 --- a/dstack/guest-agent/rpc/proto/agent_rpc_v1.proto +++ b/dstack/guest-agent/rpc/proto/agent_rpc_v1.proto @@ -189,6 +189,11 @@ message GetKeyRequest { // signing method: v0's `secp256k1_prehashed` was a mode wearing an // algorithm's name, and it does not appear here. string algorithm = 2; + // KMS root key epoch to derive from. Unset means the current epoch, which is + // what new keys should use; epoch 0 is the original root keys. Name an older + // epoch only to reach a key from before a KMS root key rotation, for example + // to migrate data or assets off it. + optional uint32 root_epoch = 3; } message GetKeyResponse { @@ -209,6 +214,9 @@ message GetKeyResponse { // unversioned surface returns. Link [0] is not: its claim encoding is new, // and it is deliberately not one the v0 claim format can produce. repeated bytes signature_chain = 3; + // KMS root key epoch the key was derived from. Link [1] is signed by that + // epoch's KMS root key. + uint32 root_epoch = 4; } // --------------------------------------------------------------------------- diff --git a/dstack/guest-agent/src/rpc_service.rs b/dstack/guest-agent/src/rpc_service.rs index a6dcd8e5fb..d79a9ed22d 100644 --- a/dstack/guest-agent/src/rpc_service.rs +++ b/dstack/guest-agent/src/rpc_service.rs @@ -2,6 +2,7 @@ // // SPDX-License-Identifier: Apache-2.0 +use std::collections::HashMap; use std::sync::{Arc, Mutex}; use std::time::{Duration, Instant}; @@ -92,11 +93,35 @@ pub struct AppState { inner: Arc, } +/// The app root key of one KMS root epoch, with its KMS signature. +#[derive(Clone)] +pub(crate) struct AppRoot { + pub epoch: u32, + pub k256_key: Vec, + pub signing_key: SigningKey, + pub kms_signature: Vec, +} + +impl AppRoot { + fn new(epoch: u32, k256_key: Vec, kms_signature: Vec) -> Result { + let signing_key = SigningKey::from_slice(&k256_key) + .context("the app root k256 key is not a valid secp256k1 scalar")?; + Ok(Self { + epoch, + k256_key, + signing_key, + kms_signature, + }) + } +} + struct AppStateInner { config: Config, keys: AppKeys, vm_config: String, cert_client: CertRequestClient, + /// Held across a fetch so concurrent requests for one epoch fetch it once. + historical_roots: tokio::sync::Mutex>, demo_cert: tokio::sync::OnceCell, /// When the last demo-cert request failed; held while one is running. demo_cert_last_failure: tokio::sync::Mutex>, @@ -300,6 +325,7 @@ impl AppState { config, keys, cert_client, + historical_roots: Default::default(), demo_cert: Default::default(), demo_cert_last_failure: Default::default(), vm_config, @@ -377,6 +403,64 @@ impl AppState { Ok(key) } + /// KMS root epoch of the app keys this agent serves. + pub(crate) fn root_epoch(&self) -> u32 { + self.inner.keys.root_epoch() + } + + /// The app root key and its KMS signature of `root_epoch`; the current + /// ones when unset. An older epoch is requested from the KMS once and kept + /// in memory. + pub(crate) async fn app_root(&self, root_epoch: Option) -> Result { + let current = self.root_epoch(); + let epoch = root_epoch.unwrap_or(current); + if epoch == current { + return Ok(AppRoot { + epoch, + k256_key: self.app_root_k256_key().to_vec(), + signing_key: self.app_root_signing_key()?.clone(), + kms_signature: self.kms_k256_signature().to_vec(), + }); + } + if epoch > current { + anyhow::bail!("root epoch {epoch} is newer than the current one, {current}"); + } + self.historical_root(epoch).await + } + + pub(crate) async fn historical_root(&self, epoch: u32) -> Result { + if let Some(root) = self + .inner + .historical_roots + .lock() + .await + .get(&epoch) + .cloned() + { + return Ok(root); + } + let keys = self + .inner + .cert_client + .get_app_key(epoch) + .await + .with_context(|| format!("failed to fetch the app keys of root epoch {epoch}"))?; + let root = AppRoot::new(epoch, keys.k256_key, keys.k256_signature)?; + let mut roots = self.inner.historical_roots.lock().await; + Ok(roots.entry(epoch).or_insert_with(|| root.clone()).clone()) + } + + #[cfg(test)] + pub(crate) async fn insert_historical_root( + &self, + epoch: u32, + k256_key: Vec, + kms_signature: Vec, + ) { + let root = AppRoot::new(epoch, k256_key, kms_signature).unwrap(); + self.inner.historical_roots.lock().await.insert(epoch, root); + } + /// The KMS root key's signature over the app root public key: the second /// link of every signature chain, produced outside this agent and passed /// through byte-for-byte on both API surfaces. @@ -1152,6 +1236,8 @@ pub(crate) mod tests { .expect("certificate carries an attestation") } + pub(crate) const HISTORICAL_K256_KEY: [u8; 32] = [0x11; 32]; + pub(crate) async fn setup_test_state() -> (AppState, tempfile::NamedTempFile) { setup_test_state_with_platform(None).await } @@ -1442,6 +1528,7 @@ pNs85uhOZE8z2jr8Pg== keys: dummy_keys, vm_config: String::new(), cert_client: dummy_cert_client, + historical_roots: Default::default(), demo_cert: Default::default(), demo_cert_last_failure: Default::default(), platform: Arc::new(TestSimulatorPlatform { diff --git a/dstack/guest-agent/src/rpc_service_v1.rs b/dstack/guest-agent/src/rpc_service_v1.rs index 354e94618a..8f178ccbf0 100644 --- a/dstack/guest-agent/src/rpc_service_v1.rs +++ b/dstack/guest-agent/src/rpc_service_v1.rs @@ -159,16 +159,19 @@ impl V1RpcHandler { Self { state } } - /// Derive the key named by `(domain, algorithm)` together with its - /// signature chain. - fn derive(&self, domain: &str, algorithm: &str) -> Result<(AppKey, Vec>)> { + /// Derive the key named by `(domain, algorithm)` under the app root of + /// `root_epoch`, together with its signature chain and that epoch. + async fn derive( + &self, + domain: &str, + algorithm: &str, + root_epoch: Option, + ) -> Result<(AppKey, Vec>, u32)> { let algorithm = Algorithm::parse(algorithm)?; - let key = AppKey::derive(self.state.app_root_k256_key(), domain, algorithm)?; - let chain = vec![ - key.claim_signature(self.state.app_root_signing_key()?)?, - self.state.kms_k256_signature().to_vec(), - ]; - Ok((key, chain)) + let root = self.state.app_root(root_epoch).await?; + let key = AppKey::derive(&root.k256_key, domain, algorithm)?; + let chain = vec![key.claim_signature(&root.signing_key)?, root.kms_signature]; + Ok((key, chain, root.epoch)) } } @@ -196,11 +199,14 @@ impl DstackGuestRpc for V1RpcHandler { } async fn get_key(self, request: GetKeyRequest) -> Result { - let (key, signature_chain) = self.derive(&request.domain, &request.algorithm)?; + let (key, signature_chain, root_epoch) = self + .derive(&request.domain, &request.algorithm, request.root_epoch) + .await?; Ok(GetKeyResponse { key: key.secret(), public_key: key.public_key(), signature_chain, + root_epoch, }) } @@ -342,9 +348,44 @@ mod tests { GetKeyRequest { domain: domain.to_string(), algorithm: algorithm.to_string(), + root_epoch: None, } } + #[tokio::test] + async fn get_key_selects_the_root_epoch() { + let (state, _guard) = state().await; + let handler = || V1RpcHandler::new(state.clone()); + let request = |root_epoch| GetKeyRequest { + root_epoch, + ..get_key_request("storage-encryption", "secp256k1") + }; + + let current = handler().get_key(request(None)).await.unwrap(); + assert_eq!(current.root_epoch, 0); + assert_eq!( + handler().get_key(request(Some(0))).await.unwrap().key, + current.key + ); + assert!(handler().get_key(request(Some(1))).await.is_err()); + } + + /// Older epochs come from the KMS once and are then served from memory; + /// the fixture has no KMS, so the cache is filled directly. + #[tokio::test] + async fn an_older_root_is_fetched_once_and_cached() { + let (state, _guard) = state().await; + let key = crate::rpc_service::tests::HISTORICAL_K256_KEY.to_vec(); + state + .insert_historical_root(1, key.clone(), vec![0xab; 65]) + .await; + let root = state.historical_root(1).await.unwrap(); + assert_eq!(root.k256_key, key); + assert_eq!(root.kms_signature, [0xab; 65]); + // An uncached epoch needs the KMS, which a local key provider lacks. + assert!(state.historical_root(2).await.is_err()); + } + /// The fixture's app root key is what `keys::tests` committed its vectors /// against. If the two drift, the vectors stop describing what this /// handler serves and every assertion built on them goes quiet. diff --git a/sdk/go/dstack/client_v1.go b/sdk/go/dstack/client_v1.go index fd0f98aced..5b27c29da6 100644 --- a/sdk/go/dstack/client_v1.go +++ b/sdk/go/dstack/client_v1.go @@ -42,6 +42,8 @@ type GetKeyV1Response struct { // key over the app root public key. See `docs/guest-api-v1.md` for the claim // encoding and the verification steps. SignatureChain [][]byte + // KMS root key epoch the key was derived from. + RootEpoch uint32 } // Response from a v1 attestation request. @@ -442,6 +444,18 @@ func (c *DstackClientV1) IssueCert(ctx context.Context, options ...IssueCertV1Op // Derivation is flat -- `a/b` is not a child of `a` -- and binds both arguments, // so these keys differ from the v0 keys of the same name. func (c *DstackClientV1) GetKey(ctx context.Context, domain string, algorithm string) (*GetKeyV1Response, error) { + return c.getKey(ctx, domain, algorithm, nil) +} + +// GetKeyAtRootEpoch is GetKey under the KMS root key epoch rootEpoch instead of +// the current one; epoch 0 is the original root keys. Name an older epoch only +// to reach a key from before a KMS root key rotation, for example to migrate +// off it. +func (c *DstackClientV1) GetKeyAtRootEpoch(ctx context.Context, domain string, algorithm string, rootEpoch uint32) (*GetKeyV1Response, error) { + return c.getKey(ctx, domain, algorithm, &rootEpoch) +} + +func (c *DstackClientV1) getKey(ctx context.Context, domain string, algorithm string, rootEpoch *uint32) (*GetKeyV1Response, error) { if algorithm == "" { return nil, fmt.Errorf("algorithm is required, use `secp256k1` or `ed25519`") } @@ -450,6 +464,9 @@ func (c *DstackClientV1) GetKey(ctx context.Context, domain string, algorithm st "domain": domain, "algorithm": algorithm, } + if rootEpoch != nil { + payload["root_epoch"] = *rootEpoch + } data, err := c.sendRPCRequest(ctx, "/v1/GetKey", payload) if err != nil { @@ -464,6 +481,7 @@ func (c *DstackClientV1) GetKey(ctx context.Context, domain string, algorithm st Key *string `json:"key"` PublicKey *string `json:"public_key"` SignatureChain *[]*string `json:"signature_chain"` + RootEpoch uint32 `json:"root_epoch"` } if err := json.Unmarshal(data, &response); err != nil { return nil, err @@ -489,7 +507,12 @@ func (c *DstackClientV1) GetKey(ctx context.Context, domain string, algorithm st } } - return &GetKeyV1Response{Key: key, PublicKey: publicKey, SignatureChain: chain}, nil + return &GetKeyV1Response{ + Key: key, + PublicKey: publicKey, + SignatureChain: chain, + RootEpoch: response.RootEpoch, + }, nil } // Attest produces a versioned attestation over reportData, which must be diff --git a/sdk/js/src/client-v1.ts b/sdk/js/src/client-v1.ts index d35de8590e..9d0f94b158 100644 --- a/sdk/js/src/client-v1.ts +++ b/sdk/js/src/client-v1.ts @@ -68,6 +68,8 @@ export interface GetKeyResponseV1 { public_key: Uint8Array /** Two links: the app root key over the v1 key claim, then the KMS root key. */ signature_chain: Uint8Array[] + /** KMS root key epoch the key was derived from. */ + root_epoch: number } export interface AttestResponseV1 { @@ -333,14 +335,19 @@ export class DstackClientV1 { * @param domain Caller-chosen domain-separation string. May be empty. * @param algorithm Exactly `secp256k1` or `ed25519`. No default, no `k256` alias. */ - async getKey(domain: string, algorithm: string): Promise { + async getKey(domain: string, algorithm: string, rootEpoch?: number): Promise { // v0 defaulted an empty algorithm to secp256k1, which let a typo hand back a // key of the wrong type under a name the caller thought meant something else. if (!algorithm) { throw new Error('algorithm is required, use "secp256k1" or "ed25519"') } - const payload = JSON.stringify({ domain, algorithm }) - const result = await send_rpc_request<{ key: string, public_key: string, signature_chain: string[] }>( + // Unset is the current KMS root key epoch; name an older one only to migrate off it. + if (rootEpoch !== undefined && (!Number.isInteger(rootEpoch) || rootEpoch < 0 || rootEpoch > 0xffffffff)) { + throw new Error('rootEpoch must be a uint32') + } + const payload = JSON.stringify( + rootEpoch === undefined ? { domain, algorithm } : { domain, algorithm, root_epoch: rootEpoch }) + const result = await send_rpc_request<{ key: string, public_key: string, signature_chain: string[], root_epoch?: number }>( this.endpoint, '/v1/GetKey', payload) throwOnRpcError(result) return Object.freeze({ @@ -348,6 +355,7 @@ export class DstackClientV1 { public_key: from_hex(result.public_key, 'public_key'), signature_chain: to_list(result.signature_chain, 'signature_chain', 'error') .map((link, i) => from_hex(link, `signature_chain[${i}]`)), + root_epoch: result.root_epoch ?? 0, __name__: 'GetKeyResponseV1' as const, }) } diff --git a/sdk/python/src/dstack_sdk/dstack_client_v1.py b/sdk/python/src/dstack_sdk/dstack_client_v1.py index e57c8bba55..750b0b9427 100644 --- a/sdk/python/src/dstack_sdk/dstack_client_v1.py +++ b/sdk/python/src/dstack_sdk/dstack_client_v1.py @@ -107,6 +107,8 @@ class GetKeyResponseV1(BaseModel): # Two links: the app root key's signature over the v1 key claim, then the # KMS root key's signature over the app root public key. signature_chain: List[HexBytes] + # KMS root key epoch the key was derived from. + root_epoch: int = 0 class GpuEvidenceBundleV1(BaseModel): @@ -237,7 +239,9 @@ async def issue_cert( result = await self._send_rpc_request("IssueCert", data) return IssueCertResponseV1(**result) - async def get_key(self, domain: str, algorithm: str) -> GetKeyResponseV1: + async def get_key( + self, domain: str, algorithm: str, root_epoch: Optional[int] = None + ) -> GetKeyResponseV1: """Derive an application key from ``(domain, algorithm)``. Both arguments are required. Derivation is flat: two domains yield @@ -245,10 +249,19 @@ async def get_key(self, domain: str, algorithm: str) -> GetKeyResponseV1: exactly ``secp256k1`` or ``ed25519`` -- there is no default and no ``k256`` alias, because in v0 a typo silently produced a key of the wrong type under a name the caller thought meant something else. + + ``root_epoch`` selects the KMS root key epoch; the default ``None`` is + the current epoch, and epoch 0 is the original root keys. Name an older + epoch only to reach a key from before a KMS root key rotation, for + example to migrate off it. """ if not algorithm: raise ValueError("algorithm is required, use `secp256k1` or `ed25519`") data: Dict[str, Any] = {"domain": domain, "algorithm": algorithm} + if root_epoch is not None: + if type(root_epoch) is not int or not 0 <= root_epoch <= 0xFFFFFFFF: + raise ValueError("root_epoch must be a uint32") + data["root_epoch"] = root_epoch result = await self._send_rpc_request("GetKey", data) return GetKeyResponseV1(**result) @@ -350,7 +363,9 @@ def issue_cert( raise NotImplementedError @call_async - def get_key(self, domain: str, algorithm: str) -> GetKeyResponseV1: + def get_key( + self, domain: str, algorithm: str, root_epoch: Optional[int] = None + ) -> GetKeyResponseV1: """Derive an application key from ``(domain, algorithm)``.""" raise NotImplementedError diff --git a/sdk/rust/src/dstack_client_v1.rs b/sdk/rust/src/dstack_client_v1.rs index efcc782365..71f9c5f33b 100644 --- a/sdk/rust/src/dstack_client_v1.rs +++ b/sdk/rust/src/dstack_client_v1.rs @@ -95,10 +95,35 @@ impl DstackClientV1 { /// default and no `k256` alias -- a typo is an error rather than a key of /// the wrong type under a name the caller misread. pub async fn get_key(&self, domain: &str, algorithm: &str) -> Result { + self.request_key(domain, algorithm, None).await + } + + /// [`Self::get_key`] under the KMS root key epoch `root_epoch` instead of + /// the current one; epoch 0 is the original root keys. Name an older epoch + /// only to reach a key from before a KMS root key rotation, for example to + /// migrate off it. + pub async fn get_key_at_root_epoch( + &self, + domain: &str, + algorithm: &str, + root_epoch: u32, + ) -> Result { + self.request_key(domain, algorithm, Some(root_epoch)).await + } + + async fn request_key( + &self, + domain: &str, + algorithm: &str, + root_epoch: Option, + ) -> Result { if algorithm.is_empty() { anyhow::bail!("algorithm is required, use `secp256k1` or `ed25519`") } - let data = json!({ "domain": domain, "algorithm": algorithm }); + let mut data = json!({ "domain": domain, "algorithm": algorithm }); + if let Some(root_epoch) = root_epoch { + data["root_epoch"] = root_epoch.into(); + } let response = self.send_rpc_request("GetKey", &data).await?; Ok(serde_json::from_value::(response)?) } diff --git a/sdk/rust/types/src/dstack_v1.rs b/sdk/rust/types/src/dstack_v1.rs index 917db972e6..1ef1d7fd13 100644 --- a/sdk/rust/types/src/dstack_v1.rs +++ b/sdk/rust/types/src/dstack_v1.rs @@ -125,6 +125,9 @@ pub struct GetKeyResponse { /// steps. Verifying is the relying party's job; this SDK does not do it. #[serde(with = "hex_vec")] pub signature_chain: Vec>, + /// KMS root key epoch the key was derived from. + #[serde(default)] + pub root_epoch: u32, } /// The value an attestation binds to: bytes, and only bytes. @@ -348,6 +351,7 @@ mod tests { key: vec![0x01; 32], public_key: vec![0x02; 33], signature_chain: vec![vec![0xaa, 0xbb], vec![0xcc]], + root_epoch: 0, }; let json = serde_json::to_string(&response).expect("serializes"); assert!(json.contains(r#""signature_chain":["aabb","cc"]"#));