Merge pull request #259 from kcalvinalvin/2025-02-06-add-utreexoroots…

…tate-for-utreexoproofindex

indexers: add utreexoRootsState to utreexoproofindex

blockchain/indexers/indexers_test.go

@@ -981,3 +981,108 @@ func TestBridgeNodePruneUndoDataGen(t *testing.T) {
 		}
 	}
 }
+
+func compareUtreexoRootsState(indexes []Indexer, blockHash *chainhash.Hash) error {
+	var err error
+	var flatMsg *wire.MsgUtreexoRoot
+	var msg *wire.MsgUtreexoRoot
+	for _, indexer := range indexes {
+		switch idxType := indexer.(type) {
+		case *FlatUtreexoProofIndex:
+			flatMsg, err = idxType.FetchMsgUtreexoRoot(blockHash)
+			if err != nil {
+				return err
+			}
+
+		case *UtreexoProofIndex:
+			msg, err = idxType.FetchMsgUtreexoRoot(blockHash)
+			if err != nil {
+				return err
+			}
+		}
+	}
+
+	if flatMsg.NumLeaves != msg.NumLeaves {
+		return fmt.Errorf("expected %v, got %v", flatMsg.NumLeaves, msg.NumLeaves)
+	}
+
+	if flatMsg.Target != msg.Target {
+		return fmt.Errorf("expected %v, got %v", flatMsg.Target, msg.Target)
+	}
+
+	if !flatMsg.BlockHash.IsEqual(&msg.BlockHash) {
+		return fmt.Errorf("expected %v, got %v", flatMsg.BlockHash, msg.BlockHash)
+	}
+
+	if len(flatMsg.Roots) != len(msg.Roots) {
+		return fmt.Errorf("expected %v, got %v", len(flatMsg.Roots), len(msg.Roots))
+	}
+	for i := range flatMsg.Roots {
+		if flatMsg.Roots[i] != msg.Roots[i] {
+			return fmt.Errorf("expected %v, got %v", flatMsg.Roots[i], msg.Roots[i])
+		}
+	}
+
+	if len(flatMsg.Proof) != len(msg.Proof) {
+		return fmt.Errorf("expected %v, got %v", len(flatMsg.Proof), len(msg.Proof))
+	}
+	for i := range flatMsg.Proof {
+		if flatMsg.Proof[i] != msg.Proof[i] {
+			return fmt.Errorf("expected %v, got %v", flatMsg.Proof[i], msg.Proof[i])
+		}
+	}
+
+	return nil
+}
+
+func TestUtreexoRootsState(t *testing.T) {
+	// Always remove the root on return.
+	defer os.RemoveAll(testDbRoot)
+
+	chain, indexes, params, _, tearDown := indexersTestChain("TestUtreexoRootsState")
+	defer tearDown()
+
+	var allSpends []*blockchain.SpendableOut
+	var nextSpends []*blockchain.SpendableOut
+
+	// Number of blocks we'll generate for the test.
+	maxHeight := int32(300)
+
+	nextBlock := btcutil.NewBlock(params.GenesisBlock)
+	for i := int32(1); i <= maxHeight; i++ {
+		newBlock, newSpendableOuts, err := blockchain.AddBlock(chain, nextBlock, nextSpends)
+		if err != nil {
+			t.Fatal(err)
+		}
+		nextBlock = newBlock
+
+		allSpends = append(allSpends, newSpendableOuts...)
+
+		var nextSpendsTmp []*blockchain.SpendableOut
+		for j := 0; j < len(allSpends); j++ {
+			randIdx := rand.Intn(len(allSpends))
+
+			spend := allSpends[randIdx]                                       // get
+			allSpends = append(allSpends[:randIdx], allSpends[randIdx+1:]...) // delete
+			nextSpendsTmp = append(nextSpendsTmp, spend)
+		}
+		nextSpends = nextSpendsTmp
+
+		err = compareUtreexoRootsState(indexes, newBlock.Hash())
+		if err != nil {
+			t.Fatal(err)
+		}
+	}
+
+	bestHash := chain.BestSnapshot().Hash
+	err := chain.InvalidateBlock(&bestHash)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	bestHash = chain.BestSnapshot().Hash
+	err = compareUtreexoRootsState(indexes, &bestHash)
+	if err != nil {
+		t.Fatal(err)
+	}
+}

blockchain/indexers/utreexoproofindex.go

@@ -6,6 +6,7 @@ package indexers
 
 import (
 	"bytes"
+	"crypto/sha256"
 	"fmt"
 	"sync"
 	"time"
@@ -66,6 +67,11 @@ type UtreexoProofIndex struct {
 	// It keeps all the elements of the forest in order to generate proofs.
 	utreexoState *UtreexoState
 
+	// utreexoRootsState is the accumulator for all the roots at each of the
+	// blocks. This is so that we can serve to peers the proof that a set of
+	// roots at a block is correct.
+	utreexoRootsState utreexo.Pollard
+
 	// The time of when the utreexo state was last flushed.
 	lastFlushTime time.Time
 }
@@ -78,6 +84,41 @@ func (idx *UtreexoProofIndex) NeedsInputs() bool {
 	return true
 }
 
+// initUtreexoRootsState creates an accumulator from all the existing roots and
+// holds it in memory so that the proofs for them can be generated.
+func (idx *UtreexoProofIndex) initUtreexoRootsState(bestHeight int32) error {
+	idx.utreexoRootsState = utreexo.NewAccumulator()
+
+	for h := int32(1); h <= bestHeight; h++ {
+		hash, err := idx.chain.BlockHashByHeight(h)
+		if err != nil {
+			return err
+		}
+		var stump utreexo.Stump
+		err = idx.db.View(func(dbTx database.Tx) error {
+			stump, err = dbFetchUtreexoState(dbTx, hash)
+			return err
+		})
+		if err != nil {
+			return err
+		}
+
+		bytes, err := blockchain.SerializeUtreexoRoots(stump.NumLeaves, stump.Roots)
+		if err != nil {
+			return err
+		}
+		rootHash := sha256.Sum256(bytes)
+
+		err = idx.utreexoRootsState.Modify(
+			[]utreexo.Leaf{{Hash: rootHash}}, nil, utreexo.Proof{})
+		if err != nil {
+			return err
+		}
+	}
+
+	return nil
+}
+
 // Init initializes the utreexo proof index. This is part of the Indexer
 // interface.
 func (idx *UtreexoProofIndex) Init(chain *blockchain.BlockChain,
@@ -93,6 +134,11 @@ func (idx *UtreexoProofIndex) Init(chain *blockchain.BlockChain,
 	idx.utreexoState = uState
 	idx.lastFlushTime = time.Now()
 
+	err = idx.initUtreexoRootsState(tipHeight)
+	if err != nil {
+		return err
+	}
+
 	// Nothing else to do if the node is an archive node.
 	if !idx.config.Pruned {
 		return nil
@@ -311,7 +357,7 @@ func (idx *UtreexoProofIndex) ConnectBlock(dbTx database.Tx, block *btcutil.Bloc
 		return err
 	}
 
-	return nil
+	return idx.updateRootsState()
 }
 
 // getUndoData returns the data needed for undo. For pruned nodes, we fetch the data from
@@ -405,7 +451,9 @@ func (idx *UtreexoProofIndex) DisconnectBlock(dbTx database.Tx, block *btcutil.B
 		}
 	}
 
-	return nil
+	// Re-initializes to the current accumulator roots, effectively disconnecting
+	// a block.
+	return idx.initUtreexoRootsState(block.Height() - 1)
 }
 
 // FetchUtreexoProof returns the Utreexo proof data for the given block hash.
@@ -579,6 +627,22 @@ func (idx *UtreexoProofIndex) VerifyAccProof(toProve []utreexo.Hash,
 	return idx.utreexoState.state.Verify(toProve, *proof, false)
 }
 
+// updateRootsState updates the roots accumulator state from the roots of the current accumulator.
+func (idx *UtreexoProofIndex) updateRootsState() error {
+	idx.mtx.Lock()
+	numLeaves := idx.utreexoState.state.GetNumLeaves()
+	roots := idx.utreexoState.state.GetRoots()
+	idx.mtx.Unlock()
+
+	bytes, err := blockchain.SerializeUtreexoRoots(numLeaves, roots)
+	if err != nil {
+		return err
+	}
+
+	rootHash := sha256.Sum256(bytes)
+	return idx.utreexoRootsState.Modify([]utreexo.Leaf{{Hash: rootHash}}, nil, utreexo.Proof{})
+}
+
 // PruneBlock is invoked when an older block is deleted after it's been
 // processed.
 //
@@ -608,6 +672,37 @@ func (idx *UtreexoProofIndex) PruneBlock(_ database.Tx, _ *chainhash.Hash, lastK
 	return idx.Flush(&bestHash, blockchain.FlushRequired, true)
 }
 
+// FetchMsgUtreexoRoot returns a complete utreexoroot bitcoin message on the requested block.
+func (idx *UtreexoProofIndex) FetchMsgUtreexoRoot(blockHash *chainhash.Hash) (*wire.MsgUtreexoRoot, error) {
+	var stump utreexo.Stump
+	err := idx.db.View(func(dbTx database.Tx) error {
+		var err error
+		stump, err = dbFetchUtreexoState(dbTx, blockHash)
+		return err
+	})
+
+	bytes, err := blockchain.SerializeUtreexoRoots(stump.NumLeaves, stump.Roots)
+	if err != nil {
+		return nil, err
+	}
+	rootHash := sha256.Sum256(bytes)
+
+	proof, err := idx.utreexoRootsState.Prove([]utreexo.Hash{rootHash})
+	if err != nil {
+		return nil, err
+	}
+
+	msg := &wire.MsgUtreexoRoot{
+		NumLeaves: stump.NumLeaves,
+		Target:    proof.Targets[0],
+		BlockHash: *blockHash,
+		Roots:     stump.Roots,
+		Proof:     proof.Proof,
+	}
+
+	return msg, nil
+}
+
 // NewUtreexoProofIndex returns a new instance of an indexer that is used to create a utreexo
 // proof index using the database passed in. The passed in maxMemoryUsage should be in bytes and
 // it determines how much memory the proof index will use up.

server.go

@@ -1386,7 +1386,7 @@ func (sp *serverPeer) OnGetUtreexoRoot(_ *peer.Peer, msg *wire.MsgGetUtreexoRoot
 	}
 
 	// Check if we're a utreexo bridge node. Ignore if we're not.
-	if sp.server.flatUtreexoProofIndex == nil {
+	if sp.server.utreexoProofIndex == nil && sp.server.flatUtreexoProofIndex == nil {
 		return
 	}
 
@@ -1401,6 +1401,15 @@ func (sp *serverPeer) OnGetUtreexoRoot(_ *peer.Peer, msg *wire.MsgGetUtreexoRoot
 		}
 	}
 
+	if sp.server.utreexoProofIndex != nil {
+		utreexoRootMsg, err = sp.server.utreexoProofIndex.FetchMsgUtreexoRoot(&msg.BlockHash)
+		if err != nil {
+			chanLog.Debugf("Unable to fetch the utreexo root for block hash %v: %v",
+				msg.BlockHash, err)
+			return
+		}
+	}
+
 	sp.QueueMessage(utreexoRootMsg, nil)
 }