mempool: add a safety check, write tests for mempoolIDs (#3487)

* mempool: add a safety check, write tests for mempoolIDs

and document 65536 limit in the mempool reactor spec

follow-up to https://github.com/tendermint/tendermint/pull/2778

* rename the test

* fixes after Ismail's review

mempool/mempool.go

@@ -382,7 +382,10 @@ func (mem *Mempool) CheckTxWithInfo(tx types.Tx, cb func(*abci.Response), txInfo
 	if !mem.cache.Push(tx) {
 		// record the sender
 		e, ok := mem.txsMap[sha256.Sum256(tx)]
-		if ok { // tx may be in cache, but not in the mempool
+		// The check is needed because tx may be in cache, but not in the mempool.
+		// E.g. after we've committed a block, txs are removed from the mempool,
+		// but not from the cache.
+		if ok {
 			memTx := e.Value.(*mempoolTx)
 			if _, loaded := memTx.senders.LoadOrStore(txInfo.PeerID, true); loaded {
 				// TODO: consider punishing peer for dups,

mempool/reactor.go

@@ -2,6 +2,7 @@ package mempool
 
 import (
 	"fmt"
+	"math"
 	"reflect"
 	"sync"
 	"time"
@@ -26,6 +27,8 @@ const (
 	// UnknownPeerID is the peer ID to use when running CheckTx when there is
 	// no peer (e.g. RPC)
 	UnknownPeerID uint16 = 0
+
+	maxActiveIDs = math.MaxUint16
 )
 
 // MempoolReactor handles mempool tx broadcasting amongst peers.
@@ -45,7 +48,8 @@ type mempoolIDs struct {
 	activeIDs map[uint16]struct{} // used to check if a given peerID key is used, the value doesn't matter
 }
 
-// Reserve searches for the next unused ID and assignes it to the peer.
+// Reserve searches for the next unused ID and assignes it to the
+// peer.
 func (ids *mempoolIDs) ReserveForPeer(peer p2p.Peer) {
 	ids.mtx.Lock()
 	defer ids.mtx.Unlock()
@@ -58,6 +62,10 @@ func (ids *mempoolIDs) ReserveForPeer(peer p2p.Peer) {
 // nextPeerID returns the next unused peer ID to use.
 // This assumes that ids's mutex is already locked.
 func (ids *mempoolIDs) nextPeerID() uint16 {
+	if len(ids.activeIDs) == maxActiveIDs {
+		panic(fmt.Sprintf("node has maximum %d active IDs and wanted to get one more", maxActiveIDs))
+	}
+
 	_, idExists := ids.activeIDs[ids.nextID]
 	for idExists {
 		ids.nextID++
@@ -88,16 +96,20 @@ func (ids *mempoolIDs) GetForPeer(peer p2p.Peer) uint16 {
 	return ids.peerMap[peer.ID()]
 }
 
+func newMempoolIDs() *mempoolIDs {
+	return &mempoolIDs{
+		peerMap:   make(map[p2p.ID]uint16),
+		activeIDs: map[uint16]struct{}{0: {}},
+		nextID:    1, // reserve unknownPeerID(0) for mempoolReactor.BroadcastTx
+	}
+}
+
 // NewMempoolReactor returns a new MempoolReactor with the given config and mempool.
 func NewMempoolReactor(config *cfg.MempoolConfig, mempool *Mempool) *MempoolReactor {
 	memR := &MempoolReactor{
 		config:  config,
 		Mempool: mempool,
-		ids: &mempoolIDs{
-			peerMap:   make(map[p2p.ID]uint16),
-			activeIDs: map[uint16]struct{}{0: {}},
-			nextID:    1, // reserve unknownPeerID(0) for mempoolReactor.BroadcastTx
-		},
+		ids:     newMempoolIDs(),
 	}
 	memR.BaseReactor = *p2p.NewBaseReactor("MempoolReactor", memR)
 	return memR

mempool/reactor_test.go

@@ -2,6 +2,7 @@ package mempool
 
 import (
 	"fmt"
+	"net"
 	"sync"
 	"testing"
 	"time"
@@ -15,6 +16,7 @@ import (
 	cfg "github.com/tendermint/tendermint/config"
 	"github.com/tendermint/tendermint/libs/log"
 	"github.com/tendermint/tendermint/p2p"
+	"github.com/tendermint/tendermint/p2p/mock"
 	"github.com/tendermint/tendermint/proxy"
 	"github.com/tendermint/tendermint/types"
 )
@@ -189,3 +191,36 @@ func TestBroadcastTxForPeerStopsWhenReactorStops(t *testing.T) {
 	// i.e. broadcastTxRoutine finishes when reactor is stopped
 	leaktest.CheckTimeout(t, 10*time.Second)()
 }
+
+func TestMempoolIDsBasic(t *testing.T) {
+	ids := newMempoolIDs()
+
+	peer := mock.NewPeer(net.IP{127, 0, 0, 1})
+
+	ids.ReserveForPeer(peer)
+	assert.EqualValues(t, 1, ids.GetForPeer(peer))
+	ids.Reclaim(peer)
+
+	ids.ReserveForPeer(peer)
+	assert.EqualValues(t, 2, ids.GetForPeer(peer))
+	ids.Reclaim(peer)
+}
+
+func TestMempoolIDsPanicsIfNodeRequestsOvermaxActiveIDs(t *testing.T) {
+	if testing.Short() {
+		return
+	}
+
+	// 0 is already reserved for UnknownPeerID
+	ids := newMempoolIDs()
+
+	for i := 0; i < maxActiveIDs-1; i++ {
+		peer := mock.NewPeer(net.IP{127, 0, 0, 1})
+		ids.ReserveForPeer(peer)
+	}
+
+	assert.Panics(t, func() {
+		peer := mock.NewPeer(net.IP{127, 0, 0, 1})
+		ids.ReserveForPeer(peer)
+	})
+}