dial: Redirect wss/ws correctly by modifying the http client

Closes #333

dial.go

@@ -70,6 +70,21 @@ func (opts *DialOptions) cloneWithDefaults(ctx context.Context) (context.Context
 	if o.HTTPHeader == nil {
 		o.HTTPHeader = http.Header{}
 	}
+	newClient := *o.HTTPClient
+	oldCheckRedirect := o.HTTPClient.CheckRedirect
+	newClient.CheckRedirect = func(req *http.Request, via []*http.Request) error {
+		switch req.URL.Scheme {
+		case "ws":
+			req.URL.Scheme = "http"
+		case "wss":
+			req.URL.Scheme = "https"
+		}
+		if oldCheckRedirect != nil {
+			return oldCheckRedirect(req, via)
+		}
+		return nil
+	}
+	o.HTTPClient = &newClient
 
 	return ctx, cancel, &o
 }

dial_test.go

@@ -304,3 +304,28 @@ type roundTripperFunc func(*http.Request) (*http.Response, error)
 func (f roundTripperFunc) RoundTrip(r *http.Request) (*http.Response, error) {
 	return f(r)
 }
+
+func TestDialRedirect(t *testing.T) {
+	t.Parallel()
+
+	ctx, cancel := context.WithTimeout(context.Background(), time.Second*5)
+	defer cancel()
+
+	_, _, err := Dial(ctx, "ws://example.com", &DialOptions{
+		HTTPClient: mockHTTPClient(func(r *http.Request) (*http.Response, error) {
+			resp := &http.Response{
+				Header: http.Header{},
+			}
+			if r.URL.Scheme != "https" {
+				resp.Header.Set("Location", "wss://example.com")
+				resp.StatusCode = http.StatusFound
+				return resp, nil
+			}
+			resp.Header.Set("Connection",  "Upgrade")
+			resp.Header.Set("Upgrade",  "meow")
+			resp.StatusCode = http.StatusSwitchingProtocols
+			return resp, nil
+		}),
+	})
+	assert.Contains(t, err, "failed to WebSocket dial: WebSocket protocol violation: Upgrade header \"meow\" does not contain websocket")
+}