Repository: pires/go-proxyproto (go)

Issue:
Mixing IPv6 and IPv4 in header
I noticed while doing a little testing that if you generate the SourceAddress and DestinationAddress, it's possible to get an IPv4 address and an IPv6 address.

Example:
```
net.LookupIP("localhost")
```

Can return `::1` as the net.IP, while it's up to the user to verify that they're sending IPv4/IPv6 addresses with v4/v6 protocols, I'm wondering if this `writeVersion1` and `writeVersion2` functions should return an error if the IPs return `nil` for the protocol being sent.

I'm willing to do the PR for it, but I just thought checking if the validation should be in the package or in the user code is the right way.

Pull request description:
Add more robust protocol mismatch handling in Header.formatVersion1
And add a bunch of invalid tests.

Fixes #6 (for v1 as v2 is already covered).

Patch:
diff --git a/v1.go b/v1.go
index 49d23c6..059faf9 100644
--- a/v1.go
+++ b/v1.go
@@ -77,26 +77,43 @@ func (header *Header) formatVersion1() ([]byte, error) {
 	// As of version 1, only "TCP4" ( \x54 \x43 \x50 \x34 ) for TCP over IPv4,
 	// and "TCP6" ( \x54 \x43 \x50 \x36 ) for TCP over IPv6 are allowed.
 	var proto string
-	if header.TransportProtocol == TCPv4 {
+	switch header.TransportProtocol {
+	case TCPv4:
 		proto = "TCP4"
-	} else if header.TransportProtocol == TCPv6 {
+	case TCPv6:
 		proto = "TCP6"
-	} else {
+	default:
 		// Unknown connection (short form)
 		return []byte("PROXY UNKNOWN\r\n"), nil
 	}
 
-	sourceAddr := header.SourceAddr.(*net.TCPAddr)
-	destAddr := header.DestinationAddr.(*net.TCPAddr)
+	sourceAddr, sourceOK := header.SourceAddr.(*net.TCPAddr)
+	destAddr, destOK := header.DestinationAddr.(*net.TCPAddr)
+	if !sourceOK || !destOK {
+		return nil, ErrInvalidAddress
+	}
+
+	sourceIP, destIP := sourceAddr.IP, destAddr.IP
+	switch header.TransportProtocol {
+	case TCPv4:
+		sourceIP = sourceIP.To4()
+		destIP = destIP.To4()
+	case TCPv6:
+		sourceIP = sourceIP.To16()
+		destIP = destIP.To16()
+	}
+	if sourceIP == nil || destIP == nil {
+		return nil, ErrInvalidAddress
+	}
 
 	buf := bytes.NewBuffer(make([]byte, 0, 108))
 	buf.Write(SIGV1)
 	buf.WriteString(separator)
 	buf.WriteString(proto)
 	buf.WriteString(separator)
-	buf.WriteString(sourceAddr.IP.String())
+	buf.WriteString(sourceIP.String())
 	buf.WriteString(separator)
-	buf.WriteString(destAddr.IP.String())
+	buf.WriteString(destIP.String())
 	buf.WriteString(separator)
 	buf.WriteString(strconv.Itoa(sourceAddr.Port))
 	buf.WriteString(separator)


Test patch:
diff --git a/header_test.go b/header_test.go
index 0312a99..8cb4977 100644
--- a/header_test.go
+++ b/header_test.go
@@ -447,27 +447,90 @@ func TestFormat(t *testing.T) {
 	if _, err := validHeader.Format(); err != nil {
 		t.Fatalf("shouldn't have thrown error %q", err.Error())
 	}
+}
 
-	invalidHeader := &Header{
-		Version:           3,
-		Command:           PROXY,
-		TransportProtocol: TCPv4,
-		SourceAddr: &net.TCPAddr{
-			IP:   net.ParseIP("10.1.1.1"),
-			Port: 1000,
+func TestFormatInvalid(t *testing.T) {
+	tests := []struct {
+		name   string
+		header *Header
+		err    error
+	}{
+		{
+			name: "invalidVersion",
+			header: &Header{
+				Version:           3,
+				Command:           PROXY,
+				TransportProtocol: TCPv4,
+				SourceAddr:        v4addr,
+				DestinationAddr:   v4addr,
+			},
+			err: ErrUnknownProxyProtocolVersion,
 		},
-		DestinationAddr: &net.TCPAddr{
-			IP:   net.ParseIP("20.2.2.2"),
-			Port: 2000,
+		{
+			name: "v2MismatchTCPv4_UDPv4",
+			header: &Header{
+				Version:           2,
+				Command:           PROXY,
+				TransportProtocol: TCPv4,
+				SourceAddr:        v4UDPAddr,
+				DestinationAddr:   v4addr,
+			},
+			err: ErrInvalidAddress,
+		},
+		{
+			name: "v2MismatchTCPv4_TCPv6",
+			header: &Header{
+				Version:           2,
+				Command:           PROXY,
+				TransportProtocol: TCPv4,
+				SourceAddr:        v4addr,
+				DestinationAddr:   v6addr,
+			},
+			err: ErrInvalidAddress,
+		},
+		{
+			name: "v2MismatchUnixStream_TCPv4",
+			header: &Header{
+				Version:           2,
+				Command:           PROXY,
+				TransportProtocol: UnixStream,
+				SourceAddr:        v4addr,
+				DestinationAddr:   unixStreamAddr,
+			},
+			err: ErrInvalidAddress,
+		},
+		{
+			name: "v1MismatchTCPv4_TCPv6",
+			header: &Header{
+				Version:           1,
+				Command:           PROXY,
+				TransportProtocol: TCPv4,
+				SourceAddr:        v6addr,
+				DestinationAddr:   v4addr,
+			},
+			err: ErrInvalidAddress,
+		},
+		{
+			name: "v1MismatchTCPv4_UDPv4",
+			header: &Header{
+				Version:           1,
+				Command:           PROXY,
+				TransportProtocol: TCPv4,
+				SourceAddr:        v4UDPAddr,
+				DestinationAddr:   v4addr,
+			},
+			err: ErrInvalidAddress,
 		},
 	}
 
-	if _, err := invalidHeader.Format(); err == nil {
-		t.Fatalf("should have thrown error %q", err.Error())
-	} else {
-		if err != ErrUnknownProxyProtocolVersion {
-			t.Fatalf("expected %q, actual %q", ErrUnknownProxyProtocolVersion.Error(), err.Error())
-		}
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			if _, err := test.header.Format(); err == nil {
+				t.Errorf("Header.Format() succeeded, want an error")
+			} else if err != test.err {
+				t.Errorf("Header.Format() = %q, want %q", err, test.err)
+			}
+		})
 	}
 }