1 files changed, 44 insertions, 11 deletions
diff --git a/vendor/golang.org/x/net/websocket/client.go b/vendor/golang.org/x/net/websocket/client.go
index 69a4ac7ee..1e64157f3 100644
--- a/vendor/golang.org/x/net/websocket/client.go
+++ b/vendor/golang.org/x/net/websocket/client.go
@@ -6,10 +6,12 @@ package websocket
 
 import (
 	"bufio"
+	"context"
 	"io"
 	"net"
 	"net/http"
 	"net/url"
+	"time"
 )
 
 // DialError is an error that occurs while dialling a websocket server.
@@ -79,28 +81,59 @@ func parseAuthority(location *url.URL) string {
 
 // DialConfig opens a new client connection to a WebSocket with a config.
 func DialConfig(config *Config) (ws *Conn, err error) {
-	var client net.Conn
+	return config.DialContext(context.Background())
+}
+
+// DialContext opens a new client connection to a WebSocket, with context support for timeouts/cancellation.
+func (config *Config) DialContext(ctx context.Context) (*Conn, error) {
 	if config.Location == nil {
 		return nil, &DialError{config, ErrBadWebSocketLocation}
 	}
 	if config.Origin == nil {
 		return nil, &DialError{config, ErrBadWebSocketOrigin}
 	}
+
 	dialer := config.Dialer
 	if dialer == nil {
 		dialer = &net.Dialer{}
 	}
-	client, err = dialWithDialer(dialer, config)
-	if err != nil {
-		goto Error
-	}
-	ws, err = NewClient(config, client)
+
+	client, err := dialWithDialer(ctx, dialer, config)
 	if err != nil {
-		client.Close()
-		goto Error
+		return nil, &DialError{config, err}
 	}
-	return
 
-Error:
-	return nil, &DialError{config, err}
+	// Cleanup the connection if we fail to create the websocket successfully
+	success := false
+	defer func() {
+		if !success {
+			_ = client.Close()
+		}
+	}()
+
+	var ws *Conn
+	var wsErr error
+	doneConnecting := make(chan struct{})
+	go func() {
+		defer close(doneConnecting)
+		ws, err = NewClient(config, client)
+		if err != nil {
+			wsErr = &DialError{config, err}
+		}
+	}()
+
+	// The websocket.NewClient() function can block indefinitely, make sure that we
+	// respect the deadlines specified by the context.
+	select {
+	case <-ctx.Done():
+		// Force the pending operations to fail, terminating the pending connection attempt
+		_ = client.SetDeadline(time.Now())
+		<-doneConnecting // Wait for the goroutine that tries to establish the connection to finish
+		return nil, &DialError{config, ctx.Err()}
+	case <-doneConnecting:
+		if wsErr == nil {
+			success = true // Disarm the deferred connection cleanup
+		}
+		return ws, wsErr
+	}
 }