zkvm: rebuild the UTXO set from the records, and stop writing under a read lock

loadUTXOs was `return nil`. The in-memory set was never rebuilt, so a restarted node came back with an empty one: it accepted a commitment it already held — reaching the opposite verdict from a running node on the same block — and reported no commitments and no UTXOs at any height. The count came from a separately persisted key, a second copy of the truth that a crash between the two writes leaves disagreeing with the records, and one removal from there drives an unsigned counter below zero. That is the same defect just removed from the nullifier count, so the UTXO set now works the same way: the records are the set, and the count is read off it. Blocks live in the same database under their raw ids, and one id in 256 begins with the UTXO prefix. The loader now counts a record as a UTXO only when it names the commitment its own key is made of; a block frame parses as a UTXO otherwise, since the accessors are bounds-safe and return zeros, and a garbage commitment would join the set. Both read paths memoized what they loaded while holding only a read lock. The memo is deleted rather than the lock upgraded: the map is loaded with the whole set at startup and maintained by the writers, so the disk read already returns what the memo would have, and a write lock on a path consensus and RPC share would serialize every reader to save a lookup. A nullifier value that is not eight bytes is now a miss rather than a panic in a consensus path; the loader already skipped those. Co-authored-by: Hanzo Dev <[email protected]>
2026-09-24 00:26:00 +00:00 · 2026-08-25 12:30:37 -07:00
parent d924766854
commit 318794d558
4 changed files with 282 additions and 104 deletions
@@ -104,7 +104,15 @@ func (ndb *NullifierDB) IsNullifierSpent(nullifier []byte) bool {
 	return err == nil
 }

-// GetNullifierHeight returns the height when a nullifier was spent
+// GetNullifierHeight returns the height when a nullifier was spent.
+//
+// A miss falls through to the records and returns what it finds without
+// memoising it. Memoising here would be a write on a path that holds only the
+// read lock, and a read lock promises every other reader that nothing is
+// changing: two callers missing at once would write the same map at the same
+// time, which is a runtime throw, not a returned error. The write lock is not
+// the answer either — it would serialise every reader of a path consensus and
+// RPC both sit on, to save a lookup the set already answers.
 func (ndb *NullifierDB) GetNullifierHeight(nullifier []byte) (uint64, error) {
 	ndb.mu.RLock()
 	defer ndb.mu.RUnlock()
@@ -119,16 +127,11 @@ func (ndb *NullifierDB) GetNullifierHeight(nullifier []byte) (uint64, error) {
 	// Load from database
 	key := makeNullifierKey(nullifier)
 	heightBytes, err := ndb.db.Get(key)
-	if err != nil {
+	if err != nil || len(heightBytes) != 8 {
 		return 0, errors.New("nullifier not found")
 	}

-	height := binary.BigEndian.Uint64(heightBytes)
-
-	// Update cache
-	ndb.nullifierCache[nullifierStr] = height
-
-	return height, nil
+	return binary.BigEndian.Uint64(heightBytes), nil
 }

 // GetNullifiersByHeight returns all nullifiers spent at a specific height
@@ -4,6 +4,7 @@
 package zkvm

 import (
+	"sync"
 	"testing"

 	"github.com/luxfi/database/memdb"
@@ -55,3 +56,54 @@ func TestSpentCountFollowsTheSet(t *testing.T) {
 	require.Error(t, ndb.RemoveNullifier([]byte("never")))
 	require.Equal(t, uint64(1), ndb.GetNullifierCount())
 }
+
+// TestNullifierReadDoesNotWriteTheSet. GetNullifierHeight used to memoise what
+// it loaded while holding only the read lock, and a read lock promises every
+// other reader that nothing is changing — so two callers missing at once wrote
+// the same map at the same time, which is a runtime throw that takes the
+// process down rather than an error anyone can handle. Consensus and RPC both
+// sit on this path.
+func TestNullifierReadDoesNotWriteTheSet(t *testing.T) {
+	ndb, err := NewNullifierDB(memdb.New(), log.NoLog{})
+	require.NoError(t, err)
+	require.NoError(t, ndb.MarkNullifierSpent([]byte("n"), 4))
+
+	// A record on disk that the in-memory set does not know about is what a
+	// memoising read would fill in.
+	ndb.nullifierCache = map[string]uint64{}
+
+	var readers sync.WaitGroup
+	for i := 0; i < 8; i++ {
+		readers.Add(1)
+		go func() {
+			defer readers.Done()
+			height, err := ndb.GetNullifierHeight([]byte("n"))
+			switch {
+			case err != nil:
+				t.Errorf("read failed: %v", err)
+			case height != 4:
+				t.Errorf("read height %d, want 4", height)
+			}
+		}()
+	}
+	readers.Wait()
+
+	require.Empty(t, ndb.nullifierCache, "a read must not change the set it reads")
+}
+
+// TestShortRecordIsNotAHeight. Blocks are stored in the same database under
+// their raw id, so the nullifier prefix can turn up over a value that is not a
+// height. loadNullifiers already passes over anything that is not eight bytes;
+// the read path has to agree with it, or reading one of those is a panic where
+// a miss was meant.
+func TestShortRecordIsNotAHeight(t *testing.T) {
+	db := memdb.New()
+	require.NoError(t, db.Put(makeNullifierKey([]byte("n")), []byte{1}))
+
+	ndb, err := NewNullifierDB(db, log.NoLog{})
+	require.NoError(t, err)
+	require.Zero(t, ndb.GetNullifierCount())
+
+	_, err = ndb.GetNullifierHeight([]byte("n"))
+	require.ErrorContains(t, err, "not found")
+}
@@ -4,7 +4,7 @@
 package zkvm

 import (
-	"encoding/binary"
+	"bytes"
 	"errors"
 	"sync"

@@ -14,12 +14,8 @@ import (
 	"github.com/luxfi/ids"
 )

-const (
-	// Database prefixes
-	utxoPrefix   = 0x10
-	utxoCountKey = "utxo_count"
-	utxoIndexKey = "utxo_index"
-)
+// Database prefix for UTXO records: utxoPrefix || commitment -> marshalled UTXO.
+const utxoPrefix = 0x10

 // UTXO represents an unspent transaction output
 type UTXO struct {
@@ -36,9 +32,12 @@ type UTXODB struct {
 	db  database.Database
 	log log.Logger

-	// Caches
-	utxoCache map[string]*UTXO // commitment -> UTXO
-	utxoCount uint64
+	// The set itself: commitment -> height when created. It is rebuilt from the
+	// records at startup, so a restarted node knows the same commitments a
+	// running one does — which is what Block.Accept consults to refuse a
+	// duplicate. Membership and the count are read off it; the bodies stay in
+	// the records.
+	utxoCache map[string]uint64

 	// Indexes
 	heightIndex map[uint64][]string // height -> commitments
@@ -51,20 +50,10 @@ func NewUTXODB(db database.Database, log log.Logger) (*UTXODB, error) {
 	udb := &UTXODB{
 		db:          db,
 		log:         log,
-		utxoCache:   make(map[string]*UTXO),
+		utxoCache:   make(map[string]uint64),
 		heightIndex: make(map[uint64][]string),
 	}

-	// Load UTXO count
-	countBytes, err := db.Get([]byte(utxoCountKey))
-	if err == database.ErrNotFound {
-		udb.utxoCount = 0
-	} else if err != nil {
-		return nil, err
-	} else {
-		udb.utxoCount = binary.BigEndian.Uint64(countBytes)
-	}
-
 	if err := udb.loadUTXOs(); err != nil {
 		return nil, err
 	}
@@ -97,20 +86,10 @@ func (udb *UTXODB) AddUTXO(utxo *UTXO) error {
 		return err
 	}

-	// Update cache
-	udb.utxoCache[commitmentStr] = utxo
-
-	// Update height index
+	// Update the set and the height index
+	udb.utxoCache[commitmentStr] = utxo.Height
 	udb.heightIndex[utxo.Height] = append(udb.heightIndex[utxo.Height], commitmentStr)

-	// Update count
-	udb.utxoCount++
-	countBytes := make([]byte, 8)
-	binary.BigEndian.PutUint64(countBytes, udb.utxoCount)
-	if err := udb.db.Put([]byte(utxoCountKey), countBytes); err != nil {
-		return err
-	}
-
 	udb.log.Debug("Added UTXO",
 		log.String("txID", utxo.TxID.String()),
 		log.Uint32("outputIndex", utxo.OutputIndex),
@@ -120,34 +99,19 @@ func (udb *UTXODB) AddUTXO(utxo *UTXO) error {
 	return nil
 }

-// GetUTXO retrieves a UTXO by commitment
+// GetUTXO retrieves a UTXO by commitment.
+//
+// The body comes from the records every time. Memoising it here would be a
+// write on a path that holds only the read lock, and a read lock promises every
+// other reader that nothing is changing: two RPC clients asking for different
+// commitments would write the same map at the same time, which is a runtime
+// throw, not a returned error. The lock is still held so that a read cannot
+// land between the record delete and the set delete a removal does together.
 func (udb *UTXODB) GetUTXO(commitment []byte) (*UTXO, error) {
 	udb.mu.RLock()
 	defer udb.mu.RUnlock()

-	commitmentStr := string(commitment)
-
-	// Check cache
-	if utxo, exists := udb.utxoCache[commitmentStr]; exists {
-		return utxo, nil
-	}
-
-	// Load from database
-	key := makeUTXOKey(commitment)
-	utxoBytes, err := udb.db.Get(key)
-	if err != nil {
-		return nil, errors.New("UTXO not found")
-	}
-
-	var utxo UTXO
-	if err := parseUTXO(utxoBytes, &utxo); err != nil {
-		return nil, err
-	}
-
-	// Update cache
-	udb.utxoCache[commitmentStr] = &utxo
-
-	return &utxo, nil
+	return udb.read(commitment)
 }

 // RemoveUTXO removes a UTXO from the set
@@ -157,15 +121,10 @@ func (udb *UTXODB) RemoveUTXO(commitment []byte) error {

 	commitmentStr := string(commitment)

-	// Get UTXO to find height
-	utxo, exists := udb.utxoCache[commitmentStr]
+	// The set holds every commitment, so absence here is absence on disk.
+	height, exists := udb.utxoCache[commitmentStr]
 	if !exists {
-		// Try loading from DB
-		var err error
-		utxo, err = udb.getUTXONoLock(commitment)
-		if err != nil {
-			return errors.New("UTXO not found")
-		}
+		return errors.New("UTXO not found")
 	}

 	// Remove from database
@@ -174,27 +133,19 @@ func (udb *UTXODB) RemoveUTXO(commitment []byte) error {
 		return err
 	}

-	// Remove from cache
+	// Remove from the set
 	delete(udb.utxoCache, commitmentStr)

 	// Update height index
-	if heightUTXOs, exists := udb.heightIndex[utxo.Height]; exists {
+	if heightUTXOs, exists := udb.heightIndex[height]; exists {
 		for i, c := range heightUTXOs {
 			if c == commitmentStr {
-				udb.heightIndex[utxo.Height] = append(heightUTXOs[:i], heightUTXOs[i+1:]...)
+				udb.heightIndex[height] = append(heightUTXOs[:i], heightUTXOs[i+1:]...)
 				break
 			}
 		}
 	}

-	// Update count
-	udb.utxoCount--
-	countBytes := make([]byte, 8)
-	binary.BigEndian.PutUint64(countBytes, udb.utxoCount)
-	if err := udb.db.Put([]byte(utxoCountKey), countBytes); err != nil {
-		return err
-	}
-
 	return nil
 }

@@ -210,19 +161,27 @@ func (udb *UTXODB) GetUTXOsByHeight(height uint64) ([]*UTXO, error) {

 	utxos := make([]*UTXO, 0, len(commitments))
 	for _, commitmentStr := range commitments {
-		if utxo, exists := udb.utxoCache[commitmentStr]; exists {
-			utxos = append(utxos, utxo)
+		utxo, err := udb.read([]byte(commitmentStr))
+		if err != nil {
+			return nil, err
 		}
+		utxos = append(utxos, utxo)
 	}

 	return utxos, nil
 }

-// GetUTXOCount returns the total number of UTXOs
+// GetUTXOCount returns the total number of UTXOs.
+//
+// It counts the set rather than reading a running total kept beside it. A total
+// is a second write, and a node that dies between the two comes back with a
+// number that disagrees with its own records — from which one removal drives an
+// unsigned counter below zero and reports 1.8e19 unspent notes forever.
+// Counting the set cannot disagree with the set.
 func (udb *UTXODB) GetUTXOCount() uint64 {
 	udb.mu.RLock()
 	defer udb.mu.RUnlock()
-	return udb.utxoCount
+	return uint64(len(udb.utxoCache))
 }

 // GetAllCommitments returns all UTXO commitments (for Merkle tree)
@@ -231,8 +190,8 @@ func (udb *UTXODB) GetAllCommitments() [][]byte {
 	defer udb.mu.RUnlock()

 	commitments := make([][]byte, 0, len(udb.utxoCache))
-	for _, utxo := range udb.utxoCache {
-		commitments = append(commitments, utxo.Commitment)
+	for commitmentStr := range udb.utxoCache {
+		commitments = append(commitments, []byte(commitmentStr))
 	}

 	return commitments
@@ -266,7 +225,7 @@ func (udb *UTXODB) PruneOldUTXOs(minHeight uint64) error {
 				continue
 			}

-			// Remove from cache
+			// Remove from the set
 			delete(udb.utxoCache, commitmentStr)
 			pruneCount++
 		}
@@ -275,30 +234,50 @@ func (udb *UTXODB) PruneOldUTXOs(minHeight uint64) error {
 		delete(udb.heightIndex, height)
 	}

-	// Update count
-	udb.utxoCount -= uint64(pruneCount)
-	countBytes := make([]byte, 8)
-	binary.BigEndian.PutUint64(countBytes, udb.utxoCount)
-	if err := udb.db.Put([]byte(utxoCountKey), countBytes); err != nil {
-		return err
-	}
-
 	udb.log.Info("Pruned old UTXOs",
 		log.Int("pruneCount", pruneCount),
 		log.Uint64("minHeight", minHeight),
-		log.Uint64("remainingUTXOs", udb.utxoCount),
+		log.Int("remainingUTXOs", len(udb.utxoCache)),
 	)

 	return nil
 }

-// loadUTXOs loads UTXOs from database to cache
+// loadUTXOs rebuilds the set and the height index from the records.
+//
+// The records share a keyspace with other writers, so a key carrying the UTXO
+// prefix counts as a UTXO only if the record under it names the commitment its
+// own key is made of. Anything else belongs to someone else and is left alone.
 func (udb *UTXODB) loadUTXOs() error {
-	return nil
+	it := udb.db.NewIteratorWithPrefix([]byte{utxoPrefix})
+	defer it.Release()
+
+	for it.Next() {
+		key := it.Key()
+		if len(key) < 2 {
+			continue
+		}
+		commitment := key[1:]
+
+		var utxo UTXO
+		if err := parseUTXO(it.Value(), &utxo); err != nil {
+			continue
+		}
+		if !bytes.Equal(utxo.Commitment, commitment) {
+			continue
+		}
+
+		commitmentStr := string(commitment)
+		udb.utxoCache[commitmentStr] = utxo.Height
+		udb.heightIndex[utxo.Height] = append(udb.heightIndex[utxo.Height], commitmentStr)
+	}
+
+	return it.Error()
 }

-// getUTXONoLock retrieves a UTXO without locking (internal use)
-func (udb *UTXODB) getUTXONoLock(commitment []byte) (*UTXO, error) {
+// read returns the record for a commitment. It touches no shared state, so it
+// is safe under either lock.
+func (udb *UTXODB) read(commitment []byte) (*UTXO, error) {
 	key := makeUTXOKey(commitment)
 	utxoBytes, err := udb.db.Get(key)
 	if err != nil {
@@ -0,0 +1,144 @@
+// Copyright (C) 2019-2025, Lux Industries Inc. All rights reserved.
+// See the file LICENSE for licensing terms.
+
+package zkvm
+
+import (
+	"sync"
+	"testing"
+
+	"github.com/luxfi/database/memdb"
+	"github.com/luxfi/ids"
+	"github.com/luxfi/log"
+	"github.com/stretchr/testify/require"
+)
+
+// TestRestartedNodeRefusesADuplicateCommitment. Block.Accept refuses a block
+// carrying a commitment the set already holds. The set used to live only in
+// memory and was never rebuilt, so that verdict depended on how recently the
+// node was started: a restarted validator accepted what a running one refused,
+// and the two ended up on different chains.
+func TestRestartedNodeRefusesADuplicateCommitment(t *testing.T) {
+	db := memdb.New()
+	utxo := &UTXO{TxID: ids.ID{7}, Commitment: []byte("commitment"), Height: 3}
+
+	live, err := NewUTXODB(db, log.NoLog{})
+	require.NoError(t, err)
+	require.NoError(t, live.AddUTXO(utxo))
+	require.ErrorContains(t, live.AddUTXO(utxo), "already exists")
+	require.Equal(t, uint64(1), live.GetUTXOCount())
+
+	restarted, err := NewUTXODB(db, log.NoLog{})
+	require.NoError(t, err)
+	require.ErrorContains(t, restarted.AddUTXO(utxo), "already exists",
+		"a restarted node must reach the same verdict as a running one")
+	require.Equal(t, uint64(1), restarted.GetUTXOCount())
+}
+
+// TestUTXOCountIsReadOffTheRecords. The count used to be a running total
+// written beside the records, so a node that died between the two writes came
+// back with a number that disagreed with its own set — and one removal from
+// there drove an unsigned counter below zero. This starts from the state that
+// used to be unrecoverable: a record on disk with no total beside it.
+func TestUTXOCountIsReadOffTheRecords(t *testing.T) {
+	db := memdb.New()
+	utxo := &UTXO{TxID: ids.ID{1}, Commitment: []byte("orphan"), Height: 1}
+
+	record, err := utxo.Marshal()
+	require.NoError(t, err)
+	require.NoError(t, db.Put(makeUTXOKey(utxo.Commitment), record))
+
+	udb, err := NewUTXODB(db, log.NoLog{})
+	require.NoError(t, err)
+	require.Equal(t, uint64(1), udb.GetUTXOCount(),
+		"the count is read off the records, so it cannot disagree with the set it describes")
+
+	require.NoError(t, udb.RemoveUTXO(utxo.Commitment))
+	require.Zero(t, udb.GetUTXOCount(),
+		"removing the last record must leave zero, not wrap to 2^64-1")
+	require.ErrorContains(t, udb.RemoveUTXO(utxo.Commitment), "not found")
+	require.Zero(t, udb.GetUTXOCount())
+}
+
+// TestRestartRebuildsTheHeightIndex. Pruning and the Merkle view read the
+// height index and the commitment list. A restart that left them empty had the
+// node reporting a UTXO set it was still holding.
+func TestRestartRebuildsTheHeightIndex(t *testing.T) {
+	db := memdb.New()
+	live, err := NewUTXODB(db, log.NoLog{})
+	require.NoError(t, err)
+	for i := 0; i < 3; i++ {
+		require.NoError(t, live.AddUTXO(&UTXO{
+			TxID: ids.ID{byte(i)}, OutputIndex: uint32(i),
+			Commitment: []byte{'c', byte(i)}, Ciphertext: []byte{'n', byte(i)}, Height: 5,
+		}))
+	}
+
+	restarted, err := NewUTXODB(db, log.NoLog{})
+	require.NoError(t, err)
+	require.Len(t, restarted.GetAllCommitments(), 3, "the commitment set is rebuilt from the records")
+
+	byHeight, err := restarted.GetUTXOsByHeight(5)
+	require.NoError(t, err)
+	require.Len(t, byHeight, 3, "and so is the height index")
+	for _, utxo := range byHeight {
+		require.Equal(t, uint64(5), utxo.Height)
+		require.Equal(t, []byte{'n', utxo.Commitment[1]}, utxo.Ciphertext,
+			"the body comes back whole, not just its key")
+	}
+
+	empty, err := restarted.GetUTXOsByHeight(6)
+	require.NoError(t, err)
+	require.Empty(t, empty, "a height nothing was created at holds nothing")
+}
+
+// TestUTXOReadDoesNotWriteTheSet. GetUTXO used to memoise what it loaded while
+// holding only the read lock, and a read lock promises every other reader that
+// nothing is changing — so two clients asking for uncached commitments wrote
+// the same map at the same time, which is a runtime throw that takes the
+// process down rather than an error anyone can handle.
+func TestUTXOReadDoesNotWriteTheSet(t *testing.T) {
+	udb, err := NewUTXODB(memdb.New(), log.NoLog{})
+	require.NoError(t, err)
+	require.NoError(t, udb.AddUTXO(&UTXO{TxID: ids.ID{1}, Commitment: []byte("c"), Height: 1}))
+
+	// A record on disk that the in-memory set does not know about is what a
+	// memoising read would fill in.
+	udb.utxoCache = map[string]uint64{}
+
+	var readers sync.WaitGroup
+	for i := 0; i < 8; i++ {
+		readers.Add(1)
+		go func() {
+			defer readers.Done()
+			utxo, err := udb.GetUTXO([]byte("c"))
+			switch {
+			case err != nil:
+				t.Errorf("read failed: %v", err)
+			case utxo.Height != 1:
+				t.Errorf("read height %d, want 1", utxo.Height)
+			}
+		}()
+	}
+	readers.Wait()
+
+	require.Empty(t, udb.utxoCache, "a read must not change the set it reads")
+}
+
+// TestForeignRecordIsNotAUTXO. Blocks are stored in the same database under
+// their raw 32-byte id, and one id in 256 begins with the UTXO prefix byte, so
+// a prefix scan alone would adopt other people's records into the set. A record
+// counts as a UTXO only when it names the commitment its own key is made of.
+func TestForeignRecordIsNotAUTXO(t *testing.T) {
+	db := memdb.New()
+
+	record, err := (&UTXO{TxID: ids.ID{1}, Commitment: []byte("orphan"), Height: 1}).Marshal()
+	require.NoError(t, err)
+	foreign := ids.ID{utxoPrefix, 1}
+	require.NoError(t, db.Put(foreign[:], record))
+
+	udb, err := NewUTXODB(db, log.NoLog{})
+	require.NoError(t, err)
+	require.Zero(t, udb.GetUTXOCount(), "someone else's key is not a commitment")
+	require.Empty(t, udb.GetAllCommitments())
+}