From 62a92741624cdd566dc5408dfcd7fda2ba8f4826 Mon Sep 17 00:00:00 2001
From: Eric Chlebek <eric@sensu.io>
Date: Fri, 19 Nov 2021 09:09:47 -0800
Subject: [PATCH] Clean up websocket close handling a bit (#4508)

* Be more tolerant of errors in transport Close()

Signed-off-by: Eric Chlebek <eric@sensu.io>

* Don't double send close message

Signed-off-by: Eric Chlebek <eric@sensu.io>
---
 backend/agentd/session.go |  9 ---------
 transport/transport.go    | 14 ++++++++------
 2 files changed, 8 insertions(+), 15 deletions(-)

diff --git a/backend/agentd/session.go b/backend/agentd/session.go
index dc51c960d7..5261da3748 100644
--- a/backend/agentd/session.go
+++ b/backend/agentd/session.go
@@ -555,15 +555,6 @@ func (s *Session) stop() {
 		}
 	}()
 
-	// Send a close message to ensure the agent closes its connection if the
-	// connection is not already closed
-	if !s.conn.Closed() {
-		if err := s.conn.SendCloseMessage(); err != nil {
-			websocketErrorCounter.WithLabelValues("send", "SendCloseMessage").Inc()
-			logger.Warning("unexpected error while sending a close message to the agent")
-		}
-	}
-
 	sessionCounter.WithLabelValues(s.cfg.Namespace).Dec()
 
 	// Gracefully wait for the send and receiver to exit, but force the websocket
diff --git a/transport/transport.go b/transport/transport.go
index 725cb70250..a99b1c8c60 100644
--- a/transport/transport.go
+++ b/transport/transport.go
@@ -146,10 +146,8 @@ func NewMessage(msgType string, payload []byte) *Message {
 	}
 }
 
-// Close attempts to send a "going away" message over the websocket connection.
-// This will cause a Write over the websocket transport, which can cause a
-// panic. We rescue potential panics and consider the connection closed,
-// returning nil, because the connection _will_ be closed. Hay!
+// Close closes the WebsocketTransport. Before closing, a closing message will
+// be sent to the other side.
 func (t *WebSocketTransport) Close() (err error) {
 	if t.Closed() {
 		return nil
@@ -162,12 +160,16 @@ func (t *WebSocketTransport) Close() (err error) {
 
 	defer func() {
 		cerr := t.Connection.Close()
-		if err == nil {
+		if err == nil && cerr != websocket.ErrCloseSent {
 			err = cerr
 		}
 	}()
 
-	return t.Connection.WriteMessage(websocket.CloseMessage, websocket.FormatCloseMessage(websocket.CloseGoingAway, "bye"))
+	if err := t.Connection.WriteMessage(websocket.CloseMessage, websocket.FormatCloseMessage(websocket.CloseGoingAway, "bye")); err != websocket.ErrCloseSent {
+		return err
+	}
+
+	return nil
 }
 
 // Closed returns true if the underlying websocket connection has been closed.