diff --git a/pkg/network/helper_test.go b/pkg/network/helper_test.go index e7cc7a5f0..eb6649b60 100644 --- a/pkg/network/helper_test.go +++ b/pkg/network/helper_test.go @@ -286,6 +286,11 @@ func (p *localPeer) SendPing(m *Message) error { _ = p.EnqueueMessage(m) return nil } +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-- diff --git a/pkg/network/peer.go b/pkg/network/peer.go index 873347b4a..da7bd850f 100644 --- a/pkg/network/peer.go +++ b/pkg/network/peer.go @@ -67,6 +67,9 @@ type Peer interface { HandleVersion(*payload.Version) error HandleVersionAck() error + // HandlePing checks ping contents against Peer's state and updates it. + HandlePing(ping *payload.Ping) error + // HandlePong checks pong contents against Peer's state and updates it. HandlePong(pong *payload.Ping) error } diff --git a/pkg/network/server.go b/pkg/network/server.go index 3a926f19c..d3f22c9fe 100644 --- a/pkg/network/server.go +++ b/pkg/network/server.go @@ -453,6 +453,16 @@ func (s *Server) handleBlockCmd(p Peer, block *block.Block) error { // handlePing processes ping request. func (s *Server) handlePing(p Peer, ping *payload.Ping) error { + err := p.HandlePing(ping) + if err != nil { + return err + } + if s.chain.BlockHeight() < ping.LastBlockIndex { + err = s.requestBlocks(p) + if err != nil { + return err + } + } return p.EnqueueP2PMessage(NewMessage(CMDPong, payload.NewPing(s.chain.BlockHeight(), s.id))) } diff --git a/pkg/network/tcp_peer.go b/pkg/network/tcp_peer.go index aec25796e..64d884feb 100644 --- a/pkg/network/tcp_peer.go +++ b/pkg/network/tcp_peer.go @@ -431,6 +431,14 @@ func (p *TCPPeer) SendPing(msg *Message) error { return p.EnqueueMessage(msg) } +// HandlePing handles a ping message received from the peer. +func (p *TCPPeer) HandlePing(ping *payload.Ping) error { + p.lock.Lock() + defer p.lock.Unlock() + p.lastBlockIndex = ping.LastBlockIndex + return nil +} + // HandlePong handles a pong message received from the peer and does appropriate // accounting of outstanding pings and timeouts. func (p *TCPPeer) HandlePong(pong *payload.Ping) error {