adjust peerRegistry overlay address handling

pkg/p2p/libp2p/peer.go

@@ -12,28 +12,28 @@ import (
 )
 
 type peerRegistry struct {
-	peers    map[string]libp2ppeer.ID
-	overlays map[libp2ppeer.ID]string
+	peers    map[string]libp2ppeer.ID        // map overlay address to underlay peer id
+	overlays map[libp2ppeer.ID]swarm.Address // map underlay peer id to overlay address
 	mu       sync.RWMutex
 }
 
 func newPeerRegistry() *peerRegistry {
 	return &peerRegistry{
 		peers:    make(map[string]libp2ppeer.ID),
-		overlays: make(map[libp2ppeer.ID]string),
+		overlays: make(map[libp2ppeer.ID]swarm.Address),
 	}
 }
 
 func (r *peerRegistry) add(peerID libp2ppeer.ID, overlay swarm.Address) {
 	r.mu.Lock()
-	r.peers[string(overlay.Bytes())] = peerID
-	r.overlays[peerID] = string(overlay.Bytes())
+	r.peers[encodePeersKey(overlay)] = peerID
+	r.overlays[peerID] = overlay
 	r.mu.Unlock()
 }
 
 func (r *peerRegistry) peerID(overlay swarm.Address) (peerID libp2ppeer.ID, found bool) {
 	r.mu.RLock()
-	peerID, found = r.peers[string(overlay.Bytes())]
+	peerID, found = r.peers[encodePeersKey(overlay)]
 	r.mu.RUnlock()
 	return peerID, found
 }
@@ -42,13 +42,17 @@ func (r *peerRegistry) overlay(peerID libp2ppeer.ID) (swarm.Address, bool) {
 	r.mu.RLock()
 	overlay, found := r.overlays[peerID]
 	r.mu.RUnlock()
-	return swarm.NewAddress([]byte(overlay)), found
+	return overlay, found
 }
 
 func (r *peerRegistry) remove(peerID libp2ppeer.ID) {
 	r.mu.Lock()
 	overlay := r.overlays[peerID]
 	delete(r.overlays, peerID)
-	delete(r.peers, overlay)
+	delete(r.peers, encodePeersKey(overlay))
 	r.mu.Unlock()
 }
+
+func encodePeersKey(overlay swarm.Address) string {
+	return string(overlay.Bytes())
+}