first draft of AllowList.registerTEEService function

Melvillian · Melvillian · commit fe8079232acf · 2025-05-24T06:48:48.000-04:00
diff --git a/.gitignore b/.gitignore
@@ -12,3 +12,6 @@ docs/
 
 # Dotenv file
 .env
+
+# Cursor.ai rules
+.cursor/
diff --git a/README.md b/README.md
@@ -50,17 +50,16 @@ This will perform a simple test to see if onchain verification of a tdx attestat
 
 ## TODOs
 
-- [] Implement TEE Device Registry
+- [X] Implement TEE Device Registry
 - [] Implement Flashtestation transaction verification
-- [] Implement TEE Device Deregistration
 
 ## Open Questions
 - Should it be Upgradeable?
-    Pros:
-    Cons:
-Should we use Automata's DAO contracts for endorsement/collateral management, or implement our own?
-    Pros:
+    Pros: 
+        - very simple to account for changes to the Automata DCAP Attestation contract, contract bugs, contract upgrades
+        - Doesn't really impact the trust model, because we already expect to have some Security Council of Unichain + Flashbots in the beginning that manages which workloadIDs to trust (via the setting of Policies)
     Cons:
+        - trust model now relies on owner (probably a security council of Unichain + Flashbots) to remain not collude. If they do collude, they can upgrade the contract to emit a malicious `Registered` event and trick users into incorrectly trusting that blocks are being verified by a trusted TEE
 
 
 ## Foundry
diff --git a/foundry.toml b/foundry.toml
@@ -11,14 +11,6 @@ gas_limit = "3000000000"
 fuzz_runs = 10_000
 bytecode_hash = "none"
 
-additional_compiler_profiles = [
-  { name = "test", via_ir = false }
-]
-
-compilation_restrictions = [
-  { paths = "test/**", via_ir = false }
-]
-
 [profile.debug]
 via_ir = false
 optimizer_runs = 200
diff --git a/script/AllowList.s.sol b/script/AllowList.s.sol
@@ -8,8 +8,7 @@ import {TD10ReportBody} from "automata-dcap-attestation/contracts/types/V4Struct
 import {TD_REPORT10_LENGTH} from "automata-dcap-attestation/contracts/types/Constants.sol";
 
 contract AllowListScript is Script {
-    AllowList public registry;
-    AutomataDcapAttestationFee public attestationFee;
+    AllowList public allowlist;
 
     function setUp() public {}
 
@@ -19,31 +18,12 @@ contract AllowListScript is Script {
     function run() public {
         vm.startBroadcast();
 
-        address initialOwner;
-        try vm.envAddress("REGISTRY_INITIAL_OWNER") returns (address owner) {
-            initialOwner = owner;
-        } catch {
-            console.log("Warning: REGISTRY_INITIAL_OWNER not found in environment variables, using msg.sender");
-            initialOwner = msg.sender;
-        }
+        allowlist = new AllowList(ETHEREUM_SEPOLIA_ATTESTATION_FEE_ADDRESS);
 
-        console.log("REGISTRY_INITIAL_OWNER:", initialOwner);
-
-        registry = new AllowList(initialOwner);
-
-        // Deploy AutomataDcapAttestationFee with deployer as owner
-        console.log("msg.sender:", msg.sender);
-        attestationFee = new AutomataDcapAttestationFee(msg.sender);
-
-        // Example: Call verifyQuoteWithAttestationFee with a sample quote
+        // Example: Call registerTEEService with a sample quote
         bytes memory sampleQuote = vm.readFileBinary("test/quote.raw");
 
-        (bool success, bytes memory output) =
-            registry.verifyQuoteWithAttestationFee(address(ETHEREUM_SEPOLIA_ATTESTATION_FEE_ADDRESS), sampleQuote);
-
-        if (!success) {
-            console.log(string(output));
-        }
+        allowlist.registerTEEService(sampleQuote);
 
         vm.stopBroadcast();
     }
diff --git a/src/AllowList.sol b/src/AllowList.sol
@@ -1,68 +1,61 @@
 // SPDX-License-Identifier: MIT
 pragma solidity 0.8.28;
 
-import "solmate/src/auth/Owned.sol";
 import {AutomataDcapAttestationFee} from "../lib/automata-dcap-attestation/evm/contracts/AutomataDcapAttestationFee.sol";
 import {TD10ReportBody} from "automata-dcap-attestation/contracts/types/V4Structs.sol";
-import {TD_REPORT10_LENGTH, TDX_TEE} from "automata-dcap-attestation/contracts/types/Constants.sol";
-import {BytesUtils} from "@automata-network/on-chain-pccs/utils/BytesUtils.sol";
+import {QuoteParser, WorkloadId} from "./utils/QuoteParser.sol";
 
 /**
  * @title AllowList
  * @dev A contract for managing trusted execution environment (TEE) identities and configurations
- * using Intel DCAP attestation
+ * using Automata's Intel DCAP attestation
  */
-contract AllowList is Owned {
-    using BytesUtils for bytes;
+contract AllowList {
 
-    // This is the number of bytes in the Output struct that come before the quoteBody.
-    // See the Output struct definition in the Automata DCAP Attestation repo:
-    // https://github.com/automata-network/automata-dcap-attestation/blob/evm-v1.0.0/evm/contracts/types/CommonStruct.sol#L113
-    // 13 bytes = quoteVersion (2 bytes) + tee (4 byte) + tcbStatus (1 byte) + fmspcBytes (6 bytes)
-    uint256 public constant SERIALIZED_OUTPUT_OFFSET = 13;
+    // Constants
 
     // Maximum size for byte arrays to prevent DoS attacks
     uint256 public constant MAX_BYTES_SIZE = 20 * 1024; // 20KB limit
 
-    // The TDX version of the quote Flashtestation's accepts
-    uint256 public constant ACCEPTED_TDX_VERSION = 4;
+    // Structs
 
     // TEE identity and status tracking
-    struct TEEDevice {
-        bytes32 publicKey; // Public key of the TEE device
-        uint64 lastActiveTime; // Timestamp of last activity
-        bool isActive; // Whether the device is currently active
+    struct RegisteredTEE {
+        uint64 registeredAt; // The most recent timestamp that the TEE last registered with a valid quote
+        WorkloadId workloadId; // The workloadID of the TEE device
+        bytes rawQuote; // The raw quote from the TEE device, which is stored to allow for future quote re-verification
     }
 
-    // Mapping from TEE identity to device information
-    mapping(bytes32 => TEEDevice) public teeDevices;
+    // Storage Variables
 
-    // State variables
-    string[] public instanceDomainNames;
+    // The address of the Automata DCAP Attestation contract, which verifies TEE quotes.
+    // This is deployed by Automata, and once set on the AllowList, it cannot be changed
+    address public attestationContract;
 
-    // Notes config and secrets locations
-    string[] public storageBackends;
-    // Maps config hash to config data and secrets for onchain DA
-    mapping(bytes32 => bytes) public artifacts;
-    // Maps identity to config hash
-    mapping(bytes32 => bytes32) public identityConfigMap;
+    // Used to track registered TEEs by the Ethereum address in the quote that originally registered them
+    mapping(address => RegisteredTEE) public registeredTEEs;
+
+    // Used to track registered TEEs by their (ethAddress, workloadId) pair,
+    // which is needed for when
+    mapping(WorkloadId => RegisteredTEE) public registeredTEEsByWorkloadId;
 
     // Events
-    event InstanceDomainRegistered(string domain, address registrar);
-    event StorageBackendSet(string location, address setter);
-    event StorageBackendRemoved(string location, address remover);
-    event ArtifactAdded(bytes32 configHash, address adder);
-    event IdentityConfigSet(bytes32 identity, bytes32 configHash, address setter);
 
+    event TEEServiceRegistered(address ethAddress, WorkloadId workloadId, bytes rawQuote, bool alreadyExists);
+
+    // Errors
+
+    error InvalidQuote(bytes output);
     error ByteSizeExceeded(uint256 size);
+    error TEEServiceAlreadyRegistered(address ethAddress, WorkloadId workloadId);
 
     /**
-     * @notice Constructor to set the the governance address, which
-     * is the only address that can register and de-register TEE devices
-     * and wipe the registry in case of a compromise
-     * @param governance The address of the governance contract
+     * Constructor to set the the Automata DCAP Attestation contract, which verifies TEE quotes
+     * @param _attestationContract The address of the attestation contract
      */
-    constructor(address governance) Owned(governance) {}
+    constructor(address _attestationContract) {
+        attestationContract = _attestationContract;
+    }
 
     /**
      * @dev Modifier to check if input bytes size is within limits
@@ -73,114 +66,77 @@ contract AllowList is Owned {
         _;
     }
 
-    function registerTEE(bytes memory quote) external onlyOwner {
-        // TODO: Implement
-    }
-
-    function deregisterTEE(bytes memory quote) external onlyOwner {
-        // TODO: Implement
-    }
-
-    function verifyFlashestationTransaction(bytes memory attestationTransaction) external {
-        // TODO: Implement
-        // 1. check signature against live builder keys
-        // 2. update liveness
-    }
-
-    function _updateLiveness() internal {
-        // TODO: Implement
-    }
-
     /**
-     * @notice Verifies a quote using AutomataDcapAttestationFee
-     * @dev TODO: move this logic into registerTEE
-     * @param attestationFeeContract The address of the AutomataDcapAttestationFee contract
-     * @param quote The DCAP quote to verify
+     * @notice Registers a TEE workload with a specific Ethereum address in the AllowList
+     * @notice The TEE must be registered with a quote whose validity is verified by the attestationContract
+     * @dev In order to mitigate DoS attacks, the quote must be less than 20KB
+     * @param rawQuote The raw quote from the TEE device. Must be a V4 TDX quote
      */
-    function verifyQuoteWithAttestationFee(address attestationFeeContract, bytes calldata quote)
-        external
-        limitBytesSize(quote)
-        returns (bool, bytes memory)
-    {
+    function registerTEEService(bytes calldata rawQuote) external limitBytesSize(rawQuote) {
         (bool success, bytes memory output) =
-            AutomataDcapAttestationFee(attestationFeeContract).verifyAndAttestOnChain(quote);
+            AutomataDcapAttestationFee(attestationContract).verifyAndAttestOnChain(rawQuote);
 
-        if (success) {
-            // check that quote is v4 and for TDX, otherwise in the next
-            // step the output will not have the byte length we expect and we'll fail
-            // to parse it, returning a unhelpful error message
-            checkTEEVersion(output);
-            checkTEEType(output);
+        if (!success) {
+            revert InvalidQuote(output);
+        }
+        // at this point we know this quote was generated by an update-to-date TEE, but we still
+        // need to make sure we can parse the workloadId and ethereum public key from the quote
 
-            // now we can safely decode the output into the TDX report body, from which we can extract
-            // the ethereum public key and compute the workloadID
-            TD10ReportBody memory td10ReportBodyStruct = parseTD10ReportBody(output);
+        // now we can safely decode the output into the TDX report body, from which we can extract
+        // the ethereum public key and compute the workloadID
+        TD10ReportBody memory td10ReportBodyStruct = QuoteParser.parseTD10ReportBody(output);
 
+        // extract the ethereum public key from the quote
+        address ethAddress = QuoteParser.extractEthereumAddress(td10ReportBodyStruct);
 
-            // TODO do flashtestations protocol, so far we've only verified that the quote is valid
-        }
+        // extract the workloadId from the quote
+        WorkloadId workloadId = QuoteParser.extractWorkloadId(td10ReportBodyStruct);
 
-        return (success, output);
+        // Register the address in the allowlist with the raw quote for future quote re-verification
+        bool previouslyRegistered = addAddress(workloadId, ethAddress, rawQuote);
+
+        emit TEEServiceRegistered(ethAddress, workloadId, rawQuote, previouslyRegistered);
     }
 
     /**
-     * @notice Parses a TD10ReportBody which contains all the data we need for
-     * registering a TEE, using the serializedOutput bytes generated by Automata's verification logic
-     * @param serializedOutput The serializedOutput bytes generated by Automata's verification logic
-     * @return report The parsed TD10ReportBody
-     * @dev Taken from Automata's DCAP Attestation repo:
-     * https://github.com/automata-network/automata-dcap-attestation/blob/evm-v1.0.0/evm/contracts/verifiers/V4QuoteVerifier.sol#L309
+     * @notice Adds a TEE to the allowlist
+     * @dev It's possible that a TEE has already registered with this ethereum address, but with a different workloadId.
+     * This is expected if the TEE gets restarted or upgraded and generates a new workloadId.
+     * It's also possible that the ethereum address and workloadId are the same, but the quote
+     * is different. This is expected if Intel releases a new set of DCAP Endorsements (i.e.
+     * a new TCB), in which case the quotes the TEE generates will be different.
+     * In both cases, we need to update the allowlist with the new quote.
+     * @param workloadId The workloadId of the TEE
+     * @param ethAddress The Ethereum address of the TEE
+     * @param rawQuote The raw quote from the TEE device
+     * @return previouslyRegistered Whether the TEE was previously registered
      */
-    function parseTD10ReportBody(bytes memory serializedOutput)
+    function addAddress(WorkloadId workloadId, address ethAddress, bytes calldata rawQuote)
         internal
-        pure
-        returns (TD10ReportBody memory report)
+        returns (bool previouslyRegistered)
     {
-        bytes memory rawReportBody = output.substring(SERIALIZED_OUTPUT_OFFSET, TD_REPORT10_LENGTH);
-
-        // note: because of the call to .substring above, we know that the length of rawReportBody is
-        // exactly TD_REPORT10_LENGTH, so we can safely call substring without checking the length
-
-        if (success) {
-            report.teeTcbSvn = bytes16(rawReportBody.substring(0, 16));
-            report.mrSeam = rawReportBody.substring(16, 48);
-            report.mrsignerSeam = rawReportBody.substring(64, 48);
-            report.seamAttributes = bytes8(rawReportBody.substring(112, 8));
-            report.tdAttributes = bytes8(rawReportBody.substring(120, 8));
-            report.xFAM = bytes8(rawReportBody.substring(128, 8));
-            report.mrTd = rawReportBody.substring(136, 48);
-            report.mrConfigId = rawReportBody.substring(184, 48);
-            report.mrOwner = rawReportBody.substring(232, 48);
-            report.mrOwnerConfig = rawReportBody.substring(280, 48);
-            report.rtMr0 = rawReportBody.substring(328, 48);
-            report.rtMr1 = rawReportBody.substring(376, 48);
-            report.rtMr2 = rawReportBody.substring(424, 48);
-            report.rtMr3 = rawReportBody.substring(472, 48);
-            report.reportData = rawReportBody.substring(520, 64);
+        // check if the TEE is already registered with the same ethereum address and workloadId,
+        // in which case this call is a no-op and possibly being called because of user error
+        if (
+            WorkloadId.unwrap(registeredTEEs[ethAddress].workloadId) == WorkloadId.unwrap(workloadId)
+                && keccak256(registeredTEEs[ethAddress].rawQuote) == keccak256(rawQuote)
+        ) {
+            revert TEEServiceAlreadyRegistered(ethAddress, workloadId);
         }
-    }
 
-    /**
-     * Parses and checks that the uint16 version from the quote header is one we accept
-     * @param rawReportBody The rawQuote bytes generated by Automata's verification logic
-     * @dev Automata currently only supports V3 and V4 TDX quotes
-     */
-    function checkTEEVersion(bytes memory rawReportBody) internal pure {
-        version = uint16(rawReportBody.substring(0, 2)); // uint16 is 2 bytes
-        if (version != ACCEPTED_TDX_VERSION) {
-            revert InvalidTEEVersion(version);
+        // this is a nice-to-have that signals to the user that the TEE was previously registered,
+        // but it's not strictly necessary
+        if (registeredTEEs[ethAddress].registeredAt > 0) {
+            previouslyRegistered = true;
         }
+
+        registeredTEEs[ethAddress] =
+            RegisteredTEE({registeredAt: uint64(block.timestamp), workloadId: workloadId, rawQuote: rawQuote});
     }
 
-    /**
-     * Parses and checks that the bytes4 tee type from the quote header is of type TDX
-     * @param rawReportBody The rawQuote bytes generated by Automata's verification logic
-     * @dev Automata currently only supports SGX and TDX TEE types
-     */
-    function checkTEEType(bytes memory rawReportBody) internal pure {
-        teeType = bytes4(rawReportBody.substring(2, 6)); // 4 bytes
-        if (teeType != TDX_TEE) {
-            revert InvalidTEEType(teeType);
-        }
+    function verifyFlashestationTransaction(bytes memory attestationTransaction) external {
+        // TODO: Implement
+        // 1. check signature against live builder keys
+        // 2. update liveness
     }
 }
diff --git a/src/utils/QuoteParser.sol b/src/utils/QuoteParser.sol
diff --git a/test/AllowList.t.sol b/test/AllowList.t.sol
