fix issue with v4-mapped IPv6 addresses

mastercactapus · Mar 12, 2023 · ab99c61 · ab99c61
1 parent ae65f2c
commit ab99c61
Show file tree

Hide file tree

Showing 3 changed files with 53 additions and 8 deletions.
diff --git a/headerv1.go b/headerv1.go
@@ -73,7 +73,11 @@ func parseV1(r *bufio.Reader) (*HeaderV1, error) {
 		if fam == "TCP4" {
 			return ip.To4() != nil
 		}
-		return ip.To4() == nil
+		if fam == "TCP6" {
+			return ip.To16() != nil
+		}
+
+		return false
 	}
 
 	srcIP := net.ParseIP(srcIPStr)

diff --git a/headerv1_test.go b/headerv1_test.go
@@ -10,12 +10,11 @@ import (
 
 func TestHeaderV1_WriteTo(t *testing.T) {
 	check := func(name string, hdr HeaderV1, exp string) {
-		t.Run(name, func(t *testing.T) {
-			buf := new(bytes.Buffer)
-			_, err := hdr.WriteTo(buf)
-			assert.NoError(t, err)
-			assert.Equal(t, exp, buf.String())
-		})
+		t.Helper()
+		buf := new(bytes.Buffer)
+		_, err := hdr.WriteTo(buf)
+		assert.NoError(t, err, name)
+		assert.Equal(t, exp, buf.String(), name)
 	}
 
 	check("blank", HeaderV1{}, "PROXY UNKNOWN\r\n")
@@ -43,5 +42,4 @@ func TestHeaderV1_WriteTo(t *testing.T) {
 	},
 		"PROXY TCP6 2001:db8:85a3::8a2e:370:7334 2002:db8:85a3::8a2e:370:7334 1234 5678\r\n",
 	)
-
 }
diff --git a/parse_test.go b/parse_test.go
@@ -3,6 +3,8 @@ package proxyprotocol
 import (
 	"bufio"
 	"bytes"
+	"net"
+	"strings"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
@@ -26,3 +28,44 @@ func TestParse_Malformed(t *testing.T) {
 			bytes.NewReader(data)))
 	assert.Error(t, err)
 }
+
+func TestParse_HeaderV1(t *testing.T) {
+	check := func(name string, hdr HeaderV1, exp string) {
+		t.Helper()
+
+		h, err := Parse(bufio.NewReader(strings.NewReader(exp)))
+		assert.NoError(t, err, name)
+		assert.Equal(t, 1, h.Version(), name+" version")
+
+		h1 := h.(*HeaderV1)
+		assert.Equal(t, hdr, *h1, name)
+	}
+
+	check("blank", HeaderV1{}, "PROXY UNKNOWN\r\n")
+	check("ipv4", HeaderV1{
+		SrcPort:  1234,
+		DestPort: 5678,
+		SrcIP:    net.ParseIP("192.168.0.1"),
+		DestIP:   net.ParseIP("192.168.0.2"),
+	},
+		"PROXY TCP4 192.168.0.1 192.168.0.2 1234 5678\r\n",
+	)
+
+	check("ipv6", HeaderV1{
+		SrcPort:  1234,
+		DestPort: 5678,
+		SrcIP:    net.ParseIP("2001:db8:85a3::8a2e:370:7334"),
+		DestIP:   net.ParseIP("2002:db8:85a3::8a2e:370:7334"),
+	},
+		"PROXY TCP6 2001:db8:85a3::8a2e:370:7334 2002:db8:85a3::8a2e:370:7334 1234 5678\r\n",
+	)
+
+	check("ipv6-mapped-ipv4", HeaderV1{
+		SrcPort:  53740,
+		DestPort: 10001,
+		SrcIP:    net.ParseIP("::ffff:192.168.0.1"),
+		DestIP:   net.ParseIP("::ffff:192.168.0.1"),
+	},
+		"PROXY TCP6 ::ffff:192.168.0.1 ::ffff:192.168.0.1 53740 10001\r\n",
+	)
+}