From 0cfed66d565338324d6e97f30eedeff95e27b7cf Mon Sep 17 00:00:00 2001 From: Spencer Solit Date: Sat, 20 Dec 2025 21:44:14 +0000 Subject: [PATCH 1/9] draft --- examples/DualDeriverPolicy.sol | 106 +++++++++ examples/TDXTD15WorkloadDeriver.sol | 123 +++++++++++ script/BlockBuilderPolicy.s.sol | 6 +- script/Interactions.s.sol | 7 +- src/BasePolicy.sol | 185 ++++++++++++++++ src/BlockBuilderPolicy.sol | 239 +++++++-------------- src/FlashtestationRegistry.sol | 5 +- src/derivers/TDXWorkloadDeriver.sol | 64 ++++++ src/interfaces/IBasePolicy.sol | 29 +++ src/interfaces/IBlockBuilderPolicy.sol | 103 +-------- src/interfaces/IPolicyCommon.sol | 51 +++++ src/interfaces/IWorkloadDeriver.sol | 15 ++ test/BlockBuilderPolicy.t.sol | 30 ++- test/Examples.t.sol | 131 +++++++++++ test/UpgradeRegression.t.sol | 117 ++++++++++ test/fixtures/LegacyBlockBuilderPolicy.sol | 188 ++++++++++++++++ test/helpers/Helper.sol | 12 +- 17 files changed, 1137 insertions(+), 274 deletions(-) create mode 100644 examples/DualDeriverPolicy.sol create mode 100644 examples/TDXTD15WorkloadDeriver.sol create mode 100644 src/BasePolicy.sol create mode 100644 src/derivers/TDXWorkloadDeriver.sol create mode 100644 src/interfaces/IBasePolicy.sol create mode 100644 src/interfaces/IPolicyCommon.sol create mode 100644 src/interfaces/IWorkloadDeriver.sol create mode 100644 test/Examples.t.sol create mode 100644 test/UpgradeRegression.t.sol create mode 100644 test/fixtures/LegacyBlockBuilderPolicy.sol 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..703b565 --- /dev/null +++ b/src/BasePolicy.sol @@ -0,0 +1,185 @@ +// 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 + 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 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..d6a70c4 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,26 @@ 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 assumes the configured deriver supports TDX report-body derivation. + 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 +195,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; } From 0283ece56555a0a23539ab6051f20b33af80a926 Mon Sep 17 00:00:00 2001 From: Spencer Solit Date: Sat, 20 Dec 2025 21:58:25 +0000 Subject: [PATCH 2/9] modify comment to say that we enforce instead of assume --- src/BlockBuilderPolicy.sol | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/BlockBuilderPolicy.sol b/src/BlockBuilderPolicy.sol index d6a70c4..145d12a 100644 --- a/src/BlockBuilderPolicy.sol +++ b/src/BlockBuilderPolicy.sol @@ -184,7 +184,8 @@ contract BlockBuilderPolicy is return (false, WorkloadId.wrap(0)); } - // NOTE: This policy assumes the configured deriver supports TDX report-body derivation. + // 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); From 4264b8b8be2f3bb650ea12aedb1357cceaecdfe7 Mon Sep 17 00:00:00 2001 From: Spencer Solit Date: Sat, 20 Dec 2025 22:11:43 +0000 Subject: [PATCH 3/9] make add and remove virtual --- src/BasePolicy.sol | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/BasePolicy.sol b/src/BasePolicy.sol index 703b565..7e6b26d 100644 --- a/src/BasePolicy.sol +++ b/src/BasePolicy.sol @@ -61,6 +61,7 @@ abstract contract BasePolicy is IBasePolicy { /// @inheritdoc IBasePolicy function addWorkloadToPolicy(WorkloadId workloadId, string calldata commitHash, string[] calldata sourceLocators) external + virtual override { _checkPolicyAuthority(); @@ -80,7 +81,7 @@ abstract contract BasePolicy is IBasePolicy { } /// @inheritdoc IBasePolicy - function removeWorkloadFromPolicy(WorkloadId workloadId) external override { + function removeWorkloadFromPolicy(WorkloadId workloadId) external virtual override { _checkPolicyAuthority(); bytes32 workloadKey = WorkloadId.unwrap(workloadId); From 88ca39ccfc5ff17f16e00465054376f6212b0c61 Mon Sep 17 00:00:00 2001 From: Frieder Erdmann Date: Tue, 6 Jan 2026 15:45:14 +0100 Subject: [PATCH 4/9] Restore policy docs after refactor --- src/BasePolicy.sol | 13 ++++++++++--- src/BlockBuilderPolicy.sol | 3 +++ src/derivers/TDXWorkloadDeriver.sol | 8 +++++++- src/interfaces/IBasePolicy.sol | 25 ++++++++++++++++++++++++- src/interfaces/IPolicyCommon.sol | 12 ++++++++++-- 5 files changed, 54 insertions(+), 7 deletions(-) diff --git a/src/BasePolicy.sol b/src/BasePolicy.sol index 7e6b26d..086c1c4 100644 --- a/src/BasePolicy.sol +++ b/src/BasePolicy.sol @@ -22,9 +22,13 @@ abstract contract BasePolicy is IBasePolicy { bytes32 quoteHash; } - // ============ Storage (DO NOT REORDER / ADD) ============ + // ============ Storage Variables (DO NOT REORDER / ADD) ============ /// @notice Mapping from workloadId to its metadata (commit hash and source locators). + /// @dev This is only updateable by governance (i.e. the policy authority). + /// 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) internal approvedWorkloads; /// @inheritdoc IBasePolicy @@ -56,6 +60,8 @@ abstract contract BasePolicy is IBasePolicy { /// @dev Cache write hook. function _setCachedWorkload(address teeAddress, CachedWorkload memory cached) internal virtual; + // ============ Functions ============ + // ============ Common policy logic ============ /// @inheritdoc IBasePolicy @@ -102,7 +108,7 @@ abstract contract BasePolicy is IBasePolicy { /// @inheritdoc IBasePolicy function isAllowedPolicy(address teeAddress) public view virtual override returns (bool allowed, WorkloadId) { - // Get full registration data + // Get full registration data and compute workload ID (, IFlashtestationRegistry.RegisteredTEE memory registration) = IFlashtestationRegistry(registry).getRegistration(teeAddress); @@ -143,6 +149,8 @@ abstract contract BasePolicy is IBasePolicy { /// @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. + /// - Stale cache entries for all other "not allowed" scenarios persist indefinitely for the same reason: + /// the caller reverts immediately and the hot path avoids extra storage writes. /// /// @param teeAddress The TEE-controlled address to check. /// @return allowed True if the TEE is registered, valid, and running an approved workload. @@ -183,4 +191,3 @@ abstract contract BasePolicy is IBasePolicy { } } } - diff --git a/src/BlockBuilderPolicy.sol b/src/BlockBuilderPolicy.sol index 145d12a..428742c 100644 --- a/src/BlockBuilderPolicy.sol +++ b/src/BlockBuilderPolicy.sol @@ -61,6 +61,8 @@ contract BlockBuilderPolicy is /// @dev This reserves 45 storage slots (out of 50 total - 5 used for approvedWorkloads, registry, nonces, cachedWorkloads, and workloadDeriver) uint256[45] __gap; + // ============ Functions ============ + /// @inheritdoc IBlockBuilderPolicy function initialize(address _initialOwner, address _registry, address deriver) external override initializer { __Ownable_init(_initialOwner); @@ -211,6 +213,7 @@ contract BlockBuilderPolicy is override returns (WorkloadId) { + // Uses TDXWorkloadDeriverLib (see the deriver for constants and derivation details). return workloadIdForReportBody(registration.parsedReportBody); } diff --git a/src/derivers/TDXWorkloadDeriver.sol b/src/derivers/TDXWorkloadDeriver.sol index a2df997..9f65e44 100644 --- a/src/derivers/TDXWorkloadDeriver.sol +++ b/src/derivers/TDXWorkloadDeriver.sol @@ -10,13 +10,20 @@ import {QuoteParser} from "../utils/QuoteParser.sol"; /// @dev Kept alongside `TDXWorkloadDeriver` so policies can reuse the exact same logic without /// having to make an external call. library TDXWorkloadDeriverLib { + // ============ TDX workload constants ============ + /// @dev See section 11.5.3 in TDX Module v1.5 Base Architecture Specification. + /// @notice Enabled FPU (always enabled). bytes8 internal constant TD_XFAM_FPU = 0x0000000000000001; + /// @notice Enabled SSE (always enabled). bytes8 internal constant TD_XFAM_SSE = 0x0000000000000002; /// @dev See section 3.4.1 in TDX Module ABI specification. + /// @notice Allows disabling of EPT violation conversion to #VE on access of PENDING pages. Needed for Linux. bytes8 internal constant TD_TDATTRS_VE_DISABLED = 0x0000000010000000; + /// @notice Enabled Supervisor Protection Keys (PKS). bytes8 internal constant TD_TDATTRS_PKS = 0x0000000040000000; + /// @notice Enabled Key Locker (KL). bytes8 internal constant TD_TDATTRS_KL = 0x0000000080000000; function workloadIdForReportBody(TD10ReportBody memory reportBody) internal pure returns (WorkloadId) { @@ -61,4 +68,3 @@ contract TDXWorkloadDeriver is IWorkloadDeriver { return TDXWorkloadDeriverLib.workloadIdForReportBody(reportBody); } } - diff --git a/src/interfaces/IBasePolicy.sol b/src/interfaces/IBasePolicy.sol index bcd01d3..7b792c0 100644 --- a/src/interfaces/IBasePolicy.sol +++ b/src/interfaces/IBasePolicy.sol @@ -6,6 +6,8 @@ import {WorkloadId} from "./IPolicyCommon.sol"; /// @notice Shared policy interface (workload allowlist + registry binding). interface IBasePolicy is IPolicyCommon { + // ============ Functions ============ + /// @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. @@ -14,6 +16,25 @@ interface IBasePolicy is IPolicyCommon { function isAllowedPolicy(address teeAddress) external view returns (bool allowed, WorkloadId workloadId); /// @notice Add a workload to a policy (governance only). + /// @notice Only the policy authority can add workloads to the policy, and it is the responsibility of the + /// authority to ensure that the workload is valid; otherwise the address associated with this workload has + /// full power to do anything whose 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 policy authority) to add a commit hash to the policy whenever + /// it adds a new workloadId. We're already relying on the authority to verify that the workloadId is valid, so + /// we can also assume the authority will not add a commit hash that is not associated with the workloadId. If + /// the authority 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 authority 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; @@ -21,9 +42,11 @@ interface IBasePolicy is IPolicyCommon { function removeWorkloadFromPolicy(WorkloadId workloadId) external; /// @notice Get metadata for an approved workload. + /// @dev This is only updateable by governance (i.e. the policy authority). 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 Address of the FlashtestationRegistry contract. function registry() external view returns (address); } - diff --git a/src/interfaces/IPolicyCommon.sol b/src/interfaces/IPolicyCommon.sol index 2736a65..a5c4af0 100644 --- a/src/interfaces/IPolicyCommon.sol +++ b/src/interfaces/IPolicyCommon.sol @@ -1,13 +1,22 @@ // SPDX-License-Identifier: MIT pragma solidity ^0.8.20; -/// @notice WorkloadID uniquely identifies a TEE workload. +/// @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 specification for more details: +/// https://github.com/flashbots/rollup-boost/blob/main/specs/flashtestations.md#workload-identity-derivation /// @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 { + // ============ Types ============ + /** * @notice Metadata associated with a workload. * @dev Used to track the source code used to build the TEE image identified by the workloadId. @@ -48,4 +57,3 @@ interface IPolicyCommon { /// @notice Emitted when the source locators array is empty. error EmptySourceLocators(); } - From ed68a40f7c3e24e2c3314776464a716c8e8614e3 Mon Sep 17 00:00:00 2001 From: Spencer Solit Date: Wed, 14 Jan 2026 03:13:58 +0000 Subject: [PATCH 5/9] docs: clarify workload deriver composition pattern Remove "(Option B)" reference and add explanation for why the deriver uses a composition pattern, making the design rationale clearer for future maintainers. --- src/BlockBuilderPolicy.sol | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/BlockBuilderPolicy.sol b/src/BlockBuilderPolicy.sol index 428742c..11ff22c 100644 --- a/src/BlockBuilderPolicy.sol +++ b/src/BlockBuilderPolicy.sol @@ -51,7 +51,8 @@ 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). + /// @notice Workload deriver used by the shared base policy logic. + /// @dev Composition pattern allows swapping derivation implementations without modifying base policy logic. IWorkloadDeriver public workloadDeriver; /// @notice Emitted when the workload deriver is set or updated. From e5930b76e524141f2dc6c505d72ac0c67555266f Mon Sep 17 00:00:00 2001 From: Spencer Solit Date: Wed, 14 Jan 2026 03:15:46 +0000 Subject: [PATCH 6/9] refactor: move WorkloadDeriverSet event to interface Move event declaration from BlockBuilderPolicy implementation to IBlockBuilderPolicy interface, following the pattern of other events in the codebase and maintaining interface as the source of truth for public API. --- src/BlockBuilderPolicy.sol | 3 --- src/interfaces/IBlockBuilderPolicy.sol | 4 ++++ 2 files changed, 4 insertions(+), 3 deletions(-) diff --git a/src/BlockBuilderPolicy.sol b/src/BlockBuilderPolicy.sol index 11ff22c..6c85a84 100644 --- a/src/BlockBuilderPolicy.sol +++ b/src/BlockBuilderPolicy.sol @@ -55,9 +55,6 @@ contract BlockBuilderPolicy is /// @dev Composition pattern allows swapping derivation implementations without modifying base policy logic. 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 45 storage slots (out of 50 total - 5 used for approvedWorkloads, registry, nonces, cachedWorkloads, and workloadDeriver) uint256[45] __gap; diff --git a/src/interfaces/IBlockBuilderPolicy.sol b/src/interfaces/IBlockBuilderPolicy.sol index 6c55000..1785642 100644 --- a/src/interfaces/IBlockBuilderPolicy.sol +++ b/src/interfaces/IBlockBuilderPolicy.sol @@ -22,6 +22,10 @@ interface IBlockBuilderPolicy is IBasePolicy { address caller, bytes32 workloadId, uint8 version, bytes32 blockContentHash, string commitHash ); + /// @notice Emitted when the workload deriver is set or updated. + /// @param deriver The address of the workload deriver contract. + event WorkloadDeriverSet(address indexed deriver); + // ============ Errors ============ /// @notice Emitted when the address is not in the approvedWorkloads mapping error UnauthorizedBlockBuilder(address caller); From daadeef14ce4da2e900d6c7e60a7934ae5ccabb7 Mon Sep 17 00:00:00 2001 From: Spencer Solit Date: Wed, 14 Jan 2026 03:19:15 +0000 Subject: [PATCH 7/9] fix: update workload derivation to simplified logic from PR #53 Remove buggy TDX workload constants and masking logic. The new approach directly hashes xFAM and tdAttributes without XOR/AND operations, trading workload constant flexibility for code simplicity and security. Updated files: - src/derivers/TDXWorkloadDeriver.sol - examples/TDXTD15WorkloadDeriver.sol - test/fixtures/LegacyBlockBuilderPolicy.sol --- examples/TDXTD15WorkloadDeriver.sol | 14 ++--------- src/derivers/TDXWorkloadDeriver.sol | 28 ++-------------------- test/fixtures/LegacyBlockBuilderPolicy.sol | 14 ++--------- 3 files changed, 6 insertions(+), 50 deletions(-) diff --git a/examples/TDXTD15WorkloadDeriver.sol b/examples/TDXTD15WorkloadDeriver.sol index 0b44fd6..e5c780a 100644 --- a/examples/TDXTD15WorkloadDeriver.sol +++ b/examples/TDXTD15WorkloadDeriver.sol @@ -66,19 +66,9 @@ library TD15ReportParser { 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( @@ -89,8 +79,8 @@ contract TDXTD15WorkloadDeriver is IWorkloadDeriver { reportBody.rtMr3, // VMM configuration reportBody.mrConfigId, - reportBody.xFAM ^ expectedXfamBits, - reportBody.tdAttributes & ~ignoredTdAttributesBitmask, + reportBody.xFAM, + reportBody.tdAttributes, // TD15 extensions bytes16(reportBody.teeTcbSvn2), reportBody.mrServiceTd diff --git a/src/derivers/TDXWorkloadDeriver.sol b/src/derivers/TDXWorkloadDeriver.sol index 9f65e44..5a4005c 100644 --- a/src/derivers/TDXWorkloadDeriver.sol +++ b/src/derivers/TDXWorkloadDeriver.sol @@ -10,31 +10,7 @@ import {QuoteParser} from "../utils/QuoteParser.sol"; /// @dev Kept alongside `TDXWorkloadDeriver` so policies can reuse the exact same logic without /// having to make an external call. library TDXWorkloadDeriverLib { - // ============ TDX workload constants ============ - - /// @dev See section 11.5.3 in TDX Module v1.5 Base Architecture Specification. - /// @notice Enabled FPU (always enabled). - bytes8 internal constant TD_XFAM_FPU = 0x0000000000000001; - /// @notice Enabled SSE (always enabled). - bytes8 internal constant TD_XFAM_SSE = 0x0000000000000002; - - /// @dev See section 3.4.1 in TDX Module ABI specification. - /// @notice Allows disabling of EPT violation conversion to #VE on access of PENDING pages. Needed for Linux. - bytes8 internal constant TD_TDATTRS_VE_DISABLED = 0x0000000010000000; - /// @notice Enabled Supervisor Protection Keys (PKS). - bytes8 internal constant TD_TDATTRS_PKS = 0x0000000040000000; - /// @notice Enabled Key Locker (KL). - 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( @@ -45,8 +21,8 @@ library TDXWorkloadDeriverLib { reportBody.rtMr3, // VMM configuration reportBody.mrConfigId, - reportBody.xFAM ^ expectedXfamBits, - reportBody.tdAttributes & ~ignoredTdAttributesBitmask + reportBody.xFAM, + reportBody.tdAttributes ) ) ); diff --git a/test/fixtures/LegacyBlockBuilderPolicy.sol b/test/fixtures/LegacyBlockBuilderPolicy.sol index 24097d0..920ea66 100644 --- a/test/fixtures/LegacyBlockBuilderPolicy.sol +++ b/test/fixtures/LegacyBlockBuilderPolicy.sol @@ -31,13 +31,6 @@ contract LegacyBlockBuilderPolicy is 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; @@ -125,9 +118,6 @@ contract LegacyBlockBuilderPolicy is 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( @@ -137,8 +127,8 @@ contract LegacyBlockBuilderPolicy is registration.parsedReportBody.rtMr2, registration.parsedReportBody.rtMr3, registration.parsedReportBody.mrConfigId, - registration.parsedReportBody.xFAM ^ expectedXfamBits, - registration.parsedReportBody.tdAttributes & ~ignoredTdAttributesBitmask + registration.parsedReportBody.xFAM, + registration.parsedReportBody.tdAttributes ) ) ); From c9cd46eada25c96b67ba08200155f25e1e24b97f Mon Sep 17 00:00:00 2001 From: Spencer Solit Date: Wed, 14 Jan 2026 04:06:48 +0000 Subject: [PATCH 8/9] Remove test for tdAttributes bit masking behavior The test_workloadId_tdAttributes_allowed_bits_ignored test verified that certain tdAttributes bits were ignored in workloadId computation. This behavior is no longer expected, so the test has been removed. --- test/BlockBuilderPolicy.t.sol | 43 ----------------------------------- 1 file changed, 43 deletions(-) diff --git a/test/BlockBuilderPolicy.t.sol b/test/BlockBuilderPolicy.t.sol index 3ee4912..bd84832 100644 --- a/test/BlockBuilderPolicy.t.sol +++ b/test/BlockBuilderPolicy.t.sol @@ -371,49 +371,6 @@ contract BlockBuilderPolicyTest is Test { assertEq(WorkloadId.unwrap(computedWorkloadIdF200), WorkloadId.unwrap(computedWorkloadId12c1)); } - // Add these test functions to BlockBuilderPolicyTest contract - - function test_workloadId_tdAttributes_allowed_bits_ignored() public { - // Register a TEE to get a baseline - _registerTEE(mockf200); - (, IFlashtestationRegistry.RegisteredTEE memory baseRegistration) = - registry.getRegistration(mockf200.teeAddress); - WorkloadId baseWorkloadId = policy.workloadIdForTDRegistration(baseRegistration); - - // Test that all combinations of allowed bits don't affect workloadId - // We test: none set, all set, and one intermediate case - bytes8[3] memory allowedBitCombos = [ - bytes8(0x00000000D0000000), // All three allowed bits set (VE_DISABLED | PKS | KL) - bytes8(0x0000000050000000), // VE_DISABLED | PKS - bytes8(0x0000000000000000) // None set - ]; - - for (uint256 i = 0; i < allowedBitCombos.length; i++) { - IFlashtestationRegistry.RegisteredTEE memory modifiedRegAllowed = baseRegistration; - // Clear the allowed bits first, then set the specific combination - modifiedRegAllowed.parsedReportBody.tdAttributes = - (baseRegistration.parsedReportBody.tdAttributes & ~bytes8(0x00000000D0000000)) | allowedBitCombos[i]; - - WorkloadId workloadId = policy.workloadIdForTDRegistration(modifiedRegAllowed); - assertEq( - WorkloadId.unwrap(baseWorkloadId), - WorkloadId.unwrap(workloadId), - "Allowed tdAttributes bits should not affect workloadId" - ); - } - - // Test that a non-allowed bit DOES change workloadId - IFlashtestationRegistry.RegisteredTEE memory modifiedReg = baseRegistration; - modifiedReg.parsedReportBody.tdAttributes = - baseRegistration.parsedReportBody.tdAttributes | bytes8(0x0000000000000001); - WorkloadId differentWorkloadId = policy.workloadIdForTDRegistration(modifiedReg); - assertNotEq( - WorkloadId.unwrap(baseWorkloadId), - WorkloadId.unwrap(differentWorkloadId), - "Non-allowed tdAttributes bits should affect workloadId" - ); - } - function test_workloadId_xfam_expected_bits_required() public { // Register a TEE to get a baseline _registerTEE(mockf200); From be98280db4a8fb72018af526d2ff1ba3909ebedc Mon Sep 17 00:00:00 2001 From: Spencer Solit Date: Wed, 28 Jan 2026 00:04:29 +0000 Subject: [PATCH 9/9] refactor: replace string error messages with custom errors in BlockBuilderPolicy Replace string-based require error messages with custom errors in `_setWorkloadDeriver` to maintain consistency with the contract's error handling pattern. Adds `InvalidWorkloadDeriver` and `DeriverMissingWorkloadIdForReportBody` custom errors to the interface. --- src/BlockBuilderPolicy.sol | 5 ++--- src/interfaces/IBlockBuilderPolicy.sol | 4 ++++ 2 files changed, 6 insertions(+), 3 deletions(-) diff --git a/src/BlockBuilderPolicy.sol b/src/BlockBuilderPolicy.sol index 6c85a84..f32ec24 100644 --- a/src/BlockBuilderPolicy.sol +++ b/src/BlockBuilderPolicy.sol @@ -77,8 +77,7 @@ contract BlockBuilderPolicy is } function _setWorkloadDeriver(address deriver) internal { - require(deriver != address(0), "InvalidWorkloadDeriver"); - require(deriver.code.length > 0, "InvalidWorkloadDeriver"); + require(deriver != address(0) && deriver.code.length > 0, InvalidWorkloadDeriver()); // Guard: this policy's `isAllowedPolicy` override assumes the configured deriver supports // `workloadIdForReportBody(TD10ReportBody)` (to avoid re-parsing raw quotes). @@ -87,7 +86,7 @@ contract BlockBuilderPolicy is TD10ReportBody memory empty; (bool ok, bytes memory ret) = deriver.staticcall(abi.encodeCall(TDXWorkloadDeriver.workloadIdForReportBody, (empty))); - require(ok && ret.length == 32, "DeriverMissingWorkloadIdForReportBody"); + require(ok && ret.length == 32, DeriverMissingWorkloadIdForReportBody()); workloadDeriver = IWorkloadDeriver(deriver); emit WorkloadDeriverSet(deriver); diff --git a/src/interfaces/IBlockBuilderPolicy.sol b/src/interfaces/IBlockBuilderPolicy.sol index 1785642..f675839 100644 --- a/src/interfaces/IBlockBuilderPolicy.sol +++ b/src/interfaces/IBlockBuilderPolicy.sol @@ -31,6 +31,10 @@ interface IBlockBuilderPolicy is IBasePolicy { error UnauthorizedBlockBuilder(address caller); /// @notice Emitted when the nonce is invalid error InvalidNonce(uint256 expected, uint256 provided); + /// @notice Emitted when the workload deriver address is invalid (zero address or no code) + error InvalidWorkloadDeriver(); + /// @notice Emitted when the workload deriver is missing the required workloadIdForReportBody method + error DeriverMissingWorkloadIdForReportBody(); // ============ Functions ============