From 2cb8ca19c19856e1264151c7a1aa540334358daf Mon Sep 17 00:00:00 2001
From: Calvin Kim <calvin@kcalvinalvin.info>
Date: Fri, 3 Jan 2025 14:17:20 +0900
Subject: [PATCH] peer, main: handle getutreexoproof message

---
 peer/peer.go |  8 +++++
 server.go    | 91 ++++++++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 99 insertions(+)

diff --git a/peer/peer.go b/peer/peer.go
index cbdf6467..3ac09cd5 100644
--- a/peer/peer.go
+++ b/peer/peer.go
@@ -143,6 +143,9 @@ type MessageListeners struct {
 	// OnUtreexoHeader is invoked when a peer receives a headers bitcoin message.
 	OnUtreexoHeader func(p *Peer, msg *wire.MsgUtreexoHeader)
 
+	// OnGetUtreexoProof is invoked when a peer receives a utreexo proof bitcoin message.
+	OnGetUtreexoProof func(p *Peer, msg *wire.MsgGetUtreexoProof)
+
 	// OnNotFound is invoked when a peer receives a notfound bitcoin
 	// message.
 	OnNotFound func(p *Peer, msg *wire.MsgNotFound)
@@ -1480,6 +1483,11 @@ out:
 				p.cfg.Listeners.OnUtreexoHeader(p, msg)
 			}
 
+		case *wire.MsgGetUtreexoProof:
+			if p.cfg.Listeners.OnGetUtreexoProof != nil {
+				p.cfg.Listeners.OnGetUtreexoProof(p, msg)
+			}
+
 		case *wire.MsgNotFound:
 			if p.cfg.Listeners.OnNotFound != nil {
 				p.cfg.Listeners.OnNotFound(p, msg)
diff --git a/server.go b/server.go
index 75e41ca1..3a8517fe 100644
--- a/server.go
+++ b/server.go
@@ -1311,6 +1311,96 @@ func (sp *serverPeer) OnGetCFCheckpt(_ *peer.Peer, msg *wire.MsgGetCFCheckpt) {
 	sp.QueueMessage(checkptMsg, nil)
 }
 
+// OnGetUtreexoProof is invoked when a peer receives a getutreexoproof bitcoin message.
+func (sp *serverPeer) OnGetUtreexoProof(_ *peer.Peer, msg *wire.MsgGetUtreexoProof) {
+	// Ignore getutreexoproof requests if not in sync.
+	if !sp.server.syncManager.IsCurrent() {
+		return
+	}
+
+	// Check if we're a utreexo node. Ignore if we're not.
+	if sp.server.utreexoProofIndex == nil && sp.server.flatUtreexoProofIndex == nil && cfg.NoUtreexo {
+		return
+	}
+
+	height, err := sp.server.chain.BlockHeightByHash(&msg.BlockHash)
+	if err != nil {
+		chanLog.Debugf("Unable to fetch height for block hash %v: %v",
+			msg.BlockHash, err)
+		return
+	}
+
+	// If we're pruned and the requested block is beyond the point where pruned blocks
+	// are able to serve blocks, just ignore the message.
+	if cfg.Prune != 0 && height >= sp.server.chain.BestSnapshot().Height-288 {
+		return
+	}
+
+	block, err := sp.server.chain.BlockByHash(&msg.BlockHash)
+	if err != nil {
+		chanLog.Debugf("Unable to fetch block with hash %v: %v",
+			msg.BlockHash, err)
+		return
+	}
+
+	// Fetch UData.
+	var udata *wire.UData
+	if !cfg.NoUtreexo {
+		udata = block.MsgBlock().UData
+	}
+	if sp.server.utreexoProofIndex != nil {
+		udata, err = sp.server.utreexoProofIndex.FetchUtreexoProof(&msg.BlockHash)
+		if err != nil {
+			chanLog.Debugf("Unable to fetch utreexo proof for block hash %v: %v",
+				msg.BlockHash, err)
+			return
+		}
+	}
+	if sp.server.flatUtreexoProofIndex != nil {
+		udata, err = sp.server.flatUtreexoProofIndex.FetchUtreexoProof(height)
+		if err != nil {
+			chanLog.Debugf("Unable to fetch utreexo proof for block hash %v: %v",
+				msg.BlockHash, err)
+			return
+		}
+	}
+
+	_, err = sp.server.chain.ReconstructUData(udata, msg.BlockHash)
+	if err != nil {
+		chanLog.Debugf("Unable to fetch utreexo proof for block hash %v: %v",
+			msg.BlockHash, err)
+		return
+	}
+
+	// Construct utreexo proof to send.
+	leafDatas := make([]wire.LeafData, 0, len(msg.LeafIndexes))
+	for _, idx := range msg.LeafIndexes {
+		if int(idx) > len(udata.LeafDatas) {
+			chanLog.Debugf("Unable to fetch utreexo proof for block hash %v: %v",
+				msg.BlockHash, err)
+			return
+		}
+		leafDatas = append(leafDatas, udata.LeafDatas[idx])
+	}
+	proofHashes := make([]utreexo.Hash, 0, len(msg.ProofIndexes))
+	for _, idx := range msg.ProofIndexes {
+		if int(idx) >= len(udata.AccProof.Proof) {
+			chanLog.Debugf("Unable to fetch utreexo proof for block hash %v: %v",
+				msg.BlockHash, err)
+			return
+		}
+		proofHashes = append(proofHashes, udata.AccProof.Proof[idx])
+	}
+	utreexoProof := wire.MsgUtreexoProof{
+		BlockHash:   msg.BlockHash,
+		ProofHashes: proofHashes,
+		LeafDatas:   leafDatas,
+	}
+
+	sp.QueueMessage(&utreexoProof, nil)
+	return
+}
+
 // enforceNodeBloomFlag disconnects the peer if the server is not configured to
 // allow bloom filters.  Additionally, if the peer has negotiated to a protocol
 // version  that is high enough to observe the bloom filter service support bit,
@@ -2547,6 +2637,7 @@ func newPeerConfig(sp *serverPeer) *peer.Config {
 			OnInv:              sp.OnInv,
 			OnHeaders:          sp.OnHeaders,
 			OnUtreexoHeader:    sp.OnUtreexoHeader,
+			OnGetUtreexoProof:  sp.OnGetUtreexoProof,
 			OnGetData:          sp.OnGetData,
 			OnGetBlocks:        sp.OnGetBlocks,
 			OnGetHeaders:       sp.OnGetHeaders,