Add option for setting additional HTTP headers in WebSocket handshake

Certain services (such as AWS IOT when using a custom authorizer) require
additional headers to be sent during the WebSocket opening handshake.  This
adds a function SetHTTPHeaders() to the ClientOptions that enables setting
additional headers.

Signed-off-by: Scott Talbert <[email protected]>

client.go

@@ -212,7 +212,7 @@ func (c *client) Connect() Token {
 			c.options.ProtocolVersion = protocolVersion
 		CONN:
 			DEBUG.Println(CLI, "about to write new connect msg")
-			c.conn, err = openConnection(broker, &c.options.TLSConfig, c.options.ConnectTimeout)
+			c.conn, err = openConnection(broker, &c.options.TLSConfig, c.options.ConnectTimeout, c.options.HTTPHeaders)
 			if err == nil {
 				DEBUG.Println(CLI, "socket connected to broker")
 				switch c.options.ProtocolVersion {
@@ -320,7 +320,7 @@ func (c *client) reconnect() {
 
 		for _, broker := range c.options.Servers {
 			DEBUG.Println(CLI, "about to write new connect msg")
-			c.conn, err = openConnection(broker, &c.options.TLSConfig, c.options.ConnectTimeout)
+			c.conn, err = openConnection(broker, &c.options.TLSConfig, c.options.ConnectTimeout, c.options.HTTPHeaders)
 			if err == nil {
 				DEBUG.Println(CLI, "socket connected to broker")
 				switch c.options.ProtocolVersion {

net.go

@@ -19,6 +19,7 @@ import (
 	"errors"
 	"fmt"
 	"net"
+	"net/http"
 	"net/url"
 	"os"
 	"reflect"
@@ -37,7 +38,7 @@ func signalError(c chan<- error, err error) {
 	}
 }
 
-func openConnection(uri *url.URL, tlsc *tls.Config, timeout time.Duration) (net.Conn, error) {
+func openConnection(uri *url.URL, tlsc *tls.Config, timeout time.Duration, headers http.Header) (net.Conn, error) {
 	switch uri.Scheme {
 	case "ws":
 		conn, err := websocket.Dial(uri.String(), "mqtt", fmt.Sprintf("http://%s", uri.Host))
@@ -50,6 +51,7 @@ func openConnection(uri *url.URL, tlsc *tls.Config, timeout time.Duration) (net.
 		config, _ := websocket.NewConfig(uri.String(), fmt.Sprintf("https://%s", uri.Host))
 		config.Protocol = []string{"mqtt"}
 		config.TlsConfig = tlsc
+		config.Header = headers
 		conn, err := websocket.DialConfig(config)
 		if err != nil {
 			return nil, err

options.go

@@ -16,6 +16,7 @@ package mqtt
 
 import (
 	"crypto/tls"
+	"net/http"
 	"net/url"
 	"strings"
 	"time"
@@ -69,6 +70,7 @@ type ClientOptions struct {
 	OnConnectionLost        ConnectionLostHandler
 	WriteTimeout            time.Duration
 	MessageChannelDepth     uint
+	HTTPHeaders             http.Header
 }
 
 // NewClientOptions will create a new ClientClientOptions type with some
@@ -106,6 +108,7 @@ func NewClientOptions() *ClientOptions {
 		OnConnectionLost:        DefaultConnectionLostHandler,
 		WriteTimeout:            0, // 0 represents timeout disabled
 		MessageChannelDepth:     100,
+		HTTPHeaders:             make(map[string][]string),
 	}
 	return o
 }
@@ -318,3 +321,10 @@ func (o *ClientOptions) SetMessageChannelDepth(s uint) *ClientOptions {
 	o.MessageChannelDepth = s
 	return o
 }
+
+// SetHTTPHeaders sets the additional HTTP headers that will be sent in the WebSocket
+// opening handshake.
+func (o *ClientOptions) SetHTTPHeaders(h http.Header) *ClientOptions {
+	o.HTTPHeaders = h
+	return o
+}

options_reader.go

@@ -16,6 +16,7 @@ package mqtt
 
 import (
 	"crypto/tls"
+	"net/http"
 	"net/url"
 	"time"
 )
@@ -135,3 +136,8 @@ func (r *ClientOptionsReader) MessageChannelDepth() uint {
 	s := r.options.MessageChannelDepth
 	return s
 }
+
+func (r *ClientOptionsReader) HTTPHeaders() http.Header {
+	h := r.options.HTTPHeaders
+	return h
+}