Repository: pires/go-proxyproto
Pull request #11 (2019-08-28 12:39:53)

Add possibility to handle proxy info differently depending on upstream IP
This introduces the possibility to decide on each new incoming connection whether the proxy header should be used, ignored, or the connection be rejected. It also allows to reject the connection when a proxy header is required but not send at all.

Without such a possibility usage of this library is prone to security attacks enabling simple faking of client IP addresses by accepting any proxy protocol header regardless of whether the upstream source is allowed to send one or not. The need for this is derived from our use case, where we have an application which accepts connections both through a proxy and direct (there's a reason for that). If we would allow the proxy header for direct connections it would allow clients to pretend they are from another IP address which we don't want.

While `NewConn` is part of the public API I decided to change it nevertheless as no official release or version of this library exists yet. In case you want to keep it as is I'd introduce another function  `NewConnWithPolicy` instead where the caller can provide the policy, and fall back to the `USE` policy in `NewConn`.

pires: I will take the weekend to look into this.
pires: I like this idea a lot! Can you please rebase and fix conflicts?

I'm totally fine with merging and breaking the API since there's no release and people can simply rely on the commit of their choosing.
Actually, I just removed the timeout parameter since it wasn't even being used, so yeah.

Patch:
diff --git a/header.go b/header.go
index 6dbaab1..7336a12 100644
--- a/header.go
+++ b/header.go
@@ -28,6 +28,7 @@ var (
 	ErrInvalidLength                        = errors.New("Invalid length")
 	ErrInvalidAddress                       = errors.New("Invalid address")
 	ErrInvalidPortNumber                    = errors.New("Invalid port number")
+	ErrSuperfluousProxyHeader               = errors.New("Upstream connection sent PROXY header but isn't allowed to send one")
 )
 
 // Header is the placeholder for proxy protocol header.
diff --git a/policy.go b/policy.go
new file mode 100644
index 0000000..c76ec67
--- /dev/null
+++ b/policy.go
@@ -0,0 +1,149 @@
+package proxyproto
+
+import (
+	"fmt"
+	"net"
+	"strings"
+)
+
+// 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) (Policy, error)
+
+// Policy defines how a connection with a PROXY header address is treated.
+type Policy int
+
+const (
+	// USE address from PROXY header
+	USE Policy = iota
+	// IGNORE address from PROXY header, but accept connection
+	IGNORE
+	// REJECT connection when PROXY header is sent
+	// Note: even though the first read on the connection returns an error if
+	// a PROXY header is present, subsequent reads do not. It is the task of
+	// the code using the connection to handle that case properly.
+	REJECT
+	// REQUIRE connection to send PROXY header, reject if not present
+	// Note: even though the first read on the connection returns an error if
+	// a PROXY header is not present, subsequent reads do not. It is the task
+	// of the code using the connection to handle that case properly.
+	REQUIRE
+)
+
+// WithPolicy adds given policy to a connection when passed as option to NewConn()
+func WithPolicy(p Policy) func(*Conn) {
+	return func(c *Conn) {
+		c.ProxyHeaderPolicy = p
+	}
+}
+
+// LaxWhiteListPolicy returns a PolicyFunc which decides whether the
+// upstream ip is allowed to send a proxy header based on a list of allowed
+// IP addresses and IP ranges. In case upstream IP is not in list the proxy
+// header will be ignored. If one of the provided IP addresses or IP ranges
+// is invalid it will return an error instead of a PolicyFunc.
+func LaxWhiteListPolicy(allowed []string) (PolicyFunc, error) {
+	allowFrom, err := parse(allowed)
+	if err != nil {
+		return nil, err
+	}
+
+	return whitelistPolicy(allowFrom, IGNORE), nil
+}
+
+// MustLaxWhiteListPolicy returns a LaxWhiteListPolicy but will panic if one
+// of the provided IP addresses or IP ranges is invalid.
+func MustLaxWhiteListPolicy(allowed []string) PolicyFunc {
+	pfunc, err := LaxWhiteListPolicy(allowed)
+	if err != nil {
+		panic(err)
+	}
+
+	return pfunc
+}
+
+// StrictWhiteListPolicy returns a PolicyFunc which decides whether the
+// upstream ip is allowed to send a proxy header based on a list of allowed
+// IP addresses and IP ranges. In case upstream IP is not in list reading on
+// the connection will be refused on the first read. Please note: subsequent
+// reads do not error. It is the task of the code using the connection to
+// handle that case properly. If one of the provided IP addresses or IP
+// ranges is invalid it will return an error instead of a PolicyFunc.
+func StrictWhiteListPolicy(allowed []string) (PolicyFunc, error) {
+	allowFrom, err := parse(allowed)
+	if err != nil {
+		return nil, err
+	}
+
+	return whitelistPolicy(allowFrom, REJECT), nil
+}
+
+// MustStrictWhiteListPolicy returns a StrictWhiteListPolicy but will panic
+// if one of the provided IP addresses or IP ranges is invalid.
+func MustStrictWhiteListPolicy(allowed []string) PolicyFunc {
+	pfunc, err := StrictWhiteListPolicy(allowed)
+	if err != nil {
+		panic(err)
+	}
+
+	return pfunc
+}
+
+func whitelistPolicy(allowed []func(net.IP) bool, def Policy) PolicyFunc {
+	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
+			return REJECT, err
+		}
+
+		for _, allowFrom := range allowed {
+			if allowFrom(upstreamIP) {
+				return USE, nil
+			}
+		}
+
+		return def, nil
+	}
+}
+
+func parse(allowed []string) ([]func(net.IP) bool, error) {
+	a := make([]func(net.IP) bool, len(allowed))
+	for i, allowFrom := range allowed {
+		if strings.LastIndex(allowFrom, "/") > 0 {
+			_, ipRange, err := net.ParseCIDR(allowFrom)
+			if err != nil {
+				return nil, fmt.Errorf("Given string %q is not a valid IP range: %v", allowFrom, err)
+			}
+
+			a[i] = ipRange.Contains
+		} else {
+			allowed := net.ParseIP(allowFrom)
+			if allowed == nil {
+				return nil, fmt.Errorf("Given string %q is not a valid IP address", allowFrom)
+			}
+
+			a[i] = allowed.Equal
+		}
+	}
+
+	return a, nil
+}
+
+func ipFromAddr(upstream net.Addr) (net.IP, error) {
+	upstreamString, _, err := net.SplitHostPort(upstream.String())
+	if err != nil {
+		return nil, err
+	}
+
+	upstreamIP := net.ParseIP(upstreamString)
+	if nil == upstreamIP {
+		return nil, fmt.Errorf("invalid IP address")
+	}
+
+	return upstreamIP, nil
+}
diff --git a/protocol.go b/protocol.go
index d577970..870d2b9 100644
--- a/protocol.go
+++ b/protocol.go
@@ -12,17 +12,35 @@ import (
 // If the connection is using the protocol, the RemoteAddr() will return
 // the correct client address.
 type Listener struct {
-	Listener net.Listener
+	Listener       net.Listener
+	Policy         PolicyFunc
+	ValidateHeader Validator
 }
 
 // Conn is used to wrap and underlying connection which
 // may be speaking the Proxy Protocol. If it is, the RemoteAddr() will
 // return the address of the client instead of the proxy address.
 type Conn struct {
-	bufReader *bufio.Reader
-	conn      net.Conn
-	header    *Header
-	once      sync.Once
+	bufReader         *bufio.Reader
+	conn              net.Conn
+	header            *Header
+	once              sync.Once
+	ProxyHeaderPolicy Policy
+	Validate          Validator
+	readErr           error
+}
+
+// Validator receives a header and decides whether it is a valid one
+// In case the header is not deemed valid it should return an error.
+type Validator func(*Header) error
+
+// ValidateHeader adds given validator for proxy headers to a connection when passed as option to NewConn()
+func ValidateHeader(v Validator) func(*Conn) {
+	return func(c *Conn) {
+		if v != nil {
+			c.Validate = v
+		}
+	}
 }
 
 // Accept waits for and returns the next connection to the listener.
@@ -32,7 +50,23 @@ func (p *Listener) Accept() (net.Conn, error) {
 	if err != nil {
 		return nil, err
 	}
-	return NewConn(conn), nil
+
+	proxyHeaderPolicy := USE
+	if p.Policy != nil {
+		proxyHeaderPolicy, err = p.Policy(conn.RemoteAddr())
+		if err != nil {
+			// can't decide the policy, we can't accept the connection
+			conn.Close()
+			return nil, err
+		}
+	}
+
+	newConn := NewConn(
+		conn,
+		WithPolicy(proxyHeaderPolicy),
+		ValidateHeader(p.ValidateHeader),
+	)
+	return newConn, nil
 }
 
 // Close closes the underlying listener.
@@ -47,11 +81,16 @@ func (p *Listener) Addr() net.Addr {
 
 // NewConn is used to wrap a net.Conn that may be speaking
 // the proxy protocol into a proxyproto.Conn
-func NewConn(conn net.Conn) *Conn {
+func NewConn(conn net.Conn, opts ...func(*Conn)) *Conn {
 	pConn := &Conn{
 		bufReader: bufio.NewReader(conn),
 		conn:      conn,
 	}
+
+	for _, opt := range opts {
+		opt(pConn)
+	}
+
 	return pConn
 }
 
@@ -59,12 +98,11 @@ func NewConn(conn net.Conn) *Conn {
 // the initial scan. If there is an error parsing the header,
 // it is returned and the socket is closed.
 func (p *Conn) Read(b []byte) (int, error) {
-	var err error
 	p.once.Do(func() {
-		err = p.readHeader()
+		p.readErr = p.readHeader()
 	})
-	if err != nil {
-		return 0, err
+	if p.readErr != nil {
+		return 0, p.readErr
 	}
 	return p.bufReader.Read(b)
 }
@@ -81,10 +119,13 @@ func (p *Conn) Close() error {
 
 // LocalAddr returns the address of the server if the proxy
 // protocol is being used, otherwise just returns the address of
-// the socket server.
+// the socket server. In case an error happens on reading the
+// proxy header the original LocalAddr is returned, not the one
+// from the proxy header even if the proxy header itself is
+// syntactically correct.
 func (p *Conn) LocalAddr() net.Addr {
-	p.once.Do(func() { p.readHeader() })
-	if p.header == nil || p.header.Command.IsLocal() {
+	p.once.Do(func() { p.readErr = p.readHeader() })
+	if p.header == nil || p.header.Command.IsLocal() || p.readErr != nil {
 		return p.conn.LocalAddr()
 	}
 
@@ -93,10 +134,13 @@ func (p *Conn) LocalAddr() net.Addr {
 
 // RemoteAddr returns the address of the client if the proxy
 // protocol is being used, otherwise just returns the address of
-// the socket peer.
+// the socket peer. In case an error happens on reading the
+// proxy header the original RemoteAddr is returned, not the one
+// from the proxy header even if the proxy header itself is
+// syntactically correct.
 func (p *Conn) RemoteAddr() net.Addr {
-	p.once.Do(func() { p.readHeader() })
-	if p.header == nil || p.header.Command.IsLocal() {
+	p.once.Do(func() { p.readErr = p.readHeader() })
+	if p.header == nil || p.header.Command.IsLocal() || p.readErr != nil {
 		return p.conn.RemoteAddr()
 	}
 
@@ -118,13 +162,36 @@ func (p *Conn) SetWriteDeadline(t time.Time) error {
 	return p.conn.SetWriteDeadline(t)
 }
 
-func (p *Conn) readHeader() (err error) {
-	p.header, err = Read(p.bufReader)
+func (p *Conn) readHeader() error {
+	header, err := Read(p.bufReader)
 	// For the purpose of this wrapper shamefully stolen from armon/go-proxyproto
 	// let's act as if there was no error when PROXY protocol is not present.
 	if err == ErrNoProxyProtocol {
-		err = nil
+		// but not if it is required that the connection has one
+		if p.ProxyHeaderPolicy == REQUIRE {
+			return err
+		}
+
+		return nil
+	}
+
+	// proxy protocol header was found
+	if err == nil && header != nil {
+		switch p.ProxyHeaderPolicy {
+		case REJECT:
+			// this connection is not allowed to send one
+			return ErrSuperfluousProxyHeader
+		case USE:
+			if p.Validate != nil {
+				err = p.Validate(header)
+				if err != nil {
+					return err
+				}
+			}
+
+			p.header = header
+		}
 	}
 
-	return
+	return err
 }


Test patch:
diff --git a/policy_test.go b/policy_test.go
new file mode 100644
index 0000000..40a9444
--- /dev/null
+++ b/policy_test.go
@@ -0,0 +1,190 @@
+package proxyproto
+
+import (
+	"net"
+	"testing"
+)
+
+type failingAddr struct{}
+
+func (f failingAddr) Network() string { return "failing" }
+func (f failingAddr) String() string  { return "failing" }
+
+func TestWhitelistPolicyReturnsErrorOnInvalidAddress(t *testing.T) {
+	var cases = []struct {
+		name   string
+		policy PolicyFunc
+	}{
+		{"strict whitelist policy", MustStrictWhiteListPolicy([]string{"10.0.0.2", "10.0.0.3", "10.0.0.4", "10.0.0.0/30"})},
+		{"lax whitelist policy", MustLaxWhiteListPolicy([]string{"10.0.0.2", "10.0.0.3", "10.0.0.4", "10.0.0.0/30"})},
+	}
+
+	for _, tc := range cases {
+		t.Run(tc.name, func(t *testing.T) {
+			_, err := tc.policy(failingAddr{})
+			if err == nil {
+				t.Fatal("Expected error, got none")
+			}
+		})
+	}
+}
+
+func TestStrictWhitelistPolicyReturnsRejectWhenUpstreamIpAddrNotInWhitelist(t *testing.T) {
+	p := MustStrictWhiteListPolicy([]string{"10.0.0.2", "10.0.0.3", "10.0.0.4", "10.0.0.0/30"})
+
+	upstream, err := net.ResolveTCPAddr("tcp", "10.0.0.5:45738")
+	if err != nil {
+		t.Fatalf("err: %v", err)
+	}
+
+	policy, err := p(upstream)
+	if err != nil {
+		t.Fatalf("err: %v", err)
+	}
+
+	if policy != REJECT {
+		t.Fatalf("Expected policy REJECT, got %v", policy)
+	}
+}
+
+func TestLaxWhitelistPolicyReturnsIgnoreWhenUpstreamIpAddrNotInWhitelist(t *testing.T) {
+	p := MustLaxWhiteListPolicy([]string{"10.0.0.2", "10.0.0.3", "10.0.0.4", "10.0.0.0/30"})
+
+	upstream, err := net.ResolveTCPAddr("tcp", "10.0.0.5:45738")
+	if err != nil {
+		t.Fatalf("err: %v", err)
+	}
+
+	policy, err := p(upstream)
+	if err != nil {
+		t.Fatalf("err: %v", err)
+	}
+
+	if policy != IGNORE {
+		t.Fatalf("Expected policy IGNORE, got %v", policy)
+	}
+}
+
+func TestWhitelistPolicyReturnsUseWhenUpstreamIpAddrInWhitelist(t *testing.T) {
+	var cases = []struct {
+		name   string
+		policy PolicyFunc
+	}{
+		{"strict whitelist policy", MustStrictWhiteListPolicy([]string{"10.0.0.2", "10.0.0.3", "10.0.0.4"})},
+		{"lax whitelist policy", MustLaxWhiteListPolicy([]string{"10.0.0.2", "10.0.0.3", "10.0.0.4"})},
+	}
+
+	upstream, err := net.ResolveTCPAddr("tcp", "10.0.0.3:45738")
+	if err != nil {
+		t.Fatalf("err: %v", err)
+	}
+
+	for _, tc := range cases {
+		t.Run(tc.name, func(t *testing.T) {
+			policy, err := tc.policy(upstream)
+			if err != nil {
+				t.Fatalf("err: %v", err)
+			}
+
+			if policy != USE {
+				t.Fatalf("Expected policy USE, got %v", policy)
+			}
+		})
+	}
+}
+
+func TestWhitelistPolicyReturnsUseWhenUpstreamIpAddrInWhitelistRange(t *testing.T) {
+	var cases = []struct {
+		name   string
+		policy PolicyFunc
+	}{
+		{"strict whitelist policy", MustStrictWhiteListPolicy([]string{"10.0.0.0/29"})},
+		{"lax whitelist policy", MustLaxWhiteListPolicy([]string{"10.0.0.0/29"})},
+	}
+
+	upstream, err := net.ResolveTCPAddr("tcp", "10.0.0.3:45738")
+	if err != nil {
+		t.Fatalf("err: %v", err)
+	}
+
+	for _, tc := range cases {
+		t.Run(tc.name, func(t *testing.T) {
+			policy, err := tc.policy(upstream)
+			if err != nil {
+				t.Fatalf("err: %v", err)
+			}
+
+			if policy != USE {
+				t.Fatalf("Expected policy USE, got %v", policy)
+			}
+		})
+	}
+}
+
+func Test_CreateWhitelistPolicyWithInvalidCidrReturnsError(t *testing.T) {
+	_, err := StrictWhiteListPolicy([]string{"20/80"})
+	if err == nil {
+		t.Error("Expected error, got none")
+	}
+}
+
+func Test_CreateWhitelistPolicyWithInvalidIpAddressReturnsError(t *testing.T) {
+	_, err := StrictWhiteListPolicy([]string{"855.222.233.11"})
+	if err == nil {
+		t.Error("Expected error, got none")
+	}
+}
+
+func Test_CreateLaxPolicyWithInvalidCidrReturnsError(t *testing.T) {
+	_, err := LaxWhiteListPolicy([]string{"20/80"})
+	if err == nil {
+		t.Error("Expected error, got none")
+	}
+}
+
+func Test_CreateLaxPolicyWithInvalidIpAddresseturnsError(t *testing.T) {
+	_, err := LaxWhiteListPolicy([]string{"855.222.233.11"})
+	if err == nil {
+		t.Error("Expected error, got none")
+	}
+}
+
+func Test_MustLaxWhiteListPolicyPanicsWithInvalidIpAddress(t *testing.T) {
+	defer func() {
+		if r := recover(); r == nil {
+			t.Error("Expected a panic, but got none")
+		}
+	}()
+
+	MustLaxWhiteListPolicy([]string{"855.222.233.11"})
+}
+
+func Test_MustLaxWhiteListPolicyPanicsWithInvalidIpRange(t *testing.T) {
+	defer func() {
+		if r := recover(); r == nil {
+			t.Error("Expected a panic, but got none")
+		}
+	}()
+
+	MustLaxWhiteListPolicy([]string{"20/80"})
+}
+
+func Test_MustStrictWhiteListPolicyPanicsWithInvalidIpAddress(t *testing.T) {
+	defer func() {
+		if r := recover(); r == nil {
+			t.Error("Expected a panic, but got none")
+		}
+	}()
+
+	MustStrictWhiteListPolicy([]string{"855.222.233.11"})
+}
+
+func Test_MustStrictWhiteListPolicyPanicsWithInvalidIpRange(t *testing.T) {
+	defer func() {
+		if r := recover(); r == nil {
+			t.Error("Expected a panic, but got none")
+		}
+	}()
+
+	MustStrictWhiteListPolicy([]string{"20/80"})
+}
diff --git a/protocol_test.go b/protocol_test.go
index 01273b3..679ed2b 100644
--- a/protocol_test.go
+++ b/protocol_test.go
@@ -6,6 +6,7 @@ package proxyproto
 
 import (
 	"bytes"
+	"fmt"
 	"net"
 	"testing"
 )
@@ -189,3 +190,306 @@ func TestParse_ipv6(t *testing.T) {
 		t.Fatalf("bad: %v", addr)
 	}
 }
+
+func TestAcceptReturnsErrorWhenPolicyFuncErrors(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")
+	policyFunc := func(upstream net.Addr) (Policy, error) { return USE, expectedErr }
+
+	pl := &Listener{Listener: l, Policy: policyFunc}
+
+	go func() {
+		conn, err := net.Dial("tcp", pl.Addr().String())
+		if err != nil {
+			t.Fatalf("err: %v", err)
+		}
+		defer conn.Close()
+	}()
+
+	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)
+	}
+}
+
+func TestReadingIsRefusedWhenProxyHeaderRequiredButMissing(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 REQUIRE, nil }
+
+	pl := &Listener{Listener: l, Policy: policyFunc}
+
+	go func() {
+		conn, err := net.Dial("tcp", pl.Addr().String())
+		if err != nil {
+			t.Fatalf("err: %v", err)
+		}
+		defer conn.Close()
+		conn.Write([]byte("ping"))
+	}()
+
+	conn, err := pl.Accept()
+	if err != nil {
+		t.Fatalf("err: %v", err)
+	}
+	defer conn.Close()
+
+	recv := make([]byte, 4)
+	_, err = conn.Read(recv)
+	if err != ErrNoProxyProtocol {
+		t.Fatalf("Expected error %v, received %v", ErrNoProxyProtocol, err)
+	}
+}
+
+func TestReadingIsRefusedWhenProxyHeaderPresentButNotAllowed(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 REJECT, nil }
+
+	pl := &Listener{Listener: l, Policy: policyFunc}
+
+	go func() {
+		conn, err := net.Dial("tcp", pl.Addr().String())
+		if err != nil {
+			t.Fatalf("err: %v", err)
+		}
+		defer conn.Close()
+		header := &Header{
+			Version:            2,
+			Command:            PROXY,
+			TransportProtocol:  TCPv4,
+			SourceAddress:      net.ParseIP("10.1.1.1"),
+			SourcePort:         1000,
+			DestinationAddress: net.ParseIP("20.2.2.2"),
+			DestinationPort:    2000,
+		}
+		header.WriteTo(conn)
+	}()
+
+	conn, err := pl.Accept()
+	if err != nil {
+		t.Fatalf("err: %v", err)
+	}
+	defer conn.Close()
+
+	recv := make([]byte, 4)
+	_, err = conn.Read(recv)
+	if err != ErrSuperfluousProxyHeader {
+		t.Fatalf("Expected error %v, received %v", ErrSuperfluousProxyHeader, err)
+	}
+}
+func TestIgnorePolicyIgnoresIpFromProxyHeader(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 IGNORE, nil }
+
+	pl := &Listener{Listener: l, Policy: policyFunc}
+
+	go func() {
+		conn, err := net.Dial("tcp", pl.Addr().String())
+		if err != nil {
+			t.Fatalf("err: %v", err)
+		}
+		defer conn.Close()
+
+		// Write out the header!
+		header := &Header{
+			Version:            2,
+			Command:            PROXY,
+			TransportProtocol:  TCPv4,
+			SourceAddress:      net.ParseIP("10.1.1.1"),
+			SourcePort:         1000,
+			DestinationAddress: net.ParseIP("20.2.2.2"),
+			DestinationPort:    2000,
+		}
+		header.WriteTo(conn)
+
+		conn.Write([]byte("ping"))
+		recv := make([]byte, 4)
+		_, err = conn.Read(recv)
+		if err != nil {
+			t.Fatalf("err: %v", err)
+		}
+		if !bytes.Equal(recv, []byte("pong")) {
+			t.Fatalf("bad: %v", recv)
+		}
+	}()
+
+	conn, err := pl.Accept()
+	if err != nil {
+		t.Fatalf("err: %v", err)
+	}
+	defer conn.Close()
+
+	recv := make([]byte, 4)
+	_, err = conn.Read(recv)
+	if err != nil {
+		t.Fatalf("err: %v", err)
+	}
+	if !bytes.Equal(recv, []byte("ping")) {
+		t.Fatalf("bad: %v", recv)
+	}
+
+	if _, err := conn.Write([]byte("pong")); err != nil {
+		t.Fatalf("err: %v", err)
+	}
+
+	// Check the remote addr
+	addr := conn.RemoteAddr().(*net.TCPAddr)
+	if addr.IP.String() != "127.0.0.1" {
+		t.Fatalf("bad: %v", addr)
+	}
+}
+
+func Test_AllOptionsAreRecognized(t *testing.T) {
+	recognizedOpt1 := false
+	opt1 := func(c *Conn) {
+		recognizedOpt1 = true
+	}
+
+	recognizedOpt2 := false
+	opt2 := func(c *Conn) {
+		recognizedOpt2 = true
+	}
+
+	server, client := net.Pipe()
+	defer func() {
+		client.Close()
+	}()
+
+	c := NewConn(server, opt1, opt2)
+	if !recognizedOpt1 {
+		t.Error("Expected option 1 recognized")
+	}
+
+	if !recognizedOpt2 {
+		t.Error("Expected option 2 recognized")
+	}
+
+	c.Close()
+}
+
+func TestReadingIsRefusedOnErrorWhenRemoteAddrRequestedFirst(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 REQUIRE, nil }
+
+	pl := &Listener{Listener: l, Policy: policyFunc}
+
+	go func() {
+		conn, err := net.Dial("tcp", pl.Addr().String())
+		if err != nil {
+			t.Fatalf("err: %v", err)
+		}
+		defer conn.Close()
+		conn.Write([]byte("ping"))
+	}()
+
+	conn, err := pl.Accept()
+	if err != nil {
+		t.Fatalf("err: %v", err)
+	}
+	defer conn.Close()
+
+	_ = conn.RemoteAddr()
+	recv := make([]byte, 4)
+	_, err = conn.Read(recv)
+	if err != ErrNoProxyProtocol {
+		t.Fatalf("Expected error %v, received %v", ErrNoProxyProtocol, err)
+	}
+}
+
+func TestReadingIsRefusedOnErrorWhenLocalAddrRequestedFirst(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 REQUIRE, nil }
+
+	pl := &Listener{Listener: l, Policy: policyFunc}
+
+	go func() {
+		conn, err := net.Dial("tcp", pl.Addr().String())
+		if err != nil {
+			t.Fatalf("err: %v", err)
+		}
+		defer conn.Close()
+		conn.Write([]byte("ping"))
+	}()
+
+	conn, err := pl.Accept()
+	if err != nil {
+		t.Fatalf("err: %v", err)
+	}
+	defer conn.Close()
+
+	_ = conn.LocalAddr()
+	recv := make([]byte, 4)
+	_, err = conn.Read(recv)
+	if err != ErrNoProxyProtocol {
+		t.Fatalf("Expected error %v, received %v", ErrNoProxyProtocol, err)
+	}
+}
+
+func Test_ConnectionErrorsWhenHeaderValidationFails(t *testing.T) {
+	l, err := net.Listen("tcp", "127.0.0.1:0")
+	if err != nil {
+		t.Fatalf("err: %v", err)
+	}
+
+	validationError := fmt.Errorf("failed to validate")
+	pl := &Listener{Listener: l, ValidateHeader: func(*Header) error { return validationError }}
+
+	go func() {
+		conn, err := net.Dial("tcp", pl.Addr().String())
+		if err != nil {
+			t.Fatalf("err: %v", err)
+		}
+		defer conn.Close()
+
+		// Write out the header!
+		header := &Header{
+			Version:            2,
+			Command:            PROXY,
+			TransportProtocol:  TCPv4,
+			SourceAddress:      net.ParseIP("10.1.1.1"),
+			SourcePort:         1000,
+			DestinationAddress: net.ParseIP("20.2.2.2"),
+			DestinationPort:    2000,
+		}
+		header.WriteTo(conn)
+	}()
+
+	conn, err := pl.Accept()
+	if err != nil {
+		t.Fatalf("err: %v", err)
+	}
+	defer conn.Close()
+
+	recv := make([]byte, 4)
+	_, err = conn.Read(recv)
+	if err != validationError {
+		t.Fatalf("expected validation error, got %v", err)
+	}
+}