Peer address parsing refactoring

pkg/proto/proto.go

@@ -6,6 +6,7 @@
 	"encoding/binary"
 	"fmt"
 	"io"
+	"math/rand/v2"
 	"net"
 	"strconv"
 	"strings"
@@ -375,23 +376,14 @@
 	return a.IP.Equal(other.IP) && a.Port == other.Port
 }
 
+// NewTCPAddrFromString creates TCPAddr from string.
+// Returns empty TCPAddr if string can't be parsed.
 func NewTCPAddrFromString(s string) TCPAddr {
-	host, port, err := net.SplitHostPort(s)
+	pi, err := NewPeerInfoFromString(s)
 	if err != nil {
-		return TCPAddr{}
+		return TCPAddr{} // return empty TCPAddr in case of error
 	}
-	ip := net.ParseIP(host)
-	if ip == nil {
-		ips, err := net.LookupIP(host)
-		if err == nil {
-			ip = ips[0]
-		}
-	}
-	p, err := strconv.ParseUint(port, 10, 16)
-	if err != nil {
-		return TCPAddr{}
-	}
-	return NewTCPAddr(ip, int(p))
+	return NewTCPAddr(pi.Addr, int(pi.Port))
 }
 
 func NewTcpAddrFromUint64(value uint64) TCPAddr {
@@ -737,26 +729,71 @@
 	return NewTCPAddr(a.Addr(), a.Port()).String()
 }
 
+func ipV4Loopback() net.IP {
+	const a, b, c, d = 127, 0, 0, 1
+	return net.IPv4(a, b, c, d)
+}
+
+func filterToIPV4(ips []net.IP) []net.IP {
+	for i := 0; i < len(ips); i++ {
+		ip := ips[i]
+		if ip.IsLoopback() {
+			// after resolving localhost, ip can be in both format, unify it to IPv4
+			ips[i] = ipV4Loopback()
+			continue
+		}
+		if ip.To4() == nil { // for now we support only IPv4
+			ips = append(ips[:i], ips[i+1:]...) // remove non-IPv4 address, assume that count of addresses is small
+		}
+	}
+	return ips
+}
+
+func resolveHostToIPv4(host string) (net.IP, error) {
+	if ip := net.ParseIP(host); ip != nil { // try to parse host as IP address
+		if ip.To4() == nil {
+			return nil, errors.Errorf("non-IPv4 address %q", host)
+		}
+		return ip, nil // host is already an IP address
+	}
+	ips, err := net.LookupIP(host) // try to resolve host
+	if err != nil {
+		return nil, errors.Wrapf(err, "failed to resolve host %q", host)
+	}
+	ips = filterToIPV4(ips)
+	if len(ips) == 0 {
+		return nil, errors.Errorf("no IPv4 addresses found for host %q", host)
+	}
+	// Select random IPv4 from the list
+	n := rand.IntN(len(ips)) // #nosec: it's ok to use math/rand/v2 here
+	return ips[n], nil
+}
+
 // PeerInfo represents the address of a single peer
 type PeerInfo struct {
 	Addr net.IP
 	Port uint16
 }
 
 func NewPeerInfoFromString(addr string) (PeerInfo, error) {
-	parts := strings.Split(addr, ":")
-	if len(parts) != 2 {
-		return PeerInfo{}, errors.Errorf("invalid addr %s", addr)
+	host, port, err := net.SplitHostPort(addr)
+	if err != nil {
+		return PeerInfo{}, errors.Wrap(err, "failed to split host and port")
 	}
-
-	ip := net.ParseIP(parts[0])
-	port, err := strconv.ParseUint(parts[1], 10, 16)
+	ip, err := resolveHostToIPv4(host)
 	if err != nil {
-		return PeerInfo{}, errors.Errorf("invalid port %s", parts[1])
+		return PeerInfo{}, errors.Wrap(err, "failed to resolve host")
+	}
+	portNum, err := strconv.ParseUint(port, 10, 16)
+	if err != nil {
+		return PeerInfo{}, errors.Errorf("invalid port %q", port)
+	}
+	if portNum == 0 {
+		return PeerInfo{}, errors.Errorf("invalid port %q", port)
 	}
 	return PeerInfo{
 		Addr: ip,
-		Port: uint16(port),
+		Port: uint16(portNum),
 	}, nil
 }
 

pkg/proto/proto_test.go

@@ -7,13 +7,15 @@ import (
 	"encoding/json"
 	"fmt"
 	"io"
+	"math"
 	"net"
 	"sort"
 	"strings"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
+
 	"github.com/wavesplatform/gowaves/pkg/crypto"
 )
 
@@ -277,10 +279,67 @@ func TestPeerInfoMarshalJSON(t *testing.T) {
 }
 
 func TestNewPeerInfoFromString(t *testing.T) {
-	rs, err := NewPeerInfoFromString("34.253.153.4:6868")
-	require.NoError(t, err)
-	assert.Equal(t, "34.253.153.4", rs.Addr.String())
-	assert.EqualValues(t, 6868, rs.Port)
+	tests := []struct {
+		in  string
+		out PeerInfo
+		err string
+	}{
+		{"34.253.153.4:6868", PeerInfo{net.IPv4(34, 253, 153, 4), 6868}, ""},
+		{
+			"34.444.153.4:6868",
+			PeerInfo{},
+			"failed to resolve host: failed to resolve host \"34.444.153.4\": lookup 34.444.153.4: no such host",
+		},
+		{
+			"jfhasjdhfkmnn:6868",
+			PeerInfo{},
+			"failed to resolve host: failed to resolve host \"jfhasjdhfkmnn\": lookup jfhasjdhfkmnn: no such host",
+		},
+		{
+			"localhost:6868",
+			PeerInfo{net.IPv4(127, 0, 0, 1), 6868},
+			"",
+		},
+		{
+			"127.0.0.1:6868",
+			PeerInfo{net.IPv4(127, 0, 0, 1), 6868},
+			"",
+		},
+		{
+			fmt.Sprintf("34.44.153.4:%d", math.MaxUint16+1),
+			PeerInfo{},
+			fmt.Sprintf("invalid port \"%d\"", math.MaxUint16+1),
+		},
+		{
+			fmt.Sprintf("34.44.153.4:%d", -42),
+			PeerInfo{},
+			fmt.Sprintf("invalid port \"%d\"", -42),
+		},
+		{"34.44.153.4:bugaga", PeerInfo{}, "invalid port \"bugaga\""},
+		{"34.44.153.4:0", PeerInfo{}, "invalid port \"0\""},
+		{"34.44.153.4:", PeerInfo{}, "invalid port \"\""},
+		{
+			"34.44.153.4",
+			PeerInfo{},
+			"failed to split host and port: address 34.44.153.4: missing port in address",
+		},
+		{
+			"34.44.153.4:42:",
+			PeerInfo{},
+			"failed to split host and port: address 34.44.153.4:42:: too many colons in address",
+		},
+	}
+	for i, tc := range tests {
+		t.Run(fmt.Sprintf("case %d", i), func(t *testing.T) {
+			rs, err := NewPeerInfoFromString(tc.in)
+			if tc.err != "" {
+				assert.EqualError(t, err, tc.err)
+			} else {
+				require.NoError(t, err)
+				assert.Equal(t, tc.out, rs)
+			}
+		})
+	}
 }
 
 func TestPeerInfoUnmarshalJSON(t *testing.T) {