Repository: gorilla/websocket
Pull request #746 (2021-12-10 09:53:06)

Dialer: add optional method NetDialTLSContext
Fixes issue: https://github.com/gorilla/websocket/issues/745

With the previous interface, NetDial and NetDialContext were used for
both TLS and non-TLS TCP connections, and afterwards TLSClientConfig was
used to do the TLS handshake.

While this API works for most cases, it prevents from using more advance
authentication methods during the TLS handshake, as this is out of the
control of the user.

This commits introduces another a new dial method, NetDialTLSContext,
which is used when dialing for TLS/TCP. The code then assumes that the
handshake is done there and TLSClientConfig is not used.

This API change is fully backwards compatible and it better aligns with
net/http.Transport API, which has these two dial flavors. See:
https://pkg.go.dev/net/http#Transport

lluiscampos: The first commit is unrelated to the issue at hand, but otherwise the CI was failing on `gofmt` for latest Go version 1.17.
garyburd: - Remove DialTLS field.  The corresponding field in net/http Transport is deprecated, so there's no need to add it here.
- Convert map of tests to slice of tests for consistency with other tests.

```
	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)
		netDialTLS        func(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,
                 …
		
```
lluiscampos: Thanks @garyburd for your feedback. I updated now the pull request following your suggestion. I updated the commit message accordingly too.

> * Remove DialTLS field.  The corresponding field in net/http Transport is deprecated, so there's no need to add it here.

As `Dialer` already had `Dial` and `DialContext` I thought it was good to add the two... but being this a new field, makes sense not to add the deprecated one.

> 
>     * Convert map of tests to slice of tests for consistency with other tests.

Roger that :+1: 



Patch:
diff --git a/README.md b/README.md
index 19aa2e7..2517a28 100644
--- a/README.md
+++ b/README.md
@@ -6,6 +6,13 @@
 Gorilla WebSocket is a [Go](http://golang.org/) implementation of the
 [WebSocket](http://www.rfc-editor.org/rfc/rfc6455.txt) protocol.
 
+
+---
+
+⚠️ **[The Gorilla WebSocket Package is looking for a new maintainer](https://github.com/gorilla/websocket/issues/370)**
+
+---
+
 ### Documentation
 
 * [API Reference](https://pkg.go.dev/github.com/gorilla/websocket?tab=doc)
@@ -30,35 +37,3 @@ The Gorilla WebSocket package passes the server tests in the [Autobahn Test
 Suite](https://github.com/crossbario/autobahn-testsuite) using the application in the [examples/autobahn
 subdirectory](https://github.com/gorilla/websocket/tree/master/examples/autobahn).
 
-### Gorilla WebSocket compared with other packages
-
-<table>
-<tr>
-<th></th>
-<th><a href="http://godoc.org/github.com/gorilla/websocket">github.com/gorilla</a></th>
-<th><a href="http://godoc.org/golang.org/x/net/websocket">golang.org/x/net</a></th>
-</tr>
-<tr>
-<tr><td colspan="3"><a href="http://tools.ietf.org/html/rfc6455">RFC 6455</a> Features</td></tr>
-<tr><td>Passes <a href="https://github.com/crossbario/autobahn-testsuite">Autobahn Test Suite</a></td><td><a href="https://github.com/gorilla/websocket/tree/master/examples/autobahn">Yes</a></td><td>No</td></tr>
-<tr><td>Receive <a href="https://tools.ietf.org/html/rfc6455#section-5.4">fragmented</a> message<td>Yes</td><td><a href="https://code.google.com/p/go/issues/detail?id=7632">No</a>, see note 1</td></tr>
-<tr><td>Send <a href="https://tools.ietf.org/html/rfc6455#section-5.5.1">close</a> message</td><td><a href="http://godoc.org/github.com/gorilla/websocket#hdr-Control_Messages">Yes</a></td><td><a href="https://code.google.com/p/go/issues/detail?id=4588">No</a></td></tr>
-<tr><td>Send <a href="https://tools.ietf.org/html/rfc6455#section-5.5.2">pings</a> and receive <a href="https://tools.ietf.org/html/rfc6455#section-5.5.3">pongs</a></td><td><a href="http://godoc.org/github.com/gorilla/websocket#hdr-Control_Messages">Yes</a></td><td>No</td></tr>
-<tr><td>Get the <a href="https://tools.ietf.org/html/rfc6455#section-5.6">type</a> of a received data message</td><td>Yes</td><td>Yes, see note 2</td></tr>
-<tr><td colspan="3">Other Features</tr></td>
-<tr><td><a href="https://tools.ietf.org/html/rfc7692">Compression Extensions</a></td><td>Experimental</td><td>No</td></tr>
-<tr><td>Read message using io.Reader</td><td><a href="http://godoc.org/github.com/gorilla/websocket#Conn.NextReader">Yes</a></td><td>No, see note 3</td></tr>
-<tr><td>Write message using io.WriteCloser</td><td><a href="http://godoc.org/github.com/gorilla/websocket#Conn.NextWriter">Yes</a></td><td>No, see note 3</td></tr>
-</table>
-
-Notes:
-
-1. Large messages are fragmented in [Chrome's new WebSocket implementation](http://www.ietf.org/mail-archive/web/hybi/current/msg10503.html).
-2. The application can get the type of a received data message by implementing
-   a [Codec marshal](http://godoc.org/golang.org/x/net/websocket#Codec.Marshal)
-   function.
-3. The go.net io.Reader and io.Writer operate across WebSocket frame boundaries.
-  Read returns when the input buffer is full or a frame boundary is
-  encountered. Each call to Write sends a single frame message. The Gorilla
-  io.Reader and io.WriteCloser operate on a single WebSocket message.
-
diff --git a/client.go b/client.go
index ecbe584..2efd835 100644
--- a/client.go
+++ b/client.go
@@ -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.
@@ -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.
@@ -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)
@@ -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


Test patch:
diff --git a/client_server_test.go b/client_server_test.go
index c03f2a9..e975e51 100644
--- a/client_server_test.go
+++ b/client_server_test.go
@@ -11,6 +11,7 @@ import (
 	"crypto/x509"
 	"encoding/base64"
 	"encoding/binary"
+	"errors"
 	"fmt"
 	"io"
 	"io/ioutil"
@@ -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()
+		}
+	}
+}