Repository: pires/go-proxyproto
Pull request #62 (2021-01-19 06:30:07)

v2: read length and skip the bytes for UNSPEC
Fixes https://github.com/pires/go-proxyproto/pull/61#issuecomment-762556537

coveralls: 
[![Coverage Status](https://coveralls.io/builds/36423797/badge)](https://coveralls.io/builds/36423797)

Coverage decreased (-0.2%) to 93.935% when pulling **17a9bd6f857b788ca34c60505cd666299efac33a on bohanyang:unspec-zero-length** into **22bc6148fd29e67df91deb259efc1980c29bc4e5 on pires:main**.

bohanyang: Just found that there're two tests with desc `TCPv4 length zero but with address and ports`, should I fix the desc or the fixture header?
pires: Nice catch! Please, fix the desc of the last one which is clearly a test for TCPv6.  

Patch:
diff --git a/v2.go b/v2.go
index b0fdd58..1dd8db6 100644
--- a/v2.go
+++ b/v2.go
@@ -10,10 +10,15 @@ import (
 )
 
 var (
-	lengthV4   = uint16(12)
-	lengthV6   = uint16(36)
-	lengthUnix = uint16(216)
-
+	lengthUnspec      = uint16(0)
+	lengthV4          = uint16(12)
+	lengthV6          = uint16(36)
+	lengthUnix        = uint16(216)
+	lengthUnspecBytes = func() []byte {
+		a := make([]byte, 2)
+		binary.BigEndian.PutUint16(a, lengthUnspec)
+		return a
+	}()
 	lengthV4Bytes = func() []byte {
 		a := make([]byte, 2)
 		binary.BigEndian.PutUint16(a, lengthV4)
@@ -82,13 +87,9 @@ func parseVersion2(reader *bufio.Reader) (header *Header, err error) {
 		return nil, ErrCantReadAddressFamilyAndProtocol
 	}
 	header.TransportProtocol = AddressFamilyAndProtocol(b14)
-	//  UNSPEC is only supported when LOCAL is set.
-	if header.TransportProtocol == UNSPEC {
-		if header.Command != LOCAL {
-			return nil, ErrUnsupportedAddressFamilyAndProtocol
-		}
-		// Ignore everything else.
-		return header, nil
+	// UNSPEC is only supported when LOCAL is set.
+	if header.TransportProtocol == UNSPEC && header.Command != LOCAL {
+		return nil, ErrUnsupportedAddressFamilyAndProtocol
 	}
 
 	// Make sure there are bytes available as specified in length
@@ -100,6 +101,12 @@ func parseVersion2(reader *bufio.Reader) (header *Header, err error) {
 		return nil, ErrInvalidLength
 	}
 
+	// Return early if the length is zero, which means that
+	// there's no address information and TLVs present for UNSPEC.
+	if length == 0 {
+		return header, nil
+	}
+
 	if _, err := reader.Peek(int(length)); err != nil {
 		return nil, ErrInvalidLength
 	}
@@ -107,39 +114,43 @@ func parseVersion2(reader *bufio.Reader) (header *Header, err error) {
 	// Length-limited reader for payload section
 	payloadReader := io.LimitReader(reader, int64(length)).(*io.LimitedReader)
 
-	// Read addresses and ports
-	if header.TransportProtocol.IsIPv4() {
-		var addr _addr4
-		if err := binary.Read(payloadReader, binary.BigEndian, &addr); err != nil {
-			return nil, ErrInvalidAddress
-		}
-		header.SourceAddr = newIPAddr(header.TransportProtocol, addr.Src[:], addr.SrcPort)
-		header.DestinationAddr = newIPAddr(header.TransportProtocol, addr.Dst[:], addr.DstPort)
-	} else if header.TransportProtocol.IsIPv6() {
-		var addr _addr6
-		if err := binary.Read(payloadReader, binary.BigEndian, &addr); err != nil {
-			return nil, ErrInvalidAddress
-		}
-		header.SourceAddr = newIPAddr(header.TransportProtocol, addr.Src[:], addr.SrcPort)
-		header.DestinationAddr = newIPAddr(header.TransportProtocol, addr.Dst[:], addr.DstPort)
-	} else if header.TransportProtocol.IsUnix() {
-		var addr _addrUnix
-		if err := binary.Read(payloadReader, binary.BigEndian, &addr); err != nil {
-			return nil, ErrInvalidAddress
-		}
+	// Read addresses and ports for protocols other than UNSPEC.
+	// Ignore address information for UNSPEC, and skip straight to read TLVs,
+	// since the length is greater than zero.
+	if header.TransportProtocol != UNSPEC {
+		if header.TransportProtocol.IsIPv4() {
+			var addr _addr4
+			if err := binary.Read(payloadReader, binary.BigEndian, &addr); err != nil {
+				return nil, ErrInvalidAddress
+			}
+			header.SourceAddr = newIPAddr(header.TransportProtocol, addr.Src[:], addr.SrcPort)
+			header.DestinationAddr = newIPAddr(header.TransportProtocol, addr.Dst[:], addr.DstPort)
+		} else if header.TransportProtocol.IsIPv6() {
+			var addr _addr6
+			if err := binary.Read(payloadReader, binary.BigEndian, &addr); err != nil {
+				return nil, ErrInvalidAddress
+			}
+			header.SourceAddr = newIPAddr(header.TransportProtocol, addr.Src[:], addr.SrcPort)
+			header.DestinationAddr = newIPAddr(header.TransportProtocol, addr.Dst[:], addr.DstPort)
+		} else if header.TransportProtocol.IsUnix() {
+			var addr _addrUnix
+			if err := binary.Read(payloadReader, binary.BigEndian, &addr); err != nil {
+				return nil, ErrInvalidAddress
+			}
 
-		network := "unix"
-		if header.TransportProtocol.IsDatagram() {
-			network = "unixgram"
-		}
+			network := "unix"
+			if header.TransportProtocol.IsDatagram() {
+				network = "unixgram"
+			}
 
-		header.SourceAddr = &net.UnixAddr{
-			Net:  network,
-			Name: parseUnixName(addr.Src[:]),
-		}
-		header.DestinationAddr = &net.UnixAddr{
-			Net:  network,
-			Name: parseUnixName(addr.Dst[:]),
+			header.SourceAddr = &net.UnixAddr{
+				Net:  network,
+				Name: parseUnixName(addr.Src[:]),
+			}
+			header.DestinationAddr = &net.UnixAddr{
+				Net:  network,
+				Name: parseUnixName(addr.Dst[:]),
+			}
 		}
 	}
 
@@ -157,9 +168,14 @@ func (header *Header) formatVersion2() ([]byte, error) {
 	buf.Write(SIGV2)
 	buf.WriteByte(header.Command.toByte())
 	buf.WriteByte(header.TransportProtocol.toByte())
-	// When UNSPEC, the receiver must ignore addresses and ports.
-	// Therefore there's no point in writing it.
-	if !header.TransportProtocol.IsUnspec() {
+	if header.TransportProtocol.IsUnspec() {
+		// For UNSPEC, write no addresses and ports but only TLVs if they are present
+		hdrLen, err := addTLVLen(lengthUnspecBytes, len(header.rawTLVs))
+		if err != nil {
+			return nil, err
+		}
+		buf.Write(hdrLen)
+	} else {
 		var addrSrc, addrDst []byte
 		if header.TransportProtocol.IsIPv4() {
 			hdrLen, err := addTLVLen(lengthV4Bytes, len(header.rawTLVs))
@@ -221,6 +237,8 @@ func (header *Header) validateLength(length uint16) bool {
 		return length >= lengthV6
 	} else if header.TransportProtocol.IsUnix() {
 		return length >= lengthUnix
+	} else if header.TransportProtocol.IsUnspec() {
+		return length >= lengthUnspec
 	}
 	return false
 }


Test patch:
diff --git a/v2_test.go b/v2_test.go
index c6acda1..57592b3 100644
--- a/v2_test.go
+++ b/v2_test.go
@@ -56,6 +56,7 @@ var (
 	}()
 	fixtureIPv4V2TLV = fixtureWithTLV(lengthV4Bytes, fixtureIPv4Address, fixtureTLV)
 	fixtureIPv6V2TLV = fixtureWithTLV(lengthV6Bytes, fixtureIPv6Address, fixtureTLV)
+	fixtureUnspecTLV = fixtureWithTLV(lengthUnspecBytes, []byte{}, fixtureTLV)
 
 	// Arbitrary bytes following proxy bytes
 	arbitraryTailBytes = []byte{'\x99', '\x97', '\x98'}
@@ -116,6 +117,11 @@ var invalidParseV2Tests = []struct {
 		reader:        newBufioReader(append(SIGV2, byte(PROXY), byte(TCPv4), invalidRune)),
 		expectedError: ErrCantReadLength,
 	},
+	{
+		desc:          "unspec but no length",
+		reader:        newBufioReader(append(SIGV2, byte(LOCAL), byte(UNSPEC))),
+		expectedError: ErrCantReadLength,
+	},
 	{
 		desc:          "TCPv4 with mismatching length",
 		reader:        newBufioReader(append(append(SIGV2, byte(PROXY), byte(TCPv4)), lengthV4Bytes...)),
@@ -132,10 +138,15 @@ var invalidParseV2Tests = []struct {
 		expectedError: ErrInvalidLength,
 	},
 	{
-		desc:          "TCPv4 length zero but with address and ports",
+		desc:          "TCPv6 with IPv6 length but IPv4 address and ports",
 		reader:        newBufioReader(append(append(append(SIGV2, byte(PROXY), byte(TCPv6)), lengthV6Bytes...), fixtureIPv4Address...)),
 		expectedError: ErrInvalidLength,
 	},
+	{
+		desc:          "unspec length greater than zero but no TLVs",
+		reader:        newBufioReader(append(append(SIGV2, byte(LOCAL), byte(UNSPEC)), fixtureUnspecTLV[:2]...)),
+		expectedError: ErrInvalidLength,
+	},
 }
 
 func TestParseV2Invalid(t *testing.T) {
@@ -166,7 +177,7 @@ var validParseAndWriteV2Tests = []struct {
 	},
 	{
 		desc:   "local unspec",
-		reader: newBufioReader(append(append(SIGV2, byte(LOCAL), byte(TCPv4)), fixtureIPv4V2...)),
+		reader: newBufioReader(append(append(SIGV2, byte(LOCAL), byte(UNSPEC)), lengthUnspecBytes...)),
 		expectedHeader: &Header{
 			Version:           2,
 			Command:           LOCAL,
@@ -221,6 +232,18 @@ var validParseAndWriteV2Tests = []struct {
 			rawTLVs:           fixtureTLV,
 		},
 	},
+	{
+		desc:   "local unspec with TLV",
+		reader: newBufioReader(append(append(SIGV2, byte(LOCAL), byte(UNSPEC)), fixtureUnspecTLV...)),
+		expectedHeader: &Header{
+			Version:           2,
+			Command:           LOCAL,
+			TransportProtocol: UNSPEC,
+			SourceAddr:        nil,
+			DestinationAddr:   nil,
+			rawTLVs:           fixtureTLV,
+		},
+	},
 	{
 		desc:   "proxy UDPv4",
 		reader: newBufioReader(append(append(SIGV2, byte(PROXY), byte(UDPv4)), fixtureIPv4V2...)),
@@ -255,7 +278,7 @@ var validParseAndWriteV2Tests = []struct {
 		},
 	},
 	{
-		desc:   "proxy unix datagram ",
+		desc:   "proxy unix datagram",
 		reader: newBufioReader(append(append(SIGV2, byte(PROXY), byte(UnixDatagram)), fixtureUnixV2...)),
 		expectedHeader: &Header{
 			Version:           2,
@@ -460,6 +483,17 @@ var tlvFormatTests = []struct {
 			rawTLVs:           make([]byte, 1<<16),
 		},
 	},
+	{
+		desc: "local unspec",
+		header: &Header{
+			Version:           2,
+			Command:           LOCAL,
+			TransportProtocol: UNSPEC,
+			SourceAddr:        nil,
+			DestinationAddr:   nil,
+			rawTLVs:           make([]byte, 1<<16),
+		},
+	},
 }
 
 func TestV2TLVFormatTooLargeTLV(t *testing.T) {