add methods for setting protocols of a peer

peerstore.go

@@ -42,6 +42,9 @@ type Peerstore interface {
 	// methods. this is a last resort.
 	Get(id peer.ID, key string) (interface{}, error)
 	Put(id peer.ID, key string, val interface{}) error
+
+	GetProtocols(peer.ID) ([]string, error)
+	SetProtocols(peer.ID, []string) error
 }
 
 // AddrBook is an interface that fits the new AddrManager. I'm patching
@@ -225,6 +228,24 @@ 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) GetProtocols(p peer.ID) ([]string, error) {
+	protos, err := ps.Get(p, "protocols")
+	if err != nil {
+		return nil, err
+	}
+
+	out, ok := protos.([]string)
+	if !ok {
+		return nil, errors.New("stored protocols array was not array of strings")
+	}
+
+	return out, nil
+}
+
 func PeerInfos(ps Peerstore, peers []peer.ID) []PeerInfo {
 	pi := make([]PeerInfo, len(peers))
 	for i, p := range peers {

peerstore_test.go

@@ -179,3 +179,30 @@ func TestAddrStreamDuplicates(t *testing.T) {
 		t.Fatal("should have received exactly ten addresses")
 	}
 }
+
+func TestPeerstoreProtoStore(t *testing.T) {
+	ps := NewPeerstore()
+	p1 := peer.ID("TESTPEER")
+
+	protos := []string{"a", "b", "c", "d"}
+
+	err := ps.SetProtocols(p1, protos)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	out, err := ps.GetProtocols(p1)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	if len(out) != len(protos) {
+		t.Fatal("got wrong number of protocols back")
+	}
+
+	for i, p := range protos {
+		if out[i] != p {
+			t.Fatal("got wrong protocol")
+		}
+	}
+}