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 +}