From 3015b65a102aeac59196f772eaac12df9d24ba96 Mon Sep 17 00:00:00 2001
From: Jeromy <why@ipfs.io>
Date: Fri, 19 Aug 2016 14:16:56 -0700
Subject: [PATCH] rework protocol storage and handling to better support upper
 layers needs

---
 peerstore.go      | 62 +++++++++++++++++++++++++++++++++++++++++------
 peerstore_test.go | 17 ++++++++++++-
 2 files changed, 70 insertions(+), 9 deletions(-)

diff --git a/peerstore.go b/peerstore.go
index 3a88d58..166be36 100644
--- a/peerstore.go
+++ b/peerstore.go
@@ -2,6 +2,7 @@ package peerstore
 
 import (
 	"errors"
+	"fmt"
 	"sync"
 	"time"
 
@@ -44,7 +45,8 @@ type Peerstore interface {
 	Put(id peer.ID, key string, val interface{}) error
 
 	GetProtocols(peer.ID) ([]string, error)
-	SetProtocols(peer.ID, []string) error
+	AddProtocols(peer.ID, ...string) error
+	SupportsProtocols(peer.ID, ...string) ([]string, error)
 }
 
 // AddrBook is an interface that fits the new AddrManager. I'm patching
@@ -192,6 +194,8 @@ func (ps *peerstore) Put(p peer.ID, key string, val interface{}) error {
 	return nil
 }
 
+var ErrNotFound = errors.New("item not found")
+
 func (ps *peerstore) Get(p peer.ID, key string) (interface{}, error) {
 	//dsk := ds.NewKey(string(p) + "/" + key)
 	//return ps.ds.Get(dsk)
@@ -200,7 +204,7 @@ func (ps *peerstore) Get(p peer.ID, key string) (interface{}, error) {
 	defer ps.dslock.Unlock()
 	i, ok := ps.ds[string(p)+"/"+key]
 	if !ok {
-		return nil, errors.New("item not found")
+		return nil, ErrNotFound
 	}
 	return i, nil
 }
@@ -228,19 +232,61 @@ func (ps *peerstore) PeerInfo(p peer.ID) PeerInfo {
 	}
 }
 
-func (ps *peerstore) SetProtocols(p peer.ID, protos []string) error {
-	return ps.Put(p, "protocols", protos)
+func (ps *peerstore) AddProtocols(p peer.ID, protos ...string) error {
+	protomap, err := ps.getProtocolMap(p)
+	if err != nil {
+		return err
+	}
+
+	for _, proto := range protos {
+		protomap[proto] = struct{}{}
+	}
+
+	return ps.Put(p, "protocols", protomap)
+}
+
+func (ps *peerstore) getProtocolMap(p peer.ID) (map[string]struct{}, error) {
+	iprotomap, err := ps.Get(p, "protocols")
+	switch err {
+	default:
+		return nil, err
+	case ErrNotFound:
+		return make(map[string]struct{}), nil
+	case nil:
+		cast, ok := iprotomap.(map[string]struct{})
+		if !ok {
+			return nil, fmt.Errorf("stored protocol set was not a map")
+		}
+
+		return cast, nil
+	}
 }
 
 func (ps *peerstore) GetProtocols(p peer.ID) ([]string, error) {
-	protos, err := ps.Get(p, "protocols")
+	pmap, err := ps.getProtocolMap(p)
 	if err != nil {
 		return nil, err
 	}
 
-	out, ok := protos.([]string)
-	if !ok {
-		return nil, errors.New("stored protocols array was not array of strings")
+	var out []string
+	for k, _ := range pmap {
+		out = append(out, k)
+	}
+
+	return out, nil
+}
+
+func (ps *peerstore) SupportsProtocols(p peer.ID, protos ...string) ([]string, error) {
+	pmap, err := ps.getProtocolMap(p)
+	if err != nil {
+		return nil, err
+	}
+
+	var out []string
+	for _, proto := range protos {
+		if _, ok := pmap[proto]; ok {
+			out = append(out, proto)
+		}
 	}
 
 	return out, nil
diff --git a/peerstore_test.go b/peerstore_test.go
index 2a0c28d..7aedb90 100644
--- a/peerstore_test.go
+++ b/peerstore_test.go
@@ -3,6 +3,7 @@ package peerstore
 import (
 	"fmt"
 	"math/rand"
+	"sort"
 	"testing"
 	"time"
 
@@ -186,7 +187,7 @@ func TestPeerstoreProtoStore(t *testing.T) {
 
 	protos := []string{"a", "b", "c", "d"}
 
-	err := ps.SetProtocols(p1, protos)
+	err := ps.AddProtocols(p1, protos...)
 	if err != nil {
 		t.Fatal(err)
 	}
@@ -200,9 +201,23 @@ func TestPeerstoreProtoStore(t *testing.T) {
 		t.Fatal("got wrong number of protocols back")
 	}
 
+	sort.Strings(out)
 	for i, p := range protos {
 		if out[i] != p {
 			t.Fatal("got wrong protocol")
 		}
 	}
+
+	supported, err := ps.SupportsProtocols(p1, "q", "w", "a", "y", "b")
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	if len(supported) != 2 {
+		t.Fatal("only expected 2 supported")
+	}
+
+	if supported[0] != "a" || supported[1] != "b" {
+		t.Fatal("got wrong supported array: ", supported)
+	}
 }