Fix h2mux request context

Author: nekohasekai

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 {

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)
 }