test and fix ChainWriter.RegisterForOperatorSets function

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"
@@ -579,6 +583,7 @@ func (w *ChainWriter) DeregisterFromOperatorSets(
 
 func (w *ChainWriter) RegisterForOperatorSets(
 	ctx context.Context,
+	registryCoordinatorAddr gethcommon.Address,
 	request RegistrationRequest,
 ) (*gethtypes.Receipt, error) {
 	if w.allocationManager == nil {
@@ -590,12 +595,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)
@@ -609,6 +629,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

@@ -7,13 +7,19 @@ import (
 	"testing"
 
 	"github.com/Layr-Labs/eigensdk-go/chainio/clients"
+	"github.com/Layr-Labs/eigensdk-go/chainio/clients/elcontracts"
+	"github.com/Layr-Labs/eigensdk-go/crypto/bls"
 	"github.com/Layr-Labs/eigensdk-go/logging"
 	"github.com/Layr-Labs/eigensdk-go/testutils"
 	"github.com/Layr-Labs/eigensdk-go/testutils/testclients"
 	"github.com/Layr-Labs/eigensdk-go/types"
+	"github.com/ethereum/go-ethereum/common"
 	"github.com/ethereum/go-ethereum/crypto"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
+
+	allocationmanager "github.com/Layr-Labs/eigensdk-go/contracts/bindings/AllocationManager"
+	regcoord "github.com/Layr-Labs/eigensdk-go/contracts/bindings/RegistryCoordinator"
 )
 
 func TestRegisterOperator(t *testing.T) {
@@ -151,3 +157,181 @@ func TestChainWriter(t *testing.T) {
 		assert.True(t, receipt.Status == 1)
 	})
 }
+
+func TestRegisterForOperatorSets(t *testing.T) {
+	eigenClients, anvilHttpEndpoint := testclients.BuildTestClients(t)
+	contractAddrs := testutils.GetContractAddressesFromContractRegistry(anvilHttpEndpoint)
+
+	avsAddress := common.HexToAddress("0xf39Fd6e51aad88F6F4ce6aB8827279cffFb92266")
+	operatorSetId := uint32(1)
+	erc20MockStrategyAddr := contractAddrs.Erc20MockStrategy
+
+	err := createOperatorSet(eigenClients, avsAddress, operatorSetId, erc20MockStrategyAddr)
+	require.NoError(t, err)
+
+	operatorAddressHex := "70997970C51812dc3A010C7d01b50e0d17dc79C8"
+	operatorPrivateKeyHex := "59c6995e998f97a5a0044966f0945389dc9e86dae88c7a8412f4603b6b78690d"
+
+	// Create operator clients
+	operatorClients, err := newTestClients(anvilHttpEndpoint, operatorPrivateKeyHex)
+	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),
+	}
+
+	registryCoordinatorAddress := contractAddrs.RegistryCoordinator
+	receipt, err := operatorClients.ElChainWriter.RegisterForOperatorSets(
+		context.Background(),
+		registryCoordinatorAddress,
+		request,
+	)
+
+	require.NoError(t, err)
+	const RECEIPT_SUCCESS_STATUS = uint64(1)
+	require.Equal(t, RECEIPT_SUCCESS_STATUS, receipt.Status)
+
+	isRegistered, err := operatorClients.ElChainReader.IsOperatorRegisteredWithOperatorSet(
+		context.Background(),
+		operatorAddress,
+		operatorSet,
+	)
+	require.NoError(t, err)
+	require.Equal(t, isRegistered, true)
+}
+
+func createOperatorSet(
+	client *clients.Clients,
+	avsAddress common.Address,
+	operatorSetId uint32,
+	erc20MockStrategyAddr common.Address,
+) error {
+	allocationManagerAddress := client.EigenlayerContractBindings.AllocationManagerAddr
+	allocationManager, err := allocationmanager.NewContractAllocationManager(
+		allocationManagerAddress,
+		client.EthHttpClient,
+	)
+	if err != nil {
+		return err
+	}
+	registryCoordinatorAddress := client.AvsRegistryContractBindings.RegistryCoordinatorAddr
+	registryCoordinator, err := regcoord.NewContractRegistryCoordinator(
+		registryCoordinatorAddress,
+		client.EthHttpClient,
+	)
+	if err != nil {
+		return err
+	}
+
+	noSendTxOpts, err := client.TxManager.GetNoSendTxOpts()
+	if err != nil {
+		return err
+	}
+
+	tx, err := allocationManager.SetAVSRegistrar(noSendTxOpts, avsAddress, registryCoordinatorAddress)
+	if err != nil {
+		return err
+	}
+
+	waitForReceipt := true
+
+	_, err = client.TxManager.Send(context.Background(), tx, waitForReceipt)
+	if err != nil {
+		return err
+	}
+
+	tx, err = registryCoordinator.EnableOperatorSets(noSendTxOpts)
+	if err != nil {
+		return err
+	}
+
+	_, err = client.TxManager.Send(context.Background(), tx, waitForReceipt)
+	if err != nil {
+		return err
+	}
+
+	operatorSetParam := regcoord.IRegistryCoordinatorOperatorSetParam{
+		MaxOperatorCount:        10,
+		KickBIPsOfOperatorStake: 100,
+		KickBIPsOfTotalStake:    1000,
+	}
+	minimumStake := big.NewInt(0)
+
+	strategyParams := regcoord.IStakeRegistryStrategyParams{
+		Strategy:   erc20MockStrategyAddr,
+		Multiplier: big.NewInt(1),
+	}
+	strategyParamsArray := []regcoord.IStakeRegistryStrategyParams{strategyParams}
+	lookAheadPeriod := uint32(0)
+	tx, err = registryCoordinator.CreateSlashableStakeQuorum(
+		noSendTxOpts,
+		operatorSetParam,
+		minimumStake,
+		strategyParamsArray,
+		lookAheadPeriod,
+	)
+	if err != nil {
+		return err
+	}
+
+	_, err = client.TxManager.Send(context.Background(), tx, waitForReceipt)
+	if err != nil {
+		return err
+	}
+
+	strategies := []common.Address{erc20MockStrategyAddr}
+	operatorSetParams := allocationmanager.IAllocationManagerTypesCreateSetParams{
+		OperatorSetId: operatorSetId,
+		Strategies:    strategies,
+	}
+	operatorSetParamsArray := []allocationmanager.IAllocationManagerTypesCreateSetParams{operatorSetParams}
+	tx, err = allocationManager.CreateOperatorSets(noSendTxOpts, avsAddress, operatorSetParamsArray)
+	if err != nil {
+		return err
+	}
+
+	_, err = client.TxManager.Send(context.Background(), tx, waitForReceipt)
+	return err
+}
+
+func newTestClients(httpEndpoint string, privateKeyHex string) (*clients.Clients, error) {
+	contractAddrs := testutils.GetContractAddressesFromContractRegistry(httpEndpoint)
+	chainioConfig := clients.BuildAllConfig{
+		EthHttpUrl:                 httpEndpoint,
+		EthWsUrl:                   httpEndpoint,
+		RegistryCoordinatorAddr:    contractAddrs.RegistryCoordinator.String(),
+		OperatorStateRetrieverAddr: contractAddrs.OperatorStateRetriever.String(),
+		AvsName:                    "exampleAvs",
+		PromMetricsIpPortAddress:   ":9090",
+	}
+	privateKey, err := crypto.HexToECDSA(privateKeyHex)
+	if err != nil {
+		return nil, err
+	}
+	testConfig := testutils.GetDefaultTestConfig()
+	logger := logging.NewTextSLogger(os.Stdout, &logging.SLoggerOptions{Level: testConfig.LogLevel})
+
+	testClients, err := clients.BuildAll(
+		chainioConfig,
+		privateKey,
+		logger,
+	)
+	if err != nil {
+		return nil, err
+	}
+	return testClients, nil
+}

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"
   }
-}
+}