Add prefixes to merkle leaves to guard against a tree shrinking attack

pkg/eigenState/avsOperators/avsOperators.go

@@ -229,7 +229,7 @@ func (a *AvsOperatorsModel) GenerateStateRoot(blockNumber uint64) (types.StateRo
 		return "", nil
 	}
 
-	fullTree, err := a.MerkleizeState(blockNumber, inputs)
+	fullTree, err := a.MerkleizeEigenState(blockNumber, inputs)
 	if err != nil {
 		a.logger.Sugar().Errorw("Failed to create merkle tree",
 			zap.Error(err),

pkg/eigenState/base/baseEigenState.go

@@ -56,8 +56,10 @@ func (b *BaseEigenState) ParseLogOutput(log *storage.TransactionLog) (map[string
 // 1. Ensures that the tree is always different for different blocks
 // 2. Allows us to have at least 1 value if there are no model changes for a block.
 func (b *BaseEigenState) InitializeMerkleTreeBaseStateWithBlock(blockNumber uint64) [][]byte {
+	blockNumberByte := []byte(fmt.Sprintf("%d", blockNumber))
 	return [][]byte{
-		[]byte(fmt.Sprintf("%d", blockNumber)),
+		types.MerkleLeafPrefix_EigenStateBlock,
+		blockNumberByte,
 	}
 }
 
@@ -104,10 +106,10 @@ type MerkleTreeInput struct {
 	Value  []byte
 }
 
-// MerkleizeState creates a merkle tree from the given inputs.
+// MerkleizeEigenState creates a merkle tree from the given inputs.
 //
 // Each input includes a SlotID and a byte representation of the state that changed
-func (b *BaseEigenState) MerkleizeState(blockNumber uint64, inputs []*MerkleTreeInput) (*merkletree.MerkleTree, error) {
+func (b *BaseEigenState) MerkleizeEigenState(blockNumber uint64, inputs []*MerkleTreeInput) (*merkletree.MerkleTree, error) {
 	om := orderedmap.New[types.SlotID, []byte]()
 
 	for _, input := range inputs {
@@ -136,7 +138,7 @@ func (b *BaseEigenState) MerkleizeState(blockNumber uint64, inputs []*MerkleTree
 }
 
 func encodeMerkleLeaf(slotID types.SlotID, value []byte) []byte {
-	return append([]byte(slotID), value...)
+	return append(types.MerkleLeafPrefix_EigenStateChange, append([]byte(slotID), value...)...)
 }
 
 func NewSlotID(txHash string, logIndex uint64) types.SlotID {

pkg/eigenState/disabledDistributionRoots/disabledDistributionRoots.go

@@ -208,7 +208,7 @@ func (ddr *DisabledDistributionRootsModel) GenerateStateRoot(blockNumber uint64)
 		return "", nil
 	}
 
-	fullTree, err := ddr.MerkleizeState(blockNumber, sortedInputs)
+	fullTree, err := ddr.MerkleizeEigenState(blockNumber, sortedInputs)
 	if err != nil {
 		ddr.logger.Sugar().Errorw("Failed to create merkle tree",
 			zap.Error(err),

pkg/eigenState/operatorShares/operatorShares.go

@@ -286,7 +286,7 @@ func (osm *OperatorSharesModel) GenerateStateRoot(blockNumber uint64) (types.Sta
 		return "", nil
 	}
 
-	fullTree, err := osm.MerkleizeState(blockNumber, inputs)
+	fullTree, err := osm.MerkleizeEigenState(blockNumber, inputs)
 	if err != nil {
 		osm.logger.Sugar().Errorw("Failed to create merkle tree",
 			zap.Error(err),

pkg/eigenState/rewardSubmissions/rewardSubmissions.go

@@ -310,7 +310,7 @@ func (rs *RewardSubmissionsModel) GenerateStateRoot(blockNumber uint64) (types.S
 		return "", nil
 	}
 
-	fullTree, err := rs.MerkleizeState(blockNumber, inputs)
+	fullTree, err := rs.MerkleizeEigenState(blockNumber, inputs)
 	if err != nil {
 		rs.logger.Sugar().Errorw("Failed to create merkle tree",
 			zap.Error(err),

pkg/eigenState/stakerDelegations/stakerDelegations.go

@@ -215,7 +215,7 @@ func (s *StakerDelegationsModel) GenerateStateRoot(blockNumber uint64) (types.St
 		return "", nil
 	}
 
-	fullTree, err := s.MerkleizeState(blockNumber, inputs)
+	fullTree, err := s.MerkleizeEigenState(blockNumber, inputs)
 	if err != nil {
 		s.logger.Sugar().Errorw("Failed to create merkle tree",
 			zap.Error(err),

pkg/eigenState/stakerShares/stakerShares.go

@@ -585,7 +585,7 @@ func (ss *StakerSharesModel) GenerateStateRoot(blockNumber uint64) (types.StateR
 		return "", nil
 	}
 
-	fullTree, err := ss.MerkleizeState(blockNumber, inputs)
+	fullTree, err := ss.MerkleizeEigenState(blockNumber, inputs)
 	if err != nil {
 		ss.logger.Sugar().Errorw("Failed to create merkle tree",
 			zap.Error(err),

pkg/eigenState/stateManager/stateManager.go

@@ -103,6 +103,7 @@ func (e *EigenStateManager) CleanupProcessedStateForBlock(blockNumber uint64) er
 func (e *EigenStateManager) GenerateStateRoot(blockNumber uint64, blockHash string) (types.StateRoot, error) {
 	sortedIndexes := e.GetSortedModelIndexes()
 	roots := [][]byte{
+		types.MerkleLeafPrefix_Block,
 		[]byte(fmt.Sprintf("%d", blockNumber)),
 		[]byte(blockHash),
 	}
@@ -168,7 +169,7 @@ func (e *EigenStateManager) encodeModelLeaf(model types.IEigenStateModel, blockN
 	if root == "" {
 		return nil, nil
 	}
-	return append([]byte(model.GetModelName()), []byte(root)...), nil
+	return append(types.MerkleLeafPrefix_EigenStateRoot, append([]byte(model.GetModelName()), []byte(root)...)...), nil
 }
 
 func (e *EigenStateManager) GetSortedModelIndexes() []int {

pkg/eigenState/submittedDistributionRoots/submittedDistributionRoots.go

@@ -275,7 +275,7 @@ func (sdr *SubmittedDistributionRootsModel) GenerateStateRoot(blockNumber uint64
 		return "", nil
 	}
 
-	fullTree, err := sdr.MerkleizeState(blockNumber, sortedInputs)
+	fullTree, err := sdr.MerkleizeEigenState(blockNumber, sortedInputs)
 	if err != nil {
 		sdr.logger.Sugar().Errorw("Failed to create merkle tree",
 			zap.Error(err),

pkg/eigenState/types/types.go

@@ -50,3 +50,12 @@ type IEigenStateModel interface {
 type StateTransitions[T any] map[uint64]func(log *storage.TransactionLog) (T, error)
 
 type SlotID string
+
+type MerkleLeafPrefix []byte
+
+var (
+	MerkleLeafPrefix_Block            MerkleLeafPrefix = []byte("0x00")
+	MerkleLeafPrefix_EigenStateRoot   MerkleLeafPrefix = []byte("0x01")
+	MerkleLeafPrefix_EigenStateBlock  MerkleLeafPrefix = []byte("0x02")
+	MerkleLeafPrefix_EigenStateChange MerkleLeafPrefix = []byte("0x03")
+)