From 1e95e4be4bae3068cbc46bbfdba74e9a363beddb Mon Sep 17 00:00:00 2001
From: Stephan Knauer <sknr@users.noreply.github.com>
Date: Sat, 11 Dec 2021 09:44:22 +0100
Subject: [PATCH 1/3] Add Shutdown method to gracefully close websocket

---
 recws.go | 39 ++++++++++++++++++++++++++++++++++-----
 1 file changed, 34 insertions(+), 5 deletions(-)

diff --git a/recws.go b/recws.go
index 1d4b27a..1e14eb4 100644
--- a/recws.go
+++ b/recws.go
@@ -16,6 +16,8 @@ import (
 	"github.com/jpillora/backoff"
 )
 
+const writeWait = time.Second
+
 // ErrNotConnected is returned when the application read/writes
 // a message and the connection is closed
 var ErrNotConnected = errors.New("websocket: not connected")
@@ -91,6 +93,17 @@ func (rc *RecConn) Close() {
 	rc.setIsConnected(false)
 }
 
+// Shutdown gracefully closes the connection by sending the websocket.CloseMessage.
+func (rc *RecConn) Shutdown() {
+	msg := websocket.FormatCloseMessage(websocket.CloseNormalClosure, "")
+	err := rc.WriteControl(websocket.CloseMessage, msg, time.Now().Add(writeWait))
+	if err != nil && err != websocket.ErrCloseSent {
+		// If close message could not be sent, then close without the handshake.
+		log.Printf("Shutdown: %v", err)
+		rc.Close()
+	}
+}
+
 // ReadMessage is a helper method for getting a reader
 // using NextReader and reading from that reader to a buffer.
 //
@@ -99,7 +112,11 @@ func (rc *RecConn) ReadMessage() (messageType int, message []byte, err error) {
 	err = ErrNotConnected
 	if rc.IsConnected() {
 		messageType, message, err = rc.Conn.ReadMessage()
-		if err != nil {
+		if websocket.IsCloseError(err, websocket.CloseNormalClosure) {
+			rc.Close()
+			return messageType, message, nil
+		}
+		if websocket.IsUnexpectedCloseError(err, websocket.CloseNormalClosure) {
 			rc.CloseAndReconnect()
 		}
 	}
@@ -117,7 +134,11 @@ func (rc *RecConn) WriteMessage(messageType int, data []byte) error {
 		rc.mu.Lock()
 		err = rc.Conn.WriteMessage(messageType, data)
 		rc.mu.Unlock()
-		if err != nil {
+		if websocket.IsCloseError(err, websocket.CloseNormalClosure) {
+			rc.Close()
+			return nil
+		}
+		if websocket.IsUnexpectedCloseError(err, websocket.CloseNormalClosure) {
 			rc.CloseAndReconnect()
 		}
 	}
@@ -137,7 +158,11 @@ func (rc *RecConn) WriteJSON(v interface{}) error {
 		rc.mu.Lock()
 		err = rc.Conn.WriteJSON(v)
 		rc.mu.Unlock()
-		if err != nil {
+		if websocket.IsCloseError(err, websocket.CloseNormalClosure) {
+			rc.Close()
+			return nil
+		}
+		if websocket.IsUnexpectedCloseError(err, websocket.CloseNormalClosure) {
 			rc.CloseAndReconnect()
 		}
 	}
@@ -156,7 +181,11 @@ func (rc *RecConn) ReadJSON(v interface{}) error {
 	err := ErrNotConnected
 	if rc.IsConnected() {
 		err = rc.Conn.ReadJSON(v)
-		if err != nil {
+		if websocket.IsCloseError(err, websocket.CloseNormalClosure) {
+			rc.Close()
+			return nil
+		}
+		if websocket.IsUnexpectedCloseError(err, websocket.CloseNormalClosure) {
 			rc.CloseAndReconnect()
 		}
 	}
@@ -421,7 +450,7 @@ func (rc *RecConn) connect() {
 			if rc.getKeepAliveTimeout() != 0 {
 				rc.keepAlive()
 			}
-		
+
 			return
 		}
 

From d31f6143c1f72605c4a11c975361f58eb1e65fda Mon Sep 17 00:00:00 2001
From: Stephan Knauer <sknr@users.noreply.github.com>
Date: Sat, 11 Dec 2021 13:54:13 +0100
Subject: [PATCH 2/3] simplify CloseError testing

---
 recws.go | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/recws.go b/recws.go
index 1e14eb4..d572630 100644
--- a/recws.go
+++ b/recws.go
@@ -116,7 +116,7 @@ func (rc *RecConn) ReadMessage() (messageType int, message []byte, err error) {
 			rc.Close()
 			return messageType, message, nil
 		}
-		if websocket.IsUnexpectedCloseError(err, websocket.CloseNormalClosure) {
+		if err != nil {
 			rc.CloseAndReconnect()
 		}
 	}
@@ -138,7 +138,7 @@ func (rc *RecConn) WriteMessage(messageType int, data []byte) error {
 			rc.Close()
 			return nil
 		}
-		if websocket.IsUnexpectedCloseError(err, websocket.CloseNormalClosure) {
+		if err != nil {
 			rc.CloseAndReconnect()
 		}
 	}
@@ -162,7 +162,7 @@ func (rc *RecConn) WriteJSON(v interface{}) error {
 			rc.Close()
 			return nil
 		}
-		if websocket.IsUnexpectedCloseError(err, websocket.CloseNormalClosure) {
+		if err != nil {
 			rc.CloseAndReconnect()
 		}
 	}
@@ -185,7 +185,7 @@ func (rc *RecConn) ReadJSON(v interface{}) error {
 			rc.Close()
 			return nil
 		}
-		if websocket.IsUnexpectedCloseError(err, websocket.CloseNormalClosure) {
+		if err != nil {
 			rc.CloseAndReconnect()
 		}
 	}

From 559c5a8d900dd70e423c8be18f0b2792367fd73a Mon Sep 17 00:00:00 2001
From: Stephan Knauer <sknr@users.noreply.github.com>
Date: Mon, 13 Dec 2021 20:47:50 +0100
Subject: [PATCH 3/3] Replace writeWait constant with param writeWait on the
 Shutdown method itself

---
 recws.go | 5 ++---
 1 file changed, 2 insertions(+), 3 deletions(-)

diff --git a/recws.go b/recws.go
index d572630..4faa7af 100644
--- a/recws.go
+++ b/recws.go
@@ -16,8 +16,6 @@ import (
 	"github.com/jpillora/backoff"
 )
 
-const writeWait = time.Second
-
 // ErrNotConnected is returned when the application read/writes
 // a message and the connection is closed
 var ErrNotConnected = errors.New("websocket: not connected")
@@ -94,7 +92,8 @@ func (rc *RecConn) Close() {
 }
 
 // Shutdown gracefully closes the connection by sending the websocket.CloseMessage.
-func (rc *RecConn) Shutdown() {
+// The writeWait param defines the duration before the deadline of the write operation is hit.
+func (rc *RecConn) Shutdown(writeWait time.Duration) {
 	msg := websocket.FormatCloseMessage(websocket.CloseNormalClosure, "")
 	err := rc.WriteControl(websocket.CloseMessage, msg, time.Now().Add(writeWait))
 	if err != nil && err != websocket.ErrCloseSent {