From 40d5750cc999ffc60ca974cb71142060d3584b5c Mon Sep 17 00:00:00 2001
From: 0xSheller <93097065+0xSheller@users.noreply.github.com>
Date: Thu, 23 Apr 2026 19:11:37 -0700
Subject: [PATCH] Fix pubkey-uniqueness check

For Issue: #338
---
 .../megapool/RocketMegapoolManager.sol        | 10 ++++-
 .../RocketMegapoolManagerInterface.sol        |  1 +
 test/megapool/megapool-tests.js               | 40 +++++++++++++++++++
 3 files changed, 49 insertions(+), 2 deletions(-)

diff --git a/contracts/contract/megapool/RocketMegapoolManager.sol b/contracts/contract/megapool/RocketMegapoolManager.sol
index 9e2ca689..a2cda999 100644
--- a/contracts/contract/megapool/RocketMegapoolManager.sol
+++ b/contracts/contract/megapool/RocketMegapoolManager.sol
@@ -40,12 +40,18 @@ contract RocketMegapoolManager is RocketBase, RocketMegapoolManagerInterface {
         setUint(setCountKey, index + 1);
         uint256 encoded = (uint256(uint160(_megapoolAddress)) << 96) | _validatorId;
         setUint(keccak256(abi.encodePacked("megapool.validator.set", index)), encoded);
-        // Add pubkey => megapool mapping and ensure uniqueness
-        bytes32 key = keccak256(abi.encodePacked("validator.megapool", _megapoolAddress, _pubkey));
+        // Add pubkey => megapool mapping and ensure uniqueness across all megapools
+        bytes32 key = keccak256(abi.encodePacked("validator.megapool", _pubkey));
         require(getAddress(key) == address(0x0), "Pubkey in use");
         setAddress(key, _megapoolAddress);
     }
 
+    /// @notice Returns the megapool that owns a given validator pubkey, or address(0) if unclaimed
+    /// @param _pubkey The validator pubkey to look up
+    function getMegapoolByPubkey(bytes calldata _pubkey) override external view returns (address) {
+        return getAddress(keccak256(abi.encodePacked("validator.megapool", _pubkey)));
+    }
+
     /// @notice Returns the last trusted member to execute a challenge
     function getLastChallenger() override external view returns (address) {
         return getAddress(challengerKey);
diff --git a/contracts/interface/megapool/RocketMegapoolManagerInterface.sol b/contracts/interface/megapool/RocketMegapoolManagerInterface.sol
index 907d3211..d816eee7 100644
--- a/contracts/interface/megapool/RocketMegapoolManagerInterface.sol
+++ b/contracts/interface/megapool/RocketMegapoolManagerInterface.sol
@@ -13,6 +13,7 @@ interface RocketMegapoolManagerInterface {
 
     function getValidatorCount() external view returns (uint256);
     function addValidator(address _megapoolAddress, uint32 _validatorId, bytes calldata _pubkey) external;
+    function getMegapoolByPubkey(bytes calldata _pubkey) external view returns (address);
     function getLastChallenger() external view returns (address);
     function getValidatorInfo(uint256 _index) external view returns (bytes memory pubkey, RocketMegapoolStorageLayout.ValidatorInfo memory validatorInfo, address megapool, uint32 validatorId);
     function stake(RocketMegapoolInterface megapool, uint32 _validatorId, uint64 _slotTimestamp, ValidatorProof calldata _proof, SlotProof calldata _slotProof) external;
diff --git a/test/megapool/megapool-tests.js b/test/megapool/megapool-tests.js
index ef2f0e0e..2cd7e393 100644
--- a/test/megapool/megapool-tests.js
+++ b/test/megapool/megapool-tests.js
@@ -409,6 +409,46 @@ export default function() {
             );
         });
 
+        it(printTitle('node', 'can not reuse pubkey across megapools'), async () => {
+            // node performs first deposit with a fresh pubkey
+            const pubkey = getValidatorPubkey();
+            const signature = getValidatorSignature();
+            const wc1 = await getMegapoolWithdrawalCredentials(node.address);
+            const depositData1 = {
+                pubkey: pubkey,
+                withdrawalCredentials: Buffer.from(wc1.substr(2), 'hex'),
+                amount: BigInt(1000000000), // gwei
+                signature: signature,
+            };
+            const depositDataRoot1 = getDepositDataRoot(depositData1);
+            const rocketNodeDeposit = await RocketNodeDeposit.deployed();
+            await rocketNodeDeposit.connect(node).deposit('4'.ether, false, pubkey, signature, depositDataRoot1, { value: '4'.ether });
+
+            // node2 (different megapool) attempts to register the same pubkey
+            const wc2 = await getMegapoolWithdrawalCredentials(node2.address);
+            const depositData2 = {
+                pubkey: pubkey,
+                withdrawalCredentials: Buffer.from(wc2.substr(2), 'hex'),
+                amount: BigInt(1000000000), // gwei
+                signature: signature,
+            };
+            const depositDataRoot2 = getDepositDataRoot(depositData2);
+            await shouldRevert(
+                rocketNodeDeposit.connect(node2).deposit('4'.ether, false, pubkey, signature, depositDataRoot2, { value: '4'.ether }),
+                'Was able to reuse pubkey from a different megapool',
+                'Pubkey in use',
+            );
+
+            // Sanity check the new getter resolves the pubkey to node's megapool
+            const rocketMegapoolManager = await RocketMegapoolManager.deployed();
+            const expectedMegapool = (await getMegapoolForNode(node)).target;
+            assert.strictEqual(
+                (await rocketMegapoolManager.getMegapoolByPubkey(pubkey)).toLowerCase(),
+                expectedMegapool.toLowerCase(),
+                'getMegapoolByPubkey did not return the owning megapool',
+            );
+        });
+
         it(printTitle('node', 'can deposit using ETH credit'), async () => {
             // Enter and exit queue to receive a 4 ETH credit
             await nodeDeposit(node, '4'.ether);