guest-agent: derive GetKey keys under a selected KMS root epoch

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<AppKeyResponse> {
+        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<String> {
         match self {
             CertRequestClient::Local { ca } => Ok(ca.pem_cert.clone()),

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;
 }
 
 // ---------------------------------------------------------------------------

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<AppStateInner>,
 }
 
+/// 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<u8>,
+    pub signing_key: SigningKey,
+    pub kms_signature: Vec<u8>,
+}
+
+impl AppRoot {
+    fn new(epoch: u32, k256_key: Vec<u8>, kms_signature: Vec<u8>) -> Result<Self> {
+        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<HashMap<u32, AppRoot>>,
     demo_cert: tokio::sync::OnceCell<String>,
     /// When the last demo-cert request failed; held while one is running.
     demo_cert_last_failure: tokio::sync::Mutex<Option<Instant>>,
@@ -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<u32>) -> Result<AppRoot> {
+        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<AppRoot> {
+        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<u8>,
+        kms_signature: Vec<u8>,
+    ) {
+        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 {

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<Vec<u8>>)> {
+    /// 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<u32>,
+    ) -> Result<(AppKey, Vec<Vec<u8>>, 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<GetKeyResponse> {
-        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.

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

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,21 +335,27 @@ 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<GetKeyResponseV1> {
+  async getKey(domain: string, algorithm: string, rootEpoch?: number): Promise<GetKeyResponseV1> {
     // 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({
       key: from_hex(result.key, 'key'),
       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,
     })
   }