From 8a58f8cf3a2d422dae3b430e80369102b846b447 Mon Sep 17 00:00:00 2001
From: David Hill <dhill@mindcry.org>
Date: Wed, 6 Apr 2016 17:56:48 -0400
Subject: [PATCH] peer: Implement sendheaders support (BIP0130).

This modifies the peer package to add support for the sendheaders
protocol message introduced by BIP0030.

NOTE: This does not add support to btcd itself. That requires the server
and sync code to make use of the new functionality exposed by these
changes.  As a result, btcd will still be using protocol version 70011.
---
 peer/peer.go      | 45 +++++++++++++++++++++++++++++++++++----------
 peer/peer_test.go |  9 ++++++++-
 server.go         |  1 +
 3 files changed, 44 insertions(+), 11 deletions(-)

diff --git a/peer/peer.go b/peer/peer.go
index 968f0fb3e..98e9f0086 100644
--- a/peer/peer.go
+++ b/peer/peer.go
@@ -26,7 +26,7 @@ import (
 
 const (
 	// MaxProtocolVersion is the max protocol version the peer supports.
-	MaxProtocolVersion = 70011
+	MaxProtocolVersion = wire.SendHeadersVersion
 
 	// outputBufferSize is the number of elements the output channels use.
 	outputBufferSize = 50
@@ -175,6 +175,10 @@ type MessageListeners struct {
 	// OnReject is invoked when a peer receives a reject bitcoin message.
 	OnReject func(p *Peer, msg *wire.MsgReject)
 
+	// OnSendHeaders is invoked when a peer receives a sendheaders bitcoin
+	// message.
+	OnSendHeaders func(p *Peer, msg *wire.MsgSendHeaders)
+
 	// OnRead is invoked when a peer receives a bitcoin message.  It
 	// consists of the number of bytes read, the message, and whether or not
 	// an error in the read occurred.  Typically, callers will opt to use
@@ -401,15 +405,16 @@ type Peer struct {
 	cfg     Config
 	inbound bool
 
-	flagsMtx        sync.Mutex // protects the peer flags below
-	na              *wire.NetAddress
-	id              int32
-	userAgent       string
-	services        wire.ServiceFlag
-	versionKnown    bool
-	protocolVersion uint32
-	versionSent     bool
-	verAckReceived  bool
+	flagsMtx             sync.Mutex // protects the peer flags below
+	na                   *wire.NetAddress
+	id                   int32
+	userAgent            string
+	services             wire.ServiceFlag
+	versionKnown         bool
+	protocolVersion      uint32
+	sendHeadersPreferred bool // peer sent a sendheaders message
+	versionSent          bool
+	verAckReceived       bool
 
 	knownInventory     *mruInventoryMap
 	prevGetBlocksMtx   sync.Mutex
@@ -725,6 +730,17 @@ func (p *Peer) StartingHeight() int32 {
 	return p.startingHeight
 }
 
+// WantsHeaders returns if the peer wants header messages instead of
+// inventory vectors for blocks.
+//
+// This function is safe for concurrent access.
+func (p *Peer) WantsHeaders() bool {
+	p.flagsMtx.Lock()
+	defer p.flagsMtx.Unlock()
+
+	return p.sendHeadersPreferred
+}
+
 // pushVersionMsg sends a version message to the connected peer using the
 // current state.
 func (p *Peer) pushVersionMsg() error {
@@ -1634,6 +1650,15 @@ out:
 				p.cfg.Listeners.OnReject(p, msg)
 			}
 
+		case *wire.MsgSendHeaders:
+			p.flagsMtx.Lock()
+			p.sendHeadersPreferred = true
+			p.flagsMtx.Unlock()
+
+			if p.cfg.Listeners.OnSendHeaders != nil {
+				p.cfg.Listeners.OnSendHeaders(p, msg)
+			}
+
 		default:
 			log.Debugf("Received unhandled message of type %v:",
 				rmsg.Command())
diff --git a/peer/peer_test.go b/peer/peer_test.go
index 317d34d57..cddcdfaf6 100644
--- a/peer/peer_test.go
+++ b/peer/peer_test.go
@@ -379,6 +379,9 @@ func TestPeerListeners(t *testing.T) {
 			OnReject: func(p *peer.Peer, msg *wire.MsgReject) {
 				ok <- msg
 			},
+			OnSendHeaders: func(p *peer.Peer, msg *wire.MsgSendHeaders) {
+				ok <- msg
+			},
 		},
 		UserAgentName:    "peer",
 		UserAgentVersion: "1.0",
@@ -496,9 +499,13 @@ func TestPeerListeners(t *testing.T) {
 		// only one version message is allowed
 		// only one verack message is allowed
 		{
-			"OnMsgReject",
+			"OnReject",
 			wire.NewMsgReject("block", wire.RejectDuplicate, "dupe block"),
 		},
+		{
+			"OnSendHeaders",
+			wire.NewMsgSendHeaders(),
+		},
 	}
 	t.Logf("Running %d tests", len(tests))
 	for _, test := range tests {
diff --git a/server.go b/server.go
index 611b8cf74..e54b77310 100644
--- a/server.go
+++ b/server.go
@@ -1453,6 +1453,7 @@ func newPeerConfig(sp *serverPeer) *peer.Config {
 		ChainParams:      sp.server.chainParams,
 		Services:         sp.server.services,
 		DisableRelayTx:   false,
+		ProtocolVersion:  70011,
 	}
 }