test and fix ChainWriter.RegisterForOperatorSets function (#416)

Co-authored-by: Madhur Shrimal <[email protected]>
Co-authored-by: Brandon Chatham <[email protected]>
Co-authored-by: ricomateo <[email protected]>
Co-authored-by: Tomás Grüner <[email protected]>

chainio/clients/elcontracts/types.go

@@ -4,6 +4,7 @@ import (
 	"math/big"
 
 	allocationmanager "github.com/Layr-Labs/eigensdk-go/contracts/bindings/AllocationManager"
+	"github.com/Layr-Labs/eigensdk-go/crypto/bls"
 
 	"github.com/ethereum/go-ethereum/common"
 )
@@ -39,8 +40,9 @@ type RegistrationRequest struct {
 	AVSAddress      common.Address
 	OperatorSetIds  []uint32
 	WaitForReceipt  bool
+	BlsKeyPair      *bls.KeyPair
+	Socket          string
 }
-
 type RemovePermissionRequest struct {
 	AccountAddress   common.Address
 	AppointeeAddress common.Address

chainio/clients/elcontracts/writer.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"errors"
 
+	"github.com/ethereum/go-ethereum/accounts/abi"
 	"github.com/ethereum/go-ethereum/accounts/abi/bind"
 
 	"math/big"
@@ -13,14 +14,17 @@ import (
 
 	"github.com/Layr-Labs/eigensdk-go/chainio/clients/eth"
 	"github.com/Layr-Labs/eigensdk-go/chainio/txmgr"
+	chainioutils "github.com/Layr-Labs/eigensdk-go/chainio/utils"
 	avsdirectory "github.com/Layr-Labs/eigensdk-go/contracts/bindings/AVSDirectory"
 	allocationmanager "github.com/Layr-Labs/eigensdk-go/contracts/bindings/AllocationManager"
 	delegationmanager "github.com/Layr-Labs/eigensdk-go/contracts/bindings/DelegationManager"
 	erc20 "github.com/Layr-Labs/eigensdk-go/contracts/bindings/IERC20"
 	rewardscoordinator "github.com/Layr-Labs/eigensdk-go/contracts/bindings/IRewardsCoordinator"
 	strategy "github.com/Layr-Labs/eigensdk-go/contracts/bindings/IStrategy"
 	permissioncontroller "github.com/Layr-Labs/eigensdk-go/contracts/bindings/PermissionController"
+	regcoord "github.com/Layr-Labs/eigensdk-go/contracts/bindings/RegistryCoordinator"
 	strategymanager "github.com/Layr-Labs/eigensdk-go/contracts/bindings/StrategyManager"
+	"github.com/Layr-Labs/eigensdk-go/crypto/bls"
 	"github.com/Layr-Labs/eigensdk-go/logging"
 	"github.com/Layr-Labs/eigensdk-go/metrics"
 	"github.com/Layr-Labs/eigensdk-go/types"
@@ -534,6 +538,7 @@ func (w *ChainWriter) DeregisterFromOperatorSets(
 
 func (w *ChainWriter) RegisterForOperatorSets(
 	ctx context.Context,
+	registryCoordinatorAddr gethcommon.Address,
 	request RegistrationRequest,
 ) (*gethtypes.Receipt, error) {
 	if w.allocationManager == nil {
@@ -545,12 +550,27 @@ func (w *ChainWriter) RegisterForOperatorSets(
 		return nil, utils.WrapError("failed to get no send tx opts", err)
 	}
 
+	pubkeyRegParams, err := getPubkeyRegistrationParams(
+		w.ethClient,
+		registryCoordinatorAddr,
+		request.OperatorAddress,
+		request.BlsKeyPair,
+	)
+	if err != nil {
+		return nil, utils.WrapError("failed to get public key registration params", err)
+	}
+
+	data, err := abiEncodeRegistrationParams(request.Socket, *pubkeyRegParams)
+	if err != nil {
+		return nil, utils.WrapError("failed to encode registration params", err)
+	}
 	tx, err := w.allocationManager.RegisterForOperatorSets(
 		noSendTxOpts,
 		request.OperatorAddress,
 		allocationmanager.IAllocationManagerTypesRegisterParams{
 			Avs:            request.AVSAddress,
 			OperatorSetIds: request.OperatorSetIds,
+			Data:           data,
 		})
 	if err != nil {
 		return nil, utils.WrapError("failed to create RegisterForOperatorSets tx", err)
@@ -564,6 +584,82 @@ func (w *ChainWriter) RegisterForOperatorSets(
 	return receipt, nil
 }
 
+func getPubkeyRegistrationParams(
+	ethClient bind.ContractBackend,
+	registryCoordinatorAddr, operatorAddress gethcommon.Address,
+	blsKeyPair *bls.KeyPair,
+) (*regcoord.IBLSApkRegistryPubkeyRegistrationParams, error) {
+	registryCoordinator, err := regcoord.NewContractRegistryCoordinator(registryCoordinatorAddr, ethClient)
+	if err != nil {
+		return nil, utils.WrapError("failed to create registry coordinator", err)
+	}
+	// params to register bls pubkey with bls apk registry
+	g1HashedMsgToSign, err := registryCoordinator.PubkeyRegistrationMessageHash(
+		&bind.CallOpts{},
+		operatorAddress,
+	)
+	if err != nil {
+		return nil, err
+	}
+	signedMsg := chainioutils.ConvertToBN254G1Point(
+		blsKeyPair.SignHashedToCurveMessage(chainioutils.ConvertBn254GethToGnark(g1HashedMsgToSign)).G1Point,
+	)
+	G1pubkeyBN254 := chainioutils.ConvertToBN254G1Point(blsKeyPair.GetPubKeyG1())
+	G2pubkeyBN254 := chainioutils.ConvertToBN254G2Point(blsKeyPair.GetPubKeyG2())
+	pubkeyRegParams := regcoord.IBLSApkRegistryPubkeyRegistrationParams{
+		PubkeyRegistrationSignature: signedMsg,
+		PubkeyG1:                    G1pubkeyBN254,
+		PubkeyG2:                    G2pubkeyBN254,
+	}
+	return &pubkeyRegParams, nil
+}
+
+func abiEncodeRegistrationParams(
+	socket string,
+	pubkeyRegistrationParams regcoord.IBLSApkRegistryPubkeyRegistrationParams,
+) ([]byte, error) {
+	registrationParamsType, err := abi.NewType("tuple", "", []abi.ArgumentMarshaling{
+		{Name: "Socket", Type: "string"},
+		{Name: "PubkeyRegParams", Type: "tuple", Components: []abi.ArgumentMarshaling{
+			{Name: "PubkeyRegistrationSignature", Type: "tuple", Components: []abi.ArgumentMarshaling{
+				{Name: "X", Type: "uint256"},
+				{Name: "Y", Type: "uint256"},
+			}},
+			{Name: "PubkeyG1", Type: "tuple", Components: []abi.ArgumentMarshaling{
+				{Name: "X", Type: "uint256"},
+				{Name: "Y", Type: "uint256"},
+			}},
+			{Name: "PubkeyG2", Type: "tuple", Components: []abi.ArgumentMarshaling{
+				{Name: "X", Type: "uint256[2]"},
+				{Name: "Y", Type: "uint256[2]"},
+			}},
+		}},
+	})
+	if err != nil {
+		return nil, err
+	}
+
+	registrationParams := struct {
+		Socket          string
+		PubkeyRegParams regcoord.IBLSApkRegistryPubkeyRegistrationParams
+	}{
+		socket,
+		pubkeyRegistrationParams,
+	}
+
+	args := abi.Arguments{
+		{Type: registrationParamsType, Name: "registrationParams"},
+	}
+
+	data, err := args.Pack(&registrationParams)
+	if err != nil {
+		return nil, err
+	}
+	// The encoder is prepending 32 bytes to the data as if it was used in a dynamic function parameter.
+	// This is not used when decoding the bytes directly, so we need to remove it.
+	return data[32:], nil
+}
+
 func (w *ChainWriter) RemovePermission(
 	ctx context.Context,
 	request RemovePermissionRequest,

chainio/clients/elcontracts/writer_test.go

@@ -10,6 +10,8 @@ import (
 	"github.com/Layr-Labs/eigensdk-go/chainio/clients"
 	"github.com/Layr-Labs/eigensdk-go/chainio/clients/elcontracts"
 	allocationmanager "github.com/Layr-Labs/eigensdk-go/contracts/bindings/AllocationManager"
+	"github.com/Layr-Labs/eigensdk-go/crypto/bls"
+
 	rewardscoordinator "github.com/Layr-Labs/eigensdk-go/contracts/bindings/IRewardsCoordinator"
 	strategy "github.com/Layr-Labs/eigensdk-go/contracts/bindings/IStrategy"
 	mockerc20 "github.com/Layr-Labs/eigensdk-go/contracts/bindings/MockERC20"
@@ -119,6 +121,106 @@ func TestRegisterOperator(t *testing.T) {
 	})
 }
 
+func TestRegisterAndDeregisterFromOperatorSets(t *testing.T) {
+	testConfig := testutils.GetDefaultTestConfig()
+	anvilC, err := testutils.StartAnvilContainer(testConfig.AnvilStateFileName)
+	require.NoError(t, err)
+
+	anvilHttpEndpoint, err := anvilC.Endpoint(context.Background(), "http")
+	require.NoError(t, err)
+	contractAddrs := testutils.GetContractAddressesFromContractRegistry(anvilHttpEndpoint)
+
+	operatorAddressHex := "70997970C51812dc3A010C7d01b50e0d17dc79C8"
+	operatorPrivateKeyHex := "59c6995e998f97a5a0044966f0945389dc9e86dae88c7a8412f4603b6b78690d"
+
+	config := elcontracts.Config{
+		DelegationManagerAddress:  contractAddrs.DelegationManager,
+		RewardsCoordinatorAddress: contractAddrs.RewardsCoordinator,
+	}
+
+	// Create operator clients
+	chainWriter, err := testclients.NewTestChainWriterFromConfig(anvilHttpEndpoint, operatorPrivateKeyHex, config)
+	require.NoError(t, err)
+
+	chainReader, err := testclients.NewTestChainReaderFromConfig(anvilHttpEndpoint, config)
+	require.NoError(t, err)
+
+	avsAddress := common.HexToAddress("0xf39Fd6e51aad88F6F4ce6aB8827279cffFb92266")
+	operatorSetId := uint32(1)
+	erc20MockStrategyAddr := contractAddrs.Erc20MockStrategy
+
+	// Create an operator set to register an operator on it
+	err = createOperatorSet(
+		anvilHttpEndpoint,
+		testutils.ANVIL_FIRST_PRIVATE_KEY,
+		avsAddress,
+		operatorSetId,
+		erc20MockStrategyAddr,
+	)
+	require.NoError(t, err)
+
+	operatorAddress := common.HexToAddress(operatorAddressHex)
+	keypair, err := bls.NewKeyPairFromString("0x01")
+	require.NoError(t, err)
+
+	request := elcontracts.RegistrationRequest{
+		OperatorAddress: operatorAddress,
+		AVSAddress:      avsAddress,
+		OperatorSetIds:  []uint32{operatorSetId},
+		WaitForReceipt:  true,
+		Socket:          "socket",
+		BlsKeyPair:      keypair,
+	}
+
+	operatorSet := allocationmanager.OperatorSet{
+		Avs: avsAddress,
+		Id:  uint32(operatorSetId),
+	}
+	t.Run("register operator for operator set", func(t *testing.T) {
+		registryCoordinatorAddress := contractAddrs.RegistryCoordinator
+		receipt, err := chainWriter.RegisterForOperatorSets(
+			context.Background(),
+			registryCoordinatorAddress,
+			request,
+		)
+
+		require.NoError(t, err)
+		require.Equal(t, gethtypes.ReceiptStatusSuccessful, receipt.Status)
+
+		isRegistered, err := chainReader.IsOperatorRegisteredWithOperatorSet(
+			context.Background(),
+			operatorAddress,
+			operatorSet,
+		)
+		require.NoError(t, err)
+		require.Equal(t, true, isRegistered)
+	})
+
+	deregistrationRequest := elcontracts.DeregistrationRequest{
+		AVSAddress:     avsAddress,
+		OperatorSetIds: []uint32{operatorSetId},
+		WaitForReceipt: true,
+	}
+
+	t.Run("deregister operator from operator set", func(t *testing.T) {
+		receipt, err := chainWriter.DeregisterFromOperatorSets(
+			context.Background(),
+			operatorAddress,
+			deregistrationRequest,
+		)
+		require.NoError(t, err)
+		require.Equal(t, gethtypes.ReceiptStatusSuccessful, receipt.Status)
+
+		isRegistered, err := chainReader.IsOperatorRegisteredWithOperatorSet(
+			context.Background(),
+			operatorAddress,
+			operatorSet,
+		)
+		require.NoError(t, err)
+		require.Equal(t, false, isRegistered)
+	})
+}
+
 func TestChainWriter(t *testing.T) {
 	clients, anvilHttpEndpoint := testclients.BuildTestClients(t)
 	contractAddrs := testutils.GetContractAddressesFromContractRegistry(anvilHttpEndpoint)

contracts/script/DeployMockAvs.s.sol

@@ -87,4 +87,4 @@ contract DeployMockAvs is DeployMockAvsRegistries {
         }
         vm.stopBroadcast();
     }
-}
+}

contracts/script/DeployMockAvsRegistries.s.sol

@@ -211,4 +211,4 @@ contract DeployMockAvsRegistries is Script, ConfigsReadWriter, EigenlayerContrac
         registry.registerContract("strategyManager", address(eigen.strategyManager));
         registry.registerContract("rewardsCoordinator", address(eigen.rewardsCoordinator));
     }
-}
+}

contracts/script/output/31337/mockAvs_deployment_output.json

@@ -7,4 +7,4 @@
     "registryCoordinator": "0x95401dc811bb5740090279Ba06cfA8fcF6113778",
     "registryCoordinatorImplementation": "0x4c5859f0F772848b2D91F1D83E2Fe57935348029"
   }
-}
+}

contracts/script/output/31337/token_and_strategy_deployment_output.json

@@ -3,4 +3,4 @@
     "erc20mock": "0xFD471836031dc5108809D173A067e8486B9047A3",
     "erc20mockstrategy": "0x1429859428C0aBc9C2C47C8Ee9FBaf82cFA0F20f"
   }
-}
+}