func (h *prunedBlockDispatcherHarness) enablePeerReplies() {
	atomic.StoreUint32(&h.shouldReply, 0)
}
enablePeerReplies allows the query peer to reply.
package chain
import (
	"encoding/binary"
	"encoding/hex"
	"fmt"
	"net"
	"sync"
	"sync/atomic"
	"testing"
	"time"
	"github.com/btcsuite/btcd/btcjson"
	"github.com/btcsuite/btcd/chaincfg"
	"github.com/btcsuite/btcd/chaincfg/chainhash"
	"github.com/btcsuite/btcd/peer"
	"github.com/btcsuite/btcd/wire"
	"github.com/lightningnetwork/lnd/ticker"
	"github.com/stretchr/testify/require"
)
var (
	addrCounter int32  
	chainParams = chaincfg.RegressionNetParams
)
func nextAddr() string {
	port := atomic.AddInt32(&addrCounter, 1)
	return fmt.Sprintf("10.0.0.1:%d", port)
}
type prunedBlockDispatcherHarness struct {
	t *testing.T
	dispatcher *PrunedBlockDispatcher
	hashes []*chainhash.Hash
	blocks map[chainhash.Hash]*wire.MsgBlock
	peerMtx       sync.Mutex
	peers         map[string]*peer.Peer
	fallbackAddrs map[string]*peer.Peer
	localConns    map[string]net.Conn  
	remoteConns   map[string]net.Conn  
	dialedPeer    chan string
	queriedPeer   chan struct{}
	blocksQueried map[chainhash.Hash]int
	shouldReply uint32  
}
func newNetworkBlockTestHarness(t *testing.T, numBlocks,
	numPeers, numWorkers uint32) *prunedBlockDispatcherHarness {
	h := &prunedBlockDispatcherHarness{
		t:             t,
		dispatcher:    &PrunedBlockDispatcher{},
		peers:         make(map[string]*peer.Peer, numPeers),
		fallbackAddrs: make(map[string]*peer.Peer, numPeers),
		localConns:    make(map[string]net.Conn, numPeers),
		remoteConns:   make(map[string]net.Conn, numPeers),
		dialedPeer:    make(chan string),
		queriedPeer:   make(chan struct{}),
		blocksQueried: make(map[chainhash.Hash]int),
		shouldReply:   0,
	}
	h.hashes, h.blocks = genBlockChain(numBlocks)
	for i := uint32(0); i < numPeers; i++ {
		h.addPeer(false)
	}
	dial := func(addr string) (net.Conn, error) {
		go func() {
			h.dialedPeer <- addr
		}()
		h.peerMtx.Lock()
		defer h.peerMtx.Unlock()
		localConn, ok := h.localConns[addr]
		if !ok {
			return nil, fmt.Errorf("local conn %v not found", addr)
		}
		remoteConn, ok := h.remoteConns[addr]
		if !ok {
			return nil, fmt.Errorf("remote conn %v not found", addr)
		}
		if p, ok := h.peers[addr]; ok {
			p.AssociateConnection(remoteConn)
		}
		if p, ok := h.fallbackAddrs[addr]; ok {
			p.AssociateConnection(remoteConn)
		}
		return localConn, nil
	}
	var err error
	h.dispatcher, err = NewPrunedBlockDispatcher(&PrunedBlockDispatcherConfig{
		ChainParams:    &chainParams,
		NumTargetPeers: int(numWorkers),
		Dial:           dial,
		GetPeers: func() ([]btcjson.GetPeerInfoResult, error) {
			h.peerMtx.Lock()
			defer h.peerMtx.Unlock()
			res := make([]btcjson.GetPeerInfoResult, 0, len(h.peers))
			for addr, peer := range h.peers {
				var rawServices [8]byte
				binary.BigEndian.PutUint64(
					rawServices[:], uint64(peer.Services()),
				)
				res = append(res, btcjson.GetPeerInfoResult{
					Addr:     addr,
					Services: hex.EncodeToString(rawServices[:]),
				})
			}
			return res, nil
		},
		GetNodeAddresses: func(*int32) ([]btcjson.GetNodeAddressesResult, error) {
			h.peerMtx.Lock()
			defer h.peerMtx.Unlock()
			res := make(
				[]btcjson.GetNodeAddressesResult, 0,
				len(h.fallbackAddrs),
			)
			for addr, peer := range h.fallbackAddrs {
				res = append(res, btcjson.GetNodeAddressesResult{
					Services: uint64(peer.Services()),
					Address:  addr,
				})
			}
			return res, nil
		},
		PeerReadyTimeout:   time.Hour,
		RefreshPeersTicker: ticker.NewForce(time.Hour),
		AllowSelfPeerConns: true,
		MaxRequestInvs:     wire.MaxInvPerMsg,
	})
	require.NoError(t, err)
	return h
}
func (h *prunedBlockDispatcherHarness) start() {
	h.t.Helper()
	err := h.dispatcher.Start()
	require.NoError(h.t, err)
	h.peerMtx.Lock()
	numPeers := len(h.peers)
	h.peerMtx.Unlock()
	for i := 0; i < numPeers; i++ {
		h.assertPeerDialed()
	}
}
func (h *prunedBlockDispatcherHarness) stop() {
	h.dispatcher.Stop()
	select {
	case <-h.dialedPeer:
		h.t.Fatal("did not consume all dialedPeer signals")
	default:
	}
	select {
	case <-h.queriedPeer:
		h.t.Fatal("did not consume all queriedPeer signals")
	default:
	}
	require.Empty(h.t, h.blocksQueried)
}
func (h *prunedBlockDispatcherHarness) addPeer(fallback bool) string {
	addr := nextAddr()
	h.peerMtx.Lock()
	defer h.peerMtx.Unlock()
	h.resetPeer(addr, fallback)
	return addr
}
func (h *prunedBlockDispatcherHarness) resetPeer(addr string, fallback bool) {
	if fallback {
		h.fallbackAddrs[addr] = h.newPeer()
	} else {
		h.peers[addr] = h.newPeer()
	}
	inConn, outConn := pipe(
		&conn{localAddr: addr, remoteAddr: "10.0.0.1:8333"},
		&conn{localAddr: "10.0.0.1:8333", remoteAddr: addr},
	)
	h.localConns[addr] = outConn
	h.remoteConns[addr] = inConn
}
func (h *prunedBlockDispatcherHarness) newPeer() *peer.Peer {
	return peer.NewInboundPeer(&peer.Config{
		ChainParams:    &chainParams,
		DisableRelayTx: true,
		Listeners: peer.MessageListeners{
			OnGetData: func(p *peer.Peer, msg *wire.MsgGetData) {
				go func() {
					h.queriedPeer <- struct{}{}
				}()
				for _, inv := range msg.InvList {
					require.Equal(h.t, wire.InvTypeBlock, inv.Type)
					block, ok := h.blocks[inv.Hash]
					require.True(h.t, ok)
					switch atomic.LoadUint32(&h.shouldReply) {
					case 1:
						continue
					case 2:
						block = produceInvalidBlock(block)
					}
					go p.QueueMessage(block, nil)
				}
			},
		},
		Services:       wire.SFNodeNetwork | wire.SFNodeWitness,
		AllowSelfConns: true,
	})
}
func (h *prunedBlockDispatcherHarness) query(blocks []*chainhash.Hash) (
	<-chan *wire.MsgBlock, <-chan error) {
	h.t.Helper()
	blockChan, errChan := h.dispatcher.Query(blocks)
	select {
	case err := <-errChan:
		require.NoError(h.t, err)
	default:
	}
	for _, block := range blocks {
		h.blocksQueried[*block]++
	}
	return blockChan, errChan
}
func (h *prunedBlockDispatcherHarness) disablePeerReplies() {
	atomic.StoreUint32(&h.shouldReply, 1)
}