feat(kms-auth)!: pass the TEE platform to policies and allow restricting it

CHANGELOG.md

@@ -21,6 +21,16 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Added
 
+- DstackKms, DstackApp: TEE platform allowlists. `IAppAuthV2` (the epoch-aware
+  policy interface) receives `AppInfoV2`, the legacy `AppInfo` plus the
+  `teeVariant` platform name the KMS reports (`dstack-tdx`, `dstack-amd-sev-snp`,
+  ...), passed through unchanged by auth-eth. Each contract allows any platform
+  until its owner calls `setAllowAnyTeeVariant(false)`; it then allows only the
+  names set by `setAllowedTeeVariants`, from the platforms `TeeVariantPolicy`
+  knows (setting an unknown one reverts). The policy covers every epoch, and a
+  restricted one refuses unknown platforms and the legacy `isAppAllowed` and
+  `isKmsAllowed`, whose callers report no platform. auth-eth requires `teeVariant` on
+  `/authorizeKms` and `/authorizeApp`.
 - KMS root and application key epochs, epoch-scoped authorization, SealedEnv
   HPKE encryption, automatic environment migration, and epoch-aware SDK key
   selection.
@@ -76,7 +86,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
   range; `remove*AtEpochs(0, latest root epoch, value)` in `DstackKms` and
   `remove*AtEpochs(0, currentAppEpoch, value)` in `DstackApp` revoke a value
   everywhere. `DstackApp` reports its setters as the new ERC-165 interface
-  `IAppEpochManagement` (`0xc65f31b6`). The `Manage.s.sol` scripts require
+  `IAppEpochManagement` (`0x8c0524c4`). The `Manage.s.sol` scripts require
   `ROOT_EPOCH` or `APP_EPOCH`, or a `FROM_*`/`TO_*` range. `DstackApp`'s
   `allowAnyDevice` stays a single switch: `setAllowAnyDevice(bool)` covers every
   app epoch, so dropping the device restriction cannot leave an older epoch

docs/app-key-epochs.md

@@ -24,9 +24,15 @@ implementation before advancing, and upgrade the KMS/auth backend and guests
 before using nonzero epochs. Once the app epoch is nonzero, the legacy
 `isAppAllowed` of both `DstackApp` and `DstackKms` refuses, since its callers
 derive epoch-zero keys, so the KMS contract calling the app policy must support
-epochs before the app epoch advances. The ERC-165 `IAppEpochAuth` capability combines
+epochs before the app epoch advances. The ERC-165 `IAppAuthV2` capability combines
 `currentAppEpoch()` and `authorizeApp(AuthorizationRequest)` in one interface; the
-policy returns the same `AuthorizationResult` struct as `DstackKms`. Legacy policies
+policy returns the same `AuthorizationResult` struct as `DstackKms`. The request
+carries `AppInfoV2`, the legacy `AppInfo` plus the `teeVariant` platform name the
+KMS reports, which `allowAnyTeeVariant` and `allowedTeeVariants` in `DstackApp` and
+`DstackKms` restrict to the platforms listed in `TeeVariantPolicy.sol`; listing a
+new platform takes a contract upgrade. Like `allowAnyDevice`, the policy covers
+every epoch: which platforms may hold an app's keys is a trust decision, not a
+property of a key generation. Legacy policies
 without this capability remain at zero; failures of an advertised getter propagate.
 
 The Ethereum auth backends query the policy through the KMS contract. The simple
@@ -73,7 +79,7 @@ writes app epochs up to the next one.
 and fail without them.
 
 App policies deployed before app epochs (v0.5.x, v0.6.0) keep working with an
-upgraded `DstackKms` and KMS: they report no `IAppEpochAuth`, so `authorizeApp`
+upgraded `DstackKms` and KMS: they report no `IAppAuthV2`, so `authorizeApp`
 asks their `isAppAllowed` and only for app epoch zero, and `currentAppEpoch`
 reports zero for them. Their owners keep the epoch-less setters until they
 upgrade, which `Manage.s.sol` detects and uses with no epoch or `APP_EPOCH=0`.
@@ -113,15 +119,15 @@ This permits staged key distribution without changing which epoch the KMS serves
 
 The original bool mappings and `allowAnyDevice` retain their storage locations and
 are used directly as row zero. New storage consumes storage-gap slots in each
-contract (four in `DstackKms`, three in `DstackApp`); no entries are copied and no reinitializer is required. Existing
+contract (seven in `DstackKms`, four in `DstackApp`); no entries are copied and no reinitializer is required. Existing
 owners, pending owners, policy flags, root registries and app epoch counters are
 preserved. Even if the pre-upgrade counter is nonzero, old global material
 approvals grant only row zero after upgrading: explicitly populate the intended
 nonzero rows. `allowAnyDevice` keeps covering every app epoch.
 
 Upgrade the KMS policy before app policies, then approve the intended rows before
 restarting guests or activating a new root. `DstackApp` advertises the optional
-ERC-165 `IAppEpochAuth` capability. `DstackKms.authorizeApp` first applies the
+ERC-165 `IAppAuthV2` capability. `DstackKms.authorizeApp` first applies the
 checks it owns (operation scope, registered root epoch, requested root epoch not
 beyond the current one, app registration, OS image at the selected root epoch,
 deployed contract), then forwards the request unchanged to the policy's

docs/onchain-governance.md

@@ -158,6 +158,7 @@ The central governance contract that manages OS image whitelisting, app registra
 | `addOsImageHashAtEpochs(uint256,uint256,bytes32)` | Whitelist an OS image hash in a range of root epochs |
 | `removeOsImageHashAtEpochs(uint256,uint256,bytes32)` | Remove an OS image from a range of root epochs; `0` through the latest revokes it |
 | `setGatewayAppId(string)` | Set the trusted Gateway app ID(s) |
+| `setAllowAnyTeeVariant(bool)` / `setAllowedTeeVariants(string[])` | Allow KMS instances and apps on any TEE platform, or only the listed ones, at every root epoch |
 | `registerApp(address)` | Register an app contract |
 | `deployAndRegisterApp(...)` | Deploy and register app in one transaction |
 | `isAppAllowed(AppBootInfo)` | Check if an app is allowed to boot |
@@ -176,6 +177,7 @@ to the next epoch. Removing a value from rows `0` through the current epoch revo
 | `addDeviceAtEpoch(uint256,bytes32)` / `addDeviceAtEpochs(uint256,uint256,bytes32)` | Whitelist a device ID |
 | `removeDeviceAtEpoch(uint256,bytes32)` / `removeDeviceAtEpochs(uint256,uint256,bytes32)` | Remove a device |
 | `setAllowAnyDevice(bool)` | Allow any device to run this app, at every app epoch |
+| `setAllowAnyTeeVariant(bool)` / `setAllowedTeeVariants(string[])` | Allow any TEE platform, or only the listed ones (e.g. `dstack-tdx`), at every app epoch |
 | `currentAppEpoch()` | The app epoch new keys are derived for |
 | `isAppAllowed(AppBootInfo)` | Check if app can boot with given config (app epoch 0 only) |
 | `disableUpgrades()` | Permanently disable contract upgrades |
@@ -189,7 +191,7 @@ app epoch 0 and the epoch-less setters `addComposeHash(bytes32)`,
 which authorize them for app epoch 0. Tell the two kinds apart by ERC-165 interface ID:
 
 ```bash
-cast call <app-id> "supportsInterface(bytes4)(bool)" 0xc65f31b6 --rpc-url $RPC_URL  # IAppEpochManagement: epoch setters
+cast call <app-id> "supportsInterface(bytes4)(bool)" 0x8c0524c4 --rpc-url $RPC_URL  # IAppEpochManagement: epoch setters
 cast call <app-id> "supportsInterface(bytes4)(bool)" 0x8fd37527 --rpc-url $RPC_URL  # IAppAuthBasicManagement: epoch-less setters
 ```
 

dstack/kms/auth-eth/README.md

@@ -188,8 +188,9 @@ this implementation keeps the API of the removed fastify + ethers server that us
 #### dstack-kms 0.6
 
 dstack-kms 0.6 posts `dstack_verifier::PolicyBootInfo`: bytes as unprefixed
-hex, plus `teeVariant` and `keyProviderInfo`, which the contract does not take
-and this backend drops. the KMS checks `appId`, `composeHash` and `instanceId`
+hex, plus `teeVariant` and `keyProviderInfo`. The legacy `/bootAuth/*` routes
+drop both. `/authorizeKms` and `/authorizeApp` require `teeVariant` and pass it
+to the contract unchanged in `AppInfoV2`; `keyProviderInfo` is always dropped. the KMS checks `appId`, `composeHash` and `instanceId`
 widths itself before calling the webhook. the `dstack-kms 0.6 bootAuth payload`
 tests pin this contract.
 

dstack/kms/auth-eth/contracts/DstackApp.sol

@@ -8,7 +8,8 @@ pragma solidity ^0.8.24;
 
 import "./IAppAuth.sol";
 import "./AuthorizationOperation.sol";
-import "./IAppEpochAuth.sol";
+import "./IAppAuthV2.sol";
+import "./TeeVariantPolicy.sol";
 import "./IAppEpochManagement.sol";
 import "@openzeppelin/contracts-upgradeable/proxy/utils/Initializable.sol";
 import "@openzeppelin/contracts-upgradeable/access/Ownable2StepUpgradeable.sol";
@@ -21,7 +22,7 @@ contract DstackApp is
     UUPSUpgradeable,
     ERC165Upgradeable,
     IAppAuth,
-    IAppEpochAuth,
+    IAppAuthV2,
     IAppEpochManagement
 {
     /// @notice Allowed compose hashes for this app. Epoch 0 only; for other epochs use the *AtEpoch getters and
@@ -48,6 +49,8 @@ contract DstackApp is
     mapping(uint256 => mapping(bytes32 => bool)) private _composeHashesByEpoch;
     mapping(uint256 => mapping(bytes32 => bool)) private _deviceIdsByEpoch;
 
+    TeeVariantPolicy.Policy private _teeVariants;
+
     event EpochPolicyChanged(bytes32 indexed policy, bytes32 indexed value, uint256 indexed epoch, bool enabled);
 
     /// @notice Legacy material events. Epoch 0 only; for other epochs use the *AtEpoch getters and EpochPolicyChanged.
@@ -58,6 +61,8 @@ contract DstackApp is
     event UpgradesDisabled();
     /// @notice Covers every app epoch; there is no EpochPolicyChanged counterpart.
     event AllowAnyDeviceSet(bool allowAny);
+    event AllowAnyTeeVariantSet(bool allowAny);
+    event AllowedTeeVariantsSet(string[] teeVariants);
     event RequireTcbUpToDateSet(bool requireUpToDate);
     /// @notice Additive audit event for reconstructing authorization policy.
     event PolicyChanged(address indexed actor, bytes32 indexed policy, bytes32 indexed value, bool enabled);
@@ -147,7 +152,7 @@ contract DstackApp is
         override(ERC165Upgradeable, IERC165)
         returns (bool)
     {
-        return interfaceId == type(IAppEpochAuth).interfaceId || interfaceId == type(IAppEpochManagement).interfaceId
+        return interfaceId == type(IAppAuthV2).interfaceId || interfaceId == type(IAppEpochManagement).interfaceId
             || interfaceId == 0x1e079198 // IAppAuth
             || super.supportsInterface(interfaceId);
     }
@@ -231,6 +236,30 @@ contract DstackApp is
         _emitPolicy("allow-any-device", bytes32(0), allowed);
     }
 
+    /// @notice Whether this app may run on any TEE platform, rather than only on
+    /// {allowedTeeVariants}. Covers every app epoch.
+    function allowAnyTeeVariant() public view returns (bool) {
+        return !_teeVariants.restricted;
+    }
+
+    /// @notice The TEE platforms, named as the KMS reports them, this app may run on
+    /// unless {allowAnyTeeVariant}.
+    function allowedTeeVariants() external view returns (string[] memory) {
+        return TeeVariantPolicy.allowedNames(_teeVariants);
+    }
+
+    function setAllowAnyTeeVariant(bool allowed) external onlyOwner {
+        _teeVariants.restricted = !allowed;
+        emit AllowAnyTeeVariantSet(allowed);
+        _emitPolicy("allow-any-tee-variant", bytes32(0), allowed);
+    }
+
+    function setAllowedTeeVariants(string[] calldata teeVariants) external onlyOwner {
+        TeeVariantPolicy.setAllowed(_teeVariants, teeVariants);
+        emit AllowedTeeVariantsSet(teeVariants);
+        _emitPolicy("allowed-tee-variants", keccak256(abi.encode(teeVariants)), true);
+    }
+
     // Set whether TCB status must be UpToDate to boot this app
     function setRequireTcbUpToDate(bool _requireUpToDate) external onlyOwner {
         requireTcbUpToDate = _requireUpToDate;
@@ -278,7 +307,8 @@ contract DstackApp is
         returns (bool isAllowed, string memory reason)
     {
         if (currentAppEpoch != 0) return (false, "App epoch requires epoch-aware authorization");
-        return _authorizeApp(bootInfo, 0);
+        // Legacy callers report no platform, so only an unrestricted policy admits them.
+        return _authorizeApp(bootInfo.tcbStatus, bootInfo.composeHash, bootInfo.deviceId, allowAnyTeeVariant(), 0);
     }
 
     /// Decide the request at the selected app epoch; the KMS fills gateway and root fields.
@@ -294,11 +324,21 @@ contract DstackApp is
             return result;
         }
         uint256 appEpoch = request.requestedAppEpoch.isSet ? request.requestedAppEpoch.value : currentAppEpoch;
-        (result.authorized, result.rejectionReason) = _authorizeApp(request.appInfo, appEpoch);
+        AppInfoV2 calldata bootInfo = request.appInfo;
+        (result.authorized, result.rejectionReason) = _authorizeApp(
+            bootInfo.tcbStatus,
+            bootInfo.composeHash,
+            bootInfo.deviceId,
+            TeeVariantPolicy.admits(_teeVariants, bootInfo.teeVariant),
+            appEpoch
+        );
     }
 
     function _authorizeApp(
-        IAppAuth.AppInfo calldata bootInfo,
+        string calldata tcbStatus,
+        bytes32 composeHash,
+        bytes32 deviceId,
+        bool teeVariantAllowed,
         uint256 appEpoch
     )
         private
@@ -307,23 +347,24 @@ contract DstackApp is
     {
         if (appEpoch > currentAppEpoch) return (false, "Requested app epoch is not active");
         // Optionally require TCB status to be up to date
-        if (
-            requireTcbUpToDate
-                && keccak256(abi.encodePacked(bootInfo.tcbStatus)) != keccak256(abi.encodePacked("UpToDate"))
-        ) {
+        if (requireTcbUpToDate && keccak256(abi.encodePacked(tcbStatus)) != keccak256(abi.encodePacked("UpToDate"))) {
             return (false, "TCB status is not up to date");
         }
 
         // Check if compose hash is allowed
-        if (!allowedComposeHashesAtEpoch(appEpoch, bootInfo.composeHash)) {
+        if (!allowedComposeHashesAtEpoch(appEpoch, composeHash)) {
             return (false, "Compose hash not allowed");
         }
 
         // Check if device is allowed (when device restriction is enabled)
-        if (!allowAnyDevice && !allowedDeviceIdsAtEpoch(appEpoch, bootInfo.deviceId)) {
+        if (!allowAnyDevice && !allowedDeviceIdsAtEpoch(appEpoch, deviceId)) {
             return (false, "Device not allowed");
         }
 
+        if (!teeVariantAllowed) {
+            return (false, "TEE variant not allowed");
+        }
+
         return (true, "");
     }
 
@@ -344,5 +385,5 @@ contract DstackApp is
     }
 
     // Add storage gap for upgradeable contracts
-    uint256[46] private __gap;
+    uint256[45] private __gap;
 }