niko_trust/internal/smt/proofs.go

package smt

// Proof encodings. All multi-byte integers are uvarints; bit strings are
// packed MSB-first; hashes are 32 raw bytes.
//
//	0x01 || nSteps || step…                                   inclusion
//	0x02 || nSteps || step… || witnessKey                     absence, leaf witness
//	0x03 || nSteps || step… || m || prefix || left || right   absence, branch witness
//	0x04                                                      absence in an empty trie
//
//	step := nPrefixBits || prefixBits || siblingSide(1 byte) || siblingHash
//
// A step describes one branch on the path from the root toward the subject
// key: prefix is that branch's matched bits (committed inside its hash),
// siblingSide is the bit selecting the sibling subtree at the branch's
// split, and siblingHash is that subtree's commitment.
//
// Verification recomposes the root from the subject side upward. Because
// every step's prefix must also match the subject key's own bits at the
// tracked depth, a proof constructed for one key verifies for no other.

import (
	"encoding/binary"
	"errors"
)

type proofStep struct {
	prefix    []byte
	prefixLen int
	sibSide   byte
	sibling   [32]byte
}

const (
	proofInclusion    = 0x01
	proofAbsentLeaf   = 0x02
	proofAbsentBranch = 0x03
	proofAbsentEmpty  = 0x04
)

// ---------------------------------------------------------------- generation

// walkDescend follows key's bits through the trie. It returns the steps
// crossed (top-down) and where the walk stopped: at a leaf, or diverging
// inside a branch's prefix.
func (t *Trie) walkDescend(key *[KeySize]byte) (steps []proofStep, leaf *node, branch *node, depth int) {
	n := t.root
	d := 0
	for n != nil && !n.leaf {
		i := 0
		for i < n.prefixLen && bit(key, d+i) == prefixBit(n.prefix, i) {
			i++
		}
		if i < n.prefixLen {
			return steps, nil, n, d
		}
		sd := d + n.prefixLen
		b := bit(key, sd)
		child := n.left
		sib := n.right
		sibSide := byte(1)
		if b == 1 {
			child = n.right
			sib = n.left
			sibSide = 0
		}
		steps = append(steps, proofStep{
			prefix:    append([]byte(nil), n.prefix...),
			prefixLen: n.prefixLen,
			sibSide:   sibSide,
			sibling:   t.hash(sib),
		})
		n = child
		d = sd + 1
	}
	if n != nil {
		return steps, n, nil, d
	}
	return steps, nil, nil, d // empty trie
}

// InclusionProof returns a proof that key is in the trie.
func (t *Trie) InclusionProof(key [KeySize]byte) ([]byte, error) {
	steps, leaf, _, _ := t.walkDescend(&key)
	if t.root == nil || leaf == nil || leaf.key != key {
		return nil, ErrAbsent
	}
	out := []byte{proofInclusion}
	out = appendUvarint(out, len(steps))
	out = appendSteps(out, steps)
	return out, nil
}

// AbsenceProof returns a proof that key is not in the trie.
func (t *Trie) AbsenceProof(key [KeySize]byte) ([]byte, error) {
	if t.root == nil {
		return []byte{proofAbsentEmpty}, nil
	}
	steps, leaf, branch, _ := t.walkDescend(&key)
	switch {
	case leaf != nil && leaf.key == key:
		return nil, ErrPresent

	case leaf != nil:
		out := []byte{proofAbsentLeaf}
		out = appendUvarint(out, len(steps))
		out = appendSteps(out, steps)
		out = append(out, leaf.key[:]...)
		return out, nil

	default:
		// Diverged inside branch.prefix at some bit; the branch itself is
		// the witness and stays intact in the tree.
		out := []byte{proofAbsentBranch}
		out = appendUvarint(out, len(steps))
		out = appendSteps(out, steps)
		out = appendUvarint(out, branch.prefixLen)
		out = append(out, branch.prefix...)
		l := t.hash(branch.left)
		r := t.hash(branch.right)
		out = append(out, l[:]...)
		out = append(out, r[:]...)
		return out, nil
	}
}

func appendSteps(out []byte, steps []proofStep) []byte {
	for _, st := range steps {
		out = appendUvarint(out, st.prefixLen)
		out = append(out, st.prefix...)
		out = append(out, st.sibSide)
		out = append(out, st.sibling[:]...)
	}
	return out
}

// ----------------------------------------------------------------- verifying

type parsedStep struct {
	prefix    []byte
	prefixLen int
	sibSide   byte
	sibling   [32]byte
	split     int // absolute bit index of the branch's split decision
}

type proofParser struct {
	b   []byte
	pos int
}

func (p *proofParser) byteVal() (byte, error) {
	if p.pos >= len(p.b) {
		return 0, ErrBadProof
	}
	v := p.b[p.pos]
	p.pos++
	return v, nil
}

func (p *proofParser) uvarint() (int, error) {
	v, n := binary.Uvarint(p.b[p.pos:])
	// Every uvarint in a proof is either a step count or a bit length; both
	// are bounded by one key's worth of bits.
	if n <= 0 || v > KeySize*8 {
		return 0, ErrBadProof
	}
	p.pos += n
	return int(v), nil
}

func (p *proofParser) bytes(n int) ([]byte, error) {
	if n < 0 || p.pos+n > len(p.b) {
		return nil, ErrBadProof
	}
	out := p.b[p.pos : p.pos+n]
	p.pos += n
	return out, nil
}

var errTrailing = errors.New("smt: trailing bytes after proof")

func (p *proofParser) done() error {
	if p.pos != len(p.b) {
		return errTrailing
	}
	return nil
}

// parseSteps reads n steps, validating each prefix against key at the
// running depth, and returns them with their split positions.
func parseSteps(p *proofParser, n int, key *[KeySize]byte) ([]parsedStep, error) {
	d := 0
	steps := make([]parsedStep, 0, n)
	for i := 0; i < n; i++ {
		plen, err := p.uvarint()
		if err != nil {
			return nil, err
		}
		var nb int
		if plen > 0 {
			nb = (plen + 7) / 8
		}
		pref, err := p.bytes(nb)
		if err != nil {
			return nil, err
		}
		side, err := p.byteVal()
		if err != nil || side > 1 {
			return nil, ErrBadProof
		}
		var sib [32]byte
		sb, err := p.bytes(32)
		if err != nil {
			return nil, err
		}
		copy(sib[:], sb)

		// The prefix must be exactly the key's bits here: the walk that
		// produced this step followed the key, so any other claim is a
		// malformed proof even before hashes are checked.
		for j := 0; j < plen; j++ {
			if prefixBit(pref, j) != bit(key, d+j) {
				return nil, ErrBadProof
			}
		}
		steps = append(steps, parsedStep{
			prefix:    append([]byte(nil), pref...),
			prefixLen: plen,
			sibSide:   side,
			sibling:   sib,
			split:     d + plen,
		})
		d = d + plen + 1
	}
	return steps, nil
}

// combine folds the subject hash upward through the steps, deepest first.
func combine(h *[32]byte, steps []parsedStep) {
	for i := len(steps) - 1; i >= 0; i-- {
		st := steps[i]
		b := byte(1) - st.sibSide // the side the subject subtree occupies
		var l, r [32]byte
		if b == 0 {
			l, r = *h, st.sibling
		} else {
			l, r = st.sibling, *h
		}
		*h = branchHash(st.prefix, st.prefixLen, &l, &r)
	}
}

// VerifyInclusion reports whether proof attests key's membership in the set
// committed to by root.
func VerifyInclusion(root [32]byte, key [KeySize]byte, proof []byte) bool {
	p := &proofParser{b: proof}
	typ, err := p.byteVal()
	if err != nil || typ != proofInclusion {
		return false
	}
	n, err := p.uvarint()
	if err != nil {
		return false
	}
	steps, err := parseSteps(p, n, &key)
	if err != nil {
		return false
	}
	if err := p.done(); err != nil {
		return false
	}
	h := leafHash(&key)
	combine(&h, steps)
	return h == root
}

// VerifyAbsence reports whether proof attests key's absence from the set
// committed to by root.
func VerifyAbsence(root [32]byte, key [KeySize]byte, proof []byte) bool {
	p := &proofParser{b: proof}
	typ, err := p.byteVal()
	if err != nil {
		return false
	}

	switch typ {
	case proofAbsentEmpty:
		return root == EmptyRoot

	case proofAbsentLeaf:
		n, err := p.uvarint()
		if err != nil {
			return false
		}
		steps, err := parseSteps(p, n, &key)
		if err != nil {
			return false
		}
		wk, err := p.bytes(KeySize)
		if err != nil {
			return false
		}
		if err := p.done(); err != nil {
			return false
		}
		var witness [KeySize]byte
		copy(witness[:], wk)
		if witness == key {
			return false // that would be proof of presence
		}
		h := leafHash(&witness)
		combine(&h, steps)
		return h == root

	case proofAbsentBranch:
		n, err := p.uvarint()
		if err != nil {
			return false
		}
		steps, err := parseSteps(p, n, &key)
		if err != nil {
			return false
		}
		m, err := p.uvarint()
		if err != nil || m == 0 || m > KeySize*8 {
			return false
		}
		pref, err := p.bytes((m + 7) / 8)
		if err != nil {
			return false
		}
		// The key must leave this prefix mid-way: matching fully would mean
		// the walk continued deeper, differing nowhere means presence of
		// nothing in particular.
		j := -1
		db := 0 // absolute bit index where the witness branch's prefix starts
		if len(steps) > 0 {
			db = steps[len(steps)-1].split + 1
		}
		for i := 0; i < m; i++ {
			if prefixBit(pref, i) != bit(&key, db+i) {
				j = i
				break
			}
		}
		if j < 0 {
			return false
		}
		var ll, lr [32]byte
		lb, err := p.bytes(32)
		if err != nil {
			return false
		}
		copy(ll[:], lb)
		rb, err := p.bytes(32)
		if err != nil {
			return false
		}
		copy(lr[:], rb)
		if err := p.done(); err != nil {
			return false
		}

		hb := branchHash(pref, m, &ll, &lr)
		// The branch hangs off the last step's split; fold it in place of a
		// leaf hash, then continue through remaining ancestor steps.
		if len(steps) == 0 {
			return hb == root
		}
		top := steps[len(steps)-1]
		b := byte(1) - top.sibSide
		var l, r [32]byte
		if b == 0 {
			l, r = hb, top.sibling
		} else {
			l, r = top.sibling, hb
		}
		h := branchHash(top.prefix, top.prefixLen, &l, &r)
		rest := steps[:len(steps)-1]
		combine(&h, rest)
		return h == root
	}
	return false
}