Return 404 for WS handler when session doesnt exist

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

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)
 
 }

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 {

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)
+}