From 316b155a42a4d00fb835438ce7e0401a64e59add Mon Sep 17 00:00:00 2001
From: Joram Wilander <jwawilander@gmail.com>
Date: Mon, 31 Oct 2016 08:59:23 -0400
Subject: PLT-3562 Switch websocket over to post-connect authentication (#4327)

* Switch websocket over to post-connect authentication

* Add ability to specify token in websocket js driver, add unit tests

* Temporarily disable client websocket tests until issues are resolved

* Minor refactoring and fix status test

* Add isAuthenticated method to WebConn and minor status updates
---
 api/websocket_test.go | 112 ++++++++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 112 insertions(+)

(limited to 'api/websocket_test.go')

diff --git a/api/websocket_test.go b/api/websocket_test.go
index b7ca4b691..144c1a39b 100644
--- a/api/websocket_test.go
+++ b/api/websocket_test.go
@@ -4,12 +4,116 @@
 package api
 
 import (
+	"encoding/json"
+	"net/http"
 	"testing"
 	"time"
 
+	"github.com/gorilla/websocket"
 	"github.com/mattermost/platform/model"
 )
 
+func TestWebSocketAuthentication(t *testing.T) {
+	th := Setup().InitBasic()
+	WebSocketClient, err := th.CreateWebSocketClient()
+	if err != nil {
+		t.Fatal(err)
+	}
+	WebSocketClient.Listen()
+
+	time.Sleep(300 * time.Millisecond)
+	if resp := <-WebSocketClient.ResponseChannel; resp.Status != model.STATUS_OK {
+		t.Fatal("should have responded OK to authentication challenge")
+	}
+
+	WebSocketClient.SendMessage("ping", nil)
+	time.Sleep(300 * time.Millisecond)
+	if resp := <-WebSocketClient.ResponseChannel; resp.Data["text"].(string) != "pong" {
+		t.Fatal("wrong response")
+	}
+
+	WebSocketClient.Close()
+
+	authToken := WebSocketClient.AuthToken
+	WebSocketClient.AuthToken = "junk"
+	if err := WebSocketClient.Connect(); err != nil {
+		t.Fatal(err)
+	}
+	WebSocketClient.Listen()
+
+	if resp := <-WebSocketClient.ResponseChannel; resp != nil {
+		t.Fatal("should have closed")
+	}
+
+	WebSocketClient.Close()
+
+	if conn, _, err := websocket.DefaultDialer.Dial(WebSocketClient.ApiUrl+"/users/websocket", nil); err != nil {
+		t.Fatal("should have connected")
+	} else {
+		req := &model.WebSocketRequest{}
+		req.Seq = 1
+		req.Action = "ping"
+		conn.WriteJSON(req)
+
+		closedAutomatically := false
+		hitNotAuthedError := false
+
+		go func() {
+			time.Sleep(10 * time.Second)
+			conn.Close()
+
+			if !closedAutomatically {
+				t.Fatal("should have closed automatically in 5 seconds")
+			}
+		}()
+
+		for {
+			if _, rawMsg, err := conn.ReadMessage(); err != nil {
+				closedAutomatically = true
+				conn.Close()
+				break
+			} else {
+				var response model.WebSocketResponse
+				if err := json.Unmarshal(rawMsg, &response); err != nil && !response.IsValid() {
+					t.Fatal("should not have failed")
+				} else {
+					if response.Error == nil || response.Error.Id != "api.web_socket_router.not_authenticated.app_error" {
+						t.Log(response.Error.Id)
+						t.Fatal("wrong error")
+						continue
+					}
+
+					hitNotAuthedError = true
+				}
+			}
+		}
+
+		if !hitNotAuthedError {
+			t.Fatal("should have received a not authenticated response")
+		}
+	}
+
+	header := http.Header{}
+	header.Set(model.HEADER_AUTH, "BEARER "+authToken)
+	if conn, _, err := websocket.DefaultDialer.Dial(WebSocketClient.ApiUrl+"/users/websocket", header); err != nil {
+		t.Fatal("should have connected")
+	} else {
+		if _, rawMsg, err := conn.ReadMessage(); err != nil {
+			t.Fatal("should not have closed automatically")
+		} else {
+			var event model.WebSocketEvent
+			if err := json.Unmarshal(rawMsg, &event); err != nil && !event.IsValid() {
+				t.Fatal("should not have failed")
+			} else if event.Event != model.WEBSOCKET_EVENT_HELLO {
+				t.Log(event.ToJson())
+				t.Fatal("should have helloed")
+			}
+		}
+
+		conn.Close()
+	}
+}
+
 func TestWebSocket(t *testing.T) {
 	th := Setup().InitBasic()
 	WebSocketClient, err := th.CreateWebSocketClient()
@@ -29,6 +133,9 @@ func TestWebSocket(t *testing.T) {
 	WebSocketClient.Listen()
 
 	time.Sleep(300 * time.Millisecond)
+	if resp := <-WebSocketClient.ResponseChannel; resp.Status != model.STATUS_OK {
+		t.Fatal("should have responded OK to authentication challenge")
+	}
 
 	WebSocketClient.SendMessage("ping", nil)
 	time.Sleep(300 * time.Millisecond)
@@ -78,6 +185,11 @@ func TestWebSocketEvent(t *testing.T) {
 
 	WebSocketClient.Listen()
 
+	time.Sleep(300 * time.Millisecond)
+	if resp := <-WebSocketClient.ResponseChannel; resp.Status != model.STATUS_OK {
+		t.Fatal("should have responded OK to authentication challenge")
+	}
+
 	omitUser := make(map[string]bool, 1)
 	omitUser["somerandomid"] = true
 	evt1 := model.NewWebSocketEvent(model.WEBSOCKET_EVENT_TYPING, "", th.BasicChannel.Id, "", omitUser)
-- 
cgit v1.2.3-1-g7c22