test(avs): update x/avs tests

imua-xyz · Feb 21, 2025 · 6e25e16 · 6e25e16
1 parent 511788b
commit 6e25e16
Show file tree

Hide file tree

Showing 2 changed files with 111 additions and 7 deletions.
diff --git 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
@@ -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
 }