Subtree Inclusion proof fuzz tests (transparency-dev#229)

* fuzz tests

* refactor HashAt

Author: phbnf

proof/proof.go

@@ -223,7 +223,7 @@ func reverse(ids []compact.NodeID) {
 	}
 }
 
-// isSubTreeValid returns whether a subtree covers a valid range.
+// isSubtreeValid returns whether a subtree covers a valid range.
 // A subtree is valid if there exist a parent tree node to:
 // - all the subtree nodes
 // - no extra node to the left of the subtree

testonly/tree.go

@@ -82,10 +82,19 @@ func (t *Tree) Hash() []byte {
 // HashAt returns the root hash at the given size.
 // Requires 0 <= size <= Size(), otherwise panics.
 func (t *Tree) HashAt(size uint64) []byte {
-	if size == 0 {
+	return t.SubtreeHashAt(0, size)
+}
+
+// SubtreeHashAt returns the root hash of the [start, end) subtree.
+// Requires 0 <= start <= end <= Size() otherwise panics.
+func (t *Tree) SubtreeHashAt(start, end uint64) []byte {
+	if start > end || end > t.size {
+		panic("invalid subtree range")
+	}
+	if start == end {
 		return t.hasher.EmptyRoot()
 	}
-	hashes := t.getNodes(compact.RangeNodes(0, size, nil))
+	hashes := t.getNodes(compact.RangeNodes(start, end, nil))
 
 	hash := hashes[len(hashes)-1]
 	for i := len(hashes) - 2; i >= 0; i-- {
@@ -105,6 +114,18 @@ func (t *Tree) InclusionProof(index, size uint64) ([][]byte, error) {
 	return nodes.Rehash(t.getNodes(nodes.IDs), t.hasher.HashChildren)
 }
 
+// SubtreeInclusionProof returns the inclusion proof for the given leaf index in the
+// [start, end) subtree.
+// It requires end <= Size(), and may panic otherwise.
+// May return and error if the subtree boundaries or the index are not valid.
+func (t *Tree) SubtreeInclusionProof(index, start, end uint64) ([][]byte, error) {
+	nodes, err := proof.SubtreeInclusion(index, start, end)
+	if err != nil {
+		return nil, err
+	}
+	return nodes.Rehash(t.getNodes(nodes.IDs), t.hasher.HashChildren)
+}
+
 // ConsistencyProof returns the consistency proof between the two given tree
 // sizes. Requires 0 <= size1 <= size2 <= Size(), otherwise may panic.
 func (t *Tree) ConsistencyProof(size1, size2 uint64) ([][]byte, error) {

testonly/tree_fuzz_test.go

@@ -5,6 +5,7 @@ package testonly
 import (
 	"bytes"
 	"math"
+	"math/bits"
 	"testing"
 
 	"github.com/google/go-cmp/cmp"
@@ -74,6 +75,45 @@ func FuzzInclusionProofAndVerify(f *testing.F) {
 	})
 }
 
+// Compute and verify inclusion proofs
+func FuzzSubtreeInclusionProofAndVerify(f *testing.F) {
+	for end := 0; end <= 8; end++ {
+		for start := 0; start <= end; start++ {
+			for index := start; index <= end; index++ {
+				f.Add(uint64(index), uint64(start), uint64(end))
+			}
+		}
+	}
+	f.Fuzz(func(t *testing.T, index, start, end uint64) {
+		if end >= math.MaxUint16 {
+			return
+		}
+		t.Logf("index=%d, start=%d, end=%d", index, start, end)
+		if start >= end {
+			return
+		}
+		if index < start {
+			return
+		}
+		if index >= end {
+			return
+		}
+		if bc := uint64(1) << bits.Len64(end-start-1); start%bc != 0 {
+			return
+		}
+		tree := newTree(genEntries(end))
+		p, err := tree.SubtreeInclusionProof(index, start, end)
+		t.Logf("proof=%v", p)
+		if err != nil {
+			t.Error(err)
+		}
+		err = proof.VerifySubtreeInclusion(tree.hasher, index, start, end, tree.LeafHash(index), p, tree.SubtreeHashAt(start, end))
+		if err != nil {
+			t.Error(err)
+		}
+	})
+}
+
 func FuzzHashAtAgainstReferenceImplementation(f *testing.F) {
 	for size := 0; size <= 8; size++ {
 		for index := 0; index <= size; index++ {