cmd/proofgen for subtree inclusion proofs

Author: phbnf

cmd/proofgen/main.go

@@ -19,6 +19,7 @@ import (
 	"encoding/json"
 	"fmt"
 	"log"
+	"math/bits"
 	"os"
 	"path/filepath"
 	"strconv"
@@ -109,6 +110,10 @@ var (
 	}
 )
 
+// =============================================================================
+// Inclusion Proofs
+// =============================================================================
+
 // inclusionProbe is a parameter set for inclusion proof verification.
 type inclusionProbe struct {
 	LeafIdx  uint64   `json:"leafIdx"`
@@ -290,6 +295,329 @@ func writeInclusionProbe(dir string, probe inclusionProbe) error {
 	return nil
 }
 
+// =============================================================================
+// Subtree Inclusion Proofs
+// =============================================================================
+
+// subtreeInclusionProbe is a parameter set for subtree inclusion proof verification.
+type subtreeInclusionProbe struct {
+	LeafIdx  uint64   `json:"leafIdx"`
+	Start    uint64   `json:"start"`
+	End      uint64   `json:"end"`
+	Root     []byte   `json:"root"`
+	LeafHash []byte   `json:"leafHash"`
+	Proof    [][]byte `json:"proof"`
+
+	Desc      string `json:"desc"`
+	WantError bool   `json:"wantErr"`
+}
+
+func bitCeil(n uint64) uint64 {
+	if n <= 1 {
+		return 1
+	}
+	return 1 << bits.Len64(n-1)
+}
+
+func toSubtreeInclusionProbe(p inclusionProbe) subtreeInclusionProbe {
+	return subtreeInclusionProbe{
+		LeafIdx:   p.LeafIdx,
+		Start:     0,
+		End:       p.TreeSize,
+		Root:      p.Root,
+		LeafHash:  p.LeafHash,
+		Proof:     p.Proof,
+		Desc:      p.Desc,
+		WantError: p.WantError,
+	}
+}
+
+func shiftSubtreeInclusionProbe(p subtreeInclusionProbe) subtreeInclusionProbe {
+	shift := bitCeil(p.End - p.Start)
+	desc := p.Desc + " - subtree"
+	leafIdx := p.LeafIdx + shift
+	if p.LeafIdx == ^uint64(0) {
+		leafIdx = ^uint64(0)
+	}
+	return subtreeInclusionProbe{
+		LeafIdx:   leafIdx,
+		Start:     p.Start + shift,
+		End:       p.End + shift,
+		Root:      p.Root,
+		LeafHash:  p.LeafHash,
+		Proof:     p.Proof,
+		Desc:      desc,
+		WantError: p.WantError,
+	}
+}
+
+func subtreeInclusionProbes(rootDir string) error {
+	for i, p := range inclusionProofs {
+		dir := filepath.Join(rootDir, strconv.Itoa(i))
+		if err := os.MkdirAll(dir, 0755); err != nil {
+			return err
+		}
+
+		leafHash := rfc6962.DefaultHasher.HashLeaf(leaves[p.leafIdx-1])
+		if err := corruptedSubtreeInclusionProbes(dir, p.leafIdx-1, p.size, p.proof, roots[p.size-1], leafHash); err != nil {
+			return err
+		}
+	}
+
+	staticDir := filepath.Join(rootDir, "additional")
+	if err := os.MkdirAll(staticDir, 0755); err != nil {
+		return err
+	}
+
+	if err := staticSubtreeInclusionProbes(staticDir); err != nil {
+		return err
+	}
+
+	singleEntryDir := filepath.Join(rootDir, "single-entry")
+	if err := os.MkdirAll(singleEntryDir, 0755); err != nil {
+		return err
+	}
+
+	if err := singleEntrySubtreeInclusionProbes(singleEntryDir); err != nil {
+		return err
+	}
+
+	return nil
+}
+
+func corruptedSubtreeInclusionProbes(dir string, leafIdx, treeSize uint64, proof [][]byte, root, leafHash []byte) error {
+	happyPath := inclusionProbe{leafIdx, treeSize, root, leafHash, proof, "happy path", false}
+	if err := writeSubtreeInclusionProbePair(dir, happyPath); err != nil {
+		return err
+	}
+
+	probes := invalidInclusionProof(leafIdx, treeSize, proof, root, leafHash)
+	for _, p := range probes {
+		if err := writeSubtreeInclusionProbePair(dir, p); err != nil {
+			return err
+		}
+	}
+
+	return nil
+}
+
+func singleEntrySubtreeInclusionProbes(dir string) error {
+	data := []byte("data")
+	hash := rfc6962.DefaultHasher.HashLeaf(data)
+	proof := [][]byte{}
+	emptyHash := []byte{}
+
+	for _, p := range []struct {
+		root    []byte
+		leaf    []byte
+		desc    string
+		wantErr bool
+	}{
+		{hash, hash, "matching root and leaf", false},
+		{hash, emptyHash, "empty leaf", true},
+		{emptyHash, hash, "empty root", true},
+		{emptyHash, emptyHash, "empty root and leaf", true},
+	} {
+		probe := inclusionProbe{0, 1, p.root, p.leaf, proof, p.desc, p.wantErr}
+		if err := writeSubtreeInclusionProbePair(dir, probe); err != nil {
+			return err
+		}
+	}
+
+	return nil
+}
+
+func staticSubtreeInclusionProbes(rootDir string) error {
+	proof := [][]byte{}
+
+	probes := []struct {
+		index, size uint64
+	}{{0, 0}, {0, 1}, {1, 0}, {2, 1}}
+	for i, p := range probes {
+		dir := filepath.Join(rootDir, strconv.Itoa(i))
+		err := os.MkdirAll(dir, 0755)
+		if err != nil {
+			return err
+		}
+
+		randomLeaf := inclusionProbe{p.index, p.size, []byte{}, sha256SomeHash, proof, "random leaf", true}
+		if err := writeSubtreeInclusionProbePair(dir, randomLeaf); err != nil {
+			return err
+		}
+
+		emptyRoot := inclusionProbe{p.index, p.size, sha256EmptyTreeHash, []byte{}, proof, "empty root", true}
+		if err := writeSubtreeInclusionProbePair(dir, emptyRoot); err != nil {
+			return err
+		}
+
+		emptyRootRandomLeaf := inclusionProbe{p.index, p.size, sha256EmptyTreeHash, sha256SomeHash, proof, "empty root and random leaf", true}
+		if err := writeSubtreeInclusionProbePair(dir, emptyRootRandomLeaf); err != nil {
+			return err
+		}
+	}
+
+	return nil
+}
+
+func errorSubtreeInclusionProbes(rootDir string) error {
+	dir := filepath.Join(rootDir, "errors")
+	if err := os.MkdirAll(dir, 0755); err != nil {
+		return err
+	}
+
+	leafHash := rfc6962.DefaultHasher.HashLeaf(leaves[0])
+
+	tests := []struct {
+		filename string
+		probe    subtreeInclusionProbe
+	}{
+		{
+			filename: "everything-zero.json",
+			probe: subtreeInclusionProbe{
+				LeafIdx:   0,
+				Start:     0,
+				End:       0,
+				Root:      sha256EmptyTreeHash,
+				LeafHash:  leafHash,
+				Proof:     nil,
+				Desc:      "empty root",
+				WantError: true,
+			},
+		},
+		{
+			filename: "start-equals-end.json",
+			probe: subtreeInclusionProbe{
+				LeafIdx:   0,
+				Start:     1,
+				End:       1,
+				Root:      sha256EmptyTreeHash,
+				LeafHash:  leafHash,
+				Proof:     nil,
+				Desc:      "empty root",
+				WantError: true,
+			},
+		},
+		{
+			filename: "invalid-subtree.json",
+			probe: subtreeInclusionProbe{
+				LeafIdx:   3,
+				Start:     3,
+				End:       5,
+				Root:      sha256EmptyTreeHash,
+				LeafHash:  leafHash,
+				Proof:     nil,
+				Desc:      "empty root",
+				WantError: true,
+			},
+		},
+		{
+			filename: "invalid-large-subtree.json",
+			probe: subtreeInclusionProbe{
+				LeafIdx:   1,
+				Start:     1,
+				End:       9223372036854775810,
+				Root:      sha256EmptyTreeHash,
+				LeafHash:  leafHash,
+				Proof:     nil,
+				Desc:      "empty root",
+				WantError: true,
+			},
+		},
+		{
+			filename: "oob-left.json",
+			probe: subtreeInclusionProbe{
+				LeafIdx:   0,
+				Start:     1,
+				End:       2,
+				Root:      sha256EmptyTreeHash,
+				LeafHash:  leafHash,
+				Proof:     nil,
+				Desc:      "empty root",
+				WantError: true,
+			},
+		},
+		{
+			filename: "oob-right.json",
+			probe: subtreeInclusionProbe{
+				LeafIdx:   3,
+				Start:     0,
+				End:       2,
+				Root:      sha256EmptyTreeHash,
+				LeafHash:  leafHash,
+				Proof:     nil,
+				Desc:      "empty root",
+				WantError: true,
+			},
+		},
+		{
+			filename: "oob-right-2.json",
+			probe: subtreeInclusionProbe{
+				LeafIdx:   3,
+				Start:     0,
+				End:       3,
+				Root:      sha256EmptyTreeHash,
+				LeafHash:  leafHash,
+				Proof:     nil,
+				Desc:      "empty root",
+				WantError: true,
+			},
+		},
+		{
+			filename: "start-larger-than-end.json",
+			probe: subtreeInclusionProbe{
+				LeafIdx:   0,
+				Start:     2,
+				End:       1,
+				Root:      sha256EmptyTreeHash,
+				LeafHash:  leafHash,
+				Proof:     nil,
+				Desc:      "empty root",
+				WantError: true,
+			},
+		},
+	}
+
+	for _, tc := range tests {
+		if err := writeSubtreeInclusionProbeWithFilename(dir, tc.filename, tc.probe); err != nil {
+			return err
+		}
+	}
+
+	return nil
+}
+
+func writeSubtreeInclusionProbeWithFilename(dir, filename string, probe subtreeInclusionProbe) error {
+	probeJson, err := json.MarshalIndent(probe, "", "  ")
+	if err != nil {
+		return fmt.Errorf("marshaling probe: %s", err)
+	}
+
+	fileLocation := filepath.Join(dir, filename)
+	if err := os.WriteFile(fileLocation, probeJson, 0644); err != nil {
+		return fmt.Errorf("writing probe: %s: %s", filename, err)
+	}
+
+	return nil
+}
+
+func writeSubtreeInclusionProbe(dir string, probe subtreeInclusionProbe) error {
+	fn := fileName(probe.Desc)
+	return writeSubtreeInclusionProbeWithFilename(dir, fn, probe)
+}
+
+func writeSubtreeInclusionProbePair(dir string, p inclusionProbe) error {
+	sp1 := toSubtreeInclusionProbe(p)
+	if err := writeSubtreeInclusionProbe(dir, sp1); err != nil {
+		return err
+	}
+	sp2 := shiftSubtreeInclusionProbe(sp1)
+	return writeSubtreeInclusionProbe(dir, sp2)
+}
+
+// =============================================================================
+// Consistency Proofs
+// =============================================================================
+
 // consistencyProbe is a parameter set for consistency proof verification.
 type consistencyProbe struct {
 	Size1 uint64   `json:"size1"`
@@ -446,6 +774,10 @@ func writeConsistencyProbe(dir string, probe consistencyProbe) error {
 	return nil
 }
 
+// =============================================================================
+// General Helpers
+// =============================================================================
+
 // extend explicitly copies |proof| slice and appends |hashes| to it.
 func extend(proof [][]byte, hashes ...[]byte) [][]byte {
 	return append(append([][]byte{}, proof...), hashes...)
@@ -469,6 +801,7 @@ func dh(h string, expLen int) []byte {
 
 func fileName(n string) string {
 	r := strings.NewReplacer(
+		" - ", "-",
 		"(", "",
 		")", "",
 		" ", "-")
@@ -481,6 +814,14 @@ func main() {
 		log.Fatalf("writing inclusion test data: %s", err)
 	}
 
+	subtreeInclusionDir := "testdata/subtreeinclusion"
+	if err := subtreeInclusionProbes(subtreeInclusionDir); err != nil {
+		log.Fatalf("writing subtree inclusion test data: %s", err)
+	}
+	if err := errorSubtreeInclusionProbes(subtreeInclusionDir); err != nil {
+		log.Fatalf("writing subtree inclusion error test data: %s", err)
+	}
+
 	consistencyDir := "testdata/consistency"
 	if err := consistencyProbes(consistencyDir); err != nil {
 		log.Fatalf("writing consistency test data: %s", err)

testdata/subtreeinclusion/errors/everything-zero.json

@@ -7,4 +7,4 @@
   "proof": null,
   "desc": "empty root",
   "wantErr": true
-}
+}

testdata/subtreeinclusion/errors/oob-left.json

@@ -7,4 +7,4 @@
   "proof": null,
   "desc": "empty root",
   "wantErr": true
-}
+}