Bumped GoTTY protocol for GoTTY@0.10

gotty-client.go

@@ -5,12 +5,15 @@ import (
 	"encoding/base64"
 	"encoding/json"
 	"fmt"
+	"io/ioutil"
 	"net/http"
 	"net/url"
 	"os"
 	"os/signal"
+	"regexp"
 	"strings"
 	"syscall"
+	"time"
 	"unicode/utf8"
 	"unsafe"
 
@@ -19,6 +22,24 @@ import (
 	"golang.org/x/crypto/ssh/terminal"
 )
 
+// GetAuthTokenURL transforms a GoTTY http URL to its AuthToken file URL
+func GetAuthTokenURL(httpURL string) (*url.URL, *http.Header, error) {
+	header := http.Header{}
+	target, err := url.Parse(httpURL)
+	if err != nil {
+		return nil, nil, err
+	}
+
+	target.Path = strings.TrimLeft(target.Path+"auth_token.js", "/")
+
+	if target.User != nil {
+		header.Add("Authorization", "Basic "+base64.StdEncoding.EncodeToString([]byte(target.User.String())))
+		target.User = nil
+	}
+
+	return target, &header, nil
+}
+
 // GetWebsocketURL transforms a GoTTY http URL to its WebSocket URL
 func GetWebsocketURL(httpURL string) (*url.URL, *http.Header, error) {
 	header := http.Header{}
@@ -46,21 +67,72 @@ func GetWebsocketURL(httpURL string) (*url.URL, *http.Header, error) {
 type Client struct {
 	Dialer    *websocket.Dialer
 	Conn      *websocket.Conn
-	Headers   *http.Header
-	Target    string
+	URL       string
 	Connected bool
 }
 
+// GetAuthToken retrieves an Auth Token from dynamic auth_token.js file
+func (c *Client) GetAuthToken() (string, error) {
+	target, header, err := GetAuthTokenURL(c.URL)
+	if err != nil {
+		return "", err
+	}
+
+	req, err := http.NewRequest("GET", target.String(), nil)
+	req.Header = *header
+	client := http.Client{}
+	resp, err := client.Do(req)
+	if err != nil {
+		return "", err
+	}
+
+	defer resp.Body.Close()
+	body, err := ioutil.ReadAll(resp.Body)
+	if err != nil {
+		return "", err
+	}
+
+	re := regexp.MustCompile("var gotty_auth_token = '(.*)'")
+	return re.FindStringSubmatch(string(body))[1], nil
+}
+
 // Connect tries to dial a websocket server
 func (c *Client) Connect() error {
-	conn, _, err := c.Dialer.Dial(c.Target, *c.Headers)
+	// Retrieve AuthToken
+	authToken, err := c.GetAuthToken()
+	if err != nil {
+		return err
+	}
+
+	// Open WebSocket connection
+	target, header, err := GetWebsocketURL(c.URL)
+	if err != nil {
+		return err
+	}
+	conn, _, err := c.Dialer.Dial(target.String(), *header)
 	if err != nil {
 		return err
 	}
 	c.Conn = conn
+
+	// Send AuthToken
+	err = c.Conn.WriteMessage(websocket.TextMessage, []byte(authToken))
+	if err != nil {
+		return err
+	}
+
+	go c.pingLoop()
+
 	return nil
 }
 
+func (c *Client) pingLoop() {
+	for {
+		c.Conn.WriteMessage(websocket.TextMessage, []byte("1"))
+		time.Sleep(30 * time.Second)
+	}
+}
+
 // Close will nicely close the dialer
 func (c *Client) Close() {
 	c.Conn.Close()
@@ -106,7 +178,7 @@ func (c *Client) termsizeLoop(done chan bool) {
 			logrus.Warnf("json.Marshal error: %v", err)
 		}
 
-		err = c.Conn.WriteMessage(websocket.TextMessage, append([]byte("1"), b...))
+		err = c.Conn.WriteMessage(websocket.TextMessage, append([]byte("2"), b...))
 		if err != nil {
 			logrus.Warnf("ws.WriteMessage failed: %v", err)
 		}
@@ -151,12 +223,19 @@ func (c *Client) readLoop(done chan bool) {
 
 		switch data[0] {
 		case '0': // data
-			fmt.Print(string(data[1:]))
-		case '1': // new title
+			buf, err := base64.StdEncoding.DecodeString(string(data[1:]))
+			if err != nil {
+				logrus.Warnf("Invalid base64 content: %q", data[1:])
+			}
+			fmt.Print(string(buf))
+		case '1': // pong
+		case '2': // new title
 			newTitle := string(data[1:])
 			fmt.Printf("\033]0;%s\007", newTitle)
-		case '2': // json prefs
+		case '3': // json prefs
 			logrus.Debugf("Unhandled protocol message: json pref: %s", string(data))
+		case '4': // autoreconnect
+			logrus.Debugf("Unhandled protocol message: autoreconnect: %s", string(data))
 		default:
 			logrus.Warnf("Unhandled protocol message: %s", string(data))
 		}
@@ -165,14 +244,8 @@ func (c *Client) readLoop(done chan bool) {
 
 // NewClient returns a GoTTY client object
 func NewClient(httpURL string) (*Client, error) {
-	target, header, err := GetWebsocketURL(httpURL)
-	if err != nil {
-		return nil, err
-	}
-
 	return &Client{
-		Dialer:  &websocket.Dialer{},
-		Target:  target.String(),
-		Headers: header,
+		Dialer: &websocket.Dialer{},
+		URL:    httpURL,
 	}, nil
 }