TUN-10621: Propagate max wait timeout

This PR addresses an issue where cloudflared prematurely closes the origin connection before the upstream-to-downstream goroutine finishes reading, causing intermittent connection drops when a client immediately closes the write-side of a connection.

When a client finishes writing data, it immediately closes its side of the connection. Under the current implementation in cloudflared's downstream-to-upstream goroutine does not wait for the second stream to complete. It unblocks the moment the first stream writes to the channel. Once this happens, the pipe returns control to `proxyTCPStream`, which prematurely closes the origin connection. Consequently, when the upstream-to-downstream goroutine attempts to read the remaining data from the origin connection, the connection is already gone, leading to unexpected failures.

We started propagating the `TimeoutAfterFirstClose` configuration/parameter. This allows the proxy to wait for a designated period, giving the second stream sufficient time to finish processing and read all remaining data before `proxyTCPStream` tears down the origin connection.

Author: macmarcelino

carrier/websocket.go

@@ -37,7 +37,7 @@ func (ws *Websocket) ServeStream(options *StartOptions, conn io.ReadWriter) erro
 	}
 	defer func() { _ = wsConn.Close() }()
 
-	stream.Pipe(wsConn, conn, ws.log)
+	stream.Pipe(wsConn, conn, stream.DefaultTimeoutAfterFirstClose, ws.log)
 	return nil
 }
 

connection/connection_test.go

@@ -151,7 +151,7 @@ func wsEchoEndpoint(w ResponseWriter, r *http.Request) error {
 	}()
 
 	originConn := &echoPipe{reader: readPipe, writer: writePipe}
-	stream.Pipe(wsConn, originConn, &log)
+	stream.Pipe(wsConn, originConn, 0, &log)
 	cancel()
 	wsConn.Close()
 	return nil
@@ -190,7 +190,7 @@ func wsFlakyEndpoint(w ResponseWriter, r *http.Request) error {
 	rInt, _ := rand.Int(rand.Reader, big.NewInt(50))
 	closedAfter := time.Millisecond * time.Duration(rInt.Int64())
 	originConn := &flakyConn{closeAt: time.Now().Add(closedAfter)}
-	stream.Pipe(wsConn, originConn, &log)
+	stream.Pipe(wsConn, originConn, 0, &log)
 	cancel()
 	wsConn.Close()
 	return nil

ingress/origin_connection.go

@@ -25,9 +25,9 @@ type OriginConnection interface {
 type streamHandlerFunc func(originConn io.ReadWriter, remoteConn net.Conn, log *zerolog.Logger)
 
 // DefaultStreamHandler is an implementation of streamHandlerFunc that
-// performs a two way io.Copy between originConn and remoteConn.
+// performs a two-way io.Copy between originConn and remoteConn.
 func DefaultStreamHandler(originConn io.ReadWriter, remoteConn net.Conn, log *zerolog.Logger) {
-	stream.Pipe(originConn, remoteConn, log)
+	stream.Pipe(originConn, remoteConn, stream.DefaultTimeoutAfterFirstClose, log)
 }
 
 // tcpConnection is an OriginConnection that directly streams to raw TCP.
@@ -38,7 +38,7 @@ type tcpConnection struct {
 }
 
 func (tc *tcpConnection) Stream(_ context.Context, tunnelConn io.ReadWriter, _ *zerolog.Logger) {
-	stream.Pipe(tunnelConn, tc, tc.logger)
+	stream.Pipe(tunnelConn, tc, stream.DefaultTimeoutAfterFirstClose, tc.logger)
 }
 
 func (tc *tcpConnection) Write(b []byte) (int, error) {

ingress/origin_connection_test.go

@@ -38,6 +38,7 @@ func TestStreamTCPConnection(t *testing.T) {
 	tcpConn := tcpConnection{
 		Conn:         cfdConn,
 		writeTimeout: 30 * time.Second,
+		logger:       TestLogger,
 	}
 
 	eyeballConn, edgeConn := net.Pipe()
@@ -157,7 +158,7 @@ func TestSocksStreamWSOverTCPConnection(t *testing.T) {
 			require.NoError(t, err)
 			defer func() { _ = wsForwarderInConn.Close() }()
 
-			stream.Pipe(wsForwarderInConn, &wsEyeball{wsForwarderOutConn}, TestLogger)
+			stream.Pipe(wsForwarderInConn, &wsEyeball{wsForwarderOutConn}, stream.DefaultTimeoutAfterFirstClose, TestLogger)
 			return nil
 		})
 

ingress/origin_service.go

@@ -99,7 +99,6 @@ type rawTCPService struct {
 	name         string
 	dialer       net.Dialer
 	writeTimeout time.Duration
-	logger       *zerolog.Logger
 }
 
 func (o *rawTCPService) String() string {
@@ -233,10 +232,12 @@ func (o *helloWorld) start(
 	if err != nil {
 		return errors.Wrap(err, "Cannot start Hello World Server")
 	}
-	go hello.StartHelloWorldServer(log, helloListener, shutdownC)
+	go func() {
+		_ = hello.StartHelloWorldServer(log, helloListener, shutdownC)
+	}()
 	o.server = helloListener
 
-	o.httpService.url = &url.URL{
+	o.url = &url.URL{
 		Scheme: "https",
 		Host:   o.server.Addr().String(),
 	}
@@ -356,7 +357,7 @@ func newHTTPTransport(service OriginService, cfg OriginRequestConfig, log *zerol
 		IdleConnTimeout:       cfg.KeepAliveTimeout.Duration,
 		TLSHandshakeTimeout:   cfg.TLSTimeout.Duration,
 		ExpectContinueTimeout: 1 * time.Second,
-		TLSClientConfig:       &tls.Config{RootCAs: originCertPool, InsecureSkipVerify: cfg.NoTLSVerify},
+		TLSClientConfig:       &tls.Config{RootCAs: originCertPool, InsecureSkipVerify: cfg.NoTLSVerify}, //nolint: gosec
 		ForceAttemptHTTP2:     cfg.Http2Origin,
 	}
 	if _, isHelloWorld := service.(*helloWorld); !isHelloWorld && cfg.OriginServerName != "" {
@@ -374,7 +375,6 @@ func newHTTPTransport(service OriginService, cfg OriginRequestConfig, log *zerol
 	// DialContext depends on which kind of origin is being used.
 	dialContext := dialer.DialContext
 	switch service := service.(type) {
-
 	// If this origin is a unix socket, enforce network type "unix".
 	case *unixSocketPath:
 		httpTransport.DialContext = func(ctx context.Context, _, _ string) (net.Conn, error) {

proxy/proxy.go

@@ -256,7 +256,7 @@ func (p *Proxy) proxyHTTPRequest(
 			reader: tr.Body,
 		}
 
-		stream.Pipe(eyeballStream, rwc, logger)
+		stream.Pipe(eyeballStream, rwc, time.Second*0, logger)
 		return nil
 	}
 
@@ -311,7 +311,7 @@ func (p *Proxy) proxyStream(
 
 // proxyTCPStream proxies private network type TCP connections as a stream towards an available origin.
 //
-// This is different than proxyStream because it's not leveraged ingress rule services and uses the
+// This is different from proxyStream because it's not leveraged ingress rule services and uses the
 // originDialer from OriginDialerService.
 func (p *Proxy) proxyTCPStream(
 	tr *tracing.TracedContext,
@@ -344,7 +344,8 @@ func (p *Proxy) proxyTCPStream(
 	connectLatency.Observe(float64(time.Since(start).Milliseconds()))
 	logger.Debug().Msg("proxy stream acknowledged")
 
-	stream.Pipe(tunnelConn, originConn, logger)
+	stream.Pipe(tunnelConn, originConn, stream.DefaultTimeoutAfterFirstClose, logger)
+
 	return nil
 }
 

proxy/proxy_test.go

@@ -223,9 +223,9 @@ func testProxyWebsocket(proxy connection.OriginProxy) func(t *testing.T) {
 			}
 			if ctx.Err() == context.DeadlineExceeded {
 				t.Errorf("Test timed out")
-				readPipe.Close()
-				writePipe.Close()
-				responseWriter.Close()
+				_ = readPipe.Close()
+				_ = writePipe.Close()
+				_ = responseWriter.Close()
 			}
 			return nil
 		})
@@ -647,7 +647,7 @@ func TestConnections(t *testing.T) {
 				ingressServiceScheme: "tcp://",
 				originService: func(t *testing.T, ln net.Listener) {
 					// closing the listener created by the test.
-					ln.Close()
+					_ = ln.Close()
 				},
 				eyeballResponseWriter: newTCPRespWriter(replayer),
 				eyeballRequestBody:    newTCPRequestBody([]byte("test2")),
@@ -756,6 +756,8 @@ func newTCPRequestBody(data []byte) *requestBody {
 	pr, pw := io.Pipe()
 	go func() {
 		_, _ = pw.Write(data)
+		// Close the write side once the payload has been sent.
+		_ = pw.Close()
 	}()
 	return &requestBody{
 		pr: pr,
@@ -801,8 +803,8 @@ func (p *pipedRequestBody) roundtrip(addr string) []byte {
 	if err != nil {
 		panic(err)
 	}
-	defer conn.Close()
-	defer resp.Body.Close()
+	defer func() { _ = conn.Close() }()
+	defer func() { _ = resp.Body.Close() }()
 
 	if resp.StatusCode != http.StatusSwitchingProtocols {
 		panic(fmt.Errorf("resp returned status code: %d", resp.StatusCode))
@@ -949,7 +951,7 @@ func runEchoTCPService(t *testing.T, l net.Listener) {
 			if err != nil {
 				panic(err)
 			}
-			defer conn.Close()
+			defer func() { _ = conn.Close() }()
 
 			for {
 				buf := make([]byte, 1024)
@@ -987,7 +989,7 @@ func runEchoWSService(t *testing.T, l net.Listener) {
 			t.Log(err)
 			return
 		}
-		defer conn.Close()
+		defer func() { _ = conn.Close() }()
 
 		for {
 			messageType, p, err := conn.ReadMessage()

stream/stream.go

@@ -2,19 +2,25 @@ package stream
 
 import (
 	"encoding/hex"
+	"errors"
 	"fmt"
 	"io"
 	"runtime/debug"
 	"sync/atomic"
 	"time"
 
 	"github.com/getsentry/sentry-go"
-	"github.com/pkg/errors"
 	"github.com/rs/zerolog"
 
 	"github.com/cloudflare/cloudflared/cfio"
 )
 
+// DefaultTimeoutAfterFirstClose controls the upper bound of how long we wait for the second stream to finish.
+// We use bidirectional streams for our communication. Since the read and write sides can be closed independently,
+// we must have a way to close the second stream once the first one finishes. We don't want to wait indefinitely,
+// since we want to prevent misbehaving clients from blocking cloudflared.
+const DefaultTimeoutAfterFirstClose = time.Second * 10
+
 type Stream interface {
 	Reader
 	WriterCloser
@@ -37,7 +43,7 @@ type nopCloseWriterAdapter struct {
 	io.ReadWriter
 }
 
-func NopCloseWriterAdapter(stream io.ReadWriter) *nopCloseWriterAdapter {
+func noopCloseWriter(stream io.ReadWriter) *nopCloseWriterAdapter {
 	return &nopCloseWriterAdapter{stream}
 }
 
@@ -72,7 +78,7 @@ func (s *bidirectionalStreamStatus) wait(maxWaitForSecondStream time.Duration) e
 
 		select {
 		case <-timer.C:
-			return fmt.Errorf("timeout waiting for second stream to finish")
+			return fmt.Errorf("timeout waiting for second stream to finish %s", maxWaitForSecondStream)
 		case <-s.doneChan:
 			return nil
 		}
@@ -85,23 +91,27 @@ func (s *bidirectionalStreamStatus) isAnyDone() bool {
 }
 
 // Pipe copies copy data to & from provided io.ReadWriters.
-func Pipe(tunnelConn, originConn io.ReadWriter, log *zerolog.Logger) {
-	_ = PipeBidirectional(NopCloseWriterAdapter(tunnelConn), NopCloseWriterAdapter(originConn), 0, log)
+func Pipe(tunnelConn, originConn io.ReadWriter, timeoutAfterFirstClose time.Duration, log *zerolog.Logger) {
+	if err := PipeBidirectional(noopCloseWriter(tunnelConn), noopCloseWriter(originConn), timeoutAfterFirstClose, log); err != nil {
+		log.Warn().Err(err).Msg("Failed to pipe bidirectional stream")
+	}
 }
 
-// PipeBidirectional copies data to two unidirectional streams. It is a special case of Pipe where it receives a concept that allows for Read and Write side to be closed independently.
-// The main difference is that when piping data from a reader to a writer, if EOF is read, then this implementation propagates the EOF signal to the destination/writer by closing the write side of the
-// Bidirectional Stream.
-// Finally, depending on once EOF is ready from one of the provided streams, the other direction of streaming data will have a configured time period to also finish, otherwise,
-// the method will return immediately  with a timeout error. It is however, the responsibility of the caller to close the associated streams in both ends in order to free all the resources/go-routines.
+// PipeBidirectional copies data between two unidirectional streams. It is a special case of Pipe that accepts streams
+// whose read and write sides can be closed independently. The main difference is that when piping data from a reader
+// to a writer, if EOF is read, this implementation propagates the EOF signal to the destination by closing the write
+// side of the bidirectional stream.
+// Finally, once EOF is received from one of the provided streams, the other direction has a configured grace period to
+// finish; otherwise, the method returns a timeout error. It is, however, the responsibility of the caller to close
+// the associated streams at both ends in order to free all resources and goroutines.
 func PipeBidirectional(downstream, upstream Stream, maxWaitForSecondStream time.Duration, log *zerolog.Logger) error {
 	status := newBiStreamStatus()
 
 	go unidirectionalStream(downstream, upstream, "upstream->downstream", status, log)
 	go unidirectionalStream(upstream, downstream, "downstream->upstream", status, log)
 
 	if err := status.wait(maxWaitForSecondStream); err != nil {
-		return errors.Wrap(err, "unable to wait for both streams while proxying")
+		return fmt.Errorf("unable to wait for both streams while proxying: %w", err)
 	}
 
 	return nil

stream/stream_test.go

@@ -30,8 +30,8 @@ func TestPipeBidirectionalFinishOneSideTimeout(t *testing.T) {
 
 func TestPipeBidirectionalClosingWriteBothSidesAlsoExists(t *testing.T) {
 	fun := func(upstream, downstream *mockedStream) {
-		downstream.CloseWrite()
-		upstream.CloseWrite()
+		_ = downstream.CloseWrite()
+		_ = upstream.CloseWrite()
 
 		downstream.writeToReader("abc")
 		upstream.writeToReader("abc")
@@ -42,7 +42,7 @@ func TestPipeBidirectionalClosingWriteBothSidesAlsoExists(t *testing.T) {
 
 func TestPipeBidirectionalClosingWriteSingleSideAlsoExists(t *testing.T) {
 	fun := func(upstream, downstream *mockedStream) {
-		downstream.CloseWrite()
+		_ = downstream.CloseWrite()
 
 		downstream.writeToReader("abc")
 		upstream.writeToReader("abc")
@@ -51,6 +51,46 @@ func TestPipeBidirectionalClosingWriteSingleSideAlsoExists(t *testing.T) {
 	testPipeBidirectionalUnblocking(t, fun, time.Millisecond*200, true)
 }
 
+// TestPipeBidirectionalReturnsWhenBothSidesFinish verifies that
+// PipeBidirectional returns as soon as both stream directions finish, without
+// waiting for the full timeout grace period to expire. This guards against a
+// regression where the second-stream wait would block for the whole timeout
+// even when the result is already available.
+func TestPipeBidirectionalReturnsWhenBothSidesFinish(t *testing.T) {
+	t.Parallel()
+
+	const (
+		timeout     = time.Second * 5
+		maxWallTime = time.Millisecond * 500
+	)
+
+	logger := zerolog.Nop()
+	downstream := newMockedStream()
+	upstream := newMockedStream()
+
+	resultCh := make(chan error, 1)
+	go func() {
+		resultCh <- PipeBidirectional(downstream, upstream, timeout, &logger)
+	}()
+
+	// Close both reader sides so both stream directions reach EOF promptly.
+	downstream.closeReader()
+	upstream.closeReader()
+
+	start := time.Now()
+	select {
+	case err := <-resultCh:
+		elapsed := time.Since(start)
+		require.NoError(t, err)
+		require.Less(t, elapsed, maxWallTime,
+			"PipeBidirectional should return as soon as both streams finish, not after the full %s timeout (took %s)",
+			timeout, elapsed,
+		)
+	case <-time.After(timeout):
+		require.Fail(t, "PipeBidirectional did not return before the timeout expired")
+	}
+}
+
 func testPipeBidirectionalUnblocking(t *testing.T, afterFun func(*mockedStream, *mockedStream), timeout time.Duration, expectTimeout bool) {
 	logger := zerolog.Nop()
 
@@ -67,9 +107,9 @@ func testPipeBidirectionalUnblocking(t *testing.T, afterFun func(*mockedStream,
 	select {
 	case err := <-resultCh:
 		if expectTimeout {
-			require.NotNil(t, err)
+			require.Error(t, err)
 		} else {
-			require.Nil(t, err)
+			require.NoError(t, err)
 		}
 
 	case <-time.After(timeout * 2):