fix: panic when write on closed channel

api/stream/client.go

@@ -24,6 +24,7 @@ type client struct {
 	conn    *websocket.Conn
 	onClose func(*client)
 	write   chan *model.MessageExternal
+	closed  chan struct{}
 	userID  uint
 	token   string
 	once    once
@@ -33,6 +34,7 @@ func newClient(conn *websocket.Conn, userID uint, token string, onClose func(*cl
 	return &client{
 		conn:    conn,
 		write:   make(chan *model.MessageExternal, 1),
+		closed:  make(chan struct{}),
 		userID:  userID,
 		token:   token,
 		onClose: onClose,
@@ -43,15 +45,15 @@ func newClient(conn *websocket.Conn, userID uint, token string, onClose func(*cl
 func (c *client) Close() {
 	c.once.Do(func() {
 		c.conn.Close()
-		close(c.write)
+		close(c.closed)
 	})
 }
 
 // NotifyClose closes the connection and notifies that the connection was closed.
 func (c *client) NotifyClose() {
 	c.once.Do(func() {
 		c.conn.Close()
-		close(c.write)
+		close(c.closed)
 		c.onClose(c)
 	})
 }
@@ -87,11 +89,9 @@ func (c *client) startWriteHandler(pingPeriod time.Duration) {
 
 	for {
 		select {
-		case message, ok := <-c.write:
-			if !ok {
-				return
-			}
-
+		case <-c.closed:
+			return
+		case message := <-c.write:
 			c.conn.SetWriteDeadline(time.Now().Add(writeWait))
 			if err := writeJSON(c.conn, message); err != nil {
 				printWebSocketError("WriteError", err)

api/stream/stream.go

@@ -84,7 +84,10 @@ func (a *API) Notify(userID uint, msg *model.MessageExternal) {
 	defer a.lock.RUnlock()
 	if clients, ok := a.clients[userID]; ok {
 		for _, c := range clients {
-			c.write <- msg
+			select {
+			case c.write <- msg:
+			case <-c.closed:
+			}
 		}
 	}
 }

api/stream/stream_test.go

@@ -99,6 +99,24 @@ func TestWritePingFails(t *testing.T) {
 	user.expectNoMessage()
 }
 
+func TestNotifyDoesNotPanicWhenClientIsClosed(t *testing.T) {
+	mode.Set(mode.TestDev)
+	defer leaktest.Check(t)()
+
+	server, api := bootTestServer(staticUserID())
+	defer server.Close()
+	defer api.Close()
+
+	ws, _, err := websocket.DefaultDialer.Dial(wsURL(server.URL), nil)
+	assert.Nil(t, err)
+	defer ws.Close()
+
+	waitForConnectedClients(api, 1)
+
+	clients(api, 1)[0].Close()
+	api.Notify(1, &model.MessageExternal{Message: "after close"})
+}
+
 func TestPing(t *testing.T) {
 	mode.Set(mode.TestDev)