limit the number of protocols we store per peer

pstoreds/ds_test.go

@@ -7,6 +7,8 @@ import (
 	"testing"
 	"time"
 
+	"github.com/stretchr/testify/require"
+
 	ds "github.com/ipfs/go-datastore"
 	badger "github.com/ipfs/go-ds-badger"
 	leveldb "github.com/ipfs/go-ds-leveldb"
@@ -27,6 +29,18 @@ func TestDsPeerstore(t *testing.T) {
 		t.Run(name, func(t *testing.T) {
 			pt.TestPeerstore(t, peerstoreFactory(t, dsFactory, DefaultOpts()))
 		})
+
+		t.Run("protobook limits", func(t *testing.T) {
+			const limit = 10
+			opts := DefaultOpts()
+			opts.MaxProtocols = limit
+			ds, close := dsFactory(t)
+			defer close()
+			ps, err := NewPeerstore(context.Background(), ds, opts)
+			require.NoError(t, err)
+			defer ps.Close()
+			pt.TestPeerstoreProtoStoreLimits(t, ps, limit)
+		})
 	}
 }
 

pstoreds/peerstore.go

@@ -6,13 +6,13 @@ import (
 	"io"
 	"time"
 
-	base32 "github.com/multiformats/go-base32"
+	"github.com/multiformats/go-base32"
 
 	ds "github.com/ipfs/go-datastore"
-	query "github.com/ipfs/go-datastore/query"
+	"github.com/ipfs/go-datastore/query"
 
-	peer "github.com/libp2p/go-libp2p-core/peer"
-	peerstore "github.com/libp2p/go-libp2p-core/peerstore"
+	"github.com/libp2p/go-libp2p-core/peer"
+	"github.com/libp2p/go-libp2p-core/peerstore"
 	pstore "github.com/libp2p/go-libp2p-peerstore"
 )
 
@@ -21,6 +21,9 @@ type Options struct {
 	// The size of the in-memory cache. A value of 0 or lower disables the cache.
 	CacheSize uint
 
+	// MaxProtocols is the maximum number of protocols we store for one peer.
+	MaxProtocols int
+
 	// Sweep interval to purge expired addresses from the datastore. If this is a zero value, GC will not run
 	// automatically, but it'll be available on demand via explicit calls.
 	GCPurgeInterval time.Duration
@@ -37,12 +40,14 @@ type Options struct {
 // DefaultOpts returns the default options for a persistent peerstore, with the full-purge GC algorithm:
 //
 // * Cache size: 1024.
+// * MaxProtocols: 1024.
 // * GC purge interval: 2 hours.
 // * GC lookahead interval: disabled.
 // * GC initial delay: 60 seconds.
 func DefaultOpts() Options {
 	return Options{
 		CacheSize:           1024,
+		MaxProtocols:        1024,
 		GCPurgeInterval:     2 * time.Hour,
 		GCLookaheadInterval: 0,
 		GCInitialDelay:      60 * time.Second,
@@ -75,7 +80,10 @@ func NewPeerstore(ctx context.Context, store ds.Batching, opts Options) (*pstore
 		return nil, err
 	}
 
-	protoBook := NewProtoBook(peerMetadata)
+	protoBook, err := NewProtoBook(peerMetadata, WithMaxProtocols(opts.MaxProtocols))
+	if err != nil {
+		return nil, err
+	}
 
 	ps := &pstoreds{
 		Metrics:        pstore.NewMetrics(),

pstoreds/protobook.go

@@ -1,10 +1,11 @@
 package pstoreds
 
 import (
+	"errors"
 	"fmt"
 	"sync"
 
-	peer "github.com/libp2p/go-libp2p-core/peer"
+	"github.com/libp2p/go-libp2p-core/peer"
 
 	pstore "github.com/libp2p/go-libp2p-core/peerstore"
 )
@@ -19,39 +20,62 @@ func (s *protoSegments) get(p peer.ID) *protoSegment {
 	return s[byte(p[len(p)-1])]
 }
 
+var errTooManyProtocols = errors.New("too many protocols")
+
+type ProtoBookOption func(*dsProtoBook) error
+
+func WithMaxProtocols(num int) ProtoBookOption {
+	return func(pb *dsProtoBook) error {
+		pb.maxProtos = num
+		return nil
+	}
+}
+
 type dsProtoBook struct {
-	segments protoSegments
-	meta     pstore.PeerMetadata
+	segments  protoSegments
+	meta      pstore.PeerMetadata
+	maxProtos int
 }
 
 var _ pstore.ProtoBook = (*dsProtoBook)(nil)
 
-func NewProtoBook(meta pstore.PeerMetadata) *dsProtoBook {
-	return &dsProtoBook{
+func NewProtoBook(meta pstore.PeerMetadata, opts ...ProtoBookOption) (*dsProtoBook, error) {
+	pb := &dsProtoBook{
 		meta: meta,
 		segments: func() (ret protoSegments) {
 			for i := range ret {
 				ret[i] = &protoSegment{}
 			}
 			return ret
 		}(),
+		maxProtos: 1024,
 	}
+
+	for _, opt := range opts {
+		if err := opt(pb); err != nil {
+			return nil, err
+		}
+	}
+	return pb, nil
 }
 
 func (pb *dsProtoBook) SetProtocols(p peer.ID, protos ...string) error {
 	if err := p.Validate(); err != nil {
 		return err
 	}
-
-	s := pb.segments.get(p)
-	s.Lock()
-	defer s.Unlock()
+	if len(protos) > pb.maxProtos {
+		return errTooManyProtocols
+	}
 
 	protomap := make(map[string]struct{}, len(protos))
 	for _, proto := range protos {
 		protomap[proto] = struct{}{}
 	}
 
+	s := pb.segments.get(p)
+	s.Lock()
+	defer s.Unlock()
+
 	return pb.meta.Put(p, "protocols", protomap)
 }
 
@@ -68,6 +92,9 @@ func (pb *dsProtoBook) AddProtocols(p peer.ID, protos ...string) error {
 	if err != nil {
 		return err
 	}
+	if len(pmap)+len(protos) > pb.maxProtos {
+		return errTooManyProtocols
+	}
 
 	for _, proto := range protos {
 		pmap[proto] = struct{}{}

pstoremem/inmem_test.go

@@ -3,6 +3,8 @@ package pstoremem
 import (
 	"testing"
 
+	"github.com/stretchr/testify/require"
+
 	pstore "github.com/libp2p/go-libp2p-core/peerstore"
 	pt "github.com/libp2p/go-libp2p-peerstore/test"
 
@@ -13,42 +15,56 @@ func TestFuzzInMemoryPeerstore(t *testing.T) {
 	// Just create and close a bunch of peerstores. If this leaks, we'll
 	// catch it in the leak check below.
 	for i := 0; i < 100; i++ {
-		ps := NewPeerstore()
+		ps, err := NewPeerstore()
+		require.NoError(t, err)
 		ps.Close()
 	}
 }
 
 func TestInMemoryPeerstore(t *testing.T) {
 	pt.TestPeerstore(t, func() (pstore.Peerstore, func()) {
-		ps := NewPeerstore()
+		ps, err := NewPeerstore()
+		require.NoError(t, err)
 		return ps, func() { ps.Close() }
 	})
 }
 
+func TestPeerstoreProtoStoreLimits(t *testing.T) {
+	const limit = 10
+	ps, err := NewPeerstore(WithMaxProtocols(limit))
+	require.NoError(t, err)
+	defer ps.Close()
+	pt.TestPeerstoreProtoStoreLimits(t, ps, limit)
+}
+
 func TestInMemoryAddrBook(t *testing.T) {
 	pt.TestAddrBook(t, func() (pstore.AddrBook, func()) {
-		ps := NewPeerstore()
+		ps, err := NewPeerstore()
+		require.NoError(t, err)
 		return ps, func() { ps.Close() }
 	})
 }
 
 func TestInMemoryKeyBook(t *testing.T) {
 	pt.TestKeyBook(t, func() (pstore.KeyBook, func()) {
-		ps := NewPeerstore()
+		ps, err := NewPeerstore()
+		require.NoError(t, err)
 		return ps, func() { ps.Close() }
 	})
 }
 
 func BenchmarkInMemoryPeerstore(b *testing.B) {
 	pt.BenchmarkPeerstore(b, func() (pstore.Peerstore, func()) {
-		ps := NewPeerstore()
+		ps, err := NewPeerstore()
+		require.NoError(b, err)
 		return ps, func() { ps.Close() }
 	}, "InMem")
 }
 
 func BenchmarkInMemoryKeyBook(b *testing.B) {
 	pt.BenchmarkKeyBook(b, func() (pstore.KeyBook, func()) {
-		ps := NewPeerstore()
+		ps, err := NewPeerstore()
+		require.NoError(b, err)
 		return ps, func() { ps.Close() }
 	})
 }

pstoremem/peerstore.go

@@ -2,10 +2,11 @@ package pstoremem
 
 import (
 	"fmt"
+	"io"
+
 	"github.com/libp2p/go-libp2p-core/peer"
 	"github.com/libp2p/go-libp2p-core/peerstore"
 	pstore "github.com/libp2p/go-libp2p-peerstore"
-	"io"
 )
 
 type pstoremem struct {
@@ -17,15 +18,26 @@ type pstoremem struct {
 	*memoryPeerMetadata
 }
 
+func WithMaxProtocols(num int) Option {
+	return func(pb *memoryProtoBook) error {
+		pb.maxProtos = num
+		return nil
+	}
+}
+
 // NewPeerstore creates an in-memory threadsafe collection of peers.
-func NewPeerstore() *pstoremem {
+func NewPeerstore(opts ...Option) (*pstoremem, error) {
+	pb, err := NewProtoBook(opts...)
+	if err != nil {
+		return nil, err
+	}
 	return &pstoremem{
 		Metrics:            pstore.NewMetrics(),
 		memoryKeyBook:      NewKeyBook(),
 		memoryAddrBook:     NewAddrBook(),
-		memoryProtoBook:    NewProtoBook(),
+		memoryProtoBook:    pb,
 		memoryPeerMetadata: NewPeerMetadata(),
-	}
+	}, nil
 }
 
 func (ps *pstoremem) Close() (err error) {

pstoremem/protobook.go

@@ -1,9 +1,10 @@
 package pstoremem
 
 import (
+	"errors"
 	"sync"
 
-	peer "github.com/libp2p/go-libp2p-core/peer"
+	"github.com/libp2p/go-libp2p-core/peer"
 
 	pstore "github.com/libp2p/go-libp2p-core/peerstore"
 )
@@ -19,17 +20,23 @@ func (s *protoSegments) get(p peer.ID) *protoSegment {
 	return s[byte(p[len(p)-1])]
 }
 
+var errTooManyProtocols = errors.New("too many protocols")
+
+type Option func(*memoryProtoBook) error
+
 type memoryProtoBook struct {
 	segments protoSegments
 
+	maxProtos int
+
 	lk       sync.RWMutex
 	interned map[string]string
 }
 
 var _ pstore.ProtoBook = (*memoryProtoBook)(nil)
 
-func NewProtoBook() *memoryProtoBook {
-	return &memoryProtoBook{
+func NewProtoBook(opts ...Option) (*memoryProtoBook, error) {
+	pb := &memoryProtoBook{
 		interned: make(map[string]string, 256),
 		segments: func() (ret protoSegments) {
 			for i := range ret {
@@ -39,7 +46,15 @@ func NewProtoBook() *memoryProtoBook {
 			}
 			return ret
 		}(),
+		maxProtos: 1024,
+	}
+
+	for _, opt := range opts {
+		if err := opt(pb); err != nil {
+			return nil, err
+		}
 	}
+	return pb, nil
 }
 
 func (pb *memoryProtoBook) internProtocol(proto string) string {
@@ -70,17 +85,19 @@ func (pb *memoryProtoBook) SetProtocols(p peer.ID, protos ...string) error {
 	if err := p.Validate(); err != nil {
 		return err
 	}
-
-	s := pb.segments.get(p)
-	s.Lock()
-	defer s.Unlock()
+	if len(protos) > pb.maxProtos {
+		return errTooManyProtocols
+	}
 
 	newprotos := make(map[string]struct{}, len(protos))
 	for _, proto := range protos {
 		newprotos[pb.internProtocol(proto)] = struct{}{}
 	}
 
+	s := pb.segments.get(p)
+	s.Lock()
 	s.protocols[p] = newprotos
+	s.Unlock()
 
 	return nil
 }
@@ -99,11 +116,13 @@ func (pb *memoryProtoBook) AddProtocols(p peer.ID, protos ...string) error {
 		protomap = make(map[string]struct{})
 		s.protocols[p] = protomap
 	}
+	if len(protomap)+len(protos) > pb.maxProtos {
+		return errTooManyProtocols
+	}
 
 	for _, proto := range protos {
 		protomap[pb.internProtocol(proto)] = struct{}{}
 	}
-
 	return nil
 }