From 88ffbe5e88a47711b649fdaec0614e2a4ef0bb52 Mon Sep 17 00:00:00 2001
From: Vladimir Shakhov <vladimir@spotify.com>
Date: Wed, 25 Apr 2018 16:12:30 +0200
Subject: [PATCH] Ability to specify auth header

---
 connection.go | 5 ++++-
 main.go       | 4 +++-
 2 files changed, 7 insertions(+), 2 deletions(-)

diff --git a/connection.go b/connection.go
index 529d17e..80d6c79 100644
--- a/connection.go
+++ b/connection.go
@@ -17,9 +17,12 @@ type session struct {
 	errChan chan error
 }
 
-func connect(url, origin string, rlConf *readline.Config) error {
+func connect(url, origin string, authHeader string, rlConf *readline.Config) error {
 	headers := make(http.Header)
 	headers.Add("Origin", origin)
+	if authHeader != "" {
+		headers.Add("Authorization", authHeader)
+	}
 
 	ws, _, err := websocket.DefaultDialer.Dial(url, headers)
 	if err != nil {
diff --git a/main.go b/main.go
index a12c446..f737a80 100644
--- a/main.go
+++ b/main.go
@@ -17,6 +17,7 @@ const Version = "0.2.1"
 var options struct {
 	origin       string
 	printVersion bool
+	authHeader   string
 }
 
 func main() {
@@ -27,6 +28,7 @@ func main() {
 	}
 	rootCmd.Flags().StringVarP(&options.origin, "origin", "o", "", "websocket origin")
 	rootCmd.Flags().BoolVarP(&options.printVersion, "version", "v", false, "print version")
+	rootCmd.Flags().StringVarP(&options.authHeader, "auth_header", "a", "", "auth header, like 'Bearer $TOKEN'")
 
 	rootCmd.Execute()
 }
@@ -67,7 +69,7 @@ func root(cmd *cobra.Command, args []string) {
 		historyFile = filepath.Join(user.HomeDir, ".ws_history")
 	}
 
-	err = connect(dest.String(), origin, &readline.Config{
+	err = connect(dest.String(), origin, options.authHeader, &readline.Config{
 		Prompt:      "> ",
 		HistoryFile: historyFile,
 	})