refactor(batching): tree entries refactor and check previous batch fo…

…r public keys

app/abci/expected_keepers.go

@@ -14,7 +14,8 @@ import (
 
 type BatchingKeeper interface {
 	GetBatchForHeight(ctx context.Context, height int64) (batchingtypes.Batch, error)
-	SetBatchSignatures(ctx context.Context, batchNum uint64, sigs batchingtypes.BatchSignatures) error
+	SetBatchSignatures(ctx context.Context, sigs batchingtypes.BatchSignatures) error
+	GetValidatorTreeEntry(ctx context.Context, batchNum uint64, valAddress sdk.ValAddress) ([]byte, error)
 }
 
 type PubKeyKeeper interface {

app/abci/handlers.go

@@ -13,6 +13,7 @@ import (
 	"cosmossdk.io/collections"
 	addresscodec "cosmossdk.io/core/address"
 	"cosmossdk.io/log"
+
 	"github.com/cosmos/cosmos-sdk/baseapp"
 	sdk "github.com/cosmos/cosmos-sdk/types"
 
@@ -75,7 +76,7 @@ func (h *Handlers) SetSEDASigner(signer utils.SEDASigner) {
 // ExtendVoteHandler handles the ExtendVote ABCI to sign a batch created
 // from the previous block.
 func (h *Handlers) ExtendVoteHandler() sdk.ExtendVoteHandler {
-	return func(ctx sdk.Context, req *abcitypes.RequestExtendVote) (*abcitypes.ResponseExtendVote, error) {
+	return func(ctx sdk.Context, _ *abcitypes.RequestExtendVote) (*abcitypes.ResponseExtendVote, error) {
 		h.logger.Debug("start extend vote handler")
 
 		// Check if there is a batch to sign at this block height.
@@ -140,7 +141,7 @@ func (h *Handlers) VerifyVoteExtensionHandler() sdk.VerifyVoteExtensionHandler {
 			return nil, err
 		}
 
-		err = h.verifyBatchSignatures(ctx, batch.BatchId, req.VoteExtension, req.ValidatorAddress)
+		err = h.verifyBatchSignatures(ctx, batch.BatchNumber, batch.BatchId, req.VoteExtension, req.ValidatorAddress)
 		if err != nil {
 			h.logger.Error("failed to verify batch signature", "req", req, "err", err)
 			return nil, err
@@ -245,7 +246,7 @@ func (h *Handlers) ProcessProposalHandler() sdk.ProcessProposalHandler {
 		for _, vote := range extendedVotes.Votes {
 			// Only consider extensions with pre-commit votes.
 			if vote.BlockIdFlag == cmttypes.BlockIDFlagCommit {
-				err = h.verifyBatchSignatures(ctx, batch.BatchId, vote.VoteExtension, vote.Validator.Address)
+				err = h.verifyBatchSignatures(ctx, batch.BatchNumber, batch.BatchId, vote.VoteExtension, vote.Validator.Address)
 				if err != nil {
 					h.logger.Error("proposal contains an invalid vote extension", "vote", vote)
 					return nil, err
@@ -306,11 +307,12 @@ func (h *Handlers) PreBlocker() sdk.PreBlocker {
 				return nil, err
 			}
 			batchSigs := batchingtypes.BatchSignatures{
+				BatchNumber:   batchNum,
 				ValidatorAddr: validator.OperatorAddress,
 				Signatures:    vote.VoteExtension,
 			}
 
-			err = h.batchingKeeper.SetBatchSignatures(ctx, batchNum, batchSigs)
+			err = h.batchingKeeper.SetBatchSignatures(ctx, batchSigs)
 			if err != nil {
 				return nil, err
 			}
@@ -325,7 +327,7 @@ func (h *Handlers) PreBlocker() sdk.PreBlocker {
 // against the validator's public key registered at the key index
 // in the pubkey module. It returns an error unless the verification
 // succeeds.
-func (h *Handlers) verifyBatchSignatures(ctx sdk.Context, batchID, voteExtension, consAddr []byte) error {
+func (h *Handlers) verifyBatchSignatures(ctx sdk.Context, batchNum uint64, batchID, voteExtension, consAddr []byte) error {
 	if len(voteExtension) == 0 || len(voteExtension) > MaxVoteExtensionLength {
 		h.logger.Error("invalid vote extension length", "len", len(voteExtension))
 		return ErrInvalidVoteExtensionLength
@@ -341,15 +343,46 @@ func (h *Handlers) verifyBatchSignatures(ctx sdk.Context, batchID, voteExtension
 	}
 
 	// Recover and verify secp256k1 public key.
-	pubKey, err := h.pubKeyKeeper.GetValidatorKeyAtIndex(ctx, valOper, utils.SEDAKeyIndexSecp256k1)
+	var expAddr []byte
+	if batchNum == collections.DefaultSequenceStart {
+		pubKey, err := h.pubKeyKeeper.GetValidatorKeyAtIndex(ctx, valOper, utils.SEDAKeyIndexSecp256k1)
+		if err != nil {
+			return err
+		}
+		expAddr, err = utils.PubKeyToEthAddress(pubKey)
+		if err != nil {
+			return err
+		}
+	} else {
+		entry, err := h.batchingKeeper.GetValidatorTreeEntry(ctx, batchNum-1, valOper)
+		if err != nil {
+			if errors.Is(err, collections.ErrNotFound) {
+				pubKey, err := h.pubKeyKeeper.GetValidatorKeyAtIndex(ctx, valOper, utils.SEDAKeyIndexSecp256k1)
+				if err != nil {
+					return err
+				}
+				expAddr, err = utils.PubKeyToEthAddress(pubKey)
+				if err != nil {
+					return err
+				}
+			} else {
+				return err
+			}
+		} else {
+			expAddr = entry[:20]
+		}
+	}
+
+	sigPubKey, err := crypto.Ecrecover(batchID, voteExtension[:65])
 	if err != nil {
 		return err
 	}
-	sigPubKey, err := crypto.Ecrecover(batchID, voteExtension[:65])
+	sigAddr, err := utils.PubKeyToEthAddress(sigPubKey)
 	if err != nil {
 		return err
 	}
-	if !bytes.Equal(pubKey, sigPubKey) {
+
+	if !bytes.Equal(expAddr, sigAddr) {
 		return ErrInvalidBatchSignature
 	}
 	return nil

app/abci/vote_extension_test.go

@@ -2,8 +2,7 @@ package abci
 
 import (
 	"bytes"
-	"crypto/ecdsa"
-	"crypto/elliptic"
+	"path/filepath"
 	"testing"
 
 	"github.com/ethereum/go-ethereum/crypto"
@@ -13,6 +12,8 @@ import (
 	"golang.org/x/crypto/sha3"
 
 	cometabci "github.com/cometbft/cometbft/abci/types"
+	"github.com/cometbft/cometbft/config"
+	"github.com/cometbft/cometbft/privval"
 
 	"cosmossdk.io/log"
 
@@ -26,23 +27,24 @@ import (
 	"github.com/sedaprotocol/seda-chain/app/params"
 	"github.com/sedaprotocol/seda-chain/app/utils"
 	batchingtypes "github.com/sedaprotocol/seda-chain/x/batching/types"
+	pubkeytypes "github.com/sedaprotocol/seda-chain/x/pubkey/types"
 )
 
-var _ utils.SEDASigner = &mockSigner{}
+func TestExtendVerifyVoteHandlers(t *testing.T) {
+	// Set up SEDA signer.
+	tmpDir := t.TempDir()
 
-type mockSigner struct {
-	PrivKey *ecdsa.PrivateKey
-}
+	privValKeyPath := filepath.Join(tmpDir, config.DefaultPrivValKeyName)
+	filePV := privval.GenFilePV(privValKeyPath, "")
+	filePV.Key.Save()
 
-func (m *mockSigner) Sign(input []byte, _ utils.SEDAKeyIndex) ([]byte, error) {
-	signature, err := crypto.Sign(input, m.PrivKey)
-	if err != nil {
-		return nil, err
-	}
-	return signature, nil
-}
+	pubKeys, err := utils.GenerateSEDAKeys(tmpDir)
+	require.NoError(t, err)
+	signer, err := utils.LoadSEDASigner(privValKeyPath)
+	require.NoError(t, err)
+
+	secp256k1PubKey := pubKeys[utils.SEDAKeyIndexSecp256k1].PubKey
 
-func TestExtendVerifyVoteHandlers(t *testing.T) {
 	// Configure address prefixes.
 	cfg := sdk.GetConfig()
 	cfg.SetBech32PrefixForAccount(params.Bech32PrefixAccAddr, params.Bech32PrefixAccPub)
@@ -69,13 +71,13 @@ func TestExtendVerifyVoteHandlers(t *testing.T) {
 		BlockHeight: 100, // created from the previous block
 	}
 
-	privKey, err := crypto.HexToECDSA("79afbf7147841fca72b45a1978dd7669470ba67abbe5c220062924380c9c364b")
-	require.NoError(t, err)
-	pubKey := elliptic.Marshal(privKey.PublicKey, privKey.PublicKey.X, privKey.PublicKey.Y)
-
-	signer := mockSigner{privKey}
-
 	mockBatchingKeeper.EXPECT().GetBatchForHeight(gomock.Any(), int64(100)).Return(mockBatch, nil).AnyTimes()
+	mockStakingKeeper.EXPECT().GetValidatorByConsAddr(gomock.Any(), gomock.Any()).Return(
+		stakingtypes.Validator{
+			OperatorAddress: "sedavaloper1ucv5709wlf9jn84ynyjzyzeavwvurmdydtn3px",
+		}, nil,
+	)
+	mockPubKeyKeeper.EXPECT().GetValidatorKeys(gomock.Any(), gomock.Any()).Return(pubkeytypes.ValidatorPubKeys{}, nil)
 
 	// Construct the handler and execute it.
 	handler := NewHandlers(
@@ -86,7 +88,7 @@ func TestExtendVerifyVoteHandlers(t *testing.T) {
 		authcodec.NewBech32Codec(sdk.GetConfig().GetBech32ValidatorAddrPrefix()),
 		logger,
 	)
-	handler.SetSEDASigner(&signer)
+	handler.SetSEDASigner(signer)
 	extendVoteHandler := handler.ExtendVoteHandler()
 	verifyVoteHandler := handler.VerifyVoteExtensionHandler()
 
@@ -98,7 +100,7 @@ func TestExtendVerifyVoteHandlers(t *testing.T) {
 	// Recover and verify public key
 	sigPubKey, err := crypto.Ecrecover(mockBatch.BatchId, evRes.VoteExtension)
 	require.NoError(t, err)
-	require.Equal(t, pubKey, sigPubKey)
+	require.Equal(t, secp256k1PubKey, sigPubKey)
 
 	testVal := sdk.ConsAddress([]byte("testval"))
 	mockStakingKeeper.EXPECT().GetValidatorByConsAddr(gomock.Any(), testVal).Return(
@@ -110,7 +112,7 @@ func TestExtendVerifyVoteHandlers(t *testing.T) {
 		gomock.Any(),
 		[]byte{230, 25, 79, 60, 174, 250, 75, 41, 158, 164, 153, 36, 34, 11, 61, 99, 153, 193, 237, 164},
 		utils.SEDAKeyIndexSecp256k1,
-	).Return(pubKey, nil)
+	).Return(secp256k1PubKey, nil)
 
 	vvRes, err := verifyVoteHandler(ctx, &cometabci.RequestVerifyVoteExtension{
 		Height:           101,

app/app.go

@@ -1002,12 +1002,10 @@ func NewApp(
 	app.SetPreBlocker(app.PreBlocker)
 
 	// Register SEDA signer and ExtendVote handler.
-	fmt.Println(cast.ToString(appOpts.Get("priv_validator_key_file")))
 	pvKeyFile := filepath.Join(homePath, cast.ToString(appOpts.Get("priv_validator_key_file")))
-	// loadPath := filepath.Join(filepath.Dir(pvKeyFile), utils.SEDAKeyFileName)
 	signer, err := utils.LoadSEDASigner(pvKeyFile)
 	if err != nil {
-		// app.Logger().Error("error loading SEDA signer - ExtendVote handler will not run", "path", loadPath) // TODO
+		app.Logger().Error("error loading SEDA signer - ExtendVote handler will not run", "err", err)
 	} else {
 		app.Logger().Info("SEDA signer successfully loaded")
 		abciHandler.SetSEDASigner(signer)

app/utils/seda_signer.go

@@ -35,8 +35,6 @@ type sedaKeys struct {
 // LoadSEDASigner loads the SEDA keys from a given file and returns
 // a SEDASigner interface.
 func LoadSEDASigner(pvKeyFilePath string) (SEDASigner, error) {
-	// TODO What if there is a rotation?
-	// TODO Can we safely assume that the file will be loaded?
 	_, err := os.ReadFile(pvKeyFilePath)
 	if err != nil {
 		return nil, fmt.Errorf("failed to read private validator key from %v: %v", pvKeyFilePath, err)
@@ -120,7 +118,6 @@ func (s *sedaKeys) ReloadIfMismatch(pubKeys []pubkeytypes.IndexedPubKey) error {
 
 // Reload reloads the signer from the key file.
 func (s *sedaKeys) reload() error {
-	// TODO merge with LoadSEDASigner??
 	keysJSONBytes, err := os.ReadFile(s.keyPath)
 	if err != nil {
 		return fmt.Errorf("failed to read SEDA keys from %v: %v", s.keyPath, err)

proto/sedachain/batching/v1/batching.proto

@@ -32,7 +32,7 @@ message Batch {
 // TreeEntries are the given batch's data result tree entries and
 // validator tree entries.
 message TreeEntries {
-  // batch_number is the identifier of the batch the tree entries from.
+  // batch_number is the identifier of the batch.
   uint64 batch_number = 1;
   // data_result_entries are the entries (unhashed leaf contents) of
   // the data result tree.
@@ -45,9 +45,11 @@ message TreeEntries {
 // BatchSignatures contains basic validator data and its batch signatures
 // under various cryptographic schemes.
 message BatchSignatures {
-  string validator_addr = 1
+  // batch_number is the identifier of the batch.
+  uint64 batch_number = 1;
+  string validator_addr = 2
       [ (cosmos_proto.scalar) = "cosmos.ValidatorAddressString" ];
-  bytes signatures = 2;
+  bytes signatures = 3;
 }
 
 // Params is a list of parameters which can be changed through governance.

proto/sedachain/batching/v1/genesis.proto

@@ -12,11 +12,13 @@ message GenesisState {
   // created batch.
   uint64 current_batch_number = 1;
   repeated Batch batches = 2 [ (gogoproto.nullable) = false ];
-  repeated TreeEntries tree_entries = 3 [ (gogoproto.nullable) = false ];
+  repeated TreeEntry tree_entries = 3 [ (gogoproto.nullable) = false ];
   repeated DataResult data_results = 4 [ (gogoproto.nullable) = false ];
   repeated BatchAssignment batch_assignments = 5
       [ (gogoproto.nullable) = false ];
-  Params params = 6 [ (gogoproto.nullable) = false ];
+  repeated BatchSignatures batch_signatures = 6
+      [ (gogoproto.nullable) = false ];
+  Params params = 7 [ (gogoproto.nullable) = false ];
 }
 
 // BatchAssignment represents a batch assignment for genesis export
@@ -25,3 +27,13 @@ message BatchAssignment {
   uint64 batch_number = 1;
   string data_request_id = 2;
 }
+
+// TreeEntry represents a tree entry for genesis export and import.
+message TreeEntry {
+  // K1 is the first part of the key. It represents the batch number.
+  uint64 k1 = 1;
+  // K2 is the second part of the key. It is empty for a data result
+  // tree entry and validator address for a validator tree entry.
+  bytes k2 = 2;
+  bytes entry = 3;
+}

scripts/local_multi_setup.sh

@@ -112,9 +112,6 @@ NODE1_ID=$($BIN tendermint show-node-id --home=$HOME/.sedad/validator1 | tail -1
 sed -i '' -E "s|persistent_peers = \"\"|persistent_peers = \"${NODE1_ID}@localhost:26656\"|g" $HOME/.sedad/validator2/config/config.toml
 sed -i '' -E "s|persistent_peers = \"\"|persistent_peers = \"${NODE1_ID}@localhost:26656\"|g" $HOME/.sedad/validator3/config/config.toml
 sed -i '' -E "s|persistent_peers = \"\"|persistent_peers = \"${NODE1_ID}@localhost:26656\"|g" $HOME/.sedad/validator4/config/config.toml
-# sed -i '' -E "s|persistent_peers = \"\"|persistent_peers = \"$($BIN tendermint show-node-id --home=$HOME/.sedad/validator1)@localhost:26656\"|g" $HOME/.sedad/validator2/config/config.toml
-# sed -i '' -E "s|persistent_peers = \"\"|persistent_peers = \"$($BIN tendermint show-node-id --home=$HOME/.sedad/validator1)@localhost:26656\"|g" $HOME/.sedad/validator3/config/config.toml
-# sed -i '' -E "s|persistent_peers = \"\"|persistent_peers = \"$($BIN tendermint show-node-id --home=$HOME/.sedad/validator1)@localhost:26656\"|g" $HOME/.sedad/validator4/config/config.toml
 
 # start all four validators
 tmux new-session -s validator1 -d