enhance: Add base context for h2 stream connections

Author: josexy

examples/dumper/main.go

@@ -132,12 +132,15 @@ func main() {
 		)
 	}
 
+	ctx, cancel := context.WithCancel(context.Background())
+
 	handler, err := mitmpgo.NewMitmProxyHandler(
 		mitmpgo.WithCACertPath(caCertPath),
 		mitmpgo.WithCAKeyPath(caKeyPath),
 		mitmpgo.WithHTTPInterceptor(httpInterceptor),
 		mitmpgo.WithWebsocketInterceptor(websocketInterceptor),
 		mitmpgo.WithErrorHandler(errHandler),
+		mitmpgo.WithStreamBaseContext(ctx),
 		// mitmpgo.WithClientCert("127.0.0.1", mitmpgo.ClientCert{
 		// 	CertPath: "certs/client.crt",
 		// 	KeyPath:  "certs/client.key",
@@ -158,6 +161,7 @@ func main() {
 
 	listenAddr := fmt.Sprintf("%s:%d", "127.0.0.1", port)
 	var closeFn func()
+
 	switch mitmMode {
 	case "socks5":
 		ln, err := net.Listen("tcp", listenAddr)
@@ -172,14 +176,15 @@ func main() {
 					return
 				}
 				go func() {
-					handler.ServeSOCKS5(context.Background(), conn)
+					handler.ServeSOCKS5(ctx, conn)
 				}()
 			}
 		}()
 	default:
 		server := &http.Server{
-			Addr:    listenAddr,
-			Handler: handler,
+			Addr:        listenAddr,
+			Handler:     handler,
+			BaseContext: func(l net.Listener) context.Context { return ctx },
 		}
 		closeFn = func() { server.Close() }
 		go func() {
@@ -197,7 +202,8 @@ func main() {
 	handler.Cleanup()
 	slog.Info("exit")
 	closeFn()
-	time.Sleep(time.Millisecond * 100)
+	cancel()
+	time.Sleep(time.Millisecond * 500)
 }
 
 func httpInterceptor(ctx context.Context, req *http.Request, invoker mitmpgo.HTTPDelegatedInvoker) (*http.Response, error) {

mitm.go

@@ -660,7 +660,7 @@ func (r *mitmProxyHandler) handleTunnelRequest(ctx context.Context, consumedRequ
 		// and finally we only need to get the [http.Request] and process the [http.ResponseWriter].
 		// Early process http2
 		if state.NegotiatedProtocol == http2.NextProtoTLS {
-			newCtx, cancel := context.WithCancel(context.Background())
+			newCtx, cancel := context.WithCancel(r.streamBaseCtx)
 			go func() {
 				connCtx.local.waitClose()
 				cancel()
@@ -673,7 +673,7 @@ func (r *mitmProxyHandler) handleTunnelRequest(ctx context.Context, consumedRequ
 		}
 	}
 
-	ctx, earlyDone, isWsUpgrade, err := r.distinguishHTTPRequest(ctx, srcConn, dstConn, tlsRequest)
+	ctx, earlyDone, isWsUpgrade, err := r.distinguishHTTPRequest(ctx, srcConn, tlsRequest)
 	if err != nil || earlyDone {
 		return
 	}
@@ -691,7 +691,7 @@ func (r *mitmProxyHandler) handleH2CRequest(ctx context.Context, rw http.Respons
 			return false, err
 		}
 		connCtx := ctx.Value(connContextKey).(*biConnContext)
-		newCtx, cancel := context.WithCancel(context.Background())
+		newCtx, cancel := context.WithCancel(r.streamBaseCtx)
 		go func() {
 			connCtx.local.waitClose()
 			cancel()
@@ -711,7 +711,7 @@ func (r *mitmProxyHandler) handleH2CRequest(ctx context.Context, rw http.Respons
 			return false, err
 		}
 		connCtx := ctx.Value(connContextKey).(*biConnContext)
-		newCtx, cancel := context.WithCancel(context.Background())
+		newCtx, cancel := context.WithCancel(r.streamBaseCtx)
 		go func() {
 			connCtx.local.waitClose()
 			cancel()
@@ -727,7 +727,7 @@ func (r *mitmProxyHandler) handleH2CRequest(ctx context.Context, rw http.Respons
 	return false, nil
 }
 
-func (r *mitmProxyHandler) distinguishHTTPRequest(ctx context.Context, srcConn, dstConn net.Conn, tlsRequest bool) (newCtx context.Context, earlyDone bool, upgrade bool, retErr error) {
+func (r *mitmProxyHandler) distinguishHTTPRequest(ctx context.Context, srcConn net.Conn, tlsRequest bool) (newCtx context.Context, earlyDone bool, upgrade bool, retErr error) {
 	reqCtx, _ := FromRequestContext(ctx)
 
 	// Read the http request for https/wss via tls tunnel
@@ -859,13 +859,18 @@ func (r *mitmProxyHandler) relayConnForWS(ctx context.Context, srcConn, dstConn
 	}
 
 	errCh := make(chan error, 2)
-	relayWSMessage := func(dir WSDirection, src, dst *websocket.Conn) {
+	relayWSMessage := func(ctx context.Context, dir WSDirection, src, dst *websocket.Conn) {
 		defer func() {
 			if fw != nil {
 				fw.close()
 			}
 		}()
 		for {
+			select {
+			case <-ctx.Done():
+				return
+			default:
+			}
 			msgType, buffer, err := readBufferFromWSConn(src)
 			if err != nil {
 				errCh <- err
@@ -885,8 +890,8 @@ func (r *mitmProxyHandler) relayConnForWS(ctx context.Context, srcConn, dstConn
 			}
 		}
 	}
-	go relayWSMessage(Send, wsSrcConn, wsDstConn)
-	go relayWSMessage(Receive, wsDstConn, wsSrcConn)
+	go relayWSMessage(ctx, Send, wsSrcConn, wsDstConn)
+	go relayWSMessage(ctx, Receive, wsDstConn, wsSrcConn)
 	err = <-errCh
 	cancel(err)
 	return

option.go

@@ -1,6 +1,7 @@
 package mitmpgo
 
 import (
+	"context"
 	"crypto/x509"
 	"net"
 	"time"
@@ -23,6 +24,8 @@ func (f OptionFunc) apply(o *options) { f(o) }
 
 // options holds all configuration parameters for the MITM proxy handler.
 type options struct {
+	streamBaseCtx context.Context // Stream Base Context for h2 connection
+
 	proxy         string      // Upstream proxy URL (e.g., "http://127.0.0.1:8080")
 	caCertPath    string      // Path to the CA certificate file for TLS interception
 	caKeyPath     string      // Path to the CA private key file for TLS interception
@@ -60,13 +63,27 @@ func newOptions(opt ...Option) *options {
 	options := &options{
 		dialer:                &net.Dialer{Timeout: 15 * time.Second},
 		wsMaxFramesPerForward: 2048,
+		streamBaseCtx:         context.Background(),
 	}
 	for _, o := range opt {
 		o.apply(options)
 	}
 	return options
 }
 
+// WithStreamBaseContext configures h2 connection stream base context.
+//
+// Example:
+//
+//	handler, err := NewMitmProxyHandler(
+//	    WithStreamBaseContext(context.Background()),
+//	)
+func WithStreamBaseContext(baseCtx context.Context) Option {
+	return OptionFunc(func(o *options) {
+		o.streamBaseCtx = baseCtx
+	})
+}
+
 // WithProxy configures an upstream proxy server for outbound connections.
 //
 // The proxy parameter should be a URL in one of these formats: