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. `AuthorizationRequest` of
+  `IAppEpochAuth` (the epoch-aware policy interface) carries `teeVariant`, the
+  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

@@ -26,7 +26,13 @@ before using nonzero epochs. Once the app epoch is nonzero, the legacy
 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
 `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
+also carries `teeVariant`, the 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
@@ -113,7 +119,7 @@ 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

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. Set the list before restricting: a restricted empty list refuses every platform |
 | `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,9 +188,10 @@ 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`
-widths itself before calling the webhook. the `dstack-kms 0.6 bootAuth payload`
+hex, plus `teeVariant` and `keyProviderInfo`. the legacy `/bootAuth/*` routes
+drop both; `/authorizeKms` and `/authorizeApp` require `teeVariant` and pass it
+to the contract unchanged, and always drop `keyProviderInfo`. the KMS checks
+`appId`, `composeHash` and `instanceId` widths itself before calling the webhook. the `dstack-kms 0.6 bootAuth payload`
 tests pin this contract.
 
 #### OpenAPI specification

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

@@ -9,6 +9,7 @@ pragma solidity ^0.8.24;
 import "./IAppAuth.sol";
 import "./AuthorizationOperation.sol";
 import "./IAppEpochAuth.sol";
+import "./TeeVariantPolicy.sol";
 import "./IAppEpochManagement.sol";
 import "@openzeppelin/contracts-upgradeable/proxy/utils/Initializable.sol";
 import "@openzeppelin/contracts-upgradeable/access/Ownable2StepUpgradeable.sol";
@@ -48,6 +49,9 @@ contract DstackApp is
     mapping(uint256 => mapping(bytes32 => bool)) private _composeHashesByEpoch;
     mapping(uint256 => mapping(bytes32 => bool)) private _deviceIdsByEpoch;
 
+    // The TEE platforms this app may run on.
+    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 +62,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);
@@ -231,6 +237,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 {
+        uint64 mask = TeeVariantPolicy.setAllowed(_teeVariants, teeVariants);
+        emit AllowedTeeVariantsSet(teeVariants);
+        _emitPolicy("allowed-tee-variants", bytes32(uint256(mask)), true);
+    }
+
     // Set whether TCB status must be UpToDate to boot this app
     function setRequireTcbUpToDate(bool _requireUpToDate) external onlyOwner {
         requireTcbUpToDate = _requireUpToDate;
@@ -278,7 +308,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, allowAnyTeeVariant(), 0);
     }
 
     /// Decide the request at the selected app epoch; the KMS fills gateway and root fields.
@@ -294,11 +325,13 @@ contract DstackApp is
             return result;
         }
         uint256 appEpoch = request.requestedAppEpoch.isSet ? request.requestedAppEpoch.value : currentAppEpoch;
-        (result.authorized, result.rejectionReason) = _authorizeApp(request.appInfo, appEpoch);
+        (result.authorized, result.rejectionReason) =
+            _authorizeApp(request.appInfo, TeeVariantPolicy.admits(_teeVariants, request.teeVariant), appEpoch);
     }
 
     function _authorizeApp(
         IAppAuth.AppInfo calldata bootInfo,
+        bool teeVariantAllowed,
         uint256 appEpoch
     )
         private
@@ -324,6 +357,10 @@ contract DstackApp is
             return (false, "Device not allowed");
         }
 
+        if (!teeVariantAllowed) {
+            return (false, "TEE variant not allowed");
+        }
+
         return (true, "");
     }
 
@@ -344,5 +381,5 @@ contract DstackApp is
     }
 
     // Add storage gap for upgradeable contracts
-    uint256[46] private __gap;
+    uint256[45] private __gap;
 }

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

@@ -9,6 +9,7 @@ pragma solidity ^0.8.24;
 import "./IAppAuth.sol";
 import "./AuthorizationOperation.sol";
 import "./IAppEpochAuth.sol";
+import "./TeeVariantPolicy.sol";
 import "@openzeppelin/contracts/utils/introspection/ERC165Checker.sol";
 import "@openzeppelin/contracts-upgradeable/proxy/utils/Initializable.sol";
 import "@openzeppelin/contracts-upgradeable/access/Ownable2StepUpgradeable.sol";
@@ -89,6 +90,9 @@ contract DstackKms is Initializable, Ownable2StepUpgradeable, UUPSUpgradeable, E
     mapping(uint256 => mapping(bytes32 => bool)) private _kmsDeviceIdsByRoot;
     mapping(uint256 => mapping(bytes32 => bool)) private _osImagesByRoot;
 
+    // The TEE platforms KMS instances and apps may run on.
+    TeeVariantPolicy.Policy private _teeVariants;
+
     event EpochPolicyChanged(bytes32 indexed policy, bytes32 indexed value, uint256 indexed epoch, bool enabled);
 
     // Events
@@ -104,6 +108,8 @@ contract DstackKms is Initializable, Ownable2StepUpgradeable, UUPSUpgradeable, E
     event OsImageHashAdded(bytes32 osImageHash);
     event OsImageHashRemoved(bytes32 osImageHash);
     event GatewayAppIdSet(string gatewayAppId);
+    event AllowAnyTeeVariantSet(bool allowAny);
+    event AllowedTeeVariantsSet(string[] teeVariants);
     // slither-disable-next-line unindexed-event-address
     event AppImplementationSet(address implementation);
     event RootEpochRegistered(
@@ -267,6 +273,30 @@ contract DstackKms is Initializable, Ownable2StepUpgradeable, UUPSUpgradeable, E
         return (keys.caPubkey, keys.k256Pubkey, keys.evidenceUrl, keys.chainHash);
     }
 
+    /// @notice Whether KMS instances and apps may run on any TEE platform, rather than only on
+    /// {allowedTeeVariants}. Covers every root epoch.
+    function allowAnyTeeVariant() public view returns (bool) {
+        return !_teeVariants.restricted;
+    }
+
+    /// @notice The TEE platforms, named as the KMS reports them, KMS instances and apps 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 {
+        uint64 mask = TeeVariantPolicy.setAllowed(_teeVariants, teeVariants);
+        emit AllowedTeeVariantsSet(teeVariants);
+        _emitPolicy("allowed-tee-variants", bytes32(uint256(mask)), true);
+    }
+
     // Function to set trusted Gateway App ID
     function setGatewayAppId(string memory appId) external onlyOwner {
         gatewayAppId = appId;
@@ -550,11 +580,14 @@ contract DstackKms is Initializable, Ownable2StepUpgradeable, UUPSUpgradeable, E
         }
         uint256 selectedRootEpoch =
             request.requestedRootEpoch.isSet ? request.requestedRootEpoch.value : _currentRootEpoch;
-        return _checkAppDeployment(request.appInfo, selectedRootEpoch);
+        return _checkAppDeployment(
+            request.appInfo, TeeVariantPolicy.admits(_teeVariants, request.teeVariant), selectedRootEpoch
+        );
     }
 
     function _checkAppDeployment(
         AppInfo calldata bootInfo,
+        bool teeVariantAllowed,
         uint256 selectedRootEpoch
     )
         private
@@ -564,6 +597,7 @@ contract DstackKms is Initializable, Ownable2StepUpgradeable, UUPSUpgradeable, E
         if (!registeredApps[bootInfo.appId]) return "App not registered";
         if (!allowedOsImagesAtEpoch(selectedRootEpoch, bootInfo.osImageHash)) return "OS image is not allowed";
         if (!isContract(bootInfo.appId)) return "App not deployed or invalid address";
+        if (!teeVariantAllowed) return "TEE variant not allowed";
         return "";
     }
 
@@ -600,7 +634,8 @@ contract DstackKms is Initializable, Ownable2StepUpgradeable, UUPSUpgradeable, E
         }
         uint256 selectedRootEpoch =
             request.requestedRootEpoch.isSet ? request.requestedRootEpoch.value : _currentRootEpoch;
-        (result.authorized, result.rejectionReason) = _kmsPolicy(request.appInfo, selectedRootEpoch);
+        (result.authorized, result.rejectionReason) =
+            _kmsPolicy(request.appInfo, TeeVariantPolicy.admits(_teeVariants, request.teeVariant), selectedRootEpoch);
     }
 
     // Legacy entry point for pre-epoch KMS binaries: reads epoch 0, and refuses
@@ -609,10 +644,19 @@ contract DstackKms is Initializable, Ownable2StepUpgradeable, UUPSUpgradeable, E
         if (_currentRootEpoch != 0) {
             return (false, "Root epoch requires epoch-aware authorization");
         }
-        return _kmsPolicy(bootInfo, 0);
+        // Legacy callers report no platform, so only an unrestricted policy admits them.
+        return _kmsPolicy(bootInfo, allowAnyTeeVariant(), 0);
     }
 
-    function _kmsPolicy(AppInfo memory bootInfo, uint256 selectedRootEpoch) private view returns (bool, string memory) {
+    function _kmsPolicy(
+        AppInfo memory bootInfo,
+        bool teeVariantAllowed,
+        uint256 selectedRootEpoch
+    )
+        private
+        view
+        returns (bool, string memory)
+    {
         // Check if the TCB status is up to date
         if (keccak256(abi.encodePacked(bootInfo.tcbStatus)) != keccak256(abi.encodePacked("UpToDate"))) {
             return (false, "TCB status is not up to date");
@@ -633,6 +677,10 @@ contract DstackKms is Initializable, Ownable2StepUpgradeable, UUPSUpgradeable, E
             return (false, "KMS is not allowed to boot on this device");
         }
 
+        if (!teeVariantAllowed) {
+            return (false, "TEE variant not allowed");
+        }
+
         return (true, "");
     }
 
@@ -650,7 +698,8 @@ contract DstackKms is Initializable, Ownable2StepUpgradeable, UUPSUpgradeable, E
         if (_appEpoch(bootInfo.appId) != 0) {
             return (false, "App epoch requires epoch-aware authorization");
         }
-        string memory rejection = _checkAppDeployment(bootInfo, 0);
+        // Legacy callers report no platform, so only an unrestricted policy admits them.
+        string memory rejection = _checkAppDeployment(bootInfo, allowAnyTeeVariant(), 0);
         if (bytes(rejection).length != 0) return (false, rejection);
         // Call the app's isAppAllowed function. The tuple is forwarded as the
         // function's return value; the slither detector trips on the named
@@ -660,7 +709,7 @@ contract DstackKms is Initializable, Ownable2StepUpgradeable, UUPSUpgradeable, E
     }
 
     // Add storage gap for upgradeable contracts
-    uint256[44] private __gap;
+    uint256[43] private __gap;
 }
 
 function isContract(address addr) view returns (bool) {

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

@@ -19,6 +19,10 @@ struct AuthorizationRequest {
     bytes32 operation;
     OptionalUint256 requestedAppEpoch;
     OptionalUint256 requestedRootEpoch;
+    /// The TEE platform the KMS attested, e.g. "dstack-tdx" or "dstack-amd-sev-snp".
+    /// Measurements, device IDs and TCB status are only meaningful per platform,
+    /// and one OS image may boot on several.
+    string teeVariant;
 }
 
 struct AuthorizationResult {

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

@@ -8,7 +8,8 @@ pragma solidity ^0.8.24;
 
 /// Allowlist management. Each allowlist setter names the app epoch rows it
 /// writes: one epoch, or `fromEpoch` through `toEpoch`, at most the next epoch.
-/// `setAllowAnyDevice` covers every app epoch.
+/// `setAllowAnyDevice`, `setAllowAnyTeeVariant` and `setAllowedTeeVariants` cover
+/// every app epoch.
 /// Replaces the epoch-less IAppAuthBasicManagement (0x8fd37527), which policies
 /// predating app epochs still report; tools tell the two apart by interface ID.
 interface IAppEpochManagement {
@@ -21,4 +22,6 @@ interface IAppEpochManagement {
     function removeDeviceAtEpoch(uint256 epoch, bytes32 deviceId) external;
     function removeDeviceAtEpochs(uint256 fromEpoch, uint256 toEpoch, bytes32 deviceId) external;
     function setAllowAnyDevice(bool allowed) external;
+    function setAllowAnyTeeVariant(bool allowed) external;
+    function setAllowedTeeVariants(string[] calldata teeVariants) external;
 }