From 2ce3c8b75f52f17da036a92a9f7d4d0ce274a34d Mon Sep 17 00:00:00 2001
From: Roman Khimov <roman@nspcc.ru>
Date: Wed, 25 Nov 2020 13:34:38 +0300
Subject: [PATCH] network: treat unsolicited addr commands as errors

See neo-project/neo#2097.
---
 pkg/network/helper_test.go |  9 +++++++++
 pkg/network/peer.go        |  9 +++++++++
 pkg/network/server.go      |  6 ++++++
 pkg/network/tcp_peer.go    | 17 +++++++++++++++++
 4 files changed, 41 insertions(+)

diff --git a/pkg/network/helper_test.go b/pkg/network/helper_test.go
index ea651d5e8..33aa4f2a6 100644
--- a/pkg/network/helper_test.go
+++ b/pkg/network/helper_test.go
@@ -227,6 +227,7 @@ type localPeer struct {
 	t              *testing.T
 	messageHandler func(t *testing.T, msg *Message)
 	pingSent       int
+	getAddrSent    int
 }
 
 func newLocalPeer(t *testing.T, s *Server) *localPeer {
@@ -323,6 +324,14 @@ 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 {
 	s := &Server{
 		ServerConfig: serverConfig,
diff --git a/pkg/network/peer.go b/pkg/network/peer.go
index da7bd850f..fb06c59d2 100644
--- a/pkg/network/peer.go
+++ b/pkg/network/peer.go
@@ -72,4 +72,13 @@ type Peer interface {
 
 	// HandlePong checks pong contents against Peer's state and updates it.
 	HandlePong(pong *payload.Ping) error
+
+	// AddGetAddrSent is to inform local peer context that a getaddr command
+	// is sent. The decision to send getaddr is server-wide, but it needs to be
+	// accounted for in peer's context, thus this method.
+	AddGetAddrSent()
+
+	// CanProcessAddr checks whether an addr command is expected to come from
+	// this peer and can be processed.
+	CanProcessAddr() bool
 }
diff --git a/pkg/network/server.go b/pkg/network/server.go
index fef50fe27..ed184e86f 100644
--- a/pkg/network/server.go
+++ b/pkg/network/server.go
@@ -680,6 +680,9 @@ func (s *Server) handleTxCmd(tx *transaction.Transaction) error {
 
 // handleAddrCmd will process received addresses.
 func (s *Server) handleAddrCmd(p Peer, addrs *payload.AddressList) error {
+	if !p.CanProcessAddr() {
+		return errors.New("unexpected addr received")
+	}
 	for _, a := range addrs.Addrs {
 		addr, err := a.GetTCPAddress()
 		if err == nil {
@@ -830,6 +833,9 @@ func (s *Server) iteratePeersWithSendMsg(msg *Message, send func(Peer, []byte) e
 		if peerOK != nil && !peerOK(peer) {
 			continue
 		}
+		if msg.Command == CMDGetAddr {
+			peer.AddGetAddrSent()
+		}
 		// Who cares about these messages anyway?
 		_ = send(peer, pkt)
 	}
diff --git a/pkg/network/tcp_peer.go b/pkg/network/tcp_peer.go
index ea383f5d9..67329c161 100644
--- a/pkg/network/tcp_peer.go
+++ b/pkg/network/tcp_peer.go
@@ -11,6 +11,7 @@ import (
 	"github.com/nspcc-dev/neo-go/pkg/io"
 	"github.com/nspcc-dev/neo-go/pkg/network/capability"
 	"github.com/nspcc-dev/neo-go/pkg/network/payload"
+	"go.uber.org/atomic"
 	"go.uber.org/zap"
 )
 
@@ -58,6 +59,9 @@ type TCPPeer struct {
 
 	wg sync.WaitGroup
 
+	// track outstanding getaddr requests.
+	getAddrSent atomic.Int32
+
 	// number of sent pings.
 	pingSent  int
 	pingTimer *time.Timer
@@ -455,3 +459,16 @@ func (p *TCPPeer) HandlePong(pong *payload.Ping) error {
 	p.lastBlockIndex = pong.LastBlockIndex
 	return nil
 }
+
+// AddGetAddrSent increments internal outstanding getaddr requests counter. The
+// peer can only send then one addr reply per getaddr request.
+func (p *TCPPeer) AddGetAddrSent() {
+	p.getAddrSent.Inc()
+}
+
+// CanProcessAddr decrements internal outstanding getaddr requests counter and
+// answers whether the addr command from the peer can be safely processed.
+func (p *TCPPeer) CanProcessAddr() bool {
+	v := p.getAddrSent.Dec()
+	return v >= 0
+}