Adds possibility to add custom headers to handshake http request

cmd/wsget/main.go

@@ -13,6 +13,7 @@ import (
 var insecure bool
 var request string
 var outputFile string
+var headers []string
 
 const (
 	LongDescription = `A command-line tool for interacting with WebSocket servers.
@@ -51,6 +52,7 @@ func main() {
 	cmd.Flags().BoolVarP(&insecure, "insecure", "k", false, "Skip SSL certificate verification")
 	cmd.Flags().StringVarP(&request, "request", "r", "", "WebSocket request that will be sent to the server")
 	cmd.Flags().StringVarP(&outputFile, "output", "o", "", "Output file for saving all request and responses")
+	cmd.Flags().StringSliceVarP(&headers, "header", "H", []string{}, "HTTP headers to attach to the request")
 
 	err := cmd.Execute()
 	if err != nil {
@@ -66,7 +68,7 @@ func run(cmd *cobra.Command, args []string) {
 		return
 	}
 
-	wsConn, err := ws.NewWS(wsURL, ws.Options{SkipSSLVerification: insecure})
+	wsConn, err := ws.NewWS(wsURL, ws.Options{SkipSSLVerification: insecure, Headers: headers})
 	if err != nil {
 		color.New(color.FgRed).Println("Unable to connect to the server: ", err)
 		return

pkg/ws/ws.go

@@ -2,6 +2,9 @@ package ws
 
 import (
 	"crypto/tls"
+	"fmt"
+	"net/http"
+	"strings"
 
 	"github.com/fatih/color"
 	"golang.org/x/net/websocket"
@@ -28,6 +31,7 @@ func (mt MessageType) String() string {
 
 const (
 	WSMessageBufferSize = 100
+	HeaderPartsNumber   = 2
 )
 
 type Message struct {
@@ -41,6 +45,7 @@ type Connection struct {
 }
 
 type Options struct {
+	Headers             []string
 	SkipSSLVerification bool
 }
 
@@ -58,6 +63,23 @@ func NewWS(url string, opts Options) (*Connection, error) {
 	}
 	cfg.TlsConfig = tlsConfig
 
+	if len(opts.Headers) > 0 {
+		Headers := make(http.Header)
+		for _, headerInput := range opts.Headers {
+			splited := strings.Split(headerInput, ":")
+			if len(splited) != HeaderPartsNumber {
+				return nil, fmt.Errorf("invalid header: %s", headerInput)
+			}
+
+			header := strings.TrimSpace(splited[0])
+			value := strings.TrimSpace(splited[1])
+
+			Headers.Add(header, value)
+		}
+
+		cfg.Header = Headers
+	}
+
 	ws, err := websocket.DialConfig(cfg)
 
 	if err != nil {