factor in

Author: phbnf

cmd/proofgen/main.go

@@ -126,15 +126,15 @@ type inclusionProbe struct {
 	WantError bool   `json:"wantErr"`
 }
 
-func inclusionProbes(rootDir string) error {
+func generateInclusionProbes(rootDir string, write func(string, inclusionProbe) error) 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 := corruptedInclusionProbes(dir, p.leafIdx-1, p.size, p.proof, roots[p.size-1], leafHash); err != nil {
+		if err := corruptedInclusionProbes(dir, p.leafIdx-1, p.size, p.proof, roots[p.size-1], leafHash, write); err != nil {
 			return err
 		}
 	}
@@ -144,7 +144,7 @@ func inclusionProbes(rootDir string) error {
 		return err
 	}
 
-	if err := staticInclusionProbes(staticDir); err != nil {
+	if err := staticInclusionProbes(staticDir, write); err != nil {
 		return err
 	}
 
@@ -153,7 +153,7 @@ func inclusionProbes(rootDir string) error {
 		return err
 	}
 
-	if err := singleEntryInclusionProbes(singleEntryDir); err != nil {
+	if err := singleEntryInclusionProbes(singleEntryDir, write); err != nil {
 		return err
 	}
 
@@ -202,23 +202,23 @@ func invalidInclusionProof(leafIdx, treeSize uint64, proof [][]byte, root, leafH
 	return ret
 }
 
-func corruptedInclusionProbes(dir string, leafIdx, treeSize uint64, proof [][]byte, root, leafHash []byte) error {
+func corruptedInclusionProbes(dir string, leafIdx, treeSize uint64, proof [][]byte, root, leafHash []byte, write func(string, inclusionProbe) error) error {
 	happyPath := inclusionProbe{leafIdx, treeSize, root, leafHash, proof, "happy path", false}
-	if err := writeInclusionProbe(dir, happyPath); err != nil {
-		return nil
+	if err := write(dir, happyPath); err != nil {
+		return err
 	}
 
 	probes := invalidInclusionProof(leafIdx, treeSize, proof, root, leafHash)
 	for _, p := range probes {
-		if err := writeInclusionProbe(dir, p); err != nil {
+		if err := write(dir, p); err != nil {
 			return err
 		}
 	}
 
 	return nil
 }
 
-func singleEntryInclusionProbes(dir string) error {
+func singleEntryInclusionProbes(dir string, write func(string, inclusionProbe) error) error {
 	data := []byte("data")
 	// Root and leaf hash for 1-entry tree are the same.
 	hash := rfc6962.DefaultHasher.HashLeaf(data)
@@ -239,15 +239,15 @@ func singleEntryInclusionProbes(dir string) error {
 	} {
 		probe := inclusionProbe{0, 1, p.root, p.leaf, proof, p.desc, p.wantErr}
 
-		if err := writeInclusionProbe(dir, probe); err != nil {
+		if err := write(dir, probe); err != nil {
 			return err
 		}
 	}
 
 	return nil
 }
 
-func staticInclusionProbes(rootDir string) error {
+func staticInclusionProbes(rootDir string, write func(string, inclusionProbe) error) error {
 	proof := [][]byte{}
 
 	probes := []struct {
@@ -261,17 +261,17 @@ func staticInclusionProbes(rootDir string) error {
 		}
 
 		randomLeaf := inclusionProbe{p.index, p.size, []byte{}, sha256SomeHash, proof, "random leaf", true}
-		if err := writeInclusionProbe(dir, randomLeaf); err != nil {
+		if err := write(dir, randomLeaf); err != nil {
 			return err
 		}
 
 		emptyRoot := inclusionProbe{p.index, p.size, sha256EmptyTreeHash, []byte{}, proof, "empty root", true}
-		if err := writeInclusionProbe(dir, emptyRoot); err != nil {
+		if err := write(dir, emptyRoot); err != nil {
 			return err
 		}
 
 		emptyRootRandomLeaf := inclusionProbe{p.index, p.size, sha256EmptyTreeHash, sha256SomeHash, proof, "empty root and random leaf", true}
-		if err := writeInclusionProbe(dir, emptyRootRandomLeaf); err != nil {
+		if err := write(dir, emptyRootRandomLeaf); err != nil {
 			return err
 		}
 	}
@@ -351,113 +351,7 @@ func shiftSubtreeInclusionProbe(p subtreeInclusionProbe) subtreeInclusionProbe {
 	}
 }
 
-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 := convertToSubtreeInclusionProbesAndWrite(dir, happyPath); err != nil {
-		return err
-	}
-
-	probes := invalidInclusionProof(leafIdx, treeSize, proof, root, leafHash)
-	for _, p := range probes {
-		if err := convertToSubtreeInclusionProbesAndWrite(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 := convertToSubtreeInclusionProbesAndWrite(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 := convertToSubtreeInclusionProbesAndWrite(dir, randomLeaf); err != nil {
-			return err
-		}
-
-		emptyRoot := inclusionProbe{p.index, p.size, sha256EmptyTreeHash, []byte{}, proof, "empty root", true}
-		if err := convertToSubtreeInclusionProbesAndWrite(dir, emptyRoot); err != nil {
-			return err
-		}
-
-		emptyRootRandomLeaf := inclusionProbe{p.index, p.size, sha256EmptyTreeHash, sha256SomeHash, proof, "empty root and random leaf", true}
-		if err := convertToSubtreeInclusionProbesAndWrite(dir, emptyRootRandomLeaf); err != nil {
-			return err
-		}
-	}
-
-	return nil
-}
 
 func errorSubtreeInclusionProbes(rootDir string) error {
 	dir := filepath.Join(rootDir, "errors")
@@ -789,12 +683,12 @@ func fileName(n string) string {
 
 func main() {
 	inclusionDir := "testdata/inclusion"
-	if err := inclusionProbes(inclusionDir); err != nil {
+	if err := generateInclusionProbes(inclusionDir, writeInclusionProbe); err != nil {
 		log.Fatalf("writing inclusion test data: %s", err)
 	}
 
 	subtreeInclusionDir := "testdata/subtreeinclusion"
-	if err := subtreeInclusionProbes(subtreeInclusionDir); err != nil {
+	if err := generateInclusionProbes(subtreeInclusionDir, convertToSubtreeInclusionProbesAndWrite); err != nil {
 		log.Fatalf("writing subtree inclusion test data: %s", err)
 	}
 	if err := errorSubtreeInclusionProbes(subtreeInclusionDir); err != nil {