Fix default port handling (#178)

at-wat · Jul 12, 2021 · d80a538 · d80a538
1 parent c43cdb1
commit d80a538
Show file tree

Hide file tree

Showing 4 changed files with 96 additions and 47 deletions.
diff --git a/client_integration_test.go b/client_integration_test.go
@@ -67,25 +67,37 @@ func ExampleClient() {
 }
 
 func TestIntegration_Connect(t *testing.T) {
-	for name, url := range urls {
-		t.Run(name, func(t *testing.T) {
-			ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
-			defer cancel()
-
-			cli, err := DialContext(ctx, url, WithTLSConfig(&tls.Config{InsecureSkipVerify: true}))
-			if err != nil {
-				t.Fatalf("Unexpected error: '%v'", err)
-			}
+	// Overwrite default port to avoid using privileged port during test.
+	defaultPorts["ws"] = 9001
+	defaultPorts["wss"] = 9443
+
+	test := func(t *testing.T, urls map[string]string) {
+		for name, url := range urls {
+			t.Run(name, func(t *testing.T) {
+				ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
+				defer cancel()
+
+				cli, err := DialContext(ctx, url, WithTLSConfig(&tls.Config{InsecureSkipVerify: true}))
+				if err != nil {
+					t.Fatalf("Unexpected error: '%v'", err)
+				}
 
-			if _, err := cli.Connect(ctx, "Client"); err != nil {
-				t.Fatalf("Unexpected error: '%v'", err)
-			}
+				if _, err := cli.Connect(ctx, "Client"); err != nil {
+					t.Fatalf("Unexpected error: '%v'", err)
+				}
 
-			if err := cli.Disconnect(ctx); err != nil {
-				t.Fatalf("Unexpected error: '%v'", err)
-			}
-		})
+				if err := cli.Disconnect(ctx); err != nil {
+					t.Fatalf("Unexpected error: '%v'", err)
+				}
+			})
+		}
 	}
+	t.Run("WithPort", func(t *testing.T) {
+		test(t, urls)
+	})
+	t.Run("WithoutPort", func(t *testing.T) {
+		test(t, urlsWithoutPort)
+	})
 }
 
 func TestIntegration_Publish(t *testing.T) {

diff --git a/conn_integration_test.go b/conn_integration_test.go
@@ -23,6 +23,10 @@ import (
 )
 
 func TestIntegration_WithTLSCertFiles(t *testing.T) {
+	// Overwrite default port to avoid using privileged port during test.
+	defaultPorts["ws"] = 9001
+	defaultPorts["wss"] = 9443
+
 	cases := map[string]struct {
 		opt         DialOption
 		expectError bool
@@ -52,31 +56,39 @@ func TestIntegration_WithTLSCertFiles(t *testing.T) {
 			false,
 		},
 	}
-	for name, c := range cases {
-		t.Run(name, func(t *testing.T) {
-			ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
-			defer cancel()
+	test := func(t *testing.T, urls map[string]string) {
+		for name, c := range cases {
+			t.Run(name, func(t *testing.T) {
+				ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
+				defer cancel()
 
-			cli, err := DialContext(ctx, urls["MQTTs"], c.opt)
+				cli, err := DialContext(ctx, urls["MQTTs"], c.opt)
 
-			if err != nil {
-				if c.expectError {
-					return
+				if err != nil {
+					if c.expectError {
+						return
+					}
+					t.Fatal(err)
 				}
-				t.Fatal(err)
-			}
-			defer cli.Close()
+				defer cli.Close()
 
-			if c.expectError {
-				t.Fatal("Expected error but succeeded")
-			}
+				if c.expectError {
+					t.Fatal("Expected error but succeeded")
+				}
 
-			if _, err := cli.Connect(ctx, "TestConnTLS", WithCleanSession(true)); err != nil {
-				t.Error(err)
-			}
-			if err := cli.Disconnect(ctx); err != nil {
-				t.Error(err)
-			}
-		})
+				if _, err := cli.Connect(ctx, "TestConnTLS", WithCleanSession(true)); err != nil {
+					t.Error(err)
+				}
+				if err := cli.Disconnect(ctx); err != nil {
+					t.Error(err)
+				}
+			})
+		}
 	}
+	t.Run("WithPort", func(t *testing.T) {
+		test(t, urls)
+	})
+	t.Run("WithoutPort", func(t *testing.T) {
+		test(t, urlsWithoutPort)
+	})
 }
diff --git a/dialer.go b/dialer.go
@@ -28,6 +28,13 @@ import (
 	"golang.org/x/net/websocket"
 )
 
+var defaultPorts = map[string]uint16{
+	"mqtt":  1883,
+	"mqtts": 8883,
+	"wss":   443,
+	"ws":    80,
+}
+
 // ErrUnsupportedProtocol means that the specified scheme in the URL is not supported.
 var ErrUnsupportedProtocol = errors.New("unsupported protocol")
 
@@ -76,15 +83,25 @@ func (d *URLDialer) DialContext(ctx context.Context) (*BaseClient, error) {
 
 // DialContext creates MQTT client using URL string.
 func DialContext(ctx context.Context, urlStr string, opts ...DialOption) (*BaseClient, error) {
+	u, err := url.Parse(urlStr)
+	if err != nil {
+		return nil, err
+	}
 	o := &DialOptions{
 		Dialer: &net.Dialer{},
 	}
+	switch u.Scheme {
+	case "tls", "ssl", "mqtts", "wss":
+		o.TLSConfig = &tls.Config{
+			ServerName: u.Hostname(),
+		}
+	}
 	for _, opt := range opts {
 		if err := opt(o); err != nil {
 			return nil, err
 		}
 	}
-	return o.dial(ctx, urlStr)
+	return o.dial(ctx, u)
 }
 
 // DialOption sets option for Dial.
@@ -155,25 +172,27 @@ func WithConnStateHandler(handler func(ConnState, error)) DialOption {
 	}
 }
 
-func (d *DialOptions) dial(ctx context.Context, urlStr string) (*BaseClient, error) {
+func (d *DialOptions) dial(ctx context.Context, u *url.URL) (*BaseClient, error) {
 	c := &BaseClient{
 		ConnState:     d.ConnState,
 		MaxPayloadLen: d.MaxPayloadLen,
 	}
 
-	u, err := url.Parse(urlStr)
-	if err != nil {
-		return nil, err
-	}
 	switch u.Scheme {
 	case "tcp", "mqtt", "tls", "ssl", "mqtts", "wss", "ws":
 	default:
 		return nil, wrapErrorf(ErrUnsupportedProtocol, "protocol %s", u.Scheme)
 	}
+	hostWithPort := u.Host
+	if u.Port() == "" {
+		if port, ok := defaultPorts[u.Scheme]; ok {
+			hostWithPort += fmt.Sprintf(":%d", port)
+		}
+	}
 
-	baseConn, err := d.Dialer.DialContext(ctx, "tcp", u.Host)
+	baseConn, err := d.Dialer.DialContext(ctx, "tcp", hostWithPort)
 	if err != nil {
-		return nil, err
+		return nil, wrapError(err, "dialing tcp")
 	}
 	switch u.Scheme {
 	case "tcp", "mqtt":
@@ -186,13 +205,13 @@ func (d *DialOptions) dial(ctx context.Context, urlStr string) (*BaseClient, err
 	case "ws":
 		wsc, err := websocket.NewConfig(u.String(), fmt.Sprintf("https://%s", u.Host))
 		if err != nil {
-			return nil, err
+			return nil, wrapError(err, "configuring websocket")
 		}
 		wsc.Protocol = append(wsc.Protocol, "mqtt")
 		wsc.TlsConfig = d.TLSConfig
 		ws, err := websocket.NewClient(wsc, baseConn)
 		if err != nil {
-			return nil, err
+			return nil, wrapError(err, "dialing websocket")
 		}
 		ws.PayloadType = websocket.BinaryFrame
 		c.Transport = ws

diff --git a/dialer_test.go b/dialer_test.go
@@ -31,6 +31,12 @@ var (
 		"WebSocket":  "ws://localhost:9001",
 		"WebSockets": "wss://localhost:9443",
 	}
+	urlsWithoutPort = map[string]string{
+		"MQTT":       "mqtt://localhost",
+		"MQTTs":      "mqtts://localhost",
+		"WebSocket":  "ws://localhost",
+		"WebSockets": "wss://localhost",
+	}
 )
 
 func TestDialOptionError(t *testing.T) {