From cc1d4b9e3572a43a9b6e723d7c1f3fb4e66ea7b4 Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Thu, 12 Nov 2026 02:53:12 -0800 Subject: [PATCH] feat(kms-auth)!: pass the TEE platform to policies and allow restricting it One OS image hash can boot on several TEE platforms (a release ships both TDX and SEV-SNP firmware), and app keys derive from the app ID alone, so an app approved on bare-metal TDX could be served on any platform the KMS enables. The KMS already sends teeVariant; auth-eth dropped it. - Append teeVariant, the platform name the KMS reports, to the AuthorizationRequest of IAppEpochAuth. auth-eth passes it through unchanged, so the KMS defines the names and each policy picks its own representation. AppInfo is unchanged, so DstackKms still forwards it as is to policies without IAppEpochAuth. - DstackApp and DstackKms allow any platform until the owner calls setAllowAnyTeeVariant(false), then only the names set by setAllowedTeeVariants. TeeVariantPolicy keeps them as one packed slot: a restricted flag and a bitmask over the platform names it hardcodes, so a check costs one read; listing an unknown name reverts, and a restricted policy refuses one. The policy covers every epoch: which platforms may hold keys is a trust decision, not a property of a key generation. A restricted policy refuses the legacy isAppAllowed/isKmsAllowed, whose callers report no platform; the 0.6.x KMS also serves SEV-SNP and Nitro, so they cannot be taken as TDX. - auth-eth requires teeVariant on /authorizeKms and /authorizeApp. - Manage.s.sol SetKmsTeeVariants/SetTeeVariants take TEE_VARIANTS=any or a list of names and refuse an empty one; Query.s.sol prints the policies. DstackKms stays under the EIP-170 limit at 24,041 bytes. BREAKING CHANGE: unreleased ABI only: authorizeKms/authorizeApp, and the IAppEpochAuth and IAppEpochManagement (now 0x8c0524c4) interface IDs. Signed-off-by: Kevin Wang --- CHANGELOG.md | 12 ++- docs/app-key-epochs.md | 10 ++- docs/onchain-governance.md | 4 +- dstack/kms/auth-eth/README.md | 7 +- dstack/kms/auth-eth/contracts/DstackApp.sol | 43 +++++++++- dstack/kms/auth-eth/contracts/DstackKms.sol | 61 +++++++++++-- .../kms/auth-eth/contracts/IAppEpochAuth.sol | 4 + .../contracts/IAppEpochManagement.sol | 5 +- .../auth-eth/contracts/TeeVariantPolicy.sol | 69 +++++++++++++++ .../kms/auth-eth/docs/formal-verification.md | 3 + dstack/kms/auth-eth/index.test.ts | 6 +- dstack/kms/auth-eth/index.ts | 3 + dstack/kms/auth-eth/openapi.json | 7 ++ dstack/kms/auth-eth/script/Manage.s.sol | 37 ++++++++ dstack/kms/auth-eth/script/Query.s.sol | 9 ++ dstack/kms/auth-eth/script/README.md | 18 ++++ dstack/kms/auth-eth/test/DstackApp.t.sol | 58 ++++++++++++- .../test/DstackAppEpoch.symbolic.t.sol | 34 +++++++- dstack/kms/auth-eth/test/DstackKms.t.sol | 86 +++++++++++++++++-- .../test/DstackKmsEpoch.symbolic.t.sol | 6 +- .../kms/auth-eth/test/LegacyAppCompat.t.sol | 24 +++++- dstack/kms/auth-eth/test/ManageScript.t.sol | 26 ++++++ 22 files changed, 501 insertions(+), 31 deletions(-) create mode 100644 dstack/kms/auth-eth/contracts/TeeVariantPolicy.sol diff --git a/CHANGELOG.md b/CHANGELOG.md index f4a360311..71b996488 100644 --- a/CHANGELOG.md +++ b/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 diff --git a/docs/app-key-epochs.md b/docs/app-key-epochs.md index f89f13c64..bbb759ecc 100644 --- a/docs/app-key-epochs.md +++ b/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 diff --git a/docs/onchain-governance.md b/docs/onchain-governance.md index cb4555244..db389a87a 100644 --- a/docs/onchain-governance.md +++ b/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 "supportsInterface(bytes4)(bool)" 0xc65f31b6 --rpc-url $RPC_URL # IAppEpochManagement: epoch setters +cast call "supportsInterface(bytes4)(bool)" 0x8c0524c4 --rpc-url $RPC_URL # IAppEpochManagement: epoch setters cast call "supportsInterface(bytes4)(bool)" 0x8fd37527 --rpc-url $RPC_URL # IAppAuthBasicManagement: epoch-less setters ``` diff --git a/dstack/kms/auth-eth/README.md b/dstack/kms/auth-eth/README.md index 5a905ceb2..0e2536d21 100644 --- a/dstack/kms/auth-eth/README.md +++ b/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 diff --git a/dstack/kms/auth-eth/contracts/DstackApp.sol b/dstack/kms/auth-eth/contracts/DstackApp.sol index 96b974c14..6266d77f1 100644 --- a/dstack/kms/auth-eth/contracts/DstackApp.sol +++ b/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; } diff --git a/dstack/kms/auth-eth/contracts/DstackKms.sol b/dstack/kms/auth-eth/contracts/DstackKms.sol index 91651a39a..d35b95350 100644 --- a/dstack/kms/auth-eth/contracts/DstackKms.sol +++ b/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) { diff --git a/dstack/kms/auth-eth/contracts/IAppEpochAuth.sol b/dstack/kms/auth-eth/contracts/IAppEpochAuth.sol index be720f49c..4134f7a0f 100644 --- a/dstack/kms/auth-eth/contracts/IAppEpochAuth.sol +++ b/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 { diff --git a/dstack/kms/auth-eth/contracts/IAppEpochManagement.sol b/dstack/kms/auth-eth/contracts/IAppEpochManagement.sol index 8af0f5b85..28ed4e2d7 100644 --- a/dstack/kms/auth-eth/contracts/IAppEpochManagement.sol +++ b/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; } diff --git a/dstack/kms/auth-eth/contracts/TeeVariantPolicy.sol b/dstack/kms/auth-eth/contracts/TeeVariantPolicy.sol new file mode 100644 index 000000000..e1780461d --- /dev/null +++ b/dstack/kms/auth-eth/contracts/TeeVariantPolicy.sol @@ -0,0 +1,69 @@ +// SPDX-FileCopyrightText: © 2026 Phala Network +// +// SPDX-License-Identifier: Apache-2.0 + +pragma solidity ^0.8.24; + +/// The TEE platforms a policy admits, named as the KMS reports them in +/// `AuthorizationRequest.teeVariant`. It covers every epoch: which platforms may hold keys +/// is a trust decision, not a property of a key generation. +library TeeVariantPolicy { + /// Zero-initialized as unrestricted, so proxies upgraded from before platform + /// policies keep their decisions. Bit `n` of `allowed` admits `names()[n]`. + struct Policy { + bool restricted; + uint64 allowed; + } + + uint256 private constant COUNT = 5; + + /// The platforms a policy can list, in the order of dstack's `TeeVariant`. + /// Append new ones (and bump COUNT); a restricted policy refuses names it + /// does not know. + function names() internal pure returns (string[COUNT] memory) { + return ["dstack-tdx", "dstack-gcp-tdx", "dstack-nitro-enclave", "dstack-amd-sev-snp", "dstack-aws-nitro-tpm"]; + } + + function admits(Policy storage policy, string calldata teeVariant) internal view returns (bool) { + if (!policy.restricted) return true; + (bool known, uint8 index) = _indexOf(teeVariant); + return known && (policy.allowed >> index) & 1 == 1; + } + + /// Returns the stored mask, a canonical value for audit events. + function setAllowed(Policy storage policy, string[] calldata teeVariants) internal returns (uint64) { + uint64 allowed = 0; + for (uint256 i = 0; i < teeVariants.length; ++i) { + (bool known, uint8 index) = _indexOf(teeVariants[i]); + require(known, "unknown TEE variant"); + allowed |= uint64(1) << index; + } + policy.allowed = allowed; + return allowed; + } + + function allowedNames(Policy storage policy) internal view returns (string[] memory list) { + string[COUNT] memory all = names(); + uint64 allowed = policy.allowed; + list = new string[](_popcount(allowed)); + uint256 n = 0; + for (uint256 i = 0; i < all.length; ++i) { + if ((allowed >> i) & 1 == 1) list[n++] = all[i]; + } + } + + function _indexOf(string calldata teeVariant) private pure returns (bool, uint8) { + bytes32 key = keccak256(bytes(teeVariant)); + string[COUNT] memory all = names(); + for (uint8 i = 0; i < all.length; ++i) { + if (keccak256(bytes(all[i])) == key) return (true, i); + } + return (false, 0); + } + + function _popcount(uint64 bits) private pure returns (uint256 count) { + for (; bits != 0; bits &= bits - 1) { + ++count; + } + } +} diff --git a/dstack/kms/auth-eth/docs/formal-verification.md b/dstack/kms/auth-eth/docs/formal-verification.md index 409330c5e..867edcc58 100644 --- a/dstack/kms/auth-eth/docs/formal-verification.md +++ b/dstack/kms/auth-eth/docs/formal-verification.md @@ -140,6 +140,9 @@ and admits one app per epoch, then one symbolic owner write where noted): - [x] `check_PolicyWrite_PreservesOtherEpochDecisions` — a write never changes an authorization decision at an epoch outside the rows it names (one epoch, or a symbolic range), including epoch 0's legacy fields. +- [x] `check_TeeVariantPolicy_NarrowsDecision` — once the TEE platform + policy is restricted, `authorizeApp` allows exactly when it would + unrestricted and the request names a listed platform. - [x] `check_AuthorizeApp_MatchesPolicy` — for both app operations, `authorizeApp` allows exactly when the epoch is active, the TCB requirement holds, that epoch admits the compose hash, and the device diff --git a/dstack/kms/auth-eth/index.test.ts b/dstack/kms/auth-eth/index.test.ts index 6eff517e3..7e0e2fee2 100644 --- a/dstack/kms/auth-eth/index.test.ts +++ b/dstack/kms/auth-eth/index.test.ts @@ -879,7 +879,7 @@ describe('GET /app/:appId', () => { describe('POST /authorizeApp', () => { const appInfo = { - appId: '0x' + '11'.repeat(20), composeHash: '0x' + '22'.repeat(32), + teeVariant: 'dstack-amd-sev-snp', appId: '0x' + '11'.repeat(20), composeHash: '0x' + '22'.repeat(32), instanceId: '', deviceId: '0x' + '33'.repeat(32), mrAggregated: '0x' + '44'.repeat(32), mrSystem: '0x' + '55'.repeat(32), osImageHash: '0x' + '66'.repeat(32), tcbStatus: 'UpToDate', advisoryIds: [], @@ -901,7 +901,7 @@ describe('POST /authorizeApp', () => { expect(mockReadContract).toHaveBeenCalledTimes(1); expect(mockReadContract).toHaveBeenCalledWith(expect.objectContaining({ functionName: 'authorizeApp', blockNumber: 98n, - args: [expect.objectContaining({ appInfo: expect.objectContaining({ appId: appInfo.appId }), requestedAppEpoch: { isSet: false, value: 0n }, requestedRootEpoch: { isSet: false, value: 0n } })], + args: [expect.objectContaining({ appInfo: expect.objectContaining({ appId: appInfo.appId }), requestedAppEpoch: { isSet: false, value: 0n }, requestedRootEpoch: { isSet: false, value: 0n }, teeVariant: appInfo.teeVariant })], })); }); @@ -954,6 +954,8 @@ describe('POST /authorizeApp', () => { const requests = [ { status: 400, body: JSON.stringify({ ...appInfo, operation: 'unknown' }) }, { status: 400, body: JSON.stringify({ ...appInfo, operation: 'app_get_key', mrAggregated: 'not-hex' }) }, + { status: 400, body: JSON.stringify({ ...appInfo, operation: 'app_get_key', teeVariant: undefined }) }, + { status: 400, body: JSON.stringify({ ...appInfo, operation: 'app_get_key', teeVariant: '' }) }, { status: 400, body: '{"appId":' }, { status: 413, body: JSON.stringify({ ...appInfo, operation: 'app_get_key', padding: 'a'.repeat(256 * 1024) }) }, ]; diff --git a/dstack/kms/auth-eth/index.ts b/dstack/kms/auth-eth/index.ts index 380f88805..164a03ec0 100644 --- a/dstack/kms/auth-eth/index.ts +++ b/dstack/kms/auth-eth/index.ts @@ -38,6 +38,7 @@ const BootInfoSchema = z.object({ }); const AuthorizationRequestSchema = BootInfoSchema.omit({ rootEpoch: true, rootEpochHash: true }).extend({ + teeVariant: z.string().min(1).max(64), operation: z.enum(['kms_bootstrap', 'kms_serve', 'kms_export_keys', 'kms_import_keys', 'app_get_key', 'app_get_cert']), requestedAppEpoch: z.number().int().min(0).max(0xffffffff).optional(), requestedRootEpoch: z.number().int().min(0).max(0xffffffff).optional(), @@ -103,6 +104,7 @@ const AUTHORIZATION_ABI = { { name: 'requestedRootEpoch', type: 'tuple', components: [ { name: 'isSet', type: 'bool' }, { name: 'value', type: 'uint256' }, ] }, + { name: 'teeVariant', type: 'string' }, ] }], outputs: [{ name: 'result', @@ -288,6 +290,7 @@ export class EthereumBackend { operation: keccak256(toBytes(appInfo.operation)), requestedAppEpoch: { isSet: appInfo.requestedAppEpoch !== undefined, value: BigInt(appInfo.requestedAppEpoch ?? 0) }, requestedRootEpoch: { isSet: appInfo.requestedRootEpoch !== undefined, value: BigInt(appInfo.requestedRootEpoch ?? 0) }, + teeVariant: appInfo.teeVariant, }], blockNumber, }); diff --git a/dstack/kms/auth-eth/openapi.json b/dstack/kms/auth-eth/openapi.json index c0af9430d..b1cb2b5c2 100644 --- a/dstack/kms/auth-eth/openapi.json +++ b/dstack/kms/auth-eth/openapi.json @@ -478,6 +478,12 @@ "maximum": 4294967295, "description": "Omit to use the current application epoch. Explicit zero selects epoch 0; future epochs are denied." }, + "teeVariant": { + "type": "string", + "minLength": 1, + "maxLength": 64, + "description": "TEE platform of the attested identity, as the KMS names it, e.g. dstack-tdx; passed to the contract unchanged." + }, "operation": { "type": "string", "enum": [ @@ -491,6 +497,7 @@ } }, "required": [ + "teeVariant", "operation" ] } diff --git a/dstack/kms/auth-eth/script/Manage.s.sol b/dstack/kms/auth-eth/script/Manage.s.sol index ff5c6bfd7..ac32439c5 100644 --- a/dstack/kms/auth-eth/script/Manage.s.sol +++ b/dstack/kms/auth-eth/script/Manage.s.sol @@ -56,6 +56,14 @@ abstract contract BaseScript is Script { return (vm.envUint(from), vm.envUint(to)); } + // TEE_VARIANTS: `any`, or the comma-separated platform names, as the KMS + // reports them, that alone are allowed. + function teeVariants() internal view returns (bool any, string[] memory names) { + names = vm.envString("TEE_VARIANTS", ","); + require(names.length > 0 && bytes(names[0]).length > 0, "set TEE_VARIANTS to any or platform names"); + any = names.length == 1 && keccak256(bytes(names[0])) == keccak256("any"); + } + function rootEpochs() internal view returns (uint256, uint256) { return epochRange("ROOT_EPOCH", "FROM_ROOT_EPOCH", "TO_ROOT_EPOCH"); } @@ -132,6 +140,20 @@ contract AddKmsDevice is BaseScript { } } +contract SetKmsTeeVariants is BaseScript { + function run() external { + (bool any, string[] memory names) = teeVariants(); + + vm.startBroadcast(vm.envUint("PRIVATE_KEY")); + // Write the list before restricting, so the restriction never applies to a stale one. + if (!any) kms.setAllowedTeeVariants(names); + kms.setAllowAnyTeeVariant(any); + vm.stopBroadcast(); + + console.log("Set KMS TEE variants to:", vm.envString("TEE_VARIANTS")); + } +} + contract SetGatewayAppId is BaseScript { function run() external { string memory gatewayId = vm.envString("GATEWAY_APP_ID"); @@ -251,6 +273,21 @@ contract SetAllowAnyDevice is BaseScript { } } +contract SetTeeVariants is BaseScript { + function run() external { + (bool any, string[] memory names) = teeVariants(); + require(address(legacyApp()) == address(0), "this app predates TEE variant policies"); + + vm.startBroadcast(vm.envUint("PRIVATE_KEY")); + // Write the list before restricting, so the restriction never applies to a stale one. + if (!any) app.setAllowedTeeVariants(names); + app.setAllowAnyTeeVariant(any); + vm.stopBroadcast(); + + console.log("Set TEE variants to:", vm.envString("TEE_VARIANTS")); + } +} + // Factory Deployment contract DeployApp is BaseScript { function run() external returns (address) { diff --git a/dstack/kms/auth-eth/script/Query.s.sol b/dstack/kms/auth-eth/script/Query.s.sol index 534b2de46..5c9c59b36 100644 --- a/dstack/kms/auth-eth/script/Query.s.sol +++ b/dstack/kms/auth-eth/script/Query.s.sol @@ -21,6 +21,13 @@ abstract contract BaseQueryScript is Script { DstackKms public kms; DstackApp public app; + function logTeeVariants(bool any, string[] memory names) internal pure { + console.log("Allow Any TEE Variant:", any); + for (uint256 i = 0; i < names.length; ++i) { + console.log(" Allowed TEE variant:", names[i]); + } + } + function setUp() public virtual { address kmsAddr = vm.envOr("KMS_CONTRACT_ADDR", address(0)); if (kmsAddr != address(0)) { @@ -202,6 +209,7 @@ contract GetKmsSettings is BaseQueryScript { console.log("Owner:", kms.owner()); console.log("Gateway App ID:", kms.gatewayAppId()); console.log("App Implementation:", kms.appImplementation()); + logTeeVariants(kms.allowAnyTeeVariant(), kms.allowedTeeVariants()); // Get KMS info (bytes memory k256Pubkey, bytes memory caPubkey, bytes memory quote, bytes memory eventlog) = kms.kmsInfo(); @@ -243,6 +251,7 @@ contract GetAppSettings is BaseQueryScript { console.log("Contract Address:", address(app)); console.log("Owner:", app.owner()); console.log("Allow Any Device:", app.allowAnyDevice()); + logTeeVariants(app.allowAnyTeeVariant(), app.allowedTeeVariants()); // Check implementation via storage bytes32 implSlot = 0x360894a13ba1a3210667c828492db98dca3e2076cc3735a920a3ca505d382bbc; diff --git a/dstack/kms/auth-eth/script/README.md b/dstack/kms/auth-eth/script/README.md index cb322fc19..9ace1fef4 100644 --- a/dstack/kms/auth-eth/script/README.md +++ b/dstack/kms/auth-eth/script/README.md @@ -51,6 +51,17 @@ KMS_CONTRACT_ADDR=0x... ROOT_EPOCH=0 DEVICE_ID=0x1234... \ forge script script/Manage.s.sol:AddKmsDevice --broadcast --rpc-url $RPC_URL ``` +#### Set KMS TEE Variants +`TEE_VARIANTS` is `any`, or the comma-separated platform names, as the KMS reports +them, that alone are allowed: `dstack-tdx`, `dstack-gcp-tdx`, `dstack-nitro-enclave`, +`dstack-amd-sev-snp` or `dstack-aws-nitro-tpm` (see `contracts/TeeVariantPolicy.sol`). +It applies to KMS instances and every app at every root epoch. A restricted policy +refuses the legacy `isKmsAllowed`/`isAppAllowed`, whose callers report no platform. +```bash +KMS_CONTRACT_ADDR=0x... TEE_VARIANTS=dstack-tdx \ +forge script script/Manage.s.sol:SetKmsTeeVariants --broadcast --rpc-url $RPC_URL +``` + #### Set Gateway App ID ```bash KMS_CONTRACT_ADDR=0x... GATEWAY_APP_ID="my-gateway" \ @@ -170,6 +181,13 @@ KMS_CONTRACT_ADDR=0x... APP_CONTRACT_ADDR=0x... ALLOW_ANY_DEVICE=true \ forge script script/Manage.s.sol:SetAllowAnyDevice --broadcast --rpc-url $RPC_URL ``` +#### Set TEE Variants +Same `TEE_VARIANTS` values as for the KMS; applies to every app epoch. +```bash +KMS_CONTRACT_ADDR=0x... APP_CONTRACT_ADDR=0x... TEE_VARIANTS=dstack-tdx,dstack-gcp-tdx \ +forge script script/Manage.s.sol:SetTeeVariants --broadcast --rpc-url $RPC_URL +``` + #### Remove Compose Hash ```bash KMS_CONTRACT_ADDR=0x... APP_CONTRACT_ADDR=0x... APP_EPOCH=0 COMPOSE_HASH=0x1234... \ diff --git a/dstack/kms/auth-eth/test/DstackApp.t.sol b/dstack/kms/auth-eth/test/DstackApp.t.sol index 0a7e0acce..30edd6bad 100644 --- a/dstack/kms/auth-eth/test/DstackApp.t.sol +++ b/dstack/kms/auth-eth/test/DstackApp.t.sol @@ -426,7 +426,7 @@ contract DstackAppTest is Test { pure returns (AuthorizationRequest memory) { - return AuthorizationRequest(info, operation, appEpoch, OptionalUint256(false, 0)); + return AuthorizationRequest(info, operation, appEpoch, OptionalUint256(false, 0), "dstack-tdx"); } function authorizeAt( @@ -442,6 +442,62 @@ contract DstackAppTest is Test { return (result.authorized, result.rejectionReason); } + function test_TeeVariantsAreUnrestrictedUntilRestrictedAndCoverEveryEpoch() public { + IAppAuth.AppInfo memory info = epochInfo(); + vm.startPrank(owner); + app.addComposeHashAtEpochs(0, 1, info.composeHash); + app.addDeviceAtEpochs(0, 1, info.deviceId); + vm.stopPrank(); + AuthorizationRequest memory request = appRequest(info, OptionalUint256(true, 0), keccak256("app_get_key")); + request.teeVariant = "dstack-amd-sev-snp"; + assertTrue(app.allowAnyTeeVariant()); + assertTrue(app.authorizeApp(request).authorized); + + string[] memory tdx = new string[](1); + tdx[0] = "dstack-tdx"; + vm.startPrank(owner); + app.setAllowedTeeVariants(tdx); + // The list applies only once the policy is restricted. + assertTrue(app.authorizeApp(request).authorized); + app.setAllowAnyTeeVariant(false); + vm.stopPrank(); + assertEq(app.allowedTeeVariants()[0], "dstack-tdx"); + assertEq(app.authorizeApp(request).rejectionReason, "TEE variant not allowed"); + (bool allowed, string memory reason) = app.isAppAllowed(info); + assertFalse(allowed); + assertEq(reason, "TEE variant not allowed"); + request.teeVariant = "dstack-tdx"; + assertTrue(app.authorizeApp(request).authorized); + + // The policy covers every app epoch. + vm.prank(owner); + app.advanceAppEpoch(0); + request.requestedAppEpoch = OptionalUint256(true, 1); + assertTrue(app.authorizeApp(request).authorized); + request.teeVariant = "dstack-amd-sev-snp"; + assertEq(app.authorizeApp(request).rejectionReason, "TEE variant not allowed"); + + // Only known platforms can be listed, and an empty list admits none. + string[] memory unknown = new string[](1); + unknown[0] = "dstack-unknown"; + vm.startPrank(owner); + vm.expectRevert("unknown TEE variant"); + app.setAllowedTeeVariants(unknown); + app.setAllowedTeeVariants(new string[](0)); + vm.stopPrank(); + assertEq(app.allowedTeeVariants().length, 0); + request.teeVariant = "dstack-tdx"; + assertEq(app.authorizeApp(request).rejectionReason, "TEE variant not allowed"); + + vm.startPrank(user); + bytes memory unauthorized = abi.encodeWithSelector(OwnableUpgradeable.OwnableUnauthorizedAccount.selector, user); + vm.expectRevert(unauthorized); + app.setAllowAnyTeeVariant(true); + vm.expectRevert(unauthorized); + app.setAllowedTeeVariants(tdx); + vm.stopPrank(); + } + function test_AuthorizeAppOwnsOperationScopeAndAppEpoch() public { IAppAuth.AppInfo memory info = epochInfo(); vm.startPrank(owner); diff --git a/dstack/kms/auth-eth/test/DstackAppEpoch.symbolic.t.sol b/dstack/kms/auth-eth/test/DstackAppEpoch.symbolic.t.sol index e9c644dfe..e5b8514cb 100644 --- a/dstack/kms/auth-eth/test/DstackAppEpoch.symbolic.t.sol +++ b/dstack/kms/auth-eth/test/DstackAppEpoch.symbolic.t.sol @@ -21,7 +21,7 @@ function _appAuth( returns (bool, string memory) { AuthorizationResult memory result = target.authorizeApp( - AuthorizationRequest(info, operation, OptionalUint256(true, epoch), OptionalUint256(false, 0)) + AuthorizationRequest(info, operation, OptionalUint256(true, epoch), OptionalUint256(false, 0), "dstack-tdx") ); return (result.authorized, result.rejectionReason); } @@ -281,6 +281,38 @@ contract DstackAppEpochSymbolicTest is Test { assert(allowedBefore == allowedAfter && reasonBefore == reasonAfter); } + // A TEE platform policy only ever narrows a decision: once restricted, a + // request passes exactly when it passes unrestricted and names a listed platform. + function check_TeeVariantPolicy_NarrowsDecision( + bool restrict, + bool listed, + bytes32 composeHash, + bytes32 deviceId, + bool upToDate, + uint256 epoch + ) + external + { + IAppAuth.AppInfo memory info = _info(composeHash, deviceId, upToDate); + (bool unrestricted,) = _appAuth(app, info, epoch, AuthorizationOperation.APP_GET_KEY); + string[] memory tdx = new string[](1); + tdx[0] = "dstack-tdx"; + vm.startPrank(OWNER); + app.setAllowedTeeVariants(tdx); + app.setAllowAnyTeeVariant(!restrict); + vm.stopPrank(); + AuthorizationResult memory result = app.authorizeApp( + AuthorizationRequest( + info, + AuthorizationOperation.APP_GET_KEY, + OptionalUint256(true, epoch), + OptionalUint256(false, 0), + listed ? "dstack-tdx" : "dstack-amd-sev-snp" + ) + ); + assert(result.authorized == (unrestricted && (!restrict || listed))); + } + function check_AuthorizeApp_MatchesPolicy( uint256 which, uint256 writeEpoch, diff --git a/dstack/kms/auth-eth/test/DstackKms.t.sol b/dstack/kms/auth-eth/test/DstackKms.t.sol index 0281634a8..31bb3ead3 100644 --- a/dstack/kms/auth-eth/test/DstackKms.t.sol +++ b/dstack/kms/auth-eth/test/DstackKms.t.sol @@ -497,11 +497,78 @@ contract DstackKmsTest is Test { info, isKms ? keccak256("kms_serve") : keccak256("app_get_key"), OptionalUint256(false, 0), - OptionalUint256(false, 0) + OptionalUint256(false, 0), + "dstack-tdx" ); return isKms ? kms.authorizeKms(request) : kms.authorizeApp(request); } + function test_TeeVariantsRestrictKmsAndApps() public { + IAppAuth.AppInfo memory info = authorizationInfo(address(kms)); + assertTrue(kms.allowAnyTeeVariant()); + string[] memory tdx = new string[](2); + tdx[0] = "dstack-tdx"; + tdx[1] = "dstack-gcp-tdx"; + vm.startPrank(owner); + kms.registerRootEpoch(ca(0xa0), k256(0xb0), "https://example.com/epoch-0"); + kms.addOsImageHashAtEpoch(0, info.osImageHash); + kms.addKmsDeviceAtEpoch(0, info.deviceId); + kms.addKmsAggregatedMrAtEpoch(0, info.mrAggregated); + address appId = kms.deployAndRegisterApp(user, false, true, info.deviceId, info.composeHash); + kms.setAllowedTeeVariants(tdx); + kms.setAllowAnyTeeVariant(false); + vm.stopPrank(); + + assertTrue(authorize(info, true).authorized); + (bool allowed, string memory reason) = kms.isKmsAllowed(info); + assertFalse(allowed); + assertEq(reason, "TEE variant not allowed"); + AuthorizationRequest memory request = AuthorizationRequest( + info, keccak256("kms_serve"), OptionalUint256(false, 0), OptionalUint256(false, 0), "dstack-tdx" + ); + request.teeVariant = "dstack-gcp-tdx"; + assertTrue(kms.authorizeKms(request).authorized); + request.teeVariant = "dstack-amd-sev-snp"; + assertEq(kms.authorizeKms(request).rejectionReason, "TEE variant not allowed"); + // A platform the contract does not know is refused once restricted. + request.teeVariant = "dstack-new-platform"; + assertEq(kms.authorizeKms(request).rejectionReason, "TEE variant not allowed"); + + info.appId = appId; + assertTrue(authorize(info, false).authorized); + (allowed, reason) = kms.isAppAllowed(info); + assertFalse(allowed); + assertEq(reason, "TEE variant not allowed"); + request = AuthorizationRequest( + info, keccak256("app_get_key"), OptionalUint256(false, 0), OptionalUint256(false, 0), "dstack-tdx" + ); + request.teeVariant = "dstack-amd-sev-snp"; + assertEq(kms.authorizeApp(request).rejectionReason, "TEE variant not allowed"); + + // Allowing any platform lifts the restriction, including for legacy callers. + vm.prank(owner); + kms.setAllowAnyTeeVariant(true); + assertTrue(kms.authorizeApp(request).authorized); + (allowed,) = kms.isAppAllowed(info); + assertTrue(allowed); + + // With the KMS unrestricted, the app's own policy still applies. + vm.startPrank(user); + DstackApp(appId).setAllowedTeeVariants(tdx); + DstackApp(appId).setAllowAnyTeeVariant(false); + vm.stopPrank(); + assertEq(kms.authorizeApp(request).rejectionReason, "TEE variant not allowed"); + (allowed, reason) = kms.isAppAllowed(info); + assertFalse(allowed); + assertEq(reason, "TEE variant not allowed"); + request.teeVariant = "dstack-gcp-tdx"; + assertTrue(kms.authorizeApp(request).authorized); + + vm.prank(user); + vm.expectRevert(abi.encodeWithSelector(OwnableUpgradeable.OwnableUnauthorizedAccount.selector, user)); + kms.setAllowAnyTeeVariant(false); + } + function test_AuthorizationOperationsFailClosedAndBootstrapBeforeRegistration() public { IAppAuth.AppInfo memory info = authorizationInfo(address(kms)); vm.startPrank(owner); @@ -510,7 +577,7 @@ contract DstackKmsTest is Test { kms.addKmsAggregatedMrAtEpoch(0, info.mrAggregated); vm.stopPrank(); AuthorizationRequest memory request = AuthorizationRequest( - info, keccak256("kms_bootstrap"), OptionalUint256(false, 0), OptionalUint256(false, 0) + info, keccak256("kms_bootstrap"), OptionalUint256(false, 0), OptionalUint256(false, 0), "dstack-tdx" ); assertTrue(kms.authorizeKms(request).authorized); assertEq(kms.authorizeKms(request).currentRootEpochHash, bytes32(0)); @@ -592,8 +659,9 @@ contract DstackKmsTest is Test { result = authorize(info, true); assertFalse(result.authorized); assertEq(result.rejectionReason, "Aggregated MR not allowed"); - AuthorizationRequest memory appRequest = - AuthorizationRequest(info, keccak256("app_get_key"), OptionalUint256(true, 1), OptionalUint256(false, 0)); + AuthorizationRequest memory appRequest = AuthorizationRequest( + info, keccak256("app_get_key"), OptionalUint256(true, 1), OptionalUint256(false, 0), "dstack-tdx" + ); result = kms.authorizeApp(appRequest); assertFalse(result.authorized); assertEq(result.rejectionReason, "Requested app epoch is not active"); @@ -678,7 +746,7 @@ contract DstackKmsTest is Test { (uint256 rootEpoch,) = kms.currentRootEpoch(); assertEq(rootEpoch, 0); AuthorizationRequest memory request = AuthorizationRequest( - info, keccak256("app_get_key"), OptionalUint256(false, 0), OptionalUint256(false, 0) + info, keccak256("app_get_key"), OptionalUint256(false, 0), OptionalUint256(false, 0), "dstack-tdx" ); AuthorizationResult memory result = kms.authorizeApp(request); assertFalse(result.authorized); @@ -716,7 +784,11 @@ contract DstackKmsTest is Test { assertEq(result.currentRootEpochHash, hash1); AuthorizationRequest memory request = AuthorizationRequest( - authorizationInfo(user), keccak256("kms_serve"), OptionalUint256(false, 0), OptionalUint256(false, 0) + authorizationInfo(user), + keccak256("kms_serve"), + OptionalUint256(false, 0), + OptionalUint256(false, 0), + "dstack-tdx" ); result = kms.authorizeApp(request); assertEq(result.rejectionReason, "Invalid authorization operation"); @@ -950,7 +1022,7 @@ contract DstackKmsTest is Test { returns (AuthorizationRequest memory) { return AuthorizationRequest( - info, keccak256("app_get_key"), OptionalUint256(true, appEpoch), OptionalUint256(true, root) + info, keccak256("app_get_key"), OptionalUint256(true, appEpoch), OptionalUint256(true, root), "dstack-tdx" ); } diff --git a/dstack/kms/auth-eth/test/DstackKmsEpoch.symbolic.t.sol b/dstack/kms/auth-eth/test/DstackKmsEpoch.symbolic.t.sol index da39b1004..f7da3cc0e 100644 --- a/dstack/kms/auth-eth/test/DstackKmsEpoch.symbolic.t.sol +++ b/dstack/kms/auth-eth/test/DstackKmsEpoch.symbolic.t.sol @@ -22,7 +22,7 @@ function _appAuth( returns (bool, string memory) { AuthorizationResult memory result = target.authorizeApp( - AuthorizationRequest(info, operation, OptionalUint256(true, epoch), OptionalUint256(false, 0)) + AuthorizationRequest(info, operation, OptionalUint256(true, epoch), OptionalUint256(false, 0), "dstack-tdx") ); return (result.authorized, result.rejectionReason); } @@ -134,7 +134,9 @@ contract DstackKmsEpochSymbolicTest is Test { pure returns (AuthorizationRequest memory) { - return AuthorizationRequest(info, operation, OptionalUint256(appSet, appEpoch), OptionalUint256(rootSet, root)); + return AuthorizationRequest( + info, operation, OptionalUint256(appSet, appEpoch), OptionalUint256(rootSet, root), "dstack-tdx" + ); } /// Calldata of one KMS policy setter: 0-5 write the row of `epoch`, 6-11 diff --git a/dstack/kms/auth-eth/test/LegacyAppCompat.t.sol b/dstack/kms/auth-eth/test/LegacyAppCompat.t.sol index 59750a092..e083b4fa3 100644 --- a/dstack/kms/auth-eth/test/LegacyAppCompat.t.sol +++ b/dstack/kms/auth-eth/test/LegacyAppCompat.t.sol @@ -76,7 +76,7 @@ contract LegacyAppCompatTest is Test { info.tcbStatus = "UpToDate"; info.advisoryIds = new string[](0); return AuthorizationRequest( - info, keccak256("app_get_key"), OptionalUint256(appSet, appEpoch), OptionalUint256(false, 0) + info, keccak256("app_get_key"), OptionalUint256(appSet, appEpoch), OptionalUint256(false, 0), "dstack-tdx" ); } @@ -118,6 +118,22 @@ contract LegacyAppCompatTest is Test { } } + function test_RestrictedKmsGatesLegacyAppsByTeeVariant() public { + string[] memory tdx = new string[](1); + tdx[0] = "dstack-tdx"; + vm.startPrank(owner); + kms.setAllowedTeeVariants(tdx); + kms.setAllowAnyTeeVariant(false); + vm.stopPrank(); + for (uint256 release; release < appReleases.length; ++release) { + address app = deployLegacyApp(release); + assertTrue(kms.authorizeApp(request(app, false, 0)).authorized); + AuthorizationRequest memory snp = request(app, false, 0); + snp.teeVariant = "dstack-amd-sev-snp"; + assertEq(kms.authorizeApp(snp).rejectionReason, "TEE variant not allowed"); + } + } + function test_InterfaceIdsTellLegacyAndEpochAppsApart() public { bytes4 basicManagement = 0x8fd37527; for (uint256 release; release < appReleases.length; ++release) { @@ -172,6 +188,12 @@ contract LegacyAppCompatTest is Test { assertTrue(upgraded.allowedComposeHashesAtEpoch(0, COMPOSE)); assertTrue(upgraded.allowedDeviceIdsAtEpoch(0, DEVICE)); assertEq(upgraded.allowAnyDevice(), anyDevice); + // Upgraded policies allow any TEE platform. + assertTrue(upgraded.allowAnyTeeVariant()); + assertEq(upgraded.allowedTeeVariants().length, 0); + AuthorizationRequest memory snp = request(app, true, 0); + snp.teeVariant = "dstack-amd-sev-snp"; + assertTrue(kms.authorizeApp(snp).authorized); assertFalse(upgraded.allowedComposeHashesAtEpoch(1, COMPOSE)); assertTrue(kms.authorizeApp(request(app, true, 0)).authorized); diff --git a/dstack/kms/auth-eth/test/ManageScript.t.sol b/dstack/kms/auth-eth/test/ManageScript.t.sol index 503e3a878..1b19bb0fd 100644 --- a/dstack/kms/auth-eth/test/ManageScript.t.sol +++ b/dstack/kms/auth-eth/test/ManageScript.t.sol @@ -51,5 +51,31 @@ contract ManageScriptTest is Test { script = addComposeHash(legacy); vm.expectRevert(bytes("this app predates app epochs: only APP_EPOCH=0 applies")); script.run(); + + // TEE variant policies take `any` or names, and exist only in current apps. + vm.setEnv("APP_EPOCH", "0"); + vm.setEnv("APP_CONTRACT_ADDR", vm.toString(current)); + vm.setEnv("TEE_VARIANTS", "dstack-tdx,dstack-gcp-tdx"); + SetTeeVariants teeScript = new SetTeeVariants(); + teeScript.setUp(); + teeScript.run(); + assertFalse(DstackApp(current).allowAnyTeeVariant()); + string[] memory names = DstackApp(current).allowedTeeVariants(); + assertEq(names.length, 2); + assertEq(names[1], "dstack-gcp-tdx"); + vm.setEnv("TEE_VARIANTS", "any"); + teeScript.run(); + assertTrue(DstackApp(current).allowAnyTeeVariant()); + // An unset or empty TEE_VARIANTS must not lock everything out. + vm.setEnv("TEE_VARIANTS", ""); + vm.expectRevert(bytes("set TEE_VARIANTS to any or platform names")); + teeScript.run(); + assertTrue(DstackApp(current).allowAnyTeeVariant()); + vm.setEnv("TEE_VARIANTS", "any"); + vm.setEnv("APP_CONTRACT_ADDR", vm.toString(legacy)); + teeScript = new SetTeeVariants(); + teeScript.setUp(); + vm.expectRevert(bytes("this app predates TEE variant policies")); + teeScript.run(); } }