websocket: allow "https" scheme on Client's Endpoint

Previously, we only check for "wss" scheme to initialize the TLS
configuration.  At some point, user may have one single address to
connect to HTTP and WebSocket server on the same domain, which require
two addresses on their side.

This commit relaxing the Endpoint rule when parsing WebSocket URI.
Scheme "https" will be considered as "wss" or secure connection with
TLS.

2 files changed, 68 insertions, 12 deletions

diff --git a/lib/websocket/client.go b/lib/websocket/client.go
index f19dcefe..455a4b78 100644
--- a/lib/websocket/client.go
+++ b/lib/websocket/client.go
@@ -28,6 +28,11 @@ const (
 		"Connection: Upgrade\r\n" +
 		"Sec-Websocket-Key: %s\r\n" +
 		"Sec-Websocket-Version: 13\r\n"
+
+	schemeWSS   = "wss"
+	schemeHTTPS = "https"
+	defTLSPort  = "443"
+	defPort     = "80"
 )
 
 var (
@@ -295,10 +300,10 @@ func (cl *Client) init() (err error) {
 }
 
 //
-// parseURI parse WebSocket connection URI from "endpoint" and get the remote
+// parseURI parse WebSocket connection URI from Endpoint and get the remote
 // URL (for checking up scheme) and remote address.
-// By default, if no port is given, it will set to 80 for URL with any scheme
-// or 443 for "wss" scheme.
+// By default, if no port is given, it will be set to 80 or 443 for "wss" or
+// "https" scheme.
 //
 // On success it will set the remote address that can be used on open().
 // On fail it will return an error.
@@ -310,24 +315,24 @@ func (cl *Client) parseURI() (err error) {
 		return err
 	}
 
+	serverAddress := cl.remoteURL.Hostname()
 	serverPort := cl.remoteURL.Port()
 
-	if len(serverPort) != 0 {
-		cl.remoteAddr = cl.remoteURL.Host
-		return nil
-	}
-
 	switch cl.remoteURL.Scheme {
-	case "wss":
-		serverPort = "443"
+	case schemeWSS, schemeHTTPS:
+		if len(serverPort) == 0 {
+			serverPort = defTLSPort
+		}
 		if cl.TLSConfig == nil {
 			cl.TLSConfig = &tls.Config{}
 		}
 	default:
-		serverPort = "80"
+		if len(serverPort) == 0 {
+			serverPort = defPort
+		}
 	}
 
-	cl.remoteAddr = cl.remoteURL.Hostname() + ":" + serverPort
+	cl.remoteAddr = serverAddress + ":" + serverPort
 
 	return nil
 }
diff --git a/lib/websocket/client_test.go b/lib/websocket/client_test.go
index 9aae9cfc..466e0932 100644
--- a/lib/websocket/client_test.go
+++ b/lib/websocket/client_test.go
@@ -5,6 +5,7 @@
 package websocket
 
 import (
+	"crypto/tls"
 	"net/http"
 	"sync"
 	"testing"
@@ -60,6 +61,56 @@ func TestConnect(t *testing.T) {
 	}
 }
 
+func TestClient_parseURI(t *testing.T) {
+	cl := &Client{}
+
+	cases := []struct {
+		endpoint         string
+		expRemoteAddress string
+		expTLSConfig     *tls.Config
+		expError         string
+	}{{
+		endpoint:         "ws://127.0.0.1:8080",
+		expRemoteAddress: "127.0.0.1:8080",
+	}, {
+		endpoint:         "wss://127.0.0.1",
+		expRemoteAddress: "127.0.0.1:443",
+		expTLSConfig:     new(tls.Config),
+	}, {
+		endpoint:         "wss://127.0.0.1:8000",
+		expRemoteAddress: "127.0.0.1:8000",
+		expTLSConfig:     new(tls.Config),
+	}, {
+		endpoint:         "http://127.0.0.1",
+		expRemoteAddress: "127.0.0.1:80",
+	}, {
+		endpoint:         "https://127.0.0.1",
+		expRemoteAddress: "127.0.0.1:443",
+		expTLSConfig:     new(tls.Config),
+	}, {
+		endpoint:         "https://127.0.0.1:8443",
+		expRemoteAddress: "127.0.0.1:8443",
+		expTLSConfig:     new(tls.Config),
+	}}
+
+	for _, c := range cases {
+		t.Log("parseURI", c.endpoint)
+
+		cl.remoteAddr = ""
+		cl.TLSConfig = nil
+		cl.Endpoint = c.endpoint
+
+		err := cl.parseURI()
+		if err != nil {
+			test.Assert(t, "error", c.expError, err.Error(), true)
+			continue
+		}
+
+		test.Assert(t, "remote address", c.expRemoteAddress, cl.remoteAddr, true)
+		test.Assert(t, "TLS config", c.expTLSConfig, cl.TLSConfig, true)
+	}
+}
+
 func TestClientPing(t *testing.T) {
 	if _testServer == nil {
 		runTestServer()