From 24426399a47e66ac80281c6d28409a18c1ef74ad Mon Sep 17 00:00:00 2001
From: Oliver Gugger <gugger@gmail.com>
Date: Fri, 14 Jul 2023 14:49:40 +0200
Subject: [PATCH] blockmanager: add test for non-connected headers

---
 blockmanager_test.go | 80 ++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 80 insertions(+)

diff --git a/blockmanager_test.go b/blockmanager_test.go
index c5f25328..45554b7c 100644
--- a/blockmanager_test.go
+++ b/blockmanager_test.go
@@ -4,14 +4,17 @@ import (
 	"encoding/binary"
 	"fmt"
 	"math/rand"
+	"reflect"
 	"strings"
 	"testing"
 	"time"
 
+	"github.com/btcsuite/btcd/blockchain"
 	"github.com/btcsuite/btcd/btcutil/gcs"
 	"github.com/btcsuite/btcd/btcutil/gcs/builder"
 	"github.com/btcsuite/btcd/chaincfg"
 	"github.com/btcsuite/btcd/chaincfg/chainhash"
+	"github.com/btcsuite/btcd/integration/rpctest"
 	"github.com/btcsuite/btcd/peer"
 	"github.com/btcsuite/btcd/txscript"
 	"github.com/btcsuite/btcd/wire"
@@ -90,6 +93,7 @@ func setupBlockManager(t *testing.T) (*blockManager, headerfs.BlockHeaderStore,
 		BlockHeaders:     hdrStore,
 		RegFilterHeaders: cfStore,
 		QueryDispatcher:  &mockDispatcher{},
+		TimeSource:       blockchain.NewMedianTime(),
 		BanPeer: func(string, banman.Reason) error {
 			return nil
 		},
@@ -865,3 +869,79 @@ func TestBlockManagerDetectBadPeers(t *testing.T) {
 		require.NoError(t, err)
 	}
 }
+
+// TestHandleHeaders checks that we handle headers correctly, and that we
+// disconnect peers that serve us bad headers (headers that don't connect to
+// each other properly).
+func TestHandleHeaders(t *testing.T) {
+	t.Parallel()
+
+	// First, we set up a block manager and a fake peer that will act as the
+	// test's remote peer.
+	bm, _, _, err := setupBlockManager(t)
+	require.NoError(t, err)
+
+	fakePeer, err := peer.NewOutboundPeer(&peer.Config{}, "fake:123")
+	require.NoError(t, err)
+
+	assertPeerDisconnected := func(shouldBeDisconnected bool) {
+		// This is quite hacky but works: We expect the peer to be
+		// disconnected, which sets the unexported "disconnected" field
+		// to 1.
+		refValue := reflect.ValueOf(fakePeer).Elem()
+		foo := refValue.FieldByName("disconnect").Int()
+
+		if shouldBeDisconnected {
+			require.EqualValues(t, 1, foo)
+		} else {
+			require.EqualValues(t, 0, foo)
+		}
+	}
+
+	// We'll want to use actual, real blocks, so we take a miner harness
+	// that we can use to generate some.
+	harness, err := rpctest.New(
+		&chaincfg.SimNetParams, nil, []string{"--txindex"}, "",
+	)
+	require.NoError(t, err)
+	t.Cleanup(func() {
+		require.NoError(t, harness.TearDown())
+	})
+
+	err = harness.SetUp(false, 0)
+	require.NoError(t, err)
+
+	// Generate 200 valid blocks that we then feed to the block manager.
+	blockHashes, err := harness.Client.Generate(200)
+	require.NoError(t, err)
+
+	hmsg := &headersMsg{
+		headers: &wire.MsgHeaders{
+			Headers: make([]*wire.BlockHeader, len(blockHashes)),
+		},
+		peer: &ServerPeer{
+			Peer: fakePeer,
+		},
+	}
+
+	for i := range blockHashes {
+		header, err := harness.Client.GetBlockHeader(blockHashes[i])
+		require.NoError(t, err)
+
+		hmsg.headers.Headers[i] = header
+	}
+
+	// Let's feed in the correct headers. This should work fine and the peer
+	// should not be disconnected.
+	bm.handleHeadersMsg(hmsg)
+	assertPeerDisconnected(false)
+
+	// Now scramble the headers and feed them in again. This should cause
+	// the peer to be disconnected.
+	rand.Shuffle(len(hmsg.headers.Headers), func(i, j int) {
+		hmsg.headers.Headers[i], hmsg.headers.Headers[j] =
+			hmsg.headers.Headers[j], hmsg.headers.Headers[i]
+	})
+	bm.handleHeadersMsg(hmsg)
+	assertPeerDisconnected(true)
+}