From e6edc7d25336823228a459cd7116a87e8fd4fb72 Mon Sep 17 00:00:00 2001
From: devin <trestin_blockchain@hotmail.com>
Date: Fri, 21 Feb 2025 12:14:53 +0800
Subject: [PATCH 01/13] refactor: improve bls signature validation logic

---
 x/avs/keeper/keeper.go | 18 +++++++++++++++++-
 x/avs/types/utils.go   |  1 +
 2 files changed, 18 insertions(+), 1 deletion(-)

diff --git a/x/avs/keeper/keeper.go b/x/avs/keeper/keeper.go
index 5d76ab605..1a9db6bdd 100644
--- a/x/avs/keeper/keeper.go
+++ b/x/avs/keeper/keeper.go
@@ -288,8 +288,24 @@ func (k Keeper) CreateAVSTask(ctx sdk.Context, params *types.TaskInfoParams) (ui
 
 func (k Keeper) RegisterBLSPublicKey(ctx sdk.Context, params *types.BlsParams) error {
 	// check bls signature to prevent rogue key attacks
-	sig := params.PubkeyRegistrationSignature
+	// a message parameter to validate that this signature is intended solely for RegisterBLSPublicKey.
+	// The message should contain (ExoCore-$chain-id-$operator-address)  marker to prevent replay attacks.
+	// Looking forward to this format: "ExoCore-exocorelocalnet_232-exo13h6xg79g82e2g2vhjwg7j4r2z2hlncelwutkjr"
+	// Note that the address of the operator must be lowercase
+	expectedMessage := "ExoCore-" + types.ChainIDWithoutRevision(ctx.ChainID()) + "-" + strings.ToLower(params.OperatorAddress.String())
+
+	if params.Message != expectedMessage {
+		return fmt.Errorf("invalid message format: %s", params.Message)
+	}
+	// check msg hash
 	msgHash := params.PubkeyRegistrationMessageHash
+
+	expectedHash := crypto.Keccak256Hash([]byte(expectedMessage))
+	if !bytes.Equal(msgHash, expectedHash.Bytes()) {
+		return fmt.Errorf("the input hash failed to validate: %s", params.Message)
+	}
+
+	sig := params.PubkeyRegistrationSignature
 	pubKey, _ := bls.PublicKeyFromBytes(params.PubKey)
 	valid, err := blst.VerifySignature(sig, [32]byte(msgHash), pubKey)
 	if err != nil || !valid {
diff --git a/x/avs/types/utils.go b/x/avs/types/utils.go
index 825c3ea27..46f6232db 100644
--- a/x/avs/types/utils.go
+++ b/x/avs/types/utils.go
@@ -44,6 +44,7 @@ type BlsParams struct {
 	PubKey                        []byte
 	PubkeyRegistrationSignature   []byte
 	PubkeyRegistrationMessageHash []byte
+	Message                       string
 }
 
 type ProofParams struct {

From 51b010f0043960e4849017e256c29161e830e5ca Mon Sep 17 00:00:00 2001
From: devin <trestin_blockchain@hotmail.com>
Date: Fri, 21 Feb 2025 12:15:20 +0800
Subject: [PATCH 02/13] test: add unit tests for BLS signature functionality

---
 x/avs/keeper/avs_test.go | 25 +++++++++++++++++++++++++
 1 file changed, 25 insertions(+)

diff --git a/x/avs/keeper/avs_test.go b/x/avs/keeper/avs_test.go
index d62e93b89..68e7b0544 100644
--- a/x/avs/keeper/avs_test.go
+++ b/x/avs/keeper/avs_test.go
@@ -1,6 +1,9 @@
 package keeper_test
 
 import (
+	testutiltx "github.com/ExocoreNetwork/exocore/testutil/tx"
+	"github.com/ethereum/go-ethereum/crypto"
+	"github.com/prysmaticlabs/prysm/v4/crypto/bls/blst"
 	"math/big"
 	"strings"
 	"time"
@@ -195,3 +198,25 @@ func (suite *AVSTestSuite) TestAddressSwitch() {
 	commonAddress := common.Address(accAddress)
 	suite.Equal(common.HexToAddress("0x8dF46478a83Ab2a429979391E9546A12AfF9E33f"), commonAddress)
 }
+
+func (suite *AVSTestSuite) TestRegisterBLSPublicKey() {
+	privateKey, err := blst.RandKey()
+	suite.NoError(err)
+	publicKey := privateKey.PublicKey()
+	operatorAddress := sdk.AccAddress(utiltx.GenerateAddress().Bytes())
+	expectedMessage := "ExoCore-" + types.ChainIDWithoutRevision(suite.Ctx.ChainID()) + "-" + strings.ToLower(operatorAddress.String())
+	expectedHash := crypto.Keccak256Hash([]byte(expectedMessage))
+	sig := privateKey.Sign(expectedHash.Bytes())
+	params := &types.BlsParams{
+		OperatorAddress:               operatorAddress,
+		AvsAddress:                    testutiltx.GenerateAddress(),
+		PubKey:                        publicKey.Marshal(),
+		PubkeyRegistrationSignature:   sig.Marshal(),
+		PubkeyRegistrationMessageHash: expectedHash.Bytes(),
+		Message:                       expectedMessage,
+	}
+ 
+	err = suite.App.AVSManagerKeeper.RegisterBLSPublicKey(suite.Ctx, params)
+	suite.NoError(err)
+
+}

From 4d83208fe421cacf79c5c7036215e52794cfb41c Mon Sep 17 00:00:00 2001
From: devin <trestin_blockchain@hotmail.com>
Date: Fri, 21 Feb 2025 12:18:17 +0800
Subject: [PATCH 03/13] rename the bls variable name

---
 x/avs/keeper/keeper.go | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/x/avs/keeper/keeper.go b/x/avs/keeper/keeper.go
index 1a9db6bdd..cc477dcf2 100644
--- a/x/avs/keeper/keeper.go
+++ b/x/avs/keeper/keeper.go
@@ -314,7 +314,7 @@ func (k Keeper) RegisterBLSPublicKey(ctx sdk.Context, params *types.BlsParams) e
 	if k.IsExistPubKeyForAVS(ctx, params.OperatorAddress.String(), params.AvsAddress.String()) {
 		return errorsmod.Wrap(types.ErrAlreadyExists, fmt.Sprintf("the operator is :%s", params.OperatorAddress))
 	}
-	bls := &types.BlsPubKeyInfo{
+	blsInfo := &types.BlsPubKeyInfo{
 		AvsAddress:      strings.ToLower(params.AvsAddress.String()),
 		OperatorAddress: strings.ToLower(params.OperatorAddress.String()),
 		PubKey:          params.PubKey,
@@ -322,10 +322,10 @@ func (k Keeper) RegisterBLSPublicKey(ctx sdk.Context, params *types.BlsParams) e
 	// check a bls key can only be used once.
 	// if operator are using multiple servers for different AVSs .
 	// In case one server is compromised, signing can continue as expected on the AVSs for which there has been no compromise.
-	if k.IsExistPubKey(ctx, bls) {
+	if k.IsExistPubKey(ctx, blsInfo) {
 		return errorsmod.Wrap(types.ErrAlreadyExists, fmt.Sprintf("the bls key is already exists:%s", bls.PubKey))
 	}
-	return k.SetOperatorPubKey(ctx, bls)
+	return k.SetOperatorPubKey(ctx, blsInfo)
 }
 
 func (k Keeper) OperatorOptAction(ctx sdk.Context, params *types.OperatorOptParams) error {

From 16b8823889270f4904530e3c2715b6f55475980d Mon Sep 17 00:00:00 2001
From: devin <trestin_blockchain@hotmail.com>
Date: Fri, 21 Feb 2025 12:32:20 +0800
Subject: [PATCH 04/13] update the pre-compiled bls key registration interface
 parameters

---
 precompiles/avs/IAVSManager.sol | 10 ++++++----
 precompiles/avs/tx.go           |  9 +++++++--
 x/avs/keeper/avs_test.go        |  6 +++---
 x/avs/keeper/keeper.go          |  6 +++---
 x/avs/types/utils.go            |  4 ++--
 5 files changed, 21 insertions(+), 14 deletions(-)

diff --git a/precompiles/avs/IAVSManager.sol b/precompiles/avs/IAVSManager.sol
index 9b62e9fa4..9450d3cc5 100644
--- a/precompiles/avs/IAVSManager.sol
+++ b/precompiles/avs/IAVSManager.sol
@@ -195,14 +195,16 @@ interface IAVSManager {
     /// @param sender The external address for calling this method.
     /// @param avsAddress The address of AVS.
     /// @param pubKey the public keys of the operator
-    /// @param pubkeyRegistrationSignature the public keys of the operator
-    /// @param pubkeyRegistrationMessageHash the public keys of the operator
+    /// @param pubKeyRegistrationSignature the public keys of the operator
+    /// @param pubKeyRegistrationMessageHash the public keys of the operator
+    /// @param msg is the message which signed using a private key
     function registerBLSPublicKey(
         address sender,
         address avsAddress,
         bytes calldata pubKey,
-        bytes calldata pubkeyRegistrationSignature,
-        bytes calldata pubkeyRegistrationMessageHash
+        bytes calldata pubKeyRegistrationSignature,
+        bytes calldata pubKeyRegistrationMessageHash,
+		string msg
     ) external returns (bool success);
 
     /// @dev operatorSubmitTask ,  this function enables a operator submit a task result.
diff --git a/precompiles/avs/tx.go b/precompiles/avs/tx.go
index 67db55608..eb23bbf7a 100644
--- a/precompiles/avs/tx.go
+++ b/precompiles/avs/tx.go
@@ -314,14 +314,19 @@ func (p Precompile) RegisterBLSPublicKey(
 	if !ok {
 		return nil, fmt.Errorf(exocmn.ErrContractInputParamOrType, 3, "[]byte", pubKeyRegistrationSignature)
 	}
-	blsParams.PubkeyRegistrationSignature = pubKeyRegistrationSignature
+	blsParams.PubKeyRegistrationSignature = pubKeyRegistrationSignature
 
 	pubKeyRegistrationMessageHash, ok := args[4].([]byte)
 	if !ok {
 		return nil, fmt.Errorf(exocmn.ErrContractInputParamOrType, 4, "[]byte", pubKeyRegistrationMessageHash)
 	}
-	blsParams.PubkeyRegistrationMessageHash = pubKeyRegistrationMessageHash
+	blsParams.PubKeyRegistrationMessageHash = pubKeyRegistrationMessageHash
 
+	msg, ok := args[5].(string)
+	if !ok {
+		return nil, fmt.Errorf(exocmn.ErrContractInputParamOrType, 5, "string", msg)
+	}
+	blsParams.Message = msg
 	err := p.avsKeeper.RegisterBLSPublicKey(ctx, blsParams)
 	if err != nil {
 		return nil, err
diff --git a/x/avs/keeper/avs_test.go b/x/avs/keeper/avs_test.go
index 68e7b0544..7b74daf54 100644
--- a/x/avs/keeper/avs_test.go
+++ b/x/avs/keeper/avs_test.go
@@ -211,11 +211,11 @@ func (suite *AVSTestSuite) TestRegisterBLSPublicKey() {
 		OperatorAddress:               operatorAddress,
 		AvsAddress:                    testutiltx.GenerateAddress(),
 		PubKey:                        publicKey.Marshal(),
-		PubkeyRegistrationSignature:   sig.Marshal(),
-		PubkeyRegistrationMessageHash: expectedHash.Bytes(),
+		PubKeyRegistrationSignature:   sig.Marshal(),
+		PubKeyRegistrationMessageHash: expectedHash.Bytes(),
 		Message:                       expectedMessage,
 	}
- 
+
 	err = suite.App.AVSManagerKeeper.RegisterBLSPublicKey(suite.Ctx, params)
 	suite.NoError(err)
 
diff --git a/x/avs/keeper/keeper.go b/x/avs/keeper/keeper.go
index cc477dcf2..e5b2aaeeb 100644
--- a/x/avs/keeper/keeper.go
+++ b/x/avs/keeper/keeper.go
@@ -298,14 +298,14 @@ func (k Keeper) RegisterBLSPublicKey(ctx sdk.Context, params *types.BlsParams) e
 		return fmt.Errorf("invalid message format: %s", params.Message)
 	}
 	// check msg hash
-	msgHash := params.PubkeyRegistrationMessageHash
+	msgHash := params.PubKeyRegistrationMessageHash
 
 	expectedHash := crypto.Keccak256Hash([]byte(expectedMessage))
 	if !bytes.Equal(msgHash, expectedHash.Bytes()) {
 		return fmt.Errorf("the input hash failed to validate: %s", params.Message)
 	}
 
-	sig := params.PubkeyRegistrationSignature
+	sig := params.PubKeyRegistrationSignature
 	pubKey, _ := bls.PublicKeyFromBytes(params.PubKey)
 	valid, err := blst.VerifySignature(sig, [32]byte(msgHash), pubKey)
 	if err != nil || !valid {
@@ -323,7 +323,7 @@ func (k Keeper) RegisterBLSPublicKey(ctx sdk.Context, params *types.BlsParams) e
 	// if operator are using multiple servers for different AVSs .
 	// In case one server is compromised, signing can continue as expected on the AVSs for which there has been no compromise.
 	if k.IsExistPubKey(ctx, blsInfo) {
-		return errorsmod.Wrap(types.ErrAlreadyExists, fmt.Sprintf("the bls key is already exists:%s", bls.PubKey))
+		return errorsmod.Wrap(types.ErrAlreadyExists, fmt.Sprintf("the bls key is already exists:%s", blsInfo.PubKey))
 	}
 	return k.SetOperatorPubKey(ctx, blsInfo)
 }
diff --git a/x/avs/types/utils.go b/x/avs/types/utils.go
index 46f6232db..0788005c5 100644
--- a/x/avs/types/utils.go
+++ b/x/avs/types/utils.go
@@ -42,8 +42,8 @@ type BlsParams struct {
 	OperatorAddress               sdk.AccAddress
 	AvsAddress                    common.Address
 	PubKey                        []byte
-	PubkeyRegistrationSignature   []byte
-	PubkeyRegistrationMessageHash []byte
+	PubKeyRegistrationSignature   []byte
+	PubKeyRegistrationMessageHash []byte
 	Message                       string
 }
 

From f0d8826591ac6ffc28f30959c402a4da6152d8d8 Mon Sep 17 00:00:00 2001
From: devin <trestin_blockchain@hotmail.com>
Date: Fri, 21 Feb 2025 12:54:25 +0800
Subject: [PATCH 05/13] remove redundant parameters

---
 precompiles/avs/IAVSManager.sol |  6 +-----
 precompiles/avs/abi.json        |  7 +------
 precompiles/avs/tx.go           | 13 +------------
 x/avs/keeper/avs_test.go        | 12 +++++-------
 x/avs/keeper/keeper.go          | 13 ++-----------
 x/avs/types/utils.go            | 10 ++++------
 6 files changed, 14 insertions(+), 47 deletions(-)

diff --git a/precompiles/avs/IAVSManager.sol b/precompiles/avs/IAVSManager.sol
index 9450d3cc5..524f036cf 100644
--- a/precompiles/avs/IAVSManager.sol
+++ b/precompiles/avs/IAVSManager.sol
@@ -196,15 +196,11 @@ interface IAVSManager {
     /// @param avsAddress The address of AVS.
     /// @param pubKey the public keys of the operator
     /// @param pubKeyRegistrationSignature the public keys of the operator
-    /// @param pubKeyRegistrationMessageHash the public keys of the operator
-    /// @param msg is the message which signed using a private key
     function registerBLSPublicKey(
         address sender,
         address avsAddress,
         bytes calldata pubKey,
-        bytes calldata pubKeyRegistrationSignature,
-        bytes calldata pubKeyRegistrationMessageHash,
-		string msg
+        bytes calldata pubKeyRegistrationSignature
     ) external returns (bool success);
 
     /// @dev operatorSubmitTask ,  this function enables a operator submit a task result.
diff --git a/precompiles/avs/abi.json b/precompiles/avs/abi.json
index 0ed98a158..40e0f054d 100644
--- a/precompiles/avs/abi.json
+++ b/precompiles/avs/abi.json
@@ -1021,12 +1021,7 @@
 			},
 			{
 				"internalType": "bytes",
-				"name": "pubkeyRegistrationSignature",
-				"type": "bytes"
-			},
-			{
-				"internalType": "bytes",
-				"name": "pubkeyRegistrationMessageHash",
+				"name": "pubKeyRegistrationSignature",
 				"type": "bytes"
 			}
 		],
diff --git a/precompiles/avs/tx.go b/precompiles/avs/tx.go
index eb23bbf7a..91e92faf1 100644
--- a/precompiles/avs/tx.go
+++ b/precompiles/avs/tx.go
@@ -315,18 +315,7 @@ func (p Precompile) RegisterBLSPublicKey(
 		return nil, fmt.Errorf(exocmn.ErrContractInputParamOrType, 3, "[]byte", pubKeyRegistrationSignature)
 	}
 	blsParams.PubKeyRegistrationSignature = pubKeyRegistrationSignature
-
-	pubKeyRegistrationMessageHash, ok := args[4].([]byte)
-	if !ok {
-		return nil, fmt.Errorf(exocmn.ErrContractInputParamOrType, 4, "[]byte", pubKeyRegistrationMessageHash)
-	}
-	blsParams.PubKeyRegistrationMessageHash = pubKeyRegistrationMessageHash
-
-	msg, ok := args[5].(string)
-	if !ok {
-		return nil, fmt.Errorf(exocmn.ErrContractInputParamOrType, 5, "string", msg)
-	}
-	blsParams.Message = msg
+	
 	err := p.avsKeeper.RegisterBLSPublicKey(ctx, blsParams)
 	if err != nil {
 		return nil, err
diff --git a/x/avs/keeper/avs_test.go b/x/avs/keeper/avs_test.go
index 7b74daf54..b8a6bf2b3 100644
--- a/x/avs/keeper/avs_test.go
+++ b/x/avs/keeper/avs_test.go
@@ -204,16 +204,14 @@ func (suite *AVSTestSuite) TestRegisterBLSPublicKey() {
 	suite.NoError(err)
 	publicKey := privateKey.PublicKey()
 	operatorAddress := sdk.AccAddress(utiltx.GenerateAddress().Bytes())
-	expectedMessage := "ExoCore-" + types.ChainIDWithoutRevision(suite.Ctx.ChainID()) + "-" + strings.ToLower(operatorAddress.String())
+	expectedMessage := "BLS Signed Message-" + types.ChainIDWithoutRevision(suite.Ctx.ChainID()) + "-" + strings.ToLower(operatorAddress.String())
 	expectedHash := crypto.Keccak256Hash([]byte(expectedMessage))
 	sig := privateKey.Sign(expectedHash.Bytes())
 	params := &types.BlsParams{
-		OperatorAddress:               operatorAddress,
-		AvsAddress:                    testutiltx.GenerateAddress(),
-		PubKey:                        publicKey.Marshal(),
-		PubKeyRegistrationSignature:   sig.Marshal(),
-		PubKeyRegistrationMessageHash: expectedHash.Bytes(),
-		Message:                       expectedMessage,
+		OperatorAddress:             operatorAddress,
+		AvsAddress:                  testutiltx.GenerateAddress(),
+		PubKey:                      publicKey.Marshal(),
+		PubKeyRegistrationSignature: sig.Marshal(),
 	}
 
 	err = suite.App.AVSManagerKeeper.RegisterBLSPublicKey(suite.Ctx, params)
diff --git a/x/avs/keeper/keeper.go b/x/avs/keeper/keeper.go
index e5b2aaeeb..882f8e1e6 100644
--- a/x/avs/keeper/keeper.go
+++ b/x/avs/keeper/keeper.go
@@ -292,22 +292,13 @@ func (k Keeper) RegisterBLSPublicKey(ctx sdk.Context, params *types.BlsParams) e
 	// The message should contain (ExoCore-$chain-id-$operator-address)  marker to prevent replay attacks.
 	// Looking forward to this format: "ExoCore-exocorelocalnet_232-exo13h6xg79g82e2g2vhjwg7j4r2z2hlncelwutkjr"
 	// Note that the address of the operator must be lowercase
-	expectedMessage := "ExoCore-" + types.ChainIDWithoutRevision(ctx.ChainID()) + "-" + strings.ToLower(params.OperatorAddress.String())
-
-	if params.Message != expectedMessage {
-		return fmt.Errorf("invalid message format: %s", params.Message)
-	}
-	// check msg hash
-	msgHash := params.PubKeyRegistrationMessageHash
+	expectedMessage := "BLS Signed Message-" + types.ChainIDWithoutRevision(ctx.ChainID()) + "-" + strings.ToLower(params.OperatorAddress.String())
 
 	expectedHash := crypto.Keccak256Hash([]byte(expectedMessage))
-	if !bytes.Equal(msgHash, expectedHash.Bytes()) {
-		return fmt.Errorf("the input hash failed to validate: %s", params.Message)
-	}
 
 	sig := params.PubKeyRegistrationSignature
 	pubKey, _ := bls.PublicKeyFromBytes(params.PubKey)
-	valid, err := blst.VerifySignature(sig, [32]byte(msgHash), pubKey)
+	valid, err := blst.VerifySignature(sig, expectedHash, pubKey)
 	if err != nil || !valid {
 		return errorsmod.Wrap(types.ErrSigNotMatchPubKey, fmt.Sprintf("the operator is :%s", params.OperatorAddress))
 	}
diff --git a/x/avs/types/utils.go b/x/avs/types/utils.go
index 0788005c5..b37cd4a5d 100644
--- a/x/avs/types/utils.go
+++ b/x/avs/types/utils.go
@@ -39,12 +39,10 @@ type TaskInfoParams struct {
 	CallerAddress         sdk.AccAddress `json:"caller_address"`
 }
 type BlsParams struct {
-	OperatorAddress               sdk.AccAddress
-	AvsAddress                    common.Address
-	PubKey                        []byte
-	PubKeyRegistrationSignature   []byte
-	PubKeyRegistrationMessageHash []byte
-	Message                       string
+	OperatorAddress             sdk.AccAddress
+	AvsAddress                  common.Address
+	PubKey                      []byte
+	PubKeyRegistrationSignature []byte
 }
 
 type ProofParams struct {

From e9107b97b7e7e5f53f5b041235212b68b5cf1f77 Mon Sep 17 00:00:00 2001
From: devin <trestin_blockchain@hotmail.com>
Date: Fri, 21 Feb 2025 12:59:41 +0800
Subject: [PATCH 06/13] update comments

---
 precompiles/avs/tx.go    | 2 +-
 x/avs/keeper/avs_test.go | 8 ++++----
 x/avs/keeper/keeper.go   | 4 ++--
 3 files changed, 7 insertions(+), 7 deletions(-)

diff --git a/precompiles/avs/tx.go b/precompiles/avs/tx.go
index 91e92faf1..59f7a0b04 100644
--- a/precompiles/avs/tx.go
+++ b/precompiles/avs/tx.go
@@ -315,7 +315,7 @@ func (p Precompile) RegisterBLSPublicKey(
 		return nil, fmt.Errorf(exocmn.ErrContractInputParamOrType, 3, "[]byte", pubKeyRegistrationSignature)
 	}
 	blsParams.PubKeyRegistrationSignature = pubKeyRegistrationSignature
-	
+
 	err := p.avsKeeper.RegisterBLSPublicKey(ctx, blsParams)
 	if err != nil {
 		return nil, err
diff --git a/x/avs/keeper/avs_test.go b/x/avs/keeper/avs_test.go
index b8a6bf2b3..f652ad41f 100644
--- a/x/avs/keeper/avs_test.go
+++ b/x/avs/keeper/avs_test.go
@@ -1,13 +1,14 @@
 package keeper_test
 
 import (
-	testutiltx "github.com/ExocoreNetwork/exocore/testutil/tx"
-	"github.com/ethereum/go-ethereum/crypto"
-	"github.com/prysmaticlabs/prysm/v4/crypto/bls/blst"
 	"math/big"
 	"strings"
 	"time"
 
+	testutiltx "github.com/ExocoreNetwork/exocore/testutil/tx"
+	"github.com/ethereum/go-ethereum/crypto"
+	"github.com/prysmaticlabs/prysm/v4/crypto/bls/blst"
+
 	"cosmossdk.io/math"
 	assetstypes "github.com/ExocoreNetwork/exocore/x/assets/types"
 	"github.com/ethereum/go-ethereum/common"
@@ -216,5 +217,4 @@ func (suite *AVSTestSuite) TestRegisterBLSPublicKey() {
 
 	err = suite.App.AVSManagerKeeper.RegisterBLSPublicKey(suite.Ctx, params)
 	suite.NoError(err)
-
 }
diff --git a/x/avs/keeper/keeper.go b/x/avs/keeper/keeper.go
index 882f8e1e6..8ffee0cda 100644
--- a/x/avs/keeper/keeper.go
+++ b/x/avs/keeper/keeper.go
@@ -289,8 +289,8 @@ func (k Keeper) CreateAVSTask(ctx sdk.Context, params *types.TaskInfoParams) (ui
 func (k Keeper) RegisterBLSPublicKey(ctx sdk.Context, params *types.BlsParams) error {
 	// check bls signature to prevent rogue key attacks
 	// a message parameter to validate that this signature is intended solely for RegisterBLSPublicKey.
-	// The message should contain (ExoCore-$chain-id-$operator-address)  marker to prevent replay attacks.
-	// Looking forward to this format: "ExoCore-exocorelocalnet_232-exo13h6xg79g82e2g2vhjwg7j4r2z2hlncelwutkjr"
+	// The message should contain (BLS Signed Message-$chain-id-$operator-address)  marker to prevent replay attacks.
+	// Looking forward to this format: "BLS Signed Message-exocorelocalnet_232-exo13h6xg79g82e2g2vhjwg7j4r2z2hlncelwutkjr"
 	// Note that the address of the operator must be lowercase
 	expectedMessage := "BLS Signed Message-" + types.ChainIDWithoutRevision(ctx.ChainID()) + "-" + strings.ToLower(params.OperatorAddress.String())
 

From 1ec0e098f12f8854b27796bc7e1dd8440ec09496 Mon Sep 17 00:00:00 2001
From: devin <trestin_blockchain@hotmail.com>
Date: Fri, 21 Feb 2025 13:11:31 +0800
Subject: [PATCH 07/13] fix some comments

---
 precompiles/avs/IAVSManager.sol | 2 +-
 x/avs/keeper/avs_test.go        | 4 +++-
 x/avs/keeper/keeper.go          | 6 +++---
 x/avs/types/types.go            | 5 ++++-
 4 files changed, 11 insertions(+), 6 deletions(-)

diff --git a/precompiles/avs/IAVSManager.sol b/precompiles/avs/IAVSManager.sol
index 524f036cf..c8cd1b81c 100644
--- a/precompiles/avs/IAVSManager.sol
+++ b/precompiles/avs/IAVSManager.sol
@@ -195,7 +195,7 @@ interface IAVSManager {
     /// @param sender The external address for calling this method.
     /// @param avsAddress The address of AVS.
     /// @param pubKey the public keys of the operator
-    /// @param pubKeyRegistrationSignature the public keys of the operator
+    /// @param pubKeyRegistrationSignature the bls signature of the operator
     function registerBLSPublicKey(
         address sender,
         address avsAddress,
diff --git a/x/avs/keeper/avs_test.go b/x/avs/keeper/avs_test.go
index f652ad41f..3757d4a51 100644
--- a/x/avs/keeper/avs_test.go
+++ b/x/avs/keeper/avs_test.go
@@ -1,6 +1,7 @@
 package keeper_test
 
 import (
+	"fmt"
 	"math/big"
 	"strings"
 	"time"
@@ -205,7 +206,8 @@ func (suite *AVSTestSuite) TestRegisterBLSPublicKey() {
 	suite.NoError(err)
 	publicKey := privateKey.PublicKey()
 	operatorAddress := sdk.AccAddress(utiltx.GenerateAddress().Bytes())
-	expectedMessage := "BLS Signed Message-" + types.ChainIDWithoutRevision(suite.Ctx.ChainID()) + "-" + strings.ToLower(operatorAddress.String())
+	expectedMessage := types.SignatureHeader + "\n" + types.ChainIDWithoutRevision(suite.Ctx.ChainID()) + "\n" + strings.ToLower(operatorAddress.String())
+	fmt.Println(expectedMessage)
 	expectedHash := crypto.Keccak256Hash([]byte(expectedMessage))
 	sig := privateKey.Sign(expectedHash.Bytes())
 	params := &types.BlsParams{
diff --git a/x/avs/keeper/keeper.go b/x/avs/keeper/keeper.go
index 8ffee0cda..800525a3f 100644
--- a/x/avs/keeper/keeper.go
+++ b/x/avs/keeper/keeper.go
@@ -292,13 +292,13 @@ func (k Keeper) RegisterBLSPublicKey(ctx sdk.Context, params *types.BlsParams) e
 	// The message should contain (BLS Signed Message-$chain-id-$operator-address)  marker to prevent replay attacks.
 	// Looking forward to this format: "BLS Signed Message-exocorelocalnet_232-exo13h6xg79g82e2g2vhjwg7j4r2z2hlncelwutkjr"
 	// Note that the address of the operator must be lowercase
-	expectedMessage := "BLS Signed Message-" + types.ChainIDWithoutRevision(ctx.ChainID()) + "-" + strings.ToLower(params.OperatorAddress.String())
+	expectedMessage := types.SignatureHeader + "\n" + types.ChainIDWithoutRevision(ctx.ChainID()) + "\n" + strings.ToLower(params.OperatorAddress.String())
 
-	expectedHash := crypto.Keccak256Hash([]byte(expectedMessage))
+	hashedMsg := crypto.Keccak256Hash([]byte(expectedMessage))
 
 	sig := params.PubKeyRegistrationSignature
 	pubKey, _ := bls.PublicKeyFromBytes(params.PubKey)
-	valid, err := blst.VerifySignature(sig, expectedHash, pubKey)
+	valid, err := blst.VerifySignature(sig, hashedMsg, pubKey)
 	if err != nil || !valid {
 		return errorsmod.Wrap(types.ErrSigNotMatchPubKey, fmt.Sprintf("the operator is :%s", params.OperatorAddress))
 	}
diff --git a/x/avs/types/types.go b/x/avs/types/types.go
index 6b6aa98db..62b95d8ba 100644
--- a/x/avs/types/types.go
+++ b/x/avs/types/types.go
@@ -48,7 +48,10 @@ var (
 	ChainIDPrefix = []byte("chain-id-prefix")
 )
 
-const InvalidTaskID = 0
+const (
+	InvalidTaskID   = 0
+	SignatureHeader = "BLS Signed Message"
+)
 
 type AVSRegisterOrDeregisterParams struct {
 	// AvsName is the name of the AVS as an arbitrary string.

From 007dfb6ae2c8a1cb647b406d501b45d8a30aaa8e Mon Sep 17 00:00:00 2001
From: MaxMustermann2 <82761650+MaxMustermann2@users.noreply.github.com>
Date: Fri, 21 Feb 2025 05:48:13 +0000
Subject: [PATCH 08/13] refactor(avs): make bls-msg fmt + update tests

---
 x/avs/keeper/avs_test.go | 129 ++++++++++++++++++++++++++++++++-------
 x/avs/keeper/keeper.go   |  15 +++--
 x/avs/types/types.go     |   4 +-
 3 files changed, 116 insertions(+), 32 deletions(-)

diff --git a/x/avs/keeper/avs_test.go b/x/avs/keeper/avs_test.go
index 3757d4a51..14a238db4 100644
--- a/x/avs/keeper/avs_test.go
+++ b/x/avs/keeper/avs_test.go
@@ -8,6 +8,7 @@ import (
 
 	testutiltx "github.com/ExocoreNetwork/exocore/testutil/tx"
 	"github.com/ethereum/go-ethereum/crypto"
+	"github.com/prysmaticlabs/prysm/v4/crypto/bls"
 	"github.com/prysmaticlabs/prysm/v4/crypto/bls/blst"
 
 	"cosmossdk.io/math"
@@ -15,7 +16,6 @@ import (
 	"github.com/ethereum/go-ethereum/common"
 
 	"github.com/ExocoreNetwork/exocore/x/avs/types"
-	avstypes "github.com/ExocoreNetwork/exocore/x/avs/types"
 	delegationtypes "github.com/ExocoreNetwork/exocore/x/delegation/types"
 	epochstypes "github.com/ExocoreNetwork/exocore/x/epochs/types"
 	operatorTypes "github.com/ExocoreNetwork/exocore/x/operator/types"
@@ -83,7 +83,7 @@ func (suite *AVSTestSuite) TestUpdateAVSInfo_Register() {
 	avsParams := &types.AVSRegisterOrDeregisterParams{
 		AvsName:               avsName,
 		AvsAddress:            common.HexToAddress(avsAddres),
-		Action:                avstypes.RegisterAction,
+		Action:                types.RegisterAction,
 		RewardContractAddress: common.HexToAddress(rewardAddress),
 		AvsOwnerAddresses:     avsOwnerAddress,
 		AssetIDs:              assetIDs,
@@ -115,7 +115,7 @@ func (suite *AVSTestSuite) TestUpdateAVSInfo_DeRegister() {
 	avsParams := &types.AVSRegisterOrDeregisterParams{
 		AvsName:              avsName,
 		AvsAddress:           common.HexToAddress(avsAddress),
-		Action:               avstypes.DeRegisterAction,
+		Action:               types.DeRegisterAction,
 		AvsOwnerAddresses:    avsOwnerAddress,
 		AssetIDs:             assetIDs,
 		MinSelfDelegation:    uint64(10),
@@ -128,10 +128,11 @@ func (suite *AVSTestSuite) TestUpdateAVSInfo_DeRegister() {
 	suite.Error(err)
 	suite.Contains(err.Error(), types.ErrUnregisterNonExistent.Error())
 
-	avsParams.Action = avstypes.RegisterAction
+	avsParams.Action = types.RegisterAction
 	err = suite.App.AVSManagerKeeper.UpdateAVSInfo(suite.Ctx, avsParams)
 	suite.NoError(err)
 	info, err := suite.App.AVSManagerKeeper.GetAVSInfo(suite.Ctx, avsAddress)
+	suite.NoError(err)
 	suite.Equal(strings.ToLower(avsAddress), info.GetInfo().AvsAddress)
 
 	epoch, _ := suite.App.EpochsKeeper.GetEpochInfo(suite.Ctx, epochstypes.DayEpochID)
@@ -143,22 +144,24 @@ func (suite *AVSTestSuite) TestUpdateAVSInfo_DeRegister() {
 		suite.Equal(epoch.CurrentEpoch, epochEnd+1)
 	}
 
-	avsParams.Action = avstypes.DeRegisterAction
+	avsParams.Action = types.DeRegisterAction
 	avsParams.CallerAddress, err = sdk.AccAddressFromBech32("exo13h6xg79g82e2g2vhjwg7j4r2z2hlncelwutkjr")
+	suite.NoError(err)
 	err = suite.App.AVSManagerKeeper.UpdateAVSInfo(suite.Ctx, avsParams)
 	suite.NoError(err)
 	info, err = suite.App.AVSManagerKeeper.GetAVSInfo(suite.Ctx, avsAddress)
 	suite.Error(err)
 	suite.Contains(err.Error(), types.ErrNoKeyInTheStore.Error())
+	suite.Nil(info)
 }
 
 func (suite *AVSTestSuite) TestUpdateAVSInfoWithOperator_Register() {
 	avsAddress := suite.avsAddress
 	operatorAddress := sdk.AccAddress(utiltx.GenerateAddress().Bytes())
 
-	operatorParams := &avstypes.OperatorOptParams{
+	operatorParams := &types.OperatorOptParams{
 		AvsAddress:      avsAddress,
-		Action:          avstypes.RegisterAction,
+		Action:          types.RegisterAction,
 		OperatorAddress: operatorAddress,
 	}
 	//  operator Not Exist
@@ -202,21 +205,103 @@ func (suite *AVSTestSuite) TestAddressSwitch() {
 }
 
 func (suite *AVSTestSuite) TestRegisterBLSPublicKey() {
-	privateKey, err := blst.RandKey()
-	suite.NoError(err)
-	publicKey := privateKey.PublicKey()
-	operatorAddress := sdk.AccAddress(utiltx.GenerateAddress().Bytes())
-	expectedMessage := types.SignatureHeader + "\n" + types.ChainIDWithoutRevision(suite.Ctx.ChainID()) + "\n" + strings.ToLower(operatorAddress.String())
-	fmt.Println(expectedMessage)
-	expectedHash := crypto.Keccak256Hash([]byte(expectedMessage))
-	sig := privateKey.Sign(expectedHash.Bytes())
-	params := &types.BlsParams{
-		OperatorAddress:             operatorAddress,
-		AvsAddress:                  testutiltx.GenerateAddress(),
-		PubKey:                      publicKey.Marshal(),
-		PubKeyRegistrationSignature: sig.Marshal(),
+	type testCase struct {
+		name          string
+		setupParams   func() (bls.SecretKey, sdk.AccAddress, *types.BlsParams)
+		expectedError error
+		errorContains string
 	}
 
-	err = suite.App.AVSManagerKeeper.RegisterBLSPublicKey(suite.Ctx, params)
-	suite.NoError(err)
+	testCases := []testCase{
+		{
+			name: "successful registration",
+			setupParams: func() (bls.SecretKey, sdk.AccAddress, *types.BlsParams) {
+				privateKey, _ := blst.RandKey()
+				operatorAddress := sdk.AccAddress(utiltx.GenerateAddress().Bytes())
+				msg := fmt.Sprintf(types.BLSMessageToSign, types.ChainIDWithoutRevision(suite.Ctx.ChainID()), operatorAddress.String())
+				hashedMsg := crypto.Keccak256Hash([]byte(msg))
+				sig := privateKey.Sign(hashedMsg.Bytes())
+
+				return privateKey, operatorAddress, &types.BlsParams{
+					OperatorAddress:             operatorAddress,
+					AvsAddress:                  testutiltx.GenerateAddress(),
+					PubKey:                      privateKey.PublicKey().Marshal(),
+					PubKeyRegistrationSignature: sig.Marshal(),
+				}
+			},
+			expectedError: nil,
+		},
+		{
+			name: "wrong chain ID",
+			setupParams: func() (bls.SecretKey, sdk.AccAddress, *types.BlsParams) {
+				privateKey, _ := blst.RandKey()
+				operatorAddress := sdk.AccAddress(utiltx.GenerateAddress().Bytes())
+				msg := fmt.Sprintf(types.BLSMessageToSign, "onemorechain_211", operatorAddress.String())
+				hashedMsg := crypto.Keccak256Hash([]byte(msg))
+				sig := privateKey.Sign(hashedMsg.Bytes())
+
+				return privateKey, operatorAddress, &types.BlsParams{
+					OperatorAddress:             operatorAddress,
+					AvsAddress:                  testutiltx.GenerateAddress(),
+					PubKey:                      privateKey.PublicKey().Marshal(),
+					PubKeyRegistrationSignature: sig.Marshal(),
+				}
+			},
+			errorContains: types.ErrSigNotMatchPubKey.Error(),
+		},
+		{
+			name: "wrong operator address in signature",
+			setupParams: func() (bls.SecretKey, sdk.AccAddress, *types.BlsParams) {
+				privateKey, _ := blst.RandKey()
+				operatorAddress := sdk.AccAddress(utiltx.GenerateAddress().Bytes())
+				anotherOperatorAddress := sdk.AccAddress(utiltx.GenerateAddress().Bytes())
+				msg := fmt.Sprintf(types.BLSMessageToSign, types.ChainIDWithoutRevision(suite.Ctx.ChainID()), anotherOperatorAddress.String())
+				hashedMsg := crypto.Keccak256Hash([]byte(msg))
+				sig := privateKey.Sign(hashedMsg.Bytes())
+
+				return privateKey, operatorAddress, &types.BlsParams{
+					OperatorAddress:             operatorAddress,
+					AvsAddress:                  testutiltx.GenerateAddress(),
+					PubKey:                      privateKey.PublicKey().Marshal(),
+					PubKeyRegistrationSignature: sig.Marshal(),
+				}
+			},
+			errorContains: types.ErrSigNotMatchPubKey.Error(),
+		},
+		{
+			name: "mismatched BLS key",
+			setupParams: func() (bls.SecretKey, sdk.AccAddress, *types.BlsParams) {
+				privateKey, _ := blst.RandKey()
+				anotherPrivateKey, _ := blst.RandKey()
+				operatorAddress := sdk.AccAddress(utiltx.GenerateAddress().Bytes())
+				msg := fmt.Sprintf(types.BLSMessageToSign, types.ChainIDWithoutRevision(suite.Ctx.ChainID()), operatorAddress.String())
+				hashedMsg := crypto.Keccak256Hash([]byte(msg))
+				sig := privateKey.Sign(hashedMsg.Bytes())
+
+				return privateKey, operatorAddress, &types.BlsParams{
+					OperatorAddress:             operatorAddress,
+					AvsAddress:                  testutiltx.GenerateAddress(),
+					PubKey:                      anotherPrivateKey.PublicKey().Marshal(),
+					PubKeyRegistrationSignature: sig.Marshal(),
+				}
+			},
+			errorContains: types.ErrSigNotMatchPubKey.Error(),
+		},
+	}
+
+	for _, tc := range testCases {
+		suite.Run(tc.name, func() {
+			_, _, params := tc.setupParams()
+			err := suite.App.AVSManagerKeeper.RegisterBLSPublicKey(suite.Ctx, params)
+
+			if tc.expectedError != nil {
+				suite.Equal(tc.expectedError, err)
+			} else if tc.errorContains != "" {
+				suite.Error(err)
+				suite.Contains(err.Error(), tc.errorContains)
+			} else {
+				suite.NoError(err)
+			}
+		})
+	}
 }
diff --git a/x/avs/keeper/keeper.go b/x/avs/keeper/keeper.go
index 800525a3f..8d0eb724c 100644
--- a/x/avs/keeper/keeper.go
+++ b/x/avs/keeper/keeper.go
@@ -288,13 +288,12 @@ func (k Keeper) CreateAVSTask(ctx sdk.Context, params *types.TaskInfoParams) (ui
 
 func (k Keeper) RegisterBLSPublicKey(ctx sdk.Context, params *types.BlsParams) error {
 	// check bls signature to prevent rogue key attacks
-	// a message parameter to validate that this signature is intended solely for RegisterBLSPublicKey.
-	// The message should contain (BLS Signed Message-$chain-id-$operator-address)  marker to prevent replay attacks.
-	// Looking forward to this format: "BLS Signed Message-exocorelocalnet_232-exo13h6xg79g82e2g2vhjwg7j4r2z2hlncelwutkjr"
-	// Note that the address of the operator must be lowercase
-	expectedMessage := types.SignatureHeader + "\n" + types.ChainIDWithoutRevision(ctx.ChainID()) + "\n" + strings.ToLower(params.OperatorAddress.String())
-
-	hashedMsg := crypto.Keccak256Hash([]byte(expectedMessage))
+	// use a templated message to
+	// (1) validate that this signature is intended solely for RegisterBLSPublicKey
+	// (2) prevent replay attacks by including the chain-id and operator-address
+	// note that the operator address is bech32 encoded and thus already lowercase
+	msg := fmt.Sprintf(types.BLSMessageToSign, types.ChainIDWithoutRevision(ctx.ChainID()), params.OperatorAddress.String())
+	hashedMsg := crypto.Keccak256Hash([]byte(msg))
 
 	sig := params.PubKeyRegistrationSignature
 	pubKey, _ := bls.PublicKeyFromBytes(params.PubKey)
@@ -307,7 +306,7 @@ func (k Keeper) RegisterBLSPublicKey(ctx sdk.Context, params *types.BlsParams) e
 	}
 	blsInfo := &types.BlsPubKeyInfo{
 		AvsAddress:      strings.ToLower(params.AvsAddress.String()),
-		OperatorAddress: strings.ToLower(params.OperatorAddress.String()),
+		OperatorAddress: params.OperatorAddress.String(),
 		PubKey:          params.PubKey,
 	}
 	// check a bls key can only be used once.
diff --git a/x/avs/types/types.go b/x/avs/types/types.go
index 62b95d8ba..b16212a1f 100644
--- a/x/avs/types/types.go
+++ b/x/avs/types/types.go
@@ -49,8 +49,8 @@ var (
 )
 
 const (
-	InvalidTaskID   = 0
-	SignatureHeader = "BLS Signed Message"
+	InvalidTaskID    = 0
+	BLSMessageToSign = "BLS12-381 Signed Message\nChainIDWithoutRevision: %s\nAccAddressBech32: %s"
 )
 
 type AVSRegisterOrDeregisterParams struct {

From 511788b5a42e368775d47375244a91ff9c1cbdf9 Mon Sep 17 00:00:00 2001
From: MaxMustermann2 <82761650+MaxMustermann2@users.noreply.github.com>
Date: Fri, 21 Feb 2025 05:55:45 +0000
Subject: [PATCH 09/13] test(avs): remove redundant return

---
 x/avs/keeper/avs_test.go | 42 ++++++++++++++++++++--------------------
 1 file changed, 21 insertions(+), 21 deletions(-)

diff --git a/x/avs/keeper/avs_test.go b/x/avs/keeper/avs_test.go
index 14a238db4..5b6d09894 100644
--- a/x/avs/keeper/avs_test.go
+++ b/x/avs/keeper/avs_test.go
@@ -8,7 +8,6 @@ import (
 
 	testutiltx "github.com/ExocoreNetwork/exocore/testutil/tx"
 	"github.com/ethereum/go-ethereum/crypto"
-	"github.com/prysmaticlabs/prysm/v4/crypto/bls"
 	"github.com/prysmaticlabs/prysm/v4/crypto/bls/blst"
 
 	"cosmossdk.io/math"
@@ -207,40 +206,40 @@ func (suite *AVSTestSuite) TestAddressSwitch() {
 func (suite *AVSTestSuite) TestRegisterBLSPublicKey() {
 	type testCase struct {
 		name          string
-		setupParams   func() (bls.SecretKey, sdk.AccAddress, *types.BlsParams)
-		expectedError error
+		setupParams   func() *types.BlsParams
 		errorContains string
 	}
 
 	testCases := []testCase{
 		{
 			name: "successful registration",
-			setupParams: func() (bls.SecretKey, sdk.AccAddress, *types.BlsParams) {
-				privateKey, _ := blst.RandKey()
+			setupParams: func() *types.BlsParams {
+				privateKey, err := blst.RandKey()
+				suite.NoError(err)
 				operatorAddress := sdk.AccAddress(utiltx.GenerateAddress().Bytes())
 				msg := fmt.Sprintf(types.BLSMessageToSign, types.ChainIDWithoutRevision(suite.Ctx.ChainID()), operatorAddress.String())
 				hashedMsg := crypto.Keccak256Hash([]byte(msg))
 				sig := privateKey.Sign(hashedMsg.Bytes())
 
-				return privateKey, operatorAddress, &types.BlsParams{
+				return &types.BlsParams{
 					OperatorAddress:             operatorAddress,
 					AvsAddress:                  testutiltx.GenerateAddress(),
 					PubKey:                      privateKey.PublicKey().Marshal(),
 					PubKeyRegistrationSignature: sig.Marshal(),
 				}
 			},
-			expectedError: nil,
 		},
 		{
 			name: "wrong chain ID",
-			setupParams: func() (bls.SecretKey, sdk.AccAddress, *types.BlsParams) {
-				privateKey, _ := blst.RandKey()
+			setupParams: func() *types.BlsParams {
+				privateKey, err := blst.RandKey()
+				suite.NoError(err)
 				operatorAddress := sdk.AccAddress(utiltx.GenerateAddress().Bytes())
 				msg := fmt.Sprintf(types.BLSMessageToSign, "onemorechain_211", operatorAddress.String())
 				hashedMsg := crypto.Keccak256Hash([]byte(msg))
 				sig := privateKey.Sign(hashedMsg.Bytes())
 
-				return privateKey, operatorAddress, &types.BlsParams{
+				return &types.BlsParams{
 					OperatorAddress:             operatorAddress,
 					AvsAddress:                  testutiltx.GenerateAddress(),
 					PubKey:                      privateKey.PublicKey().Marshal(),
@@ -251,15 +250,16 @@ func (suite *AVSTestSuite) TestRegisterBLSPublicKey() {
 		},
 		{
 			name: "wrong operator address in signature",
-			setupParams: func() (bls.SecretKey, sdk.AccAddress, *types.BlsParams) {
-				privateKey, _ := blst.RandKey()
+			setupParams: func() *types.BlsParams {
+				privateKey, err := blst.RandKey()
+				suite.NoError(err)
 				operatorAddress := sdk.AccAddress(utiltx.GenerateAddress().Bytes())
 				anotherOperatorAddress := sdk.AccAddress(utiltx.GenerateAddress().Bytes())
 				msg := fmt.Sprintf(types.BLSMessageToSign, types.ChainIDWithoutRevision(suite.Ctx.ChainID()), anotherOperatorAddress.String())
 				hashedMsg := crypto.Keccak256Hash([]byte(msg))
 				sig := privateKey.Sign(hashedMsg.Bytes())
 
-				return privateKey, operatorAddress, &types.BlsParams{
+				return &types.BlsParams{
 					OperatorAddress:             operatorAddress,
 					AvsAddress:                  testutiltx.GenerateAddress(),
 					PubKey:                      privateKey.PublicKey().Marshal(),
@@ -270,15 +270,17 @@ func (suite *AVSTestSuite) TestRegisterBLSPublicKey() {
 		},
 		{
 			name: "mismatched BLS key",
-			setupParams: func() (bls.SecretKey, sdk.AccAddress, *types.BlsParams) {
-				privateKey, _ := blst.RandKey()
-				anotherPrivateKey, _ := blst.RandKey()
+			setupParams: func() *types.BlsParams {
+				privateKey, err := blst.RandKey()
+				suite.NoError(err)
+				anotherPrivateKey, err := blst.RandKey()
+				suite.NoError(err)
 				operatorAddress := sdk.AccAddress(utiltx.GenerateAddress().Bytes())
 				msg := fmt.Sprintf(types.BLSMessageToSign, types.ChainIDWithoutRevision(suite.Ctx.ChainID()), operatorAddress.String())
 				hashedMsg := crypto.Keccak256Hash([]byte(msg))
 				sig := privateKey.Sign(hashedMsg.Bytes())
 
-				return privateKey, operatorAddress, &types.BlsParams{
+				return &types.BlsParams{
 					OperatorAddress:             operatorAddress,
 					AvsAddress:                  testutiltx.GenerateAddress(),
 					PubKey:                      anotherPrivateKey.PublicKey().Marshal(),
@@ -291,12 +293,10 @@ func (suite *AVSTestSuite) TestRegisterBLSPublicKey() {
 
 	for _, tc := range testCases {
 		suite.Run(tc.name, func() {
-			_, _, params := tc.setupParams()
+			params := tc.setupParams()
 			err := suite.App.AVSManagerKeeper.RegisterBLSPublicKey(suite.Ctx, params)
 
-			if tc.expectedError != nil {
-				suite.Equal(tc.expectedError, err)
-			} else if tc.errorContains != "" {
+			if tc.errorContains != "" {
 				suite.Error(err)
 				suite.Contains(err.Error(), tc.errorContains)
 			} else {

From 6e25e169d8052e2e7efe6d2e52f981b81f9e6497 Mon Sep 17 00:00:00 2001
From: MaxMustermann2 <82761650+MaxMustermann2@users.noreply.github.com>
Date: Fri, 21 Feb 2025 06:17:27 +0000
Subject: [PATCH 10/13] test(avs): update x/avs tests

---
 x/avs/keeper/avs_test.go | 95 ++++++++++++++++++++++++++++++++++++++++
 x/avs/keeper/keeper.go   | 23 +++++++---
 2 files changed, 111 insertions(+), 7 deletions(-)

diff --git a/x/avs/keeper/avs_test.go b/x/avs/keeper/avs_test.go
index 5b6d09894..b75cf1c88 100644
--- a/x/avs/keeper/avs_test.go
+++ b/x/avs/keeper/avs_test.go
@@ -229,6 +229,61 @@ func (suite *AVSTestSuite) TestRegisterBLSPublicKey() {
 				}
 			},
 		},
+		{
+			name: "reuse BLS key - same operator + avs",
+			setupParams: func() *types.BlsParams {
+				privateKey, err := blst.RandKey()
+				suite.NoError(err)
+				operatorAddress := sdk.AccAddress(utiltx.GenerateAddress().Bytes())
+				msg := fmt.Sprintf(types.BLSMessageToSign, types.ChainIDWithoutRevision(suite.Ctx.ChainID()), operatorAddress.String())
+				hashedMsg := crypto.Keccak256Hash([]byte(msg))
+				sig := privateKey.Sign(hashedMsg.Bytes())
+
+				params := types.BlsParams{
+					OperatorAddress:             operatorAddress,
+					AvsAddress:                  testutiltx.GenerateAddress(),
+					PubKey:                      privateKey.PublicKey().Marshal(),
+					PubKeyRegistrationSignature: sig.Marshal(),
+				}
+				err = suite.App.AVSManagerKeeper.RegisterBLSPublicKey(suite.Ctx, &params)
+				suite.NoError(err)
+				return &params
+			},
+			errorContains: types.ErrAlreadyExists.Error() + "a key has already been set for this operator and avs",
+		},
+		{
+			name: "reuse BLS key - different operator + same avs",
+			setupParams: func() *types.BlsParams {
+				privateKey, err := blst.RandKey()
+				suite.NoError(err)
+				operatorAddress := sdk.AccAddress(utiltx.GenerateAddress().Bytes())
+				msg := fmt.Sprintf(types.BLSMessageToSign, types.ChainIDWithoutRevision(suite.Ctx.ChainID()), operatorAddress.String())
+				hashedMsg := crypto.Keccak256Hash([]byte(msg))
+				sig := privateKey.Sign(hashedMsg.Bytes())
+
+				params := types.BlsParams{
+					OperatorAddress:             operatorAddress,
+					AvsAddress:                  testutiltx.GenerateAddress(),
+					PubKey:                      privateKey.PublicKey().Marshal(),
+					PubKeyRegistrationSignature: sig.Marshal(),
+				}
+				err = suite.App.AVSManagerKeeper.RegisterBLSPublicKey(suite.Ctx, &params)
+				suite.NoError(err)
+
+				anotherOperatorAddress := sdk.AccAddress(utiltx.GenerateAddress().Bytes())
+				msg = fmt.Sprintf(types.BLSMessageToSign, types.ChainIDWithoutRevision(suite.Ctx.ChainID()), anotherOperatorAddress.String())
+				hashedMsg = crypto.Keccak256Hash([]byte(msg))
+				sig = privateKey.Sign(hashedMsg.Bytes())
+
+				return &types.BlsParams{
+					OperatorAddress:             anotherOperatorAddress,
+					AvsAddress:                  params.AvsAddress,
+					PubKey:                      privateKey.PublicKey().Marshal(),
+					PubKeyRegistrationSignature: sig.Marshal(),
+				}
+			},
+			errorContains: types.ErrAlreadyExists.Error() + "this BLS key is already in use",
+		},
 		{
 			name: "wrong chain ID",
 			setupParams: func() *types.BlsParams {
@@ -289,6 +344,46 @@ func (suite *AVSTestSuite) TestRegisterBLSPublicKey() {
 			},
 			errorContains: types.ErrSigNotMatchPubKey.Error(),
 		},
+		{
+			name: "invalid public key format",
+			setupParams: func() *types.BlsParams {
+				operatorAddress := sdk.AccAddress(utiltx.GenerateAddress().Bytes())
+				return &types.BlsParams{
+					OperatorAddress: operatorAddress,
+					AvsAddress:      testutiltx.GenerateAddress(),
+					PubKey:          []byte("invalid"),
+				}
+			},
+			errorContains: types.ErrSigNotMatchPubKey.Error(),
+		},
+		{
+			name: "invalid signature format",
+			setupParams: func() *types.BlsParams {
+				privateKey, err := blst.RandKey()
+				suite.NoError(err)
+				operatorAddress := sdk.AccAddress(utiltx.GenerateAddress().Bytes())
+				return &types.BlsParams{
+					OperatorAddress:             operatorAddress,
+					AvsAddress:                  testutiltx.GenerateAddress(),
+					PubKey:                      privateKey.PublicKey().Marshal(),
+					PubKeyRegistrationSignature: []byte("invalid"),
+				}
+			},
+			errorContains: types.ErrSigNotMatchPubKey.Error(),
+		},
+		{
+			name: "empty operator address",
+			setupParams: func() *types.BlsParams {
+				privateKey, err := blst.RandKey()
+				suite.NoError(err)
+				return &types.BlsParams{
+					AvsAddress:                  testutiltx.GenerateAddress(),
+					PubKey:                      privateKey.PublicKey().Marshal(),
+					PubKeyRegistrationSignature: []byte{},
+				}
+			},
+			errorContains: types.ErrSigNotMatchPubKey.Error(),
+		},
 	}
 
 	for _, tc := range testCases {
diff --git a/x/avs/keeper/keeper.go b/x/avs/keeper/keeper.go
index 8d0eb724c..a4eb4f672 100644
--- a/x/avs/keeper/keeper.go
+++ b/x/avs/keeper/keeper.go
@@ -299,21 +299,30 @@ func (k Keeper) RegisterBLSPublicKey(ctx sdk.Context, params *types.BlsParams) e
 	pubKey, _ := bls.PublicKeyFromBytes(params.PubKey)
 	valid, err := blst.VerifySignature(sig, hashedMsg, pubKey)
 	if err != nil || !valid {
-		return errorsmod.Wrap(types.ErrSigNotMatchPubKey, fmt.Sprintf("the operator is :%s", params.OperatorAddress))
+		return errorsmod.Wrap(types.ErrSigNotMatchPubKey, fmt.Sprintf("the operator is %s", params.OperatorAddress))
 	}
+	// check that a key has not already been set for this operator and avs
 	if k.IsExistPubKeyForAVS(ctx, params.OperatorAddress.String(), params.AvsAddress.String()) {
-		return errorsmod.Wrap(types.ErrAlreadyExists, fmt.Sprintf("the operator is :%s", params.OperatorAddress))
+		return errorsmod.Wrap(
+			types.ErrAlreadyExists,
+			"a key has already been set for this operator and avs",
+		)
 	}
 	blsInfo := &types.BlsPubKeyInfo{
 		AvsAddress:      strings.ToLower(params.AvsAddress.String()),
 		OperatorAddress: params.OperatorAddress.String(),
 		PubKey:          params.PubKey,
 	}
-	// check a bls key can only be used once.
-	// if operator are using multiple servers for different AVSs .
-	// In case one server is compromised, signing can continue as expected on the AVSs for which there has been no compromise.
+	// verify that the BLS key is not already in use by
+	// (1) even the same operator on a different AVS
+	// (2) different operators on the same AVS
+	// this design decision is made to ensure that the same BLS key
+	// is not reused across different AVSs by the same operator
 	if k.IsExistPubKey(ctx, blsInfo) {
-		return errorsmod.Wrap(types.ErrAlreadyExists, fmt.Sprintf("the bls key is already exists:%s", blsInfo.PubKey))
+		return errorsmod.Wrap(
+			types.ErrAlreadyExists,
+			"this BLS key is already in use",
+		)
 	}
 	return k.SetOperatorPubKey(ctx, blsInfo)
 }
@@ -417,7 +426,7 @@ func (k Keeper) GetAVSEpochInfo(ctx sdk.Context, addr string) (*epochstypes.Epoc
 	}
 	avsInfo := avsInfoResp.Info
 	// Epoch information must be available because it is checked when setting AVS information.
-	// Therefore, we don’t need to check it here.
+	// Therefore, we don't need to check it here.
 	epochInfo, _ := k.epochsKeeper.GetEpochInfo(ctx, avsInfo.EpochIdentifier)
 	return &epochInfo, nil
 }

From d9c5be2f6ebdbe7871f36e3b4e5346442af26586 Mon Sep 17 00:00:00 2001
From: MaxMustermann2 <82761650+MaxMustermann2@users.noreply.github.com>
Date: Fri, 21 Feb 2025 06:32:16 +0000
Subject: [PATCH 11/13] fix(avs): check bls key format

---
 precompiles/avs/tx.go    |  1 +
 x/avs/keeper/avs_test.go | 46 ++++++++++++++++++++++++++++------------
 x/avs/keeper/keeper.go   |  5 ++++-
 3 files changed, 38 insertions(+), 14 deletions(-)

diff --git a/precompiles/avs/tx.go b/precompiles/avs/tx.go
index 59f7a0b04..c5b459cc3 100644
--- a/precompiles/avs/tx.go
+++ b/precompiles/avs/tx.go
@@ -316,6 +316,7 @@ func (p Precompile) RegisterBLSPublicKey(
 	}
 	blsParams.PubKeyRegistrationSignature = pubKeyRegistrationSignature
 
+	// validates key format by itself so we don't need to do it before this.
 	err := p.avsKeeper.RegisterBLSPublicKey(ctx, blsParams)
 	if err != nil {
 		return nil, err
diff --git a/x/avs/keeper/avs_test.go b/x/avs/keeper/avs_test.go
index b75cf1c88..edc3fbe69 100644
--- a/x/avs/keeper/avs_test.go
+++ b/x/avs/keeper/avs_test.go
@@ -10,6 +10,7 @@ import (
 	"github.com/ethereum/go-ethereum/crypto"
 	"github.com/prysmaticlabs/prysm/v4/crypto/bls/blst"
 
+	errorsmod "cosmossdk.io/errors"
 	"cosmossdk.io/math"
 	assetstypes "github.com/ExocoreNetwork/exocore/x/assets/types"
 	"github.com/ethereum/go-ethereum/common"
@@ -249,7 +250,10 @@ func (suite *AVSTestSuite) TestRegisterBLSPublicKey() {
 				suite.NoError(err)
 				return &params
 			},
-			errorContains: types.ErrAlreadyExists.Error() + "a key has already been set for this operator and avs",
+			errorContains: errorsmod.Wrap(
+				types.ErrAlreadyExists,
+				"a key has already been set for this operator and avs",
+			).Error(),
 		},
 		{
 			name: "reuse BLS key - different operator + same avs",
@@ -282,7 +286,10 @@ func (suite *AVSTestSuite) TestRegisterBLSPublicKey() {
 					PubKeyRegistrationSignature: sig.Marshal(),
 				}
 			},
-			errorContains: types.ErrAlreadyExists.Error() + "this BLS key is already in use",
+			errorContains: errorsmod.Wrap(
+				types.ErrAlreadyExists,
+				"this BLS key is already in use",
+			).Error(),
 		},
 		{
 			name: "wrong chain ID",
@@ -328,17 +335,19 @@ func (suite *AVSTestSuite) TestRegisterBLSPublicKey() {
 			setupParams: func() *types.BlsParams {
 				privateKey, err := blst.RandKey()
 				suite.NoError(err)
-				anotherPrivateKey, err := blst.RandKey()
-				suite.NoError(err)
 				operatorAddress := sdk.AccAddress(utiltx.GenerateAddress().Bytes())
 				msg := fmt.Sprintf(types.BLSMessageToSign, types.ChainIDWithoutRevision(suite.Ctx.ChainID()), operatorAddress.String())
 				hashedMsg := crypto.Keccak256Hash([]byte(msg))
 				sig := privateKey.Sign(hashedMsg.Bytes())
-
+				// generate a different private key
+				anotherPrivateKey, err := blst.RandKey()
+				suite.NoError(err)
+				anotherPublicKey := anotherPrivateKey.PublicKey()
 				return &types.BlsParams{
-					OperatorAddress:             operatorAddress,
-					AvsAddress:                  testutiltx.GenerateAddress(),
-					PubKey:                      anotherPrivateKey.PublicKey().Marshal(),
+					OperatorAddress: operatorAddress,
+					AvsAddress:      testutiltx.GenerateAddress(),
+					// provide a different public key than the one which signed it, so that verification fails
+					PubKey:                      anotherPublicKey.Marshal(),
 					PubKeyRegistrationSignature: sig.Marshal(),
 				}
 			},
@@ -347,14 +356,20 @@ func (suite *AVSTestSuite) TestRegisterBLSPublicKey() {
 		{
 			name: "invalid public key format",
 			setupParams: func() *types.BlsParams {
+				privateKey, err := blst.RandKey()
+				suite.NoError(err)
 				operatorAddress := sdk.AccAddress(utiltx.GenerateAddress().Bytes())
+				msg := fmt.Sprintf(types.BLSMessageToSign, types.ChainIDWithoutRevision(suite.Ctx.ChainID()), operatorAddress.String())
+				hashedMsg := crypto.Keccak256Hash([]byte(msg))
+				sig := privateKey.Sign(hashedMsg.Bytes())
 				return &types.BlsParams{
-					OperatorAddress: operatorAddress,
-					AvsAddress:      testutiltx.GenerateAddress(),
-					PubKey:          []byte("invalid"),
+					OperatorAddress:             operatorAddress,
+					AvsAddress:                  testutiltx.GenerateAddress(),
+					PubKey:                      []byte("invalid"),
+					PubKeyRegistrationSignature: sig.Marshal(),
 				}
 			},
-			errorContains: types.ErrSigNotMatchPubKey.Error(),
+			errorContains: types.ErrParsePubKey.Error(),
 		},
 		{
 			name: "invalid signature format",
@@ -374,12 +389,17 @@ func (suite *AVSTestSuite) TestRegisterBLSPublicKey() {
 		{
 			name: "empty operator address",
 			setupParams: func() *types.BlsParams {
+				operatorAddress := sdk.AccAddress{}
 				privateKey, err := blst.RandKey()
 				suite.NoError(err)
+				msg := fmt.Sprintf(types.BLSMessageToSign, types.ChainIDWithoutRevision(suite.Ctx.ChainID()), operatorAddress.String())
+				hashedMsg := crypto.Keccak256Hash([]byte(msg))
+				sig := privateKey.Sign(hashedMsg.Bytes())
 				return &types.BlsParams{
+					OperatorAddress:             operatorAddress,
 					AvsAddress:                  testutiltx.GenerateAddress(),
 					PubKey:                      privateKey.PublicKey().Marshal(),
-					PubKeyRegistrationSignature: []byte{},
+					PubKeyRegistrationSignature: sig.Marshal(),
 				}
 			},
 			errorContains: types.ErrSigNotMatchPubKey.Error(),
diff --git a/x/avs/keeper/keeper.go b/x/avs/keeper/keeper.go
index a4eb4f672..2169836d1 100644
--- a/x/avs/keeper/keeper.go
+++ b/x/avs/keeper/keeper.go
@@ -296,7 +296,10 @@ func (k Keeper) RegisterBLSPublicKey(ctx sdk.Context, params *types.BlsParams) e
 	hashedMsg := crypto.Keccak256Hash([]byte(msg))
 
 	sig := params.PubKeyRegistrationSignature
-	pubKey, _ := bls.PublicKeyFromBytes(params.PubKey)
+	pubKey, err := bls.PublicKeyFromBytes(params.PubKey)
+	if err != nil {
+		return errorsmod.Wrap(types.ErrParsePubKey, fmt.Sprintf("the operator is %s", params.OperatorAddress))
+	}
 	valid, err := blst.VerifySignature(sig, hashedMsg, pubKey)
 	if err != nil || !valid {
 		return errorsmod.Wrap(types.ErrSigNotMatchPubKey, fmt.Sprintf("the operator is %s", params.OperatorAddress))

From 271d6d6e6c90bf79565a61980c9bbceccaa11482 Mon Sep 17 00:00:00 2001
From: MaxMustermann2 <82761650+MaxMustermann2@users.noreply.github.com>
Date: Fri, 21 Feb 2025 06:56:50 +0000
Subject: [PATCH 12/13] test(avs): fix test type

---
 x/avs/keeper/avs_test.go | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/x/avs/keeper/avs_test.go b/x/avs/keeper/avs_test.go
index edc3fbe69..2ed3bd855 100644
--- a/x/avs/keeper/avs_test.go
+++ b/x/avs/keeper/avs_test.go
@@ -402,7 +402,7 @@ func (suite *AVSTestSuite) TestRegisterBLSPublicKey() {
 					PubKeyRegistrationSignature: sig.Marshal(),
 				}
 			},
-			errorContains: types.ErrSigNotMatchPubKey.Error(),
+			errorContains: types.ErrInvalidAddr.Error(),
 		},
 	}
 

From 2a46eee9a91273235814ac800ac3796b9321010d Mon Sep 17 00:00:00 2001
From: devin <trestin_blockchain@hotmail.com>
Date: Mon, 24 Feb 2025 10:06:50 +0800
Subject: [PATCH 13/13] fix comments

---
 precompiles/avs/query_test.go |  3 +--
 x/avs/keeper/avs_test.go      | 15 +++++----------
 2 files changed, 6 insertions(+), 12 deletions(-)

diff --git a/precompiles/avs/query_test.go b/precompiles/avs/query_test.go
index a68956b5c..9c1ce6937 100644
--- a/precompiles/avs/query_test.go
+++ b/precompiles/avs/query_test.go
@@ -61,8 +61,7 @@ var baseTestCases = []avsTestCases{
 
 func (suite *AVSManagerPrecompileSuite) TestGetOptedInOperatorAccAddrs() {
 	method := suite.precompile.Methods[avsManagerPrecompile.MethodGetOptInOperators]
-	operatorAddress, avsAddress, slashContract :=
-		sdk.AccAddress(utiltx.GenerateAddress().Bytes()).String(), suite.Address, "0xDF907c29719154eb9872f021d21CAE6E5025d7aB"
+	operatorAddress, avsAddress, slashContract := sdk.AccAddress(utiltx.GenerateAddress().Bytes()).String(), suite.Address, "0xDF907c29719154eb9872f021d21CAE6E5025d7aB"
 
 	operatorOptIn := func() {
 		optedInfo := &types.OptedInfo{
diff --git a/x/avs/keeper/avs_test.go b/x/avs/keeper/avs_test.go
index 39348862d..dd1605a1e 100644
--- a/x/avs/keeper/avs_test.go
+++ b/x/avs/keeper/avs_test.go
@@ -6,8 +6,8 @@ import (
 	"strings"
 	"time"
 
-	testutiltx "github.com/ExocoreNetwork/exocore/testutil/tx"
 	"github.com/ethereum/go-ethereum/crypto"
+	testutiltx "github.com/imua-xyz/imuachain/testutil/tx"
 	"github.com/prysmaticlabs/prysm/v4/crypto/bls/blst"
 
 	errorsmod "cosmossdk.io/errors"
@@ -15,18 +15,14 @@ import (
 	"github.com/ethereum/go-ethereum/common"
 	assetstypes "github.com/imua-xyz/imuachain/x/assets/types"
 
-	"github.com/ExocoreNetwork/exocore/x/avs/types"
-	delegationtypes "github.com/ExocoreNetwork/exocore/x/delegation/types"
-	epochstypes "github.com/ExocoreNetwork/exocore/x/epochs/types"
-	operatorTypes "github.com/ExocoreNetwork/exocore/x/operator/types"
-
-	sdk "github.com/cosmos/cosmos-sdk/types"
-	utiltx "github.com/evmos/evmos/v16/testutil/tx"
 	"github.com/imua-xyz/imuachain/x/avs/types"
-	avstypes "github.com/imua-xyz/imuachain/x/avs/types"
 	delegationtypes "github.com/imua-xyz/imuachain/x/delegation/types"
 	epochstypes "github.com/imua-xyz/imuachain/x/epochs/types"
 	operatorTypes "github.com/imua-xyz/imuachain/x/operator/types"
+
+	sdk "github.com/cosmos/cosmos-sdk/types"
+	utiltx "github.com/evmos/evmos/v16/testutil/tx"
+	avstypes "github.com/imua-xyz/imuachain/x/avs/types"
 )
 
 func (suite *AVSTestSuite) TestAVS() {
@@ -162,7 +158,6 @@ func (suite *AVSTestSuite) TestUpdateAVSInfo_DeRegister() {
 		suite.Equal(epoch.CurrentEpoch, epochEnd+1)
 	}
 
-
 	avsParams.Action = avstypes.DeRegisterAction
 	avsParams.CallerAddress, err = sdk.AccAddressFromBech32(avsOwnerAddress[0])