torchwood: add HashReaderOverlay

Author: FiloSottile

NEWS.md

@@ -11,6 +11,9 @@
 - Added `NewCosignatureVerifierFromKey` to produce a verifier from a
   `crypto.PublicKey` instead of a vkey.
 
+- Added `HashReaderOverlay`, a `tlog.HashReader` that extends a tree with the
+  stored hashes of appended records, held in memory.
+
 ### mpt
 
 - New package with a Merkle Patricia Trie implementation.

cmd/spicy/spicy.go

@@ -169,18 +169,19 @@ func main() {
 		log.Fatalf("could not parse latest checkpoint: %v", err)
 	}
 
-	hashes := make(map[int64]tlog.Hash)
-	hashReader := tlog.HashReaderFunc(func(indexes []int64) ([]tlog.Hash, error) {
+	edgeHashes := make(map[int64]tlog.Hash)
+	edgeReader := tlog.HashReaderFunc(func(indexes []int64) ([]tlog.Hash, error) {
 		list := make([]tlog.Hash, 0, len(indexes))
 		for _, id := range indexes {
-			h, ok := hashes[id]
+			h, ok := edgeHashes[id]
 			if !ok {
-				return nil, fmt.Errorf("index %d not in hashes", id)
+				return nil, fmt.Errorf("index %d not in edge", id)
 			}
 			list = append(list, h)
 		}
 		return list, nil
 	})
+	hashReader := torchwood.NewHashReaderOverlay(c.N, edgeReader)
 
 	edge, err := os.ReadFile(filepath.Join(*assetsFlag, "edge"))
 	if err != nil {
@@ -210,38 +211,34 @@ func main() {
 		if err != nil {
 			log.Fatalf("malformed edge file: %v", err)
 		}
-		hashes[idx[i]] = hash
+		edgeHashes[idx[i]] = hash
 	}
 
 	fmt.Fprintf(os.Stderr, "Log loaded.\n")
 	fmt.Fprintf(os.Stderr, "  - Name: %s\n", c.Origin)
 	fmt.Fprintf(os.Stderr, "  - Current size: %d\n", c.N)
 	fmt.Fprintf(os.Stderr, "  - Assets directory: %s\n", *assetsFlag)
 
-	for i, path := range flag.Args() {
+	for _, path := range flag.Args() {
 		if _, err := os.Stat(path + ".spicy"); err == nil {
 			log.Fatalf("spicy signature already exists for %q", path)
 		}
 		f, err := os.ReadFile(path)
 		if err != nil {
 			log.Fatalf("could not read %q: %v", path, err)
 		}
-		n := c.N + int64(i)
-		hh, err := tlog.StoredHashes(n, f, hashReader)
-		if err != nil {
+		n := hashReader.Size()
+		if err := hashReader.AppendRecordHash(tlog.RecordHash(f)); err != nil {
 			log.Fatalf("could not append %q: %v", path, err)
 		}
-		for k, h := range hh {
-			hashes[tlog.StoredHashIndex(0, n)+int64(k)] = h
-		}
 		entryPath := filepath.Join(*assetsFlag, strconv.FormatInt(n, 10))
 		if err := os.WriteFile(entryPath, f, 0644); err != nil {
 			log.Fatalf("could not copy %q to assets: %v", path, err)
 		}
 		fmt.Fprintf(os.Stderr, "  + %q is now entry %d\n", path, n)
 	}
 
-	N := c.N + int64(len(flag.Args()))
+	N := hashReader.Size()
 	th, err := tlog.TreeHash(N, hashReader)
 	if err != nil {
 		log.Fatalf("could not compute tree hash: %v", err)
@@ -255,8 +252,12 @@ func main() {
 		log.Fatalf("could not sign new checkpoint: %v", err)
 	}
 	newEdge := fmt.Sprintf("size %d\n", N)
-	for _, idx := range torchwood.RightEdge(N) {
-		newEdge += fmt.Sprintf("%s\n", hashes[idx])
+	newEdgeHashes, err := hashReader.ReadHashes(torchwood.RightEdge(N))
+	if err != nil {
+		log.Fatalf("could not read new edge hashes: %v", err)
+	}
+	for _, h := range newEdgeHashes {
+		newEdge += fmt.Sprintf("%s\n", h)
 	}
 
 	if err := os.WriteFile(filepath.Join(*assetsFlag, "latest"), newCheckpoint, 0644); err != nil {

overlay.go

@@ -0,0 +1,101 @@
+package torchwood
+
+import (
+	"fmt"
+
+	"golang.org/x/mod/sumdb/tlog"
+)
+
+// A HashReaderOverlay is a [tlog.HashReader] for a tree that extends an
+// original tree with appended records, whose stored hashes are held in memory.
+//
+// It can be used to stage the hashes of new records before they are persisted,
+// and to compute tree hashes, subtree hashes, proofs, and tiles that span both
+// the original tree and the appended records.
+//
+// The underlying HashReader is responsible for serving any requested hashes in
+// the original tree. Which hashes are requested depends on the caller:
+// [HashReaderOverlay.AppendRecordHash], [tlog.StoredHashes], [tlog.TreeHash],
+// and [SubtreeHash] only need the [RightEdge] of the original tree (as long as
+// the tree or subtree extends past the original tree), while
+// [tlog.ReadTileData] might need more (such as the full contents of the
+// rightmost tiles).
+type HashReaderOverlay struct {
+	r      tlog.HashReader
+	size   int64
+	base   int64 // tlog.StoredHashCount of the original tree size
+	hashes []tlog.Hash
+}
+
+// NewHashReaderOverlay returns a [HashReaderOverlay] over an original tree of
+// size n, whose hashes are served by r. If n is 0, r may be nil.
+//
+// n must not be negative or greater than 2^62, the largest supported tree
+// size; NewHashReaderOverlay panics otherwise.
+func NewHashReaderOverlay(n int64, r tlog.HashReader) *HashReaderOverlay {
+	if n < 0 || n > maxN {
+		panic("tlog: tree size out of range in NewHashReaderOverlay")
+	}
+	return &HashReaderOverlay{r: r, size: n, base: tlog.StoredHashCount(n)}
+}
+
+// Size returns the size of the extended tree: the size of the original tree
+// plus the number of appended leaf records.
+func (o *HashReaderOverlay) Size() int64 {
+	return o.size
+}
+
+// AppendRecordHash extends the tree by one record with the given record hash
+// (as produced by [tlog.RecordHash]), computing and adding the record's stored
+// hashes to the overlay.
+//
+// The hashes of completed subtrees may be requested from the overlay, which
+// needs at most the [RightEdge] of the original tree from the underlying
+// HashReader.
+func (o *HashReaderOverlay) AppendRecordHash(rh tlog.Hash) error {
+	if o.size >= maxN {
+		return fmt.Errorf("tlog: tree size out of range in AppendRecordHash")
+	}
+	hashes, err := tlog.StoredHashesForRecordHash(o.size, rh, o)
+	if err != nil {
+		return err
+	}
+	o.hashes = append(o.hashes, hashes...)
+	o.size++
+	return nil
+}
+
+// ReadHashes implements [tlog.HashReader]. Any hashes in the original tree are
+// requested from the underlying HashReader in a single ReadHashes call.
+func (o *HashReaderOverlay) ReadHashes(indexes []int64) ([]tlog.Hash, error) {
+	list := make([]tlog.Hash, len(indexes))
+	var origIndexes []int64
+	var origPositions []int
+	for i, id := range indexes {
+		switch {
+		case id < 0 || id >= o.base+int64(len(o.hashes)):
+			return nil, fmt.Errorf("tlog: index %d out of range for tree of size %d in HashReaderOverlay", id, o.size)
+		case id < o.base:
+			origIndexes = append(origIndexes, id)
+			origPositions = append(origPositions, i)
+		default:
+			list[i] = o.hashes[id-o.base]
+		}
+	}
+	if len(origIndexes) > 0 {
+		if o.r == nil {
+			return nil, fmt.Errorf("tlog: no underlying HashReader for original tree in HashReaderOverlay")
+		}
+		hashes, err := o.r.ReadHashes(origIndexes)
+		if err != nil {
+			return nil, err
+		}
+		if len(hashes) != len(origIndexes) {
+			return nil, fmt.Errorf("tlog: ReadHashes(%d indexes) = %d hashes", len(origIndexes), len(hashes))
+		}
+		for i, h := range hashes {
+			list[origPositions[i]] = h
+		}
+	}
+	return list, nil
+}

overlay_test.go

@@ -0,0 +1,202 @@
+package torchwood_test
+
+import (
+	"fmt"
+	"slices"
+	"testing"
+
+	"filippo.io/torchwood"
+	"golang.org/x/mod/sumdb/tlog"
+)
+
+// TestHashReaderOverlayRightEdge checks that appending records, computing tree
+// and subtree hashes, and proving new records only need the right edge of the
+// original tree.
+func TestHashReaderOverlayRightEdge(t *testing.T) {
+	const N = 130
+	leaves, full := buildTree(t, N)
+	for n0 := int64(0); n0 <= N; n0++ {
+		edge := make(map[int64]tlog.Hash)
+		if n0 > 0 {
+			idx := torchwood.RightEdge(n0)
+			hashes, err := full.ReadHashes(idx)
+			fatalIfErr(t, err)
+			for i, id := range idx {
+				edge[id] = hashes[i]
+			}
+		}
+		r := tlog.HashReaderFunc(func(indexes []int64) ([]tlog.Hash, error) {
+			list := make([]tlog.Hash, 0, len(indexes))
+			for _, id := range indexes {
+				h, ok := edge[id]
+				if !ok {
+					return nil, fmt.Errorf("index %d not on the right edge of tree of size %d", id, n0)
+				}
+				list = append(list, h)
+			}
+			return list, nil
+		})
+
+		overlay := torchwood.NewHashReaderOverlay(n0, r)
+		for n := n0; n < N; n++ {
+			if overlay.Size() != n {
+				t.Fatalf("n0=%d: Size() = %d, want %d", n0, overlay.Size(), n)
+			}
+			fatalIfErr(t, overlay.AppendRecordHash(leaves[n]))
+
+			got, err := tlog.TreeHash(n+1, overlay)
+			fatalIfErr(t, err)
+			want, err := tlog.TreeHash(n+1, full)
+			fatalIfErr(t, err)
+			if got != want {
+				t.Fatalf("n0=%d: TreeHash(%d) = %v, want %v", n0, n+1, got, want)
+			}
+		}
+
+		th, err := tlog.TreeHash(N, overlay)
+		fatalIfErr(t, err)
+		for i := n0; i < N; i++ {
+			proof, err := tlog.ProveRecord(N, i, overlay)
+			fatalIfErr(t, err)
+			if err := tlog.CheckRecord(proof, N, th, i, leaves[i]); err != nil {
+				t.Fatalf("n0=%d: CheckRecord(%d): %v", n0, i, err)
+			}
+		}
+
+		for start := int64(0); start < N; start++ {
+			// Subtrees wholly within the original tree may not be on its
+			// right edge, so only check those extending past it.
+			for end := max(start+1, n0+1); end <= N; end++ {
+				if !torchwood.ValidSubtree(start, end) {
+					continue
+				}
+				got, err := torchwood.SubtreeHash(start, end, overlay)
+				fatalIfErr(t, err)
+				want, err := torchwood.SubtreeHash(start, end, full)
+				fatalIfErr(t, err)
+				if got != want {
+					t.Fatalf("n0=%d: SubtreeHash(%d, %d) = %v, want %v", n0, start, end, got, want)
+				}
+			}
+		}
+	}
+}
+
+// TestHashReaderOverlayFullTree checks subtree hashes and tile contents
+// against a reader for the whole tree, with the underlying reader restricted
+// to the original tree.
+func TestHashReaderOverlayFullTree(t *testing.T) {
+	const N = 130
+	leaves, full := buildTree(t, N)
+	for n0 := int64(0); n0 <= N; n0++ {
+		base := tlog.StoredHashCount(n0)
+		r := tlog.HashReaderFunc(func(indexes []int64) ([]tlog.Hash, error) {
+			for _, id := range indexes {
+				if id >= base {
+					return nil, fmt.Errorf("index %d beyond the original tree of size %d", id, n0)
+				}
+			}
+			return full.ReadHashes(indexes)
+		})
+
+		overlay := torchwood.NewHashReaderOverlay(n0, r)
+		for n := n0; n < N; n++ {
+			fatalIfErr(t, overlay.AppendRecordHash(leaves[n]))
+		}
+
+		for start := int64(0); start < N; start++ {
+			for end := start + 1; end <= N; end++ {
+				if !torchwood.ValidSubtree(start, end) {
+					continue
+				}
+				got, err := torchwood.SubtreeHash(start, end, overlay)
+				fatalIfErr(t, err)
+				want, err := torchwood.SubtreeHash(start, end, full)
+				fatalIfErr(t, err)
+				if got != want {
+					t.Fatalf("n0=%d: SubtreeHash(%d, %d) = %v, want %v", n0, start, end, got, want)
+				}
+			}
+		}
+
+		for _, tile := range tlog.NewTiles(2, n0, N) {
+			got, err := tlog.ReadTileData(tile, overlay)
+			fatalIfErr(t, err)
+			want, err := tlog.ReadTileData(tile, full)
+			fatalIfErr(t, err)
+			if !slices.Equal(got, want) {
+				t.Fatalf("n0=%d: ReadTileData(%v) = %x, want %x", n0, tile, got, want)
+			}
+		}
+	}
+}
+
+func TestHashReaderOverlayBatchedReads(t *testing.T) {
+	const n0 = 100
+	leaves, full := buildTree(t, 130)
+	var reads int
+	r := tlog.HashReaderFunc(func(indexes []int64) ([]tlog.Hash, error) {
+		reads++
+		return full.ReadHashes(indexes)
+	})
+	overlay := torchwood.NewHashReaderOverlay(n0, r)
+	for n := int64(n0); n < 130; n++ {
+		fatalIfErr(t, overlay.AppendRecordHash(leaves[n]))
+	}
+
+	// Interleave original tree and overlay indexes.
+	base := tlog.StoredHashCount(n0)
+	indexes := []int64{0, base, 1, base + 1, 2, base + 2, 3}
+	reads = 0
+	got, err := overlay.ReadHashes(indexes)
+	fatalIfErr(t, err)
+	if reads != 1 {
+		t.Errorf("ReadHashes made %d underlying reads, want 1", reads)
+	}
+	want, err := full.ReadHashes(indexes)
+	fatalIfErr(t, err)
+	if !slices.Equal(got, want) {
+		t.Errorf("ReadHashes(%v) = %v, want %v", indexes, got, want)
+	}
+}
+
+func TestHashReaderOverlayErrors(t *testing.T) {
+	leaves, full := buildTree(t, 10)
+
+	// Requests beyond the extended tree are rejected.
+	overlay := torchwood.NewHashReaderOverlay(10, full)
+	if _, err := overlay.ReadHashes([]int64{tlog.StoredHashCount(10)}); err == nil {
+		t.Errorf("expected error for index beyond the tree size")
+	}
+	if _, err := overlay.ReadHashes([]int64{-1}); err == nil {
+		t.Errorf("expected error for negative index")
+	}
+
+	// A nil underlying reader works from size 0...
+	overlay = torchwood.NewHashReaderOverlay(0, nil)
+	for n := int64(0); n < 10; n++ {
+		fatalIfErr(t, overlay.AppendRecordHash(leaves[n]))
+	}
+	got, err := tlog.TreeHash(10, overlay)
+	fatalIfErr(t, err)
+	want, err := tlog.TreeHash(10, full)
+	fatalIfErr(t, err)
+	if got != want {
+		t.Errorf("TreeHash(10) = %v, want %v", got, want)
+	}
+
+	// ...but fails requests for a non-empty original tree.
+	overlay = torchwood.NewHashReaderOverlay(10, nil)
+	if _, err := tlog.TreeHash(10, overlay); err == nil {
+		t.Errorf("expected error for nil underlying reader")
+	}
+	// Record 10 completes no subtrees, so it doesn't need the original tree.
+	fatalIfErr(t, overlay.AppendRecordHash(leaves[0]))
+	// Record 11 completes the subtree [10, 12) and needs the hash of [8, 10).
+	if err := overlay.AppendRecordHash(leaves[1]); err == nil {
+		t.Errorf("expected error appending a record that completes a subtree")
+	}
+	if overlay.Size() != 11 {
+		t.Errorf("Size() = %d, want 11 after the failed append", overlay.Size())
+	}
+}