From 6e64bcad480afc6426f6f3b320478e98808b2c4d Mon Sep 17 00:00:00 2001
From: Nicolas De Loof <nicolas.deloof@gmail.com>
Date: Wed, 3 Mar 2021 16:03:48 +0100
Subject: [PATCH] support both square brackets and raw syntax for IPV6

Signed-off-by: Nicolas De Loof <nicolas.deloof@gmail.com>
---
 nat/nat.go      |  13 +++---
 nat/nat_test.go | 102 +++++++++++++++++++++++++++++++++++-------------
 2 files changed, 82 insertions(+), 33 deletions(-)

diff --git a/nat/nat.go b/nat/nat.go
index bb7e4e336..0f17a5ec5 100644
--- a/nat/nat.go
+++ b/nat/nat.go
@@ -175,13 +175,16 @@ func splitParts(rawport string) (string, string, string) {
 // ParsePortSpec parses a port specification string into a slice of PortMappings
 func ParsePortSpec(rawPort string) ([]PortMapping, error) {
 	var proto string
-	rawIP, hostPort, containerPort := splitParts(rawPort)
+	ip, hostPort, containerPort := splitParts(rawPort)
 	proto, containerPort = SplitProtoPort(containerPort)
 
-	// Strip [] from IPV6 addresses
-	ip, _, err := net.SplitHostPort(rawIP + ":")
-	if err != nil {
-		return nil, fmt.Errorf("Invalid ip address %v: %s", rawIP, err)
+	if ip != "" && ip[0] == '[' {
+		// Strip [] from IPV6 addresses
+		rawIP, _, err := net.SplitHostPort(ip + ":")
+		if err != nil {
+			return nil, fmt.Errorf("Invalid ip address %v: %s", ip, err)
+		}
+		ip = rawIP
 	}
 	if ip != "" && net.ParseIP(ip) == nil {
 		return nil, fmt.Errorf("Invalid ip address: %s", ip)
diff --git a/nat/nat_test.go b/nat/nat_test.go
index 12e61ea12..348877e7a 100644
--- a/nat/nat_test.go
+++ b/nat/nat_test.go
@@ -201,42 +201,88 @@ func TestParsePortSpecFull(t *testing.T) {
 }
 
 func TestPartPortSpecIPV6(t *testing.T) {
-	portMappings, err := ParsePortSpec("[2001:4860:0:2001::68]::333")
-	if err != nil {
-		t.Fatalf("expected nil error, got: %v", err)
+	type test struct {
+		name     string
+		spec     string
+		expected []PortMapping
 	}
-
-	expected := []PortMapping{
+	cases := []test{
 		{
-			Port: "333/tcp",
-			Binding: PortBinding{
-				HostIP:   "2001:4860:0:2001::68",
-				HostPort: "",
+			name: "square angled IPV6 without host port",
+			spec: "[2001:4860:0:2001::68]::333",
+			expected: []PortMapping{
+				{
+					Port: "333/tcp",
+					Binding: PortBinding{
+						HostIP:   "2001:4860:0:2001::68",
+						HostPort: "",
+					},
+				},
 			},
 		},
-	}
-	if !reflect.DeepEqual(expected, portMappings) {
-		t.Fatalf("wrong port mappings: got=%v, want=%v", portMappings, expected)
-	}
-}
-
-func TestPartPortSpecIPV6WithHostPort(t *testing.T) {
-	portMappings, err := ParsePortSpec("[::1]:80:80")
-	if err != nil {
-		t.Fatalf("expected nil error, got: %v", err)
-	}
-
-	expected := []PortMapping{
 		{
-			Port: "80/tcp",
-			Binding: PortBinding{
-				HostIP:   "::1",
-				HostPort: "80",
+			name: "square angled IPV6 with host port",
+			spec: "[::1]:80:80",
+			expected: []PortMapping{
+				{
+					Port: "80/tcp",
+					Binding: PortBinding{
+						HostIP:   "::1",
+						HostPort: "80",
+					},
+				},
+			},
+		},
+		{
+			name: "IPV6 without host port",
+			spec: "2001:4860:0:2001::68::333",
+			expected: []PortMapping{
+				{
+					Port: "333/tcp",
+					Binding: PortBinding{
+						HostIP:   "2001:4860:0:2001::68",
+						HostPort: "",
+					},
+				},
+			},
+		},
+		{
+			name: "IPV6 with host port",
+			spec: "::1:80:80",
+			expected: []PortMapping{
+				{
+					Port: "80/tcp",
+					Binding: PortBinding{
+						HostIP:   "::1",
+						HostPort: "80",
+					},
+				},
+			},
+		},
+		{
+			name: ":: IPV6, without host port",
+			spec: "::::80",
+			expected: []PortMapping{
+				{
+					Port: "80/tcp",
+					Binding: PortBinding{
+						HostIP:   "::",
+						HostPort: "",
+					},
+				},
 			},
 		},
 	}
-	if !reflect.DeepEqual(expected, portMappings) {
-		t.Fatalf("wrong port mappings: got=%v, want=%v", portMappings, expected)
+	for _, c := range cases {
+		t.Run(c.name, func(t *testing.T) {
+			portMappings, err := ParsePortSpec(c.spec)
+			if err != nil {
+				t.Fatalf("expected nil error, got: %v", err)
+			}
+			if !reflect.DeepEqual(c.expected, portMappings) {
+				t.Fatalf("wrong port mappings: got=%v, want=%v", portMappings, c.expected)
+			}
+		})
 	}
 }