kms: Allow upgrade base image

kms/rpc/proto/kms_rpc.proto

@@ -4,10 +4,16 @@ import "google/protobuf/empty.proto";
 
 package kms;
 
+message GetAppKeyRequest {
+  // When upgradable is true, the disk encryption key is derived without the rootfs hash
+  // so that the app can upgrade the base image
+  bool upgradable = 1;
+}
+
 // The kms public RPC service.
 service KMS {
   // Request the app key given the app id and tdx quote
-  rpc GetAppKey(google.protobuf.Empty) returns (AppKeyResponse) {
+  rpc GetAppKey(GetAppKeyRequest) returns (AppKeyResponse) {
     // Retrieves the app key and certificate given the app id and tdx quote
   }
   // Request the app environment encryption public key given the app id

kms/src/main_service.rs

@@ -3,7 +3,7 @@ use std::sync::Arc;
 use anyhow::{bail, Context, Result};
 use kms_rpc::{
     kms_server::{KmsRpc, KmsServer},
-    AppId, AppKeyResponse, PublicKeyResponse,
+    AppId, AppKeyResponse, GetAppKeyRequest, PublicKeyResponse,
 };
 use ra_rpc::{CallContext, RpcCall};
 use ra_tls::{
@@ -111,7 +111,7 @@ impl RpcHandler {
 }
 
 impl KmsRpc for RpcHandler {
-    async fn get_app_key(self) -> Result<AppKeyResponse> {
+    async fn get_app_key(self, request: GetAppKeyRequest) -> Result<AppKeyResponse> {
         let attest = self.ensure_attested()?;
         let app_id = attest.decode_app_id().context("Failed to decode app ID")?;
         let instance_id = attest
@@ -133,17 +133,18 @@ impl KmsRpc for RpcHandler {
             &[app_id.as_bytes(), "app-key".as_bytes()],
         )
         .context("Failed to derive app key")?;
-
-        let app_disk_key = derive_ecdsa_key_pair(
-            &state.root_ca.key,
-            &[
-                rootfs_hash.as_bytes(),
-                app_id.as_bytes(),
-                instance_id.as_bytes(),
-                "app-disk-crypt-key".as_bytes(),
-            ],
-        )
-        .context("Failed to derive app disk key")?;
+        let mut context_data = if request.upgradable {
+            vec![]
+        } else {
+            vec![rootfs_hash.as_bytes()]
+        };
+        context_data.extend(vec![
+            app_id.as_bytes(),
+            instance_id.as_bytes(),
+            "app-disk-crypt-key".as_bytes(),
+        ]);
+        let app_disk_key = derive_ecdsa_key_pair(&state.root_ca.key, &context_data)
+            .context("Failed to derive app disk key")?;
 
         let env_crypt_key = {
             let secret = derive_dh_secret(

tdxctl/src/fde_setup.rs

@@ -7,6 +7,7 @@ use std::{
 
 use anyhow::{bail, Context, Result};
 use fs_err as fs;
+use kms_rpc::GetAppKeyRequest;
 use ra_rpc::client::RaClient;
 use serde::{Deserialize, Serialize};
 use tracing::{info, warn};
@@ -96,8 +97,7 @@ struct InstanceInfo {
 
 impl InstanceInfo {
     fn is_bootstrapped(&self) -> bool {
-        self.bootstrapped
-            .unwrap_or_else(|| !self.instance_id.is_empty())
+        self.bootstrapped.unwrap_or(!self.instance_id.is_empty())
     }
 }
 
@@ -228,7 +228,7 @@ impl SetupFdeArgs {
             )?;
             let kms_client = kms_rpc::kms_client::KmsClient::new(ra_client);
             let response = kms_client
-                .get_app_key()
+                .get_app_key(GetAppKeyRequest { upgradable: true })
                 .await
                 .context("Failed to get app key")?;
             let keys_json =