Repository: pires/go-proxyproto
Pull request #104 (2023-10-29 13:34:05)

header.go prevent otherHeader.rawTLVs being nil
If header.Version != otherHeader.Version then otherHeader.rawTLVs is not present.

This resulted in a crash at my side.  Switching position of the 2 ifs solves it.

pires: Thank you for contributing.

Can you, please, add a test that covers for your scenario, ie that triggers the panic?
coveralls: 
[![Coverage Status](https://coveralls.io/builds/63852699/badge)](https://coveralls.io/builds/63852699)

coverage: 95.017%. remained the same
when pulling **7095f70d0d8969d24541970a363026174d0ac525 on hilbix:tino**
into **e5b291b295b4a6e0425a134c8b3bfa6b2e692866 on pires:main**.


Patch:
diff --git a/go.mod b/go.mod
index b887d98..0148202 100644
--- a/go.mod
+++ b/go.mod
@@ -2,6 +2,6 @@ module github.com/pires/go-proxyproto
 
 go 1.18
 
-require golang.org/x/net v0.17.0
+require golang.org/x/net v0.23.0
 
-require golang.org/x/text v0.13.0 // indirect
+require golang.org/x/text v0.14.0 // indirect
diff --git a/go.sum b/go.sum
index 66be8d9..345dfc9 100644
--- a/go.sum
+++ b/go.sum
@@ -1,4 +1,4 @@
-golang.org/x/net v0.17.0 h1:pVaXccu2ozPjCXewfr1S7xza/zcXTity9cCdXQYSjIM=
-golang.org/x/net v0.17.0/go.mod h1:NxSsAGuq816PNPmqtQdLE42eU2Fs7NoRIZrHJAlaCOE=
-golang.org/x/text v0.13.0 h1:ablQoSUd0tRdKxZewP80B+BaqeKJuVhuRxj/dkrun3k=
-golang.org/x/text v0.13.0/go.mod h1:TvPlkZtksWOMsz7fbANvkp4WM8x/WCo/om8BMLbz+aE=
+golang.org/x/net v0.23.0 h1:7EYJ93RZ9vYSZAIb2x3lnuvqO5zneoD6IvWjuhfxjTs=
+golang.org/x/net v0.23.0/go.mod h1:JKghWKKOSdJwpW2GEx0Ja7fmaKnMsbu+MWVZTokSYmg=
+golang.org/x/text v0.14.0 h1:ScX5w1eTa3QqT8oi6+ziP7dTV1S2+ALU0bI+0zXKWiQ=
+golang.org/x/text v0.14.0/go.mod h1:18ZOQIKpY8NJVqYksKHtTdi31H5itFRjB5/qKTNYzSU=
diff --git a/header.go b/header.go
index 81ebeb3..209c2cc 100644
--- a/header.go
+++ b/header.go
@@ -155,11 +155,11 @@ func (header *Header) EqualsTo(otherHeader *Header) bool {
 	if otherHeader == nil {
 		return false
 	}
-	// TLVs only exist for version 2
-	if header.Version == 2 && !bytes.Equal(header.rawTLVs, otherHeader.rawTLVs) {
+	if header.Version != otherHeader.Version || header.Command != otherHeader.Command || header.TransportProtocol != otherHeader.TransportProtocol {
 		return false
 	}
-	if header.Version != otherHeader.Version || header.Command != otherHeader.Command || header.TransportProtocol != otherHeader.TransportProtocol {
+	// TLVs only exist for version 2
+	if header.Version == 2 && !bytes.Equal(header.rawTLVs, otherHeader.rawTLVs) {
 		return false
 	}
 	// Return early for header with LOCAL command, which contains no address information
diff --git a/policy.go b/policy.go
index 6d505be..ebef8b9 100644
--- a/policy.go
+++ b/policy.go
@@ -14,6 +14,21 @@ import (
 // In case an error is returned the connection is denied.
 type PolicyFunc func(upstream net.Addr) (Policy, error)
 
+// ConnPolicyFunc can be used to decide whether to trust the PROXY info
+// based on connection policy options. If set, the connecting addresses
+// (remote and local) are passed in as argument.
+//
+// See below for the different policies.
+//
+// In case an error is returned the connection is denied.
+type ConnPolicyFunc func(connPolicyOptions ConnPolicyOptions) (Policy, error)
+
+// ConnPolicyOptions contains the remote and local addresses of a connection.
+type ConnPolicyOptions struct {
+	Upstream   net.Addr
+	Downstream net.Addr
+}
+
 // Policy defines how a connection with a PROXY header address is treated.
 type Policy int
 
@@ -170,3 +185,22 @@ func ipFromAddr(upstream net.Addr) (net.IP, error) {
 
 	return upstreamIP, nil
 }
+
+// IgnoreProxyHeaderNotOnInterface retuns a ConnPolicyFunc which can be used to
+// decide whether to use or ignore PROXY headers depending on the connection
+// being made on a specific interface. This policy can be used when the server
+// is bound to multiple interfaces but wants to allow on only one interface.
+func IgnoreProxyHeaderNotOnInterface(allowedIP net.IP) ConnPolicyFunc {
+	return func(connOpts ConnPolicyOptions) (Policy, error) {
+		ip, err := ipFromAddr(connOpts.Downstream)
+		if err != nil {
+			return REJECT, err
+		}
+
+		if allowedIP.Equal(ip) {
+			return USE, nil
+		}
+
+		return IGNORE, nil
+	}
+}
diff --git a/protocol.go b/protocol.go
index 4ce16a2..3b40306 100644
--- a/protocol.go
+++ b/protocol.go
@@ -22,9 +22,14 @@ var DefaultReadHeaderTimeout = 10 * time.Second
 // connections in order to prevent blocking operations. If no ReadHeaderTimeout
 // is set, a default of 200ms will be used. This can be disabled by setting the
 // timeout to < 0.
+//
+// Only one of Policy or ConnPolicy should be provided. If both are provided then
+// a panic would occur during accept.
 type Listener struct {
-	Listener          net.Listener
+	Listener net.Listener
+	// Deprecated: use ConnPolicyFunc instead. This will be removed in future release.
 	Policy            PolicyFunc
+	ConnPolicy        ConnPolicyFunc
 	ValidateHeader    Validator
 	ReadHeaderTimeout time.Duration
 }
@@ -58,6 +63,15 @@ func ValidateHeader(v Validator) func(*Conn) {
 	}
 }
 
+// SetReadHeaderTimeout sets the readHeaderTimeout for a connection when passed as option to NewConn()
+func SetReadHeaderTimeout(t time.Duration) func(*Conn) {
+	return func(c *Conn) {
+		if t >= 0 {
+			c.readHeaderTimeout = t
+		}
+	}
+}
+
 // Accept waits for and returns the next connection to the listener.
 func (p *Listener) Accept() (net.Conn, error) {
 	// Get the underlying connection
@@ -67,8 +81,18 @@ func (p *Listener) Accept() (net.Conn, error) {
 	}
 
 	proxyHeaderPolicy := USE
-	if p.Policy != nil {
-		proxyHeaderPolicy, err = p.Policy(conn.RemoteAddr())
+	if p.Policy != nil && p.ConnPolicy != nil {
+		panic("only one of policy or connpolicy must be provided.")
+	}
+	if p.Policy != nil || p.ConnPolicy != nil {
+		if p.Policy != nil {
+			proxyHeaderPolicy, err = p.Policy(conn.RemoteAddr())
+		} else {
+			proxyHeaderPolicy, err = p.ConnPolicy(ConnPolicyOptions{
+				Upstream:   conn.RemoteAddr(),
+				Downstream: conn.LocalAddr(),
+			})
+		}
 		if err != nil {
 			// can't decide the policy, we can't accept the connection
 			conn.Close()


Test patch:
diff --git a/policy_test.go b/policy_test.go
index c8b2624..a888bdd 100644
--- a/policy_test.go
+++ b/policy_test.go
@@ -211,3 +211,41 @@ func TestSkipProxyHeaderForCIDR(t *testing.T) {
 		t.Errorf("Expected a REJECT policy for the %s address", upstream)
 	}
 }
+
+func TestIgnoreProxyHeaderNotOnInterface(t *testing.T) {
+	downstream, err := net.ResolveTCPAddr("tcp", "10.0.0.3:45738")
+	if err != nil {
+		t.Fatalf("err: %v", err)
+	}
+
+	var cases = []struct {
+		name              string
+		policy            ConnPolicyFunc
+		downstreamAddress net.Addr
+		expectedPolicy    Policy
+		expectError       bool
+	}{
+		{"ignore header for requests non on interface", IgnoreProxyHeaderNotOnInterface(net.ParseIP("192.0.2.1")), downstream, IGNORE, false},
+		{"use headers for requests on interface", IgnoreProxyHeaderNotOnInterface(net.ParseIP("10.0.0.3")), downstream, USE, false},
+		{"invalid address should return error", IgnoreProxyHeaderNotOnInterface(net.ParseIP("10.0.0.3")), failingAddr{}, REJECT, true},
+	}
+
+	for _, tc := range cases {
+		t.Run(tc.name, func(t *testing.T) {
+			policy, err := tc.policy(ConnPolicyOptions{
+				Downstream: tc.downstreamAddress,
+			})
+			if !tc.expectError && err != nil {
+				t.Fatalf("err: %v", err)
+			}
+			if tc.expectError && err == nil {
+				t.Fatal("Expected error, got none")
+			}
+
+			if policy != tc.expectedPolicy {
+				t.Fatalf("Expected policy %v, got %v", tc.expectedPolicy, policy)
+			}
+		})
+	}
+
+}
diff --git a/protocol_test.go b/protocol_test.go
index b1fb2c8..fd976d1 100644
--- a/protocol_test.go
+++ b/protocol_test.go
@@ -675,6 +675,82 @@ func TestAcceptReturnsErrorWhenPolicyFuncErrors(t *testing.T) {
 	}
 }
 
+func TestPanicIfPolicyAndConnPolicySet(t *testing.T) {
+	l, err := net.Listen("tcp", "127.0.0.1:0")
+	if err != nil {
+		t.Fatalf("err: %v", err)
+	}
+
+	connPolicyFunc := func(connopts ConnPolicyOptions) (Policy, error) { return USE, nil }
+	policyFunc := func(upstream net.Addr) (Policy, error) { return USE, nil }
+
+	pl := &Listener{Listener: l, ConnPolicy: connPolicyFunc, Policy: policyFunc}
+
+	cliResult := make(chan error)
+	go func() {
+		conn, err := net.Dial("tcp", pl.Addr().String())
+		if err != nil {
+			cliResult <- err
+			return
+		}
+		defer conn.Close()
+
+		close(cliResult)
+	}()
+
+	defer func() {
+		if r := recover(); r != nil {
+			fmt.Printf("accept did panic as expected with error, %v", r)
+		}
+	}()
+	conn, err := pl.Accept()
+	if err != nil {
+		t.Fatalf("Expected the accept to panic but did not and error is returned, got %v", err)
+	}
+
+	if conn != nil {
+		t.Fatalf("xpected the accept to panic but did not, got %v", conn)
+	}
+	t.Fatalf("expected the accept to panic but did not")
+}
+
+func TestAcceptReturnsErrorWhenConnPolicyFuncErrors(t *testing.T) {
+	l, err := net.Listen("tcp", "127.0.0.1:0")
+	if err != nil {
+		t.Fatalf("err: %v", err)
+	}
+
+	expectedErr := fmt.Errorf("failure")
+	connPolicyFunc := func(connopts ConnPolicyOptions) (Policy, error) { return USE, expectedErr }
+
+	pl := &Listener{Listener: l, ConnPolicy: connPolicyFunc}
+
+	cliResult := make(chan error)
+	go func() {
+		conn, err := net.Dial("tcp", pl.Addr().String())
+		if err != nil {
+			cliResult <- err
+			return
+		}
+		defer conn.Close()
+
+		close(cliResult)
+	}()
+
+	conn, err := pl.Accept()
+	if err != expectedErr {
+		t.Fatalf("Expected error %v, got %v", expectedErr, err)
+	}
+
+	if conn != nil {
+		t.Fatalf("Expected no connection, got %v", conn)
+	}
+	err = <-cliResult
+	if err != nil {
+		t.Fatalf("client error: %v", err)
+	}
+}
+
 func TestReadingIsRefusedWhenProxyHeaderRequiredButMissing(t *testing.T) {
 	l, err := net.Listen("tcp", "127.0.0.1:0")
 	if err != nil {
@@ -979,6 +1055,63 @@ func TestSkipProxyProtocolPolicy(t *testing.T) {
 		t.Fatalf("err: %v", err)
 	}
 
+	connPolicyFunc := func(connopts ConnPolicyOptions) (Policy, error) { return SKIP, nil }
+
+	pl := &Listener{
+		Listener:   l,
+		ConnPolicy: connPolicyFunc,
+	}
+
+	cliResult := make(chan error)
+	ping := []byte("ping")
+	go func() {
+		conn, err := net.Dial("tcp", pl.Addr().String())
+		if err != nil {
+			cliResult <- err
+			return
+		}
+		defer conn.Close()
+
+		if _, err := conn.Write(ping); err != nil {
+			cliResult <- err
+			return
+		}
+
+		close(cliResult)
+	}()
+
+	conn, err := pl.Accept()
+	if err != nil {
+		t.Fatalf("err: %v", err)
+	}
+	defer conn.Close()
+
+	_, ok := conn.(*net.TCPConn)
+	if !ok {
+		t.Fatal("err: should be a tcp connection")
+	}
+	_ = conn.LocalAddr()
+	recv := make([]byte, 4)
+	if _, err = conn.Read(recv); err != nil {
+		t.Fatalf("Unexpected read error: %v", err)
+	}
+
+	if !bytes.Equal(ping, recv) {
+		t.Fatalf("Unexpected %s data while expected %s", recv, ping)
+	}
+
+	err = <-cliResult
+	if err != nil {
+		t.Fatalf("client error: %v", err)
+	}
+}
+
+func TestSkipProxyProtocolConnPolicy(t *testing.T) {
+	l, err := net.Listen("tcp", "127.0.0.1:0")
+	if err != nil {
+		t.Fatalf("err: %v", err)
+	}
+
 	policyFunc := func(upstream net.Addr) (Policy, error) { return SKIP, nil }
 
 	pl := &Listener{