From 7a49c557e011e785d36f898d82532f3d29459063 Mon Sep 17 00:00:00 2001
From: Marcos Lilljedahl <marcosnils@gmail.com>
Date: Mon, 31 Jul 2017 18:20:27 -0300
Subject: [PATCH] Return 404 for WS handler when session doesnt exist

---
 api.go                |  4 +---
 handlers/bootstrap.go |  8 ++++----
 handlers/ping.go      |  2 +-
 handlers/ws.go        | 14 ++++++++++++++
 4 files changed, 20 insertions(+), 8 deletions(-)

diff --git a/api.go b/api.go
index 313e2bf..fcb762d 100644
--- a/api.go
+++ b/api.go
@@ -41,8 +41,6 @@ func main() {
 		}
 	}()
 
-	server := handlers.Broadcast.GetHandler()
-
 	r := mux.NewRouter()
 	corsRouter := mux.NewRouter()
 
@@ -77,7 +75,7 @@ func main() {
 		http.ServeFile(rw, r, "www/sdk.js")
 	})
 
-	corsRouter.Handle("/sessions/{sessionId}/ws/", server)
+	corsRouter.HandleFunc("/sessions/{sessionId}/ws/", handlers.WebSocket)
 	r.Handle("/metrics", promhttp.Handler())
 
 	// Generic routes
diff --git a/handlers/bootstrap.go b/handlers/bootstrap.go
index 9e8e90b..74a8282 100644
--- a/handlers/bootstrap.go
+++ b/handlers/bootstrap.go
@@ -12,7 +12,7 @@ import (
 )
 
 var core pwd.PWDApi
-var Broadcast pwd.BroadcastApi
+var broadcast pwd.BroadcastApi
 
 func Bootstrap() {
 	c, err := client.NewEnvClient()
@@ -22,18 +22,18 @@ func Bootstrap() {
 
 	d := docker.NewDocker(c)
 
-	Broadcast, err = pwd.NewBroadcast(WS, WSError)
+	broadcast, err = pwd.NewBroadcast(WS, WSError)
 	if err != nil {
 		log.Fatal(err)
 	}
 
-	t := pwd.NewScheduler(Broadcast, d)
+	t := pwd.NewScheduler(broadcast, d)
 
 	s, err := storage.NewFileStorage(config.SessionsFile)
 
 	if err != nil && !os.IsNotExist(err) {
 		log.Fatal("Error decoding sessions from disk ", err)
 	}
-	core = pwd.NewPWD(d, t, Broadcast, s)
+	core = pwd.NewPWD(d, t, broadcast, s)
 
 }
diff --git a/handlers/ping.go b/handlers/ping.go
index 5e621f9..49cfdd1 100644
--- a/handlers/ping.go
+++ b/handlers/ping.go
@@ -21,7 +21,7 @@ func Ping(rw http.ResponseWriter, req *http.Request) {
 		return
 	}
 
-	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Microsecond)
+	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
 	defer cancel()
 
 	if _, err := c.Info(ctx); err != nil && err == context.DeadlineExceeded {
diff --git a/handlers/ws.go b/handlers/ws.go
index 746d2bc..c8b938b 100644
--- a/handlers/ws.go
+++ b/handlers/ws.go
@@ -3,6 +3,7 @@ package handlers
 import (
 	"fmt"
 	"log"
+	"net/http"
 
 	"github.com/googollee/go-socket.io"
 	"github.com/gorilla/mux"
@@ -52,3 +53,16 @@ func WS(so socketio.Socket) {
 func WSError(so socketio.Socket) {
 	log.Println("error ws")
 }
+
+func WebSocket(rw http.ResponseWriter, req *http.Request) {
+	vars := mux.Vars(req)
+	sessionId := vars["sessionId"]
+
+	session := core.SessionGet(sessionId)
+	if session == nil {
+		rw.WriteHeader(http.StatusNotFound)
+		return
+	}
+
+	broadcast.GetHandler().ServeHTTP(rw, req)
+}