From 1f9a429579af7cc256f4521605243e45a6c5564e Mon Sep 17 00:00:00 2001
From: quest <quest@mac.com>
Date: Tue, 21 Jul 2020 23:18:09 -0700
Subject: [PATCH] peer: prevent last block height going backwards

---
 peer/peer.go      |  7 ++++-
 peer/peer_test.go | 67 +++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 73 insertions(+), 1 deletion(-)

diff --git a/peer/peer.go b/peer/peer.go
index 4f4879a79..5047a3e74 100644
--- a/peer/peer.go
+++ b/peer/peer.go
@@ -530,10 +530,15 @@ func (p *Peer) String() string {
 // This function is safe for concurrent access.
 func (p *Peer) UpdateLastBlockHeight(newHeight int32) {
 	p.statsMtx.Lock()
+	defer p.statsMtx.Unlock()
+
+	if newHeight <= p.lastBlock {
+		return
+	}
+
 	log.Tracef("Updating last block height of peer %v from %v to %v",
 		p.addr, p.lastBlock, newHeight)
 	p.lastBlock = newHeight
-	p.statsMtx.Unlock()
 }
 
 // UpdateLastAnnouncedBlock updates meta-data about the last block hash this
diff --git a/peer/peer_test.go b/peer/peer_test.go
index bba78b70e..a930b2f98 100644
--- a/peer/peer_test.go
+++ b/peer/peer_test.go
@@ -917,3 +917,70 @@ func TestDuplicateVersionMsg(t *testing.T) {
 		t.Fatal("peer did not disconnect")
 	}
 }
+
+// TestUpdateLastBlockHeight ensures the last block height is set properly
+// during the initial version negotiation and is only allowed to advance to
+// higher values via the associated update function.
+func TestUpdateLastBlockHeight(t *testing.T) {
+	// Create a pair of peers that are connected to each other using a fake
+	// connection and the remote peer starting at height 100.
+	const remotePeerHeight = 100
+	verack := make(chan struct{})
+	peerCfg := peer.Config{
+		Listeners: peer.MessageListeners{
+			OnVerAck: func(p *peer.Peer, msg *wire.MsgVerAck) {
+				verack <- struct{}{}
+			},
+		},
+		UserAgentName:          "peer",
+		UserAgentVersion:       "1.0",
+		ChainParams:            &chaincfg.MainNetParams,
+		Services:               0,
+		TstAllowSelfConnection: true,
+	}
+	remotePeerCfg := peerCfg
+	remotePeerCfg.NewestBlock = func() (*chainhash.Hash, int32, error) {
+		return &chainhash.Hash{}, remotePeerHeight, nil
+	}
+	inConn, outConn := pipe(
+		&conn{laddr: "10.0.0.1:9108", raddr: "10.0.0.2:9108"},
+		&conn{laddr: "10.0.0.2:9108", raddr: "10.0.0.1:9108"},
+	)
+	localPeer, err := peer.NewOutboundPeer(&peerCfg, inConn.laddr)
+	if err != nil {
+		t.Fatalf("NewOutboundPeer: unexpected err: %v\n", err)
+	}
+	localPeer.AssociateConnection(outConn)
+	inPeer := peer.NewInboundPeer(&remotePeerCfg)
+	inPeer.AssociateConnection(inConn)
+
+	// Wait for the veracks from the initial protocol version negotiation.
+	for i := 0; i < 2; i++ {
+		select {
+		case <-verack:
+		case <-time.After(time.Second):
+			t.Fatal("verack timeout")
+		}
+	}
+
+	// Ensure the latest block height starts at the value reported by the remote
+	// peer via its version message.
+	if height := localPeer.LastBlock(); height != remotePeerHeight {
+		t.Fatalf("wrong starting height - got %d, want %d", height,
+			remotePeerHeight)
+	}
+
+	// Ensure the latest block height is not allowed to go backwards.
+	localPeer.UpdateLastBlockHeight(remotePeerHeight - 1)
+	if height := localPeer.LastBlock(); height != remotePeerHeight {
+		t.Fatalf("height allowed to go backwards - got %d, want %d", height,
+			remotePeerHeight)
+	}
+
+	// Ensure the latest block height is allowed to advance.
+	localPeer.UpdateLastBlockHeight(remotePeerHeight + 1)
+	if height := localPeer.LastBlock(); height != remotePeerHeight+1 {
+		t.Fatalf("height not allowed to advance - got %d, want %d", height,
+			remotePeerHeight+1)
+	}
+}