Merge pull request #498 from Dstack-TEE/feature/app-tcb-toggle

feat(kms): optional TCB UpToDate requirement for apps

Author: Leechael

kms/auth-eth/contracts/DstackApp.sol

@@ -33,23 +33,50 @@ contract DstackApp is
     // Mapping of allowed device IDs for this app
     mapping(bytes32 => bool) public allowedDeviceIds;
 
+    // Whether to require TCB status to be UpToDate
+    bool public requireTcbUpToDate;
+
     // Additional events specific to DstackApp
     event UpgradesDisabled();
     event AllowAnyDeviceSet(bool allowAny);
+    event RequireTcbUpToDateSet(bool requireUpToDate);
 
     /// @custom:oz-upgrades-unsafe-allow constructor
     constructor() {
         _disableInitializers();
     }
 
-    // Initialize the contract
+    // Old initialize — preserved for upgrade compatibility
+    function initialize(
+        address initialOwner,
+        bool _disableUpgrades,
+        bool _allowAnyDevice,
+        bytes32 initialDeviceId,
+        bytes32 initialComposeHash
+    ) public initializer {
+        _initializeCommon(initialOwner, _disableUpgrades, _allowAnyDevice, initialDeviceId, initialComposeHash);
+    }
+
+    // New initialize — includes requireTcbUpToDate
     function initialize(
         address initialOwner,
         bool _disableUpgrades,
+        bool _requireTcbUpToDate,
         bool _allowAnyDevice,
         bytes32 initialDeviceId,
         bytes32 initialComposeHash
     ) public initializer {
+        requireTcbUpToDate = _requireTcbUpToDate;
+        _initializeCommon(initialOwner, _disableUpgrades, _allowAnyDevice, initialDeviceId, initialComposeHash);
+    }
+
+    function _initializeCommon(
+        address initialOwner,
+        bool _disableUpgrades,
+        bool _allowAnyDevice,
+        bytes32 initialDeviceId,
+        bytes32 initialComposeHash
+    ) internal {
         require(initialOwner != address(0), "invalid owner address");
 
         _upgradesDisabled = _disableUpgrades;
@@ -72,6 +99,10 @@ contract DstackApp is
         __ERC165_init();
     }
 
+    function version() public pure returns (uint256) {
+        return 2;
+    }
+
     /**
      * @dev See {IERC165-supportsInterface}.
      * @notice Returns true if this contract implements the interface defined by interfaceId
@@ -114,6 +145,12 @@ contract DstackApp is
         emit AllowAnyDeviceSet(_allowAnyDevice);
     }
 
+    // Set whether TCB status must be UpToDate to boot this app
+    function setRequireTcbUpToDate(bool _requireUpToDate) external onlyOwner {
+        requireTcbUpToDate = _requireUpToDate;
+        emit RequireTcbUpToDateSet(_requireUpToDate);
+    }
+
     // Add a device ID to allowed list
     function addDevice(bytes32 deviceId) external onlyOwner {
         allowedDeviceIds[deviceId] = true;
@@ -130,6 +167,15 @@ contract DstackApp is
     function isAppAllowed(
         IAppAuth.AppBootInfo calldata bootInfo
     ) external view override returns (bool isAllowed, string memory reason) {
+        // Optionally require TCB status to be up to date
+        if (
+            requireTcbUpToDate &&
+            keccak256(abi.encodePacked(bootInfo.tcbStatus)) !=
+            keccak256(abi.encodePacked("UpToDate"))
+        ) {
+            return (false, "TCB status is not up to date");
+        }
+
         // Check if compose hash is allowed
         if (!allowedComposeHashes[bootInfo.composeHash]) {
             return (false, "Compose hash not allowed");
@@ -150,5 +196,5 @@ contract DstackApp is
     }
 
     // Add storage gap for upgradeable contracts
-    uint256[50] private __gap;
+    uint256[49] private __gap;
 }

kms/auth-eth/contracts/DstackKms.sol

@@ -144,30 +144,47 @@ contract DstackKms is
     function deployAndRegisterApp(
         address initialOwner,
         bool disableUpgrades,
+        bool requireTcbUpToDate,
         bool allowAnyDevice,
         bytes32 initialDeviceId,
         bytes32 initialComposeHash
-    ) external returns (address appId) {
+    ) public returns (address appId) {
         require(appImplementation != address(0), "DstackApp implementation not set");
         require(initialOwner != address(0), "Invalid owner address");
 
-        // Prepare initialization data
         bytes memory initData = abi.encodeWithSelector(
-            bytes4(keccak256("initialize(address,bool,bool,bytes32,bytes32)")),
+            bytes4(keccak256("initialize(address,bool,bool,bool,bytes32,bytes32)")),
             initialOwner,
             disableUpgrades,
+            requireTcbUpToDate,
             allowAnyDevice,
             initialDeviceId,
             initialComposeHash
         );
 
-        // Deploy proxy contract
         appId = address(new ERC1967Proxy(appImplementation, initData));
-        // Register to KMS
         registerApp(appId);
         emit AppDeployedViaFactory(appId, msg.sender);
     }
 
+    // Backward compatible factory method for old SDK callers (pre TCB flag)
+    function deployAndRegisterApp(
+        address initialOwner,
+        bool disableUpgrades,
+        bool allowAnyDevice,
+        bytes32 initialDeviceId,
+        bytes32 initialComposeHash
+    ) external returns (address appId) {
+        return deployAndRegisterApp(
+            initialOwner,
+            disableUpgrades,
+            false,
+            allowAnyDevice,
+            initialDeviceId,
+            initialComposeHash
+        );
+    }
+
     // Function to register an aggregated MR measurement
     function addKmsAggregatedMr(bytes32 mrAggregated) external onlyOwner {
         kmsAllowedAggregatedMrs[mrAggregated] = true;

kms/auth-eth/foundry-cast-cheatsheet.md

@@ -84,6 +84,10 @@ npx hardhat kms:deploy-impl --network test
 cast send $KMS_CONTRACT_ADDRESS "upgradeTo(address)" "NEW_IMPL_ADDRESS" \
   --private-key $PRIVATE_KEY --rpc-url $RPC_URL --gas-limit 500000
 
+# Note: Existing KMS proxy deployments can be upgraded in-place using the steps above.
+# This release only adds optional app boot TCB checks in DstackApp and keeps the KMS
+# storage layout unchanged, so no initializer is required for the KMS upgrade.
+
 # Verify upgrade success
 cast storage $KMS_CONTRACT_ADDRESS 0x360894a13ba1a3210667c828492db98dca3e2076cc3735a920a3ca505d382bbc --rpc-url $RPC_URL
 # Should show the new implementation address
@@ -181,18 +185,18 @@ cast send $KMS_CONTRACT_ADDRESS "removeKmsDevice(bytes32)" \
 
 ```bash
 # kms:create-app - Deploy and register DstackApp in single transaction
-cast send $KMS_CONTRACT_ADDRESS "deployAndRegisterApp(address,bool,bool,bytes32,bytes32)" \
-  "$DEPLOYER_ADDRESS" false true \
+cast send $KMS_CONTRACT_ADDRESS "deployAndRegisterApp(address,bool,bool,bool,bytes32,bytes32)" \
+  "$DEPLOYER_ADDRESS" false false true \
   "0x1234567890abcdef1234567890abcdef1234567890abcdef1234567890abcdef" \
   "0xfedcba0987654321fedcba0987654321fedcba0987654321fedcba0987654321" \
   --private-key $PRIVATE_KEY --rpc-url $RPC_URL
-# Parameters: (owner, disableUpgrades, allowAnyDevice, initialDeviceId, initialComposeHash)
+# Parameters: (owner, disableUpgrades, requireTcbUpToDate, allowAnyDevice, initialDeviceId, initialComposeHash)
 # Use 0x0000...0000 for empty device/hash values
-# To decode return: cast abi-decode "deployAndRegisterApp(address,bool,bool,bytes32,bytes32)(address,address)" RETURN_DATA
+# To decode return: cast abi-decode "deployAndRegisterApp(address,bool,bool,bool,bytes32,bytes32)(address,address)" RETURN_DATA
 
 # Example with no initial data:
-cast send $KMS_CONTRACT_ADDRESS "deployAndRegisterApp(address,bool,bool,bytes32,bytes32)" \
-  "$DEPLOYER_ADDRESS" false true \
+cast send $KMS_CONTRACT_ADDRESS "deployAndRegisterApp(address,bool,bool,bool,bytes32,bytes32)" \
+  "$DEPLOYER_ADDRESS" false false true \
   "0x0000000000000000000000000000000000000000000000000000000000000000" \
   "0x0000000000000000000000000000000000000000000000000000000000000000" \
   --private-key $PRIVATE_KEY --rpc-url $RPC_URL
@@ -379,7 +383,7 @@ cast abi-decode "kmsAllowedAggregatedMrs(bytes32)(bool)" RETURN_DATA
 cast abi-decode "isAppAllowed((address,bytes32,address,bytes32,bytes32,bytes32,bytes32,string,string[]))(bool,string)" RETURN_DATA
 
 # Decode factory deployment response
-cast abi-decode "deployAndRegisterApp(address,bool,bool,bytes32,bytes32)(address,address)" RETURN_DATA
+cast abi-decode "deployAndRegisterApp(address,bool,bool,bool,bytes32,bytes32)(address,address)" RETURN_DATA
 ```
 
 ### Get Contract Information
@@ -459,7 +463,7 @@ cast send $KMS_CONTRACT_ADDRESS "addKmsAggregatedMr(bytes32)" "0x..." \
   --private-key $PRIVATE_KEY --rpc-url $RPC_URL
 
 # 3. Users can now deploy apps via factory immediately!
-cast send $KMS_CONTRACT_ADDRESS "deployAndRegisterApp(address,bool,bool,bytes32,bytes32)" \
+cast send $KMS_CONTRACT_ADDRESS "deployAndRegisterApp(address,bool,bool,bool,bytes32,bytes32)" \
   "$USER_ADDRESS" false true "0x..." "0x..." \
   --private-key $USER_PRIVATE_KEY --rpc-url $RPC_URL
 ```
@@ -485,7 +489,7 @@ cast send $KMS_CONTRACT_ADDRESS "addKmsAggregatedMr(bytes32)" "0x..." \
   --private-key $PRIVATE_KEY --rpc-url $RPC_URL
 
 # 5. Users can now deploy apps via factory
-cast send $KMS_CONTRACT_ADDRESS "deployAndRegisterApp(address,bool,bool,bytes32,bytes32)" \
+cast send $KMS_CONTRACT_ADDRESS "deployAndRegisterApp(address,bool,bool,bool,bytes32,bytes32)" \
   "$USER_ADDRESS" false true "0x..." "0x..." \
   --private-key $USER_PRIVATE_KEY --rpc-url $RPC_URL
 ```
@@ -547,7 +551,7 @@ mycast send $APP_AUTH_ADDRESS "addDevice(bytes32)" \
   "0x1234567890abcdef1234567890abcdef1234567890abcdef1234567890abcdef"
 
 # Example: Factory deployment
-mycast send $KMS_CONTRACT_ADDRESS "deployAndRegisterApp(address,bool,bool,bytes32,bytes32)" \
+mycast send $KMS_CONTRACT_ADDRESS "deployAndRegisterApp(address,bool,bool,bool,bytes32,bytes32)" \
   "$DEPLOYER_ADDRESS" false true \
   "0x1234567890abcdef1234567890abcdef1234567890abcdef1234567890abcdef" \
   "0xfedcba0987654321fedcba0987654321fedcba0987654321fedcba0987654321"
@@ -557,4 +561,4 @@ mycast storage $KMS_CONTRACT_ADDRESS 0x360894a13ba1a3210667c828492db98dca3e2076c
 
 # Example: Upgrade contract
 mycast send $KMS_CONTRACT_ADDRESS "upgradeTo(address)" "NEW_IMPL_ADDRESS" --gas-limit 500000
-``` 
+``` 

kms/auth-eth/hardhat.config.ts

@@ -255,6 +255,7 @@ task("kms:get-app-implementation", "Get current DstackApp implementation address
 
 task("app:deploy", "Deploy DstackApp with a UUPS proxy")
   .addFlag("allowAnyDevice", "Allow any device to boot this app")
+  .addFlag("requireTcbUpToDate", "Require TCB status to be UpToDate")
   .addOptionalParam("device", "Initial device ID", "", types.string)
   .addOptionalParam("hash", "Initial compose hash", "", types.string)
   .setAction(async (taskArgs, hre) => {
@@ -278,14 +279,14 @@ task("app:deploy", "Deploy DstackApp with a UUPS proxy")
       console.log("Initial compose hash:", composeHash === "0x0000000000000000000000000000000000000000000000000000000000000000" ? "none" : composeHash);
     }
 
-    // Use standard deployment - all cases use the same 6-parameter initializer
     const appContract = await deployContract(hre, "DstackApp", [
-      deployerAddress, 
-      false, 
-      taskArgs.allowAnyDevice,
+      deployerAddress,
+      false,                          // _disableUpgrades
+      taskArgs.requireTcbUpToDate,    // _requireTcbUpToDate
+      taskArgs.allowAnyDevice,        // _allowAnyDevice
       deviceId,
       composeHash
-    ]);
+    ], false, "initialize(address,bool,bool,bool,bytes32,bytes32)");
 
     if (!appContract) {
       return;
@@ -344,6 +345,7 @@ task("app:deploy", "Deploy DstackApp with a UUPS proxy")
 
 task("kms:create-app", "Create DstackApp via KMS factory method (single transaction)")
   .addFlag("allowAnyDevice", "Allow any device to boot this app")
+  .addFlag("requireTcbUpToDate", "Require TCB status to be UpToDate")
   .addOptionalParam("device", "Initial device ID", "", types.string)
   .addOptionalParam("hash", "Initial compose hash", "", types.string)
   .setAction(async (taskArgs, hre) => {
@@ -362,10 +364,11 @@ task("kms:create-app", "Create DstackApp via KMS factory method (single transact
     console.log("Initial compose hash:", composeHash === "0x0000000000000000000000000000000000000000000000000000000000000000" ? "none" : composeHash);
     console.log("Using factory method for single-transaction deployment...");
 
-    // Single transaction deployment via factory
-    const tx = await kmsContract.deployAndRegisterApp(
+    // Single transaction deployment via factory (explicit signature to disambiguate overloads)
+    const tx = await kmsContract["deployAndRegisterApp(address,bool,bool,bool,bytes32,bytes32)"](
       deployerAddress,  // deployer owns the contract
       false,           // disableUpgrades
+      taskArgs.requireTcbUpToDate,
       taskArgs.allowAnyDevice,
       deviceId,
       composeHash

kms/auth-eth/lib/deployment-helpers.ts

@@ -62,14 +62,15 @@ export async function accountBalance(ethers: any, address: string) {
 export async function estimateDeploymentCost(
     hre: HardhatRuntimeEnvironment,
     contractName: string,
-    initializerArgs: any[] = []
+    initializerArgs: any[] = [],
+    initializer?: string
 ) {
     console.log("Estimating deployment costs...");
     const factory = await hre.ethers.getContractFactory(contractName);
 
     // Get the data for initialize function
     const initData = factory.interface.encodeFunctionData(
-        "initialize",
+        initializer || "initialize",
         initializerArgs
     );
 

kms/auth-eth/scripts/deploy.ts

@@ -6,7 +6,7 @@ import { HardhatRuntimeEnvironment } from "hardhat/types";
 import * as helpers from "../lib/deployment-helpers";
 
 // This function should be called directly by Hardhat tasks
-export async function deployContract(hre: HardhatRuntimeEnvironment, contractName: string, initializerArgs: any[] = [], quiet: boolean = false) {
+export async function deployContract(hre: HardhatRuntimeEnvironment, contractName: string, initializerArgs: any[] = [], quiet: boolean = false, initializer?: string) {
   try {
     function log(...msgs: any[]) {
       if (!quiet) {
@@ -29,7 +29,8 @@ export async function deployContract(hre: HardhatRuntimeEnvironment, contractNam
       await helpers.estimateDeploymentCost(
         hre,
         contractName,
-        initializerArgs
+        initializerArgs,
+        initializer
       );
 
       // Prompt for confirmation
@@ -43,7 +44,7 @@ export async function deployContract(hre: HardhatRuntimeEnvironment, contractNam
     log("Deploying proxy...");
     const contract = await hre.upgrades.deployProxy(contractFactory,
       initializerArgs,
-      { kind: 'uups' }
+      { kind: 'uups', ...(initializer ? { initializer } : {}) }
     );
     log("Waiting for deployment...");
     await contract.waitForDeployment();