torchwood: add SubtreeHash, ValidSubtree, CoverInterval, and subtree proofs

Author: FiloSottile

subtree.go

@@ -0,0 +1,290 @@
+package torchwood
+
+import (
+	"fmt"
+	"math/bits"
+
+	"golang.org/x/mod/sumdb/tlog"
+)
+
+// A SubtreeProof is a verifiable proof that a particular tree head contains a
+// particular subtree. A [tlog.TreeProof] is a special case of a SubtreeProof
+// where the subtree has start index 0. A [tlog.RecordProof] is a special case
+// of a SubtreeProof where the subtree has size 1.
+//
+// draft-ietf-plants-merkle-tree-certs-04 calls this a "Subtree Consistency
+// Proof".
+type SubtreeProof []tlog.Hash
+
+// ProveSubtree returns the proof that the tree of size t contains the subtree
+// [start, end), where start and end are record indexes.
+func ProveSubtree(t, start, end int64, r tlog.HashReader) (SubtreeProof, error) {
+	if t < 0 || end > t || !ValidSubtree(start, end) {
+		return nil, fmt.Errorf("tlog: invalid inputs in ProveSubtree")
+	}
+	indexes := subtreeProofIndex(0, t, start, end, true, nil)
+	if len(indexes) == 0 {
+		return SubtreeProof{}, nil
+	}
+	hashes, err := r.ReadHashes(indexes)
+	if err != nil {
+		return nil, err
+	}
+	if len(hashes) != len(indexes) {
+		return nil, fmt.Errorf("tlog: ReadHashes(%d indexes) = %d hashes", len(indexes), len(hashes))
+	}
+
+	p, hashes := subtreeProof(0, t, start, end, true, hashes)
+	if len(hashes) != 0 {
+		panic("tlog: bad index math in ProveSubtree")
+	}
+	return p, nil
+}
+
+// subtreeProofIndex builds the list of indexes needed to construct the proof
+// that the subtree [start, end) is contained in the node with leaves [lo, hi).
+// It appends those indexes to need and returns the result. See
+// draft-ietf-plants-merkle-tree-certs-04, Section 4.4. b reports whether the
+// verifier already knows the hash of the subtree portion in [lo, hi); it starts
+// true and becomes false past the first straddled node.
+func subtreeProofIndex(lo, hi, start, end int64, b bool, need []int64) []int64 {
+	if !(lo <= start && start < end && end <= hi) {
+		panic("tlog: bad math in subtreeProofIndex")
+	}
+	if lo == start && hi == end {
+		if b {
+			return need
+		}
+		return subTreeIndex(lo, hi, need)
+	}
+	k, _ := maxpow2(hi - lo)
+	switch {
+	case end <= lo+k: // subtree in the left child
+		need = subtreeProofIndex(lo, lo+k, start, end, b, need)
+		need = subTreeIndex(lo+k, hi, need)
+	case lo+k <= start: // subtree in the right child
+		need = subTreeIndex(lo, lo+k, need)
+		need = subtreeProofIndex(lo+k, hi, start, end, b, need)
+	default: // subtree straddles the split, which implies start == lo
+		if start != lo {
+			panic("tlog: bad math in subtreeProofIndex")
+		}
+		need = subtreeProofIndex(lo+k, hi, lo+k, end, false, need)
+		need = subTreeIndex(lo, lo+k, need)
+	}
+	return need
+}
+
+// subtreeProof constructs the proof that the subtree [start, end) is contained
+// in the node with leaves [lo, hi). It returns any leftover hashes as well.
+func subtreeProof(lo, hi, start, end int64, b bool, hashes []tlog.Hash) (SubtreeProof, []tlog.Hash) {
+	if !(lo <= start && start < end && end <= hi) {
+		panic("tlog: bad math in subtreeProof")
+	}
+	if lo == start && hi == end {
+		if b {
+			// The verifier knows the subtree hash, so we don't need to send it.
+			return SubtreeProof{}, hashes
+		}
+		th, hashes := subTreeHash(lo, hi, hashes)
+		return SubtreeProof{th}, hashes
+	}
+	k, _ := maxpow2(hi - lo)
+	var p SubtreeProof
+	var th tlog.Hash
+	switch {
+	case end <= lo+k: // subtree in the left child
+		p, hashes = subtreeProof(lo, lo+k, start, end, b, hashes)
+		th, hashes = subTreeHash(lo+k, hi, hashes)
+	case lo+k <= start: // subtree in the right child
+		th, hashes = subTreeHash(lo, lo+k, hashes)
+		p, hashes = subtreeProof(lo+k, hi, start, end, b, hashes)
+	default: // subtree straddles the split, which implies start == lo
+		if start != lo {
+			panic("tlog: bad math in subtreeProof")
+		}
+		p, hashes = subtreeProof(lo+k, hi, lo+k, end, false, hashes)
+		th, hashes = subTreeHash(lo, lo+k, hashes)
+	}
+	return append(p, th), hashes
+}
+
+// CheckSubtree verifies that p is a valid proof that the tree of size t
+// with hash th contains the subtree [start, end) with hash sh.
+func CheckSubtree(p SubtreeProof, t int64, th tlog.Hash, start, end int64, sh tlog.Hash) error {
+	if t < 0 || end > t || !ValidSubtree(start, end) {
+		return fmt.Errorf("tlog: invalid inputs in CheckSubtree")
+	}
+	sh2, th2, err := runSubtreeProof(p, 0, t, start, end, true, sh)
+	if err != nil {
+		return err
+	}
+	if sh2 == sh && th2 == th {
+		return nil
+	}
+	return errProofFailed
+}
+
+// runSubtreeProof runs the proof p that the subtree [start, end) is contained in
+// the node with leaves [lo, hi). Running the proof means constructing and
+// returning the implied hashes of both the subtree and the node.
+func runSubtreeProof(p SubtreeProof, lo, hi, start, end int64, b bool, sh tlog.Hash) (subtreeHash, nodeHash tlog.Hash, err error) {
+	if !(lo <= start && start < end && end <= hi) {
+		panic("tlog: bad math in runSubtreeProof")
+	}
+	if lo == start && hi == end {
+		if b {
+			// The verifier knows the subtree hash, and it is the node hash.
+			if len(p) != 0 {
+				return tlog.Hash{}, tlog.Hash{}, errProofFailed
+			}
+			return sh, sh, nil
+		}
+		if len(p) != 1 {
+			return tlog.Hash{}, tlog.Hash{}, errProofFailed
+		}
+		return p[0], p[0], nil
+	}
+	if len(p) == 0 {
+		return tlog.Hash{}, tlog.Hash{}, errProofFailed
+	}
+	k, _ := maxpow2(hi - lo)
+	switch {
+	case end <= lo+k: // subtree in the left child
+		sh2, nh, err := runSubtreeProof(p[:len(p)-1], lo, lo+k, start, end, b, sh)
+		if err != nil {
+			return tlog.Hash{}, tlog.Hash{}, err
+		}
+		return sh2, tlog.NodeHash(nh, p[len(p)-1]), nil
+	case lo+k <= start: // subtree in the right child
+		sh2, nh, err := runSubtreeProof(p[:len(p)-1], lo+k, hi, start, end, b, sh)
+		if err != nil {
+			return tlog.Hash{}, tlog.Hash{}, err
+		}
+		return sh2, tlog.NodeHash(p[len(p)-1], nh), nil
+	default: // subtree straddles the split, which implies start == lo
+		if start != lo {
+			panic("tlog: bad math in runSubtreeProof")
+		}
+		left := p[len(p)-1]
+		sh2, nh, err := runSubtreeProof(p[:len(p)-1], lo+k, hi, lo+k, end, false, sh)
+		if err != nil {
+			return tlog.Hash{}, tlog.Hash{}, err
+		}
+		return tlog.NodeHash(left, sh2), tlog.NodeHash(left, nh), nil
+	}
+}
+
+// A RecordInSubtreeProof is a verifiable proof that a particular subtree
+// contains a particular record. A [tlog.RecordProof] is a special case of a
+// RecordInSubtreeProof where the subtree has start index 0.
+//
+// draft-ietf-plants-merkle-tree-certs-04 calls this a "Subtree Inclusion Proof".
+type RecordInSubtreeProof []tlog.Hash
+
+// ProveRecordInSubtree returns the proof that the subtree [start, end) contains
+// the record with index n.
+func ProveRecordInSubtree(start, end, n int64, r tlog.HashReader) (RecordInSubtreeProof, error) {
+	if !ValidSubtree(start, end) || n < start || n >= end {
+		return nil, fmt.Errorf("tlog: invalid inputs in ProveRecordInSubtree")
+	}
+	indexes := leafProofIndex(start, end, n, nil)
+	if len(indexes) == 0 {
+		return RecordInSubtreeProof{}, nil
+	}
+	hashes, err := r.ReadHashes(indexes)
+	if err != nil {
+		return nil, err
+	}
+	if len(hashes) != len(indexes) {
+		return nil, fmt.Errorf("tlog: ReadHashes(%d indexes) = %d hashes", len(indexes), len(hashes))
+	}
+
+	p, hashes := leafProof(start, end, n, hashes)
+	if len(hashes) != 0 {
+		panic("tlog: bad index math in ProveRecordInSubtree")
+	}
+	return p, nil
+}
+
+// CheckRecordInSubtree verifies that p is a valid proof that the subtree
+// [start, end) with hash sh contains the record with index n and hash h.
+func CheckRecordInSubtree(p RecordInSubtreeProof, start, end int64, sh tlog.Hash, n int64, h tlog.Hash) error {
+	if !ValidSubtree(start, end) || n < start || n >= end {
+		return fmt.Errorf("tlog: invalid inputs in CheckRecordInSubtree")
+	}
+	sh2, err := runRecordProof(p, start, end, n, h)
+	if err != nil {
+		return err
+	}
+	if sh2 == sh {
+		return nil
+	}
+	return errProofFailed
+}
+
+// SubtreeHash computes the hash for the subtree [start, end) using the
+// HashReader to obtain previously stored hashes. SubtreeHash makes a single
+// call to ReadHash requesting at most 1 + log₂(end - start) hashes.
+//
+// A [tlog.TreeHash] is a special case of a SubtreeHash where start is 0.
+//
+// See draft-ietf-plants-merkle-tree-certs-04, Section 4.
+func SubtreeHash(start, end int64, r tlog.HashReader) (tlog.Hash, error) {
+	if !ValidSubtree(start, end) {
+		return tlog.Hash{}, fmt.Errorf("tlog: invalid inputs in SubtreeHash")
+	}
+	indexes := subTreeIndex(start, end, nil)
+	hashes, err := r.ReadHashes(indexes)
+	if err != nil {
+		return tlog.Hash{}, err
+	}
+	if len(hashes) != len(indexes) {
+		return tlog.Hash{}, fmt.Errorf("tlog: ReadHashes(%d indexes) = %d hashes", len(indexes), len(hashes))
+	}
+	hash, hashes := subTreeHash(start, end, hashes)
+	if len(hashes) != 0 {
+		panic("tlog: bad index math in SubtreeHash")
+	}
+	return hash, nil
+}
+
+// ValidSubtree reports whether [start, end) is a valid subtree.
+//
+// See draft-ietf-plants-merkle-tree-certs-04, Section 4.
+func ValidSubtree(start, end int64) bool {
+	if start < 0 || end <= start {
+		return false
+	}
+	return start&(bitCeil(end-start)-1) == 0
+}
+
+// CoverInterval returns leftStart and mid for the two subtrees [leftStart, mid)
+// and [mid, end) that cover the interval [start, end) as efficiently as
+// possible.
+//
+// The subtrees are adjacent and the second ends at end, but the first may begin
+// before start. See draft-ietf-plants-merkle-tree-certs-04, Section 4.5.
+//
+// It is an error if start < 0, end <= start, or [start, end) is already a
+// subtree.
+func CoverInterval(start, end int64) (leftStart, mid int64, err error) {
+	if start < 0 || end <= start {
+		return 0, 0, fmt.Errorf("tlog: invalid interval in CoverInterval")
+	}
+	if ValidSubtree(start, end) {
+		return 0, 0, fmt.Errorf("tlog: interval is already a subtree in CoverInterval")
+	}
+	last := end - 1
+	split := bits.Len64(uint64(start^last)) - 1
+	mask := int64(1)<<split - 1
+	mid = last &^ mask
+	leftSplit := bits.Len64(uint64(^start & mask))
+	leftStart = start &^ (int64(1)<<leftSplit - 1)
+	return leftStart, mid, nil
+}
+
+// bitCeil returns the smallest power of two not smaller than n, for n >= 1.
+func bitCeil(n int64) int64 {
+	return 1 << bits.Len64(uint64(n-1))
+}