diff --git a/internal/websocket/handle_connection.go b/internal/websocket/handle_connection.go index 2a0f414..03f0114 100644 --- a/internal/websocket/handle_connection.go +++ b/internal/websocket/handle_connection.go @@ -20,6 +20,22 @@ var upgrader = websocket.Upgrader{ CheckOrigin: func(r *http.Request) bool { return true }, } +// initialize a new WebsocketClient +func newWebsocketClient(ws *websocket.Conn, address string) models.WebsocketClient { + wsclient := models.WebsocketClient{ + Conn: ws, + LastPong: time.Now(), + RealIP: address, + } + + wsclient.Conn.SetPongHandler(func(appData string) error { + wsclient.LastPong = time.Now() + return nil + }) + + return wsclient +} + // Upgrades HTTP requests to WebSocket connections and manages the client lifecycle func HandleConnection(w http.ResponseWriter, r *http.Request) { // Upgrade initial GET request to a WebSocket @@ -31,11 +47,8 @@ func HandleConnection(w http.ResponseWriter, r *http.Request) { defer ws.Close() // Register the new client - client := models.WebsocketClient{ - Conn: ws, - LastPong: time.Now(), - RealIP: r.RemoteAddr, - } + client := newWebsocketClient(ws, r.RemoteAddr) + mu.Lock() Clients[ws.LocalAddr()] = &client mu.Unlock()