Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Dialer: add optional method NetDialTLSContext #746

Merged
merged 1 commit into from
Jan 4, 2022
Merged

Dialer: add optional method NetDialTLSContext #746

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
45 changes: 37 additions & 8 deletions client.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -56,9 +56,15 @@ type Dialer struct {
NetDial func(network, addr string) (net.Conn, error)

// NetDialContext specifies the dial function for creating TCP connections. If
// NetDialContext is nil, net.DialContext is used.
// NetDialContext is nil, NetDial is used.
NetDialContext func(ctx context.Context, network, addr string) (net.Conn, error)

// NetDialTLSContext specifies the dial function for creating TLS/TCP connections. If
// NetDialTLSContext is nil, NetDialContext is used.
// If NetDialTLSContext is set, Dial assumes the TLS handshake is done there and
// TLSClientConfig is ignored.
NetDialTLSContext func(ctx context.Context, network, addr string) (net.Conn, error)

// Proxy specifies a function to return a proxy for a given
// Request. If the function returns a non-nil error, the
// request is aborted with the provided error.
Expand All @@ -67,6 +73,8 @@ type Dialer struct {

// TLSClientConfig specifies the TLS configuration to use with tls.Client.
// If nil, the default configuration is used.
// If either NetDialTLS or NetDialTLSContext are set, Dial assumes the TLS handshake
// is done there and TLSClientConfig is ignored.
TLSClientConfig *tls.Config

// HandshakeTimeout specifies the duration for the handshake to complete.
Expand Down Expand Up @@ -239,13 +247,32 @@ func (d *Dialer) DialContext(ctx context.Context, urlStr string, requestHeader h
// Get network dial function.
var netDial func(network, add string) (net.Conn, error)

if d.NetDialContext != nil {
netDial = func(network, addr string) (net.Conn, error) {
return d.NetDialContext(ctx, network, addr)
switch u.Scheme {
case "http":
if d.NetDialContext != nil {
netDial = func(network, addr string) (net.Conn, error) {
return d.NetDialContext(ctx, network, addr)
}
} else if d.NetDial != nil {
netDial = d.NetDial
}
} else if d.NetDial != nil {
netDial = d.NetDial
} else {
case "https":
if d.NetDialTLSContext != nil {
netDial = func(network, addr string) (net.Conn, error) {
return d.NetDialTLSContext(ctx, network, addr)
}
} else if d.NetDialContext != nil {
netDial = func(network, addr string) (net.Conn, error) {
return d.NetDialContext(ctx, network, addr)
}
} else if d.NetDial != nil {
netDial = d.NetDial
}
default:
return nil, nil, errMalformedURL
}

if netDial == nil {
netDialer := &net.Dialer{}
netDial = func(network, addr string) (net.Conn, error) {
return netDialer.DialContext(ctx, network, addr)
Expand Down Expand Up @@ -306,7 +333,9 @@ func (d *Dialer) DialContext(ctx context.Context, urlStr string, requestHeader h
}
}()

if u.Scheme == "https" {
if u.Scheme == "https" && d.NetDialTLSContext == nil {
// If NetDialTLSContext is set, assume that the TLS handshake has already been done

cfg := cloneTLSConfig(d.TLSClientConfig)
if cfg.ServerName == "" {
cfg.ServerName = hostNoPort
Expand Down
178 changes: 178 additions & 0 deletions client_server_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,7 @@ import (
"crypto/x509"
"encoding/base64"
"encoding/binary"
"errors"
"fmt"
"io"
"io/ioutil"
Expand Down Expand Up @@ -920,3 +921,180 @@ func TestEmptyTracingDialWithContext(t *testing.T) {
defer ws.Close()
sendRecv(t, ws)
}

// TestNetDialConnect tests selection of dial method between NetDial, NetDialContext, NetDialTLS or NetDialTLSContext
func TestNetDialConnect(t *testing.T) {

upgrader := Upgrader{}
handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
if IsWebSocketUpgrade(r) {
c, err := upgrader.Upgrade(w, r, http.Header{"X-Test-Host": {r.Host}})
if err != nil {
t.Fatal(err)
}
c.Close()
} else {
w.Header().Set("X-Test-Host", r.Host)
}
})

server := httptest.NewServer(handler)
defer server.Close()

tlsServer := httptest.NewTLSServer(handler)
defer tlsServer.Close()

testUrls := map[*httptest.Server]string{
server: "ws://" + server.Listener.Addr().String() + "/",
tlsServer: "wss://" + tlsServer.Listener.Addr().String() + "/",
}

cas := rootCAs(t, tlsServer)
tlsConfig := &tls.Config{
RootCAs: cas,
ServerName: "example.com",
InsecureSkipVerify: false,
}

tests := []struct {
name string
server *httptest.Server // server to use
netDial func(network, addr string) (net.Conn, error)
netDialContext func(ctx context.Context, network, addr string) (net.Conn, error)
netDialTLSContext func(ctx context.Context, network, addr string) (net.Conn, error)
tlsClientConfig *tls.Config
}{

{
name: "HTTP server, all NetDial* defined, shall use NetDialContext",
server: server,
netDial: func(network, addr string) (net.Conn, error) {
return nil, errors.New("NetDial should not be called")
},
netDialContext: func(_ context.Context, network, addr string) (net.Conn, error) {
return net.Dial(network, addr)
},
netDialTLSContext: func(_ context.Context, network, addr string) (net.Conn, error) {
return nil, errors.New("NetDialTLSContext should not be called")
},
tlsClientConfig: nil,
},
{
name: "HTTP server, all NetDial* undefined",
server: server,
netDial: nil,
netDialContext: nil,
netDialTLSContext: nil,
tlsClientConfig: nil,
},
{
name: "HTTP server, NetDialContext undefined, shall fallback to NetDial",
server: server,
netDial: func(network, addr string) (net.Conn, error) {
return net.Dial(network, addr)
},
netDialContext: nil,
netDialTLSContext: func(ctx context.Context, network, addr string) (net.Conn, error) {
return nil, errors.New("NetDialTLSContext should not be called")
},
tlsClientConfig: nil,
},
{
name: "HTTPS server, all NetDial* defined, shall use NetDialTLSContext",
server: tlsServer,
netDial: func(network, addr string) (net.Conn, error) {
return nil, errors.New("NetDial should not be called")
},
netDialContext: func(ctx context.Context, network, addr string) (net.Conn, error) {
return nil, errors.New("NetDialContext should not be called")
},
netDialTLSContext: func(ctx context.Context, network, addr string) (net.Conn, error) {
netConn, err := net.Dial(network, addr)
if err != nil {
return nil, err
}
tlsConn := tls.Client(netConn, tlsConfig)
err = tlsConn.Handshake()
if err != nil {
return nil, err
}
return tlsConn, nil
},
tlsClientConfig: nil,
},
{
name: "HTTPS server, NetDialTLSContext undefined, shall fallback to NetDialContext and do handshake",
server: tlsServer,
netDial: func(network, addr string) (net.Conn, error) {
return nil, errors.New("NetDial should not be called")
},
netDialContext: func(ctx context.Context, network, addr string) (net.Conn, error) {
return net.Dial(network, addr)
},
netDialTLSContext: nil,
tlsClientConfig: tlsConfig,
},
{
name: "HTTPS server, NetDialTLSContext and NetDialContext undefined, shall fallback to NetDial and do handshake",
server: tlsServer,
netDial: func(network, addr string) (net.Conn, error) {
return net.Dial(network, addr)
},
netDialContext: nil,
netDialTLSContext: nil,
tlsClientConfig: tlsConfig,
},
{
name: "HTTPS server, all NetDial* undefined",
server: tlsServer,
netDial: nil,
netDialContext: nil,
netDialTLSContext: nil,
tlsClientConfig: tlsConfig,
},
{
name: "HTTPS server, all NetDialTLSContext defined, dummy TlsClientConfig defined, shall not do handshake",
server: tlsServer,
netDial: func(network, addr string) (net.Conn, error) {
return nil, errors.New("NetDial should not be called")
},
netDialContext: func(ctx context.Context, network, addr string) (net.Conn, error) {
return nil, errors.New("NetDialContext should not be called")
},
netDialTLSContext: func(ctx context.Context, network, addr string) (net.Conn, error) {
netConn, err := net.Dial(network, addr)
if err != nil {
return nil, err
}
tlsConn := tls.Client(netConn, tlsConfig)
err = tlsConn.Handshake()
if err != nil {
return nil, err
}
return tlsConn, nil
},
tlsClientConfig: &tls.Config{
RootCAs: nil,
ServerName: "badserver.com",
InsecureSkipVerify: false,
},
},
}

for _, tc := range tests {
dialer := Dialer{
NetDial: tc.netDial,
NetDialContext: tc.netDialContext,
NetDialTLSContext: tc.netDialTLSContext,
TLSClientConfig: tc.tlsClientConfig,
}

// Test websocket dial
c, _, err := dialer.Dial(testUrls[tc.server], nil)
if err != nil {
t.Errorf("FAILED %s, err: %s", tc.name, err.Error())
} else {
c.Close()
}
}
}