feat: add UAM to SocketRegistry (#532)

ypatil12 · web-flow · commit 510bdb18afb1 · 2025-08-19T12:45:48.000-04:00
**Motivation:**

In the `SocketRegistry`, only the operator can update the socket, and
not an admin or appointee.

**Modifications:**

- Inherit `PermissionControllerMixin` in `SocketRegistry`
- Update constructor params for `TaskAVSRegistrarBase` and
`AVSRegistrarWithSocket`

**Result:**

Clearer code
diff --git a/docs/middlewareV2/AVSRegistrar.md b/docs/middlewareV2/AVSRegistrar.md
@@ -302,9 +302,12 @@ function getOperatorSocket(
 
 ```solidity
 /**
- * @notice Update the socket URL for the calling operator
- * @param operator The operator address (must be msg.sender)
- * @param socket The new socket URL
+ * @notice Updates the socket for an operator.
+ * @param operator The operator to set the socket for.
+ * @param socket The socket to set for the operator.
+ * @dev This function can only be called by the operator themselves.
+ * @dev Reverts for:
+ *      - InvalidPermissions: The caller does not have permission to call this function (via core `PermissionController`)
  */
 function updateSocket(address operator, string memory socket) external;
 ```
@@ -317,7 +320,7 @@ Allows an operator to update their socket URL after registration. The operator d
 - Emits `OperatorSocketSet` event
 
 *Requirements:*
-- `msg.sender` MUST be the `operator`
+- Caller MUST be authorized, either as the operator themselves or an admin/appointee (see the Core[`PermissionController.md`](https://github.com/Layr-Labs/eigenlayer-contracts/blob/main/docs/permissions/PermissionController.md))
 
 ---
 
diff --git a/src/avs/task/TaskAVSRegistrarBase.sol b/src/avs/task/TaskAVSRegistrarBase.sol
@@ -5,6 +5,8 @@ import {OwnableUpgradeable} from "@openzeppelin-upgrades/contracts/access/Ownabl
 import {Initializable} from "@openzeppelin-upgrades/contracts/proxy/utils/Initializable.sol";
 import {IAllocationManager} from
     "eigenlayer-contracts/src/contracts/interfaces/IAllocationManager.sol";
+import {IPermissionController} from
+    "eigenlayer-contracts/src/contracts/interfaces/IPermissionController.sol";
 import {IKeyRegistrar} from "eigenlayer-contracts/src/contracts/interfaces/IKeyRegistrar.sol";
 import {AVSRegistrarWithSocket} from
     "../../middlewareV2/registrar/presets/AVSRegistrarWithSocket.sol";
@@ -26,11 +28,13 @@ abstract contract TaskAVSRegistrarBase is
      * @dev Constructor that passes parameters to parent
      * @param _allocationManager The AllocationManager contract address
      * @param _keyRegistrar The KeyRegistrar contract address
+     * @param _permissionController The PermissionController contract address
      */
     constructor(
         IAllocationManager _allocationManager,
-        IKeyRegistrar _keyRegistrar
-    ) AVSRegistrarWithSocket(_allocationManager, _keyRegistrar) {
+        IKeyRegistrar _keyRegistrar,
+        IPermissionController _permissionController
+    ) AVSRegistrarWithSocket(_allocationManager, _keyRegistrar, _permissionController) {
         _disableInitializers();
     }
 
diff --git a/src/interfaces/ISocketRegistryV2.sol b/src/interfaces/ISocketRegistryV2.sol
@@ -1,20 +1,12 @@
 // SPDX-License-Identifier: BUSL-1.1
 pragma solidity ^0.8.27;
 
-interface ISocketRegistryErrors {
-    /// @notice Thrown when the caller is not the operator
-    error CallerNotOperator();
-
-    /// @notice Thrown when the data length mismatch
-    error DataLengthMismatch();
-}
-
 interface ISocketRegistryEvents {
     /// @notice Emitted when an operator socket is set
     event OperatorSocketSet(address indexed operator, string socket);
 }
 
-interface ISocketRegistryV2 is ISocketRegistryErrors, ISocketRegistryEvents {
+interface ISocketRegistryV2 is ISocketRegistryEvents {
     /**
      * @notice Gets the socket for an operator.
      * @param operator The operator to get the socket for.
@@ -30,7 +22,7 @@ interface ISocketRegistryV2 is ISocketRegistryErrors, ISocketRegistryEvents {
      * @param socket The socket to set for the operator.
      * @dev This function can only be called by the operator themselves.
      * @dev Reverts for:
-     *      - CallerNotOperator: The caller is not the operator
+     *      - InvalidPermissions: The caller does not have permission to call this function (via core `PermissionController`)
      */
     function updateSocket(address operator, string memory socket) external;
 }
diff --git a/src/middlewareV2/registrar/modules/SocketRegistry.sol b/src/middlewareV2/registrar/modules/SocketRegistry.sol
@@ -7,12 +7,20 @@ import {
     OperatorSetLib,
     OperatorSet
 } from "eigenlayer-contracts/src/contracts/libraries/OperatorSetLib.sol";
+import {IPermissionController} from
+    "eigenlayer-contracts/src/contracts/interfaces/IPermissionController.sol";
+import {PermissionControllerMixin} from
+    "eigenlayer-contracts/src/contracts/mixins/PermissionControllerMixin.sol";
 
 /// @notice A module that allows for the setting and removal of operator sockets
 /// @dev This contract assumes a single socket per operator
-abstract contract SocketRegistry is SocketRegistryStorage {
+abstract contract SocketRegistry is SocketRegistryStorage, PermissionControllerMixin {
     using OperatorSetLib for OperatorSet;
 
+    constructor(
+        IPermissionController _permissionController
+    ) PermissionControllerMixin(_permissionController) {}
+
     /// @inheritdoc ISocketRegistryV2
     function getOperatorSocket(
         address operator
@@ -21,8 +29,7 @@ abstract contract SocketRegistry is SocketRegistryStorage {
     }
 
     /// @inheritdoc ISocketRegistryV2
-    function updateSocket(address operator, string memory socket) external {
-        require(msg.sender == operator, CallerNotOperator());
+    function updateSocket(address operator, string memory socket) external checkCanCall(operator) {
         _setOperatorSocket(operator, socket);
     }
 
diff --git a/src/middlewareV2/registrar/modules/SocketRegistryStorage.sol b/src/middlewareV2/registrar/modules/SocketRegistryStorage.sol
@@ -2,6 +2,8 @@
 pragma solidity ^0.8.12;
 
 import {ISocketRegistryV2} from "../../../interfaces/ISocketRegistryV2.sol";
+import {IPermissionController} from
+    "eigenlayer-contracts/src/contracts/interfaces/IPermissionController.sol";
 
 /**
  * @title Storage variables for the `SocketRegistry` contract.
diff --git a/src/middlewareV2/registrar/presets/AVSRegistrarWithSocket.sol b/src/middlewareV2/registrar/presets/AVSRegistrarWithSocket.sol
@@ -3,6 +3,8 @@ pragma solidity ^0.8.27;
 
 import {IAllocationManager} from
     "eigenlayer-contracts/src/contracts/interfaces/IAllocationManager.sol";
+import {IPermissionController} from
+    "eigenlayer-contracts/src/contracts/interfaces/IPermissionController.sol";
 import {IKeyRegistrar} from "eigenlayer-contracts/src/contracts/interfaces/IKeyRegistrar.sol";
 
 import {IAVSRegistrarWithSocket} from "../../../interfaces/IAVSRegistrarWithSocket.sol";
@@ -12,8 +14,9 @@ import {SocketRegistry} from "../modules/SocketRegistry.sol";
 contract AVSRegistrarWithSocket is AVSRegistrar, SocketRegistry, IAVSRegistrarWithSocket {
     constructor(
         IAllocationManager _allocationManager,
-        IKeyRegistrar _keyRegistrar
-    ) AVSRegistrar(_allocationManager, _keyRegistrar) {}
+        IKeyRegistrar _keyRegistrar,
+        IPermissionController _permissionController
+    ) AVSRegistrar(_allocationManager, _keyRegistrar) SocketRegistry(_permissionController) {}
 
     function initialize(
         address avs
diff --git a/test/mocks/MockTaskAVSRegistrar.sol b/test/mocks/MockTaskAVSRegistrar.sol
@@ -4,19 +4,22 @@ pragma solidity ^0.8.27;
 import {IAllocationManager} from
     "eigenlayer-contracts/src/contracts/interfaces/IAllocationManager.sol";
 import {IKeyRegistrar} from "eigenlayer-contracts/src/contracts/interfaces/IKeyRegistrar.sol";
-
+import {IPermissionController} from
+    "eigenlayer-contracts/src/contracts/interfaces/IPermissionController.sol";
 import {TaskAVSRegistrarBase} from "../../src/avs/task/TaskAVSRegistrarBase.sol";
 
 contract MockTaskAVSRegistrar is TaskAVSRegistrarBase {
     /**
      * @dev Constructor that passes parameters to parent TaskAVSRegistrarBase
      * @param _allocationManager The AllocationManager contract address
      * @param _keyRegistrar The KeyRegistrar contract address
+     * @param _permissionController The PermissionController contract address
      */
     constructor(
         IAllocationManager _allocationManager,
-        IKeyRegistrar _keyRegistrar
-    ) TaskAVSRegistrarBase(_allocationManager, _keyRegistrar) {}
+        IKeyRegistrar _keyRegistrar,
+        IPermissionController _permissionController
+    ) TaskAVSRegistrarBase(_allocationManager, _keyRegistrar, _permissionController) {}
 
     /**
      * @dev Initializer that calls parent initializer
diff --git a/test/unit/TaskAVSRegistrarBaseUnit.t.sol b/test/unit/TaskAVSRegistrarBaseUnit.t.sol
@@ -1,7 +1,6 @@
 // SPDX-License-Identifier: BUSL-1.1
 pragma solidity ^0.8.27;
 
-import {Test} from "forge-std/Test.sol";
 import {ProxyAdmin} from "@openzeppelin/contracts/proxy/transparent/ProxyAdmin.sol";
 import {TransparentUpgradeableProxy} from
     "@openzeppelin/contracts/proxy/transparent/TransparentUpgradeableProxy.sol";
@@ -17,12 +16,11 @@ import {ITaskAVSRegistrarBaseTypes} from "../../src/interfaces/ITaskAVSRegistrar
 import {ITaskAVSRegistrarBaseErrors} from "../../src/interfaces/ITaskAVSRegistrarBase.sol";
 import {ITaskAVSRegistrarBaseEvents} from "../../src/interfaces/ITaskAVSRegistrarBase.sol";
 import {MockTaskAVSRegistrar} from "../mocks/MockTaskAVSRegistrar.sol";
-import {AllocationManagerMock} from "../mocks/AllocationManagerMock.sol";
-import {KeyRegistrarMock} from "../mocks/KeyRegistrarMock.sol";
+import {MockEigenLayerDeployer} from "./middlewareV2/MockDeployer.sol";
 
 // Base test contract with common setup
 contract TaskAVSRegistrarBaseUnitTests is
-    Test,
+    MockEigenLayerDeployer,
     ITaskAVSRegistrarBaseTypes,
     ITaskAVSRegistrarBaseErrors,
     ITaskAVSRegistrarBaseEvents
@@ -32,10 +30,6 @@ contract TaskAVSRegistrarBaseUnitTests is
     address public owner = address(0x4);
     address public nonOwner = address(0x5);
 
-    // Mock contracts
-    AllocationManagerMock public allocationManager;
-    KeyRegistrarMock public keyRegistrar;
-
     // Test operator set IDs
     uint32 public constant AGGREGATOR_OPERATOR_SET_ID = 1;
     uint32 public constant EXECUTOR_OPERATOR_SET_ID_1 = 2;
@@ -44,20 +38,19 @@ contract TaskAVSRegistrarBaseUnitTests is
 
     // Contract under test
     MockTaskAVSRegistrar public registrar;
-    ProxyAdmin public proxyAdmin;
 
     function setUp() public virtual {
-        // Deploy mock contracts
-        allocationManager = new AllocationManagerMock();
-        keyRegistrar = new KeyRegistrarMock();
+        // Deploy mock EigenLayer contracts
+        _deployMockEigenLayer();
 
         // Create initial valid config
         AvsConfig memory initialConfig = _createValidAvsConfig();
 
         // Deploy the registrar with proxy pattern
-        proxyAdmin = new ProxyAdmin();
         MockTaskAVSRegistrar registrarImpl = new MockTaskAVSRegistrar(
-            IAllocationManager(address(allocationManager)), IKeyRegistrar(address(keyRegistrar))
+            IAllocationManager(address(allocationManagerMock)),
+            IKeyRegistrar(address(keyRegistrarMock)),
+            permissionController
         );
         TransparentUpgradeableProxy proxy = new TransparentUpgradeableProxy(
             address(registrarImpl),
@@ -139,7 +132,9 @@ contract TaskAVSRegistrarBaseUnitTests_Constructor is TaskAVSRegistrarBaseUnitTe
         // Deploy new registrar with proxy pattern
         ProxyAdmin newProxyAdmin = new ProxyAdmin();
         MockTaskAVSRegistrar newRegistrarImpl = new MockTaskAVSRegistrar(
-            IAllocationManager(address(allocationManager)), IKeyRegistrar(address(keyRegistrar))
+            IAllocationManager(address(allocationManagerMock)),
+            IKeyRegistrar(address(keyRegistrarMock)),
+            permissionController
         );
         TransparentUpgradeableProxy newProxy = new TransparentUpgradeableProxy(
             address(newRegistrarImpl),
@@ -169,7 +164,9 @@ contract TaskAVSRegistrarBaseUnitTests_Constructor is TaskAVSRegistrarBaseUnitTe
         // Deploy implementation
         ProxyAdmin newProxyAdmin = new ProxyAdmin();
         MockTaskAVSRegistrar newRegistrarImpl = new MockTaskAVSRegistrar(
-            IAllocationManager(address(allocationManager)), IKeyRegistrar(address(keyRegistrar))
+            IAllocationManager(address(allocationManagerMock)),
+            IKeyRegistrar(address(keyRegistrarMock)),
+            permissionController
         );
 
         // Expect event during initialization
@@ -190,7 +187,9 @@ contract TaskAVSRegistrarBaseUnitTests_Constructor is TaskAVSRegistrarBaseUnitTe
         // Deploy implementation
         ProxyAdmin newProxyAdmin = new ProxyAdmin();
         MockTaskAVSRegistrar newRegistrarImpl = new MockTaskAVSRegistrar(
-            IAllocationManager(address(allocationManager)), IKeyRegistrar(address(keyRegistrar))
+            IAllocationManager(address(allocationManagerMock)),
+            IKeyRegistrar(address(keyRegistrarMock)),
+            permissionController
         );
 
         // Expect revert during initialization
@@ -208,7 +207,9 @@ contract TaskAVSRegistrarBaseUnitTests_Constructor is TaskAVSRegistrarBaseUnitTe
         // Deploy implementation
         ProxyAdmin newProxyAdmin = new ProxyAdmin();
         MockTaskAVSRegistrar newRegistrarImpl = new MockTaskAVSRegistrar(
-            IAllocationManager(address(allocationManager)), IKeyRegistrar(address(keyRegistrar))
+            IAllocationManager(address(allocationManagerMock)),
+            IKeyRegistrar(address(keyRegistrarMock)),
+            permissionController
         );
 
         // Expect revert during initialization
@@ -226,7 +227,9 @@ contract TaskAVSRegistrarBaseUnitTests_Constructor is TaskAVSRegistrarBaseUnitTe
         // Deploy implementation
         ProxyAdmin newProxyAdmin = new ProxyAdmin();
         MockTaskAVSRegistrar newRegistrarImpl = new MockTaskAVSRegistrar(
-            IAllocationManager(address(allocationManager)), IKeyRegistrar(address(keyRegistrar))
+            IAllocationManager(address(allocationManagerMock)),
+            IKeyRegistrar(address(keyRegistrarMock)),
+            permissionController
         );
 
         // Expect revert during initialization
@@ -244,7 +247,9 @@ contract TaskAVSRegistrarBaseUnitTests_Constructor is TaskAVSRegistrarBaseUnitTe
         // Deploy implementation
         ProxyAdmin newProxyAdmin = new ProxyAdmin();
         MockTaskAVSRegistrar newRegistrarImpl = new MockTaskAVSRegistrar(
-            IAllocationManager(address(allocationManager)), IKeyRegistrar(address(keyRegistrar))
+            IAllocationManager(address(allocationManagerMock)),
+            IKeyRegistrar(address(keyRegistrarMock)),
+            permissionController
         );
 
         // Expect revert during initialization
@@ -416,7 +421,9 @@ contract TaskAVSRegistrarBaseUnitTests_Upgradeable is TaskAVSRegistrarBaseUnitTe
     function test_Implementation_CannotBeInitialized() public {
         // Deploy a new implementation
         MockTaskAVSRegistrar newImpl = new MockTaskAVSRegistrar(
-            IAllocationManager(address(allocationManager)), IKeyRegistrar(address(keyRegistrar))
+            IAllocationManager(address(allocationManagerMock)),
+            IKeyRegistrar(address(keyRegistrarMock)),
+            permissionController
         );
 
         // Try to initialize the implementation directly, should revert
@@ -438,7 +445,9 @@ contract TaskAVSRegistrarBaseUnitTests_Upgradeable is TaskAVSRegistrarBaseUnitTe
 
         // Deploy new implementation (could have new functions/logic)
         MockTaskAVSRegistrar newImpl = new MockTaskAVSRegistrar(
-            IAllocationManager(address(allocationManager)), IKeyRegistrar(address(keyRegistrar))
+            IAllocationManager(address(allocationManagerMock)),
+            IKeyRegistrar(address(keyRegistrarMock)),
+            permissionController
         );
 
         // Upgrade proxy to new implementation
@@ -463,7 +472,9 @@ contract TaskAVSRegistrarBaseUnitTests_Upgradeable is TaskAVSRegistrarBaseUnitTe
 
         // Deploy new implementation
         MockTaskAVSRegistrar newImpl = new MockTaskAVSRegistrar(
-            IAllocationManager(address(allocationManager)), IKeyRegistrar(address(keyRegistrar))
+            IAllocationManager(address(allocationManagerMock)),
+            IKeyRegistrar(address(keyRegistrarMock)),
+            permissionController
         );
 
         // Try to upgrade from non-owner, should revert
@@ -511,7 +522,9 @@ contract TaskAVSRegistrarBaseUnitTests_Upgradeable is TaskAVSRegistrarBaseUnitTe
 
         // Deploy new implementation
         MockTaskAVSRegistrar newImpl = new MockTaskAVSRegistrar(
-            IAllocationManager(address(allocationManager)), IKeyRegistrar(address(keyRegistrar))
+            IAllocationManager(address(allocationManagerMock)),
+            IKeyRegistrar(address(keyRegistrarMock)),
+            permissionController
         );
 
         // Upgrade
@@ -540,8 +553,9 @@ contract TaskAVSRegistrarBaseUnitTests_Upgradeable is TaskAVSRegistrarBaseUnitTe
         TransparentUpgradeableProxy uninitializedProxy = new TransparentUpgradeableProxy(
             address(
                 new MockTaskAVSRegistrar(
-                    IAllocationManager(address(allocationManager)),
-                    IKeyRegistrar(address(keyRegistrar))
+                    IAllocationManager(address(allocationManagerMock)),
+                    IKeyRegistrar(address(keyRegistrarMock)),
+                    permissionController
                 )
             ),
             address(new ProxyAdmin()),
@@ -564,7 +578,9 @@ contract TaskAVSRegistrarBaseUnitTests_Upgradeable is TaskAVSRegistrarBaseUnitTe
     function test_DisableInitializers_InImplementation() public {
         // This test verifies that the implementation contract has initializers disabled
         MockTaskAVSRegistrar impl = new MockTaskAVSRegistrar(
-            IAllocationManager(address(allocationManager)), IKeyRegistrar(address(keyRegistrar))
+            IAllocationManager(address(allocationManagerMock)),
+            IKeyRegistrar(address(keyRegistrarMock)),
+            permissionController
         );
 
         // Try to initialize the implementation, should revert
diff --git a/test/unit/middlewareV2/AVSRegistrarSocketUnit.t.sol b/test/unit/middlewareV2/AVSRegistrarSocketUnit.t.sol
