@inco/lightning 0.8.0-devnet-7 → 0.8.0-devnet-8

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@inco/lightning",
-  "version": "0.8.0-devnet-7",
+  "version": "0.8.0-devnet-8",
   "repository": "https://github.com/Inco-fhevm/inco-monorepo",
   "files": [
     "src/",

package/src/IncoLightning.sol

@@ -54,6 +54,7 @@ contract IncoLightning is
     /// @param owner The address that will own this contract and can authorize upgrades
     function initialize(address owner) public initializer {
         __Ownable_init(owner);
+        _setAcceptedVersion(2, true);
     }
 
     /// @notice Withdraws accumulated protocol fees to the owner
@@ -63,6 +64,18 @@ contract IncoLightning is
         _withdrawFeesTo(owner());
     }
 
+    /// @notice Adds a version to the accepted input versions whitelist.
+    /// @param version The version number to accept (must be non-negative).
+    function addAcceptedVersion(uint16 version) external onlyOwner {
+        _setAcceptedVersion(version, true);
+    }
+
+    /// @notice Removes a version from the accepted input versions whitelist.
+    /// @param version The version number to remove.
+    function removeAcceptedVersion(uint16 version) external onlyOwner {
+        _setAcceptedVersion(version, false);
+    }
+
     /// @notice Required for CreateX deterministic deployment
     /// @dev Empty fallback allows the contract to be deployed via CreateX's create2 mechanism
     fallback() external {}

package/src/interfaces/IIncoLightning.sol

@@ -21,4 +21,8 @@ interface IIncoLightning is
 
     function withdrawFees() external;
 
+    function addAcceptedVersion(uint16 version) external;
+
+    function removeAcceptedVersion(uint16 version) external;
+
 }

package/src/lightning-parts/EncryptedInput.sol

@@ -24,15 +24,67 @@ error ExternalHandleDoesNotMatchComputedHandle(
     address aclAddress,
     address userAddress,
     address contractAddress,
-    int32 version
+    uint16 version
 );
 
+/// @title EncryptedInputStorage
+/// @notice Storage for EncryptedInput contract using ERC-7201 namespaced storage pattern
+/// @dev Stores a whitelist of accepted versions for encrypted inputs.
+abstract contract EncryptedInputStorage {
+
+    /// @notice Storage for version whitelist, using ERC-7201 namespaced storage.
+    struct StorageForEncryptedInput {
+        /// @notice Mapping of accepted version numbers for encrypted inputs.
+        /// @dev Only versions in this mapping are accepted for new inputs.
+        mapping(uint16 => bool) acceptedVersions;
+    }
+
+    /// @notice Storage slot location using keccak256 of a unique namespace string
+    bytes32 private constant ENCRYPTED_INPUT_STORAGE_LOCATION = keccak256("inco.storage.EncryptedInput");
+
+    /// @notice Retrieves the storage struct from its dedicated slot
+    /// @dev Uses assembly to directly access the storage slot for gas efficiency
+    /// @return $ Reference to the StorageForEncryptedInput struct
+    function _getEncryptedInputStorage() internal pure returns (StorageForEncryptedInput storage $) {
+        bytes32 loc = ENCRYPTED_INPUT_STORAGE_LOCATION;
+        assembly {
+            $.slot := loc
+        }
+    }
+
+}
+
 /// @title EncryptedInput
 /// @notice Handles the submission of client-encrypted values to create on-chain encrypted handles.
 /// @dev Users encrypt values client-side and submit them with a prepended handle as a checksum.
 /// The contract verifies the handle matches the on-chain computation to ensure context consistency.
 /// This is a paid operation - users must send ETH to cover processing costs.
-abstract contract EncryptedInput is IEncryptedInput, BaseAccessControlList, HandleGeneration, Fee {
+abstract contract EncryptedInput is
+    IEncryptedInput,
+    EncryptedInputStorage,
+    BaseAccessControlList,
+    HandleGeneration,
+    Fee
+{
+
+    /// @notice Sets whether a version is accepted.
+    /// @param version The version number (must be non-negative).
+    /// @param accepted Whether the version should be accepted.
+    function _setAcceptedVersion(uint16 version, bool accepted) internal {
+        _getEncryptedInputStorage().acceptedVersions[version] = accepted;
+        if (accepted) {
+            emit VersionAccepted(version);
+        } else {
+            emit VersionRemoved(version);
+        }
+    }
+
+    /// @notice Checks whether a version is accepted.
+    /// @param version The version number to check.
+    /// @return True if the version is in the whitelist.
+    function isAcceptedVersion(uint16 version) public view returns (bool) {
+        return _getEncryptedInputStorage().acceptedVersions[version];
+    }
 
     /// @notice Emitted when a new encrypted input is submitted.
     /// @param result The generated handle for the encrypted value.
@@ -44,7 +96,7 @@ abstract contract EncryptedInput is IEncryptedInput, BaseAccessControlList, Hand
         bytes32 indexed result,
         address indexed contractAddress,
         address indexed user,
-        int32 version,
+        uint16 version,
         bytes ciphertext,
         uint256 eventId
     );
@@ -111,10 +163,13 @@ abstract contract EncryptedInput is IEncryptedInput, BaseAccessControlList, Hand
     function _newInput(bytes calldata input, address user, ETypes inputType) private returns (bytes32 handle) {
         // Since there is no sensible way to handle abi.decode errors (https://github.com/argotorg/solidity/issues/10381)
         // at least fail early on a conservative minimum length
-        require(input.length >= 68, "Input too short, should be at least 68 bytes");
+        require(input.length >= 68, InputLengthTooShort(input.length));
         // Parse the version from the first 4 bytes.
         bytes4 prefix = bytes4(input[:4]);
-        int32 version = int32(uint32(prefix));
+        uint16 version = uint16(uint32(prefix));
+
+        // Reject versions not in the whitelist
+        require(isAcceptedVersion(version), InvalidInputVersion(version));
         // Remove external handle prepended to input as a checksum
         (bytes32 externalHandle, bytes memory ciphertext) = abi.decode(input[4:], (bytes32, bytes));
         handle = getInputHandle(ciphertext, user, msg.sender, version, inputType);

package/src/lightning-parts/TEELifecycle.sol

@@ -23,8 +23,6 @@ contract TEELifecycleStorage {
 
     struct StorageForTeeLifecycle {
         IQuoteVerifier quoteVerifier;
-        BootstrapResult verifiedBootstrapResult;
-        bool bootstrapComplete;
         // @notice The list of approved TEE versions, each item is the MR_AGGREGATED 32 bytes
         // @dev The index of the TEE version is the version number
         bytes32[] approvedTeeVersions;
@@ -70,8 +68,6 @@ abstract contract TEELifecycle is
     // Events
     // @notice A new MR_AGGREGATED has been approved
     event NewTEEVersionApproved(uint256 indexed version, bytes32 indexed mrAggregated);
-    // @notice A previously approved TEE version has been removed
-    event TEEVersionRemoved(bytes32 indexed mrAggregated);
     event NewCovalidatorAdded(address covalidatorAddress, bytes quote);
     event BootstrapStageComplete(address indexed newEoaSigner, BootstrapResult bootstrapResult);
     // @notice Emitted to prove that an EOA has upgraded their TDX to a new
@@ -213,28 +209,6 @@ abstract contract TEELifecycle is
         emit NewTEEVersionApproved($.approvedTeeVersions.length - 1, newMrAggregated);
     }
 
-    /**
-     * @notice Removes a previously approved TEE version from the contract state
-     * @param mrAggregated - The MR_AGGREGATED bytes of the TEE version to remove
-     */
-    function removeApprovedTeeVersion(bytes32 mrAggregated) public onlyOwner {
-        StorageForTeeLifecycle storage $ = getTeeLifecycleStorage();
-        bool found = false;
-        for (uint256 i = 0; i < $.approvedTeeVersions.length; i++) {
-            if ($.approvedTeeVersions[i] == mrAggregated) {
-                // Shift all elements after index i to the left to preserve insertion order
-                for (uint256 j = i; j < $.approvedTeeVersions.length - 1; j++) {
-                    $.approvedTeeVersions[j] = $.approvedTeeVersions[j + 1];
-                }
-                $.approvedTeeVersions.pop();
-                found = true;
-                break;
-            }
-        }
-        require(found, TEEVersionNotFound());
-        emit TEEVersionRemoved(mrAggregated);
-    }
-
     function _isApprovedTeeVersion(bytes32 newMrAggregated) internal view returns (bool) {
         StorageForTeeLifecycle storage $ = getTeeLifecycleStorage();
         for (uint256 i = 0; i < $.approvedTeeVersions.length; i++) {
@@ -283,6 +257,40 @@ abstract contract TEELifecycle is
         return getTeeLifecycleStorage().networkPubkey.length > 0;
     }
 
+    /**
+     * @notice Resets the contract state to the initial state.
+     * @dev This function performs a destructive operation and should be used carefully. In particular:
+     * 1) This function will remove all signers and reset the network pubkey and approved TEE versions, requiring a full bootstrap to restore normal operation.
+     * 2) The contract effectively enters a maintenance mode where any operations that depend on signers or the network key will fail until re-bootstrap is completed.
+     * 3) This function is intended to be called only as part of a coordinated key rotation or recovery process, with off-chain components updated accordingly.
+     * @dev This function is only callable by the owner.
+     */
+    function reset() public onlyOwner {
+        getTeeLifecycleStorage().networkPubkey = bytes("");
+        getTeeLifecycleStorage().approvedTeeVersions = new bytes32[](0);
+
+        // Directly access SignatureVerifier storage to reset signers and threshold
+        // We bypass the normal functions because:
+        // 1. removeSigner is external and has threshold validation
+        // 2. setThreshold requires newThreshold > 0
+        // We emit RemovedSignatureVerifier and ThresholdChanged so off-chain
+        // systems tracking signer changes via events stay in sync.
+        StorageForSigVerifier storage sigStorage = getSigVerifierStorage();
+        uint256 oldThreshold = sigStorage.threshold;
+        if (oldThreshold > 0) {
+            emit ThresholdChanged(oldThreshold, 0);
+        }
+        sigStorage.threshold = 0;
+        // In theory, this could revert if there are many signers. But in
+        // practice, this should not be a problem.
+        while (sigStorage.signers.length > 0) {
+            address signer = sigStorage.signers[sigStorage.signers.length - 1];
+            sigStorage.isSigner[signer] = false;
+            sigStorage.signers.pop();
+            emit RemovedSignatureVerifier(signer);
+        }
+    }
+
     /**
      * @notice From https://github.com/automata-network/automata-dcap-attestation/blob/evm-v1.0.0/evm/contracts/AttestationEntrypointBase.sol#L103
      * @notice full on-chain verification for an attestation
@@ -381,7 +389,7 @@ abstract contract TEELifecycle is
 
     /**
      * @notice Computes the MR_AGGREGATED from the TD10 report body. It is defined as:
-     * KECCAK256(mrTd ++ rtMr0 ++ rtMr1 ++ rtMr2 ++ rtMr3)
+     * as KECCAK256(abi.encode(MRTD, RTMR0, RTMR1, RTMR2))
      * @param mrTd - The MR_TD bytes
      * @param rtMr0 - The RTMR0 bytes
      * @param rtMr1 - The RTMR1 bytes
@@ -393,7 +401,7 @@ abstract contract TEELifecycle is
         pure
         returns (bytes32)
     {
-        bytes memory message = abi.encodePacked(mrTd, rtMr0, rtMr1, rtMr2);
+        bytes memory message = abi.encode(mrTd, rtMr0, rtMr1, rtMr2);
         return keccak256(message);
     }
 

package/src/lightning-parts/interfaces/IEncryptedInput.sol

@@ -5,6 +5,12 @@ import {euint256, ebool, eaddress} from "../../Types.sol";
 
 interface IEncryptedInput {
 
+    error InvalidInputVersion(uint16 version);
+    error InputLengthTooShort(uint256 length);
+
+    event VersionAccepted(uint16 version);
+    event VersionRemoved(uint16 version);
+
     function newEuint256(bytes calldata ciphertext, address user) external payable returns (euint256 newValue);
     function newEbool(bytes calldata ciphertext, address user) external payable returns (ebool newValue);
     function newEaddress(bytes calldata ciphertext, address user) external payable returns (eaddress newValue);

package/src/lightning-parts/primitives/HandleGeneration.sol

@@ -41,7 +41,7 @@ contract HandleGeneration is IHandleGeneration, HandleMetadata {
         bytes memory ciphertext,
         address user,
         address contractAddress,
-        int32 version,
+        uint16 version,
         ETypes inputType
     ) internal view returns (bytes32 generatedHandle) {
         return getInputHandle(ciphertext, address(this), user, contractAddress, version, inputType);
@@ -61,7 +61,7 @@ contract HandleGeneration is IHandleGeneration, HandleMetadata {
         address executorAddress,
         address user,
         address contractAddress,
-        int32 version,
+        uint16 version,
         ETypes inputType
     ) internal view returns (bytes32 generatedHandle) {
         // Here we ensure that our hashing scheme is binary-compatible between IncoLightning and IncoFhevm, this helps

package/src/lightning-parts/primitives/test/SignatureVerifier.t.sol

@@ -17,7 +17,7 @@ contract TestSignatureVerifier is TestUtils, SignatureVerifier {
     // Helper function to create sorted signatures
     // Takes a digest and an array of private keys, generates signatures,
     // and returns them sorted by signer address in ascending order
-    function getSortedSignatures(bytes32 digest, uint256[] memory privKeys) internal returns (bytes[] memory) {
+    function getSortedSignatures(bytes32 digest, uint256[] memory privKeys) internal pure returns (bytes[] memory) {
         bytes[] memory signatures = new bytes[](privKeys.length);
         address[] memory signers = new address[](privKeys.length);
 

package/src/lightning-parts/test/HandleMetadata.t.sol

@@ -22,6 +22,7 @@ import {VerifierAddressGetter} from "../primitives/VerifierAddressGetter.sol";
 import {FEE} from "../Fee.sol";
 import {HandleAlreadyExists} from "../../Errors.sol";
 import {ExternalHandleDoesNotMatchComputedHandle} from "../EncryptedInput.sol";
+import {IEncryptedInput} from "../interfaces/IEncryptedInput.sol";
 
 contract TestHandleMetadata is
     EIP712,
@@ -32,7 +33,9 @@ contract TestHandleMetadata is
     EncryptedInput
 {
 
-    constructor() EIP712("", "") VerifierAddressGetter(address(0)) {}
+    constructor() EIP712("", "") VerifierAddressGetter(address(0)) {
+        _setAcceptedVersion(2, true);
+    }
 
     function testTypeAssignment() public pure {
         bytes32 someHandle = bytes32(keccak256("someHandle"));
@@ -105,16 +108,16 @@ contract TestHandleMetadata is
     {
         // We need a single word here to get correct encoding
         bytes memory ciphertext = abi.encode(word);
+        uint16 version = 2; // version - X-Wing
         bytes32 handle = getInputHandle(
             ciphertext,
             address(this),
             user,
             contractAddress,
-            0,
-            /* unspecified */
+            version, // version - X-Wing
             inputType
         );
-        input = abi.encodePacked(int32(0), abi.encode(handle, ciphertext));
+        input = abi.encodePacked(uint32(version), abi.encode(handle, ciphertext));
     }
 
     // ============ Tests for HandleGeneration functions ============
@@ -155,6 +158,53 @@ contract TestHandleMetadata is
         return newInputNotPaying(input, user, inputType);
     }
 
+    function testNewInputUnwhitelistedVersion() public {
+        address self = address(this);
+        bytes32 ciphertextData = keccak256(abi.encodePacked("unwhitelisted_version"));
+        bytes memory ciphertext = abi.encode(ciphertextData);
+        uint16 unwhitelistedVersion = 0;
+        bytes32 handle =
+            getInputHandle(ciphertext, address(this), self, address(this), unwhitelistedVersion, ETypes.Uint256);
+        bytes memory badInput = abi.encodePacked(uint32(unwhitelistedVersion), abi.encode(handle, ciphertext));
+
+        vm.expectRevert(abi.encodeWithSelector(IEncryptedInput.InvalidInputVersion.selector, unwhitelistedVersion));
+        this.exposedNewInputNotPaying(badInput, self, ETypes.Uint256);
+    }
+
+    function testNewInputVersionAddedThenRemoved() public {
+        address self = address(this);
+        bytes32 ciphertextData = keccak256(abi.encodePacked("version_toggle"));
+        bytes memory ciphertext = abi.encode(ciphertextData);
+        uint16 version = 3;
+
+        bytes32 handle = getInputHandle(ciphertext, address(this), self, address(this), version, ETypes.Uint256);
+        bytes memory input = abi.encodePacked(uint32(version), abi.encode(handle, ciphertext));
+
+        // Version 3 is not accepted yet
+        vm.expectRevert(abi.encodeWithSelector(IEncryptedInput.InvalidInputVersion.selector, version));
+        this.exposedNewInputNotPaying(input, self, ETypes.Uint256);
+
+        // Whitelist version 3
+        _setAcceptedVersion(3, true);
+
+        // Now it should succeed
+        bytes32 resultHandle = this.exposedNewInputNotPaying(input, self, ETypes.Uint256);
+        assert(typeOf(resultHandle) == ETypes.Uint256);
+
+        // Remove version 3
+        _setAcceptedVersion(3, false);
+
+        // Use different ciphertext to avoid HandleAlreadyExists
+        bytes32 ciphertextData2 = keccak256(abi.encodePacked("version_toggle_2"));
+        bytes memory ciphertext2 = abi.encode(ciphertextData2);
+        bytes32 handle2 = getInputHandle(ciphertext2, address(this), self, address(this), version, ETypes.Uint256);
+        bytes memory input2 = abi.encodePacked(uint32(version), abi.encode(handle2, ciphertext2));
+
+        // Should revert again
+        vm.expectRevert(abi.encodeWithSelector(IEncryptedInput.InvalidInputVersion.selector, version));
+        this.exposedNewInputNotPaying(input2, self, ETypes.Uint256);
+    }
+
     function testNewInputNotPaying() public {
         address self = address(this);
         bytes32 ciphertextData = keccak256(abi.encodePacked("notpaying_ciphertext"));
@@ -172,10 +222,10 @@ contract TestHandleMetadata is
         returns (bytes memory input)
     {
         bytes memory ciphertext = abi.encode(word);
-        int32 version = 0; // unspecified
+        uint16 version = 2; // version - X-Wing
         // For external calls via this., msg.sender is address(this)
         bytes32 handle = getInputHandle(ciphertext, address(this), user, address(this), version, inputType);
-        input = abi.encodePacked(version, abi.encode(handle, ciphertext));
+        input = abi.encodePacked(uint32(version), abi.encode(handle, ciphertext));
     }
 
     // ============ Tests for EncryptedInput error branches ============
@@ -184,7 +234,7 @@ contract TestHandleMetadata is
         address self = address(this);
         // Input less than 64 bytes should revert
         bytes memory shortInput = hex"deadbeef";
-        vm.expectRevert("Input too short, should be at least 68 bytes");
+        vm.expectRevert(abi.encodeWithSelector(IEncryptedInput.InputLengthTooShort.selector, shortInput.length));
         this.exposedNewInputNotPaying(shortInput, self, ETypes.Uint256);
     }
 
@@ -194,8 +244,8 @@ contract TestHandleMetadata is
         bytes memory ciphertext = abi.encode(ciphertextData);
         // Create input with wrong handle (just a random bytes32)
         bytes32 wrongHandle = bytes32(uint256(12345));
-        int32 version = 0; // unspecified
-        bytes memory badInput = abi.encodePacked(version, abi.encode(wrongHandle, ciphertext));
+        uint16 version = 2; // version - X-Wing
+        bytes memory badInput = abi.encodePacked(uint32(version), abi.encode(wrongHandle, ciphertext));
 
         // Compute the expected handle for the error message
         bytes32 expectedHandle = getInputHandle(ciphertext, address(this), self, address(this), version, ETypes.Uint256);

package/src/test/FakeIncoInfra/FakeIncoInfraBase.sol

@@ -19,16 +19,16 @@ contract FakeIncoInfraBase is TestUtils, KVStore, HandleGeneration {
     {
         // We need a single word here to get correct encoding
         bytes memory ciphertext = abi.encode(word);
+        uint16 version = 2; // version - X-Wing
         bytes32 handle = getInputHandle(
             ciphertext,
             address(inco),
             user,
             contractAddress,
-            0,
-            /* unspecified */
+            version, // version - X-Wing
             inputType
         );
-        input = abi.encodePacked(int32(0), abi.encode(handle, ciphertext));
+        input = abi.encodePacked(uint32(version), abi.encode(handle, ciphertext));
     }
 
     function fakePrepareEuint256Ciphertext(uint256 value, address userAddress, address contractAddress)

package/src/test/FakeIncoInfra/MockRemoteAttestation.sol

@@ -60,7 +60,8 @@ contract MockRemoteAttestation is TestUtils {
             hex"010101010101010101010101010101010101010101010101010101010101010101010101010101010101010101010101";
         // See DEFAULT_MR_AGGREGATED in attestation/src/remote_attestation.rs to
         // see the calculation of the default value.
-        mrAggregated = hex"c3a67bac251d4946d7b17481d39631676042fe3afab06e70c22105ad8383c19f";
+        // Note: This uses abi.encode (not encodePacked) to avoid hash collision vulnerabilities.
+        mrAggregated = hex"3d48a1faa8620d86ae037f4fd6746987733d085314b3cd5d5d074ade8bab6ebd";
         bootstrapResult = BootstrapResult({networkPubkey: networkPubkey});
 
         quote = createQuote(mrtd, signer);

package/src/test/TEELifecycle/TEELifecycleMockTest.t.sol

@@ -24,7 +24,8 @@ contract TEELifecycleMockTest is MockRemoteAttestation, TEELifecycle {
         hex"010101010101010101010101010101010101010101010101010101010101010101010101010101010101010101010101";
     // See DEFAULT_MR_AGGREGATED in attestation/src/remote_attestation.rs to
     // see the calculation of the default value.
-    bytes32 testMrAggregated = hex"c3a67bac251d4946d7b17481d39631676042fe3afab06e70c22105ad8383c19f";
+    // Note: This uses abi.encode (not encodePacked) to avoid hash collision vulnerabilities.
+    bytes32 testMrAggregated = hex"3d48a1faa8620d86ae037f4fd6746987733d085314b3cd5d5d074ade8bab6ebd";
 
     function setUp() public {
         getTeeLifecycleStorage().quoteVerifier = new FakeQuoteVerifier();
@@ -133,61 +134,6 @@ contract TEELifecycleMockTest is MockRemoteAttestation, TEELifecycle {
         vm.stopPrank();
     }
 
-    function testRemoveApprovedTeeVersionPreservesOrder() public {
-        bytes32 mrAggregated1 = hex"1111111111111111111111111111111111111111111111111111111111111111";
-        bytes32 mrAggregated2 = hex"2222222222222222222222222222222222222222222222222222222222222222";
-        bytes32 mrAggregated3 = hex"3333333333333333333333333333333333333333333333333333333333333333";
-
-        vm.startPrank(this.owner());
-
-        // Add three versions
-        this.approveNewTeeVersion(mrAggregated1);
-        this.approveNewTeeVersion(mrAggregated2);
-        this.approveNewTeeVersion(mrAggregated3);
-
-        // Verify all exist in order
-        assertEq(this.approvedTeeVersions(0), mrAggregated1);
-        assertEq(this.approvedTeeVersions(1), mrAggregated2);
-        assertEq(this.approvedTeeVersions(2), mrAggregated3);
-
-        // Remove the middle one (mrAggregated2)
-        this.removeApprovedTeeVersion(mrAggregated2);
-
-        // Verify insertion order is preserved: mrAggregated1 stays at 0, mrAggregated3 shifts to 1
-        assertEq(this.approvedTeeVersions(0), mrAggregated1);
-        assertEq(this.approvedTeeVersions(1), mrAggregated3);
-
-        // Verify index 2 is now out of bounds
-        vm.expectRevert(TEELifecycle.IndexOutOfBounds.selector);
-        this.approvedTeeVersions(2);
-
-        vm.stopPrank();
-    }
-
-    function testRemoveApprovedTeeVersionNotFound() public {
-        bytes32 nonExistentMrAggregated = hex"aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa";
-
-        vm.startPrank(this.owner());
-        vm.expectRevert(TEELifecycle.TEEVersionNotFound.selector);
-        this.removeApprovedTeeVersion(nonExistentMrAggregated);
-        vm.stopPrank();
-    }
-
-    function testRemoveApprovedTeeVersionOnlyOwner() public {
-        bytes32 mrAggregated = hex"1111111111111111111111111111111111111111111111111111111111111111";
-
-        vm.startPrank(this.owner());
-        this.approveNewTeeVersion(mrAggregated);
-        vm.stopPrank();
-
-        // Try to remove as non-owner
-        address nonOwner = address(0x1234);
-        vm.startPrank(nonOwner);
-        vm.expectRevert();
-        this.removeApprovedTeeVersion(mrAggregated);
-        vm.stopPrank();
-    }
-
     // Helper function to create a successful bootstrap result
     function successfulBootstrapResult()
         internal
@@ -449,4 +395,83 @@ contract TEELifecycleMockTest is MockRemoteAttestation, TEELifecycle {
         );
     }
 
+    // ============ Tests for reset() ============
+
+    function testReset_OnlyOwner() public {
+        address nonOwner = address(0x1234);
+        vm.startPrank(nonOwner);
+        vm.expectRevert();
+        this.reset();
+        vm.stopPrank();
+    }
+
+    function testReset_WithMultipleSigners() public {
+        // Complete bootstrap with first signer
+        (
+            BootstrapResult memory bootstrapResult,,
+            address bootstrapPartyAddress,
+            bytes memory quote,
+            bytes memory signature,
+            bytes32 mrAggregated
+        ) = successfulBootstrapResult();
+
+        vm.startPrank(this.owner());
+        this.approveNewTeeVersion(mrAggregated);
+        this.verifyBootstrapResult(bootstrapResult, quote, signature);
+
+        // Add a second node
+        (uint256 newNodePrivkey, address newNodeAddress) = getLabeledKeyPair("newNode");
+        AddNodeResult memory addNodeResult = AddNodeResult({networkPubkey: testNetworkPubkey});
+        bytes memory addNodeSignature = signAddNodeResult(addNodeResult, newNodePrivkey);
+        bytes memory addNodeQuote = createQuote(testMrtd, newNodeAddress);
+        this.verifyAddNodeResult(mrAggregated, addNodeResult, addNodeQuote, addNodeSignature);
+
+        // Verify state before reset
+        assertTrue(this.isBootstrapComplete(), "Bootstrap should be complete before reset");
+        assertEq(this.networkPubkey(), testNetworkPubkey, "Network pubkey should be set before reset");
+        assertEq(this.approvedTeeVersions(0), mrAggregated, "Approved TEE version should exist before reset");
+        assertEq(this.getSignersCount(), 2, "Should have 2 signers before reset");
+        assertTrue(this.isSigner(bootstrapPartyAddress), "First signer should exist");
+        assertTrue(this.isSigner(newNodeAddress), "Second signer should exist");
+
+        // Call reset
+        this.reset();
+
+        // Verify all state has been cleared
+        assertFalse(this.isBootstrapComplete(), "Bootstrap should not be complete after reset");
+        assertEq(this.networkPubkey().length, 0, "Network pubkey should be empty after reset");
+        assertEq(this.getSignersCount(), 0, "Should have 0 signers after reset");
+        assertFalse(this.isSigner(bootstrapPartyAddress), "First signer should be removed");
+        assertFalse(this.isSigner(newNodeAddress), "Second signer should be removed");
+        assertEq(this.getThreshold(), 0, "Threshold should be 0 after reset");
+
+        // Verify approved TEE versions array is empty
+        vm.expectRevert(TEELifecycle.IndexOutOfBounds.selector);
+        this.approvedTeeVersions(0);
+
+        vm.stopPrank();
+    }
+
+    function testReset_AllowsNewBootstrap() public {
+        // Complete bootstrap
+        (BootstrapResult memory bootstrapResult,,, bytes memory quote, bytes memory signature, bytes32 mrAggregated) =
+            successfulBootstrapResult();
+
+        vm.startPrank(this.owner());
+        this.approveNewTeeVersion(mrAggregated);
+        this.verifyBootstrapResult(bootstrapResult, quote, signature);
+        assertTrue(this.isBootstrapComplete(), "Bootstrap should be complete");
+
+        // Reset the contract
+        this.reset();
+        assertFalse(this.isBootstrapComplete(), "Bootstrap should not be complete after reset");
+
+        // Should be able to bootstrap again
+        this.approveNewTeeVersion(mrAggregated);
+        this.verifyBootstrapResult(bootstrapResult, quote, signature);
+        assertTrue(this.isBootstrapComplete(), "Should be able to bootstrap again after reset");
+
+        vm.stopPrank();
+    }
+
 }

package/src/test/TestDeploy.t.sol

@@ -5,6 +5,7 @@ import {Test} from "forge-std/Test.sol";
 import {TrivialEncryption} from "../lightning-parts/TrivialEncryption.sol";
 import {ETypes} from "../Types.sol";
 import {UUPSUpgradeable} from "@openzeppelin/contracts/proxy/utils/UUPSUpgradeable.sol";
+import {OwnableUpgradeable} from "@openzeppelin/contracts-upgradeable/access/OwnableUpgradeable.sol";
 import {inco} from "../Lib.sol";
 import {IncoTest} from "./IncoTest.sol";
 
@@ -26,6 +27,7 @@ contract TestDeploy is Test, IncoTest {
         vm.expectEmit(false, false, true, false, address(inco));
         emit TrivialEncryption.TrivialEncrypt(bytes32(uint256(1)), bytes32(uint256(1)), ETypes.Bool, 0);
         inco.asEbool(true);
+        assertTrue(inco.isAcceptedVersion(2));
     }
 
     function testUpgrade() public {
@@ -35,4 +37,30 @@ contract TestDeploy is Test, IncoTest {
         assertEq(ReturnTwo(address(inco)).getTwo(), 2);
     }
 
+    function testAddAcceptedVersion() public {
+        assertFalse(inco.isAcceptedVersion(42));
+        vm.prank(owner);
+        inco.addAcceptedVersion(42);
+        assertTrue(inco.isAcceptedVersion(42));
+    }
+
+    function testAddAcceptedVersionNotOwner() public {
+        vm.expectRevert(abi.encodeWithSelector(OwnableUpgradeable.OwnableUnauthorizedAccount.selector, alice));
+        vm.prank(alice);
+        inco.addAcceptedVersion(42);
+    }
+
+    function testRemoveAcceptedVersion() public {
+        assertFalse(inco.isAcceptedVersion(42));
+        vm.prank(owner);
+        inco.removeAcceptedVersion(42); // removing a non-existent version should be no-op
+        assertFalse(inco.isAcceptedVersion(42));
+    }
+
+    function testRemoveAcceptedVersionNotOwner() public {
+        vm.expectRevert(abi.encodeWithSelector(OwnableUpgradeable.OwnableUnauthorizedAccount.selector, alice));
+        vm.prank(alice);
+        inco.removeAcceptedVersion(42);
+    }
+
 }