Merge pull request #5 from ipfs/feat/update-proto

rework protocol methods

p2p/host/peerstore/peerinfo_test.go

@@ -0,0 +1,57 @@
+package peerstore
+
+import (
+	peer "github.com/ipfs/go-libp2p-peer"
+	ma "github.com/jbenet/go-multiaddr"
+	"testing"
+)
+
+func mustAddr(t *testing.T, s string) ma.Multiaddr {
+	addr, err := ma.NewMultiaddr(s)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	return addr
+}
+
+func TestPeerInfoMarshal(t *testing.T) {
+	a := mustAddr(t, "/ip4/1.2.3.4/tcp/4536")
+	b := mustAddr(t, "/ip4/1.2.3.8/udp/7777")
+	id, err := peer.IDB58Decode("QmaCpDMGvV2BGHeYERUEnRQAwe3N8SzbUtfsmvsqQLuvuJ")
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	pi := &PeerInfo{
+		ID:    id,
+		Addrs: []ma.Multiaddr{a, b},
+	}
+
+	data, err := pi.MarshalJSON()
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	pi2 := new(PeerInfo)
+	if err := pi2.UnmarshalJSON(data); err != nil {
+		t.Fatal(err)
+	}
+
+	if pi2.ID != pi.ID {
+		t.Fatal("ids didnt match after marshal")
+	}
+
+	if !pi.Addrs[0].Equal(pi2.Addrs[0]) {
+		t.Fatal("wrong addrs")
+	}
+
+	if !pi.Addrs[1].Equal(pi2.Addrs[1]) {
+		t.Fatal("wrong addrs")
+	}
+
+	lgbl := pi2.Loggable()
+	if lgbl["peerID"] != id.Pretty() {
+		t.Fatal("loggables gave wrong peerID output")
+	}
+}

p2p/host/peerstore/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
@@ -66,7 +68,7 @@ type AddrBook interface {
 	// This is used when we receive the best estimate of the validity of an address.
 	SetAddrs(p peer.ID, addrs []ma.Multiaddr, ttl time.Duration)
 
-	// Addresses returns all known (and valid) addresses for a given
+	// Addresses returns all known (and valid) addresses for a given peer
 	Addrs(p peer.ID) []ma.Multiaddr
 
 	// AddrStream returns a channel that gets all addresses for a given
@@ -170,6 +172,9 @@ type peerstore struct {
 	// TODO: use a datastore for this
 	ds     map[string]interface{}
 	dslock sync.Mutex
+
+	// lock for protocol information, separate from datastore lock
+	protolock sync.Mutex
 }
 
 // NewPeerstore creates a threadsafe collection of peers.
@@ -192,6 +197,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 +207,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 +235,67 @@ 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 {
+	ps.protolock.Lock()
+	defer ps.protolock.Unlock()
+	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")
+	ps.protolock.Lock()
+	defer ps.protolock.Unlock()
+	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) {
+	ps.protolock.Lock()
+	defer ps.protolock.Unlock()
+	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

p2p/host/peerstore/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,45 @@ 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)
+	}
+}
+
+func TestBasicPeerstore(t *testing.T) {
+	ps := NewPeerstore()
+
+	var pids []peer.ID
+	addrs := getAddrs(t, 10)
+	for i, a := range addrs {
+		p := peer.ID(fmt.Sprint(i))
+		pids = append(pids, p)
+		ps.AddAddr(p, a, PermanentAddrTTL)
+	}
+
+	peers := ps.Peers()
+	if len(peers) != 10 {
+		t.Fatal("expected ten peers")
+	}
+
+	pinfo := ps.PeerInfo(pids[0])
+	if !pinfo.Addrs[0].Equal(addrs[0]) {
+		t.Fatal("stored wrong address")
+	}
 }