From 2e0cf56e86656e791e20fb31566881e8fb748649 Mon Sep 17 00:00:00 2001
From: Corey Hulen <corey@hulen.com>
Date: Tue, 14 Feb 2017 17:25:40 -0500
Subject: Fixing websocket auth then sending challenge (#5406)

---
 api/websocket.go        |  6 ++++-
 api/websocket_test.go   | 63 +++++++++++++++++++++++++++++++++++++++++++++++++
 app/websocket_router.go |  2 ++
 3 files changed, 70 insertions(+), 1 deletion(-)

diff --git a/api/websocket.go b/api/websocket.go
index 5f390395d..5c0858910 100644
--- a/api/websocket.go
+++ b/api/websocket.go
@@ -34,7 +34,11 @@ func connect(c *Context, w http.ResponseWriter, r *http.Request) {
 	}
 
 	wc := app.NewWebConn(ws, c.Session, c.T, c.Locale)
-	app.HubRegister(wc)
+
+	if len(c.Session.UserId) > 0 {
+		app.HubRegister(wc)
+	}
+
 	go wc.WritePump()
 	wc.ReadPump()
 }
diff --git a/api/websocket_test.go b/api/websocket_test.go
index 6b8937d81..ab2959b03 100644
--- a/api/websocket_test.go
+++ b/api/websocket_test.go
@@ -251,6 +251,69 @@ func TestWebSocketEvent(t *testing.T) {
 	}
 }
 
+func TestCreateDirectChannelWithSocket(t *testing.T) {
+	th := Setup().InitBasic()
+	Client := th.BasicClient
+	user2 := th.BasicUser2
+
+	users := make([]*model.User, 0)
+	users = append(users, user2)
+
+	for i := 0; i < 10; i++ {
+		users = append(users, th.CreateUser(Client))
+	}
+
+	WebSocketClient, err := th.CreateWebSocketClient()
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer WebSocketClient.Close()
+	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")
+	}
+
+	wsr := <-WebSocketClient.EventChannel
+	if wsr.Event != model.WEBSOCKET_EVENT_HELLO {
+		t.Fatal("missing hello")
+	}
+
+	stop := make(chan bool)
+	count := 0
+
+	go func() {
+		for {
+			select {
+			case wsr := <-WebSocketClient.EventChannel:
+				if wsr.Event == model.WEBSOCKET_EVENT_DIRECT_ADDED {
+					count = count + 1
+				}
+
+			case <-stop:
+				return
+			}
+		}
+	}()
+
+	for _, user := range users {
+		time.Sleep(100 * time.Millisecond)
+		if _, err := Client.CreateDirectChannel(user.Id); err != nil {
+			t.Fatal("failed to create DM channel")
+		}
+	}
+
+	time.Sleep(5000 * time.Millisecond)
+
+	stop <- true
+
+	if count != len(users) {
+		t.Fatal("We didn't get the proper amount of direct_added messages")
+	}
+
+}
+
 func TestWebsocketOriginSecurity(t *testing.T) {
 	Setup().InitBasic()
 	url := "ws://localhost" + utils.Cfg.ServiceSettings.ListenAddress
diff --git a/app/websocket_router.go b/app/websocket_router.go
index 984b9d17e..78e36229a 100644
--- a/app/websocket_router.go
+++ b/app/websocket_router.go
@@ -58,6 +58,8 @@ func (wr *WebSocketRouter) ServeWebSocket(conn *WebConn, r *model.WebSocketReque
 			conn.SessionToken = session.Token
 			conn.UserId = session.UserId
 
+			HubRegister(conn)
+
 			resp := model.NewWebSocketResponse(model.STATUS_OK, r.Seq, nil)
 			resp.DoPreComputeJson()
 			conn.Send <- resp
-- 
cgit v1.2.3-1-g7c22