feat: bls key ownership check

contracts/interfaces/IBLSKeyChecker.sol

@@ -0,0 +1,21 @@
+// SPDX-License-Identifier: MIT
+pragma solidity ^0.8.12;
+
+interface IBLSKeyChecker {
+    struct BLSKeyWithProof {
+        uint256[2][] blsG1PublicKeys;
+        uint256[2][2] aggG2PublicKey;
+        uint256[2] signature;
+        bytes32 salt;
+        uint256 expiry;
+    }
+
+    function isSaltSpent(address operator, bytes32 salt) external view returns (bool);
+
+    function calculateKeyWithProofHash(address operator, bytes32 salt, uint256 expiry)
+        external
+        view
+        returns (bytes32);
+
+    function domainSeparator() external view returns (bytes32);
+}

contracts/interfaces/ILagrangeCommittee.sol

@@ -1,7 +1,9 @@
 // SPDX-License-Identifier: MIT
 pragma solidity ^0.8.12;
 
-interface ILagrangeCommittee {
+import {IBLSKeyChecker} from "./IBLSKeyChecker.sol";
+
+interface ILagrangeCommittee is IBLSKeyChecker {
     struct UnsubscribedParam {
         uint32 chainID;
         uint256 blockNumber;
@@ -59,15 +61,15 @@ interface ILagrangeCommittee {
         uint96 maxWeight
     ) external;
 
-    function addOperator(address operator, address signAddress, uint256[2][] memory blsPubKeys) external;
+    function addOperator(address operator, address signAddress, BLSKeyWithProof memory blsKeyWithProof) external;
 
     function removeOperator(address operator) external;
 
     function unsubscribeByAdmin(address[] memory operators, uint32 chainID) external;
 
-    function addBlsPubKeys(address operator, uint256[2][] memory additionalBlsPubKeys) external;
+    function addBlsPubKeys(address operator, BLSKeyWithProof memory blsKeyWithProof) external;
 
-    function updateBlsPubKey(address operator, uint32 index, uint256[2] memory blsPubKey) external;
+    function updateBlsPubKey(address operator, uint32 index, BLSKeyWithProof memory blsKeyWithProof) external;
 
     function removeBlsPubKeys(address operator, uint32[] memory indices) external;
 

contracts/interfaces/ILagrangeService.sol

@@ -2,6 +2,7 @@
 pragma solidity ^0.8.12;
 
 import {ISignatureUtils} from "eigenlayer-contracts/src/contracts/interfaces/ISignatureUtils.sol";
+import {IBLSKeyChecker} from "./IBLSKeyChecker.sol";
 
 interface ILagrangeService {
     function addOperatorsToWhitelist(address[] calldata operators) external;
@@ -10,13 +11,13 @@ interface ILagrangeService {
 
     function register(
         address signAddress,
-        uint256[2][] memory blsPubKeys,
+        IBLSKeyChecker.BLSKeyWithProof memory blsKeyWithProof,
         ISignatureUtils.SignatureWithSaltAndExpiry memory operatorSignature
     ) external;
 
-    function addBlsPubKeys(uint256[2][] memory additionalBlsPubKeys) external;
+    function addBlsPubKeys(IBLSKeyChecker.BLSKeyWithProof memory blsKeyWithProof) external;
 
-    function updateBlsPubKey(uint32 index, uint256[2] memory blsPubKey) external;
+    function updateBlsPubKey(uint32 index, IBLSKeyChecker.BLSKeyWithProof memory blsKeyWithProof) external;
 
     function removeBlsPubKeys(uint32[] memory indices) external;
 

contracts/library/BLSKeyChecker.sol

@@ -0,0 +1,80 @@
+// SPDX-License-Identifier: MIT
+pragma solidity ^0.8.12;
+
+import {BN254} from "eigenlayer-middleware/libraries/BN254.sol";
+
+import "../interfaces/IBLSKeyChecker.sol";
+
+abstract contract BLSKeyChecker is IBLSKeyChecker {
+    using BN254 for BN254.G1Point;
+
+    uint256 internal constant PAIRING_EQUALITY_CHECK_GAS = 120000;
+
+    bytes32 public constant DOMAIN_TYPEHASH =
+        keccak256("EIP712Domain(string name,uint256 chainId,address verifyingContract)");
+
+    bytes32 public constant BLS_KEY_WITH_PROOF_TYPEHASH =
+        keccak256("BLSKeyWithProof(address operator,bytes32 salt,uint256 expiry)");
+
+    struct SaltStorage {
+        mapping(address => mapping(bytes32 => bool)) operatorSalts;
+    }
+
+    // keccak256(abi.encode(uint256(keccak256("lagrange.blskeychecker.storage")) - 1)) & ~bytes32(uint256(0xff))
+    bytes32 private constant SaltStorageLocation = 0x51615ea63289f14fdd891b383e2929b2f73c675cf292e602b5fceb059f7a4700;
+
+    function _getSaltStorage() private pure returns (SaltStorage storage $) {
+        assembly {
+            $.slot := SaltStorageLocation
+        }
+    }
+
+    function _validateBLSKeyWithProof(address operator, BLSKeyWithProof memory keyWithProof) internal {
+        require(
+            keyWithProof.expiry >= block.timestamp, "BLSKeyChecker.checkBLSKeyWithProof: operator signature expired"
+        );
+        SaltStorage storage $ = _getSaltStorage();
+        require(!$.operatorSalts[operator][keyWithProof.salt], "BLSKeyChecker.checkBLSKeyWithProof: salt already spent");
+
+        $.operatorSalts[operator][keyWithProof.salt] = true;
+
+        BN254.G1Point memory aggG1 = BN254.G1Point(0, 0);
+        for (uint256 i = 0; i < keyWithProof.blsG1PublicKeys.length; i++) {
+            aggG1 = aggG1.plus(BN254.G1Point(keyWithProof.blsG1PublicKeys[i][0], keyWithProof.blsG1PublicKeys[i][1]));
+        }
+
+        BN254.G2Point memory aggG2 = BN254.G2Point(keyWithProof.aggG2PublicKey[0], keyWithProof.aggG2PublicKey[1]);
+
+        // check the pairing equation e(g1, aggG2) == e(aggG1, -g2)
+        // to ensure that the aggregated G2 public key is correct
+        (bool pairing, bool valid) =
+            BN254.safePairing(BN254.generatorG1(), aggG2, aggG1, BN254.negGeneratorG2(), PAIRING_EQUALITY_CHECK_GAS);
+        require(pairing && valid, "BLSKeyChecker.checkBLSKeyWithProof: invalid BLS key");
+
+        BN254.G1Point memory sig = BN254.G1Point(keyWithProof.signature[0], keyWithProof.signature[1]);
+
+        BN254.G1Point memory msgHash =
+            BN254.hashToG1(calculateKeyWithProofHash(operator, keyWithProof.salt, keyWithProof.expiry));
+
+        // check the BLS signature
+        (pairing, valid) = BN254.safePairing(sig, BN254.negGeneratorG2(), msgHash, aggG2, PAIRING_EQUALITY_CHECK_GAS);
+        require(pairing && valid, "BLSKeyChecker.checkBLSKeyWithProof: invalid BLS signature");
+    }
+
+    function isSaltSpent(address operator, bytes32 salt) public view returns (bool) {
+        SaltStorage storage $ = _getSaltStorage();
+        return $.operatorSalts[operator][salt];
+    }
+
+    function calculateKeyWithProofHash(address operator, bytes32 salt, uint256 expiry) public view returns (bytes32) {
+        bytes32 structHash = keccak256(abi.encode(BLS_KEY_WITH_PROOF_TYPEHASH, operator, salt, expiry));
+        bytes32 digestHash = keccak256(abi.encodePacked("\x19\x01", domainSeparator(), structHash));
+
+        return digestHash;
+    }
+
+    function domainSeparator() public view returns (bytes32) {
+        return
+            keccak256(abi.encode(DOMAIN_TYPEHASH, keccak256("Lagrange State Committee"), block.chainid, address(this)));
+    }
+}

contracts/mock/BLSKeyCheckerMock.sol

@@ -0,0 +1,16 @@
+// SPDX-License-Identifier: MIT
+pragma solidity ^0.8.12;
+
+import {BN254} from "eigenlayer-middleware/libraries/BN254.sol";
+
+import "../interfaces/IBLSKeyChecker.sol";
+import "../library/BLSKeyChecker.sol";
+
+contract BLSKeyCheckerMock is BLSKeyChecker {
+    constructor() {}
+
+    function checkBLSKeyWithProof(address operator, BLSKeyWithProof calldata keyWithProof) external returns (bool) {
+        _validateBLSKeyWithProof(operator, keyWithProof);
+        return true;
+    }
+}

contracts/protocol/LagrangeCommittee.sol

@@ -8,8 +8,9 @@ import "@openzeppelin/contracts/utils/math/SafeCast.sol";
 import "../interfaces/ILagrangeCommittee.sol";
 import "../interfaces/ILagrangeService.sol";
 import "../interfaces/IVoteWeigher.sol";
+import "../library/BLSKeyChecker.sol";
 
-contract LagrangeCommittee is Initializable, OwnableUpgradeable, ILagrangeCommittee {
+contract LagrangeCommittee is BLSKeyChecker, Initializable, OwnableUpgradeable, ILagrangeCommittee {
     ILagrangeService public immutable service;
     IVoteWeigher public immutable voteWeigher;
 
@@ -60,12 +61,12 @@ contract LagrangeCommittee is Initializable, OwnableUpgradeable, ILagrangeCommit
     }
 
     // Adds a new operator to the committee
-    function addOperator(address operator, address signAddress, uint256[2][] calldata blsPubKeys)
+    function addOperator(address operator, address signAddress, BLSKeyWithProof calldata blsKeyWithProof)
         external
         onlyService
     {
-        _validateBlsPubKeys(blsPubKeys);
-        _registerOperator(operator, signAddress, blsPubKeys);
+        _validateBlsPubKeys(blsKeyWithProof.blsG1PublicKeys);
+        _registerOperator(operator, signAddress, blsKeyWithProof);
     }
 
     // Removes an operator from the committee
@@ -86,18 +87,22 @@ contract LagrangeCommittee is Initializable, OwnableUpgradeable, ILagrangeCommit
     }
 
     // Adds additional BLS public keys to an operator
-    function addBlsPubKeys(address operator, uint256[2][] calldata additionalBlsPubKeys) external onlyService {
-        _validateBlsPubKeys(additionalBlsPubKeys);
-        _addBlsPubKeys(operator, additionalBlsPubKeys);
+    function addBlsPubKeys(address operator, BLSKeyWithProof calldata blsKeyWithProof) external onlyService {
+        _validateBlsPubKeys(blsKeyWithProof.blsG1PublicKeys);
+        _addBlsPubKeys(operator, blsKeyWithProof);
     }
 
     // Updates an operator's BLS public key for the given index
-    function updateBlsPubKey(address operator, uint32 index, uint256[2] calldata blsPubKey) external onlyService {
-        require(blsPubKey[0] != 0 && blsPubKey[1] != 0, "Invalid BLS Public Key.");
+    function updateBlsPubKey(address operator, uint32 index, BLSKeyWithProof calldata blsKeyWithProof)
+        external
+        onlyService
+    {
+        _validateBLSKeyWithProof(operator, blsKeyWithProof);
+        require(blsKeyWithProof.blsG1PublicKeys.length == 1, "Length should be 1 for update");
         uint256[2][] storage _blsPubKeys = operatorsStatus[operator].blsPubKeys;
         require(_blsPubKeys.length > index, "Invalid index");
-        _checkBlsPubKeyDuplicate(_blsPubKeys, blsPubKey);
-        _blsPubKeys[index] = blsPubKey;
+        _checkBlsPubKeyDuplicate(_blsPubKeys, blsKeyWithProof.blsG1PublicKeys[0]);
+        _blsPubKeys[index] = blsKeyWithProof.blsG1PublicKeys[0];
     }
 
     // Removes BLS public keys from an operator for the given indices
@@ -395,6 +400,13 @@ contract LagrangeCommittee is Initializable, OwnableUpgradeable, ILagrangeCommit
         } else if (_lastEpoch == 0) {
             return 0;
         } else {
+            if (_blockNumber >= _lastEpochBlock - committeeParam.duration) {
+                _epochNumber = _lastEpoch;
+                while (_blockNumber < _getUpdatedBlock(_chainID, _epochNumber)) {
+                    _epochNumber--;
+                }
+                return _epochNumber;
+            }
             // binary search
             uint256 _low = 0;
             uint256 _high = _lastEpoch;
@@ -410,24 +422,30 @@ contract LagrangeCommittee is Initializable, OwnableUpgradeable, ILagrangeCommit
         }
     }
 
-    function _registerOperator(address _operator, address _signAddress, uint256[2][] memory _blsPubKeys) internal {
+    function _registerOperator(address _operator, address _signAddress, BLSKeyWithProof memory _blsKeyWithProof)
+        internal
+    {
+        _validateBLSKeyWithProof(_operator, _blsKeyWithProof);
+
         delete operatorsStatus[_operator];
         OperatorStatus storage _opStatus = operatorsStatus[_operator];
         _opStatus.signAddress = _signAddress;
-        uint256 _length = _blsPubKeys.length;
+        uint256 _length = _blsKeyWithProof.blsG1PublicKeys.length;
         for (uint256 i; i < _length; i++) {
-            _checkBlsPubKeyDuplicate(_opStatus.blsPubKeys, _blsPubKeys[i]);
-            _opStatus.blsPubKeys.push(_blsPubKeys[i]);
+            _checkBlsPubKeyDuplicate(_opStatus.blsPubKeys, _blsKeyWithProof.blsG1PublicKeys[i]);
+            _opStatus.blsPubKeys.push(_blsKeyWithProof.blsG1PublicKeys[i]);
         }
     }
 
-    function _addBlsPubKeys(address _operator, uint256[2][] memory _additionalBlsPubKeys) internal {
+    function _addBlsPubKeys(address _operator, BLSKeyWithProof memory _blsKeyWithProof) internal {
+        _validateBLSKeyWithProof(_operator, _blsKeyWithProof);
+
         OperatorStatus storage _opStatus = operatorsStatus[_operator];
         require(_opStatus.blsPubKeys.length != 0, "Operator is not registered.");
-        uint256 _length = _additionalBlsPubKeys.length;
+        uint256 _length = _blsKeyWithProof.blsG1PublicKeys.length;
         for (uint256 i; i < _length; i++) {
-            _checkBlsPubKeyDuplicate(_opStatus.blsPubKeys, _additionalBlsPubKeys[i]);
-            _opStatus.blsPubKeys.push(_additionalBlsPubKeys[i]);
+            _checkBlsPubKeyDuplicate(_opStatus.blsPubKeys, _blsKeyWithProof.blsG1PublicKeys[i]);
+            _opStatus.blsPubKeys.push(_blsKeyWithProof.blsG1PublicKeys[i]);
         }
     }
 
@@ -562,7 +580,6 @@ contract LagrangeCommittee is Initializable, OwnableUpgradeable, ILagrangeCommit
 
     function _validateBlsPubKeys(uint256[2][] memory _blsPubKeys) internal pure {
         require(_blsPubKeys.length != 0, "Empty BLS Public Keys.");
-        // TODO: need to add validation for blsPubKeys with signatures
         uint256 _length = _blsPubKeys.length;
         for (uint256 i; i < _length; i++) {
             require(_blsPubKeys[i][0] != 0 && _blsPubKeys[i][1] != 0, "Invalid BLS Public Key.");

contracts/protocol/LagrangeService.sol

@@ -10,6 +10,7 @@ import {IStakeManager} from "../interfaces/IStakeManager.sol";
 import {ILagrangeCommittee} from "../interfaces/ILagrangeCommittee.sol";
 import {ILagrangeService} from "../interfaces/ILagrangeService.sol";
 import {IVoteWeigher} from "../interfaces/IVoteWeigher.sol";
+import {IBLSKeyChecker} from "../interfaces/IBLSKeyChecker.sol";
 
 contract LagrangeService is Initializable, OwnableUpgradeable, ILagrangeService {
     mapping(address => bool) public operatorWhitelist;
@@ -69,30 +70,33 @@ contract LagrangeService is Initializable, OwnableUpgradeable, ILagrangeService
     /// Add the operator to the service.
     function register(
         address signAddress,
-        uint256[2][] calldata blsPubKeys,
+        IBLSKeyChecker.BLSKeyWithProof calldata blsKeyWithProof,
         ISignatureUtils.SignatureWithSaltAndExpiry memory operatorSignature
     ) external onlyWhitelisted {
         address _operator = msg.sender;
-        committee.addOperator(_operator, signAddress, blsPubKeys);
+        committee.addOperator(_operator, signAddress, blsKeyWithProof);
         uint32 serveUntilBlock = type(uint32).max;
         stakeManager.lockStakeUntil(_operator, serveUntilBlock);
         avsDirectory.registerOperatorToAVS(_operator, operatorSignature);
         emit OperatorRegistered(_operator, serveUntilBlock);
     }
 
     /// Add extra BlsPubKeys
-    function addBlsPubKeys(uint256[2][] calldata additionalBlsPubKeys) external onlyWhitelisted {
+    function addBlsPubKeys(IBLSKeyChecker.BLSKeyWithProof calldata blsKeyWithProof) external onlyWhitelisted {
         address _operator = msg.sender;
-        committee.addBlsPubKeys(_operator, additionalBlsPubKeys);
+        committee.addBlsPubKeys(_operator, blsKeyWithProof);
         uint32 serveUntilBlock = type(uint32).max;
         stakeManager.lockStakeUntil(_operator, serveUntilBlock);
         emit OperatorRegistered(_operator, serveUntilBlock);
     }
 
     /// Update the BlsPubKey for the given index
-    function updateBlsPubKey(uint32 index, uint256[2] calldata blsPubKey) external onlyWhitelisted {
+    function updateBlsPubKey(uint32 index, IBLSKeyChecker.BLSKeyWithProof memory blsKeyWithProof)
+        external
+        onlyWhitelisted
+    {
         address _operator = msg.sender;
-        committee.updateBlsPubKey(_operator, index, blsPubKey);
+        committee.updateBlsPubKey(_operator, index, blsKeyWithProof);
     }
 
     /// Remove the BlsPubKeys for the given indices