network: limit message number from the same sender

pkg/config/application_config.go

@@ -30,4 +30,6 @@ type ApplicationConfiguration struct {
 	Oracle            OracleConfiguration     `yaml:"Oracle"`
 	P2PNotary         P2PNotary               `yaml:"P2PNotary"`
 	StateRoot         StateRoot               `yaml:"StateRoot"`
+	// ExtensiblePoolSize is the maximum amount of the extensible payloads from a single sender.
+	ExtensiblePoolSize int `yaml:"ExtensiblePoolSize"`
 }

pkg/network/extpool/pool.go

@@ -1,6 +1,7 @@
 package extpool
 
 import (
+	"container/list"
 	"errors"
 	"sync"
 
@@ -12,14 +13,23 @@ import (
 // Pool represents pool of extensible payloads.
 type Pool struct {
 	lock     sync.RWMutex
-	verified map[util.Uint256]*payload.Extensible
+	verified map[util.Uint256]*list.Element
+	senders  map[util.Uint160]*list.List
+	// singleCap represents maximum number of payloads from the single sender.
+	singleCap int
 	chain     blockchainer.Blockchainer
 }
 
 // New returns new payload pool using provided chain.
-func New(bc blockchainer.Blockchainer) *Pool {
+func New(bc blockchainer.Blockchainer, capacity int) *Pool {
+	if capacity <= 0 {
+		panic("invalid capacity")
+	}
+
 	return &Pool{
-		verified: make(map[util.Uint256]*payload.Extensible),
+		verified:  make(map[util.Uint256]*list.Element),
+		senders:   make(map[util.Uint160]*list.List),
+		singleCap: capacity,
 		chain:     bc,
 	}
 }
@@ -44,7 +54,17 @@ func (p *Pool) Add(e *payload.Extensible) (bool, error) {
 	if _, ok := p.verified[h]; ok {
 		return false, nil
 	}
-	p.verified[h] = e
+
+	lst, ok := p.senders[e.Sender]
+	if ok && lst.Len() >= p.singleCap {
+		value := lst.Remove(lst.Front())
+		delete(p.verified, value.(*payload.Extensible).Hash())
+	} else if !ok {
+		lst = list.New()
+		p.senders[e.Sender] = lst
+	}
+
+	p.verified[h] = lst.PushBack(e)
 	return true, nil
 }
 
@@ -72,7 +92,11 @@ func (p *Pool) Get(h util.Uint256) *payload.Extensible {
 	p.lock.RLock()
 	defer p.lock.RUnlock()
 
-	return p.verified[h]
+	elem, ok := p.verified[h]
+	if !ok {
+		return nil
+	}
+	return elem.Value.(*payload.Extensible)
 }
 
 const extensibleVerifyMaxGAS = 2000000
@@ -81,13 +105,27 @@ const extensibleVerifyMaxGAS = 2000000
 func (p *Pool) RemoveStale(index uint32) {
 	p.lock.Lock()
 	defer p.lock.Unlock()
-	for h, e := range p.verified {
+
+	for s, lst := range p.senders {
+		for elem := lst.Front(); elem != nil; {
+			e := elem.Value.(*payload.Extensible)
+			h := e.Hash()
+			old := elem
+			elem = elem.Next()
+
 			if e.ValidBlockEnd <= index || !p.chain.IsExtensibleAllowed(e.Sender) {
 				delete(p.verified, h)
+				lst.Remove(old)
 				continue
 			}
 			if err := p.chain.VerifyWitness(e.Sender, e, &e.Witness, extensibleVerifyMaxGAS); err != nil {
 				delete(p.verified, h)
+				lst.Remove(old)
+				continue
+			}
+		}
+		if lst.Len() == 0 {
+			delete(p.senders, s)
 		}
 	}
 }

pkg/network/extpool/pool_test.go

@@ -16,7 +16,7 @@ func TestAddGet(t *testing.T) {
 	bc := newTestChain()
 	bc.height = 10
 
-	p := New(bc)
+	p := New(bc, 100)
 	t.Run("invalid witness", func(t *testing.T) {
 		ep := &payload.Extensible{ValidBlockEnd: 100, Sender: util.Uint160{0x42}}
 		p.testAdd(t, false, errVerification, ep)
@@ -41,11 +41,52 @@ func TestAddGet(t *testing.T) {
 	})
 }
 
+func TestCapacityLimit(t *testing.T) {
+	bc := newTestChain()
+	bc.height = 10
+
+	t.Run("invalid capacity", func(t *testing.T) {
+		require.Panics(t, func() { New(bc, 0) })
+	})
+
+	p := New(bc, 3)
+
+	first := &payload.Extensible{ValidBlockEnd: 11}
+	p.testAdd(t, true, nil, first)
+
+	for _, height := range []uint32{12, 13} {
+		ep := &payload.Extensible{ValidBlockEnd: height}
+		p.testAdd(t, true, nil, ep)
+	}
+
+	require.NotNil(t, p.Get(first.Hash()))
+
+	ok, err := p.Add(&payload.Extensible{ValidBlockEnd: 14})
+	require.True(t, ok)
+	require.NoError(t, err)
+
+	require.Nil(t, p.Get(first.Hash()))
+}
+
+// This test checks that sender count is updated
+// when oldest payload is removed during `Add`.
+func TestDecreaseSenderOnEvict(t *testing.T) {
+	bc := newTestChain()
+	bc.height = 10
+
+	p := New(bc, 2)
+	senders := []util.Uint160{{1}, {2}, {3}}
+	for i := uint32(11); i < 17; i++ {
+		ep := &payload.Extensible{Sender: senders[i%3], ValidBlockEnd: i}
+		p.testAdd(t, true, nil, ep)
+	}
+}
+
 func TestRemoveStale(t *testing.T) {
 	bc := newTestChain()
 	bc.height = 10
 
-	p := New(bc)
+	p := New(bc, 100)
 	eps := []*payload.Extensible{
 		{ValidBlockEnd: 11},                             // small height
 		{ValidBlockEnd: 12},                             // good

pkg/network/server.go

@@ -33,6 +33,7 @@ const (
 	defaultMinPeers           = 5
 	defaultAttemptConnPeers   = 20
 	defaultMaxPeers           = 100
+	defaultExtensiblePoolSize = 20
 	maxBlockBatch             = 200
 	minPoolCount              = 30
 )
@@ -121,6 +122,12 @@ func newServerFromConstructors(config ServerConfig, chain blockchainer.Blockchai
 		return nil, errors.New("logger is a required parameter")
 	}
 
+	if config.ExtensiblePoolSize <= 0 {
+		config.ExtensiblePoolSize = defaultExtensiblePoolSize
+		log.Info("ExtensiblePoolSize is not set or wrong, using default value",
+			zap.Int("ExtensiblePoolSize", config.ExtensiblePoolSize))
+	}
+
 	s := &Server{
 		ServerConfig:      config,
 		chain:             chain,
@@ -132,7 +139,7 @@ func newServerFromConstructors(config ServerConfig, chain blockchainer.Blockchai
 		unregister:        make(chan peerDrop),
 		peers:             make(map[Peer]bool),
 		syncReached:       atomic.NewBool(false),
-		extensiblePool:    extpool.New(chain),
+		extensiblePool:    extpool.New(chain, config.ExtensiblePoolSize),
 		log:               log,
 		transactions:      make(chan *transaction.Transaction, 64),
 	}

pkg/network/server_config.go

@@ -78,6 +78,9 @@ type (
 
 		// StateRootCfg is stateroot module configuration.
 		StateRootCfg config.StateRoot
+
+		// ExtensiblePoolSize is size of the pool for extensible payloads from a single sender.
+		ExtensiblePoolSize int
 	}
 )
 
@@ -112,5 +115,6 @@ func NewServerConfig(cfg config.Config) ServerConfig {
 		OracleCfg:          appConfig.Oracle,
 		P2PNotaryCfg:       appConfig.P2PNotary,
 		StateRootCfg:       appConfig.StateRoot,
+		ExtensiblePoolSize: appConfig.ExtensiblePoolSize,
 	}
 }