protocol: readHeader() respects previously called conn.SetReadDeadline(t)

Author: pires

protocol.go

@@ -150,8 +150,13 @@ func (p *Listener) Addr() net.Addr {
 	return p.Listener.Addr()
 }
 
-// NewConn is used to wrap a net.Conn that may be speaking
-// the proxy protocol into a proxyproto.Conn.
+// NewConn is used to wrap a net.Conn that may be speaking the PROXY protocol
+// into a proxyproto.Conn.
+//
+// NOTE: NewConn may interfere with previously set ReadDeadline on the provided net.Conn,
+// because it sets a temporary deadline when detecting and reading the PROXY protocol header.
+// If you need to enforce a specific ReadDeadline on the connection, be sure to call Conn.SetReadDeadline
+// again after NewConn returns, to restore your desired deadline.
 func NewConn(conn net.Conn, opts ...func(*Conn)) *Conn {
 	// For v1 the header length is at most 108 bytes.
 	// For v2 the header length is at most 52 bytes plus the length of the TLVs.
@@ -176,18 +181,20 @@ func NewConn(conn net.Conn, opts ...func(*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) {
-	p.once.Do(func() {
-		p.readErr = p.readHeader()
-	})
-	if p.readErr != nil {
-		return 0, p.readErr
+	// Ensure header processing runs at most once and surface any errors.
+	if err := p.ensureHeaderProcessed(); err != nil {
+		return 0, err
 	}
 
 	return p.reader.Read(b)
 }
 
 // Write wraps original conn.Write.
 func (p *Conn) Write(b []byte) (int, error) {
+	// Ensure header processing has completed before writing.
+	if err := p.ensureHeaderProcessed(); err != nil {
+		return 0, err
+	}
 	return p.conn.Write(b)
 }
 
@@ -199,7 +206,8 @@ func (p *Conn) Close() error {
 // 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() })
+	// Ensure header processing runs at most once.
+	_ = p.ensureHeaderProcessed()
 	return p.header
 }
 
@@ -210,7 +218,8 @@ func (p *Conn) ProxyHeader() *Header {
 // from the proxy header even if the proxy header itself is
 // syntactically correct.
 func (p *Conn) LocalAddr() net.Addr {
-	p.once.Do(func() { p.readErr = p.readHeader() })
+	// Ensure header processing runs at most once.
+	_ = p.ensureHeaderProcessed()
 	if p.header == nil || p.header.Command.IsLocal() || p.readErr != nil {
 		return p.conn.LocalAddr()
 	}
@@ -225,7 +234,8 @@ func (p *Conn) LocalAddr() net.Addr {
 // from the proxy header even if the proxy header itself is
 // syntactically correct.
 func (p *Conn) RemoteAddr() net.Addr {
-	p.once.Do(func() { p.readErr = p.readHeader() })
+	// Ensure header processing runs at most once.
+	_ = p.ensureHeaderProcessed()
 	if p.header == nil || p.header.Command.IsLocal() || p.readErr != nil {
 		return p.conn.RemoteAddr()
 	}
@@ -291,11 +301,25 @@ func (p *Conn) SetWriteDeadline(t time.Time) error {
 // readHeader reads the proxy protocol header from the connection.
 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.
+	// apply a temporary deadline without extending a user-configured
+	// deadline. If the user has no deadline, we use now + timeout.
 	if p.readHeaderTimeout > 0 {
-		if err := p.conn.SetReadDeadline(time.Now().Add(p.readHeaderTimeout)); err != nil {
+		var (
+			storedDeadline time.Time
+			hasDeadline    bool
+		)
+		if t := p.readDeadline.Load(); t != nil {
+			storedDeadline = t.(time.Time)
+			hasDeadline = !storedDeadline.IsZero()
+		}
+
+		headerDeadline := time.Now().Add(p.readHeaderTimeout)
+		if hasDeadline && storedDeadline.Before(headerDeadline) {
+			// Clamp to the user's earlier deadline to avoid extending it.
+			headerDeadline = storedDeadline
+		}
+
+		if err := p.conn.SetReadDeadline(headerDeadline); err != nil {
 			return err
 		}
 	}
@@ -304,7 +328,7 @@ func (p *Conn) readHeader() error {
 
 	// 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.
+	// SetReadDeadline override, we can restore the user's deadline (if any).
 	// 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 {
@@ -352,8 +376,23 @@ func (p *Conn) readHeader() error {
 	return err
 }
 
+// ensureHeaderProcessed runs header processing once.
+func (p *Conn) ensureHeaderProcessed() error {
+	p.once.Do(func() {
+		p.readErr = p.readHeader()
+	})
+	if p.readErr != nil {
+		return p.readErr
+	}
+	return nil
+}
+
 // ReadFrom implements the io.ReaderFrom ReadFrom method.
 func (p *Conn) ReadFrom(r io.Reader) (int64, error) {
+	// Ensure header processing has completed before reading/writing.
+	if err := p.ensureHeaderProcessed(); err != nil {
+		return 0, err
+	}
 	if rf, ok := p.conn.(io.ReaderFrom); ok {
 		return rf.ReadFrom(r)
 	}
@@ -362,9 +401,9 @@ func (p *Conn) ReadFrom(r io.Reader) (int64, error) {
 
 // WriteTo implements io.WriterTo.
 func (p *Conn) WriteTo(w io.Writer) (int64, error) {
-	p.once.Do(func() { p.readErr = p.readHeader() })
-	if p.readErr != nil {
-		return 0, p.readErr
+	// Ensure header processing has completed before reading/writing.
+	if err := p.ensureHeaderProcessed(); err != nil {
+		return 0, err
 	}
 
 	b := make([]byte, p.bufReader.Buffered())

protocol_test.go

@@ -223,6 +223,117 @@ func TestUseWithReadHeaderTimeout(t *testing.T) {
 	}
 }
 
+func TestReadHeaderTimeoutRespectsEarlierDeadline(t *testing.T) {
+	const (
+		headerTimeout = 200 * time.Millisecond
+		userTimeout   = 60 * time.Millisecond
+		tolerance     = 100 * time.Millisecond
+	)
+
+	l, err := net.Listen("tcp", testLocalhostRandomPort)
+	if err != nil {
+		t.Fatalf("err: %v", err)
+	}
+
+	pl := &Listener{
+		Listener:          l,
+		ReadHeaderTimeout: headerTimeout,
+		Policy: func(_ net.Addr) (Policy, error) {
+			// Use REQUIRE so a timeout is surfaced as ErrNoProxyProtocol.
+			return REQUIRE, nil
+		},
+	}
+
+	type dialResult struct {
+		conn net.Conn
+		err  error
+	}
+
+	dialResultCh := make(chan dialResult, 1)
+	go func() {
+		conn, err := net.Dial("tcp", pl.Addr().String())
+		dialResultCh <- dialResult{conn: conn, err: err}
+	}()
+
+	conn, err := pl.Accept()
+	if err != nil {
+		t.Fatalf("err: %v", err)
+	}
+	t.Cleanup(func() {
+		if closeErr := conn.Close(); closeErr != nil {
+			t.Errorf("failed to close connection: %v", closeErr)
+		}
+	})
+
+	result := <-dialResultCh
+	if result.err != nil {
+		t.Fatalf("client error: %v", result.err)
+	}
+	t.Cleanup(func() {
+		if closeErr := result.conn.Close(); closeErr != nil {
+			t.Errorf("failed to close client connection: %v", closeErr)
+		}
+	})
+
+	// Set a shorter user deadline than the readHeaderTimeout and do not send data.
+	if err := conn.SetReadDeadline(time.Now().Add(userTimeout)); err != nil {
+		t.Fatalf("err: %v", err)
+	}
+
+	start := time.Now()
+	recv := make([]byte, 1)
+	_, err = conn.Read(recv)
+	elapsed := time.Since(start)
+
+	// The read should honor the earlier user deadline instead of waiting
+	// for the longer readHeaderTimeout.
+	if !errors.Is(err, ErrNoProxyProtocol) {
+		t.Fatalf("expected ErrNoProxyProtocol, got: %v", err)
+	}
+	if elapsed > userTimeout+tolerance {
+		t.Fatalf("read exceeded user deadline: elapsed=%v timeout=%v", elapsed, userTimeout)
+	}
+}
+
+func TestDeadlineSettersAfterHeaderProcessed(t *testing.T) {
+	conn, peer := net.Pipe()
+	t.Cleanup(func() {
+		if closeErr := conn.Close(); closeErr != nil {
+			t.Errorf("failed to close connection: %v", closeErr)
+		}
+	})
+	t.Cleanup(func() {
+		if closeErr := peer.Close(); closeErr != nil {
+			t.Errorf("failed to close peer connection: %v", closeErr)
+		}
+	})
+
+	proxyConn := NewConn(conn)
+
+	// Ensure header processing completes by sending a non-PROXY byte
+	// and reading it through the proxy connection.
+	go func() {
+		if _, err := peer.Write([]byte("x")); err != nil {
+			t.Errorf("failed to write peer data: %v", err)
+		}
+	}()
+	buf := make([]byte, 1)
+	if _, err := proxyConn.Read(buf); err != nil {
+		t.Fatalf("read failed: %v", err)
+	}
+
+	deadline := time.Now().Add(time.Second)
+	if err := proxyConn.SetDeadline(deadline); err != nil {
+		t.Fatalf("unexpected SetDeadline error: %v", err)
+	}
+	if err := proxyConn.SetReadDeadline(deadline); err != nil {
+		t.Fatalf("unexpected SetReadDeadline error: %v", err)
+	}
+	if err := proxyConn.SetWriteDeadline(deadline); err != nil {
+		t.Fatalf("unexpected SetWriteDeadline error: %v", err)
+	}
+}
+
 func TestReadHeaderTimeoutIsReset(t *testing.T) {
 	const timeout = time.Millisecond * 250