Repository: pires/go-proxyproto
Pull request #52 (2020-11-07 17:08:45)

Add Conn.ProxyHeader
This function returns the full proxy protocol header, if available.
Sometimes the local/remote addresses aren't enough and the caller is
interested in e.g. TLVs.


Patch:
diff --git a/protocol.go b/protocol.go
index dea810d..30477bb 100644
--- a/protocol.go
+++ b/protocol.go
@@ -117,6 +117,13 @@ func (p *Conn) Close() error {
 	return p.conn.Close()
 }
 
+// ProxyHeader returns the proxy protocol header, if any. If an error occurs
+// while reading the proxy header, nil is returned.
+func (p *Conn) ProxyHeader() *Header {
+	p.once.Do(func() { p.readErr = p.readHeader() })
+	return p.header
+}
+
 // LocalAddr returns the address of the server if the proxy
 // protocol is being used, otherwise just returns the address of
 // the socket server. In case an error happens on reading the


Test patch:
diff --git a/protocol_test.go b/protocol_test.go
index 429cee7..daadd37 100644
--- a/protocol_test.go
+++ b/protocol_test.go
@@ -67,6 +67,19 @@ func TestParse_ipv4(t *testing.T) {
 
 	pl := &Listener{Listener: l}
 
+	header := &Header{
+		Version:           2,
+		Command:           PROXY,
+		TransportProtocol: TCPv4,
+		SourceAddr: &net.TCPAddr{
+			IP:   net.ParseIP("10.1.1.1"),
+			Port: 1000,
+		},
+		DestinationAddr: &net.TCPAddr{
+			IP:   net.ParseIP("20.2.2.2"),
+			Port: 2000,
+		},
+	}
 	go func() {
 		conn, err := net.Dial("tcp", pl.Addr().String())
 		if err != nil {
@@ -75,19 +88,6 @@ func TestParse_ipv4(t *testing.T) {
 		defer conn.Close()
 
 		// Write out the header!
-		header := &Header{
-			Version:           2,
-			Command:           PROXY,
-			TransportProtocol: TCPv4,
-			SourceAddr: &net.TCPAddr{
-				IP:   net.ParseIP("10.1.1.1"),
-				Port: 1000,
-			},
-			DestinationAddr: &net.TCPAddr{
-				IP:   net.ParseIP("20.2.2.2"),
-				Port: 2000,
-			},
-		}
 		header.WriteTo(conn)
 
 		conn.Write([]byte("ping"))
@@ -128,6 +128,11 @@ func TestParse_ipv4(t *testing.T) {
 	if addr.Port != 1000 {
 		t.Fatalf("bad: %v", addr)
 	}
+
+	h := conn.(*Conn).ProxyHeader()
+	if !h.EqualsTo(header) {
+		t.Errorf("bad: %v", h)
+	}
 }
 
 func TestParse_ipv6(t *testing.T) {
@@ -138,6 +143,20 @@ func TestParse_ipv6(t *testing.T) {
 
 	pl := &Listener{Listener: l}
 
+	header := &Header{
+		Version:           2,
+		Command:           PROXY,
+		TransportProtocol: TCPv6,
+		SourceAddr: &net.TCPAddr{
+			IP:   net.ParseIP("ffff::ffff"),
+			Port: 1000,
+		},
+		DestinationAddr: &net.TCPAddr{
+			IP:   net.ParseIP("ffff::ffff"),
+			Port: 2000,
+		},
+	}
+
 	go func() {
 		conn, err := net.Dial("tcp", pl.Addr().String())
 		if err != nil {
@@ -146,19 +165,6 @@ func TestParse_ipv6(t *testing.T) {
 		defer conn.Close()
 
 		// Write out the header!
-		header := &Header{
-			Version:           2,
-			Command:           PROXY,
-			TransportProtocol: TCPv6,
-			SourceAddr: &net.TCPAddr{
-				IP:   net.ParseIP("ffff::ffff"),
-				Port: 1000,
-			},
-			DestinationAddr: &net.TCPAddr{
-				IP:   net.ParseIP("ffff::ffff"),
-				Port: 2000,
-			},
-		}
 		header.WriteTo(conn)
 
 		conn.Write([]byte("ping"))
@@ -199,6 +205,11 @@ func TestParse_ipv6(t *testing.T) {
 	if addr.Port != 1000 {
 		t.Fatalf("bad: %v", addr)
 	}
+
+	h := conn.(*Conn).ProxyHeader()
+	if !h.EqualsTo(header) {
+		t.Errorf("bad: %v", h)
+	}
 }
 
 func TestAcceptReturnsErrorWhenPolicyFuncErrors(t *testing.T) {