Update preconf registry 2

.gitignore

@@ -1,4 +1,5 @@
 /Node/target
 .vscode
 .env
-tools/tx_spammer/venv
+tools/tx_spammer/venv
+.DS_Store

SmartContracts/src/avs/PreconfRegistry.sol

@@ -1,6 +1,7 @@
 // SPDX-License-Identifier: UNLICENSED
 pragma solidity 0.8.25;
 
+import {PreconfConstants} from "./libraries/PreconfConstants.sol";
 import {BLS12381} from "../libraries/BLS12381.sol";
 import {BLSSignatureChecker} from "./utils/BLSSignatureChecker.sol";
 import {IPreconfRegistry} from "../interfaces/IPreconfRegistry.sol";
@@ -15,23 +16,26 @@ contract PreconfRegistry is IPreconfRegistry, ISignatureUtils, BLSSignatureCheck
     uint256 internal nextPreconferIndex;
 
     // Maps the preconfer's address to an index that may change over the lifetime of a preconfer
-    mapping(address => uint256) public preconferToIndex;
+    mapping(address preconfer => uint256 index) internal preconferToIndex;
 
-    // Maps the preconfer's address to an incrementing nonce used for validator signatures
-    mapping(address => uint256) public preconferToNonce;
+    // Maps an index to the preconfer's address
+    // We need this mapping to deregister a preconfer in O(1) time.
+    // While it may also be done by just using the above map and sending a "witness" that is calculated offchain,
+    // we ideally do not want the node to maintain historical state.
+    mapping(uint256 index => address preconfer) internal indexToPreconfer;
 
-    // Maps the preconfer's ecsda and one associated BLS public key hash to the timestamp
-    // at which the key hash was added
-    mapping(address => mapping(bytes32 => uint256)) public preconferToPubKeyHashToTimestamp;
+    // Maps a validator's BLS pub key hash to the validator's details
+    mapping(bytes32 publicKeyHash => Validator) internal validators;
 
     constructor(IServiceManager _preconfServiceManager) {
         preconfServiceManager = _preconfServiceManager;
         nextPreconferIndex = 1;
     }
 
     /**
-     * @notice Registers a preconfer by giving them a non-zero registry index
-     * @param operatorSignature The signature of the preconfer in the format expected by Eigenlayer registry
+     * @notice Registers a preconfer in the registry by giving it a non-zero index
+     * @dev This function internally accesses Eigenlayer via the AVS service manager
+     * @param operatorSignature The signature of the operator in the format expected by Eigenlayer
      */
     function registerPreconfer(SignatureWithSaltAndExpiry calldata operatorSignature) external {
         // Preconfer must not have registered already
@@ -42,100 +46,183 @@ contract PreconfRegistry is IPreconfRegistry, ISignatureUtils, BLSSignatureCheck
         uint256 _nextPreconferIndex = nextPreconferIndex;
 
         preconferToIndex[msg.sender] = _nextPreconferIndex;
-        nextPreconferIndex = _nextPreconferIndex + 1;
+        indexToPreconfer[_nextPreconferIndex] = msg.sender;
 
-        emit PreconferRegistered(msg.sender, _nextPreconferIndex);
+        unchecked {
+            nextPreconferIndex = _nextPreconferIndex + 1;
+        }
+
+        emit PreconferRegistered(msg.sender);
 
         preconfServiceManager.registerOperatorToAVS(msg.sender, operatorSignature);
     }
 
     /**
-     * @notice Deregisters a preconfer from the registry
-     * @dev The preconfer that has the last index must be provided as a witness to save gas
-     * @param lastIndexWitness The address of the preconfer that has the last index
+     * @notice Deregisters a preconfer from the registry by setting its index to zero
+     * @dev It assigns the index of the last preconfer to the preconfer being removed and
+     * decrements the global index counter.
      */
-    function deregisterPreconfer(address lastIndexWitness) external {
+    function deregisterPreconfer() external {
         // Preconfer must have registered already
         if (preconferToIndex[msg.sender] == 0) {
             revert PreconferNotRegistered();
         }
 
-        // Ensure that provided witness is the preconfer that has the last index
-        uint256 _nextPreconferIndex = nextPreconferIndex - 1;
-        if (preconferToIndex[lastIndexWitness] != _nextPreconferIndex) {
-            revert LastIndexWitnessIncorrect();
-        }
+        unchecked {
+            uint256 _nextPreconferIndex = nextPreconferIndex - 1;
 
-        // Update to the decremented index to account for the removed preconfer
-        nextPreconferIndex = _nextPreconferIndex;
+            // Update to the decremented index to account for the removed preconfer
+            nextPreconferIndex = _nextPreconferIndex;
 
-        // Remove the preconfer and exchange its index with the last preconfer
-        uint256 removedPreconferIndex = preconferToIndex[msg.sender];
-        preconferToIndex[msg.sender] = 0;
-        preconferToIndex[lastIndexWitness] = removedPreconferIndex;
+            uint256 removedPreconferIndex = preconferToIndex[msg.sender];
+            address lastPreconfer = indexToPreconfer[_nextPreconferIndex];
+
+            // Remove the preconfer and exchange its index with the last preconfer
+            preconferToIndex[msg.sender] = 0;
+            preconferToIndex[lastPreconfer] = removedPreconferIndex;
+            indexToPreconfer[removedPreconferIndex] = lastPreconfer;
+        }
 
         emit PreconferDeregistered(msg.sender);
 
         preconfServiceManager.deregisterOperatorFromAVS(msg.sender);
     }
 
     /**
-     * @notice Associates a batch of validators with a preconfer
-     * @param pubkeys The public keys of the validators
-     * @param signatures The BLS signatures of the validators
+     * @notice Assigns a validator to a preconfer
+     * @dev The function allows different validators to be assigned to different preconfers, but
+     * generally, it will be called by a preconfer to assign validators to itself.
+     * @param addValidatorParams Contains the public key, signature, expiry, and preconfer
      */
-    function addValidators(BLS12381.G1Point[] calldata pubkeys, BLS12381.G2Point[] calldata signatures) external {
-        if (pubkeys.length != signatures.length) {
-            revert ArrayLengthMismatch();
-        }
+    function addValidators(AddValidatorParam[] calldata addValidatorParams) external {
+        for (uint256 i; i < addValidatorParams.length; ++i) {
+            // Revert if preconfer is not registered
+            if (preconferToIndex[addValidatorParams[i].preconfer] == 0) {
+                revert PreconferNotRegistered();
+            }
+
+            bytes memory message =
+                _createMessage(ValidatorOp.ADD, addValidatorParams[i].signatureExpiry, addValidatorParams[i].preconfer);
 
-        uint256 preconferNonce = preconferToNonce[msg.sender];
-        for (uint256 i; i < pubkeys.length; ++i) {
             // Revert if any signature is invalid
-            if (!verifySignature(_createMessage(preconferNonce), signatures[i], pubkeys[i])) {
+            if (!verifySignature(message, addValidatorParams[i].signature, addValidatorParams[i].pubkey)) {
                 revert InvalidValidatorSignature();
             }
 
+            // Revert if the signature has expired
+            if (block.timestamp > addValidatorParams[i].signatureExpiry) {
+                revert ValidatorSignatureExpired();
+            }
+
             // Point compress the public key just how it is done on the consensus layer
-            uint256[2] memory compressedPubKey = pubkeys[i].compress();
+            uint256[2] memory compressedPubKey = addValidatorParams[i].pubkey.compress();
             // Use the hash for ease of mapping
             bytes32 pubKeyHash = keccak256(abi.encodePacked(compressedPubKey));
 
-            preconferToPubKeyHashToTimestamp[msg.sender][pubKeyHash] = block.timestamp;
-
-            emit ValidatorAdded(msg.sender, compressedPubKey);
-
-            unchecked {
-                ++preconferNonce;
+            Validator memory validator = validators[pubKeyHash];
+
+            // Update the validator if it has no preconfer assigned, or if it has stopped proposing
+            // for the former preconfer
+            if (
+                validator.preconfer == address(0)
+                    || (validator.stopProposingAt != 0 && block.timestamp > validator.stopProposingAt)
+            ) {
+                unchecked {
+                    validators[pubKeyHash] = Validator({
+                        preconfer: addValidatorParams[i].preconfer,
+                        // The delay is crucial in order to not contradict the lookahead
+                        startProposingAt: uint40(block.timestamp + PreconfConstants.TWO_EPOCHS),
+                        stopProposingAt: uint40(0)
+                    });
+                }
+            } else {
+                // Validator is already proposing for a preconfer
+                revert ValidatorAlreadyActive();
             }
-        }
 
-        preconferToNonce[msg.sender] = preconferNonce;
+            emit ValidatorAdded(pubKeyHash, addValidatorParams[i].preconfer);
+        }
     }
 
     /**
-     * @notice Removes a batch of validators for a preconfer
-     * @param validatorPubKeyHashes The hashes of the public keys of the validators
+     * @notice Unassigns a validator from a preconfer
+     * @dev Instead of removing the validator immediately, we delay the removal by two epochs,
+     * & set the `stopProposingAt` timestamp.
+     * @param removeValidatorParams Contains the public key, signature and expiry
      */
-    function removeValidators(bytes32[] memory validatorPubKeyHashes) external {
-        for (uint256 i; i < validatorPubKeyHashes.length; ++i) {
-            if (preconferToPubKeyHashToTimestamp[msg.sender][validatorPubKeyHashes[i]] == 0) {
-                revert InvalidValidatorPubKeyHash();
+    function removeValidators(RemoveValidatorParam[] calldata removeValidatorParams) external {
+        for (uint256 i; i < removeValidatorParams.length; ++i) {
+            // Point compress the public key just how it is done on the consensus layer
+            uint256[2] memory compressedPubKey = removeValidatorParams[i].pubkey.compress();
+            // Use the hash for ease of mapping
+            bytes32 pubKeyHash = keccak256(abi.encodePacked(compressedPubKey));
+
+            Validator memory validator = validators[pubKeyHash];
+
+            // Revert if the validator is not active (or already removed, but waiting to stop proposing)
+            if (validator.preconfer == address(0) || validator.stopProposingAt != 0) {
+                revert ValidatorAlreadyInactive();
+            }
+
+            bytes memory message =
+                _createMessage(ValidatorOp.REMOVE, removeValidatorParams[i].signatureExpiry, validator.preconfer);
+
+            // Revert if any signature is invalid
+            if (!verifySignature(message, removeValidatorParams[i].signature, removeValidatorParams[i].pubkey)) {
+                revert InvalidValidatorSignature();
             }
-            preconferToPubKeyHashToTimestamp[msg.sender][validatorPubKeyHashes[i]] = 0;
-            emit ValidatorRemoved(msg.sender, validatorPubKeyHashes[i]);
+
+            // Revert if the signature has expired
+            if (block.timestamp > removeValidatorParams[i].signatureExpiry) {
+                revert ValidatorSignatureExpired();
+            }
+
+            unchecked {
+                // We also need to delay the removal by two epochs to avoid contradicting the lookahead
+                validators[pubKeyHash].stopProposingAt = uint40(block.timestamp + PreconfConstants.TWO_EPOCHS);
+            }
+
+            emit ValidatorRemoved(pubKeyHash, validator.preconfer);
         }
     }
 
+    //=======
+    // Views
+    //=======
+
+    function getMessageToSign(ValidatorOp validatorOp, uint256 expiry, address preconfer)
+        external
+        view
+        returns (bytes memory)
+    {
+        return _createMessage(validatorOp, expiry, preconfer);
+    }
+
+    function getNextPreconferIndex() external view returns (uint256) {
+        return nextPreconferIndex;
+    }
+
+    function getPreconferIndex(address preconfer) external view returns (uint256) {
+        return preconferToIndex[preconfer];
+    }
+
+    function getPreconferAtIndex(uint256 index) external view returns (address) {
+        return indexToPreconfer[index];
+    }
+
+    function getValidator(bytes32 pubKeyHash) external view returns (Validator memory) {
+        return validators[pubKeyHash];
+    }
+
     //=========
     // Helpers
     //=========
 
-    /**
-     * @notice Returns the message to be signed by the preconfer
-     * @param nonce The nonce of the preconfer
-     */
-    function _createMessage(uint256 nonce) internal view returns (bytes memory) {
-        return abi.encodePacked(block.chainid, msg.sender, nonce);
+    function _createMessage(ValidatorOp validatorOp, uint256 expiry, address preconfer)
+        internal
+        view
+        returns (bytes memory)
+    {
+        return abi.encodePacked(block.chainid, validatorOp, expiry, preconfer);
     }
 }

SmartContracts/src/avs/libraries/PreconfConstants.sol

@@ -0,0 +1,6 @@
+// SPDX-License-Identifier: UNLICENSED
+pragma solidity 0.8.25;
+
+library PreconfConstants {
+    uint256 internal constant TWO_EPOCHS = 768;
+}