write first pair of tests for AllowList.sol

Author: Melvillian

foundry.toml

@@ -5,7 +5,7 @@ solc_version = '0.8.28'
 optimizer_runs = 44444444
 via_ir = true
 ffi = true
-fs_permissions = [{ access = "read", path = "foundry-out/" }, {access = "read", path = "test/quote.raw"}]
+fs_permissions = [{ access = "read", path = "foundry-out/" }, {access = "read", path = "test/raw_tdx_quotes/"}, {access = "read", path = "script/quote.raw/"}]
 evm_version = "prague"
 gas_limit = "3000000000"
 fuzz_runs = 10_000

script/AllowList.s.sol

@@ -20,8 +20,10 @@ contract AllowListScript is Script {
 
         allowlist = new AllowList(ETHEREUM_SEPOLIA_ATTESTATION_FEE_ADDRESS);
 
-        // Example: Call registerTEEService with a sample quote
-        bytes memory sampleQuote = vm.readFileBinary("test/quote.raw");
+        // Example: Call registerTEEService with a sample quote whose user data is simply `00`
+        bytes memory sampleQuote = vm.readFileBinary(
+            "test/raw_tdx_quotes/bf42a348f49c9f8ab2ef750ddaffd294c45d8adf947e4d1a72158dcdbd6997c2ca7decaa1ad42648efebdfefe79cbc1b63eb2499fe2374648162fd8f5245f446/quote.bin"
+        );
 
         allowlist.registerTEEService(sampleQuote);
 

src/AllowList.sol

@@ -1,36 +1,35 @@
 // SPDX-License-Identifier: MIT
 pragma solidity 0.8.28;
 
-import {AutomataDcapAttestationFee} from "../lib/automata-dcap-attestation/evm/contracts/AutomataDcapAttestationFee.sol";
-import {TD10ReportBody} from "automata-dcap-attestation/contracts/types/V4Structs.sol";
+import {IAttestation} from "./interfaces/IAttestation.sol";
 import {QuoteParser, WorkloadId} from "./utils/QuoteParser.sol";
+import {TD10ReportBody} from "automata-dcap-attestation/contracts/types/V4Structs.sol";
+
+// TEE identity and status tracking
+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
+}
 
 /**
  * @title AllowList
  * @dev A contract for managing trusted execution environment (TEE) identities and configurations
  * using Automata's Intel DCAP attestation
  */
 contract AllowList {
-
     // Constants
 
     // Maximum size for byte arrays to prevent DoS attacks
     uint256 public constant MAX_BYTES_SIZE = 20 * 1024; // 20KB limit
 
     // Structs
 
-    // TEE identity and status tracking
-    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
-    }
-
     // Storage Variables
 
     // 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;
+    IAttestation public attestationContract;
 
     // Used to track registered TEEs by the Ethereum address in the quote that originally registered them
     mapping(address => RegisteredTEE) public registeredTEEs;
@@ -54,7 +53,7 @@ contract AllowList {
      * @param _attestationContract The address of the attestation contract
      */
     constructor(address _attestationContract) {
-        attestationContract = _attestationContract;
+        attestationContract = IAttestation(_attestationContract);
     }
 
     /**
@@ -73,8 +72,7 @@ contract AllowList {
      * @param rawQuote The raw quote from the TEE device. Must be a V4 TDX quote
      */
     function registerTEEService(bytes calldata rawQuote) external limitBytesSize(rawQuote) {
-        (bool success, bytes memory output) =
-            AutomataDcapAttestationFee(attestationContract).verifyAndAttestOnChain(rawQuote);
+        (bool success, bytes memory output) = attestationContract.verifyAndAttestOnChain(rawQuote);
 
         if (!success) {
             revert InvalidQuote(output);

src/interfaces/IAttestation.sol

@@ -0,0 +1,6 @@
+// SPDX-License-Identifier: MIT
+pragma solidity 0.8.28;
+
+interface IAttestation {
+    function verifyAndAttestOnChain(bytes calldata rawQuote) external returns (bool, bytes memory);
+}

src/utils/QuoteParser.sol

@@ -120,7 +120,6 @@ library QuoteParser {
         );
     }
 
-
     /**
      * 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

test/AllowList.t.sol

@@ -2,16 +2,51 @@
 pragma solidity 0.8.28;
 
 import {Test, console} from "forge-std/Test.sol";
-import {AllowList} from "../src/AllowList.sol";
+import {AllowList, RegisteredTEE} from "../src/AllowList.sol";
+import {QuoteParser, WorkloadId} from "../src/utils/QuoteParser.sol";
+import {MockAutomataDcapAttestationFee} from "./mocks/MockAutomataDcapAttestationFee.sol";
 
 contract AllowListTest is Test {
-    AllowList public registry;
+    AllowList public allowlist;
+    MockAutomataDcapAttestationFee public attestationContract;
 
     function setUp() public {
-        registry = new AllowList(address(this));
+        // deploy a fresh set of test contracts before each test
+        attestationContract = new MockAutomataDcapAttestationFee();
+        allowlist = new AllowList(address(attestationContract));
     }
 
-    function test_registerTEEService() public {}
+    function test_succesful_registerTEEService() public {
+        attestationContract.setSuccess(true);
+        bytes memory output = vm.readFileBinary("test/raw_tdx_quotes/bf42a348f49c9f8ab2ef750ddaffd294c45d8adf947e4d1a72158dcdbd6997c2ca7decaa1ad42648efebdfefe79cbc1b63eb2499fe2374648162fd8f5245f446/output.bin");
+        attestationContract.setOutput(output);
 
-    function testFuzz_registerTEEService(bytes memory _quote) public {}
+        address expectedAddress = 0xf200f222043C5bC6c70AA6e35f5C5FDe079F3a03;
+        uint64 expectedRegisteredAt = uint64(block.timestamp);
+        // note: this is taken directly from the output of QuoteParser.extractWorkloadId, so it's not
+        // a good test of the QuoteParser.extractWorkloadId function, but it's a good regression test
+        WorkloadId expectedWorkloadId = WorkloadId.wrap(0x3db2a20cfbe18050a9996d00a94ad82346bc340a438906d04c5a1f960f31a416);
+
+        bytes memory quote = vm.readFileBinary("test/raw_tdx_quotes/bf42a348f49c9f8ab2ef750ddaffd294c45d8adf947e4d1a72158dcdbd6997c2ca7decaa1ad42648efebdfefe79cbc1b63eb2499fe2374648162fd8f5245f446/quote.bin");
+        vm.expectEmit(address(allowlist));
+        emit AllowList.TEEServiceRegistered(expectedAddress, expectedWorkloadId, quote, false);
+        allowlist.registerTEEService(quote);
+
+        (uint64 registeredAt, WorkloadId workloadId, bytes memory rawQuote) =
+            allowlist.registeredTEEs(expectedAddress);
+        vm.assertEq(registeredAt, expectedRegisteredAt, "Registered at mismatch");
+        vm.assertEq(WorkloadId.unwrap(workloadId), WorkloadId.unwrap(expectedWorkloadId), "Workload ID mismatch");
+        vm.assertEq(rawQuote, quote, "Raw quote mismatch");
+    }
+
+    function test_reverts_with_invalid_quote_registerTEEService() public {
+        attestationContract.setSuccess(false);
+        // don't bother setting the output, since it should revert before it's used
+
+        vm.expectPartialRevert(AllowList.InvalidQuote.selector); // the "partial" just means we don't care about the bytes argument to InvalidQuote(bytes)
+        bytes memory quote = vm.readFileBinary("test/raw_tdx_quotes/bf42a348f49c9f8ab2ef750ddaffd294c45d8adf947e4d1a72158dcdbd6997c2ca7decaa1ad42648efebdfefe79cbc1b63eb2499fe2374648162fd8f5245f446/quote.bin");
+        allowlist.registerTEEService(quote);
+    }
+
+    function testFuzz_registerTEEService(bytes memory _quote) public { /** TODO: fuzz things that are fuzzable **/}
 }

test/mocks/MockAutomataDcapAttestationFee.sol

@@ -0,0 +1,29 @@
+// SPDX-License-Identifier: MIT
+pragma solidity 0.8.28;
+
+/**
+ * @title MockAutomataDcapAttestationFee
+ * @dev A mock implementation of the AutomataDcapAttestationFee contract for testing
+ * @dev This mock allows us to control the output of the verifyAndAttestOnChain function,
+ *      and skip dealing with the vast complexity of Automata's DCAP Attestation contract
+ * @dev This is useful for testing the AllowList contract with different quote and output values
+ *      without having to deploy the AutomataDcapAttestationFee contract
+ */
+contract MockAutomataDcapAttestationFee {
+    bool public success;
+    bytes public output;
+
+    constructor() {}
+
+    function verifyAndAttestOnChain(bytes calldata /* rawQuote */ ) external view returns (bool, bytes memory) {
+        return (success, output);
+    }
+
+    function setSuccess(bool _success) public {
+        success = _success;
+    }
+
+    function setOutput(bytes memory _output) public {
+        output = _output;
+    }
+}

test/raw_tdx_quotes/README.md

@@ -0,0 +1,14 @@
+# Raw TDX Quotes (For Testing)
+
+This directory contains TDX quote and serialized Output values that we use for testing
+
+## File Structure
+
+Each directory within `raw_tdx_quotes/` is named by the 64-byte hex uncompressed public key that was used for the `TD10ReportBody.reportData` field. Each of these public key directories contains a 
+
+- `quote.bin`: which is a TDX quote binary file whose `TD10ReportBody.reportData` is the public key named by its parent directory
+- `output.bin`: which is the binary output returned by `AutomataDcapAttestationFee.verifyAndAttestOnChain` when called on the Ethereum Sepolia network
+
+### Helpers
+
+- `hex2bin.py`: a simple python script for writing string hex data (such as `0xdeadbeef`) in its binary form to a file. We use this to take the 0x-prefixed hex string data given to us in events from running Forge scripts and write it in file form that can be read with the Forge cheatcode `vm.readFileBinary`