Rename mr_image to os_image_hash

Author: kvinwang

docs/deployment.md

@@ -87,7 +87,7 @@ cd dstack/kms/dstack-app/
 ```
 It will create a template `.env` file. Edit the `.env` file and set the required variables.
 Especially the `KMS_CONTRACT_ADDR` variable set to the address of the KmsAuth Proxy contract deployed in the previous step.
-The `IMAGE_DOWNLOAD_URL` variable should be set to the URL of the dstack OS image used to verify the mr_image.
+The `IMAGE_DOWNLOAD_URL` variable should be set to the URL of the dstack OS image used to verify the os_image_hash.
 ```
 # .env
 VMM_RPC=unix:../../vmm-data/vmm.sock

dstack-types/src/lib.rs

@@ -125,7 +125,7 @@ pub struct SysConfig {
 #[derive(Deserialize, Serialize, Debug, Clone)]
 pub struct VmConfig {
     #[serde(with = "hex_bytes")]
-    pub mr_image: Vec<u8>,
+    pub os_image_hash: Vec<u8>,
     pub cpu_count: u32,
     pub memory_size: u64,
 }

dstack-util/src/system_setup.rs

@@ -326,8 +326,8 @@ impl<'a> Stage0<'a> {
             .await
             .context("Failed to get app key")?;
 
-        extend_rtmr3("mr-image", &response.mr_image)
-            .context("Failed to extend mr-image to RTMR3")?;
+        extend_rtmr3("os-image-hash", &response.os_image_hash)
+            .context("Failed to extend os-image-hash to RTMR3")?;
 
         let keys = AppKeys {
             ca_cert: tmp_ca.ca_cert,

guest-agent/rpc/proto/agent_rpc.proto

@@ -192,7 +192,7 @@ message WorkerInfo {
   // MR Aggregated
   bytes mr_aggregated = 9;
   // MR Image
-  bytes mr_image = 10;
+  bytes os_image_hash = 10;
   // MR Key Provider
   bytes mr_key_provider = 11;
   // Key provider info

guest-agent/src/http_routes.rs

@@ -32,7 +32,7 @@ async fn index(state: &State<AppState>) -> Result<RawHtml<String>, String> {
         instance_id,
         device_id,
         mr_aggregated: _,
-        mr_image: _,
+        os_image_hash: _,
         mr_key_provider: _,
         key_provider_info,
         compose_hash: _,

guest-agent/src/rpc_service.rs

@@ -356,7 +356,7 @@ impl WorkerRpc for ExternalRpcHandler {
             "rtmr2": hex::encode(app_info.rtmr2),
             "rtmr3": hex::encode(app_info.rtmr3),
             "mr_aggregated": hex::encode(app_info.mr_aggregated),
-            "mr_image": hex::encode(&app_info.mr_image),
+            "os_image_hash": hex::encode(&app_info.os_image_hash),
             "mr_key_provider": hex::encode(app_info.mr_key_provider),
             "compose_hash": hex::encode(&app_info.compose_hash),
             "device_id": hex::encode(&app_info.device_id),
@@ -370,7 +370,7 @@ impl WorkerRpc for ExternalRpcHandler {
             instance_id: app_info.instance_id,
             device_id: app_info.device_id,
             mr_aggregated: app_info.mr_aggregated.to_vec(),
-            mr_image: app_info.mr_image.clone(),
+            os_image_hash: app_info.os_image_hash.clone(),
             mr_key_provider: app_info.mr_key_provider.to_vec(),
             key_provider_info: String::from_utf8(app_info.key_provider_info).unwrap_or_default(),
             compose_hash: app_info.compose_hash.clone(),

kms/auth-eth/contracts/IAppAuth.sol

@@ -9,7 +9,7 @@ interface IAppAuth {
         bytes32 deviceId;
         bytes32 mrAggregated;
         bytes32 mrSystem;
-        bytes32 mrImage;
+        bytes32 osImageHash;
         string tcbStatus;
         string[] advisoryIds;
     }

kms/auth-eth/contracts/KmsAuth.sol

@@ -43,10 +43,7 @@ contract KmsAuth is
     mapping(bytes32 => bool) public kmsAllowedDeviceIds;
 
     // Mapping of allowed image measurements
-    mapping(bytes32 => bool) public appAllowedImages;
-
-    // Mapping of allowed KMS compose hashes
-    mapping(bytes32 => bool) public appAllowedSystemMrs;
+    mapping(bytes32 => bool) public allowedOsImages;
 
     // Sequence number for app IDs - per user
     mapping(address => uint256) public nextAppSequence;
@@ -58,10 +55,8 @@ contract KmsAuth is
     event KmsAggregatedMrRemoved(bytes32 mrAggregated);
     event KmsDeviceAdded(bytes32 deviceId);
     event KmsDeviceRemoved(bytes32 deviceId);
-    event AppImageMrAdded(bytes32 mrImage);
-    event AppImageMrRemoved(bytes32 mrImage);
-    event AppSystemMrAdded(bytes32 mrSystem);
-    event AppSystemMrRemoved(bytes32 mrSystem);
+    event OsImageHashAdded(bytes32 osImageHash);
+    event OsImageHashRemoved(bytes32 osImageHash);
     event GatewayAppIdSet(string gatewayAppId);
 
     /// @custom:oz-upgrades-unsafe-allow constructor
@@ -150,27 +145,15 @@ contract KmsAuth is
     }
 
     // Function to register an image measurement
-    function addAppImageMr(bytes32 mrImage) external onlyOwner {
-        appAllowedImages[mrImage] = true;
-        emit AppImageMrAdded(mrImage);
+    function addOsImageHash(bytes32 osImageHash) external onlyOwner {
+        allowedOsImages[osImageHash] = true;
+        emit OsImageHashAdded(osImageHash);
     }
 
     // Function to deregister an image measurement
-    function removeAppImageMr(bytes32 mrImage) external onlyOwner {
-        appAllowedImages[mrImage] = false;
-        emit AppImageMrRemoved(mrImage);
-    }
-
-    // Function to register a system MR measurement
-    function addAppSystemMr(bytes32 mrSystem) external onlyOwner {
-        appAllowedSystemMrs[mrSystem] = true;
-        emit AppSystemMrAdded(mrSystem);
-    }
-
-    // Function to deregister a system MR measurement
-    function removeAppSystemMr(bytes32 mrSystem) external onlyOwner {
-        appAllowedSystemMrs[mrSystem] = false;
-        emit AppSystemMrRemoved(mrSystem);
+    function removeOsImageHash(bytes32 osImageHash) external onlyOwner {
+        allowedOsImages[osImageHash] = false;
+        emit OsImageHashRemoved(osImageHash);
     }
 
     // Function to check if KMS is allowed to boot
@@ -185,6 +168,11 @@ contract KmsAuth is
             return (false, "TCB status is not up to date");
         }
 
+        // Check if the OS image is allowed
+        if (!allowedOsImages[bootInfo.osImageHash]) {
+            return (false, "OS image is not allowed");
+        }
+
         // Check if the aggregated MR is allowed
         if (!kmsAllowedAggregatedMrs[bootInfo.mrAggregated]) {
             return (false, "Aggregated MR not allowed");
@@ -208,11 +196,8 @@ contract KmsAuth is
         }
 
         // Check aggregated MR and image measurements
-        if (
-            !appAllowedSystemMrs[bootInfo.mrSystem] &&
-            !appAllowedImages[bootInfo.mrImage]
-        ) {
-            return (false, "Neither system MR nor image is allowed");
+        if (!allowedOsImages[bootInfo.osImageHash]) {
+            return (false, "OS image is not allowed");
         }
 
         // Ask the app controller if the app is allowed to boot

kms/auth-eth/hardhat.config.ts

@@ -142,41 +142,23 @@ task("kms:remove", "Remove a Aggregated MR of an KMS instance")
 
 // Image Management Tasks
 task("kms:add-image", "Add an image measurement")
-  .addPositionalParam("mrImage", "Image measurement")
-  .setAction(async ({ mrImage }, { ethers }) => {
+  .addPositionalParam("osImageHash", "Image measurement")
+  .setAction(async ({ osImageHash }, { ethers }) => {
     const kmsAuth = await getKmsAuth(ethers);
-    const tx = await kmsAuth.addAppImageMr(mrImage);
+    const tx = await kmsAuth.addOsImageHash(osImageHash);
     await waitTx(tx);
     console.log("Image added successfully");
   });
 
 task("kms:remove-image", "Remove an image measurement")
-  .addPositionalParam("mrImage", "Image measurement")
-  .setAction(async ({ mrImage }, { ethers }) => {
+  .addPositionalParam("osImageHash", "Image measurement")
+  .setAction(async ({ osImageHash }, { ethers }) => {
     const kmsAuth = await getKmsAuth(ethers);
-    const tx = await kmsAuth.removeAppImageMr(mrImage);
+    const tx = await kmsAuth.removeOsImageHash(osImageHash);
     await waitTx(tx);
     console.log("Image removed successfully");
   });
 
-task("kms:add-system", "Add a system measurement")
-  .addPositionalParam("mrSystem", "System measurement")
-  .setAction(async ({ mrSystem }, { ethers }) => {
-    const kmsAuth = await getKmsAuth(ethers);
-    const tx = await kmsAuth.addAppSystemMr(mrSystem);
-    await waitTx(tx);
-    console.log("System measurement added successfully");
-  });
-
-task("kms:remove-system", "Remove a system measurement")
-  .addPositionalParam("mrSystem", "System measurement")
-  .setAction(async ({ mrSystem }, { ethers }) => {
-    const kmsAuth = await getKmsAuth(ethers);
-    const tx = await kmsAuth.removeAppSystemMr(mrSystem);
-    await waitTx(tx);
-    console.log("System measurement removed successfully");
-  });
-
 task("kms:add-device", "Add a device ID of an KMS instance")
   .addPositionalParam("deviceId", "Device ID")
   .setAction(async ({ deviceId }, { ethers }) => {

kms/auth-eth/run-tests.sh

@@ -0,0 +1 @@
+npm run test