Repository: pires/go-proxyproto (go)

Issue:
ReadHeaderTimeout is setting a hard timeout on connections regardless of header sending
I think this setting is supposed to timeout connections that are not sending the headers and terminate them after this amount of time. However, I set the setting to 5 seconds, and the result is that all connections are terminated after 5 seconds, even if they correctly sent the headers. The connection starts up, but is then abruptly ended after the 5 seconds are passed.

I think what is missing in #74 is a call to `SetReadDeadline` to reset the timeout after the proxy header was sent successfully.


Pull request description:
Reset read deadline after parsing header
Can you, please, help review @unmarshal, as it is related to your PR #74?

Fixes #75

Patch:
diff --git a/protocol.go b/protocol.go
index 0f493ba..bfc24e9 100644
--- a/protocol.go
+++ b/protocol.go
@@ -8,10 +8,19 @@ import (
 	"time"
 )
 
+// DefaultReadHeaderTimeout is how long header processing waits for header to
+// be read from the wire, if Listener.ReaderHeaderTimeout is not set.
+// It's kept as a global variable so to make it easier to find and override,
+// e.g. go build -ldflags -X "github.com/pires/go-proxyproto.DefaultReadHeaderTimeout=1s"
+var DefaultReadHeaderTimeout = 200 * time.Millisecond
+
 // Listener is used to wrap an underlying listener,
 // whose connections may be using the HAProxy Proxy Protocol.
 // If the connection is using the protocol, the RemoteAddr() will return
-// the correct client address.
+// the correct client address. ReadHeaderTimeout will be applied to all
+// 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.
 type Listener struct {
 	Listener          net.Listener
 	Policy            PolicyFunc
@@ -21,7 +30,8 @@ type Listener struct {
 
 // 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.
+// return the address of the client instead of the proxy address. Each connection
+// will have its own readHeaderTimeout and readDeadline set by the Accept() call.
 type Conn struct {
 	bufReader         *bufio.Reader
 	conn              net.Conn
@@ -30,6 +40,8 @@ type Conn struct {
 	ProxyHeaderPolicy Policy
 	Validate          Validator
 	readErr           error
+	readHeaderTimeout time.Duration
+	readDeadline      time.Time
 }
 
 // Validator receives a header and decides whether it is a valid one
@@ -53,10 +65,6 @@ func (p *Listener) Accept() (net.Conn, error) {
 		return nil, err
 	}
 
-	if d := p.ReadHeaderTimeout; d != 0 {
-		conn.SetReadDeadline(time.Now().Add(d))
-	}
-
 	proxyHeaderPolicy := USE
 	if p.Policy != nil {
 		proxyHeaderPolicy, err = p.Policy(conn.RemoteAddr())
@@ -72,6 +80,15 @@ func (p *Listener) Accept() (net.Conn, error) {
 		WithPolicy(proxyHeaderPolicy),
 		ValidateHeader(p.ValidateHeader),
 	)
+
+	// If the ReadHeaderTimeout for the listener is unset, use the default timeout.
+	if p.ReadHeaderTimeout == 0 {
+		p.ReadHeaderTimeout = DefaultReadHeaderTimeout
+	}
+
+	// Set the readHeaderTimeout of the new conn to the value of the listener
+	newConn.readHeaderTimeout = p.ReadHeaderTimeout
+
 	return newConn, nil
 }
 
@@ -110,6 +127,7 @@ func (p *Conn) Read(b []byte) (int, error) {
 	if p.readErr != nil {
 		return 0, p.readErr
 	}
+
 	return p.bufReader.Read(b)
 }
 
@@ -197,11 +215,16 @@ func (p *Conn) UDPConn() (conn *net.UDPConn, ok bool) {
 
 // SetDeadline wraps original conn.SetDeadline
 func (p *Conn) SetDeadline(t time.Time) error {
+	p.readDeadline = t
 	return p.conn.SetDeadline(t)
 }
 
 // SetReadDeadline wraps original conn.SetReadDeadline
 func (p *Conn) SetReadDeadline(t time.Time) error {
+	// Set a local var that tells us the desired deadline. This is
+	// needed in order to reset the read deadline to the one that is
+	// desired by the user, rather than an empty deadline.
+	p.readDeadline = t
 	return p.conn.SetReadDeadline(t)
 }
 
@@ -211,7 +234,28 @@ func (p *Conn) SetWriteDeadline(t time.Time) error {
 }
 
 func (p *Conn) readHeader() error {
+	// If the connection's readHeaderTimeout is more than 0,
+	// push our deadline back to now plus the timeout. This should only
+	// run on the connection, as we don't want to override the previous
+	// read deadline the user may have used.
+	if p.readHeaderTimeout > 0 {
+		p.conn.SetReadDeadline(time.Now().Add(p.readHeaderTimeout))
+	}
+
 	header, err := Read(p.bufReader)
+
+	// If the connection's readHeaderTimeout is more than 0, undo the change to the
+	// deadline that we made above. Because we retain the readDeadline as part of our
+	// SetReadDeadline override, we know the user's desired deadline so we use that.
+	// Therefore, we check whether the error is a net.Timeout and if it is, we decide
+	// the proxy proto does not exist and set the error accordingly.
+	if p.readHeaderTimeout > 0 {
+		p.conn.SetReadDeadline(p.readDeadline)
+		if netErr, ok := err.(net.Error); ok && netErr.Timeout() {
+			err = ErrNoProxyProtocol
+		}
+	}
+
 	// 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 {


Test patch:
diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml
index 8af161a..122a182 100644
--- a/.github/workflows/test.yml
+++ b/.github/workflows/test.yml
@@ -10,7 +10,7 @@ jobs:
     runs-on: ubuntu-latest
     strategy:
       matrix:
-        go: [ '1.15', '1.14' ]
+        go: [ '1.15' ]
     steps:
     - uses: actions/checkout@v2
     - name: Set up Go
diff --git a/protocol_test.go b/protocol_test.go
index 3f06815..4d1d8bb 100644
--- a/protocol_test.go
+++ b/protocol_test.go
@@ -9,6 +9,7 @@ import (
 	"context"
 	"crypto/tls"
 	"crypto/x509"
+	"errors"
 	"fmt"
 	"io"
 	"io/ioutil"
@@ -63,7 +64,117 @@ func TestPassthrough(t *testing.T) {
 	}
 }
 
-func TestReadHeaderTimeout(t *testing.T) {
+// TestRequiredWithReadHeaderTimeout will iterate through 3 different timeouts to see
+// whether using a REQUIRE policy for a listener would cause an error if the timeout
+// is triggerred without a proxy protocol header being defined.
+func TestRequiredWithReadHeaderTimeout(t *testing.T) {
+	for _, duration := range []int{100, 200, 400} {
+		t.Run(fmt.Sprint(duration), func(t *testing.T) {
+			start := time.Now()
+
+			l, err := net.Listen("tcp", "127.0.0.1:0")
+
+			if err != nil {
+				t.Fatalf("err: %v", err)
+			}
+
+			pl := &Listener{
+				Listener:          l,
+				ReadHeaderTimeout: time.Millisecond * time.Duration(duration),
+				Policy: func(upstream net.Addr) (Policy, error) {
+					return REQUIRE, nil
+				},
+			}
+
+			ctx, cancel := context.WithCancel(context.Background())
+			defer cancel()
+
+			go func() {
+				conn, err := net.Dial("tcp", pl.Addr().String())
+				if err != nil {
+					t.Fatalf("err: %v", err)
+				}
+				defer conn.Close()
+
+				<-ctx.Done()
+			}()
+
+			conn, err := pl.Accept()
+			if err != nil {
+				t.Fatalf("err: %v", err)
+			}
+			defer conn.Close()
+
+			// Read blocks forever if there is no ReadHeaderTimeout and the policy is not REQUIRE
+			recv := make([]byte, 4)
+			_, err = conn.Read(recv)
+
+			if err != nil && !errors.Is(err, ErrNoProxyProtocol) && time.Since(start)-pl.ReadHeaderTimeout > 10*time.Millisecond {
+				t.Fatal("proxy proto should not be found and time should be close to read timeout")
+			}
+		})
+	}
+}
+
+// TestUseWithReadHeaderTimeout will iterate through 3 different timeouts to see
+// whether using a USE policy for a listener would not cause an error if the timeout
+// is triggerred without a proxy protocol header being defined.
+func TestUseWithReadHeaderTimeout(t *testing.T) {
+	for _, duration := range []int{100, 200, 400} {
+		t.Run(fmt.Sprint(duration), func(t *testing.T) {
+			start := time.Now()
+
+			l, err := net.Listen("tcp", "127.0.0.1:0")
+
+			if err != nil {
+				t.Fatalf("err: %v", err)
+			}
+
+			pl := &Listener{
+				Listener:          l,
+				ReadHeaderTimeout: time.Millisecond * time.Duration(duration),
+				Policy: func(upstream net.Addr) (Policy, error) {
+					return USE, nil
+				},
+			}
+
+			ctx, cancel := context.WithCancel(context.Background())
+			defer cancel()
+
+			go func() {
+				conn, err := net.Dial("tcp", pl.Addr().String())
+				if err != nil {
+					t.Fatalf("err: %v", err)
+				}
+				defer conn.Close()
+
+				<-ctx.Done()
+			}()
+
+			conn, err := pl.Accept()
+			if err != nil {
+				t.Fatalf("err: %v", err)
+			}
+			defer conn.Close()
+
+			// 2 times the ReadHeaderTimeout because the first timeout
+			// should occur (the one set on the listener) and allow for the second to follow up
+			conn.SetDeadline(time.Now().Add(pl.ReadHeaderTimeout * 2))
+
+			// Read blocks forever if there is no ReadHeaderTimeout
+			recv := make([]byte, 4)
+			_, err = conn.Read(recv)
+
+			if err != nil && !errors.Is(err, ErrNoProxyProtocol) && (time.Since(start)-(pl.ReadHeaderTimeout*2)) > 10*time.Millisecond {
+				t.Fatal("proxy proto should not be found and time should be close to read timeout")
+			}
+		})
+	}
+}
+
+func TestReadHeaderTimeoutIsReset(t *testing.T) {
+	const timeout = time.Millisecond * 250
+
 	l, err := net.Listen("tcp", "127.0.0.1:0")
 	if err != nil {
 		t.Fatalf("err: %v", err)
@@ -71,12 +182,111 @@ func TestReadHeaderTimeout(t *testing.T) {
 
 	pl := &Listener{
 		Listener:          l,
-		ReadHeaderTimeout: 1 * time.Millisecond,
+		ReadHeaderTimeout: timeout,
+	}
+
+	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 {
+			t.Fatalf("err: %v", err)
+		}
+		defer conn.Close()
+
+		// Write out the header!
+		header.WriteTo(conn)
+
+		// Sleep here longer than the configured timeout.
+		time.Sleep(timeout * 2)
+
+		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()
+
+	// Set our deadlines higher than our ReadHeaderTimeout
+	conn.SetReadDeadline(time.Now().Add(timeout * 3))
+	conn.SetWriteDeadline(time.Now().Add(timeout * 3))
+
+	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() != "10.1.1.1" {
+		t.Fatalf("bad: %v", addr)
+	}
+	if addr.Port != 1000 {
+		t.Fatalf("bad: %v", addr)
 	}
 
-	ctx, cancel := context.WithCancel(context.Background())
-	defer cancel()
+	h := conn.(*Conn).ProxyHeader()
+	if !h.EqualsTo(header) {
+		t.Errorf("bad: %v", h)
+	}
+}
 
+// TestReadHeaderTimeoutIsEmpty ensures the default is set if it is empty.
+// Because the default is 200ms and we wait longer than that to send a message,
+// we expect the actual address and port to be returned,
+// rather than the ProxyHeader we defined.
+func TestReadHeaderTimeoutIsEmpty(t *testing.T) {
+	l, err := net.Listen("tcp", "127.0.0.1:0")
+	if err != nil {
+		t.Fatalf("err: %v", err)
+	}
+
+	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 {
@@ -84,7 +294,13 @@ func TestReadHeaderTimeout(t *testing.T) {
 		}
 		defer conn.Close()
 
-		<-ctx.Done()
+		// Sleep here longer than the configured timeout.
+		time.Sleep(250 * time.Millisecond)
+
+		// Write out the header!
+		header.WriteTo(conn)
+
+		conn.Write([]byte("ping"))
 	}()
 
 	conn, err := pl.Accept()
@@ -93,10 +309,85 @@ func TestReadHeaderTimeout(t *testing.T) {
 	}
 	defer conn.Close()
 
-	// Read blocks forever if there is no ReadHeaderTimeout
 	recv := make([]byte, 4)
 	_, err = conn.Read(recv)
+	if err != nil {
+		t.Fatalf("err: %v", err)
+	}
 
+	// Check the remote addr
+	addr := conn.RemoteAddr().(*net.TCPAddr)
+	if addr.IP.String() == "10.1.1.1" {
+		t.Fatalf("bad: %v", addr)
+	}
+	if addr.Port == 1000 {
+		t.Fatalf("bad: %v", addr)
+	}
+}
+
+// TestReadHeaderTimeoutIsNegative does the same as above except
+// with a negative timeout. Therefore, we expect the right ProxyHeader
+// to be returned.
+func TestReadHeaderTimeoutIsNegative(t *testing.T) {
+	l, err := net.Listen("tcp", "127.0.0.1:0")
+	if err != nil {
+		t.Fatalf("err: %v", err)
+	}
+
+	pl := &Listener{
+		Listener:          l,
+		ReadHeaderTimeout: -1,
+	}
+
+	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 {
+			t.Fatalf("err: %v", err)
+		}
+		defer conn.Close()
+
+		// Sleep here longer than the configured timeout.
+		time.Sleep(250 * time.Millisecond)
+
+		// Write out the header!
+		header.WriteTo(conn)
+
+		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 != nil {
+		t.Fatalf("err: %v", err)
+	}
+
+	// Check the remote addr
+	addr := conn.RemoteAddr().(*net.TCPAddr)
+	if addr.IP.String() != "10.1.1.1" {
+		t.Fatalf("bad: %v", addr)
+	}
+	if addr.Port != 1000 {
+		t.Fatalf("bad: %v", addr)
+	}
 }
 
 func TestParse_ipv4(t *testing.T) {
diff --git a/tlv_test.go b/tlv_test.go
index 1f7285c..6373cb2 100644
--- a/tlv_test.go
+++ b/tlv_test.go
@@ -13,30 +13,6 @@ var (
 	fixturePartialLenTLV = []byte{byte(PP2_TYPE_MIN_CUSTOM) + 3, 0x00, 0x02, 0x00}
 )
 
-func checkTLVs(t *testing.T, name string, raw []byte, expected []PP2Type) []TLV {
-	header, err := parseVersion2(bufio.NewReader(bytes.NewReader(raw)))
-	if err != nil {
-		t.Fatalf("%s: Unexpected error reading header %#v", name, err)
-	}
-
-	tlvs, err := header.TLVs()
-	if err != nil {
-		t.Fatalf("%s: Unexpected error splitting TLVS %#v", name, err)
-	}
-
-	if len(tlvs) != len(expected) {
-		t.Fatalf("%s: Expected %d TLVs, actual %d", name, len(expected), len(tlvs))
-	}
-
-	for i, et := range expected {
-		if at := tlvs[i].Type; at != et {
-			t.Fatalf("%s: Expected type %X, actual %X", name, et, at)
-		}
-	}
-
-	return tlvs
-}
-
 var invalidTLVTests = []struct {
 	name          string
 	reader        *bufio.Reader