add new connpolicy which can validate upstream and downstream addresses

Author: kmala

policy.go

@@ -6,14 +6,28 @@ import (
 	"strings"
 )
 
-// PolicyFunc can be used to decide whether to trust the PROXY info based on
-// upstream/downstream IP. If set, the connecting addresses(remote and local)
-// are passed in as arguments.
+// PolicyFunc can be used to decide whether to trust the PROXY info from
+// upstream. If set, the connecting address is passed in as an argument.
 //
 // See below for the different policies.
 //
 // In case an error is returned the connection is denied.
-type PolicyFunc func(upstream net.Addr, downstream net.Addr) (Policy, error)
+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
@@ -44,7 +58,7 @@ const (
 // Kubernetes pods local traffic. The def is a policy to use when an upstream
 // address doesn't match the skipHeaderCIDR.
 func SkipProxyHeaderForCIDR(skipHeaderCIDR *net.IPNet, def Policy) PolicyFunc {
-	return func(upstream net.Addr, downstream net.Addr) (Policy, error) {
+	return func(upstream net.Addr) (Policy, error) {
 		ip, err := ipFromAddr(upstream)
 		if err != nil {
 			return def, err
@@ -58,25 +72,6 @@ func SkipProxyHeaderForCIDR(skipHeaderCIDR *net.IPNet, def Policy) PolicyFunc {
 	}
 }
 
-// IgnoreProxyHeaderNotOnInterface retuns a PolicyFunc 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) PolicyFunc {
-	return func(upstream net.Addr, downstream net.Addr) (Policy, error) {
-		ip, err := ipFromAddr(downstream)
-		if err != nil {
-			return REJECT, err
-		}
-
-		if allowedIP.Equal(ip) {
-			return USE, nil
-		}
-
-		return IGNORE, nil
-	}
-}
-
 // WithPolicy adds given policy to a connection when passed as option to NewConn()
 func WithPolicy(p Policy) func(*Conn) {
 	return func(c *Conn) {
@@ -137,7 +132,7 @@ func MustStrictWhiteListPolicy(allowed []string) PolicyFunc {
 }
 
 func whitelistPolicy(allowed []func(net.IP) bool, def Policy) PolicyFunc {
-	return func(upstream net.Addr, downstream net.Addr) (Policy, error) {
+	return func(upstream net.Addr) (Policy, error) {
 		upstreamIP, err := ipFromAddr(upstream)
 		if err != nil {
 			// something is wrong with the source IP, better reject the connection
@@ -190,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
+	}
+}

policy_test.go

@@ -21,7 +21,7 @@ func TestWhitelistPolicyReturnsErrorOnInvalidAddress(t *testing.T) {
 
 	for _, tc := range cases {
 		t.Run(tc.name, func(t *testing.T) {
-			_, err := tc.policy(failingAddr{}, nil)
+			_, err := tc.policy(failingAddr{})
 			if err == nil {
 				t.Fatal("Expected error, got none")
 			}
@@ -37,7 +37,7 @@ func TestStrictWhitelistPolicyReturnsRejectWhenUpstreamIpAddrNotInWhitelist(t *t
 		t.Fatalf("err: %v", err)
 	}
 
-	policy, err := p(upstream, nil)
+	policy, err := p(upstream)
 	if err != nil {
 		t.Fatalf("err: %v", err)
 	}
@@ -55,7 +55,7 @@ func TestLaxWhitelistPolicyReturnsIgnoreWhenUpstreamIpAddrNotInWhitelist(t *test
 		t.Fatalf("err: %v", err)
 	}
 
-	policy, err := p(upstream, nil)
+	policy, err := p(upstream)
 	if err != nil {
 		t.Fatalf("err: %v", err)
 	}
@@ -81,7 +81,7 @@ func TestWhitelistPolicyReturnsUseWhenUpstreamIpAddrInWhitelist(t *testing.T) {
 
 	for _, tc := range cases {
 		t.Run(tc.name, func(t *testing.T) {
-			policy, err := tc.policy(upstream, nil)
+			policy, err := tc.policy(upstream)
 			if err != nil {
 				t.Fatalf("err: %v", err)
 			}
@@ -109,7 +109,7 @@ func TestWhitelistPolicyReturnsUseWhenUpstreamIpAddrInWhitelistRange(t *testing.
 
 	for _, tc := range cases {
 		t.Run(tc.name, func(t *testing.T) {
-			policy, err := tc.policy(upstream, nil)
+			policy, err := tc.policy(upstream)
 			if err != nil {
 				t.Fatalf("err: %v", err)
 			}
@@ -194,7 +194,7 @@ func TestSkipProxyHeaderForCIDR(t *testing.T) {
 	f := SkipProxyHeaderForCIDR(cidr, REJECT)
 
 	upstream, _ := net.ResolveTCPAddr("tcp", "192.0.2.255:12345")
-	policy, err := f(upstream, nil)
+	policy, err := f(upstream)
 	if err != nil {
 		t.Fatalf("err: %v", err)
 	}
@@ -203,7 +203,7 @@ func TestSkipProxyHeaderForCIDR(t *testing.T) {
 	}
 
 	upstream, _ = net.ResolveTCPAddr("tcp", "8.8.8.8:12345")
-	policy, err = f(upstream, nil)
+	policy, err = f(upstream)
 	if err != nil {
 		t.Fatalf("err: %v", err)
 	}
@@ -220,7 +220,7 @@ func TestIgnoreProxyHeaderNotOnInterface(t *testing.T) {
 
 	var cases = []struct {
 		name              string
-		policy            PolicyFunc
+		policy            ConnPolicyFunc
 		downstreamAddress net.Addr
 		expectedPolicy    Policy
 		expectError       bool
@@ -232,7 +232,9 @@ func TestIgnoreProxyHeaderNotOnInterface(t *testing.T) {
 
 	for _, tc := range cases {
 		t.Run(tc.name, func(t *testing.T) {
-			policy, err := tc.policy(nil, tc.downstreamAddress)
+			policy, err := tc.policy(ConnPolicyOptions{
+				Downstream: tc.downstreamAddress,
+			})
 			if !tc.expectError && err != nil {
 				t.Fatalf("err: %v", err)
 			}

protocol.go

@@ -22,9 +22,13 @@ 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
 	Policy            PolicyFunc
+	ConnPolicy        ConnPolicyFunc
 	ValidateHeader    Validator
 	ReadHeaderTimeout time.Duration
 }
@@ -67,8 +71,18 @@ func (p *Listener) Accept() (net.Conn, error) {
 	}
 
 	proxyHeaderPolicy := USE
-	if p.Policy != nil {
-		proxyHeaderPolicy, err = p.Policy(conn.RemoteAddr(), conn.LocalAddr())
+	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()