diff --git a/common/websockets.go b/common/websockets.go index 1d1728cb..b9f336ec 100644 --- a/common/websockets.go +++ b/common/websockets.go @@ -67,6 +67,9 @@ func RouteWebsockets(w http.ResponseWriter, r *http.Request, user User) RouteErr } break } + if conn == nil { + panic("conn must not be nil") + } messages := bytes.Split(message, []byte("\r")) for _, msg := range messages { @@ -89,6 +92,7 @@ func RouteWebsockets(w http.ResponseWriter, r *http.Request, user User) RouteErr } } conn.Close() + DebugLog("Closing connection for user " + strconv.Itoa(user.ID)) return nil } @@ -98,6 +102,7 @@ func wsPageResponses(wsUser *WSUser, conn *websocket.Conn, page string) { page = Config.DefaultPath } + DebugLog("Entering page " + page) switch page { // Live Topic List is an experimental feature // TODO: Optimise this to reduce the amount of contention @@ -120,7 +125,10 @@ func wsPageResponses(wsUser *WSUser, conn *websocket.Conn, page string) { default: return } - wsUser.SetPageForSocket(conn, page) + err := wsUser.SetPageForSocket(conn, page) + if err != nil { + LogError(err) + } } // TODO: Use a map instead of a switch to make this more modular? @@ -129,6 +137,7 @@ func wsLeavePage(wsUser *WSUser, conn *websocket.Conn, page string) { page = Config.DefaultPath } + DebugLog("Leaving page " + page) switch page { case "/topics/": wsUser.FinalizePage("/topics/", func() { @@ -141,7 +150,10 @@ func wsLeavePage(wsUser *WSUser, conn *websocket.Conn, page string) { delete(adminStatsWatchers, conn) adminStatsMutex.Unlock() } - wsUser.SetPageForSocket(conn, "") + err := wsUser.SetPageForSocket(conn, "") + if err != nil { + LogError(err) + } } // TODO: Abstract this diff --git a/common/ws_user.go b/common/ws_user.go index 50f2c325..d71f2687 100644 --- a/common/ws_user.go +++ b/common/ws_user.go @@ -8,6 +8,7 @@ import ( ) var ErrNoneOnPage = errors.New("This user isn't on that page") +var ErrInvalidSocket = errors.New("That's not a valid WebSocket Connection") type WSUser struct { User *User @@ -84,6 +85,9 @@ func (wsUser *WSUser) RemoveSocket(conn *websocket.Conn) { wsUser.Lock() if len(wsUser.Sockets) < 6 { for i, socket := range wsUser.Sockets { + if socket == nil { + continue + } if socket.conn == conn { wsUser.Sockets[i] = nil wsUser.Unlock() @@ -104,20 +108,32 @@ func (wsUser *WSUser) RemoveSocket(conn *websocket.Conn) { wsUser.Unlock() } -func (wsUser *WSUser) SetPageForSocket(conn *websocket.Conn, page string) { +func (wsUser *WSUser) SetPageForSocket(conn *websocket.Conn, page string) error { + if conn == nil { + return ErrInvalidSocket + } + wsUser.Lock() for _, socket := range wsUser.Sockets { + if socket == nil { + continue + } if socket.conn == conn { socket.Page = page } } wsUser.Unlock() + + return nil } func (wsUser *WSUser) InPage(page string) bool { wsUser.Lock() defer wsUser.Unlock() for _, socket := range wsUser.Sockets { + if socket == nil { + continue + } if socket.Page == page { return true } @@ -129,6 +145,9 @@ func (wsUser *WSUser) FinalizePage(page string, handle func()) { wsUser.Lock() defer wsUser.Unlock() for _, socket := range wsUser.Sockets { + if socket == nil { + continue + } if socket.Page == page { return }