From 47b4beb173f72748cef72a2f4e1156b875521dff Mon Sep 17 00:00:00 2001
From: byung2 <heebyung@gmail.com>
Date: Tue, 10 Apr 2018 21:05:19 +0900
Subject: [PATCH] Add --ws-origin option to allow cross origin requests to WS
 endpoint

---
 cmd/gotty-client/main.go | 11 ++++++++++-
 gotty-client.go          |  4 ++++
 2 files changed, 14 insertions(+), 1 deletion(-)

diff --git a/cmd/gotty-client/main.go b/cmd/gotty-client/main.go
index 08f471a..3763bd3 100644
--- a/cmd/gotty-client/main.go
+++ b/cmd/gotty-client/main.go
@@ -27,7 +27,7 @@ func main() {
 	app.BashComplete = func(c *cli.Context) {
 		for _, command := range []string{
 			"--debug", "--skip-tls-verify", "--use-proxy-from-env",
-			"--v2", "--detach-keys", "--help",
+			"--v2", "--detach-keys", "--ws-origin", "--help",
 			"--generate-bash-completion", "--version",
 			"http://user:pass@host:1234/path/\\\\?arg=abcdef\\\\&arg=ghijkl",
 			"https://user:pass@host:1234/path/\\\\?arg=abcdef\\\\&arg=ghijkl",
@@ -63,6 +63,11 @@ func main() {
 			Usage:  "For Gotty 2.0",
 			EnvVar: "GOTTY_CLIENT_GOTTY2",
 		},
+		cli.StringFlag{
+			Name:   "ws-origin, w",
+			Usage:  "WebSocket Origin URL",
+			EnvVar: "GOTTY_CLIENT_WS_ORIGIN",
+		},
 	}
 
 	app.Action = action
@@ -102,6 +107,10 @@ func action(c *cli.Context) error {
 		client.V2 = true
 	}
 
+	if wsOrigin := c.String("ws-origin"); wsOrigin != "" {
+		client.WSOrigin = wsOrigin
+	}
+
 	if detachKey := c.String("detach-keys"); detachKey != "" {
 		escapeKeys, err := term.ToBytes(detachKey)
 		if err != nil {
diff --git a/gotty-client.go b/gotty-client.go
index 89e9188..4e96b3b 100644
--- a/gotty-client.go
+++ b/gotty-client.go
@@ -134,6 +134,7 @@ type Client struct {
 	EscapeKeys      []byte
 	V2              bool
 	message         *gottyMessageType
+	WSOrigin        string
 }
 
 type querySingleType struct {
@@ -207,6 +208,9 @@ func (c *Client) Connect() error {
 	if err != nil {
 		return err
 	}
+	if c.WSOrigin != "" {
+		header.Add("Origin", c.WSOrigin)
+	}
 	logrus.Debugf("Connecting to websocket: %q", target.String())
 	if c.SkipTLSVerify {
 		c.Dialer.TLSClientConfig = &tls.Config{InsecureSkipVerify: true}