From c2a62a74af4e10ebe654bde3e4c80b284d151eac Mon Sep 17 00:00:00 2001
From: Steven Allen <steven@stebalien.com>
Date: Fri, 31 May 2019 10:55:48 -0700
Subject: [PATCH] feat(peer): implement AddrInfosFromP2pAddrs and SplitAddr

* SplitAddr is a simpler way to split an address into a multiaddr and an ID.
* AddrInfosFromP2pAddrs converts a set of multiaddrs into a set of AddrInfos.
---
 peer/addrinfo.go      |  46 ++++++++++++--
 peer/addrinfo_test.go | 135 ++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 175 insertions(+), 6 deletions(-)
 create mode 100644 peer/addrinfo_test.go

diff --git a/peer/addrinfo.go b/peer/addrinfo.go
index cc9633b..6423d84 100644
--- a/peer/addrinfo.go
+++ b/peer/addrinfo.go
@@ -21,18 +21,51 @@ func (pi AddrInfo) String() string {
 
 var ErrInvalidAddr = fmt.Errorf("invalid p2p multiaddr")
 
-func AddrInfoFromP2pAddr(m ma.Multiaddr) (*AddrInfo, error) {
+// AddrInfosFromP2pAddrs converts a set of Multiaddrs to a set of AddrInfos.
+func AddrInfosFromP2pAddrs(maddrs ...ma.Multiaddr) ([]AddrInfo, error) {
+	m := make(map[ID][]ma.Multiaddr)
+	for _, maddr := range maddrs {
+		transport, id := SplitAddr(maddr)
+		if id == "" {
+			return nil, ErrInvalidAddr
+		}
+		if transport == nil {
+			if _, ok := m[id]; !ok {
+				m[id] = nil
+			}
+		} else {
+			m[id] = append(m[id], transport)
+		}
+	}
+	ais := make([]AddrInfo, 0, len(m))
+	for id, maddrs := range m {
+		ais = append(ais, AddrInfo{ID: id, Addrs: maddrs})
+	}
+	return ais, nil
+}
+
+// SplitAddr splits a p2p Multiaddr into a transport multiaddr and a peer ID.
+//
+// * Returns a nil transport if the address only contains a /p2p part.
+// * Returns a empty peer ID if the address doesn't contain a /p2p part.
+func SplitAddr(m ma.Multiaddr) (transport ma.Multiaddr, id ID) {
 	if m == nil {
-		return nil, ErrInvalidAddr
+		return nil, ""
 	}
 
 	transport, p2ppart := ma.SplitLast(m)
 	if p2ppart == nil || p2ppart.Protocol().Code != ma.P_P2P {
-		return nil, ErrInvalidAddr
+		return m, ""
 	}
-	id, err := IDFromBytes(p2ppart.RawValue())
-	if err != nil {
-		return nil, err
+	id = ID(p2ppart.RawValue()) // already validated by the multiaddr library.
+	return transport, id
+}
+
+// AddrInfoFromP2pAddr converts a Multiaddr to an AddrInfo.
+func AddrInfoFromP2pAddr(m ma.Multiaddr) (*AddrInfo, error) {
+	transport, id := SplitAddr(m)
+	if id == "" {
+		return nil, ErrInvalidAddr
 	}
 	info := &AddrInfo{ID: id}
 	if transport != nil {
@@ -41,6 +74,7 @@ func AddrInfoFromP2pAddr(m ma.Multiaddr) (*AddrInfo, error) {
 	return info, nil
 }
 
+// AddrInfoToP2pAddr converts an AddrInfo to a list of Multiaddrs.
 func AddrInfoToP2pAddrs(pi *AddrInfo) ([]ma.Multiaddr, error) {
 	var addrs []ma.Multiaddr
 	p2ppart, err := ma.NewComponent("p2p", IDB58Encode(pi.ID))
diff --git a/peer/addrinfo_test.go b/peer/addrinfo_test.go
new file mode 100644
index 0000000..b902fdf
--- /dev/null
+++ b/peer/addrinfo_test.go
@@ -0,0 +1,135 @@
+package peer_test
+
+import (
+	"testing"
+
+	ma "github.com/multiformats/go-multiaddr"
+
+	. "github.com/libp2p/go-libp2p-core/peer"
+)
+
+var (
+	testID                         ID
+	maddrFull, maddrTpt, maddrPeer ma.Multiaddr
+)
+
+func init() {
+	var err error
+	testID, err = IDB58Decode("QmS3zcG7LhYZYSJMhyRZvTddvbNUqtt8BJpaSs6mi1K5Va")
+	if err != nil {
+		panic(err)
+	}
+	maddrPeer = ma.StringCast("/p2p/" + IDB58Encode(testID))
+	maddrTpt = ma.StringCast("/ip4/127.0.0.1/tcp/1234")
+	maddrFull = maddrTpt.Encapsulate(maddrPeer)
+}
+
+func TestSplitAddr(t *testing.T) {
+	tpt, id := SplitAddr(maddrFull)
+	if !tpt.Equal(maddrTpt) {
+		t.Fatal("expected transport")
+	}
+	if id != testID {
+		t.Fatalf("%s != %s", id, testID)
+	}
+
+	tpt, id = SplitAddr(maddrPeer)
+	if tpt != nil {
+		t.Fatal("expected no transport")
+	}
+	if id != testID {
+		t.Fatalf("%s != %s", id, testID)
+	}
+
+	tpt, id = SplitAddr(maddrTpt)
+	if !tpt.Equal(maddrTpt) {
+		t.Fatal("expected a transport")
+	}
+	if id != "" {
+		t.Fatal("expected no peer ID")
+	}
+}
+
+func TestAddrInfoFromP2pAddr(t *testing.T) {
+	ai, err := AddrInfoFromP2pAddr(maddrFull)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if len(ai.Addrs) != 1 || !ai.Addrs[0].Equal(maddrTpt) {
+		t.Fatal("expected transport")
+	}
+	if ai.ID != testID {
+		t.Fatalf("%s != %s", ai.ID, testID)
+	}
+
+	ai, err = AddrInfoFromP2pAddr(maddrPeer)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if len(ai.Addrs) != 0 {
+		t.Fatal("expected transport")
+	}
+	if ai.ID != testID {
+		t.Fatalf("%s != %s", ai.ID, testID)
+	}
+
+	_, err = AddrInfoFromP2pAddr(maddrTpt)
+	if err != ErrInvalidAddr {
+		t.Fatalf("wrong error: %s", err)
+	}
+}
+
+func TestAddrInfosFromP2pAddrs(t *testing.T) {
+	infos, err := AddrInfosFromP2pAddrs()
+	if err != nil {
+		t.Fatal(err)
+	}
+	if len(infos) != 0 {
+		t.Fatal("expected no addrs")
+	}
+	infos, err = AddrInfosFromP2pAddrs(nil)
+	if err == nil {
+		t.Fatal("expected nil multiaddr to fail")
+	}
+
+	addrs := []ma.Multiaddr{
+		ma.StringCast("/ip4/128.199.219.111/tcp/4001/ipfs/QmSoLV4Bbm51jM9C4gDYZQ9Cy3U6aXMJDAbzgu2fzaDs64"),
+		ma.StringCast("/ip4/104.236.76.40/tcp/4001/ipfs/QmSoLV4Bbm51jM9C4gDYZQ9Cy3U6aXMJDAbzgu2fzaDs64"),
+
+		ma.StringCast("/ipfs/QmSoLer265NRgSp2LA3dPaeykiS1J6DifTC88f5uVQKNAd"),
+		ma.StringCast("/ip4/178.62.158.247/tcp/4001/ipfs/QmSoLer265NRgSp2LA3dPaeykiS1J6DifTC88f5uVQKNAd"),
+
+		ma.StringCast("/ipfs/QmSoLPppuBtQSGwKDZT2M73ULpjvfd3aZ6ha4oFGL1KrGM"),
+	}
+	expected := map[string][]ma.Multiaddr{
+		"QmSoLV4Bbm51jM9C4gDYZQ9Cy3U6aXMJDAbzgu2fzaDs64": {
+			ma.StringCast("/ip4/128.199.219.111/tcp/4001"),
+			ma.StringCast("/ip4/104.236.76.40/tcp/4001"),
+		},
+		"QmSoLer265NRgSp2LA3dPaeykiS1J6DifTC88f5uVQKNAd": {
+			ma.StringCast("/ip4/178.62.158.247/tcp/4001"),
+		},
+		"QmSoLPppuBtQSGwKDZT2M73ULpjvfd3aZ6ha4oFGL1KrGM": nil,
+	}
+	infos, err = AddrInfosFromP2pAddrs(addrs...)
+	if err != nil {
+		t.Fatal(err)
+	}
+	for _, info := range infos {
+		exaddrs, ok := expected[info.ID.Pretty()]
+		if !ok {
+			t.Fatalf("didn't expect peer %s", info.ID)
+		}
+		if len(info.Addrs) != len(exaddrs) {
+			t.Fatalf("got %d addrs, expected %d", len(info.Addrs), len(exaddrs))
+		}
+		// AddrInfosFromP2pAddrs preserves order. I'd like to keep this
+		// guarantee for now.
+		for i, addr := range info.Addrs {
+			if !exaddrs[i].Equal(addr) {
+				t.Fatalf("expected %s, got %s", exaddrs[i], addr)
+			}
+		}
+		delete(expected, info.ID.Pretty())
+	}
+}