4 files changed, 129 insertions, 1 deletions
diff --git a/api4/api.go b/api4/api.go
index c967537ee..dffed60e4 100644
--- a/api4/api.go
+++ b/api4/api.go
@@ -92,7 +92,6 @@ var BaseRoutes *Routes
 func InitRouter() {
 	app.Srv.Router = mux.NewRouter()
 	app.Srv.Router.NotFoundHandler = http.HandlerFunc(Handle404)
-	app.Srv.WebSocketRouter = app.NewWebSocketRouter()
 }
 
 func InitApi(full bool) {
@@ -174,6 +173,7 @@ func InitApi(full bool) {
 	InitBrand()
 	InitCommand()
 	InitStatus()
+	InitWebSocket()
 
 	app.Srv.Router.Handle("/api/v4/{anything:.*}", http.HandlerFunc(Handle404))
 
diff --git a/api4/apitestlib.go b/api4/apitestlib.go
index c6c1dfa94..87a3976f5 100644
--- a/api4/apitestlib.go
+++ b/api4/apitestlib.go
@@ -20,6 +20,7 @@ import (
 	"github.com/mattermost/platform/model"
 	"github.com/mattermost/platform/store"
 	"github.com/mattermost/platform/utils"
+	"github.com/mattermost/platform/wsapi"
 
 	s3 "github.com/minio/minio-go"
 )
@@ -55,9 +56,11 @@ func SetupEnterprise() *TestHelper {
 		app.NewServer()
 		app.InitStores()
 		InitRouter()
+		wsapi.InitRouter()
 		app.StartServer()
 		utils.InitHTML()
 		InitApi(true)
+		wsapi.InitApi()
 		utils.EnableDebugLogForTest()
 		app.Srv.Store.MarkSystemRanUnitTests()
 
@@ -85,8 +88,10 @@ func Setup() *TestHelper {
 		app.NewServer()
 		app.InitStores()
 		InitRouter()
+		wsapi.InitRouter()
 		app.StartServer()
 		InitApi(true)
+		wsapi.InitApi()
 		utils.EnableDebugLogForTest()
 		app.Srv.Store.MarkSystemRanUnitTests()
 
@@ -167,6 +172,10 @@ func (me *TestHelper) CreateClient() *model.Client4 {
 	return model.NewAPIv4Client("http://localhost" + utils.Cfg.ServiceSettings.ListenAddress)
 }
 
+func (me *TestHelper) CreateWebSocketClient() (*model.WebSocketClient, *model.AppError) {
+	return model.NewWebSocketClient4("ws://localhost"+utils.Cfg.ServiceSettings.ListenAddress, me.Client.AuthToken)
+}
+
 func (me *TestHelper) CreateUser() *model.User {
 	return me.CreateUserWithClient(me.Client)
 }
diff --git a/api4/websocket.go b/api4/websocket.go
new file mode 100644
index 000000000..c70327222
--- /dev/null
+++ b/api4/websocket.go
@@ -0,0 +1,46 @@
+// Copyright (c) 2017 Mattermost, Inc. All Rights Reserved.
+// See License.txt for license information.
+
+package api4
+
+import (
+	"net/http"
+
+	l4g "github.com/alecthomas/log4go"
+	"github.com/gorilla/websocket"
+	"github.com/mattermost/platform/app"
+	"github.com/mattermost/platform/model"
+	"github.com/mattermost/platform/utils"
+)
+
+func InitWebSocket() {
+	l4g.Debug(utils.T("api.web_socket.init.debug"))
+
+	BaseRoutes.ApiRoot.Handle("/websocket", ApiHandlerTrustRequester(connectWebSocket)).Methods("GET")
+}
+
+func connectWebSocket(c *Context, w http.ResponseWriter, r *http.Request) {
+	originChecker := utils.GetOriginChecker(r)
+
+	upgrader := websocket.Upgrader{
+		ReadBufferSize:  model.SOCKET_MAX_MESSAGE_SIZE_KB,
+		WriteBufferSize: model.SOCKET_MAX_MESSAGE_SIZE_KB,
+		CheckOrigin:     originChecker,
+	}
+
+	ws, err := upgrader.Upgrade(w, r, nil)
+	if err != nil {
+		l4g.Error(utils.T("api.web_socket.connect.error"), err)
+		c.Err = model.NewLocAppError("connect", "api.web_socket.connect.upgrade.app_error", nil, "")
+		return
+	}
+
+	wc := app.NewWebConn(ws, c.Session, c.T, "")
+
+	if len(c.Session.UserId) > 0 {
+		app.HubRegister(wc)
+	}
+
+	go wc.WritePump()
+	wc.ReadPump()
+}
diff --git a/api4/websocket_test.go b/api4/websocket_test.go
new file mode 100644
index 000000000..6018bf7da
--- /dev/null
+++ b/api4/websocket_test.go
@@ -0,0 +1,73 @@
+// Copyright (c) 2017 Mattermost, Inc. All Rights Reserved.
+// See License.txt for license information.
+
+package api4
+
+import (
+	"testing"
+	"time"
+
+	"github.com/mattermost/platform/model"
+)
+
+func TestWebSocket(t *testing.T) {
+	th := Setup().InitBasic()
+	defer TearDown()
+	WebSocketClient, err := th.CreateWebSocketClient()
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer WebSocketClient.Close()
+
+	time.Sleep(300 * time.Millisecond)
+
+	// Test closing and reconnecting
+	WebSocketClient.Close()
+	if err := WebSocketClient.Connect(); 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.SendMessage("", nil)
+	time.Sleep(300 * time.Millisecond)
+	if resp := <-WebSocketClient.ResponseChannel; resp.Error.Id != "api.web_socket_router.no_action.app_error" {
+		t.Fatal("should have been no action response")
+	}
+
+	WebSocketClient.SendMessage("junk", nil)
+	time.Sleep(300 * time.Millisecond)
+	if resp := <-WebSocketClient.ResponseChannel; resp.Error.Id != "api.web_socket_router.bad_action.app_error" {
+		t.Fatal("should have been bad action response")
+	}
+
+	req := &model.WebSocketRequest{}
+	req.Seq = 0
+	req.Action = "ping"
+	WebSocketClient.Conn.WriteJSON(req)
+	time.Sleep(300 * time.Millisecond)
+	if resp := <-WebSocketClient.ResponseChannel; resp.Error.Id != "api.web_socket_router.bad_seq.app_error" {
+		t.Fatal("should have been bad action response")
+	}
+
+	WebSocketClient.UserTyping("", "")
+	time.Sleep(300 * time.Millisecond)
+	if resp := <-WebSocketClient.ResponseChannel; resp.Error.Id != "api.websocket_handler.invalid_param.app_error" {
+		t.Fatal("should have been invalid param response")
+	} else {
+		if resp.Error.DetailedError != "" {
+			t.Fatal("detailed error not cleared")
+		}
+	}
+}