feat: add merkle package from cometbft (#2)

go.work

@@ -0,0 +1,6 @@
+go 1.21.5
+
+use (
+	.
+	./merkle
+)

go.work.sum

@@ -0,0 +1,3 @@
+github.com/stretchr/objx v0.5.0/go.mod h1:Yh+to48EsGEfYuaHDzXPcE3xhTkx73EhmCGUpEOglKo=
+golang.org/x/mod v0.14.0/go.mod h1:hTbmBsO62+eylJbnUtE2MGJUyE7QWk4xUqPFrRgJ+7c=
+golang.org/x/tools v0.16.0/go.mod h1:kYVVN6I1mBNoB1OX+noeBjbRk4IUEPa7JJ+TJMEooJ0=

merkle/README.md

@@ -0,0 +1,6 @@
+# Merkle Tree
+
+For smaller static data structures that don't require immutable snapshots or mutability;
+for instance the transactions and validation signatures of a block can be hashed using this simple merkle tree logic.
+
+This is a forked copy of github.com/cometbft/cometbft/crypto/merkle

merkle/doc.go

@@ -0,0 +1,30 @@
+/*
+Package merkle computes a deterministic minimal height Merkle tree hash.
+If the number of items is not a power of two, some leaves
+will be at different levels. Tries to keep both sides of
+the tree the same size, but the left may be one greater.
+
+Use this for short deterministic trees, such as the validator list.
+For larger datasets, use IAVLTree.
+
+Be aware that the current implementation by itself does not prevent
+second pre-image attacks. Hence, use this library with caution.
+Otherwise you might run into similar issues as, e.g., in early Bitcoin:
+https://bitcointalk.org/?topic=102395
+
+	              *
+	             / \
+	           /     \
+	         /         \
+	       /             \
+	      *               *
+	     / \             / \
+	    /   \           /   \
+	   /     \         /     \
+	  *       *       *       h6
+	 / \     / \     / \
+	h0  h1  h2  h3  h4  h5
+
+TODO(ismail): add 2nd pre-image protection or clarify further on how we use this and why this secure.
+*/
+package merkle

merkle/go.mod

@@ -0,0 +1,15 @@
+module github.com/celestiaorg/go-square/merkle
+
+go 1.21.5
+
+require (
+	github.com/stretchr/testify v1.8.4
+	google.golang.org/protobuf v1.31.0
+)
+
+require (
+	github.com/davecgh/go-spew v1.1.1 // indirect
+	github.com/google/go-cmp v0.5.9 // indirect
+	github.com/pmezard/go-difflib v1.0.0 // indirect
+	gopkg.in/yaml.v3 v3.0.1 // indirect
+)

merkle/go.sum

@@ -0,0 +1,18 @@
+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/golang/protobuf v1.5.0/go.mod h1:FsONVRAS9T7sI+LIUmWTfcYkHO4aIWwzhcaSAoJOfIk=
+github.com/google/go-cmp v0.5.5/go.mod h1:v8dTdLbMG2kIc/vJvl+f65V22dbkXbowE6jgT/gNBxE=
+github.com/google/go-cmp v0.5.9 h1:O2Tfq5qg4qc4AmwVlvv0oLiVAGB7enBSJ2x2DqQFi38=
+github.com/google/go-cmp v0.5.9/go.mod h1:17dUlkBOakJ0+DkrSSNjCkIjxS6bF9zb3elmeNGIjoY=
+github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
+github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
+github.com/stretchr/testify v1.8.4 h1:CcVxjf3Q8PM0mHUKJCdn+eZZtm5yQwehR5yeSVQQcUk=
+github.com/stretchr/testify v1.8.4/go.mod h1:sz/lmYIOXD/1dqDmKjjqLyZ2RngseejIcXlSw2iwfAo=
+golang.org/x/xerrors v0.0.0-20191204190536-9bdfabe68543/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
+google.golang.org/protobuf v1.26.0-rc.1/go.mod h1:jlhhOSvTdKEhbULTjvd4ARK9grFBp09yW+WbY/TyQbw=
+google.golang.org/protobuf v1.31.0 h1:g0LDEJHgrBl9N9r17Ru3sqWhkIx2NB67okBHPwC7hs8=
+google.golang.org/protobuf v1.31.0/go.mod h1:HV8QOd/L58Z+nl8r43ehVNZIU/HEI6OcFqwMG9pJV4I=
+gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405 h1:yhCVgyC4o1eVCa2tZl7eS0r+SDo693bJlVdllGtEeKM=
+gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
+gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA=
+gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=

merkle/hash.go

@@ -0,0 +1,31 @@
+package merkle
+
+import (
+	"crypto/sha256"
+)
+
+// TODO: make these have a large predefined capacity
+var (
+	leafPrefix  = []byte{0}
+	innerPrefix = []byte{1}
+)
+
+// returns empty sha256 hash
+func emptyHash() []byte {
+	return hash([]byte{})
+}
+
+// returns sha256(0x00 || leaf)
+func leafHash(leaf []byte) []byte {
+	return hash(append(leafPrefix, leaf...))
+}
+
+// returns sha256(0x01 || left || right)
+func innerHash(left []byte, right []byte) []byte {
+	return hash(append(innerPrefix, append(left, right...)...))
+}
+
+func hash(bz []byte) []byte {
+	h := sha256.Sum256(bz)
+	return h[:]
+}

merkle/proof.go

@@ -0,0 +1,252 @@
+package merkle
+
+import (
+	"bytes"
+	"crypto/sha256"
+	"errors"
+	"fmt"
+
+	wire "github.com/celestiaorg/go-square/merkle/proto/gen/merkle/v1"
+)
+
+const (
+	// MaxAunts is the maximum number of aunts that can be included in a Proof.
+	// This corresponds to a tree of size 2^100, which should be sufficient for all conceivable purposes.
+	// This maximum helps prevent Denial-of-Service attacks by limiting the size of the proofs.
+	MaxAunts = 100
+)
+
+// Proof represents a Merkle proof.
+// NOTE: The convention for proofs is to include leaf hashes but to
+// exclude the root hash.
+// This convention is implemented across IAVL range proofs as well.
+// Keep this consistent unless there's a very good reason to change
+// everything.  This also affects the generalized proof system as
+// well.
+type Proof struct {
+	Total    int64    `json:"total"`     // Total number of items.
+	Index    int64    `json:"index"`     // Index of item to prove.
+	LeafHash []byte   `json:"leaf_hash"` // Hash of item value.
+	Aunts    [][]byte `json:"aunts"`     // Hashes from leaf's sibling to a root's child.
+}
+
+// ProofsFromByteSlices computes inclusion proof for given items.
+// proofs[0] is the proof for items[0].
+func ProofsFromByteSlices(items [][]byte) (rootHash []byte, proofs []*Proof) {
+	trails, rootSPN := trailsFromByteSlices(items)
+	rootHash = rootSPN.Hash
+	proofs = make([]*Proof, len(items))
+	for i, trail := range trails {
+		proofs[i] = &Proof{
+			Total:    int64(len(items)),
+			Index:    int64(i),
+			LeafHash: trail.Hash,
+			Aunts:    trail.FlattenAunts(),
+		}
+	}
+	return
+}
+
+// Verify that the Proof proves the root hash.
+// Check sp.Index/sp.Total manually if needed
+func (sp *Proof) Verify(rootHash []byte, leaf []byte) error {
+	if rootHash == nil {
+		return fmt.Errorf("invalid root hash: cannot be nil")
+	}
+	if sp.Total < 0 {
+		return errors.New("proof total must be positive")
+	}
+	if sp.Index < 0 {
+		return errors.New("proof index cannot be negative")
+	}
+	leafHash := leafHash(leaf)
+	if !bytes.Equal(sp.LeafHash, leafHash) {
+		return fmt.Errorf("invalid leaf hash: wanted %X got %X", leafHash, sp.LeafHash)
+	}
+	computedHash, err := sp.computeRootHash()
+	if err != nil {
+		return fmt.Errorf("compute root hash: %w", err)
+	}
+	if !bytes.Equal(computedHash, rootHash) {
+		return fmt.Errorf("invalid root hash: wanted %X got %X", rootHash, computedHash)
+	}
+	return nil
+}
+
+// Compute the root hash given a leaf hash.  Panics in case of errors.
+func (sp *Proof) ComputeRootHash() []byte {
+	computedHash, err := sp.computeRootHash()
+	if err != nil {
+		panic(fmt.Errorf("ComputeRootHash errored %w", err))
+	}
+	return computedHash
+}
+
+// Compute the root hash given a leaf hash.
+func (sp *Proof) computeRootHash() ([]byte, error) {
+	return computeHashFromAunts(
+		sp.Index,
+		sp.Total,
+		sp.LeafHash,
+		sp.Aunts,
+	)
+}
+
+// String implements the stringer interface for Proof.
+// It is a wrapper around StringIndented.
+func (sp *Proof) String() string {
+	return sp.StringIndented("")
+}
+
+// StringIndented generates a canonical string representation of a Proof.
+func (sp *Proof) StringIndented(indent string) string {
+	return fmt.Sprintf(`Proof{
+%s  Aunts: %X
+%s}`,
+		indent, sp.Aunts,
+		indent)
+}
+
+// ValidateBasic performs basic validation.
+// NOTE: it expects the LeafHash and the elements of Aunts to be of size tmhash.Size,
+// and it expects at most MaxAunts elements in Aunts.
+func (sp *Proof) ValidateBasic() error {
+	if sp.Total < 0 {
+		return errors.New("negative Total")
+	}
+	if sp.Index < 0 {
+		return errors.New("negative Index")
+	}
+	if len(sp.LeafHash) != sha256.Size {
+		return fmt.Errorf("expected LeafHash size to be %d, got %d", sha256.Size, len(sp.LeafHash))
+	}
+	if len(sp.Aunts) > MaxAunts {
+		return fmt.Errorf("expected no more than %d aunts, got %d", MaxAunts, len(sp.Aunts))
+	}
+	for i, auntHash := range sp.Aunts {
+		if len(auntHash) != sha256.Size {
+			return fmt.Errorf("expected Aunts#%d size to be %d, got %d", i, sha256.Size, len(auntHash))
+		}
+	}
+	return nil
+}
+
+func (sp *Proof) ToProto() *wire.Proof {
+	if sp == nil {
+		return nil
+	}
+	pb := new(wire.Proof)
+
+	pb.Total = sp.Total
+	pb.Index = sp.Index
+	pb.LeafHash = sp.LeafHash
+	pb.Aunts = sp.Aunts
+
+	return pb
+}
+
+func ProofFromProto(pb *wire.Proof) (*Proof, error) {
+	if pb == nil {
+		return nil, errors.New("nil proof")
+	}
+
+	sp := new(Proof)
+
+	sp.Total = pb.Total
+	sp.Index = pb.Index
+	sp.LeafHash = pb.LeafHash
+	sp.Aunts = pb.Aunts
+
+	return sp, sp.ValidateBasic()
+}
+
+// Use the leafHash and innerHashes to get the root merkle hash.
+// If the length of the innerHashes slice isn't exactly correct, the result is nil.
+// Recursive impl.
+func computeHashFromAunts(index, total int64, leafHash []byte, innerHashes [][]byte) ([]byte, error) {
+	if index >= total || index < 0 || total <= 0 {
+		return nil, fmt.Errorf("invalid index %d and/or total %d", index, total)
+	}
+	switch total {
+	case 0:
+		panic("Cannot call computeHashFromAunts() with 0 total")
+	case 1:
+		if len(innerHashes) != 0 {
+			return nil, fmt.Errorf("unexpected inner hashes")
+		}
+		return leafHash, nil
+	default:
+		if len(innerHashes) == 0 {
+			return nil, fmt.Errorf("expected at least one inner hash")
+		}
+		numLeft := getSplitPoint(total)
+		if index < numLeft {
+			leftHash, err := computeHashFromAunts(index, numLeft, leafHash, innerHashes[:len(innerHashes)-1])
+			if err != nil {
+				return nil, err
+			}
+
+			return innerHash(leftHash, innerHashes[len(innerHashes)-1]), nil
+		}
+		rightHash, err := computeHashFromAunts(index-numLeft, total-numLeft, leafHash, innerHashes[:len(innerHashes)-1])
+		if err != nil {
+			return nil, err
+		}
+		return innerHash(innerHashes[len(innerHashes)-1], rightHash), nil
+	}
+}
+
+// ProofNode is a helper structure to construct merkle proof.
+// The node and the tree is thrown away afterwards.
+// Exactly one of node.Left and node.Right is nil, unless node is the root, in which case both are nil.
+// node.Parent.Hash = hash(node.Hash, node.Right.Hash) or
+// hash(node.Left.Hash, node.Hash), depending on whether node is a left/right child.
+type ProofNode struct {
+	Hash   []byte
+	Parent *ProofNode
+	Left   *ProofNode // Left sibling  (only one of Left,Right is set)
+	Right  *ProofNode // Right sibling (only one of Left,Right is set)
+}
+
+// FlattenAunts will return the inner hashes for the item corresponding to the leaf,
+// starting from a leaf ProofNode.
+func (spn *ProofNode) FlattenAunts() [][]byte {
+	// Nonrecursive impl.
+	innerHashes := [][]byte{}
+	for spn != nil {
+		switch {
+		case spn.Left != nil:
+			innerHashes = append(innerHashes, spn.Left.Hash)
+		case spn.Right != nil:
+			innerHashes = append(innerHashes, spn.Right.Hash)
+		default:
+			break
+		}
+		spn = spn.Parent
+	}
+	return innerHashes
+}
+
+// trails[0].Hash is the leaf hash for items[0].
+// trails[i].Parent.Parent....Parent == root for all i.
+func trailsFromByteSlices(items [][]byte) (trails []*ProofNode, root *ProofNode) {
+	// Recursive impl.
+	switch len(items) {
+	case 0:
+		return []*ProofNode{}, &ProofNode{emptyHash(), nil, nil, nil}
+	case 1:
+		trail := &ProofNode{leafHash(items[0]), nil, nil, nil}
+		return []*ProofNode{trail}, trail
+	default:
+		k := getSplitPoint(int64(len(items)))
+		lefts, leftRoot := trailsFromByteSlices(items[:k])
+		rights, rightRoot := trailsFromByteSlices(items[k:])
+		rootHash := innerHash(leftRoot.Hash, rightRoot.Hash)
+		root := &ProofNode{rootHash, nil, nil, nil}
+		leftRoot.Parent = root
+		leftRoot.Right = rightRoot
+		rightRoot.Parent = root
+		rightRoot.Left = leftRoot
+		return append(lefts, rights...), root
+	}
+}