From 6be79e969eb01e7a53b7c10fc9e1dc73ca2e3b55 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E4=B8=96=E7=95=8C?= <i@sekai.icu>
Date: Sun, 31 Dec 2023 15:51:48 +0800
Subject: [PATCH] Fix h2mux request context

---
 h2mux.go      | 15 +++++++++++++--
 h2mux_conn.go |  8 +++++++-
 2 files changed, 20 insertions(+), 3 deletions(-)

diff --git a/h2mux.go b/h2mux.go
index cb827a9..f328648 100644
--- a/h2mux.go
+++ b/h2mux.go
@@ -204,10 +204,21 @@ func (s *h2MuxClientSession) OpenContext(ctx context.Context) (net.Conn, error)
 		Body:   pipeInReader,
 		URL:    &url.URL{Scheme: "https", Host: "localhost"},
 	}
-	request = request.WithContext(ctx)
-	conn := newLateHTTPConn(pipeInWriter)
+	connCtx, cancel := context.WithCancel(context.Background())
+	request = request.WithContext(connCtx)
+	conn := newLateHTTPConn(pipeInWriter, cancel)
+	requestDone := make(chan struct{})
+	go func() {
+		select {
+		case <-requestDone:
+			return
+		case <-ctx.Done():
+			cancel()
+		}
+	}()
 	go func() {
 		response, err := s.transport.RoundTrip(request)
+		close(requestDone)
 		if err != nil {
 			conn.setup(nil, err)
 		} else if response.StatusCode != 200 {
diff --git a/h2mux_conn.go b/h2mux_conn.go
index 31c3286..f4f0ef0 100644
--- a/h2mux_conn.go
+++ b/h2mux_conn.go
@@ -1,6 +1,7 @@
 package mux
 
 import (
+	"context"
 	"io"
 	"net"
 	"os"
@@ -16,6 +17,7 @@ type httpConn struct {
 	writer io.Writer
 	create chan struct{}
 	err    error
+	cancel context.CancelFunc
 }
 
 func newHTTPConn(reader io.Reader, writer io.Writer) *httpConn {
@@ -25,10 +27,11 @@ func newHTTPConn(reader io.Reader, writer io.Writer) *httpConn {
 	}
 }
 
-func newLateHTTPConn(writer io.Writer) *httpConn {
+func newLateHTTPConn(writer io.Writer, cancel context.CancelFunc) *httpConn {
 	return &httpConn{
 		create: make(chan struct{}),
 		writer: writer,
+		cancel: cancel,
 	}
 }
 
@@ -55,6 +58,9 @@ func (c *httpConn) Write(b []byte) (n int, err error) {
 }
 
 func (c *httpConn) Close() error {
+	if c.cancel != nil {
+		c.cancel()
+	}
 	return common.Close(c.reader, c.writer)
 }