Ensure leafs are 32-bytes (#51)

Author: fasmat

go.mod

@@ -5,14 +5,15 @@ go 1.23.2
 require (
 	github.com/minio/sha256-simd v1.0.1
 	github.com/stretchr/testify v1.9.0
+	golang.org/x/exp v0.0.0-20241009180824-f66d83c29e7c
 )
 
 require (
 	github.com/davecgh/go-spew v1.1.1 // indirect
-	github.com/klauspost/cpuid/v2 v2.2.3 // indirect
+	github.com/klauspost/cpuid/v2 v2.2.8 // indirect
 	github.com/kr/pretty v0.3.1 // indirect
 	github.com/pmezard/go-difflib v1.0.0 // indirect
-	golang.org/x/sys v0.25.0 // indirect
+	golang.org/x/sys v0.26.0 // indirect
 	gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c // indirect
 	gopkg.in/yaml.v3 v3.0.1 // indirect
 )

go.sum

@@ -1,8 +1,8 @@
 github.com/creack/pty v1.1.9/go.mod h1:oKZEueFk5CKHvIhNR5MUki03XCEU+Q6VDXinZuGJ33E=
 github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c=
 github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
-github.com/klauspost/cpuid/v2 v2.2.3 h1:sxCkb+qR91z4vsqw4vGGZlDgPz3G7gjaLyK3V8y70BU=
-github.com/klauspost/cpuid/v2 v2.2.3/go.mod h1:RVVoqg1df56z8g3pUjL/3lE5UfnlrJX8tyFgg4nqhuY=
+github.com/klauspost/cpuid/v2 v2.2.8 h1:+StwCXwm9PdpiEkPyzBXIy+M9KUb4ODm0Zarf1kS5BM=
+github.com/klauspost/cpuid/v2 v2.2.8/go.mod h1:Lcz8mBdAVJIBVzewtcLocK12l3Y+JytZYpaMropDUws=
 github.com/kr/pretty v0.2.1/go.mod h1:ipq/a2n7PKx3OHsz4KJII5eveXtPO4qwEXGdVfWzfnI=
 github.com/kr/pretty v0.3.1 h1:flRD4NNwYAUpkphVc1HcthR4KEIFJ65n8Mw5qdRn3LE=
 github.com/kr/pretty v0.3.1/go.mod h1:hoEshYVHaxMs3cyo3Yncou5ZscifuDolrwPKZanG3xk=
@@ -19,9 +19,11 @@ github.com/rogpeppe/go-internal v1.9.0 h1:73kH8U+JUqXU8lRuOHeVHaa/SZPifC7BkcraZV
 github.com/rogpeppe/go-internal v1.9.0/go.mod h1:WtVeX8xhTBvf0smdhujwtBcq4Qrzq/fJaraNFVN+nFs=
 github.com/stretchr/testify v1.9.0 h1:HtqpIVDClZ4nwg75+f6Lvsy/wHu+3BoSGCbBAcpTsTg=
 github.com/stretchr/testify v1.9.0/go.mod h1:r2ic/lqez/lEtzL7wO/rwa5dbSLXVDPFyf8C91i36aY=
-golang.org/x/sys v0.0.0-20220704084225-05e143d24a9e/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
-golang.org/x/sys v0.25.0 h1:r+8e+loiHxRqhXVl6ML1nO3l1+oFoWbnlu2Ehimmi34=
-golang.org/x/sys v0.25.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA=
+golang.org/x/exp v0.0.0-20241009180824-f66d83c29e7c h1:7dEasQXItcW1xKJ2+gg5VOiBnqWrJc+rq0DPKyvvdbY=
+golang.org/x/exp v0.0.0-20241009180824-f66d83c29e7c/go.mod h1:NQtJDoLvd6faHhE7m4T/1IY708gDefGGjR/iUW8yQQ8=
+golang.org/x/sys v0.5.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
+golang.org/x/sys v0.26.0 h1:KHjCJyddX0LoSTb3J+vWpupP9p0oznkqVk/IfjymZbo=
+golang.org/x/sys v0.26.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA=
 gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
 gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c h1:Hei/4ADfdWqJk1ZMxUNpqntNwaWcugrBjAiHlqqRiVk=
 gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c/go.mod h1:JHkPIbrfpd72SG/EVd6muEfDQjcINNoR0C8j2r3qZ4Q=

iterators.go

@@ -2,26 +2,23 @@ package merkle
 
 import (
 	"errors"
-	"sort"
+	"slices"
+
+	"golang.org/x/exp/maps"
 )
 
 var noMoreItems = errors.New("no more items")
 
 type Set map[uint64]bool
 
 func (s Set) AsSortedSlice() []uint64 {
-	var ret []uint64
-	for key, value := range s {
-		if value {
-			ret = append(ret, key)
-		}
-	}
-	sort.Slice(ret, func(i, j int) bool { return ret[i] < ret[j] })
+	ret := maps.Keys(s)
+	slices.Sort(ret)
 	return ret
 }
 
 func SetOf(members ...uint64) Set {
-	ret := make(Set)
+	ret := make(Set, len(members))
 	for _, member := range members {
 		ret[member] = true
 	}

merkle.go

@@ -104,6 +104,9 @@ type Tree struct {
 // AddLeaf incorporates a new leaf to the state of the tree. It updates the state required to eventually determine the
 // root of the tree and also updates the proof, if applicable.
 func (t *Tree) AddLeaf(value []byte) error {
+	if len(value) != NodeSize {
+		return fmt.Errorf("expected node size %d, got %d", NodeSize, len(value))
+	}
 	n := node{
 		value:        value,
 		OnProvenPath: t.leavesToProve.Pop(),

merkle_test.go

@@ -466,49 +466,64 @@ func TestTree_GetParkedNodes(t *testing.T) {
 	tree, err := NewTreeBuilder().Build()
 	r.NoError(err)
 
-	r.NoError(tree.AddLeaf([]byte{0}))
-	r.EqualValues(
-		[][]byte{{0}},
-		tree.GetParkedNodes(nil))
+	leaf0 := [32]byte{}
+	r.NoError(tree.AddLeaf(leaf0[:]))
+	r.EqualValues([][]byte{leaf0[:]}, tree.GetParkedNodes(nil))
 
-	r.NoError(tree.AddLeaf([]byte{1}))
+	leaf1 := [32]byte{}
+	leaf1[0] = 1
+	r.NoError(tree.AddLeaf(leaf1[:]))
 	r.EqualValues(
-		[][]byte{{}, decode(r, "b413f47d13ee2fe6c845b2ee141af81de858df4ec549a58b7970bb96645bc8d2")},
-		tree.GetParkedNodes(nil))
+		[][]byte{{}, decode(r, "cb592844121d926f1ca3ad4e1d6fb9d8e260ed6e3216361f7732e975a0e8bbf6")},
+		tree.GetParkedNodes(nil),
+	)
 
-	r.NoError(tree.AddLeaf([]byte{2}))
+	leaf2 := [32]byte{}
+	leaf2[0] = 2
+	r.NoError(tree.AddLeaf(leaf2[:]))
 	r.EqualValues(
-		[][]byte{{2}, decode(r, "b413f47d13ee2fe6c845b2ee141af81de858df4ec549a58b7970bb96645bc8d2")},
-		tree.GetParkedNodes(nil))
+		[][]byte{leaf2[:], decode(r, "cb592844121d926f1ca3ad4e1d6fb9d8e260ed6e3216361f7732e975a0e8bbf6")},
+		tree.GetParkedNodes(nil),
+	)
 
-	r.NoError(tree.AddLeaf([]byte{3}))
+	leaf3 := [32]byte{}
+	leaf3[0] = 3
+	r.NoError(tree.AddLeaf(leaf3[:]))
 	r.EqualValues(
-		[][]byte{{}, {}, decode(r, "7699a4fdd6b8b6908a344f73b8f05c8e1400f7253f544602c442ff5c65504b24")},
-		tree.GetParkedNodes(nil))
+		[][]byte{{}, {}, decode(r, "ba94ffe7edabf26ef12736f8eb5ce74d15bedb6af61444ae2906e926b1a95084")},
+		tree.GetParkedNodes(nil),
+	)
 }
 
 func TestTree_SetParkedNodes(t *testing.T) {
 	r := require.New(t)
 
+	node0 := [32]byte{}
+	leaf1 := [32]byte{}
+	leaf1[0] = 1
 	tree, err := NewTreeBuilder().Build()
 	r.NoError(err)
-	r.NoError(tree.SetParkedNodes([][]byte{{0}}))
-	r.NoError(tree.AddLeaf([]byte{1}))
-	parkedNodes := [][]byte{{}, decode(r, "b413f47d13ee2fe6c845b2ee141af81de858df4ec549a58b7970bb96645bc8d2")}
+	r.NoError(tree.SetParkedNodes([][]byte{node0[:]}))
+	r.NoError(tree.AddLeaf(leaf1[:]))
+	parkedNodes := [][]byte{{}, decode(r, "cb592844121d926f1ca3ad4e1d6fb9d8e260ed6e3216361f7732e975a0e8bbf6")}
 	r.EqualValues(parkedNodes, tree.GetParkedNodes(nil))
 
+	leaf2 := [32]byte{}
+	leaf2[0] = 2
 	tree, err = NewTreeBuilder().Build()
 	r.NoError(err)
 	r.NoError(tree.SetParkedNodes(parkedNodes))
-	r.NoError(tree.AddLeaf([]byte{2}))
-	parkedNodes = [][]byte{{2}, decode(r, "b413f47d13ee2fe6c845b2ee141af81de858df4ec549a58b7970bb96645bc8d2")}
+	r.NoError(tree.AddLeaf(leaf2[:]))
+	parkedNodes = [][]byte{leaf2[:], decode(r, "cb592844121d926f1ca3ad4e1d6fb9d8e260ed6e3216361f7732e975a0e8bbf6")}
 	r.EqualValues(parkedNodes, tree.GetParkedNodes(nil))
 
+	leaf3 := [32]byte{}
+	leaf3[0] = 3
 	tree, err = NewTreeBuilder().Build()
 	r.NoError(err)
 	r.NoError(tree.SetParkedNodes(parkedNodes))
-	r.NoError(tree.AddLeaf([]byte{3}))
-	parkedNodes = [][]byte{{}, {}, decode(r, "7699a4fdd6b8b6908a344f73b8f05c8e1400f7253f544602c442ff5c65504b24")}
+	r.NoError(tree.AddLeaf(leaf3[:]))
+	parkedNodes = [][]byte{{}, {}, decode(r, "ba94ffe7edabf26ef12736f8eb5ce74d15bedb6af61444ae2906e926b1a95084")}
 	r.EqualValues(parkedNodes, tree.GetParkedNodes(nil))
 }
 

validation.go

@@ -4,14 +4,18 @@ import (
 	"bytes"
 	"errors"
 	"fmt"
-	"sort"
+
+	"golang.org/x/exp/slices"
 )
 
 const MaxUint = ^uint(0)
 
 // ValidatePartialTree uses leafIndices, leaves and proof to calculate the merkle root of the tree and then compares it
 // to expectedRoot.
-func ValidatePartialTree(leafIndices []uint64, leaves, proof [][]byte, expectedRoot []byte,
+func ValidatePartialTree(
+	leafIndices []uint64,
+	leaves, proof [][]byte,
+	expectedRoot []byte,
 	hash HashFunc,
 ) (bool, error) {
 	v, err := newValidator(leafIndices, leaves, proof, hash, false)
@@ -25,7 +29,10 @@ func ValidatePartialTree(leafIndices []uint64, leaves, proof [][]byte, expectedR
 // ValidatePartialTree uses leafIndices, leaves and proof to calculate the merkle root of the tree and then compares it
 // to expectedRoot. Additionally, it reconstructs the parked nodes when each proven leaf was originally added to the
 // tree and returns a list of snapshots. This method is ~15% slower than ValidatePartialTree.
-func ValidatePartialTreeWithParkingSnapshots(leafIndices []uint64, leaves, proof [][]byte, expectedRoot []byte,
+func ValidatePartialTreeWithParkingSnapshots(
+	leafIndices []uint64,
+	leaves, proof [][]byte,
+	expectedRoot []byte,
 	hash HashFunc,
 ) (bool, []ParkingSnapshot, error) {
 	v, err := newValidator(leafIndices, leaves, proof, hash, true)
@@ -38,22 +45,20 @@ func ValidatePartialTreeWithParkingSnapshots(leafIndices []uint64, leaves, proof
 
 func newValidator(
 	leafIndices []uint64,
-	leaves,
-	proof [][]byte,
+	leaves, proof [][]byte,
 	hash HashFunc,
 	storeSnapshots bool,
 ) (*Validator, error) {
 	if len(leafIndices) != len(leaves) {
-		return nil, fmt.Errorf("number of leaves (%d) must equal number of indices (%d)", len(leaves),
-			len(leafIndices))
+		return nil, fmt.Errorf("number of leaves (%d) must equal number of indices (%d)", len(leaves), len(leafIndices))
 	}
 	if len(leaves) == 0 {
 		return nil, errors.New("at least one leaf is required for validation")
 	}
-	if !sort.SliceIsSorted(leafIndices, func(i, j int) bool { return leafIndices[i] < leafIndices[j] }) {
+	if !slices.IsSorted(leafIndices) {
 		return nil, errors.New("leafIndices are not sorted")
 	}
-	if len(SetOf(leafIndices...)) != len(leafIndices) {
+	if len(slices.Compact(leafIndices)) != len(leafIndices) {
 		return nil, errors.New("leafIndices contain duplicates")
 	}
 	proofNodes := &proofIterator{proof}