From 60a56da7b91a60838627a21582ae17b77e8a9d16 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=D0=9A=D0=B8=D1=80=D0=B8=D0=BB=D0=BB=20=D0=A1=D1=8B=D1=81?=
 =?UTF-8?q?=D0=BE=D0=B5=D0=B2?= <k.sysoev@me.com>
Date: Mon, 30 Oct 2023 20:26:39 +0800
Subject: [PATCH] Adds possibility to add custom headers to handshake http
 request

---
 cmd/wsget/main.go |  4 +++-
 pkg/ws/ws.go      | 22 ++++++++++++++++++++++
 2 files changed, 25 insertions(+), 1 deletion(-)

diff --git a/cmd/wsget/main.go b/cmd/wsget/main.go
index 7db3eb4..cb1fb84 100644
--- a/cmd/wsget/main.go
+++ b/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
diff --git a/pkg/ws/ws.go b/pkg/ws/ws.go
index b0ff2a1..a0e932b 100644
--- a/pkg/ws/ws.go
+++ b/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 {