Dstack-TEE · kvinwang · Feb 9, 2026
diff --git a/kms/auth-eth/contracts/DstackApp.sol b/kms/auth-eth/contracts/DstackApp.sol
@@ -33,9 +33,13 @@ 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() {
@@ -46,13 +50,15 @@ contract DstackApp is
     function initialize(
         address initialOwner,
         bool _disableUpgrades,
+        bool _requireTcbUpToDate,
         bool _allowAnyDevice,
         bytes32 initialDeviceId,
         bytes32 initialComposeHash
     ) public initializer {
         require(initialOwner != address(0), "invalid owner address");
 
         _upgradesDisabled = _disableUpgrades;
+        requireTcbUpToDate = _requireTcbUpToDate;
         allowAnyDevice = _allowAnyDevice;
 
         // Add initial device if provided
@@ -114,6 +120,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 +142,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 +171,5 @@ contract DstackApp is
     }
 
     // Add storage gap for upgradeable contracts
-    uint256[50] private __gap;
+    uint256[49] private __gap;
 }
diff --git a/kms/auth-eth/contracts/DstackKms.sol b/kms/auth-eth/contracts/DstackKms.sol
@@ -144,18 +144,38 @@ contract DstackKms is
     function deployAndRegisterApp(
         address initialOwner,
         bool disableUpgrades,
+        bool requireTcbUpToDate,
         bool allowAnyDevice,
         bytes32 initialDeviceId,
         bytes32 initialComposeHash
     ) external returns (address appId) {
+        return _deployAndRegisterApp(
+            initialOwner,
+            disableUpgrades,
+            requireTcbUpToDate,
+            allowAnyDevice,
+            initialDeviceId,
+            initialComposeHash
+        );
+    }
+
+    function _deployAndRegisterApp(
+        address initialOwner,
+        bool disableUpgrades,
+        bool requireTcbUpToDate,
+        bool allowAnyDevice,
+        bytes32 initialDeviceId,
+        bytes32 initialComposeHash
+    ) internal 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
@@ -168,6 +188,24 @@ contract DstackKms is
         emit AppDeployedViaFactory(appId, msg.sender);
     }
 
+    // Backward compatible factory method (pre TCB requirement flag)
+    function deployAndRegisterApp(
+        address initialOwner,
+        bool disableUpgrades,
+        bool allowAnyDevice,
+        bytes32 initialDeviceId,
+        bytes32 initialComposeHash
+    ) external returns (address appId) {
+        return _deployAndRegisterApp(
+            initialOwner,
+            disableUpgrades,
+            false, // requireTcbUpToDate (default)
+            allowAnyDevice,
+            initialDeviceId,
+            initialComposeHash
+        );
+    }
+
     // Function to register an aggregated MR measurement
     function addKmsAggregatedMr(bytes32 mrAggregated) external onlyOwner {
         kmsAllowedAggregatedMrs[mrAggregated] = true;

diff --git a/kms/auth-eth/foundry-cast-cheatsheet.md b/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
-``` 
+``` 
diff --git a/kms/auth-eth/hardhat.config.ts b/kms/auth-eth/hardhat.config.ts
@@ -333,6 +333,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) => {
@@ -355,6 +356,7 @@ task("kms:create-app", "Create DstackApp via KMS factory method (single transact
     const tx = await kmsContract.deployAndRegisterApp(
       deployerAddress,  // deployer owns the contract
       false,           // disableUpgrades
+      taskArgs.requireTcbUpToDate,
       taskArgs.allowAnyDevice,
       deviceId,
       composeHash

diff --git a/kms/auth-eth/test/DstackApp.test.ts b/kms/auth-eth/test/DstackApp.test.ts
@@ -20,6 +20,7 @@ describe("DstackApp", function () {
     appAuth = await deployContract(hre, "DstackApp", [
       owner.address, 
       false,  // _disableUpgrades
+      false,  // _requireTcbUpToDate
       true,   // _allowAnyDevice
       ethers.ZeroHash,  // initialDeviceId (empty)
       ethers.ZeroHash   // initialComposeHash (empty)
@@ -91,6 +92,26 @@ describe("DstackApp", function () {
       expect(isAllowed).to.be.true;
     });
 
+    it("Should reject outdated TCB when required", async function () {
+      await appAuth.setRequireTcbUpToDate(true);
+
+      const bootInfo = {
+        appId: appId,
+        composeHash,
+        instanceId,
+        deviceId,
+        mrAggregated,
+        mrSystem,
+        osImageHash,
+        tcbStatus: "OutOfDate",
+        advisoryIds: []
+      };
+
+      const [isAllowed, reason] = await appAuth.isAppAllowed(bootInfo);
+      expect(isAllowed).to.be.false;
+      expect(reason).to.equal("TCB status is not up to date");
+    });
+
     it("Should reject unallowed compose hash", async function () {
       const bootInfo = {
         tcbStatus: "UpToDate",
@@ -138,7 +159,7 @@ describe("DstackApp", function () {
       const contractFactory = await ethers.getContractFactory("DstackApp");
       appAuthWithData = await hre.upgrades.deployProxy(
         contractFactory,
-        [owner.address, false, false, testDevice, testHash],
+        [owner.address, false, false, false, testDevice, testHash],
         { 
           kind: 'uups'
         }
@@ -237,7 +258,7 @@ describe("DstackApp", function () {
       const contractFactory = await ethers.getContractFactory("DstackApp");
       const appAuthEmpty = await hre.upgrades.deployProxy(
         contractFactory,
-        [owner.address, false, false, ethers.ZeroHash, ethers.ZeroHash],
+        [owner.address, false, false, false, ethers.ZeroHash, ethers.ZeroHash],
         {
           kind: 'uups'
         }

diff --git a/kms/auth-eth/test/setup.ts b/kms/auth-eth/test/setup.ts
@@ -33,6 +33,7 @@ beforeAll(async () => {
   const appAuth = await deployContract(hre, "DstackApp", [
     owner.address, 
     false,  // _disableUpgrades
+    false,  // _requireTcbUpToDate
     true,   // _allowAnyDevice
     ethers.ZeroHash,  // initialDeviceId (empty)
     ethers.ZeroHash   // initialComposeHash (empty)