Merge pull request #10 from tendermint/pex-reactor-fixes-#9

Pex reactor fixes #9

Dockerfile

@@ -0,0 +1,13 @@
+FROM golang:latest
+
+RUN curl https://glide.sh/get | sh
+
+RUN mkdir -p /go/src/github.com/tendermint/go-p2p
+WORKDIR /go/src/github.com/tendermint/go-p2p
+
+COPY glide.yaml /go/src/github.com/tendermint/go-p2p/
+COPY glide.lock /go/src/github.com/tendermint/go-p2p/
+
+RUN glide install
+
+COPY . /go/src/github.com/tendermint/go-p2p

addrbook.go

@@ -15,7 +15,7 @@ import (
 	"time"
 
 	. "github.com/tendermint/go-common"
-	"github.com/tendermint/go-crypto"
+	crypto "github.com/tendermint/go-crypto"
 )
 
 const (
@@ -73,7 +73,12 @@ const (
 	serializationVersion = 1
 )
 
-/* AddrBook - concurrency safe peer address manager */
+const (
+	bucketTypeNew = 0x01
+	bucketTypeOld = 0x02
+)
+
+// AddrBook - concurrency safe peer address manager.
 type AddrBook struct {
 	BaseService
 
@@ -91,11 +96,7 @@ type AddrBook struct {
 	nNew              int
 }
 
-const (
-	bucketTypeNew = 0x01
-	bucketTypeOld = 0x02
-)
-
+// NewAddrBook creates a new address book.
 // Use Start to begin processing asynchronous address updates.
 func NewAddrBook(filePath string, routabilityStrict bool) *AddrBook {
 	am := &AddrBook{
@@ -125,6 +126,7 @@ func (a *AddrBook) init() {
 	}
 }
 
+// OnStart implements Service.
 func (a *AddrBook) OnStart() error {
 	a.BaseService.OnStart()
 	a.loadFromFile(a.filePath)
@@ -133,6 +135,7 @@ func (a *AddrBook) OnStart() error {
 	return nil
 }
 
+// OnStop implements Service.
 func (a *AddrBook) OnStop() {
 	a.BaseService.OnStop()
 }
@@ -254,15 +257,21 @@ func (a *AddrBook) MarkAttempt(addr *NetAddress) {
 	ka.markAttempt()
 }
 
+// MarkBad currently just ejects the address. In the future, consider
+// blacklisting.
 func (a *AddrBook) MarkBad(addr *NetAddress) {
+	a.RemoveAddress(addr)
+}
+
+// RemoveAddress removes the address from the book.
+func (a *AddrBook) RemoveAddress(addr *NetAddress) {
 	a.mtx.Lock()
 	defer a.mtx.Unlock()
 	ka := a.addrLookup[addr.String()]
 	if ka == nil {
 		return
 	}
-	// We currently just eject the address.
-	// In the future, consider blacklisting.
+	log.Info("Remove address from book", "addr", addr)
 	a.removeFromAllBuckets(ka)
 }
 
@@ -309,6 +318,10 @@ type addrBookJSON struct {
 }
 
 func (a *AddrBook) saveToFile(filePath string) {
+	log.Info("Saving AddrBook to file", "size", a.Size())
+
+	a.mtx.Lock()
+	defer a.mtx.Unlock()
 	// Compile Addrs
 	addrs := []*knownAddress{}
 	for _, ka := range a.addrLookup {
@@ -386,7 +399,6 @@ out:
 	for {
 		select {
 		case <-dumpAddressTicker.C:
-			log.Info("Saving AddrBook to file", "size", a.Size())
 			a.saveToFile(a.filePath)
 		case <-a.Quit:
 			break out

addrbook_test.go

@@ -148,3 +148,19 @@ func randIPv4Address(t *testing.T) *NetAddress {
 		}
 	}
 }
+
+func TestAddrBookRemoveAddress(t *testing.T) {
+	fname := createTempFileName("addrbook_test")
+	book := NewAddrBook(fname, true)
+
+	addr := randIPv4Address(t)
+	book.AddAddress(addr, addr)
+	assert.Equal(t, 1, book.Size())
+
+	book.RemoveAddress(addr)
+	assert.Equal(t, 0, book.Size())
+
+	nonExistingAddr := randIPv4Address(t)
+	book.RemoveAddress(nonExistingAddr)
+	assert.Equal(t, 0, book.Size())
+}

connection_test.go

@@ -39,7 +39,6 @@ func TestMConnectionSend(t *testing.T) {
 	assert.True(mconn.Send(0x01, msg))
 	// Note: subsequent Send/TrySend calls could pass because we are reading from
 	// the send queue in a separate goroutine.
-	assert.False(mconn.CanSend(0x01), "CanSend should return false because queue is full")
 	server.Read(make([]byte, len(msg)))
 	assert.True(mconn.CanSend(0x01))
 

pex_reactor.go

@@ -2,58 +2,86 @@ package p2p
 
 import (
 	"bytes"
-	"errors"
 	"fmt"
 	"math/rand"
 	"reflect"
 	"time"
 
-	. "github.com/tendermint/go-common"
+	cmn "github.com/tendermint/go-common"
 	wire "github.com/tendermint/go-wire"
 )
 
-var pexErrInvalidMessage = errors.New("Invalid PEX message")
-
 const (
-	PexChannel               = byte(0x00)
-	ensurePeersPeriodSeconds = 30
+	// PexChannel is a channel for PEX messages
+	PexChannel = byte(0x00)
+
+	// period to ensure peers connected
+	defaultEnsurePeersPeriod = 30 * time.Second
 	minNumOutboundPeers      = 10
 	maxPexMessageSize        = 1048576 // 1MB
+
+	// maximum messages one peer can send to us during `msgCountByPeerFlushInterval`
+	defaultMaxMsgCountByPeer    = 1000
+	msgCountByPeerFlushInterval = 1 * time.Hour
 )
 
-/*
-PEXReactor handles PEX (peer exchange) and ensures that an
-adequate number of peers are connected to the switch.
-*/
+// PEXReactor handles PEX (peer exchange) and ensures that an
+// adequate number of peers are connected to the switch.
+//
+// It uses `AddrBook` (address book) to store `NetAddress`es of the peers.
+//
+// ## Preventing abuse
+//
+// For now, it just limits the number of messages from one peer to
+// `defaultMaxMsgCountByPeer` messages per `msgCountByPeerFlushInterval` (1000
+// msg/hour).
+//
+// NOTE [2017-01-17]:
+//   Limiting is fine for now. Maybe down the road we want to keep track of the
+//   quality of peer messages so if peerA keeps telling us about peers we can't
+//   connect to then maybe we should care less about peerA. But I don't think
+//   that kind of complexity is priority right now.
 type PEXReactor struct {
 	BaseReactor
 
-	sw   *Switch
-	book *AddrBook
+	sw                *Switch
+	book              *AddrBook
+	ensurePeersPeriod time.Duration
+
+	// tracks message count by peer, so we can prevent abuse
+	msgCountByPeer    *cmn.CMap
+	maxMsgCountByPeer uint16
 }
 
-func NewPEXReactor(book *AddrBook) *PEXReactor {
-	pexR := &PEXReactor{
-		book: book,
+// NewPEXReactor creates new PEX reactor.
+func NewPEXReactor(b *AddrBook) *PEXReactor {
+	r := &PEXReactor{
+		book:              b,
+		ensurePeersPeriod: defaultEnsurePeersPeriod,
+		msgCountByPeer:    cmn.NewCMap(),
+		maxMsgCountByPeer: defaultMaxMsgCountByPeer,
 	}
-	pexR.BaseReactor = *NewBaseReactor(log, "PEXReactor", pexR)
-	return pexR
+	r.BaseReactor = *NewBaseReactor(log, "PEXReactor", r)
+	return r
 }
 
-func (pexR *PEXReactor) OnStart() error {
-	pexR.BaseReactor.OnStart()
-	pexR.book.Start()
-	go pexR.ensurePeersRoutine()
+// OnStart implements BaseService
+func (r *PEXReactor) OnStart() error {
+	r.BaseReactor.OnStart()
+	r.book.Start()
+	go r.ensurePeersRoutine()
+	go r.flushMsgCountByPeer()
 	return nil
 }
 
-func (pexR *PEXReactor) OnStop() {
-	pexR.BaseReactor.OnStop()
-	pexR.book.Stop()
+// OnStop implements BaseService
+func (r *PEXReactor) OnStop() {
+	r.BaseReactor.OnStop()
+	r.book.Stop()
 }
 
-// Implements Reactor
-func (pexR *PEXReactor) GetChannels() []*ChannelDescriptor {
+// GetChannels implements Reactor
+func (r *PEXReactor) GetChannels() []*ChannelDescriptor {
 	return []*ChannelDescriptor{
 		&ChannelDescriptor{
 			ID:                PexChannel,
@@ -63,37 +91,45 @@ func (pexR *PEXReactor) GetChannels() []*ChannelDescriptor {
 	}
 }
 
-// Implements Reactor
-func (pexR *PEXReactor) AddPeer(peer *Peer) {
-	// Add the peer to the address book
-	netAddr, err := NewNetAddressString(peer.ListenAddr)
-	if err != nil {
-		// this should never happen
-		log.Error("Error in AddPeer: invalid peer address", "addr", peer.ListenAddr, "error", err)
+// AddPeer implements Reactor by adding peer to the address book (if inbound)
+// or by requesting more addresses (if outbound).
+func (r *PEXReactor) AddPeer(p *Peer) {
+	if p.IsOutbound() {
+		// For outbound peers, the address is already in the books.
+		// Either it was added in DialSeeds or when we
+		// received the peer's address in r.Receive
+		if r.book.NeedMoreAddrs() {
+			r.RequestPEX(p)
+		}
+	} else { // For inbound connections, the peer is its own source
+		addr, err := NewNetAddressString(p.ListenAddr)
+		if err != nil {
+			// this should never happen
+			log.Error("Error in AddPeer: invalid peer address", "addr", p.ListenAddr, "error", err)
+			return
+		}
+		r.book.AddAddress(addr, addr)
+	}
+}
+
+// RemovePeer implements Reactor.
+func (r *PEXReactor) RemovePeer(p *Peer, reason interface{}) {
+	// If we aren't keeping track of local temp data for each peer here, then we
+	// don't have to do anything.
+}
+
+// Receive implements Reactor by handling incoming PEX messages.
+func (r *PEXReactor) Receive(chID byte, src *Peer, msgBytes []byte) {
+	srcAddr := src.Connection().RemoteAddress
+	srcAddrStr := srcAddr.String()
+
+	r.IncrementMsgCountForPeer(srcAddrStr)
+	if r.ReachedMaxMsgCountForPeer(srcAddrStr) {
+		log.Warn("Maximum number of messages reached for peer", "peer", srcAddrStr)
+		// TODO remove src from peers?
 		return
 	}
 
-	if peer.IsOutbound() {
-		if pexR.book.NeedMoreAddrs() {
-			pexR.RequestPEX(peer)
-		}
-	} else {
-		// For inbound connections, the peer is its own source
-		// (For outbound peers, the address is already in the books)
-		pexR.book.AddAddress(netAddr, netAddr)
-	}
-}
-
-// Implements Reactor
-func (pexR *PEXReactor) RemovePeer(peer *Peer, reason interface{}) {
-	// TODO
-}
-
-// Implements Reactor
-// Handles incoming PEX messages.
-func (pexR *PEXReactor) Receive(chID byte, src *Peer, msgBytes []byte) {
-
-	// decode message
 	_, msg, err := DecodeMessage(msgBytes)
 	if err != nil {
 		log.Warn("Error decoding message", "error", err)
@@ -104,87 +140,127 @@ func (pexR *PEXReactor) Receive(chID byte, src *Peer, msgBytes []byte) {
 	switch msg := msg.(type) {
 	case *pexRequestMessage:
 		// src requested some peers.
-		// TODO: prevent abuse.
-		pexR.SendAddrs(src, pexR.book.GetSelection())
+		r.SendAddrs(src, r.book.GetSelection())
 	case *pexAddrsMessage:
 		// We received some peer addresses from src.
-		// TODO: prevent abuse.
 		// (We don't want to get spammed with bad peers)
-		srcAddr := src.Connection().RemoteAddress
 		for _, addr := range msg.Addrs {
 			if addr != nil {
-				pexR.book.AddAddress(addr, srcAddr)
+				r.book.AddAddress(addr, srcAddr)
 			}
 		}
 	default:
-		log.Warn(Fmt("Unknown message type %v", reflect.TypeOf(msg)))
+		log.Warn(fmt.Sprintf("Unknown message type %v", reflect.TypeOf(msg)))
 	}
 }
 
-// Asks peer for more addresses.
-func (pexR *PEXReactor) RequestPEX(peer *Peer) {
-	peer.Send(PexChannel, struct{ PexMessage }{&pexRequestMessage{}})
+// RequestPEX asks peer for more addresses.
+func (r *PEXReactor) RequestPEX(p *Peer) {
+	p.Send(PexChannel, struct{ PexMessage }{&pexRequestMessage{}})
 }
 
-func (pexR *PEXReactor) SendAddrs(peer *Peer, addrs []*NetAddress) {
-	peer.Send(PexChannel, struct{ PexMessage }{&pexAddrsMessage{Addrs: addrs}})
+// SendAddrs sends addrs to the peer.
+func (r *PEXReactor) SendAddrs(p *Peer, addrs []*NetAddress) {
+	p.Send(PexChannel, struct{ PexMessage }{&pexAddrsMessage{Addrs: addrs}})
+}
+
+// SetEnsurePeersPeriod sets period to ensure peers connected.
+func (r *PEXReactor) SetEnsurePeersPeriod(d time.Duration) {
+	r.ensurePeersPeriod = d
+}
+
+// SetMaxMsgCountByPeer sets maximum messages one peer can send to us during 'msgCountByPeerFlushInterval'.
+func (r *PEXReactor) SetMaxMsgCountByPeer(v uint16) {
+	r.maxMsgCountByPeer = v
+}
+
+// ReachedMaxMsgCountForPeer returns true if we received too many
+// messages from peer with address `addr`.
+// NOTE: assumes the value in the CMap is non-nil
+func (r *PEXReactor) ReachedMaxMsgCountForPeer(addr string) bool {
+	return r.msgCountByPeer.Get(addr).(uint16) >= r.maxMsgCountByPeer
+}
+
+// Increment or initialize the msg count for the peer in the CMap
+func (r *PEXReactor) IncrementMsgCountForPeer(addr string) {
+	var count uint16
+	countI := r.msgCountByPeer.Get(addr)
+	if countI != nil {
+		count = countI.(uint16)
+	}
+	count++
+	r.msgCountByPeer.Set(addr, count)
 }
 
 // Ensures that sufficient peers are connected. (continuous)
-func (pexR *PEXReactor) ensurePeersRoutine() {
+func (r *PEXReactor) ensurePeersRoutine() {
 	// Randomize when routine starts
-	time.Sleep(time.Duration(rand.Int63n(500*ensurePeersPeriodSeconds)) * time.Millisecond)
+	ensurePeersPeriodMs := r.ensurePeersPeriod.Nanoseconds() / 1e6
+	time.Sleep(time.Duration(rand.Int63n(ensurePeersPeriodMs)) * time.Millisecond)
 
 	// fire once immediately.
-	pexR.ensurePeers()
+	r.ensurePeers()
+
 	// fire periodically
-	timer := NewRepeatTimer("pex", ensurePeersPeriodSeconds*time.Second)
-FOR_LOOP:
+	ticker := time.NewTicker(r.ensurePeersPeriod)
+
 	for {
 		select {
-		case <-timer.Ch:
-			pexR.ensurePeers()
-		case <-pexR.Quit:
-			break FOR_LOOP
+		case <-ticker.C:
+			r.ensurePeers()
+		case <-r.Quit:
+			ticker.Stop()
+			return
 		}
 	}
-
-	// Cleanup
-	timer.Stop()
 }
 
-// Ensures that sufficient peers are connected. (once)
-func (pexR *PEXReactor) ensurePeers() {
-	numOutPeers, _, numDialing := pexR.Switch.NumPeers()
+// ensurePeers ensures that sufficient peers are connected. (once)
+//
+// Old bucket / New bucket are arbitrary categories to denote whether an
+// address is vetted or not, and this needs to be determined over time via a
+// heuristic that we haven't perfected yet, or, perhaps is manually edited by
+// the node operator. It should not be used to compute what addresses are
+// already connected or not.
+//
+// TODO Basically, we need to work harder on our good-peer/bad-peer marking.
+// What we're currently doing in terms of marking good/bad peers is just a
+// placeholder. It should not be the case that an address becomes old/vetted
+// upon a single successful connection.
+func (r *PEXReactor) ensurePeers() {
+	numOutPeers, _, numDialing := r.Switch.NumPeers()
 	numToDial := minNumOutboundPeers - (numOutPeers + numDialing)
 	log.Info("Ensure peers", "numOutPeers", numOutPeers, "numDialing", numDialing, "numToDial", numToDial)
 	if numToDial <= 0 {
 		return
 	}
-	toDial := NewCMap()
+
+	toDial := make(map[string]*NetAddress)
 
 	// Try to pick numToDial addresses to dial.
-	// TODO: improve logic.
 	for i := 0; i < numToDial; i++ {
-		newBias := MinInt(numOutPeers, 8)*10 + 10
+		// The purpose of newBias is to first prioritize old (more vetted) peers
+		// when we have few connections, but to allow for new (less vetted) peers
+		// if we already have many connections. This algorithm isn't perfect, but
+		// it somewhat ensures that we prioritize connecting to more-vetted
+		// peers.
+		newBias := cmn.MinInt(numOutPeers, 8)*10 + 10
 		var picked *NetAddress
 		// Try to fetch a new peer 3 times.
 		// This caps the maximum number of tries to 3 * numToDial.
 		for j := 0; j < 3; j++ {
-			try := pexR.book.PickAddress(newBias)
+			try := r.book.PickAddress(newBias)
 			if try == nil {
 				break
 			}
-			alreadySelected := toDial.Has(try.IP.String())
-			alreadyDialing := pexR.Switch.IsDialing(try)
-			alreadyConnected := pexR.Switch.Peers().Has(try.IP.String())
+			_, alreadySelected := toDial[try.IP.String()]
+			alreadyDialing := r.Switch.IsDialing(try)
+			alreadyConnected := r.Switch.Peers().Has(try.IP.String())
 			if alreadySelected || alreadyDialing || alreadyConnected {
-				/*
-					log.Info("Cannot dial address", "addr", try,
-						"alreadySelected", alreadySelected,
-						"alreadyDialing", alreadyDialing,
-						"alreadyConnected", alreadyConnected)
-				*/
+				// log.Info("Cannot dial address", "addr", try,
+				// 	"alreadySelected", alreadySelected,
+				// 	"alreadyDialing", alreadyDialing,
+				//  "alreadyConnected", alreadyConnected)
 				continue
 			} else {
 				log.Info("Will dial address", "addr", try)
@@ -195,26 +271,40 @@ func (pexR *PEXReactor) ensurePeers() {
 		if picked == nil {
 			continue
 		}
-		toDial.Set(picked.IP.String(), picked)
+		toDial[picked.IP.String()] = picked
 	}
 
 	// Dial picked addresses
-	for _, item := range toDial.Values() {
+	for _, item := range toDial {
 		go func(picked *NetAddress) {
-			_, err := pexR.Switch.DialPeerWithAddress(picked, false)
+			_, err := r.Switch.DialPeerWithAddress(picked, false)
 			if err != nil {
-				pexR.book.MarkAttempt(picked)
+				r.book.MarkAttempt(picked)
 			}
-		}(item.(*NetAddress))
+		}(item)
 	}
 
 	// If we need more addresses, pick a random peer and ask for more.
-	if pexR.book.NeedMoreAddrs() {
-		if peers := pexR.Switch.Peers().List(); len(peers) > 0 {
+	if r.book.NeedMoreAddrs() {
+		if peers := r.Switch.Peers().List(); len(peers) > 0 {
 			i := rand.Int() % len(peers)
 			peer := peers[i]
 			log.Info("No addresses to dial. Sending pexRequest to random peer", "peer", peer)
-			pexR.RequestPEX(peer)
+			r.RequestPEX(peer)
+		}
+	}
+}
+
+func (r *PEXReactor) flushMsgCountByPeer() {
+	ticker := time.NewTicker(msgCountByPeerFlushInterval)
+
+	for {
+		select {
+		case <-ticker.C:
+			r.msgCountByPeer.Clear()
+		case <-r.Quit:
+			ticker.Stop()
+			return
 		}
 	}
 }
@@ -227,6 +317,8 @@ const (
 	msgTypeAddrs   = byte(0x02)
 )
 
+// PexMessage is a primary type for PEX messages. Underneath, it could contain
+// either pexRequestMessage, or pexAddrsMessage messages.
 type PexMessage interface{}
 
 var _ = wire.RegisterInterface(
@@ -235,6 +327,7 @@ var _ = wire.RegisterInterface(
 	wire.ConcreteType{&pexAddrsMessage{}, msgTypeAddrs},
 )
 
+// DecodeMessage implements interface registered above.
 func DecodeMessage(bz []byte) (msgType byte, msg PexMessage, err error) {
 	msgType = bz[0]
 	n := new(int)

pex_reactor_test.go

@@ -0,0 +1,163 @@
+package p2p
+
+import (
+	"io/ioutil"
+	"math/rand"
+	"os"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+	cmn "github.com/tendermint/go-common"
+	wire "github.com/tendermint/go-wire"
+)
+
+func TestPEXReactorBasic(t *testing.T) {
+	assert, require := assert.New(t), require.New(t)
+
+	dir, err := ioutil.TempDir("", "pex_reactor")
+	require.Nil(err)
+	defer os.RemoveAll(dir)
+	book := NewAddrBook(dir+"addrbook.json", true)
+
+	r := NewPEXReactor(book)
+
+	assert.NotNil(r)
+	assert.NotEmpty(r.GetChannels())
+}
+
+func TestPEXReactorAddRemovePeer(t *testing.T) {
+	assert, require := assert.New(t), require.New(t)
+
+	dir, err := ioutil.TempDir("", "pex_reactor")
+	require.Nil(err)
+	defer os.RemoveAll(dir)
+	book := NewAddrBook(dir+"addrbook.json", true)
+
+	r := NewPEXReactor(book)
+
+	size := book.Size()
+	peer := createRandomPeer(false)
+
+	r.AddPeer(peer)
+	assert.Equal(size+1, book.Size())
+
+	r.RemovePeer(peer, "peer not available")
+	assert.Equal(size+1, book.Size())
+
+	outboundPeer := createRandomPeer(true)
+
+	r.AddPeer(outboundPeer)
+	assert.Equal(size+1, book.Size(), "outbound peers should not be added to the address book")
+
+	r.RemovePeer(outboundPeer, "peer not available")
+	assert.Equal(size+1, book.Size())
+}
+
+func TestPEXReactorRunning(t *testing.T) {
+	require := require.New(t)
+
+	N := 3
+	switches := make([]*Switch, N)
+
+	dir, err := ioutil.TempDir("", "pex_reactor")
+	require.Nil(err)
+	defer os.RemoveAll(dir)
+	book := NewAddrBook(dir+"addrbook.json", false)
+
+	// create switches
+	for i := 0; i < N; i++ {
+		switches[i] = makeSwitch(i, "127.0.0.1", "123.123.123", func(i int, sw *Switch) *Switch {
+			r := NewPEXReactor(book)
+			r.SetEnsurePeersPeriod(250 * time.Millisecond)
+			sw.AddReactor("pex", r)
+			return sw
+		})
+	}
+
+	// fill the address book and add listeners
+	for _, s := range switches {
+		addr, _ := NewNetAddressString(s.NodeInfo().ListenAddr)
+		book.AddAddress(addr, addr)
+		s.AddListener(NewDefaultListener("tcp", s.NodeInfo().ListenAddr, true))
+	}
+
+	// start switches
+	for _, s := range switches {
+		_, err := s.Start() // start switch and reactors
+		require.Nil(err)
+	}
+
+	time.Sleep(1 * time.Second)
+
+	// check peers are connected after some time
+	for _, s := range switches {
+		outbound, inbound, _ := s.NumPeers()
+		if outbound+inbound == 0 {
+			t.Errorf("%v expected to be connected to at least one peer", s.NodeInfo().ListenAddr)
+		}
+	}
+
+	// stop them
+	for _, s := range switches {
+		s.Stop()
+	}
+}
+
+func TestPEXReactorReceive(t *testing.T) {
+	assert, require := assert.New(t), require.New(t)
+
+	dir, err := ioutil.TempDir("", "pex_reactor")
+	require.Nil(err)
+	defer os.RemoveAll(dir)
+	book := NewAddrBook(dir+"addrbook.json", true)
+
+	r := NewPEXReactor(book)
+
+	peer := createRandomPeer(false)
+
+	size := book.Size()
+	netAddr, _ := NewNetAddressString(peer.ListenAddr)
+	addrs := []*NetAddress{netAddr}
+	msg := wire.BinaryBytes(struct{ PexMessage }{&pexAddrsMessage{Addrs: addrs}})
+	r.Receive(PexChannel, peer, msg)
+	assert.Equal(size+1, book.Size())
+
+	msg = wire.BinaryBytes(struct{ PexMessage }{&pexRequestMessage{}})
+	r.Receive(PexChannel, peer, msg)
+}
+
+func TestPEXReactorAbuseFromPeer(t *testing.T) {
+	assert, require := assert.New(t), require.New(t)
+
+	dir, err := ioutil.TempDir("", "pex_reactor")
+	require.Nil(err)
+	defer os.RemoveAll(dir)
+	book := NewAddrBook(dir+"addrbook.json", true)
+
+	r := NewPEXReactor(book)
+	r.SetMaxMsgCountByPeer(5)
+
+	peer := createRandomPeer(false)
+
+	msg := wire.BinaryBytes(struct{ PexMessage }{&pexRequestMessage{}})
+	for i := 0; i < 10; i++ {
+		r.Receive(PexChannel, peer, msg)
+	}
+
+	assert.True(r.ReachedMaxMsgCountForPeer(peer.ListenAddr))
+}
+
+func createRandomPeer(outbound bool) *Peer {
+	addr := cmn.Fmt("%v.%v.%v.%v:46656", rand.Int()%256, rand.Int()%256, rand.Int()%256, rand.Int()%256)
+	netAddr, _ := NewNetAddressString(addr)
+	return &Peer{
+		Key: cmn.RandStr(12),
+		NodeInfo: &NodeInfo{
+			ListenAddr: addr,
+		},
+		outbound: outbound,
+		mconn:    &MConnection{RemoteAddress: netAddr},
+	}
+}

switch.go

@@ -531,10 +531,12 @@ func makeSwitch(i int, network, version string, initSwitch func(int, *Switch) *S
 	// TODO: let the config be passed in?
 	s := initSwitch(i, NewSwitch(cfg.NewMapConfig(nil)))
 	s.SetNodeInfo(&NodeInfo{
-		PubKey:  privKey.PubKey().(crypto.PubKeyEd25519),
-		Moniker: Fmt("switch%d", i),
-		Network: network,
-		Version: version,
+		PubKey:     privKey.PubKey().(crypto.PubKeyEd25519),
+		Moniker:    Fmt("switch%d", i),
+		Network:    network,
+		Version:    version,
+		RemoteAddr: Fmt("%v:%v", network, rand.Intn(64512)+1023),
+		ListenAddr: Fmt("%v:%v", network, rand.Intn(64512)+1023),
 	})
 	s.SetNodePrivKey(privKey)
 	return s