diff --git a/examples/DualDeriverPolicy.sol b/examples/DualDeriverPolicy.sol new file mode 100644 index 0000000..066e947 --- /dev/null +++ b/examples/DualDeriverPolicy.sol @@ -0,0 +1,106 @@ +// SPDX-License-Identifier: MIT +pragma solidity 0.8.28; + +import {BasePolicy} from "../src/BasePolicy.sol"; +import {IFlashtestationRegistry} from "../src/interfaces/IFlashtestationRegistry.sol"; +import {IWorkloadDeriver} from "../src/interfaces/IWorkloadDeriver.sol"; +import {WorkloadId} from "../src/interfaces/IPolicyCommon.sol"; +import {HEADER_LENGTH} from "automata-dcap-attestation/contracts/types/Constants.sol"; + +/// @notice Example policy that supports two workload derivation strategies during a migration. +/// @dev Demonstrates how to inherit `BasePolicy` and override `isAllowedPolicy` to accept +/// workloads derived by either an "old" deriver or a "new" deriver. +/// +/// WARNING: This contract is an UNAUDITED EXAMPLE and is NOT intended for production use. +/// It is provided to illustrate one possible migration pattern; do not deploy without +/// a thorough security review +/// +/// Quote format detection: +/// - You *can* sometimes distinguish formats by length (e.g. TD10 vs a hypothetical TD15 report body), +/// but real quote payload sizes can vary and some derivers accept both "raw report body" and "quote-like" blobs. +/// - This example uses a light length-based hint to decide which deriver to try first, but still falls back to +/// trying both with `try/catch` to remain robust. +contract DualDeriverPolicy is BasePolicy { + address public immutable OWNER; + IWorkloadDeriver public immutable OLD_DERIVER; + IWorkloadDeriver public immutable NEW_DERIVER; + + // For the example TD15 report-body format: 584 (TD10) + 16 + 48. + uint256 internal constant TD_REPORT15_LENGTH = 648; + + constructor(address owner_, address registry_, IWorkloadDeriver oldDeriver_, IWorkloadDeriver newDeriver_) { + OWNER = owner_; + OLD_DERIVER = oldDeriver_; + NEW_DERIVER = newDeriver_; + _basePolicyInit(registry_); + } + + function _checkPolicyAuthority() internal view override { + require(msg.sender == OWNER, "NotOwner"); + } + + // Not used by this example (we override `isAllowedPolicy`), but required by BasePolicy. + function _workloadDeriver() internal view override returns (IWorkloadDeriver) { + return NEW_DERIVER; + } + + // This example policy does not use caching, so the cache hooks are implemented as no-ops. + function _getCachedWorkload(address) internal pure override returns (CachedWorkload memory) { + return CachedWorkload({workloadId: WorkloadId.wrap(0), quoteHash: bytes32(0)}); + } + + function _setCachedWorkload(address, CachedWorkload memory) internal override {} + + /// @inheritdoc BasePolicy + function isAllowedPolicy(address teeAddress) + public + view + override + returns (bool allowed, WorkloadId) + { + (, IFlashtestationRegistry.RegisteredTEE memory registration) = + IFlashtestationRegistry(registry).getRegistration(teeAddress); + if (!registration.isValid) return (false, WorkloadId.wrap(0)); + + bytes memory rawQuote = registration.rawQuote; + + // Hint: if it looks like a TD15 report body or a quote containing one, try NEW first. + bool preferNew = rawQuote.length == TD_REPORT15_LENGTH || rawQuote.length == HEADER_LENGTH + TD_REPORT15_LENGTH; + + (bool okNew, WorkloadId idNew) = (false, WorkloadId.wrap(0)); + (bool okOld, WorkloadId idOld) = (false, WorkloadId.wrap(0)); + + if (preferNew) { + try NEW_DERIVER.workloadIdForQuote(rawQuote) returns (WorkloadId id) { + okNew = true; + idNew = id; + } catch {} + + try OLD_DERIVER.workloadIdForQuote(rawQuote) returns (WorkloadId id) { + okOld = true; + idOld = id; + } catch {} + } else { + try OLD_DERIVER.workloadIdForQuote(rawQuote) returns (WorkloadId id) { + okOld = true; + idOld = id; + } catch {} + + try NEW_DERIVER.workloadIdForQuote(rawQuote) returns (WorkloadId id) { + okNew = true; + idNew = id; + } catch {} + } + + // Prefer NEW workload IDs during migration if both are approved. + if (okNew && bytes(approvedWorkloads[WorkloadId.unwrap(idNew)].commitHash).length > 0) { + return (true, idNew); + } + if (okOld && bytes(approvedWorkloads[WorkloadId.unwrap(idOld)].commitHash).length > 0) { + return (true, idOld); + } + + return (false, WorkloadId.wrap(0)); + } +} + diff --git a/examples/TDXTD15WorkloadDeriver.sol b/examples/TDXTD15WorkloadDeriver.sol new file mode 100644 index 0000000..0b44fd6 --- /dev/null +++ b/examples/TDXTD15WorkloadDeriver.sol @@ -0,0 +1,123 @@ +// SPDX-License-Identifier: MIT +pragma solidity 0.8.28; + +// NOTE: The parsing approach and offsets here are taken from Automata's DCAP attestation +// https://github.com/automata-network/automata-dcap-attestation/tree/main + +import {BytesUtils} from "@automata-network/on-chain-pccs/utils/BytesUtils.sol"; +import {HEADER_LENGTH} from "automata-dcap-attestation/contracts/types/Constants.sol"; + +import {IWorkloadDeriver} from "../src/interfaces/IWorkloadDeriver.sol"; +import {WorkloadId} from "../src/interfaces/IPolicyCommon.sol"; + +/// @notice "TDX 1.5" report body example with two additional fields. +/// @dev This is NOT a canonical flashtestations type; it is an example format to demonstrate how to swap derivation logic. +struct TD15ReportBody { + bytes16 teeTcbSvn; + bytes mrSeam; // 48 bytes + bytes mrsignerSeam; // 48 bytes + bytes8 seamAttributes; + bytes8 tdAttributes; + bytes8 xFAM; + bytes mrTd; // 48 bytes + bytes mrConfigId; // 48 bytes + bytes mrOwner; // 48 bytes + bytes mrOwnerConfig; // 48 bytes + bytes rtMr0; // 48 bytes + bytes rtMr1; // 48 bytes + bytes rtMr2; // 48 bytes + bytes rtMr3; // 48 bytes + bytes reportData; // 64 bytes + bytes16 teeTcbSvn2; + bytes mrServiceTd; // 48 bytes +} + +library TD15ReportParser { + using BytesUtils for bytes; + + // 584-byte TD10 report body + 16 + 48 extra bytes. + uint256 internal constant TD_REPORT15_LENGTH = 648; + + function parse(bytes memory reportBytes) internal pure returns (bool success, TD15ReportBody memory report) { + success = reportBytes.length == TD_REPORT15_LENGTH; + if (!success) return (false, report); + + report.teeTcbSvn = bytes16(reportBytes.substring(0, 16)); + report.mrSeam = reportBytes.substring(16, 48); + report.mrsignerSeam = reportBytes.substring(64, 48); + report.seamAttributes = bytes8(reportBytes.substring(112, 8)); + report.tdAttributes = bytes8(reportBytes.substring(120, 8)); + report.xFAM = bytes8(reportBytes.substring(128, 8)); + report.mrTd = reportBytes.substring(136, 48); + report.mrConfigId = reportBytes.substring(184, 48); + report.mrOwner = reportBytes.substring(232, 48); + report.mrOwnerConfig = reportBytes.substring(280, 48); + report.rtMr0 = reportBytes.substring(328, 48); + report.rtMr1 = reportBytes.substring(376, 48); + report.rtMr2 = reportBytes.substring(424, 48); + report.rtMr3 = reportBytes.substring(472, 48); + report.reportData = reportBytes.substring(520, 64); + report.teeTcbSvn2 = bytes16(reportBytes.substring(584, 16)); + report.mrServiceTd = reportBytes.substring(600, 48); + } +} + +/// @notice Example deriver that expects a TD15 report body and hashes the two additional fields. +contract TDXTD15WorkloadDeriver is IWorkloadDeriver { + using BytesUtils for bytes; + + // Same constants as the current TDX deriver. + bytes8 internal constant TD_XFAM_FPU = 0x0000000000000001; + bytes8 internal constant TD_XFAM_SSE = 0x0000000000000002; + bytes8 internal constant TD_TDATTRS_VE_DISABLED = 0x0000000010000000; + bytes8 internal constant TD_TDATTRS_PKS = 0x0000000040000000; + bytes8 internal constant TD_TDATTRS_KL = 0x0000000080000000; + + error InvalidTD15ReportLength(uint256 length); + + function workloadIdForReportBody(TD15ReportBody memory reportBody) public pure returns (WorkloadId) { + bytes8 expectedXfamBits = TD_XFAM_FPU | TD_XFAM_SSE; + bytes8 ignoredTdAttributesBitmask = TD_TDATTRS_VE_DISABLED | TD_TDATTRS_PKS | TD_TDATTRS_KL; + + return WorkloadId.wrap( + keccak256( + bytes.concat( + reportBody.mrTd, + reportBody.rtMr0, + reportBody.rtMr1, + reportBody.rtMr2, + reportBody.rtMr3, + // VMM configuration + reportBody.mrConfigId, + reportBody.xFAM ^ expectedXfamBits, + reportBody.tdAttributes & ~ignoredTdAttributesBitmask, + // TD15 extensions + bytes16(reportBody.teeTcbSvn2), + reportBody.mrServiceTd + ) + ) + ); + } + + /// @inheritdoc IWorkloadDeriver + /// @dev Accepts either: + /// - the raw TD15 report body bytes (length == TD_REPORT15_LENGTH), or + /// - a quote-like blob where the report body starts at HEADER_LENGTH. + function workloadIdForQuote(bytes calldata rawQuote) external pure returns (WorkloadId) { + bytes memory raw = rawQuote; + + // Try raw report body first. + (bool ok, TD15ReportBody memory report) = TD15ReportParser.parse(raw); + if (ok) return workloadIdForReportBody(report); + + // Try treating `rawQuote` as a quote with a header prefix. + if (raw.length >= HEADER_LENGTH + TD15ReportParser.TD_REPORT15_LENGTH) { + bytes memory reportBytes = raw.substring(HEADER_LENGTH, TD15ReportParser.TD_REPORT15_LENGTH); + (ok, report) = TD15ReportParser.parse(reportBytes); + if (ok) return workloadIdForReportBody(report); + } + + revert InvalidTD15ReportLength(raw.length); + } +} + diff --git a/script/BlockBuilderPolicy.s.sol b/script/BlockBuilderPolicy.s.sol index 429e3ed..34aaa53 100644 --- a/script/BlockBuilderPolicy.s.sol +++ b/script/BlockBuilderPolicy.s.sol @@ -5,6 +5,7 @@ import {Script, console} from "forge-std/Script.sol"; import {Upgrades} from "openzeppelin-foundry-upgrades/Upgrades.sol"; import {BlockBuilderPolicy} from "../src/BlockBuilderPolicy.sol"; import {FlashtestationRegistry} from "../src/FlashtestationRegistry.sol"; +import {TDXWorkloadDeriver} from "../src/derivers/TDXWorkloadDeriver.sol"; /// @title BlockBuilderPolicyScript /// @notice Deploy the block builder policy contract, which is a simple contract that allows an organization @@ -31,8 +32,11 @@ contract BlockBuilderPolicyScript is Script { address owner = vm.envAddress("OWNER_BLOCK_BUILDER_POLICY"); console.log("owner", owner); + address deriver = address(new TDXWorkloadDeriver()); + console.log("TDXWorkloadDeriver deployed at:", deriver); + address policy = Upgrades.deployUUPSProxy( - "BlockBuilderPolicy.sol", abi.encodeCall(BlockBuilderPolicy.initialize, (owner, registry)) + "BlockBuilderPolicy.sol", abi.encodeCall(BlockBuilderPolicy.initialize, (owner, registry, deriver)) ); console.log("BlockBuilderPolicy deployed at:", policy); vm.stopBroadcast(); diff --git a/script/Interactions.s.sol b/script/Interactions.s.sol index 56682e6..f4a5595 100644 --- a/script/Interactions.s.sol +++ b/script/Interactions.s.sol @@ -3,7 +3,7 @@ pragma solidity 0.8.28; import {Script, console} from "forge-std/Script.sol"; import {BlockBuilderPolicy} from "../src/BlockBuilderPolicy.sol"; -import {WorkloadId} from "../src/interfaces/IBlockBuilderPolicy.sol"; +import {WorkloadId} from "../src/interfaces/IPolicyCommon.sol"; import {FlashtestationRegistry} from "../src/FlashtestationRegistry.sol"; import {IFlashtestationRegistry} from "../src/interfaces/IFlashtestationRegistry.sol"; import {DeploymentUtils} from "./utils/DeploymentUtils.sol"; @@ -112,7 +112,10 @@ contract RegisterTEEScript is Script, DeploymentUtils { console.logAddress(registryAddress); FlashtestationRegistry registry = FlashtestationRegistry(registryAddress); - registry.registerTEEService(vm.readFileBinary(pathToAttestationQuote), bytes("") /* currently not used */ ); + registry.registerTEEService( + vm.readFileBinary(pathToAttestationQuote), + bytes("") /* currently not used */ + ); // fetch the TEE-related data we just added, so the caller of this script can use // the outputs in future scripts (like Interactions.s.sol:AddWorkloadToPolicyScript) diff --git a/src/BasePolicy.sol b/src/BasePolicy.sol new file mode 100644 index 0000000..7e6b26d --- /dev/null +++ b/src/BasePolicy.sol @@ -0,0 +1,186 @@ +// SPDX-License-Identifier: MIT +pragma solidity 0.8.28; + +import {IBasePolicy} from "./interfaces/IBasePolicy.sol"; +import {IFlashtestationRegistry} from "./interfaces/IFlashtestationRegistry.sol"; +import {IWorkloadDeriver} from "./interfaces/IWorkloadDeriver.sol"; +import {WorkloadId} from "./interfaces/IPolicyCommon.sol"; + +/// @notice Shared policy implementation holding common workload/registry logic. +/// @dev IMPORTANT: This contract is designed to be inherited by upgradeable policies without +/// breaking `BlockBuilderPolicy` storage layout. It MUST NOT introduce additional storage +/// beyond `approvedWorkloads` and `registry` (in that order). +abstract contract BasePolicy is IBasePolicy { + // ============ Types ============ + + /** + * @notice Cached workload information for gas optimization. + * @dev Stored in the derived contract; BasePolicy only defines the type + logic hooks. + */ + struct CachedWorkload { + WorkloadId workloadId; + bytes32 quoteHash; + } + + // ============ Storage (DO NOT REORDER / ADD) ============ + + /// @notice Mapping from workloadId to its metadata (commit hash and source locators). + mapping(bytes32 workloadId => WorkloadMetadata) internal approvedWorkloads; + + /// @inheritdoc IBasePolicy + address public override registry; + + // ============ Initialization ============ + + /// @dev Shared initializer helper for derived contracts. + function _basePolicyInit(address _registry) internal { + if (_registry == address(0)) revert InvalidRegistry(); + registry = _registry; + emit RegistrySet(_registry); + } + + // ============ Required hooks (implemented by derived policy) ============ + + /// @dev Access-control hook (typically maps to `OwnableUpgradeable._checkOwner()`). + /// @dev `BasePolicy` is intentionally not `OwnableUpgradeable` to preserve `BlockBuilderPolicy`'s + /// storage layout and keep the authorization mechanism flexible for downstream policies. + function _checkPolicyAuthority() internal view virtual; + + /// @dev Workload deriver hook + /// @return The configured workload deriver (or address(0) if not configured). + function _workloadDeriver() internal view virtual returns (IWorkloadDeriver); + + /// @dev Cache read hook. + function _getCachedWorkload(address teeAddress) internal view virtual returns (CachedWorkload memory); + + /// @dev Cache write hook. + function _setCachedWorkload(address teeAddress, CachedWorkload memory cached) internal virtual; + + // ============ Common policy logic ============ + + /// @inheritdoc IBasePolicy + function addWorkloadToPolicy(WorkloadId workloadId, string calldata commitHash, string[] calldata sourceLocators) + external + virtual + override + { + _checkPolicyAuthority(); + + require(bytes(commitHash).length > 0, EmptyCommitHash()); + require(sourceLocators.length > 0, EmptySourceLocators()); + + bytes32 workloadKey = WorkloadId.unwrap(workloadId); + + // Check if workload already exists + require(bytes(approvedWorkloads[workloadKey].commitHash).length == 0, WorkloadAlreadyInPolicy()); + + // Store the workload metadata + approvedWorkloads[workloadKey] = WorkloadMetadata({commitHash: commitHash, sourceLocators: sourceLocators}); + + emit WorkloadAddedToPolicy(workloadKey); + } + + /// @inheritdoc IBasePolicy + function removeWorkloadFromPolicy(WorkloadId workloadId) external virtual override { + _checkPolicyAuthority(); + + bytes32 workloadKey = WorkloadId.unwrap(workloadId); + + // Check if workload exists + require(bytes(approvedWorkloads[workloadKey].commitHash).length > 0, WorkloadNotInPolicy()); + + // Remove the workload metadata + delete approvedWorkloads[workloadKey]; + + emit WorkloadRemovedFromPolicy(workloadKey); + } + + /// @inheritdoc IBasePolicy + function getWorkloadMetadata(WorkloadId workloadId) external view override returns (WorkloadMetadata memory) { + return approvedWorkloads[WorkloadId.unwrap(workloadId)]; + } + + /// @inheritdoc IBasePolicy + function isAllowedPolicy(address teeAddress) public view virtual override returns (bool allowed, WorkloadId) { + // Get full registration data + (, IFlashtestationRegistry.RegisteredTEE memory registration) = + IFlashtestationRegistry(registry).getRegistration(teeAddress); + + // Invalid Registrations means the attestation used to register the TEE is no longer valid + // and so we cannot trust any input from the TEE. + if (!registration.isValid) { + return (false, WorkloadId.wrap(0)); + } + + IWorkloadDeriver deriver = _workloadDeriver(); + if (address(deriver) == address(0)) { + // Treat missing deriver as "not allowed" to avoid bricking callers during upgrades. + return (false, WorkloadId.wrap(0)); + } + WorkloadId workloadId = deriver.workloadIdForQuote(registration.rawQuote); + + // Check if the workload exists in our approved workloads mapping + if (bytes(approvedWorkloads[WorkloadId.unwrap(workloadId)].commitHash).length > 0) { + return (true, workloadId); + } + + return (false, WorkloadId.wrap(0)); + } + + /// @notice Cached variant of `isAllowedPolicy` for gas-sensitive call paths. + /// @dev Why this exists: + /// - `isAllowedPolicy` is `view` but can be expensive because workloadId derivation can be costly. + /// - Some policies (notably `BlockBuilderPolicy`) need an O(1)-ish authorization check on the hot path + /// (e.g. `verifyBlockBuilderProof`), so we cache the derived workloadId keyed by the TEE's quoteHash. + /// + /// @dev How it is expected to be used: + /// - Derived policies call `_cachedIsAllowedPolicy(teeAddress)` from their hot path. + /// - Cache storage lives in the derived policy; `BasePolicy` only implements the algorithm and calls the + /// `_getCachedWorkload` / `_setCachedWorkload` hooks. + /// - The function checks registry validity + quoteHash; on cache hit it avoids recomputing derivation. + /// - On cache miss (or quote change), it falls back to `isAllowedPolicy` and updates the cache if allowed. + /// + /// @dev Cache cleanup: + /// - This function does not proactively delete stale cache entries. On invalid registrations it returns + /// `(false, 0)` and callers typically revert, so cleanup is intentionally skipped to keep the hot path cheap. + /// + /// @param teeAddress The TEE-controlled address to check. + /// @return allowed True if the TEE is registered, valid, and running an approved workload. + /// @return workloadId The approved workloadId for this TEE (or 0 if not allowed). + function _cachedIsAllowedPolicy(address teeAddress) internal returns (bool, WorkloadId) { + // Get the current registration status (fast path) + (bool isValid, bytes32 quoteHash) = IFlashtestationRegistry(registry).getRegistrationStatus(teeAddress); + if (!isValid) { + return (false, WorkloadId.wrap(0)); + } + + // Now, check if we have a cached workload for this TEE + CachedWorkload memory cached = _getCachedWorkload(teeAddress); + bytes32 cachedWorkloadKey = WorkloadId.unwrap(cached.workloadId); + + // Check if we've already fetched and computed the workloadId for this TEE + if (cachedWorkloadKey != 0 && cached.quoteHash == quoteHash) { + // Cache hit - verify the workload is still a part of this policy's approved workloads + if (bytes(approvedWorkloads[cachedWorkloadKey].commitHash).length > 0) { + return (true, cached.workloadId); + } + // The workload is no longer approved, so the policy is no longer valid for this TEE + return (false, WorkloadId.wrap(0)); + } else { + // Cache miss or quote changed - use the view function to get the result + // + // Correctness note: even if a downstream policy "disables" caching by implementing the cache hooks as + // no-ops (e.g. always returning zero and never persisting), this function still behaves correctly: + // it simply falls back to the canonical `isAllowedPolicy` check every time. + (bool allowed, WorkloadId workloadId) = isAllowedPolicy(teeAddress); + + if (allowed) { + // Update cache with the new workload ID + _setCachedWorkload(teeAddress, CachedWorkload({workloadId: workloadId, quoteHash: quoteHash})); + } + + return (allowed, workloadId); + } + } +} + diff --git a/src/BlockBuilderPolicy.sol b/src/BlockBuilderPolicy.sol index 766575f..145d12a 100644 --- a/src/BlockBuilderPolicy.sol +++ b/src/BlockBuilderPolicy.sol @@ -6,20 +6,14 @@ import "@openzeppelin/contracts-upgradeable/proxy/utils/UUPSUpgradeable.sol"; import "@openzeppelin/contracts-upgradeable/access/OwnableUpgradeable.sol"; import {EIP712Upgradeable} from "@openzeppelin/contracts-upgradeable/utils/cryptography/EIP712Upgradeable.sol"; import {ECDSA} from "@openzeppelin/contracts/utils/cryptography/ECDSA.sol"; -import {FlashtestationRegistry} from "./FlashtestationRegistry.sol"; import {IFlashtestationRegistry} from "./interfaces/IFlashtestationRegistry.sol"; -import {IBlockBuilderPolicy, WorkloadId} from "./interfaces/IBlockBuilderPolicy.sol"; - -/** - * @notice Cached workload information for gas optimization - * @dev Stores computed workloadId and associated quoteHash to avoid expensive recomputation - */ -struct CachedWorkload { - /// @notice The computed workload identifier - WorkloadId workloadId; - /// @notice The keccak256 hash of the raw quote used to compute this workloadId - bytes32 quoteHash; -} +import {IBasePolicy} from "./interfaces/IBasePolicy.sol"; +import {IBlockBuilderPolicy} from "./interfaces/IBlockBuilderPolicy.sol"; +import {IWorkloadDeriver} from "./interfaces/IWorkloadDeriver.sol"; +import {WorkloadId} from "./interfaces/IPolicyCommon.sol"; +import {TD10ReportBody} from "automata-dcap-attestation/contracts/types/V4Structs.sol"; +import {BasePolicy} from "./BasePolicy.sol"; +import {TDXWorkloadDeriver} from "./derivers/TDXWorkloadDeriver.sol"; /** * @title BlockBuilderPolicy @@ -33,6 +27,7 @@ struct CachedWorkload { * is allowed under any workload in a Policy, and the FlashtestationRegistry will handle the rest */ contract BlockBuilderPolicy is + BasePolicy, Initializable, UUPSUpgradeable, OwnableUpgradeable, @@ -47,34 +42,8 @@ contract BlockBuilderPolicy is bytes32 public constant VERIFY_BLOCK_BUILDER_PROOF_TYPEHASH = keccak256("VerifyBlockBuilderProof(uint8 version,bytes32 blockContentHash,uint256 nonce)"); - // ============ TDX workload constants ============ - - /// @dev See section 11.5.3 in TDX Module v1.5 Base Architecture Specification https://www.intel.com/content/www/us/en/content-details/733575/intel-tdx-module-v1-5-base-architecture-specification.html - /// @notice Enabled FPU (always enabled) - bytes8 constant TD_XFAM_FPU = 0x0000000000000001; - /// @notice Enabled SSE (always enabled) - bytes8 constant TD_XFAM_SSE = 0x0000000000000002; - - /// @dev See section 3.4.1 in TDX Module ABI specification https://cdrdv2.intel.com/v1/dl/getContent/733579 - /// @notice Allows disabling of EPT violation conversion to #VE on access of PENDING pages. Needed for Linux - bytes8 constant TD_TDATTRS_VE_DISABLED = 0x0000000010000000; - /// @notice Enabled Supervisor Protection Keys (PKS) - bytes8 constant TD_TDATTRS_PKS = 0x0000000040000000; - /// @notice Enabled Key Locker (KL) - bytes8 constant TD_TDATTRS_KL = 0x0000000080000000; - // ============ Storage Variables ============ - /// @notice Mapping from workloadId to its metadata (commit hash and source locators) - /// @dev This is only updateable by governance (i.e. the owner) of the Policy contract - /// Adding and removing a workload is O(1). - /// This means the critical `_cachedIsAllowedPolicy` function is O(1) since we can directly check if a workloadId exists - /// in the mapping - mapping(bytes32 workloadId => WorkloadMetadata) private approvedWorkloads; - - /// @inheritdoc IBlockBuilderPolicy - address public registry; - /// @inheritdoc IBlockBuilderPolicy mapping(address teeAddress => uint256 permitNonce) public nonces; @@ -82,25 +51,72 @@ contract BlockBuilderPolicy is /// @dev Maps teeAddress to cached workload information for gas optimization mapping(address teeAddress => CachedWorkload) private cachedWorkloads; + /// @notice Workload deriver used by the shared base policy logic (Option B). + IWorkloadDeriver public workloadDeriver; + + /// @notice Emitted when the workload deriver is set or updated. + event WorkloadDeriverSet(address indexed deriver); + /// @dev Storage gap to allow for future storage variable additions in upgrades - /// @dev This reserves 46 storage slots (out of 50 total - 4 used for approvedWorkloads, registry, nonces, and cachedWorkloads) - uint256[46] __gap; + /// @dev This reserves 45 storage slots (out of 50 total - 5 used for approvedWorkloads, registry, nonces, cachedWorkloads, and workloadDeriver) + uint256[45] __gap; /// @inheritdoc IBlockBuilderPolicy - function initialize(address _initialOwner, address _registry) external override initializer { + function initialize(address _initialOwner, address _registry, address deriver) external override initializer { __Ownable_init(_initialOwner); __UUPSUpgradeable_init(); __EIP712_init("BlockBuilderPolicy", "1"); - require(_registry != address(0), InvalidRegistry()); + _basePolicyInit(_registry); + _setWorkloadDeriver(deriver); + } - registry = _registry; - emit RegistrySet(_registry); + /// @notice Set the workload deriver (governance only). + /// @dev This is needed for proxy upgrade flows where a new storage variable is introduced. + function setWorkloadDeriver(address deriver) external override onlyOwner { + _setWorkloadDeriver(deriver); + } + + function _setWorkloadDeriver(address deriver) internal { + require(deriver != address(0), "InvalidWorkloadDeriver"); + require(deriver.code.length > 0, "InvalidWorkloadDeriver"); + + // Guard: this policy's `isAllowedPolicy` override assumes the configured deriver supports + // `workloadIdForReportBody(TD10ReportBody)` (to avoid re-parsing raw quotes). + // + // Solidity "type safety" is not enforced at runtime, so we proactively verify the method exists and succeeds. + TD10ReportBody memory empty; + (bool ok, bytes memory ret) = + deriver.staticcall(abi.encodeCall(TDXWorkloadDeriver.workloadIdForReportBody, (empty))); + require(ok && ret.length == 32, "DeriverMissingWorkloadIdForReportBody"); + + workloadDeriver = IWorkloadDeriver(deriver); + emit WorkloadDeriverSet(deriver); } /// @notice Restricts upgrades to owner only /// @param newImplementation The address of the new implementation contract function _authorizeUpgrade(address newImplementation) internal override onlyOwner {} + // ============ BasePolicy hooks ============ + + function _checkPolicyAuthority() internal view override { + _checkOwner(); + } + + function _workloadDeriver() internal view override returns (IWorkloadDeriver) { + return workloadDeriver; + } + + function _getCachedWorkload(address teeAddress) internal view override returns (CachedWorkload memory) { + return cachedWorkloads[teeAddress]; + } + + function _setCachedWorkload(address teeAddress, CachedWorkload memory cached) internal override { + cachedWorkloads[teeAddress] = cached; + } + + // ============ Block-builder proof verification ============ + /// @inheritdoc IBlockBuilderPolicy function verifyBlockBuilderProof(uint8 version, bytes32 blockContentHash) external override { _verifyBlockBuilderProof(msg.sender, version, blockContentHash); @@ -152,21 +168,27 @@ contract BlockBuilderPolicy is emit BlockBuilderProofVerified(teeAddress, workloadKey, version, blockContentHash, commitHash); } - /// @inheritdoc IBlockBuilderPolicy - function isAllowedPolicy(address teeAddress) public view override returns (bool allowed, WorkloadId) { - // Get full registration data and compute workload ID + /// @inheritdoc BasePolicy + /// @dev Override to avoid re-parsing `registration.rawQuote` on cache misses. We already have the parsed report body + /// stored in the registry registration. + function isAllowedPolicy(address teeAddress) + public + view + override(BasePolicy, IBasePolicy) + returns (bool allowed, WorkloadId) + { (, IFlashtestationRegistry.RegisteredTEE memory registration) = - FlashtestationRegistry(registry).getRegistration(teeAddress); + IFlashtestationRegistry(registry).getRegistration(teeAddress); - // Invalid Registrations means the attestation used to register the TEE is no longer valid - // and so we cannot trust any input from the TEE if (!registration.isValid) { return (false, WorkloadId.wrap(0)); } - WorkloadId workloadId = workloadIdForTDRegistration(registration); + // NOTE: This policy enforces that the configured deriver supports TDX report-body derivation + // via checks in _setWorkloadDeriver + WorkloadId workloadId = + TDXWorkloadDeriver(address(workloadDeriver)).workloadIdForReportBody(registration.parsedReportBody); - // Check if the workload exists in our approved workloads mapping if (bytes(approvedWorkloads[WorkloadId.unwrap(workloadId)].commitHash).length > 0) { return (true, workloadId); } @@ -174,118 +196,22 @@ contract BlockBuilderPolicy is return (false, WorkloadId.wrap(0)); } - /// @notice isAllowedPolicy but with caching to reduce gas costs - /// @dev This function is only used by the verifyBlockBuilderProof function, which needs to be as efficient as possible - /// because it is called onchain for every flashblock. The workloadId is cached to avoid expensive recomputation - /// @dev A careful reader will notice that this function does not delete stale cache entries. It overwrites them - /// if the underlying TEE registration is still valid. But for stale cache entries in every other scenario, the - /// cache entry persists indefinitely. This is because every other instance results in a return value of (false, 0) - /// to the caller (which is always the verifyBlockBuilderProof function) and it immediately reverts. This is an unfortunate - /// consequence of our need to make this function as gas-efficient as possible, otherwise we would try to cleanup - /// stale cache entries - /// @param teeAddress The TEE-controlled address - /// @return True if the TEE is using an approved workload in the policy - /// @return The workloadId of the TEE that is using an approved workload in the policy, or 0 if - /// the TEE is not using an approved workload in the policy - function _cachedIsAllowedPolicy(address teeAddress) private returns (bool, WorkloadId) { - // Get the current registration status (fast path) - (bool isValid, bytes32 quoteHash) = FlashtestationRegistry(registry).getRegistrationStatus(teeAddress); - if (!isValid) { - return (false, WorkloadId.wrap(0)); - } - - // Now, check if we have a cached workload for this TEE - CachedWorkload memory cached = cachedWorkloads[teeAddress]; - - // Check if we've already fetched and computed the workloadId for this TEE - bytes32 cachedWorkloadId = WorkloadId.unwrap(cached.workloadId); - if (cachedWorkloadId != 0 && cached.quoteHash == quoteHash) { - // Cache hit - verify the workload is still a part of this policy's approved workloads - if (bytes(approvedWorkloads[cachedWorkloadId].commitHash).length > 0) { - return (true, cached.workloadId); - } else { - // The workload is no longer approved, so the policy is no longer valid for this TEE\ - return (false, WorkloadId.wrap(0)); - } - } else { - // Cache miss or quote changed - use the view function to get the result - (bool allowed, WorkloadId workloadId) = isAllowedPolicy(teeAddress); - - if (allowed) { - // Update cache with the new workload ID - cachedWorkloads[teeAddress] = CachedWorkload({workloadId: workloadId, quoteHash: quoteHash}); - } - - return (allowed, workloadId); - } + /// @notice Derive a workloadId from a parsed report body via the configured deriver. + /// @dev This is `view` because it performs an external call to the configured deriver contract. + /// @dev We intentionally call `workloadIdForReportBody` directly to avoid re-parsing the quote when the report body + /// is already available. + function workloadIdForReportBody(TD10ReportBody memory reportBody) public view returns (WorkloadId) { + return TDXWorkloadDeriver(address(workloadDeriver)).workloadIdForReportBody(reportBody); } /// @inheritdoc IBlockBuilderPolicy function workloadIdForTDRegistration(IFlashtestationRegistry.RegisteredTEE memory registration) public - pure + view override returns (WorkloadId) { - // We expect FPU and SSE xfam bits to be set, and anything else should be handled by explicitly allowing the workloadid - bytes8 expectedXfamBits = TD_XFAM_FPU | TD_XFAM_SSE; - - // We don't mind VE_DISABLED, PKS, and KL tdattributes bits being set either way, anything else requires explicitly allowing the workloadid - bytes8 ignoredTdAttributesBitmask = TD_TDATTRS_VE_DISABLED | TD_TDATTRS_PKS | TD_TDATTRS_KL; - - return WorkloadId.wrap( - keccak256( - bytes.concat( - registration.parsedReportBody.mrTd, - registration.parsedReportBody.rtMr0, - registration.parsedReportBody.rtMr1, - registration.parsedReportBody.rtMr2, - registration.parsedReportBody.rtMr3, - // VMM configuration - registration.parsedReportBody.mrConfigId, - registration.parsedReportBody.xFAM ^ expectedXfamBits, - registration.parsedReportBody.tdAttributes & ~ignoredTdAttributesBitmask - ) - ) - ); - } - - /// @inheritdoc IBlockBuilderPolicy - function addWorkloadToPolicy(WorkloadId workloadId, string calldata commitHash, string[] calldata sourceLocators) - external - override - onlyOwner - { - require(bytes(commitHash).length > 0, EmptyCommitHash()); - require(sourceLocators.length > 0, EmptySourceLocators()); - - bytes32 workloadKey = WorkloadId.unwrap(workloadId); - - // Check if workload already exists - require(bytes(approvedWorkloads[workloadKey].commitHash).length == 0, WorkloadAlreadyInPolicy()); - - // Store the workload metadata - approvedWorkloads[workloadKey] = WorkloadMetadata({commitHash: commitHash, sourceLocators: sourceLocators}); - - emit WorkloadAddedToPolicy(workloadKey); - } - - /// @inheritdoc IBlockBuilderPolicy - function removeWorkloadFromPolicy(WorkloadId workloadId) external override onlyOwner { - bytes32 workloadKey = WorkloadId.unwrap(workloadId); - - // Check if workload exists - require(bytes(approvedWorkloads[workloadKey].commitHash).length > 0, WorkloadNotInPolicy()); - - // Remove the workload metadata - delete approvedWorkloads[workloadKey]; - - emit WorkloadRemovedFromPolicy(workloadKey); - } - - /// @inheritdoc IBlockBuilderPolicy - function getWorkloadMetadata(WorkloadId workloadId) external view override returns (WorkloadMetadata memory) { - return approvedWorkloads[WorkloadId.unwrap(workloadId)]; + return workloadIdForReportBody(registration.parsedReportBody); } /// @inheritdoc IBlockBuilderPolicy diff --git a/src/FlashtestationRegistry.sol b/src/FlashtestationRegistry.sol index 36d3635..26f2085 100644 --- a/src/FlashtestationRegistry.sol +++ b/src/FlashtestationRegistry.sol @@ -5,8 +5,9 @@ import "@openzeppelin/contracts-upgradeable/proxy/utils/Initializable.sol"; import "@openzeppelin/contracts-upgradeable/proxy/utils/UUPSUpgradeable.sol"; import "@openzeppelin/contracts-upgradeable/access/OwnableUpgradeable.sol"; import {ECDSA} from "@openzeppelin/contracts/utils/cryptography/ECDSA.sol"; -import {ReentrancyGuardTransientUpgradeable} from - "@openzeppelin/contracts-upgradeable/utils/ReentrancyGuardTransientUpgradeable.sol"; +import { + ReentrancyGuardTransientUpgradeable +} from "@openzeppelin/contracts-upgradeable/utils/ReentrancyGuardTransientUpgradeable.sol"; import {EIP712Upgradeable} from "@openzeppelin/contracts-upgradeable/utils/cryptography/EIP712Upgradeable.sol"; import {IAttestation} from "./interfaces/IAttestation.sol"; import {IFlashtestationRegistry} from "./interfaces/IFlashtestationRegistry.sol"; diff --git a/src/derivers/TDXWorkloadDeriver.sol b/src/derivers/TDXWorkloadDeriver.sol new file mode 100644 index 0000000..a2df997 --- /dev/null +++ b/src/derivers/TDXWorkloadDeriver.sol @@ -0,0 +1,64 @@ +// SPDX-License-Identifier: MIT +pragma solidity 0.8.28; + +import {TD10ReportBody} from "automata-dcap-attestation/contracts/types/V4Structs.sol"; +import {IWorkloadDeriver} from "../interfaces/IWorkloadDeriver.sol"; +import {WorkloadId} from "../interfaces/IPolicyCommon.sol"; +import {QuoteParser} from "../utils/QuoteParser.sol"; + +/// @notice Pure TDX workload-id derivation helpers. +/// @dev Kept alongside `TDXWorkloadDeriver` so policies can reuse the exact same logic without +/// having to make an external call. +library TDXWorkloadDeriverLib { + /// @dev See section 11.5.3 in TDX Module v1.5 Base Architecture Specification. + bytes8 internal constant TD_XFAM_FPU = 0x0000000000000001; + bytes8 internal constant TD_XFAM_SSE = 0x0000000000000002; + + /// @dev See section 3.4.1 in TDX Module ABI specification. + bytes8 internal constant TD_TDATTRS_VE_DISABLED = 0x0000000010000000; + bytes8 internal constant TD_TDATTRS_PKS = 0x0000000040000000; + bytes8 internal constant TD_TDATTRS_KL = 0x0000000080000000; + + function workloadIdForReportBody(TD10ReportBody memory reportBody) internal pure returns (WorkloadId) { + // We expect FPU and SSE xfam bits to be set, and anything else should be handled by explicitly allowing + // the workloadId (by governance). + bytes8 expectedXfamBits = TD_XFAM_FPU | TD_XFAM_SSE; + + // We don't mind VE_DISABLED, PKS, and KL tdattributes bits being set either way; anything else requires + // explicitly allowing the workloadId. + bytes8 ignoredTdAttributesBitmask = TD_TDATTRS_VE_DISABLED | TD_TDATTRS_PKS | TD_TDATTRS_KL; + + return WorkloadId.wrap( + keccak256( + bytes.concat( + reportBody.mrTd, + reportBody.rtMr0, + reportBody.rtMr1, + reportBody.rtMr2, + reportBody.rtMr3, + // VMM configuration + reportBody.mrConfigId, + reportBody.xFAM ^ expectedXfamBits, + reportBody.tdAttributes & ~ignoredTdAttributesBitmask + ) + ) + ); + } +} + +/// @notice Workload deriver that matches the current onchain TDX derivation logic. +contract TDXWorkloadDeriver is IWorkloadDeriver { + /// @notice Pure helper to derive a workload ID from a parsed TDX report body. + /// @dev Makes it easy to compute workload IDs pre-registration (e.g. governance approvals). + function workloadIdForReportBody(TD10ReportBody memory reportBody) public pure returns (WorkloadId) { + return TDXWorkloadDeriverLib.workloadIdForReportBody(reportBody); + } + + /// @inheritdoc IWorkloadDeriver + function workloadIdForQuote(bytes calldata rawQuote) external pure returns (WorkloadId) { + bytes memory raw = rawQuote; + TD10ReportBody memory reportBody = QuoteParser.parseV4Quote(raw); + return TDXWorkloadDeriverLib.workloadIdForReportBody(reportBody); + } +} + diff --git a/src/interfaces/IBasePolicy.sol b/src/interfaces/IBasePolicy.sol new file mode 100644 index 0000000..bcd01d3 --- /dev/null +++ b/src/interfaces/IBasePolicy.sol @@ -0,0 +1,29 @@ +// SPDX-License-Identifier: MIT +pragma solidity ^0.8.20; + +import {IPolicyCommon} from "./IPolicyCommon.sol"; +import {WorkloadId} from "./IPolicyCommon.sol"; + +/// @notice Shared policy interface (workload allowlist + registry binding). +interface IBasePolicy is IPolicyCommon { + /// @notice Check if this TEE-controlled address has a valid registry registration and + /// whether its workload is approved under this policy. + /// @param teeAddress The TEE-controlled address. + /// @return allowed True if the TEE is using an approved workload in the policy. + /// @return workloadId The workloadId of the TEE that is using an approved workload, or 0 if not allowed. + function isAllowedPolicy(address teeAddress) external view returns (bool allowed, WorkloadId workloadId); + + /// @notice Add a workload to a policy (governance only). + function addWorkloadToPolicy(WorkloadId workloadId, string calldata commitHash, string[] calldata sourceLocators) + external; + + /// @notice Remove a workload from a policy (governance only). + function removeWorkloadFromPolicy(WorkloadId workloadId) external; + + /// @notice Get metadata for an approved workload. + function getWorkloadMetadata(WorkloadId workloadId) external view returns (WorkloadMetadata memory); + + /// @notice Address of the FlashtestationRegistry contract. + function registry() external view returns (address); +} + diff --git a/src/interfaces/IBlockBuilderPolicy.sol b/src/interfaces/IBlockBuilderPolicy.sol index f05cfb1..6c55000 100644 --- a/src/interfaces/IBlockBuilderPolicy.sol +++ b/src/interfaces/IBlockBuilderPolicy.sol @@ -2,42 +2,16 @@ pragma solidity ^0.8.20; import {IFlashtestationRegistry} from "./IFlashtestationRegistry.sol"; - -/// @notice WorkloadID uniquely identifies a TEE workload. A workload is roughly equivalent to a version of an application's -/// code, can be reproduced from source code, and is derived from a combination of the TEE's measurement registers. -/// The TDX platform provides several registers that capture cryptographic hashes of code, data, and configuration -/// loaded into the TEE's environment. This means that whenever a TEE device changes anything about its compute stack -/// (e.g. user code, firmware, OS, etc), the workloadID will change. -/// See the [Flashtestation's specification](https://github.com/flashbots/rollup-boost/blob/main/specs/flashtestations.md#workload-identity-derivation) for more details -type WorkloadId is bytes32; +import {IBasePolicy} from "./IBasePolicy.sol"; +import {WorkloadId} from "./IPolicyCommon.sol"; /** * @title IBlockBuilderPolicy * @dev Interface exposing errors, events, and external/public functions of BlockBuilderPolicy */ -interface IBlockBuilderPolicy { - // ============ Types ============ - - /** - * @notice Metadata associated with a workload - * @dev Used to track the source code used to build the TEE image identified by the workloadId - */ - struct WorkloadMetadata { - string commitHash; - string[] sourceLocators; - } - +interface IBlockBuilderPolicy is IBasePolicy { // ============ Events ============ - /// @notice Emitted when a workload is added to the policy - /// @param workloadId The workload identifier - event WorkloadAddedToPolicy(bytes32 indexed workloadId); - /// @notice Emitted when a workload is removed from the policy - /// @param workloadId The workload identifier - event WorkloadRemovedFromPolicy(bytes32 indexed workloadId); - /// @notice Emitted when the registry is set in the initializer - /// @param registry The address of the registry - event RegistrySet(address indexed registry); /// @notice Emitted when a block builder proof is successfully verified /// @param caller The address that called the verification function (TEE address) /// @param workloadId The workload identifier of the TEE @@ -49,28 +23,22 @@ interface IBlockBuilderPolicy { ); // ============ Errors ============ - - /// @notice Emitted when the registry is the 0x0 address - error InvalidRegistry(); - /// @notice Emitted when a workload to be added is already in the policy - error WorkloadAlreadyInPolicy(); - /// @notice Emitted when a workload to be removed is not in the policy - error WorkloadNotInPolicy(); /// @notice Emitted when the address is not in the approvedWorkloads mapping error UnauthorizedBlockBuilder(address caller); /// @notice Emitted when the nonce is invalid error InvalidNonce(uint256 expected, uint256 provided); - /// @notice Emitted when the commit hash is empty - error EmptyCommitHash(); - /// @notice Emitted when the source locators array is empty - error EmptySourceLocators(); // ============ Functions ============ /// @notice Initializer to set the FlashtestationRegistry contract which verifies TEE quotes and the initial owner of the contract /// @param _initialOwner The address of the initial owner of the contract /// @param _registry The address of the registry contract - function initialize(address _initialOwner, address _registry) external; + /// @param _workloadDeriver Address of the workload deriver used for workloadId computation + function initialize(address _initialOwner, address _registry, address _workloadDeriver) external; + + /// @notice Set the workload deriver (governance only). + /// @dev Useful for proxy upgrade flows introducing the deriver variable. + function setWorkloadDeriver(address _workloadDeriver) external; /// @notice Verify a block builder proof with a Flashtestation Transaction /// @param version The version of the flashtestation's protocol used to generate the block builder proof @@ -102,14 +70,6 @@ interface IBlockBuilderPolicy { bytes calldata eip712Sig ) external; - /// @notice Check if this TEE-controlled address has registered a valid TEE workload with the registry, and - /// if the workload is approved under this policy - /// @param teeAddress The TEE-controlled address - /// @return allowed True if the TEE is using an approved workload in the policy - /// @return workloadId The workloadId of the TEE that is using an approved workload in the policy, or 0 if - /// the TEE is not using an approved workload in the policy - function isAllowedPolicy(address teeAddress) external view returns (bool, WorkloadId workloadId); - /// @notice Application specific mapping of registration data to a workload identifier /// @dev Think of the workload identifier as the version of the application for governance. /// The workloadId verifiably maps to a version of source code that builds the TEE VM image @@ -117,44 +77,9 @@ interface IBlockBuilderPolicy { /// @return workloadId The computed workload identifier function workloadIdForTDRegistration(IFlashtestationRegistry.RegisteredTEE memory registration) external - pure + view returns (WorkloadId); - /// @notice Add a workload to a policy (governance only) - /// @notice Only the owner of this contract can add workloads to the policy - /// and it is the responsibility of the owner to ensure that the workload is valid - /// otherwise the address associated with this workload has full power to do anything - /// who's authorization is based on this policy - /// @dev The commitHash solves the following problem; The only way for a smart contract like BlockBuilderPolicy - /// to verify that a TEE (identified by its workloadId) is running a specific piece of code (for instance, - /// op-rbuilder) is to reproducibly build that workload onchain. This is prohibitively expensive, so instead - /// we rely on a permissioned multisig (the owner of this contract) to add a commit hash to the policy whenever - /// it adds a new workloadId. We're already relying on the owner to verify that the workloadId is valid, so - /// we can also assume the owner will not add a commit hash that is not associated with the workloadId. If - /// the owner did act maliciously, this can easily be determined offchain by an honest actor building the - /// TEE image from the given commit hash, deriving the image's workloadId, and then comparing it to the - /// workloadId stored on the policy that is associated with the commit hash. If the workloadId is different, - /// this can be used to prove that the owner acted maliciously. In the honest case, this Policy serves as a - /// source of truth for which source code of build software (i.e. the commit hash) is used to build the TEE image - /// identified by the workloadId. - /// @param workloadId The workload identifier - /// @param commitHash The 40-character hexadecimal commit hash of the git repository - /// whose source code is used to build the TEE image identified by the workloadId - /// @param sourceLocators An array of URIs pointing to the source code - function addWorkloadToPolicy(WorkloadId workloadId, string calldata commitHash, string[] calldata sourceLocators) - external; - - /// @notice Remove a workload from a policy (governance only) - /// @param workloadId The workload identifier - function removeWorkloadFromPolicy(WorkloadId workloadId) external; - - /// @notice Mapping from workloadId to its metadata (commit hash and source locators) - /// @dev This is only updateable by governance (i.e. the owner) of the Policy contract - /// Adding and removing a workload is O(1) - /// @param workloadId The workload identifier to query - /// @return The metadata associated with the workload - function getWorkloadMetadata(WorkloadId workloadId) external view returns (WorkloadMetadata memory); - /// @notice Computes the digest for the EIP-712 signature /// @param structHash The struct hash for the EIP-712 signature /// @return The digest for the EIP-712 signature @@ -165,10 +90,7 @@ interface IBlockBuilderPolicy { /// @param blockContentHash The hash of the block content /// @param nonce The nonce to use for the EIP-712 signature /// @return The struct hash for the EIP-712 signature - function computeStructHash(uint8 version, bytes32 blockContentHash, uint256 nonce) - external - pure - returns (bytes32); + function computeStructHash(uint8 version, bytes32 blockContentHash, uint256 nonce) external pure returns (bytes32); /// @notice Returns the domain separator for the EIP-712 signature /// @dev This is useful for when both onchain and offchain users want to compute the domain separator @@ -178,9 +100,6 @@ interface IBlockBuilderPolicy { // ============ Auto-generated getters for public state ============ - /// @notice Address of the FlashtestationRegistry contract that verifies TEE quotes - function registry() external view returns (address); - /// @notice Tracks nonces for EIP-712 signatures to prevent replay attacks function nonces(address teeAddress) external view returns (uint256); diff --git a/src/interfaces/IPolicyCommon.sol b/src/interfaces/IPolicyCommon.sol new file mode 100644 index 0000000..2736a65 --- /dev/null +++ b/src/interfaces/IPolicyCommon.sol @@ -0,0 +1,51 @@ +// SPDX-License-Identifier: MIT +pragma solidity ^0.8.20; + +/// @notice WorkloadID uniquely identifies a TEE workload. +/// @dev Shared type used across policies and derivers. +type WorkloadId is bytes32; + +/// @notice Common policy surface shared across policy implementations. +/// @dev This is intentionally small and reusable (no block-builder-specific concerns). +interface IPolicyCommon { + /** + * @notice Metadata associated with a workload. + * @dev Used to track the source code used to build the TEE image identified by the workloadId. + */ + struct WorkloadMetadata { + string commitHash; + string[] sourceLocators; + } + + // ============ Events ============ + + /// @notice Emitted when a workload is added to the policy. + /// @param workloadId The workload identifier. + event WorkloadAddedToPolicy(bytes32 indexed workloadId); + + /// @notice Emitted when a workload is removed from the policy. + /// @param workloadId The workload identifier. + event WorkloadRemovedFromPolicy(bytes32 indexed workloadId); + + /// @notice Emitted when the registry is set (initialization). + /// @param registry The address of the registry. + event RegistrySet(address indexed registry); + + // ============ Errors ============ + + /// @notice Emitted when the registry is the 0x0 address. + error InvalidRegistry(); + + /// @notice Emitted when a workload to be added is already in the policy. + error WorkloadAlreadyInPolicy(); + + /// @notice Emitted when a workload to be removed is not in the policy. + error WorkloadNotInPolicy(); + + /// @notice Emitted when the commit hash is empty. + error EmptyCommitHash(); + + /// @notice Emitted when the source locators array is empty. + error EmptySourceLocators(); +} + diff --git a/src/interfaces/IWorkloadDeriver.sol b/src/interfaces/IWorkloadDeriver.sol new file mode 100644 index 0000000..c286b56 --- /dev/null +++ b/src/interfaces/IWorkloadDeriver.sol @@ -0,0 +1,15 @@ +// SPDX-License-Identifier: MIT +pragma solidity ^0.8.20; + +import {TD10ReportBody} from "automata-dcap-attestation/contracts/types/V4Structs.sol"; +import {WorkloadId} from "./IPolicyCommon.sol"; + +/// @notice Computes a workload ID from a parsed TDX report body. +/// @dev policies delegate workload derivation to an injected deriver. +interface IWorkloadDeriver { + /// @notice Derive a workload ID from a raw attestation quote. + /// @dev Policies can pass `registration.rawQuote` from the registry. + /// @dev The concrete deriver may parse the quote and internally call a report-body helper. + function workloadIdForQuote(bytes calldata rawQuote) external pure returns (WorkloadId); +} + diff --git a/test/BlockBuilderPolicy.t.sol b/test/BlockBuilderPolicy.t.sol index 71a86c1..3ee4912 100644 --- a/test/BlockBuilderPolicy.t.sol +++ b/test/BlockBuilderPolicy.t.sol @@ -6,7 +6,10 @@ import {UnsafeUpgrades} from "openzeppelin-foundry-upgrades/Upgrades.sol"; import "@openzeppelin/contracts-upgradeable/access/OwnableUpgradeable.sol"; import "@openzeppelin/contracts/utils/cryptography/ECDSA.sol"; import {BlockBuilderPolicy} from "../src/BlockBuilderPolicy.sol"; -import {IBlockBuilderPolicy, WorkloadId} from "../src/interfaces/IBlockBuilderPolicy.sol"; +import {IBlockBuilderPolicy} from "../src/interfaces/IBlockBuilderPolicy.sol"; +import {IPolicyCommon} from "../src/interfaces/IPolicyCommon.sol"; +import {WorkloadId} from "../src/interfaces/IPolicyCommon.sol"; +import {TDXWorkloadDeriver} from "../src/derivers/TDXWorkloadDeriver.sol"; import {FlashtestationRegistry} from "../src/FlashtestationRegistry.sol"; import {IFlashtestationRegistry} from "../src/interfaces/IFlashtestationRegistry.sol"; import {MockQuote} from "../test/FlashtestationRegistry.t.sol"; @@ -28,6 +31,7 @@ contract BlockBuilderPolicyTest is Test { FlashtestationRegistry public registry; MockAutomataDcapAttestationFee public attestationContract; BlockBuilderPolicy public policy; + TDXWorkloadDeriver public deriver; Upgrader public upgrader = new Upgrader(); address public owner = address(this); @@ -93,7 +97,8 @@ contract BlockBuilderPolicyTest is Test { ) }); - WorkloadId arbitraryWorkloadId = WorkloadId.wrap(0x1dd337a1486a84a7d4200553584996abec87a87473d445262d5562f84ec456a8); + WorkloadId arbitraryWorkloadId = + WorkloadId.wrap(0x1dd337a1486a84a7d4200553584996abec87a87473d445262d5562f84ec456a8); WorkloadId wrongWorkloadId = WorkloadId.wrap(0x20ab431377d40de192f7c754ac0f1922de05ab2f73e74204f0b3ab73a8856876); using ECDSA for bytes32; @@ -106,9 +111,11 @@ contract BlockBuilderPolicyTest is Test { abi.encodeCall(FlashtestationRegistry.initialize, (owner, address(attestationContract))) ); registry = FlashtestationRegistry(registryProxy); + deriver = new TDXWorkloadDeriver(); address policyImplementation = address(new BlockBuilderPolicy()); address policyProxy = UnsafeUpgrades.deployUUPSProxy( - policyImplementation, abi.encodeCall(BlockBuilderPolicy.initialize, (owner, address(registry))) + policyImplementation, + abi.encodeCall(BlockBuilderPolicy.initialize, (owner, address(registry), address(deriver))) ); policy = BlockBuilderPolicy(policyProxy); } @@ -127,11 +134,13 @@ contract BlockBuilderPolicyTest is Test { abi.encodeCall(FlashtestationRegistry.initialize, (owner, address(attestationContract))) ); registry = FlashtestationRegistry(registryProxy); + deriver = new TDXWorkloadDeriver(); address policyImplementation = address(new BlockBuilderPolicy()); vm.expectRevert(abi.encodeWithSelector(OwnableUpgradeable.OwnableInvalidOwner.selector, address(0x0))); UnsafeUpgrades.deployUUPSProxy( - policyImplementation, abi.encodeCall(BlockBuilderPolicy.initialize, (address(0), address(registry))) + policyImplementation, + abi.encodeCall(BlockBuilderPolicy.initialize, (address(0), address(registry), address(deriver))) ); } @@ -143,11 +152,12 @@ contract BlockBuilderPolicyTest is Test { abi.encodeCall(FlashtestationRegistry.initialize, (owner, address(attestationContract))) ); registry = FlashtestationRegistry(registryProxy); + deriver = new TDXWorkloadDeriver(); address policyImplementation = address(new BlockBuilderPolicy()); - vm.expectRevert(abi.encodeWithSelector(IBlockBuilderPolicy.InvalidRegistry.selector)); + vm.expectRevert(abi.encodeWithSelector(IPolicyCommon.InvalidRegistry.selector)); UnsafeUpgrades.deployUUPSProxy( - policyImplementation, abi.encodeCall(BlockBuilderPolicy.initialize, (owner, address(0))) + policyImplementation, abi.encodeCall(BlockBuilderPolicy.initialize, (owner, address(0), address(deriver))) ); } @@ -185,17 +195,17 @@ contract BlockBuilderPolicyTest is Test { function test_addWorkloadToPolicy_reverts_if_duplicate() public { policy.addWorkloadToPolicy(mockf200.workloadId, mockf200.commitHash, mockf200.sourceLocators); - vm.expectRevert(IBlockBuilderPolicy.WorkloadAlreadyInPolicy.selector); + vm.expectRevert(IPolicyCommon.WorkloadAlreadyInPolicy.selector); policy.addWorkloadToPolicy(mockf200.workloadId, mockf200.commitHash, mockf200.sourceLocators); } function test_addWorkloadToPolicy_reverts_if_empty_commit_hash() public { - vm.expectRevert(abi.encodeWithSelector(IBlockBuilderPolicy.EmptyCommitHash.selector, 0)); + vm.expectRevert(IPolicyCommon.EmptyCommitHash.selector); policy.addWorkloadToPolicy(mockf200.workloadId, "", mockf200.sourceLocators); } function test_addWorkloadToPolicy_reverts_if_empty_source_locators() public { - vm.expectRevert(abi.encodeWithSelector(IBlockBuilderPolicy.EmptySourceLocators.selector, 0)); + vm.expectRevert(IPolicyCommon.EmptySourceLocators.selector); policy.addWorkloadToPolicy(mockf200.workloadId, mockf200.commitHash, new string[](0)); } @@ -258,7 +268,7 @@ contract BlockBuilderPolicyTest is Test { } function test_removeWorkloadFromPolicy_reverts_if_not_present() public { - vm.expectRevert(IBlockBuilderPolicy.WorkloadNotInPolicy.selector); + vm.expectRevert(IPolicyCommon.WorkloadNotInPolicy.selector); policy.removeWorkloadFromPolicy(mockf200.workloadId); } diff --git a/test/Examples.t.sol b/test/Examples.t.sol new file mode 100644 index 0000000..9b45ebe --- /dev/null +++ b/test/Examples.t.sol @@ -0,0 +1,131 @@ +// SPDX-License-Identifier: MIT +pragma solidity 0.8.28; + +import {Test} from "forge-std/Test.sol"; +import {DualDeriverPolicy} from "../examples/DualDeriverPolicy.sol"; +import {IFlashtestationRegistry} from "../src/interfaces/IFlashtestationRegistry.sol"; +import {WorkloadId} from "../src/interfaces/IPolicyCommon.sol"; +import {TDXTD15WorkloadDeriver} from "../examples/TDXTD15WorkloadDeriver.sol"; +import {TDXWorkloadDeriver} from "../src/derivers/TDXWorkloadDeriver.sol"; + +contract MockRegistryAlwaysValid { + bytes32 public quoteHash; + bytes public rawQuote; + + constructor(bytes32 quoteHash_, bytes memory rawQuote_) { + quoteHash = quoteHash_; + rawQuote = rawQuote_; + } + + function getRegistration(address) + external + view + returns (bool isValid, IFlashtestationRegistry.RegisteredTEE memory registeredTEE) + { + IFlashtestationRegistry.RegisteredTEE memory reg; + reg.isValid = true; + reg.quoteHash = quoteHash; + reg.rawQuote = rawQuote; + return (true, reg); + } + + function getRegistrationStatus(address) external view returns (bool isValid, bytes32 quoteHash_) { + return (true, quoteHash); + } +} + +contract ExamplesTest is Test { + function test_example_dual_deriver_policy_allows_td10_quote_via_old_deriver() public { + TDXWorkloadDeriver oldDeriver = new TDXWorkloadDeriver(); + TDXTD15WorkloadDeriver newDeriver = new TDXTD15WorkloadDeriver(); + + bytes memory td10Quote = + vm.readFileBinary("test/raw_tdx_quotes/0x46f6b3ACF1dD8Ac0085e30192741336c4aF6EdAF/quote.bin"); + MockRegistryAlwaysValid registry = new MockRegistryAlwaysValid(keccak256(td10Quote), td10Quote); + + DualDeriverPolicy policy = new DualDeriverPolicy(address(this), address(registry), oldDeriver, newDeriver); + + WorkloadId idToApprove = oldDeriver.workloadIdForQuote(td10Quote); + string[] memory locators = new string[](1); + locators[0] = "ipfs://example"; + policy.addWorkloadToPolicy(idToApprove, "deadbeefdeadbeefdeadbeefdeadbeefdeadbeef", locators); + + (bool allowed, WorkloadId workloadId) = policy.isAllowedPolicy(address(0xBEEF)); + assertTrue(allowed); + assertEq(WorkloadId.unwrap(workloadId), WorkloadId.unwrap(idToApprove)); + } + + function test_example_dual_deriver_policy_allows_td15_report_via_new_deriver() public { + TDXWorkloadDeriver oldDeriver = new TDXWorkloadDeriver(); + TDXTD15WorkloadDeriver newDeriver = new TDXTD15WorkloadDeriver(); + + // Build a synthetic TD15 report body (648 bytes) and store it as the registry's "rawQuote" for the example. + bytes memory td15Report = new bytes(648); + // Set mrConfigId (48 bytes) to a non-zero value for determinism. + for (uint256 i = 0; i < 48; i++) { + td15Report[184 + i] = bytes1(uint8(i + 1)); + } + // Set mrServiceTd (48 bytes) to a different non-zero value. + for (uint256 i = 0; i < 48; i++) { + td15Report[600 + i] = bytes1(uint8(0xAA)); + } + + MockRegistryAlwaysValid registry = new MockRegistryAlwaysValid(keccak256(td15Report), td15Report); + DualDeriverPolicy policy = new DualDeriverPolicy(address(this), address(registry), oldDeriver, newDeriver); + + WorkloadId workloadIdToApprove = newDeriver.workloadIdForQuote(td15Report); + string[] memory locators = new string[](1); + locators[0] = "ipfs://example"; + policy.addWorkloadToPolicy(workloadIdToApprove, "deadbeefdeadbeefdeadbeefdeadbeefdeadbeef", locators); + + (bool allowed, WorkloadId workloadId) = policy.isAllowedPolicy(address(0xBEEF)); + assertTrue(allowed); + assertEq(WorkloadId.unwrap(workloadId), WorkloadId.unwrap(workloadIdToApprove)); + } + + function test_example_dual_deriver_policy_rejects_non_owner_updates() public { + TDXWorkloadDeriver oldDeriver = new TDXWorkloadDeriver(); + TDXTD15WorkloadDeriver newDeriver = new TDXTD15WorkloadDeriver(); + + bytes memory td15Report = new bytes(648); + for (uint256 i = 0; i < 48; i++) { + td15Report[184 + i] = bytes1(uint8(i + 1)); + } + for (uint256 i = 0; i < 48; i++) { + td15Report[600 + i] = bytes1(uint8(0xAA)); + } + + MockRegistryAlwaysValid registry = new MockRegistryAlwaysValid(keccak256(td15Report), td15Report); + DualDeriverPolicy policy = new DualDeriverPolicy(address(this), address(registry), oldDeriver, newDeriver); + + WorkloadId workloadIdToApprove = newDeriver.workloadIdForQuote(td15Report); + string[] memory locators = new string[](1); + locators[0] = "ipfs://example"; + + vm.prank(address(0x123)); + vm.expectRevert(bytes("NotOwner")); + policy.addWorkloadToPolicy(workloadIdToApprove, "deadbeefdeadbeefdeadbeefdeadbeefdeadbeef", locators); + } + + function test_example_td15_derivation_hashes_extra_fields() public { + TDXTD15WorkloadDeriver deriver = new TDXTD15WorkloadDeriver(); + + bytes memory td15ReportA = new bytes(648); + bytes memory td15ReportB = new bytes(648); + + // Make base measurements identical... + for (uint256 i = 0; i < 48; i++) { + td15ReportA[184 + i] = bytes1(uint8(i + 1)); // mrConfigId + td15ReportB[184 + i] = bytes1(uint8(i + 1)); + td15ReportA[600 + i] = bytes1(uint8(0x11)); // mrServiceTd + td15ReportB[600 + i] = bytes1(uint8(0x11)); + } + // ...but flip teeTcbSvn2 (16 bytes at offset 584). + td15ReportB[584] = bytes1(uint8(0xFF)); + + WorkloadId idA = deriver.workloadIdForQuote(td15ReportA); + WorkloadId idB = deriver.workloadIdForQuote(td15ReportB); + assertNotEq(WorkloadId.unwrap(idA), WorkloadId.unwrap(idB)); + } +} + diff --git a/test/UpgradeRegression.t.sol b/test/UpgradeRegression.t.sol new file mode 100644 index 0000000..af5cecc --- /dev/null +++ b/test/UpgradeRegression.t.sol @@ -0,0 +1,117 @@ +// SPDX-License-Identifier: MIT +pragma solidity 0.8.28; + +import {Test} from "forge-std/Test.sol"; +import {UnsafeUpgrades} from "openzeppelin-foundry-upgrades/Upgrades.sol"; +import "@openzeppelin/contracts/utils/cryptography/ECDSA.sol"; + +import {BlockBuilderPolicy} from "../src/BlockBuilderPolicy.sol"; +import {FlashtestationRegistry} from "../src/FlashtestationRegistry.sol"; +import {IFlashtestationRegistry} from "../src/interfaces/IFlashtestationRegistry.sol"; +import {IBlockBuilderPolicy} from "../src/interfaces/IBlockBuilderPolicy.sol"; +import {TDXWorkloadDeriver} from "../src/derivers/TDXWorkloadDeriver.sol"; +import {WorkloadId} from "../src/interfaces/IPolicyCommon.sol"; + +import {LegacyBlockBuilderPolicy} from "./fixtures/LegacyBlockBuilderPolicy.sol"; +import {MockAutomataDcapAttestationFee} from "./mocks/MockAutomataDcapAttestationFee.sol"; +import {Upgrader} from "./helpers/Upgrader.sol"; +import {Helper} from "./helpers/Helper.sol"; + +contract UpgradeRegressionTest is Test { + using ECDSA for bytes32; + + MockAutomataDcapAttestationFee public attestationContract; + FlashtestationRegistry public registry; + Upgrader public upgrader = new Upgrader(); + + address public owner = address(this); + + struct QuoteFixture { + bytes output; + bytes quote; + address teeAddress; + uint256 privateKey; + } + + QuoteFixture mock46f6 = QuoteFixture({ + output: vm.readFileBinary("test/raw_tdx_quotes/0x46f6b3ACF1dD8Ac0085e30192741336c4aF6EdAF/output.bin"), + quote: vm.readFileBinary("test/raw_tdx_quotes/0x46f6b3ACF1dD8Ac0085e30192741336c4aF6EdAF/quote.bin"), + teeAddress: 0x46f6b3ACF1dD8Ac0085e30192741336c4aF6EdAF, + privateKey: 0x92e4b5ed61db615b26da2271da5b47c42d691b3164561cfb4edbc85ca6ca61a8 + }); + + function setUp() public { + attestationContract = new MockAutomataDcapAttestationFee(); + address registryImplementation = address(new FlashtestationRegistry()); + address registryProxy = UnsafeUpgrades.deployUUPSProxy( + registryImplementation, + abi.encodeCall(FlashtestationRegistry.initialize, (owner, address(attestationContract))) + ); + registry = FlashtestationRegistry(registryProxy); + } + + function _registerTEE(QuoteFixture memory q) internal { + attestationContract.setQuoteResult(q.quote, true, q.output); + vm.prank(q.teeAddress); + registry.registerTEEService(q.quote, bytes("")); + } + + function test_upgrade_preserves_storage_layout_and_behavior() public { + // 1) Deploy legacy policy behind a proxy + address legacyImpl = address(new LegacyBlockBuilderPolicy()); + address policyProxy = UnsafeUpgrades.deployUUPSProxy( + legacyImpl, abi.encodeCall(LegacyBlockBuilderPolicy.initialize, (owner, address(registry), address(0))) + ); + LegacyBlockBuilderPolicy legacy = LegacyBlockBuilderPolicy(policyProxy); + + // 2) Register TEE and approve its workload in the legacy policy + _registerTEE(mock46f6); + (, IFlashtestationRegistry.RegisteredTEE memory registration) = registry.getRegistration(mock46f6.teeAddress); + WorkloadId workloadId = legacy.workloadIdForTDRegistration(registration); + + string[] memory locators = new string[](1); + locators[0] = "ipfs://example"; + string memory commitHash = "1234567890abcdef1234567890abcdef12345678"; + legacy.addWorkloadToPolicy(workloadId, commitHash, locators); + + // 3) Exercise nonce storage in legacy via permit flow + bytes32 blockContentHash = Helper.computeFlashtestationBlockContentHash(); + bytes32 structHash = legacy.computeStructHash(1, blockContentHash, 0); + bytes32 digest = legacy.getHashedTypeDataV4(structHash); + (uint8 v, bytes32 r, bytes32 s) = vm.sign(mock46f6.privateKey, digest); + bytes memory signature = abi.encodePacked(r, s, v); + legacy.permitVerifyBlockBuilderProof(1, blockContentHash, 0, signature); + assertEq(legacy.nonces(mock46f6.teeAddress), 1); + + // 4) Upgrade to the new implementation + address newImpl = address(new BlockBuilderPolicy()); + upgrader.upgradeProxy(policyProxy, newImpl, bytes(""), owner); + + BlockBuilderPolicy policy = BlockBuilderPolicy(policyProxy); + + // 5) Configure the new deriver (new storage variable added in gap slot) + TDXWorkloadDeriver deriver = new TDXWorkloadDeriver(); + policy.setWorkloadDeriver(address(deriver)); + + // 6) Storage assertions + assertEq(policy.registry(), address(registry)); + assertEq(policy.nonces(mock46f6.teeAddress), 1, "nonce should be preserved across upgrade"); + + IBlockBuilderPolicy.WorkloadMetadata memory meta = policy.getWorkloadMetadata(workloadId); + assertEq(meta.commitHash, commitHash, "workload metadata should be preserved across upgrade"); + + // 7) Behavioral assertion: isAllowedPolicy still works after upgrade + (bool allowed, WorkloadId derivedId) = policy.isAllowedPolicy(mock46f6.teeAddress); + assertTrue(allowed); + assertEq(WorkloadId.unwrap(derivedId), WorkloadId.unwrap(workloadId)); + + // 8) Permit flow still works post-upgrade (nonce increments from preserved value) + structHash = policy.computeStructHash(1, blockContentHash, 1); + digest = policy.getHashedTypeDataV4(structHash); + (v, r, s) = vm.sign(mock46f6.privateKey, digest); + signature = abi.encodePacked(r, s, v); + policy.permitVerifyBlockBuilderProof(1, blockContentHash, 1, signature); + assertEq(policy.nonces(mock46f6.teeAddress), 2); + } +} + diff --git a/test/fixtures/LegacyBlockBuilderPolicy.sol b/test/fixtures/LegacyBlockBuilderPolicy.sol new file mode 100644 index 0000000..24097d0 --- /dev/null +++ b/test/fixtures/LegacyBlockBuilderPolicy.sol @@ -0,0 +1,188 @@ +// SPDX-License-Identifier: MIT +pragma solidity 0.8.28; + +import "@openzeppelin/contracts-upgradeable/proxy/utils/Initializable.sol"; +import "@openzeppelin/contracts-upgradeable/proxy/utils/UUPSUpgradeable.sol"; +import "@openzeppelin/contracts-upgradeable/access/OwnableUpgradeable.sol"; +import {EIP712Upgradeable} from "@openzeppelin/contracts-upgradeable/utils/cryptography/EIP712Upgradeable.sol"; +import {ECDSA} from "@openzeppelin/contracts/utils/cryptography/ECDSA.sol"; +import {FlashtestationRegistry} from "../../src/FlashtestationRegistry.sol"; +import {IFlashtestationRegistry} from "../../src/interfaces/IFlashtestationRegistry.sol"; +import {IBlockBuilderPolicy} from "../../src/interfaces/IBlockBuilderPolicy.sol"; +import {WorkloadId} from "../../src/interfaces/IPolicyCommon.sol"; + +/// @notice Legacy (pre-refactor) BlockBuilderPolicy used for upgrade regression tests. +/// @dev Intentionally mirrors the historical storage layout: +/// approvedWorkloads (slot 0), registry (slot 1), nonces (slot 2), cachedWorkloads (slot 3), gap. +contract LegacyBlockBuilderPolicy is + Initializable, + UUPSUpgradeable, + OwnableUpgradeable, + EIP712Upgradeable, + IBlockBuilderPolicy +{ + using ECDSA for bytes32; + + struct CachedWorkload { + WorkloadId workloadId; + bytes32 quoteHash; + } + + bytes32 public constant VERIFY_BLOCK_BUILDER_PROOF_TYPEHASH = + keccak256("VerifyBlockBuilderProof(uint8 version,bytes32 blockContentHash,uint256 nonce)"); + + // TDX workload constants (legacy) + bytes8 constant TD_XFAM_FPU = 0x0000000000000001; + bytes8 constant TD_XFAM_SSE = 0x0000000000000002; + bytes8 constant TD_TDATTRS_VE_DISABLED = 0x0000000010000000; + bytes8 constant TD_TDATTRS_PKS = 0x0000000040000000; + bytes8 constant TD_TDATTRS_KL = 0x0000000080000000; + + // ===== Storage ===== + mapping(bytes32 workloadId => WorkloadMetadata) private approvedWorkloads; + address public registry; + mapping(address teeAddress => uint256 permitNonce) public nonces; + mapping(address teeAddress => CachedWorkload) private cachedWorkloads; + uint256[46] __gap; + + function initialize(address _initialOwner, address _registry, address) external initializer { + __Ownable_init(_initialOwner); + __UUPSUpgradeable_init(); + __EIP712_init("BlockBuilderPolicy", "1"); + if (_registry == address(0)) revert InvalidRegistry(); + registry = _registry; + emit RegistrySet(_registry); + } + + function setCachedWorkload(address teeAddress, WorkloadId workloadId, bytes32 quoteHash) external onlyOwner { + cachedWorkloads[teeAddress] = CachedWorkload({workloadId: workloadId, quoteHash: quoteHash}); + } + + function _authorizeUpgrade(address) internal override onlyOwner {} + + function verifyBlockBuilderProof(uint8 version, bytes32 blockContentHash) external override { + _verifyBlockBuilderProof(msg.sender, version, blockContentHash); + } + + function permitVerifyBlockBuilderProof( + uint8 version, + bytes32 blockContentHash, + uint256 nonce, + bytes calldata eip712Sig + ) external override { + bytes32 digest = getHashedTypeDataV4(computeStructHash(version, blockContentHash, nonce)); + address teeAddress = digest.recover(eip712Sig); + + uint256 expectedNonce = nonces[teeAddress]; + if (nonce != expectedNonce) revert InvalidNonce(expectedNonce, nonce); + nonces[teeAddress]++; + + _verifyBlockBuilderProof(teeAddress, version, blockContentHash); + } + + function _verifyBlockBuilderProof(address teeAddress, uint8 version, bytes32 blockContentHash) internal { + (bool allowed, WorkloadId workloadId) = _cachedIsAllowedPolicy(teeAddress); + if (!allowed) revert UnauthorizedBlockBuilder(teeAddress); + + bytes32 workloadKey = WorkloadId.unwrap(workloadId); + string memory commitHash = approvedWorkloads[workloadKey].commitHash; + emit BlockBuilderProofVerified(teeAddress, workloadKey, version, blockContentHash, commitHash); + } + + function isAllowedPolicy(address teeAddress) public view override returns (bool allowed, WorkloadId) { + (, IFlashtestationRegistry.RegisteredTEE memory registration) = + FlashtestationRegistry(registry).getRegistration(teeAddress); + if (!registration.isValid) return (false, WorkloadId.wrap(0)); + + WorkloadId workloadId = workloadIdForTDRegistration(registration); + if (bytes(approvedWorkloads[WorkloadId.unwrap(workloadId)].commitHash).length > 0) return (true, workloadId); + return (false, WorkloadId.wrap(0)); + } + + function _cachedIsAllowedPolicy(address teeAddress) private returns (bool, WorkloadId) { + (bool isValid, bytes32 quoteHash) = FlashtestationRegistry(registry).getRegistrationStatus(teeAddress); + if (!isValid) return (false, WorkloadId.wrap(0)); + + CachedWorkload memory cached = cachedWorkloads[teeAddress]; + bytes32 cachedWorkloadId = WorkloadId.unwrap(cached.workloadId); + if (cachedWorkloadId != 0 && cached.quoteHash == quoteHash) { + if (bytes(approvedWorkloads[cachedWorkloadId].commitHash).length > 0) { + return (true, cached.workloadId); + } + return (false, WorkloadId.wrap(0)); + } + + (bool allowed, WorkloadId workloadId) = isAllowedPolicy(teeAddress); + if (allowed) { + cachedWorkloads[teeAddress] = CachedWorkload({workloadId: workloadId, quoteHash: quoteHash}); + } + return (allowed, workloadId); + } + + function workloadIdForTDRegistration(IFlashtestationRegistry.RegisteredTEE memory registration) + public + pure + override + returns (WorkloadId) + { + bytes8 expectedXfamBits = TD_XFAM_FPU | TD_XFAM_SSE; + bytes8 ignoredTdAttributesBitmask = TD_TDATTRS_VE_DISABLED | TD_TDATTRS_PKS | TD_TDATTRS_KL; + + return WorkloadId.wrap( + keccak256( + bytes.concat( + registration.parsedReportBody.mrTd, + registration.parsedReportBody.rtMr0, + registration.parsedReportBody.rtMr1, + registration.parsedReportBody.rtMr2, + registration.parsedReportBody.rtMr3, + registration.parsedReportBody.mrConfigId, + registration.parsedReportBody.xFAM ^ expectedXfamBits, + registration.parsedReportBody.tdAttributes & ~ignoredTdAttributesBitmask + ) + ) + ); + } + + function addWorkloadToPolicy(WorkloadId workloadId, string calldata commitHash, string[] calldata sourceLocators) + external + override + onlyOwner + { + if (bytes(commitHash).length == 0) revert EmptyCommitHash(); + if (sourceLocators.length == 0) revert EmptySourceLocators(); + + bytes32 workloadKey = WorkloadId.unwrap(workloadId); + if (bytes(approvedWorkloads[workloadKey].commitHash).length != 0) revert WorkloadAlreadyInPolicy(); + approvedWorkloads[workloadKey] = WorkloadMetadata({commitHash: commitHash, sourceLocators: sourceLocators}); + emit WorkloadAddedToPolicy(workloadKey); + } + + function removeWorkloadFromPolicy(WorkloadId workloadId) external override onlyOwner { + bytes32 workloadKey = WorkloadId.unwrap(workloadId); + if (bytes(approvedWorkloads[workloadKey].commitHash).length == 0) revert WorkloadNotInPolicy(); + delete approvedWorkloads[workloadKey]; + emit WorkloadRemovedFromPolicy(workloadKey); + } + + function getWorkloadMetadata(WorkloadId workloadId) external view override returns (WorkloadMetadata memory) { + return approvedWorkloads[WorkloadId.unwrap(workloadId)]; + } + + function setWorkloadDeriver(address) external pure override { + revert("Legacy: no deriver"); + } + + function getHashedTypeDataV4(bytes32 structHash) public view returns (bytes32) { + return _hashTypedDataV4(structHash); + } + + function computeStructHash(uint8 version, bytes32 blockContentHash, uint256 nonce) public pure returns (bytes32) { + return keccak256(abi.encode(VERIFY_BLOCK_BUILDER_PROOF_TYPEHASH, version, blockContentHash, nonce)); + } + + function domainSeparator() external view returns (bytes32) { + return _domainSeparatorV4(); + } +} + diff --git a/test/helpers/Helper.sol b/test/helpers/Helper.sol index f66ed8b..f6fa249 100644 --- a/test/helpers/Helper.sol +++ b/test/helpers/Helper.sol @@ -312,11 +312,7 @@ library Helper { }); } - function withQuoteVersion(OutputBuilder memory builder, uint16 value) - internal - pure - returns (OutputBuilder memory) - { + function withQuoteVersion(OutputBuilder memory builder, uint16 value) internal pure returns (OutputBuilder memory) { builder.quoteVersion = value; return builder; } @@ -326,11 +322,7 @@ library Helper { return builder; } - function withTcbStatus(OutputBuilder memory builder, TCBStatus value) - internal - pure - returns (OutputBuilder memory) - { + function withTcbStatus(OutputBuilder memory builder, TCBStatus value) internal pure returns (OutputBuilder memory) { builder.tcbStatus = value; return builder; }