github.com/nspcc-dev/neo-go@v0.105.2-0.20240517133400-6be757af3eba/pkg/network/helper_test.go

package network

import (
	"context"
	"errors"
	"fmt"
	"net"
	"sync"
	"sync/atomic"
	"testing"
	"time"

	"github.com/nspcc-dev/neo-go/internal/fakechain"
	"github.com/nspcc-dev/neo-go/pkg/config"
	"github.com/nspcc-dev/neo-go/pkg/io"
	"github.com/nspcc-dev/neo-go/pkg/network/payload"
	"github.com/stretchr/testify/require"
	"go.uber.org/zap/zaptest"
)

type testDiscovery struct {
	sync.Mutex
	bad          []string
	connected    []string
	unregistered []string
	backfill     []string
}

func newTestDiscovery([]string, time.Duration, Transporter) Discoverer { return new(testDiscovery) }

func (d *testDiscovery) BackFill(addrs ...string) {
	d.Lock()
	defer d.Unlock()
	d.backfill = append(d.backfill, addrs...)
}
func (d *testDiscovery) PoolCount() int { return 0 }
func (d *testDiscovery) RegisterSelf(p AddressablePeer) {
	d.Lock()
	defer d.Unlock()
	d.bad = append(d.bad, p.ConnectionAddr())
}
func (d *testDiscovery) GetFanOut() int {
	d.Lock()
	defer d.Unlock()
	return (len(d.connected) + len(d.backfill)) * 2 / 3
}
func (d *testDiscovery) NetworkSize() int {
	d.Lock()
	defer d.Unlock()
	return len(d.connected) + len(d.backfill)
}
func (d *testDiscovery) RegisterGood(AddressablePeer) {}
func (d *testDiscovery) RegisterConnected(p AddressablePeer) {
	d.Lock()
	defer d.Unlock()
	d.connected = append(d.connected, p.ConnectionAddr())
}
func (d *testDiscovery) UnregisterConnected(p AddressablePeer, force bool) {
	d.Lock()
	defer d.Unlock()
	d.unregistered = append(d.unregistered, p.ConnectionAddr())
}
func (d *testDiscovery) UnconnectedPeers() []string {
	d.Lock()
	defer d.Unlock()
	return d.unregistered
}
func (d *testDiscovery) RequestRemote(n int) {}
func (d *testDiscovery) BadPeers() []string {
	d.Lock()
	defer d.Unlock()
	return d.bad
}
func (d *testDiscovery) GoodPeers() []AddressWithCapabilities { return []AddressWithCapabilities{} }

var defaultMessageHandler = func(t *testing.T, msg *Message) {}

type localPeer struct {
	netaddr        net.TCPAddr
	server         *Server
	version        *payload.Version
	lastBlockIndex uint32
	handshaked     int32 // TODO: use atomic.Bool after #2626.
	isFullNode     bool
	t              *testing.T
	messageHandler func(t *testing.T, msg *Message)
	pingSent       int
	getAddrSent    int
	droppedWith    atomic.Value
}

func newLocalPeer(t *testing.T, s *Server) *localPeer {
	naddr, _ := net.ResolveTCPAddr("tcp", "0.0.0.0:0")
	return &localPeer{
		t:              t,
		server:         s,
		netaddr:        *naddr,
		messageHandler: defaultMessageHandler,
	}
}

func (p *localPeer) ConnectionAddr() string {
	return p.netaddr.String()
}
func (p *localPeer) RemoteAddr() net.Addr {
	return &p.netaddr
}
func (p *localPeer) PeerAddr() net.Addr {
	return &p.netaddr
}
func (p *localPeer) StartProtocol() {}
func (p *localPeer) Disconnect(err error) {
	if p.droppedWith.Load() == nil {
		p.droppedWith.Store(err)
	}
	fmt.Println("peer dropped:", err)
	p.server.unregister <- peerDrop{p, err}
}

func (p *localPeer) BroadcastPacket(_ context.Context, m []byte) error {
	if len(m) == 0 {
		return errors.New("empty msg")
	}
	msg := &Message{}
	r := io.NewBinReaderFromBuf(m)
	for r.Len() > 0 {
		err := msg.Decode(r)
		if err == nil {
			p.messageHandler(p.t, msg)
		}
	}
	return nil
}
func (p *localPeer) EnqueueP2PMessage(msg *Message) error {
	return p.EnqueueHPMessage(msg)
}
func (p *localPeer) EnqueueP2PPacket(m []byte) error {
	return p.BroadcastPacket(context.TODO(), m)
}
func (p *localPeer) BroadcastHPPacket(ctx context.Context, m []byte) error {
	return p.BroadcastPacket(ctx, m)
}
func (p *localPeer) EnqueueHPMessage(msg *Message) error {
	p.messageHandler(p.t, msg)
	return nil
}
func (p *localPeer) EnqueueHPPacket(m []byte) error {
	return p.BroadcastPacket(context.TODO(), m)
}
func (p *localPeer) Version() *payload.Version {
	return p.version
}
func (p *localPeer) LastBlockIndex() uint32 {
	return p.lastBlockIndex
}
func (p *localPeer) HandleVersion(v *payload.Version) error {
	p.version = v
	return nil
}
func (p *localPeer) SendVersion() error {
	m, err := p.server.getVersionMsg(nil)
	if err != nil {
		return err
	}
	_ = p.EnqueueHPMessage(m)
	return nil
}
func (p *localPeer) SendVersionAck(m *Message) error {
	_ = p.EnqueueHPMessage(m)
	return nil
}
func (p *localPeer) HandleVersionAck() error {
	atomic.StoreInt32(&p.handshaked, 1)
	return nil
}
func (p *localPeer) SetPingTimer() {
	p.pingSent++
}
func (p *localPeer) HandlePing(ping *payload.Ping) error {
	p.lastBlockIndex = ping.LastBlockIndex
	return nil
}

func (p *localPeer) HandlePong(pong *payload.Ping) error {
	p.lastBlockIndex = pong.LastBlockIndex
	p.pingSent--
	return nil
}

func (p *localPeer) Handshaked() bool {
	return atomic.LoadInt32(&p.handshaked) != 0
}

func (p *localPeer) IsFullNode() bool {
	return p.isFullNode
}

func (p *localPeer) AddGetAddrSent() {
	p.getAddrSent++
}
func (p *localPeer) CanProcessAddr() bool {
	p.getAddrSent--
	return p.getAddrSent >= 0
}

func newTestServer(t *testing.T, serverConfig ServerConfig) *Server {
	return newTestServerWithCustomCfg(t, serverConfig, nil)
}

func newTestServerWithCustomCfg(t *testing.T, serverConfig ServerConfig, protocolCfg func(*config.Blockchain)) *Server {
	if len(serverConfig.Addresses) == 0 {
		// Normally it will be done by ApplicationConfiguration.GetAddresses().
		serverConfig.Addresses = []config.AnnounceableAddress{{Address: ":0"}}
	}
	s, err := newServerFromConstructors(serverConfig, fakechain.NewFakeChainWithCustomCfg(protocolCfg), new(fakechain.FakeStateSync), zaptest.NewLogger(t),
		newFakeTransp, newTestDiscovery)
	require.NoError(t, err)
	return s
}