Use websocket for log streaming

client/client.go

@@ -6,12 +6,13 @@ import (
 	"encoding/json"
 	"fmt"
 	"io"
+	"log"
 	"net/http"
 	"net/url"
 	"path"
-	"strconv"
 	"strings"
 
+	"github.com/gorilla/websocket"
 	"github.com/ubclaunchpad/inertia/common"
 )
 
@@ -177,15 +178,31 @@ func (c *Client) Reset() (*http.Response, error) {
 }
 
 // Logs get logs of given container
-func (c *Client) Logs(stream bool, container string) (*http.Response, error) {
+func (c *Client) Logs(container string) (*http.Response, error) {
 	reqContent := map[string]string{
-		common.Stream:    strconv.FormatBool(stream),
 		common.Container: container,
 	}
 
 	return c.get("/logs", reqContent)
 }
 
+// LogsWebSocket opens a websocket connection to given container's logs
+func (c *Client) LogsWebSocket(container string) (SocketReader, error) {
+	host, err := url.Parse("https://" + c.RemoteVPS.GetIPAndPort())
+	if err != nil {
+		return nil, err
+	}
+	url := &url.URL{Scheme: "ws", Host: host.Host, Path: "/logs"}
+
+	header := http.Header{}
+	header.Set("Authorization", "Bearer "+c.Daemon.Token)
+	socket, _, err := websocket.DefaultDialer.Dial(url.String(), header)
+	if err != nil {
+		log.Fatal("dial:", err)
+	}
+	return socket, nil
+}
+
 // AddUser adds an authorized user for access to Inertia Web
 func (c *Client) AddUser(username, password string, admin bool) (*http.Response, error) {
 	reqContent := &common.UserRequest{
@@ -222,11 +239,7 @@ func (c *Client) get(endpoint string, queries map[string]string) (*http.Response
 
 	// Add query strings
 	if queries != nil {
-		q := req.URL.Query()
-		for k, v := range queries {
-			q.Add(k, v)
-		}
-		req.URL.RawQuery = q.Encode()
+		encodeQuery(req.URL, queries)
 	}
 
 	client := buildHTTPSClient()
@@ -289,3 +302,11 @@ func buildHTTPSClient() *http.Client {
 	}
 	return &http.Client{Transport: tr}
 }
+
+func encodeQuery(url *url.URL, queries map[string]string) {
+	q := url.Query()
+	for k, v := range queries {
+		q.Add(k, v)
+	}
+	url.RawQuery = q.Encode()
+}

client/client_test.go

@@ -323,15 +323,14 @@ func TestLogs(t *testing.T) {
 		defer req.Body.Close()
 		q := req.URL.Query()
 		assert.Equal(t, "docker-compose", q.Get(common.Container))
-		assert.Equal(t, "true", q.Get(common.Stream))
 
 		// Check auth
 		assert.Equal(t, "Bearer "+fakeAuth, req.Header.Get("Authorization"))
 	}))
 	defer testServer.Close()
 
 	d := getMockClient(testServer)
-	resp, err := d.Logs(true, "docker-compose")
+	resp, err := d.Logs("docker-compose")
 	assert.Nil(t, err)
 	assert.Equal(t, http.StatusOK, resp.StatusCode)
 }

client/socket.go

@@ -0,0 +1,7 @@
+package client
+
+// SocketReader is an interface to a websocket connection
+type SocketReader interface {
+	ReadMessage() (messageType int, p []byte, err error)
+	Close() error
+}

deploy.go

@@ -189,13 +189,13 @@ var deploymentLogsCmd = &cobra.Command{
 			container = args[0]
 		}
 
-		resp, err := deployment.Logs(stream, container)
-		if err != nil {
-			log.Fatal(err)
-		}
-		defer resp.Body.Close()
-
 		if !stream {
+			resp, err := deployment.Logs(container)
+			if err != nil {
+				log.Fatal(err)
+			}
+			defer resp.Body.Close()
+
 			body, err := ioutil.ReadAll(resp.Body)
 			if err != nil {
 				log.Fatal(err)
@@ -212,11 +212,16 @@ var deploymentLogsCmd = &cobra.Command{
 					resp.StatusCode, body)
 			}
 		} else {
-			reader := bufio.NewReader(resp.Body)
+			socket, err := deployment.LogsWebSocket(container)
+			if err != nil {
+				log.Fatal(err)
+			}
+			defer socket.Close()
+
 			for {
-				line, err := reader.ReadBytes('\n')
+				_, line, err := socket.ReadMessage()
 				if err != nil {
-					break
+					log.Fatal(err)
 				}
 				fmt.Print(string(line))
 			}