go/mailpit
1
0
Fork 0
mirror of https://github.com/axllent/mailpit.git synced 2025-03-17 21:18:19 +02:00
Code Issues Releases Activity

Fix: Correctly close websockets on client disconnect (#207)

Browse Source
This commit is contained in:
Ralph Slooten 2023-11-11 15:32:57 +13:00
parent ffe6167d96
commit 74236258db
2 changed files with 27 additions and 10 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

25
server/websockets/client.go
View File

@ -53,7 +53,24 @@ type Client struct {
send chan []byte
}
// writePump pumps messages from the hub to the websocket connection.
// ReadPump is used here solely to monitor the connection, not to actually receive messages.
func (c *Client) readPump() {
defer func() {
c.hub.unregister <- c
}()
for {
_, _, err := c.conn.ReadMessage()
if err != nil {
if websocket.IsUnexpectedCloseError(err, websocket.CloseGoingAway, websocket.CloseAbnormalClosure) {
logger.Log().Errorf("[websocket] error: %v", err)
}
break
}
}
}
// WritePump pumps messages from the hub to the websocket connection.
//
// A goroutine running writePump is started for each connection. The
// application ensures that there is at most one writer to a connection by
@ -62,7 +79,7 @@ func (c *Client) writePump() {
ticker := time.NewTicker(pingPeriod)
defer func() {
ticker.Stop()
_ = c.conn.Close()
c.hub.unregister <- c
}()
for {
select {
@ -122,8 +139,8 @@ func ServeWs(hub *Hub, w http.ResponseWriter, r *http.Request) {
client := &Client{hub: hub, conn: conn, send: make(chan []byte, 256)}
client.hub.register <- client
// Allow collection of memory referenced by the caller by doing all work in
// new goroutines.
// Allow collection of memory referenced by the caller by doing all work in new goroutines.
go client.readPump()
go client.writePump()
}

12
server/websockets/hub.go
View File

@ -1,7 +1,4 @@
// Copyright 2013 The Gorilla WebSocket Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.
// Package websockets is used to broadcast messages to connected clients
package websockets
import (
@ -47,14 +44,17 @@ func (h *Hub) Run() {
for {
select {
case client := <-h.register:
h.Clients[client] = true
if _, ok := h.Clients[client]; !ok {
logger.Log().Debugf("[websocket] client %s connected", client.conn.RemoteAddr().String())
h.Clients[client] = true
}
case client := <-h.unregister:
if _, ok := h.Clients[client]; ok {
logger.Log().Debugf("[websocket] client %s disconnected", client.conn.RemoteAddr().String())
delete(h.Clients, client)
close(client.send)
}
case message := <-h.Broadcast:
// logger.Log().Debugf("[broadcast] %s", message)
for client := range h.Clients {
select {
case client.send <- message: