From a1da47fb7cef608b17622881c07ff62a5b1f44ee Mon Sep 17 00:00:00 2001
From: Gavin Liu <gavinpub@foxmail.com>
Date: Thu, 14 Feb 2019 12:54:09 +0800
Subject: [PATCH] Add flag to specify sec-websocket-protocal

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

diff --git a/connection.go b/connection.go
index 9a6265e..fad5fbc 100644
--- a/connection.go
+++ b/connection.go
@@ -18,15 +18,16 @@ type session struct {
 	errChan chan error
 }
 
-func connect(url, origin string, rlConf *readline.Config, allowInsecure bool) error {
+func connect(url, origin string, subProtocals string, rlConf *readline.Config, allowInsecure bool) error {
 	headers := make(http.Header)
 	headers.Add("Origin", origin)
 
 	dialer := websocket.Dialer{
 		Proxy: http.ProxyFromEnvironment,
-		TLSClientConfig:&tls.Config{
+		TLSClientConfig: &tls.Config{
 			InsecureSkipVerify: allowInsecure,
 		},
+		Subprotocols: []string{subProtocals},
 	}
 	ws, _, err := dialer.Dial(url, headers)
 	if err != nil {
diff --git a/main.go b/main.go
index 7fdc26b..ef1083b 100644
--- a/main.go
+++ b/main.go
@@ -17,7 +17,8 @@ const Version = "0.2.1"
 var options struct {
 	origin       string
 	printVersion bool
-	insecure bool
+	insecure     bool
+	subProtocals string
 }
 
 func main() {
@@ -29,6 +30,7 @@ func main() {
 	rootCmd.Flags().StringVarP(&options.origin, "origin", "o", "", "websocket origin")
 	rootCmd.Flags().BoolVarP(&options.printVersion, "version", "v", false, "print version")
 	rootCmd.Flags().BoolVarP(&options.insecure, "insecure", "k", false, "skip ssl certificate check")
+	rootCmd.Flags().StringVarP(&options.subProtocals, "subprotocal", "s", "", "sec-websocket-protocal field")
 
 	rootCmd.Execute()
 }
@@ -69,7 +71,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.subProtocals, &readline.Config{
 		Prompt:      "> ",
 		HistoryFile: historyFile,
 	}, options.insecure)