From 00793e624419f91e63168d02875342cafa27fe1f Mon Sep 17 00:00:00 2001
From: Cole Brown <bigswim@gmail.com>
Date: Wed, 13 Jun 2018 19:27:14 -0400
Subject: [PATCH] Move Peers from Peerstore to AddrBook

---
 addr_manager_ds.go | 30 ++++++++++++++++++++++++++++++
 peerstore.go       | 12 ++++++------
 2 files changed, 36 insertions(+), 6 deletions(-)

diff --git a/addr_manager_ds.go b/addr_manager_ds.go
index 1c94fe0..ed39709 100644
--- a/addr_manager_ds.go
+++ b/addr_manager_ds.go
@@ -39,6 +39,11 @@ func peerAddressKey(p *peer.ID, addr *ma.Multiaddr) (ds.Key, error) {
 	return ds.NewKey(p.Pretty()).ChildString(hash.B58String()), nil
 }
 
+func peerIDFromKey(key ds.Key) (peer.ID, error) {
+	idstring := key.Parent().Type()
+	return peer.IDB58Decode(idstring)
+}
+
 func (mgr *DatastoreAddrManager) AddAddr(p peer.ID, addr ma.Multiaddr, ttl time.Duration) {
 	mgr.AddAddrs(p, []ma.Multiaddr{addr}, ttl)
 }
@@ -110,6 +115,31 @@ func (mgr *DatastoreAddrManager) Addrs(p peer.ID) []ma.Multiaddr {
 	return addrs
 }
 
+func (mgr *DatastoreAddrManager) Peers() []peer.ID {
+	q := query.Query{}
+	results, err := mgr.ds.Query(q)
+	if err != nil {
+		log.Error(err)
+		return []peer.ID{}
+	}
+
+	idset := make(map[peer.ID]struct{})
+	for result := range results.Next() {
+		key := ds.RawKey(result.Key)
+		id, err := peerIDFromKey(key)
+		if err != nil {
+			continue
+		}
+		idset[id] = struct{}{}
+	}
+
+	ids := make([]peer.ID, 0, len(idset))
+	for id := range idset {
+		ids = append(ids, id)
+	}
+	return ids
+}
+
 func (mgr *DatastoreAddrManager) AddrStream(ctx context.Context, p peer.ID) <-chan ma.Multiaddr {
 	initial := mgr.Addrs(p)
 	return mgr.subsManager.AddrStream(ctx, p, initial)
diff --git a/peerstore.go b/peerstore.go
index abfdeae..c47d8aa 100644
--- a/peerstore.go
+++ b/peerstore.go
@@ -30,9 +30,6 @@ type Peerstore interface {
 	KeyBook
 	Metrics
 
-	// Peers returns a list of all peer.IDs in this Peerstore
-	Peers() []peer.ID
-
 	// PeerInfo returns a peer.PeerInfo struct for given peer.ID.
 	// This is a small slice of the information Peerstore has on
 	// that peer, useful to other services.
@@ -83,6 +80,9 @@ type AddrBook interface {
 
 	// ClearAddresses removes all previously stored addresses
 	ClearAddrs(p peer.ID)
+
+	// Peers returns all of the peer IDs stored in the AddrBook
+	Peers() []peer.ID
 }
 
 // KeyBook tracks the Public keys of Peers.
@@ -179,7 +179,7 @@ func (kb *keybook) AddPrivKey(p peer.ID, sk ic.PrivKey) error {
 type peerstore struct {
 	*keybook
 	*metrics
-	AddrManager
+	AddrBook
 
 	// store other data, like versions
 	//ds ds.ThreadSafeDatastore
@@ -231,7 +231,7 @@ func (ps *peerstore) Peers() []peer.ID {
 	for _, p := range ps.keybook.Peers() {
 		set[p] = struct{}{}
 	}
-	for _, p := range ps.AddrManager.Peers() {
+	for _, p := range ps.AddrBook.Peers() {
 		set[p] = struct{}{}
 	}
 
@@ -245,7 +245,7 @@ func (ps *peerstore) Peers() []peer.ID {
 func (ps *peerstore) PeerInfo(p peer.ID) PeerInfo {
 	return PeerInfo{
 		ID:    p,
-		Addrs: ps.AddrManager.Addrs(p),
+		Addrs: ps.AddrBook.Addrs(p),
 	}
 }