SetDialTLS should override dial func in EnableH2C

imroc · Dec 1, 2022 · fe9bec8 · fe9bec8
1 parent d673df0
commit fe9bec8
Show file tree

Hide file tree

Showing 3 changed files with 15 additions and 15 deletions.
diff --git a/internal/http2/go115.go b/internal/http2/go115.go
@@ -11,24 +11,18 @@ import (
 	"context"
 	"crypto/tls"
 	reqtls "github.com/imroc/req/v3/pkg/tls"
-	"net"
 )
 
 // dialTLSWithContext uses tls.Dialer, added in Go 1.15, to open a TLS
 // connection.
-func (t *Transport) dialTLSWithContext(ctx context.Context, network, addr string, cfg *tls.Config) (tlsCn reqtls.Conn, err error) {
-	var conn net.Conn
-	if t.DialTLSContext != nil {
-		conn, err = t.DialTLSContext(ctx, network, addr)
-	} else {
-		dialer := &tls.Dialer{
-			Config: cfg,
-		}
-		conn, err = dialer.DialContext(ctx, network, addr)
+func (t *Transport) dialTLSWithContext(ctx context.Context, network, addr string, cfg *tls.Config) (reqtls.Conn, error) {
+	dialer := &tls.Dialer{
+		Config: cfg,
 	}
+	conn, err := dialer.DialContext(ctx, network, addr)
 	if err != nil {
-		return
+		return nil, err
 	}
-	tlsCn = conn.(reqtls.Conn)
-	return
+	tlsCn := conn.(reqtls.Conn)
+	return tlsCn, nil
 }
diff --git a/internal/http2/transport.go b/internal/http2/transport.go
@@ -73,6 +73,7 @@ const (
 // for concurrent use by multiple goroutines.
 type Transport struct {
 	*transport.Options
+
 	// DialTLS specifies an optional dial function for creating
 	// TLS connections for requests.
 	//
@@ -560,6 +561,11 @@ func (t *Transport) dialTLS(ctx context.Context) func(string, string, *tls.Confi
 	if t.DialTLS != nil {
 		return t.DialTLS
 	}
+	if t.DialTLSContext != nil {
+		return func(network string, addr string, cfg *tls.Config) (net.Conn, error) {
+			return t.DialTLSContext(ctx, network, addr)
+		}
+	}
 	return func(network, addr string, cfg *tls.Config) (net.Conn, error) {
 		tlsCn, err := t.dialTLSWithContext(ctx, network, addr, cfg)
 		if err != nil {

diff --git a/transport.go b/transport.go
@@ -367,7 +367,7 @@ func (t *Transport) EnableForceHTTP2() *Transport {
 func (t *Transport) EnableH2C() *Transport {
 	t.Options.EnableH2C = true
 	t.t2.AllowHTTP = true
-	t.t2.DialTLS = func(network, addr string, cfg *tls.Config) (net.Conn, error) {
+	t.DialTLSContext = func(ctx context.Context, network, addr string) (net.Conn, error) {
 		return net.Dial(network, addr)
 	}
 	return t
@@ -377,7 +377,7 @@ func (t *Transport) EnableH2C() *Transport {
 func (t *Transport) DisableH2C() *Transport {
 	t.Options.EnableH2C = false
 	t.t2.AllowHTTP = false
-	t.t2.DialTLS = nil
+	t.t2.DialTLSContext = nil
 	return t
 }