WIP: tree entries refactor

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

@@ -140,7 +140,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 +245,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 +306,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 +326,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 +342,45 @@ 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, 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
+			}
+		}
+		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

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,12 @@ 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 +26,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;
+}

x/batching/keeper/batch.go

@@ -56,7 +56,7 @@ func (k Keeper) SetNewBatch(ctx context.Context, batch types.Batch, dataEntries
 	}
 	batch.BatchNumber = newBatchNum
 
-	err = k.setTreeEntries(ctx, newBatchNum, dataEntries, valEntries)
+	err = k.setTreeEntriesForBatch(ctx, newBatchNum, dataEntries, valEntries)
 	if err != nil {
 		return err
 	}
@@ -132,9 +132,14 @@ type validatorTreeEntry struct {
 	entry      []byte
 }
 
-// setTreeEntries stores the data result entries and validator entries
-// using the given batch number as the key.
-func (k Keeper) setTreeEntries(ctx context.Context, batchNum uint64, dataEntries [][]byte, valEntries []validatorTreeEntry) error {
+// setTreeEntry stores a tree entry.
+func (k Keeper) setTreeEntry(ctx context.Context, entry types.TreeEntry) error {
+	return k.treeEntries.Set(ctx, collections.Join(entry.K1, entry.K2), entry.Entry)
+}
+
+// setTreeEntriesForBatch stores the data result entries and validator
+// entries for a given batch.
+func (k Keeper) setTreeEntriesForBatch(ctx context.Context, batchNum uint64, dataEntries [][]byte, valEntries []validatorTreeEntry) error {
 	for _, entry := range dataEntries {
 		err := k.treeEntries.Set(ctx, collections.Join(batchNum, []byte{}), entry)
 		if err != nil {
@@ -150,13 +155,15 @@ func (k Keeper) setTreeEntries(ctx context.Context, batchNum uint64, dataEntries
 	return nil
 }
 
+// GetValidatorTreeEntry returns the tree entry of a given validator
+// for a specified batch
 func (k Keeper) GetValidatorTreeEntry(ctx context.Context, batchNum uint64, valAddress sdk.ValAddress) ([]byte, error) {
 	return k.treeEntries.Get(ctx, collections.Join(batchNum, valAddress.Bytes()))
 }
 
-// GetTreeEntries returns the tree entries corresponding to the given
-// batch number.
-func (k Keeper) GetTreeEntries(ctx context.Context, batchNum uint64) (types.TreeEntries, error) {
+// GetTreeEntriesForBatch returns the tree entries corresponding to
+// the givenbatch number.
+func (k Keeper) GetTreeEntriesForBatch(ctx context.Context, batchNum uint64) (types.TreeEntries, error) {
 	rng := collections.NewPrefixedPairRange[uint64, []byte](batchNum)
 	itr, err := k.treeEntries.Iterate(ctx, rng)
 	if err != nil {
@@ -189,49 +196,34 @@ func (k Keeper) GetTreeEntries(ctx context.Context, batchNum uint64) (types.Tree
 
 // IterateBatches iterates over the tree entries and performs a given
 // callback function.
-func (k Keeper) IterateTreeEntries(ctx sdk.Context, callback func(types.TreeEntries) (stop bool)) error {
+func (k Keeper) IterateTreeEntries(ctx sdk.Context, callback func(types.TreeEntry) (stop bool)) error {
 	iter, err := k.treeEntries.Iterate(ctx, nil)
 	if err != nil {
 		return err
 	}
 	defer iter.Close()
 
-	keys, err := iter.Keys()
-	if err != nil {
-		return err
-	}
-
-	var dataEntries, valEntries [][]byte
-	for i, key := range keys {
-		value, err := k.treeEntries.Get(ctx, key)
+	for ; iter.Valid(); iter.Next() {
+		kv, err := iter.KeyValue()
 		if err != nil {
 			return err
 		}
-		if len(key.K2()) == 0 {
-			dataEntries = append(dataEntries, value)
-		} else {
-			valEntries = append(valEntries, value)
-		}
 
-		if i == len(keys)-1 || key.K1() != keys[i+1].K1() {
-			// Execute callback and initialize the entries.
-			if callback(types.TreeEntries{
-				BatchNumber:       key.K1(),
-				DataResultEntries: dataEntries,
-				ValidatorEntries:  valEntries,
-			}) {
-				break
-			}
-			dataEntries, valEntries = [][]byte{}, [][]byte{}
+		if callback(types.TreeEntry{
+			K1:    kv.Key.K1(),
+			K2:    kv.Key.K2(),
+			Entry: kv.Value,
+		}) {
+			break
 		}
 	}
 	return nil
 }
 
 // GetAllTreeEntries retrieves all tree entries from the store.
-func (k Keeper) GetAllTreeEntries(ctx sdk.Context) ([]types.TreeEntries, error) {
-	var entries []types.TreeEntries
-	err := k.IterateTreeEntries(ctx, func(entry types.TreeEntries) bool {
+func (k Keeper) GetAllTreeEntries(ctx sdk.Context) ([]types.TreeEntry, error) {
+	var entries []types.TreeEntry
+	err := k.IterateTreeEntries(ctx, func(entry types.TreeEntry) bool {
 		entries = append(entries, entry)
 		return false
 	})
@@ -242,12 +234,12 @@ func (k Keeper) GetAllTreeEntries(ctx sdk.Context) ([]types.TreeEntries, error)
 }
 
 // SetBatchSignatures stores a validator's signatures of a batch.
-func (k Keeper) SetBatchSignatures(ctx context.Context, batchNum uint64, sigs types.BatchSignatures) error {
+func (k Keeper) SetBatchSignatures(ctx context.Context, sigs types.BatchSignatures) error {
 	valAddr, err := k.validatorAddressCodec.StringToBytes(sigs.ValidatorAddr)
 	if err != nil {
 		return err
 	}
-	return k.batchSignatures.Set(ctx, collections.Join(batchNum, valAddr), sigs)
+	return k.batchSignatures.Set(ctx, collections.Join(sigs.BatchNumber, valAddr), sigs)
 }
 
 // GetBatchSignatures retrieves the batch signatures by a given
@@ -287,3 +279,26 @@ func (k Keeper) GetBatchSigsForBatch(ctx context.Context, batchNum uint64) ([]ty
 	}
 	return sigs, err
 }
+
+// GetAllBatchSignatures returns all batch signatures in the store.
+func (k Keeper) GetAllBatchSignatures(ctx context.Context) ([]types.BatchSignatures, error) {
+	itr, err := k.batchSignatures.Iterate(ctx, nil)
+	if err != nil {
+		return nil, err
+	}
+	defer itr.Close()
+
+	kvs, err := itr.KeyValues()
+	if err != nil {
+		return nil, err
+	}
+	if len(kvs) == 0 {
+		return nil, collections.ErrNotFound
+	}
+
+	sigs := make([]types.BatchSignatures, len(kvs))
+	for i, kv := range kvs {
+		sigs[i] = kv.Value
+	}
+	return sigs, err
+}