Merge pull request #147 from libp2p/feat/supports-no-alloc

Implement the FirstSupportedProtocol API.
libp2p · May 14, 2020 · 1d296f2 · 1d296f2
2 parents 104663b + 371ef2d
commit 1d296f2
Show file tree

Hide file tree

Showing 3 changed files with 52 additions and 0 deletions.
diff --git a/p2p/host/peerstore/pstoreds/protobook.go b/p2p/host/peerstore/pstoreds/protobook.go
@@ -122,6 +122,28 @@ func (pb *dsProtoBook) SupportsProtocols(p peer.ID, protos ...string) ([]string,
 	return res, nil
 }
 
+func (pb *dsProtoBook) FirstSupportedProtocol(p peer.ID, protos ...string) (string, error) {
+	if err := p.Validate(); err != nil {
+		return "", err
+	}
+
+	s := pb.segments.get(p)
+	s.RLock()
+	defer s.RUnlock()
+
+	pmap, err := pb.getProtocolMap(p)
+	if err != nil {
+		return "", err
+	}
+	for _, proto := range protos {
+		if _, ok := pmap[proto]; ok {
+			return proto, nil
+		}
+	}
+
+	return "", nil
+}
+
 func (pb *dsProtoBook) RemoveProtocols(p peer.ID, protos ...string) error {
 	if err := p.Validate(); err != nil {
 		return err

diff --git a/p2p/host/peerstore/pstoremem/protobook.go b/p2p/host/peerstore/pstoremem/protobook.go
@@ -163,3 +163,20 @@ func (pb *memoryProtoBook) SupportsProtocols(p peer.ID, protos ...string) ([]str
 
 	return out, nil
 }
+
+func (pb *memoryProtoBook) FirstSupportedProtocol(p peer.ID, protos ...string) (string, error) {
+	if err := p.Validate(); err != nil {
+		return "", err
+	}
+
+	s := pb.segments.get(p)
+	s.RLock()
+	defer s.RUnlock()
+
+	for _, proto := range protos {
+		if _, ok := s.protocols[p][proto]; ok {
+			return proto, nil
+		}
+	}
+	return "", nil
+}
diff --git a/p2p/host/peerstore/test/peerstore_suite.go b/p2p/host/peerstore/test/peerstore_suite.go
@@ -14,6 +14,7 @@ import (
 	ma "github.com/multiformats/go-multiaddr"
 
 	pstore "github.com/libp2p/go-libp2p-core/peerstore"
+	"github.com/stretchr/testify/require"
 )
 
 var peerstoreSuite = map[string]func(pstore.Peerstore) func(*testing.T){
@@ -242,6 +243,18 @@ func testPeerstoreProtoStore(ps pstore.Peerstore) func(t *testing.T) {
 			t.Fatal("got wrong supported array: ", supported)
 		}
 
+		b, err := ps.FirstSupportedProtocol(p1, "q", "w", "a", "y", "b")
+		require.NoError(t, err)
+		require.Equal(t, "a", b)
+
+		b, err = ps.FirstSupportedProtocol(p1, "q", "x", "z")
+		require.NoError(t, err)
+		require.Empty(t, b)
+
+		b, err = ps.FirstSupportedProtocol(p1, "a")
+		require.NoError(t, err)
+		require.Equal(t, "a", b)
+
 		protos = []string{"other", "yet another", "one more"}
 		err = ps.SetProtocols(p1, protos...)
 		if err != nil {