From 426d729c513205073e4fdc71c6d2dd5f87fa65b2 Mon Sep 17 00:00:00 2001
From: Will Scott <will@cypherpunk.email>
Date: Tue, 25 Feb 2020 10:24:47 -0800
Subject: [PATCH] cleanup addrToIP logic

---
 p2p/host/autonat/svc.go      | 28 +++++++++++-----------------
 p2p/host/autonat/svc_test.go | 10 +++++-----
 2 files changed, 16 insertions(+), 22 deletions(-)

diff --git a/p2p/host/autonat/svc.go b/p2p/host/autonat/svc.go
index 141f4ee7..1073e190 100644
--- a/p2p/host/autonat/svc.go
+++ b/p2p/host/autonat/svc.go
@@ -3,7 +3,7 @@ package autonat
 import (
 	"bytes"
 	"context"
-	"fmt"
+	"errors"
 	"math/rand"
 	"net"
 	"strings"
@@ -106,31 +106,25 @@ func (as *AutoNATService) handleStream(s network.Stream) {
 }
 
 // Optimistically extract the net.IP host from a multiaddress.
-func addrToIP(addr ma.Multiaddr) net.IP {
+func addrToIP(addr ma.Multiaddr) (net.IP, error) {
 	n, ip, err := manet.DialArgs(addr)
 	if err != nil {
-		return nil
+		return nil, err
 	}
 
-	// if no port:
-	if n == "ip" || n == "ip4" || n == "ip6" {
-		// Strip v6 zone if it's there.
-		if strings.Contains(ip, "%") {
-			ip = ip[:strings.Index(ip, "%")]
-		}
-		return net.ParseIP(ip)
+	if strings.HasPrefix(n, "tcp") || strings.HasPrefix(n, "udp") {
+		ip, _, err = net.SplitHostPort(ip)
+	} else if !strings.HasPrefix(n, "ip") {
+		return nil, errors.New("non-ip multiaddr")
 	}
-
-	ip, _, err = net.SplitHostPort(ip)
 	if err != nil {
-		fmt.Printf("failed to split: %v", err)
-		return nil
+		return nil, err
 	}
 	// Strip v6 zone if it's there.
 	if strings.Contains(ip, "%") {
 		ip = ip[:strings.Index(ip, "%")]
 	}
-	return net.ParseIP(ip)
+	return net.ParseIP(ip), nil
 }
 
 func (as *AutoNATService) handleDial(p peer.ID, obsaddr ma.Multiaddr, mpi *pb.Message_PeerInfo) *pb.Message_DialResponse {
@@ -158,7 +152,7 @@ func (as *AutoNATService) handleDial(p peer.ID, obsaddr ma.Multiaddr, mpi *pb.Me
 	if !as.skipDial(obsaddr) {
 		addrs = append(addrs, obsaddr)
 		seen[obsaddr.String()] = struct{}{}
-		obsHost = addrToIP(obsaddr)
+		obsHost, _ = addrToIP(obsaddr)
 	}
 
 	for _, maddr := range mpi.GetAddrs() {
@@ -172,7 +166,7 @@ func (as *AutoNATService) handleDial(p peer.ID, obsaddr ma.Multiaddr, mpi *pb.Me
 			continue
 		}
 
-		if !bytes.Equal(obsHost, addrToIP(addr)) {
+		if ip, err := addrToIP(addr); err != nil || !bytes.Equal(obsHost, ip) {
 			continue
 		}
 
diff --git a/p2p/host/autonat/svc_test.go b/p2p/host/autonat/svc_test.go
index f2aa91a6..9be52fbf 100644
--- a/p2p/host/autonat/svc_test.go
+++ b/p2p/host/autonat/svc_test.go
@@ -139,26 +139,26 @@ func TestAutoNATServiceDialRateLimiter(t *testing.T) {
 
 func TestAddrToIP(t *testing.T) {
 	addr, _ := ma.NewMultiaddr("/ip4/127.0.0.1/tcp/0")
-	if !addrToIP(addr).Equal(net.IPv4(127, 0, 0, 1)) {
+	if ip, err := addrToIP(addr); err != nil || !ip.Equal(net.IPv4(127, 0, 0, 1)) {
 		t.Fatal("addrToIP of ipv4 localhost incorrect!")
 	}
 
 	addr, _ = ma.NewMultiaddr("/ip4/192.168.0.1/tcp/6")
-	if !addrToIP(addr).Equal(net.IPv4(192, 168, 0, 1)) {
+	if ip, err := addrToIP(addr); err != nil || !ip.Equal(net.IPv4(192, 168, 0, 1)) {
 		t.Fatal("addrToIP of ipv4 incorrect!")
 	}
 
 	addr, _ = ma.NewMultiaddr("/ip6/::ffff:127.0.0.1/tcp/111")
-	if !addrToIP(addr).Equal(net.ParseIP("::ffff:127.0.0.1")) {
+	if ip, err := addrToIP(addr); err != nil || !ip.Equal(net.ParseIP("::ffff:127.0.0.1")) {
 		t.Fatal("addrToIP of ipv6 incorrect!")
 	}
 	addr, _ = ma.NewMultiaddr("/ip6zone/eth0/ip6/fe80::1")
-	if !addrToIP(addr).Equal(net.ParseIP("fe80::1")) {
+	if ip, err := addrToIP(addr); err != nil || !ip.Equal(net.ParseIP("fe80::1")) {
 		t.Fatal("addrToIP of ip6zone incorrect!")
 	}
 
 	addr, _ = ma.NewMultiaddr("/unix/a/b/c/d")
-	if addrToIP(addr) != nil {
+	if _, err := addrToIP(addr); err == nil {
 		t.Fatal("invalid addrToIP populates")
 	}
 }