Update

vendor/github.com/pires/go-proxyproto/protocol.go

@@ -11,15 +11,30 @@ import (
 	"time"
 )
 
+// readBufferSize is the size used for bufio.Reader's internal buffer.
+//
+// This is kept low to reduce per-connection memory overhead. If the header is
+// larger than readBufferSize, the header will be decoded with multiple Read
+// calls. 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. We use 256 bytes to
+// accommodate for the most common cases.
+const readBufferSize = 256
+
 var (
 	// 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"
+	// e.g. go build -ldflags -X "github.com/pires/go-proxyproto.DefaultReadHeaderTimeout=1s".
 	DefaultReadHeaderTimeout = 10 * time.Second
 
-	// ErrInvalidUpstream should be returned when an upstream connection address
-	// is not trusted, and therefore is invalid.
+	// ErrInvalidUpstream should be returned (possibly wrapped) by a policy
+	// function when an upstream connection address is not trusted or cannot be
+	// classified. Listener.Accept closes that connection and keeps listening;
+	// a policy error that does not wrap ErrInvalidUpstream is returned by
+	// Accept itself, which typically stops the caller's accept loop. All
+	// built-in policies wrap address-classification failures in
+	// ErrInvalidUpstream so a single unclassifiable peer cannot stop the
+	// listener.
 	ErrInvalidUpstream = fmt.Errorf("proxyproto: upstream connection address not trusted for PROXY information")
 )
 
@@ -31,15 +46,41 @@ var (
 // is set, a default of 10s will be used. This can be disabled by setting the
 // timeout to < 0.
 //
+// When neither Policy nor ConnPolicy is set, DefaultPolicy applies: REQUIRE,
+// so connections that do not open with a PROXY header fail their first
+// Read/Write with ErrNoProxyProtocol. Headers are still honored from ANY peer
+// under REQUIRE; a listener reachable by untrusted clients should set a
+// trusted-source policy (e.g. TrustProxyHeaderFrom or TrustProxyHeaderFromRanges).
+//
+// Listener is stream-oriented (TCP, Unix stream): the header is read once at
+// the start of the byte stream. It cannot implement the PROXY protocol over
+// UDP, where the spec requires a header parsed independently in each datagram;
+// use ParseUDPDatagram and Header.FormatUDPDatagram for that.
+//
+// Note that ReadHeaderTimeout only bounds how long a single slow connection can
+// hold a goroutine and file descriptor during header detection; it is not a
+// connection count or accept-rate limit. Deployments exposed to untrusted
+// clients should keep ReadHeaderTimeout low and enforce connection/rate limits
+// upstream (or around Accept).
+//
 // Only one of Policy or ConnPolicy should be provided. If both are provided then
 // a panic would occur during accept.
 type Listener struct {
+	// Listener is the underlying listener.
 	Listener net.Listener
 	// Deprecated: use ConnPolicyFunc instead. This will be removed in future release.
-	Policy            PolicyFunc
-	ConnPolicy        ConnPolicyFunc
-	ValidateHeader    Validator
+	Policy PolicyFunc
+	// ConnPolicy is the policy function for accepted connections.
+	ConnPolicy ConnPolicyFunc
+	// ValidateHeader is the validator function for the proxy header.
+	ValidateHeader Validator
+	// ReadHeaderTimeout is the timeout for reading the proxy header.
 	ReadHeaderTimeout time.Duration
+	// ReadBufferSize is the read buffer size for accepted connections. When > 0,
+	// each accepted connection uses this size for proxy header detection; 0 means default.
+	// See the sizing note on WithBufferSize: values below 107 bytes break v1
+	// header parsing.
+	ReadBufferSize int
 }
 
 // Conn is used to wrap and underlying connection which
@@ -47,12 +88,13 @@ type Listener struct {
 // 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 {
-	readDeadline      atomic.Value // time.Time
-	once              sync.Once
-	readErr           error
-	conn              net.Conn
-	bufReader         *bufio.Reader
-	reader            io.Reader
+	readDeadline atomic.Value // time.Time
+	once         sync.Once
+	readErr      error
+	conn         net.Conn
+	bufReader    *bufio.Reader
+	// bufferSize is set when the client overrides via WithBufferSize; nil means use default.
+	bufferSize        *int
 	header            *Header
 	ProxyHeaderPolicy Policy
 	Validate          Validator
@@ -63,7 +105,7 @@ type Conn struct {
 // 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()
+// 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 {
@@ -72,7 +114,9 @@ func ValidateHeader(v Validator) func(*Conn) {
 	}
 }
 
-// SetReadHeaderTimeout sets the readHeaderTimeout for a connection when passed as option to NewConn()
+// SetReadHeaderTimeout sets the readHeaderTimeout for a connection when passed as option to NewConn().
+// A value of 0 disables the header read timeout; negative values are ignored,
+// leaving the connection's current timeout (the NewConn default) in place.
 func SetReadHeaderTimeout(t time.Duration) func(*Conn) {
 	return func(c *Conn) {
 		if t >= 0 {
@@ -81,16 +125,38 @@ func SetReadHeaderTimeout(t time.Duration) func(*Conn) {
 	}
 }
 
+// WithBufferSize sets the size of the read buffer used for proxy header detection.
+// Values <= 0 are ignored and the default (256 bytes) is used. Values < 16 are
+// effectively 16 due to bufio's minimum. The default is tuned for typical proxy
+// protocol header lengths.
+//
+// The buffer must be able to hold an entire v1 header line (up to 107 bytes):
+// v1 parsing requires the line to be available without refilling the buffer
+// (the slow-loris defense), so a smaller buffer rejects every v1 connection
+// with ErrCantReadVersion1Header even from well-behaved senders. v2 parsing
+// refills freely and works with any size.
+func WithBufferSize(length int) func(*Conn) {
+	return func(c *Conn) {
+		if length <= 0 {
+			return
+		}
+		p := new(int)
+		*p = length
+		c.bufferSize = p
+		c.bufReader = bufio.NewReaderSize(c.conn, length)
+	}
+}
+
 // Accept waits for and returns the next valid connection to the listener.
 func (p *Listener) Accept() (net.Conn, error) {
 	for {
-		// Get the underlying connection
+		// Get the underlying connection.
 		conn, err := p.Listener.Accept()
 		if err != nil {
 			return nil, err
 		}
 
-		proxyHeaderPolicy := USE
+		proxyHeaderPolicy := DefaultPolicy
 		if p.Policy != nil && p.ConnPolicy != nil {
 			panic("only one of policy or connpolicy must be provided.")
 		}
@@ -104,35 +170,42 @@ func (p *Listener) Accept() (net.Conn, error) {
 				})
 			}
 			if err != nil {
-				// can't decide the policy, we can't accept the connection
-				conn.Close()
+				// can't decide the policy, we can't accept the connection.
+				if closeErr := conn.Close(); closeErr != nil {
+					return nil, closeErr
+				}
 
 				if errors.Is(err, ErrInvalidUpstream) {
-					// keep listening for other connections
+					// keep listening for other connections.
 					continue
 				}
 
 				return nil, err
 			}
-			// Handle a connection as a regular one
+			// Handle a connection as a regular one.
 			if proxyHeaderPolicy == SKIP {
 				return conn, nil
 			}
 		}
 
-		newConn := NewConn(
-			conn,
+		opts := []func(*Conn){
 			WithPolicy(proxyHeaderPolicy),
 			ValidateHeader(p.ValidateHeader),
-		)
-
-		// If the ReadHeaderTimeout for the listener is unset, use the default timeout.
-		if p.ReadHeaderTimeout == 0 {
-			p.ReadHeaderTimeout = DefaultReadHeaderTimeout
 		}
+		if p.ReadBufferSize > 0 {
+			opts = append(opts, WithBufferSize(p.ReadBufferSize))
+		}
+		newConn := NewConn(conn, opts...)
 
-		// Set the readHeaderTimeout of the new conn to the value of the listener
-		newConn.readHeaderTimeout = p.ReadHeaderTimeout
+		// Set the readHeaderTimeout of the new conn to the value of the listener,
+		// falling back to the default when unset. Read into a local rather than
+		// writing back to the shared Listener: mutating p here races with
+		// concurrent Accept calls and would silently rewrite the caller's struct.
+		readHeaderTimeout := p.ReadHeaderTimeout
+		if readHeaderTimeout == 0 {
+			readHeaderTimeout = DefaultReadHeaderTimeout
+		}
+		newConn.readHeaderTimeout = readHeaderTimeout
 
 		return newConn, nil
 	}
@@ -148,19 +221,40 @@ 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.
+//
+// Conn is stream-oriented; see the note on Listener about the PROXY protocol
+// over UDP datagrams.
+//
+// By default the returned Conn applies DefaultPolicy (REQUIRE, so the peer
+// must open with a PROXY header; override with the WithPolicy option) and
+// DefaultReadHeaderTimeout (10s) while detecting the PROXY protocol header, so
+// a client that connects but never sends data cannot make header detection
+// block forever.
+//
+// The timeout bounds header detection only, not the first Read end-to-end:
+// under a non-REQUIRE policy, when no header is present Read falls through to
+// a normal read of the underlying connection, which can still block on a
+// silent client (pinning a goroutine and file descriptor). For an end-to-end
+// bound, set a read deadline on the connection, or keep the REQUIRE policy,
+// which makes the first Read fail when no header arrives within the timeout.
+//
+// Override the timeout with the SetReadHeaderTimeout option; pass
+// SetReadHeaderTimeout(0) to disable it entirely.
+//
+// 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.
-	// We use 256 bytes to be safe.
-	const bufSize = 256
-	br := bufio.NewReaderSize(conn, bufSize)
+	br := bufio.NewReaderSize(conn, readBufferSize)
 
 	pConn := &Conn{
-		bufReader: br,
-		reader:    io.MultiReader(br, conn),
-		conn:      conn,
+		bufReader:         br,
+		conn:              conn,
+		ProxyHeaderPolicy: DefaultPolicy,
+		readHeaderTimeout: DefaultReadHeaderTimeout,
 	}
 
 	for _, opt := range opts {
@@ -170,26 +264,53 @@ func NewConn(conn net.Conn, opts ...func(*Conn)) *Conn {
 	return pConn
 }
 
-// Read is check for the proxy protocol header when doing
-// the initial scan. If there is an error parsing the header,
-// it is returned and the socket is closed.
+// Read checks for the proxy protocol header on the first call, then reads
+// from the connection. If there is an error processing the header, it is
+// returned by this and every subsequent call. The connection is NOT closed by
+// this package; the caller should close it.
 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)
+	// Drain the buffer if it exists and has data.
+	if p.bufReader != nil {
+		if p.bufReader.Buffered() > 0 {
+			n, err := p.bufReader.Read(b)
+
+			// Did we empty the buffer?
+			// Buffering a net.Conn means the buffer doesn't return io.EOF until the connection returns io.EOF.
+			// Therefore, we use Buffered() == 0 to detect if we are done with the buffer.
+			if p.bufReader.Buffered() == 0 {
+				// Garbage collect the buffer.
+				p.bufReader = nil
+			}
+
+			// Return immediately. Do not touch p.conn.
+			// If err is EOF here, it means the connection is actually closed,
+			// so we should return that error to the user anyway.
+			return n, err
+		}
+		// If buffer was empty to begin with (shouldn't happen with the >0 check
+		// but good for safety), clear it.
+		p.bufReader = nil
+	}
+
+	// From now on, read directly from the underlying connection.
+	return p.conn.Read(b)
 }
 
-// Write wraps original conn.Write
+// 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)
 }
 
-// Close wraps original conn.Close
+// Close wraps original conn.Close.
 func (p *Conn) Close() error {
 	return p.conn.Close()
 }
@@ -197,7 +318,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
 }
 
@@ -208,8 +330,9 @@ 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() })
-	if p.header == nil || p.header.Command.IsLocal() || p.readErr != nil {
+	// Ensure header processing runs at most once.
+	_ = p.ensureHeaderProcessed()
+	if p.header == nil || p.header.Command.IsLocal() || p.readErr != nil || p.header.DestinationAddr == nil {
 		return p.conn.LocalAddr()
 	}
 
@@ -223,8 +346,9 @@ 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() })
-	if p.header == nil || p.header.Command.IsLocal() || p.readErr != nil {
+	// Ensure header processing runs at most once.
+	_ = p.ensureHeaderProcessed()
+	if p.header == nil || p.header.Command.IsLocal() || p.readErr != nil || p.header.SourceAddr == nil {
 		return p.conn.RemoteAddr()
 	}
 
@@ -266,13 +390,13 @@ func (p *Conn) UDPConn() (conn *net.UDPConn, ok bool) {
 	return
 }
 
-// SetDeadline wraps original conn.SetDeadline
+// SetDeadline wraps original conn.SetDeadline.
 func (p *Conn) SetDeadline(t time.Time) error {
 	p.readDeadline.Store(t)
 	return p.conn.SetDeadline(t)
 }
 
-// SetReadDeadline wraps original conn.SetReadDeadline
+// 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
@@ -281,18 +405,33 @@ func (p *Conn) SetReadDeadline(t time.Time) error {
 	return p.conn.SetReadDeadline(t)
 }
 
-// SetWriteDeadline wraps original conn.SetWriteDeadline
+// SetWriteDeadline wraps original conn.SetWriteDeadline.
 func (p *Conn) SetWriteDeadline(t time.Time) error {
 	return p.conn.SetWriteDeadline(t)
 }
 
+// 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
 		}
 	}
@@ -301,7 +440,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 {
@@ -309,11 +448,17 @@ func (p *Conn) readHeader() error {
 		if t == nil {
 			t = time.Time{}
 		}
-		if err := p.conn.SetReadDeadline(t.(time.Time)); err != nil {
-			return err
-		}
+		// Restore the user's deadline on a best-effort basis. This must not
+		// discard a header we already parsed: some connections (notably
+		// net.Pipe) return an error from SetReadDeadline once the peer has
+		// closed, which can happen right after the peer sends the header and
+		// closes. Only surface a restore failure when the read produced neither
+		// a header nor an error of its own.
+		restoreErr := p.conn.SetReadDeadline(t.(time.Time))
 		if netErr, ok := err.(net.Error); ok && netErr.Timeout() {
 			err = ErrNoProxyProtocol
+		} else if err == nil && header == nil {
+			err = restoreErr
 		}
 	}
 
@@ -349,24 +494,44 @@ func (p *Conn) readHeader() error {
 	return err
 }
 
-// ReadFrom implements the io.ReaderFrom ReadFrom method
+// 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)
 	}
 	return io.Copy(p.conn, r)
 }
 
-// WriteTo implements io.WriterTo
+// 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
+	}
+
+	// If the buffer has been drained (or cleared), copy directly from conn.
+	if p.bufReader == nil {
+		return io.Copy(w, p.conn)
 	}
 
 	b := make([]byte, p.bufReader.Buffered())
 	if _, err := p.bufReader.Read(b); err != nil {
-		return 0, err // this should never as we read buffered data
+		return 0, err // this should never happen as we read buffered data.
 	}
 
 	var n int64