martian: add 1m graceful close timeout for bicopy

This will prevent the proxy from leaking connections with defective servers.

This cannot be tested due to #1013. Enable the test again once that is fixed.
For now I tested it manually.

Contributes to #800.

Author: Choraden

internal/martian/closewriter.go renamed to internal/martian/close.go

@@ -43,3 +43,11 @@ func asCloseWriter(w io.Writer) (closeWriter, bool) {
 
 	return reflectx.LookupImpl[closeWriter](reflect.ValueOf(w))
 }
+
+func asCloser(w any) (io.Closer, bool) {
+	if c, ok := w.(io.Closer); ok {
+		return c, ok
+	}
+
+	return reflectx.LookupImpl[io.Closer](reflect.ValueOf(w))
+}

internal/martian/copy.go

@@ -21,6 +21,7 @@ import (
 	"context"
 	"io"
 	"sync"
+	"time"
 
 	"github.com/saucelabs/forwarder/internal/martian/log"
 )
@@ -43,13 +44,35 @@ var copyBufPool = sync.Pool{
 	},
 }
 
+var bicopyGracefulTimeout = 1 * time.Minute
+
 func bicopy(ctx context.Context, cc ...copier) {
+	ctx, cancel := context.WithCancel(ctx)
+	defer cancel()
+
 	donec := make(chan struct{}, len(cc))
 	for i := range cc {
 		go cc[i].copy(ctx, donec)
 	}
-	for range cc {
+
+	for i := range cc {
 		<-donec
+		if i == 0 {
+			// Forcibly close all tunnels 1 minute after the first tunnel finished.
+			go gracefulCloseAfter(ctx, bicopyGracefulTimeout, cc...)
+		}
+	}
+}
+
+func gracefulCloseAfter(ctx context.Context, d time.Duration, cc ...copier) {
+	select {
+	case <-ctx.Done():
+		return
+	case <-time.After(d):
+		log.Infof(ctx, "forcibly closing tunnel after %v of graceful period", d)
+	}
+	for i := range cc {
+		cc[i].close(ctx)
 	}
 }
 
@@ -67,6 +90,13 @@ func (c copier) copy(ctx context.Context, donec chan<- struct{}) {
 	if _, err := io.CopyBuffer(c.dst, c.src, buf); err != nil && !isClosedConnError(err) {
 		log.Errorf(ctx, "failed to copy %s tunnel: %v", c.name, err)
 	}
+	c.closeWriter(ctx)
+
+	log.Debugf(ctx, "%s tunnel finished copying", c.name)
+	donec <- struct{}{}
+}
+
+func (c copier) closeWriter(ctx context.Context) {
 	var closeErr error
 	if cw, ok := asCloseWriter(c.dst); ok {
 		closeErr = cw.CloseWrite()
@@ -78,7 +108,15 @@ func (c copier) copy(ctx context.Context, donec chan<- struct{}) {
 	if closeErr != nil {
 		log.Infof(ctx, "failed to close write side of %s tunnel: %v", c.name, closeErr)
 	}
+}
 
-	log.Debugf(ctx, "%s tunnel finished copying", c.name)
-	donec <- struct{}{}
+func (c copier) close(ctx context.Context) {
+	cc, ok := asCloser(c.dst)
+	if !ok {
+		log.Errorf(ctx, "cannot close %s tunnel (%T)", c.name, c.dst)
+		return
+	}
+	if err := cc.Close(); err != nil && !isClosedConnError(err) {
+		log.Infof(ctx, "failed to close %s tunnel: %v", c.name, err)
+	}
 }

internal/martian/proxy_test.go

@@ -2045,3 +2045,64 @@ func TestReadHeaderConnectionReset(t *testing.T) {
 		t.Fatalf("conn.Read(): got %v, want io.EOF", err)
 	}
 }
+
+func TestTunnelGracefulClose(t *testing.T) {
+	t.Parallel()
+
+	t.Skip("panic: close of closed channel, See #1013")
+
+	l, err := net.Listen("tcp", "localhost:0")
+	if err != nil {
+		t.Fatalf("net.Listen(): got %v, want no error", err)
+	}
+
+	done := make(chan struct{})
+
+	// Malicious server that hangs indefinitely.
+	go func() {
+		t.Logf("Waiting for server side connection")
+		conn, err := l.Accept()
+		if err != nil {
+			t.Errorf("Got error while accepting connection on destination listener: %v", err)
+			return
+		}
+		defer conn.Close()
+		t.Logf("Accepted server side connection")
+
+		<-done
+	}()
+
+	bicopyGracefulTimeout = 100 * time.Millisecond
+	h := testHelper{
+		Proxy: func(p *Proxy) {
+			tr := new(ProxyTrace)
+			tr.WroteResponse = func(info WroteResponseInfo) {
+				close(done)
+			}
+			p.Trace = tr
+		},
+	}
+	conn, cancel := h.proxyConn(t)
+	defer cancel()
+	defer conn.Close()
+
+	req, err := http.NewRequest(http.MethodConnect, "//"+l.Addr().String(), http.NoBody)
+	if err != nil {
+		t.Fatalf("http.NewRequest(): got %v, want no error", err)
+	}
+
+	if err := req.Write(conn); err != nil {
+		t.Fatalf("req.Write(): got %v, want no error", err)
+	}
+	res, err := http.ReadResponse(bufio.NewReader(conn), req)
+	if err != nil {
+		t.Fatalf("http.ReadResponse(): got %v, want no error", err)
+	}
+	res.Body.Close()
+	conn.Close()
+	select {
+	case <-done:
+	case <-time.After(5 * time.Second):
+		t.Fatalf("timed out waiting for tunnel to close all connections")
+	}
+}