Detect and drop broken conns in the fastproxy pool

Co-authored-by: Romain <rtribotte@users.noreply.github.com>

pkg/proxy/fast/connpool.go

@@ -1,42 +1,309 @@
 package fast
 
 import (
+	"bufio"
+	"errors"
 	"fmt"
+	"io"
 	"net"
+	"net/http"
+	"net/http/httputil"
+	"strings"
+	"sync"
+	"sync/atomic"
 	"time"
 
 	"github.com/rs/zerolog/log"
+	"github.com/valyala/fasthttp"
 )
 
+// rwWithUpgrade contains a ResponseWriter and an upgradeHandler,
+// used to upgrade the connection (e.g. Websockets).
+type rwWithUpgrade struct {
+	RW      http.ResponseWriter
+	Upgrade upgradeHandler
+}
+
 // conn is an enriched net.Conn.
 type conn struct {
 	net.Conn
 
+	RWCh  chan rwWithUpgrade
+	ErrCh chan error
+
+	br *bufio.Reader
+
 	idleAt      time.Time // the last time it was marked as idle.
 	idleTimeout time.Duration
+
+	responseHeaderTimeout time.Duration
+
+	expectedResponse atomic.Bool
+	broken           atomic.Bool
+	upgraded         atomic.Bool
+
+	closeMu  sync.Mutex
+	closed   bool
+	closeErr error
+
+	bufferPool        *pool[[]byte]
+	limitedReaderPool *pool[*io.LimitedReader]
 }
 
-func (c *conn) isExpired() bool {
+// Read reads data from the connection.
+// Overrides conn Read to use the buffered reader.
+func (c *conn) Read(b []byte) (n int, err error) {
+	return c.br.Read(b)
+}
+
+// Close closes the connection.
+// Ensures that connection is closed only once,
+// to avoid duplicate close error.
+func (c *conn) Close() error {
+	c.closeMu.Lock()
+	defer c.closeMu.Unlock()
+
+	if c.closed {
+		return c.closeErr
+	}
+
+	c.closed = true
+	c.closeErr = c.Conn.Close()
+
+	return c.closeErr
+}
+
+// isStale returns whether the connection is in an invalid state (i.e. expired/broken).
+func (c *conn) isStale() bool {
 	expTime := c.idleAt.Add(c.idleTimeout)
-	return c.idleTimeout > 0 && time.Now().After(expTime)
+	return c.idleTimeout > 0 && time.Now().After(expTime) || c.broken.Load()
+}
+
+// isUpgraded returns whether this connection has been upgraded (e.g. Websocket).
+// An upgraded connection should not be reused and putted back in the connection pool.
+func (c *conn) isUpgraded() bool {
+	return c.upgraded.Load()
+}
+
+// readLoop handles the successive HTTP response read operations on the connection,
+// and watches for unsolicited bytes or connection errors when idle.
+func (c *conn) readLoop() {
+	defer c.Close()
+
+	for {
+		_, err := c.br.Peek(1)
+		if err != nil {
+			select {
+			// An error occurred while a response was expected to be handled.
+			case <-c.RWCh:
+				c.ErrCh <- err
+			// An error occurred on an idle connection.
+			default:
+				c.broken.Store(true)
+			}
+			return
+		}
+
+		// Unsolicited response received on an idle connection.
+		if !c.expectedResponse.Load() {
+			c.broken.Store(true)
+			return
+		}
+
+		r := <-c.RWCh
+		if err = c.handleResponse(r); err != nil {
+			c.ErrCh <- err
+			return
+		}
+
+		c.expectedResponse.Store(false)
+		c.ErrCh <- nil
+	}
+}
+
+func (c *conn) handleResponse(r rwWithUpgrade) error {
+	res := fasthttp.AcquireResponse()
+	defer fasthttp.ReleaseResponse(res)
+
+	res.Header.SetNoDefaultContentType(true)
+
+	for {
+		var (
+			timer      *time.Timer
+			errTimeout atomic.Pointer[timeoutError]
+		)
+		if c.responseHeaderTimeout > 0 {
+			timer = time.AfterFunc(c.responseHeaderTimeout, func() {
+				errTimeout.Store(&timeoutError{errors.New("timeout awaiting response headers")})
+				c.Close() // This close call is needed to interrupt the read operation below when the timeout is over.
+			})
+		}
+
+		res.Header.SetNoDefaultContentType(true)
+		if err := res.Header.Read(c.br); err != nil {
+			if c.responseHeaderTimeout > 0 {
+				if errT := errTimeout.Load(); errT != nil {
+					return errT
+				}
+			}
+			return err
+		}
+
+		if timer != nil {
+			timer.Stop()
+		}
+
+		fixPragmaCacheControl(&res.Header)
+
+		resCode := res.StatusCode()
+		is1xx := 100 <= resCode && resCode <= 199
+		// treat 101 as a terminal status, see issue 26161
+		is1xxNonTerminal := is1xx && resCode != http.StatusSwitchingProtocols
+		if is1xxNonTerminal {
+			removeConnectionHeaders(&res.Header)
+			h := r.RW.Header()
+
+			for _, header := range hopHeaders {
+				res.Header.Del(header)
+			}
+
+			res.Header.VisitAll(func(key, value []byte) {
+				r.RW.Header().Add(string(key), string(value))
+			})
+
+			r.RW.WriteHeader(res.StatusCode())
+			// Clear headers, it's not automatically done by ResponseWriter.WriteHeader() for 1xx responses
+			for k := range h {
+				delete(h, k)
+			}
+
+			res.Reset()
+			res.Header.Reset()
+			res.Header.SetNoDefaultContentType(true)
+
+			continue
+		}
+		break
+	}
+
+	announcedTrailers := res.Header.Peek("Trailer")
+
+	// Deal with 101 Switching Protocols responses: (WebSocket, h2c, etc)
+	if res.StatusCode() == http.StatusSwitchingProtocols {
+		r.Upgrade(r.RW, res, c)
+		c.upgraded.Store(true) // As the connection has been upgraded, it cannot be added back to the pool.
+		return nil
+	}
+
+	removeConnectionHeaders(&res.Header)
+
+	for _, header := range hopHeaders {
+		res.Header.Del(header)
+	}
+
+	if len(announcedTrailers) > 0 {
+		res.Header.Add("Trailer", string(announcedTrailers))
+	}
+
+	res.Header.VisitAll(func(key, value []byte) {
+		r.RW.Header().Add(string(key), string(value))
+	})
+
+	r.RW.WriteHeader(res.StatusCode())
+
+	if res.Header.ContentLength() == 0 {
+		return nil
+	}
+
+	// When a body is not allowed for a given status code the body is ignored.
+	// The connection will be marked as broken by the next Peek in the readloop.
+	if !isBodyAllowedForStatus(res.StatusCode()) {
+		return nil
+	}
+
+	// Chunked response, Content-Length is set to -1 by FastProxy when "Transfer-Encoding: chunked" header is received.
+	if res.Header.ContentLength() == -1 {
+		cbr := httputil.NewChunkedReader(c.br)
+
+		b := c.bufferPool.Get()
+		if b == nil {
+			b = make([]byte, bufferSize)
+		}
+		defer c.bufferPool.Put(b)
+
+		if _, err := io.CopyBuffer(&writeFlusher{r.RW}, cbr, b); err != nil {
+			return err
+		}
+
+		res.Header.Reset()
+		res.Header.SetNoDefaultContentType(true)
+		if err := res.Header.ReadTrailer(c.br); err != nil {
+			return err
+		}
+
+		if res.Header.Len() > 0 {
+			var announcedTrailersKey []string
+			if len(announcedTrailers) > 0 {
+				announcedTrailersKey = strings.Split(string(announcedTrailers), ",")
+			}
+
+			res.Header.VisitAll(func(key, value []byte) {
+				for _, s := range announcedTrailersKey {
+					if strings.EqualFold(s, strings.TrimSpace(string(key))) {
+						r.RW.Header().Add(string(key), string(value))
+						return
+					}
+				}
+
+				r.RW.Header().Add(http.TrailerPrefix+string(key), string(value))
+			})
+		}
+
+		return nil
+	}
+
+	brl := c.limitedReaderPool.Get()
+	if brl == nil {
+		brl = &io.LimitedReader{}
+	}
+	defer c.limitedReaderPool.Put(brl)
+
+	brl.R = c.br
+	brl.N = int64(res.Header.ContentLength())
+
+	b := c.bufferPool.Get()
+	if b == nil {
+		b = make([]byte, bufferSize)
+	}
+	defer c.bufferPool.Put(b)
+
+	if _, err := io.CopyBuffer(r.RW, brl, b); err != nil {
+		return err
+	}
+
+	return nil
 }
 
 // connPool is a net.Conn pool implementation using channels.
 type connPool struct {
-	dialer          func() (net.Conn, error)
-	idleConns       chan *conn
-	idleConnTimeout time.Duration
-	ticker          *time.Ticker
-	doneCh          chan struct{}
+	dialer                func() (net.Conn, error)
+	idleConns             chan *conn
+	idleConnTimeout       time.Duration
+	responseHeaderTimeout time.Duration
+	ticker                *time.Ticker
+	bufferPool            pool[[]byte]
+	limitedReaderPool     pool[*io.LimitedReader]
+	doneCh                chan struct{}
 }
 
 // newConnPool creates a new connPool.
-func newConnPool(maxIdleConn int, idleConnTimeout time.Duration, dialer func() (net.Conn, error)) *connPool {
+func newConnPool(maxIdleConn int, idleConnTimeout, responseHeaderTimeout time.Duration, dialer func() (net.Conn, error)) *connPool {
 	c := &connPool{
-		dialer:          dialer,
-		idleConns:       make(chan *conn, maxIdleConn),
-		idleConnTimeout: idleConnTimeout,
-		doneCh:          make(chan struct{}),
+		dialer:                dialer,
+		idleConns:             make(chan *conn, maxIdleConn),
+		idleConnTimeout:       idleConnTimeout,
+		responseHeaderTimeout: responseHeaderTimeout,
+		doneCh:                make(chan struct{}),
 	}
 
 	if idleConnTimeout > 0 {
@@ -72,22 +339,28 @@ func (c *connPool) AcquireConn() (*conn, error) {
 			return nil, err
 		}
 
-		if !co.isExpired() {
+		if !co.isStale() {
 			return co, nil
 		}
 
-		// As the acquired conn is expired we can close it
+		// As the acquired conn is stale we can close it
 		// without putting it again into the pool.
 		if err := co.Close(); err != nil {
 			log.Debug().
 				Err(err).
-				Msg("Unexpected error while releasing the connection")
+				Msg("Unexpected error while closing the connection")
 		}
 	}
 }
 
 // ReleaseConn releases the given net.Conn to the pool.
 func (c *connPool) ReleaseConn(co *conn) {
+	// An upgraded connection cannot be safely reused for another roundTrip,
+	// thus we are not putting it back to the pool.
+	if co.isUpgraded() {
+		return
+	}
+
 	co.idleAt = time.Now()
 	c.releaseConn(co)
 }
@@ -97,7 +370,7 @@ func (c *connPool) cleanIdleConns() {
 	for {
 		select {
 		case co := <-c.idleConns:
-			if !co.isExpired() {
+			if !co.isStale() {
 				c.releaseConn(co)
 				return
 			}
@@ -105,7 +378,7 @@ func (c *connPool) cleanIdleConns() {
 			if err := co.Close(); err != nil {
 				log.Debug().
 					Err(err).
-					Msg("Unexpected error while releasing the connection")
+					Msg("Unexpected error while closing the connection")
 			}
 
 		default:
@@ -155,9 +428,33 @@ func (c *connPool) askForNewConn(errCh chan<- error) {
 		return
 	}
 
-	c.releaseConn(&conn{
-		Conn:        co,
-		idleAt:      time.Now(),
-		idleTimeout: c.idleConnTimeout,
-	})
+	newConn := &conn{
+		Conn:                  co,
+		br:                    bufio.NewReaderSize(co, bufioSize),
+		idleAt:                time.Now(),
+		idleTimeout:           c.idleConnTimeout,
+		responseHeaderTimeout: c.responseHeaderTimeout,
+		RWCh:                  make(chan rwWithUpgrade),
+		ErrCh:                 make(chan error),
+		bufferPool:            &c.bufferPool,
+		limitedReaderPool:     &c.limitedReaderPool,
+	}
+	go newConn.readLoop()
+
+	c.releaseConn(newConn)
+}
+
+// isBodyAllowedForStatus reports whether a given response status code
+// permits a body. See RFC 7230, section 3.3.
+// From https://github.com/golang/go/blame/master/src/net/http/transfer.go#L459
+func isBodyAllowedForStatus(status int) bool {
+	switch {
+	case status >= 100 && status <= 199:
+		return false
+	case status == 204:
+		return false
+	case status == 304:
+		return false
+	}
+	return true
 }