From 68ed4d307803969385c7c547db0f386b99825531 Mon Sep 17 00:00:00 2001
From: Matheus Degiovani <opensource@matheusd.com>
Date: Thu, 29 Aug 2024 08:56:09 -0300
Subject: [PATCH] p2p: Move headers requested check earlier in the call

This moves the check for requested headers earlier in the call stack for
processing headers messages.

This ensures misbehaving peers are disconnected earlier, before spending
cpu processing for block hash calculation.
---
 p2p/peering.go | 22 ++++++++++++++--------
 1 file changed, 14 insertions(+), 8 deletions(-)

diff --git a/p2p/peering.go b/p2p/peering.go
index f334f4cb9..ad07faf6c 100644
--- a/p2p/peering.go
+++ b/p2p/peering.go
@@ -1122,6 +1122,15 @@ func (rp *RemotePeer) receivedHeaders(ctx context.Context, msg *wire.MsgHeaders)
 	const opf = "remotepeer(%v).receivedHeaders"
 	rp.requestedHeadersMu.Lock()
 
+	// Ensure the wallet requested headers from this peer.
+	if !rp.sendheaders && rp.requestedHeaders == nil {
+		op := errors.Opf(opf, rp.raddr)
+		err := errors.E(op, errors.Protocol, "received unrequested headers")
+		rp.Disconnect(err)
+		rp.requestedHeadersMu.Unlock()
+		return
+	}
+
 	// Ensure the remote peer sent as many headers as it could. It can only
 	// send fewer than 2k headers when the last one is >= their advertised
 	// height (their tip height). This handles cases where a peer might
@@ -1137,12 +1146,12 @@ func (rp *RemotePeer) receivedHeaders(ctx context.Context, msg *wire.MsgHeaders)
 		return
 	}
 
+	// Sanity check the headers connect to each other in sequence.
 	var prevHash chainhash.Hash
 	var prevHeight uint32
 	for i, h := range msg.Headers {
 		hash := h.BlockHash()
 
-		// Sanity check the headers connect to each other in sequence.
 		if i > 0 && (!prevHash.IsEqual(&h.PrevBlock) || h.Height != prevHeight+1) {
 			op := errors.Opf(opf, rp.raddr)
 			err := errors.E(op, errors.Protocol, "received out-of-sequence headers")
@@ -1155,6 +1164,7 @@ func (rp *RemotePeer) receivedHeaders(ctx context.Context, msg *wire.MsgHeaders)
 		prevHeight = h.Height
 	}
 
+	// Track the height of the last received header.
 	if prevHeight > 0 {
 		rp.lastHeightMu.Lock()
 		if int32(prevHeight) > rp.lastHeight {
@@ -1163,6 +1173,7 @@ func (rp *RemotePeer) receivedHeaders(ctx context.Context, msg *wire.MsgHeaders)
 		rp.lastHeightMu.Unlock()
 	}
 
+	// Async headers.
 	if rp.sendheaders {
 		rp.requestedHeadersMu.Unlock()
 		select {
@@ -1171,13 +1182,8 @@ func (rp *RemotePeer) receivedHeaders(ctx context.Context, msg *wire.MsgHeaders)
 		}
 		return
 	}
-	if rp.requestedHeaders == nil {
-		op := errors.Opf(opf, rp.raddr)
-		err := errors.E(op, errors.Protocol, "received unrequested headers")
-		rp.Disconnect(err)
-		rp.requestedHeadersMu.Unlock()
-		return
-	}
+
+	// Headers as a response to getheaders.
 	c := rp.requestedHeaders
 	rp.requestedHeaders = nil
 	rp.requestedHeadersMu.Unlock()