lib/websocket: cleanup the channel gracefulClose during Close

When calling Close, we initialize the internal channel gracefulClose
so the client can check the close response from server on another
goroutine serve() and signal back to Close to continue the closing
process.

This channel is never closed and nil-ed again after Close which
may cause resource leaks.

While at it, use sendClose to minimize duplicate code.

1 files changed, 25 insertions, 24 deletions

diff --git a/lib/websocket/client.go b/lib/websocket/client.go
index 1fca4c64..1e7256c9 100644
--- a/lib/websocket/client.go
+++ b/lib/websocket/client.go
@@ -170,31 +170,27 @@ type Client struct {
 // with status normal to server and wait for response for as long as 10
 // seconds.
 func (cl *Client) Close() (err error) {
-	cl.Lock()
-	defer cl.Unlock()
+	cl.gracefulClose = make(chan bool, 1)
+	defer func() {
+		close(cl.gracefulClose)
+		cl.gracefulClose = nil
+	}()
 
-	if cl.conn == nil {
-		return
+	err = cl.sendClose(StatusNormal, nil)
+	if errors.Is(err, ErrConnClosed) {
+		return nil
 	}
 
-	var (
-		logp          = `websocket: Close`
-		packet []byte = NewFrameClose(true, StatusNormal, nil)
+	cl.Lock()
+	defer cl.Unlock()
 
-		timer *time.Timer
-		wait  bool
+	var (
+		timer = time.NewTimer(defaultTimeout)
+		wait  = true
 	)
 
-	cl.gracefulClose = make(chan bool, 1)
-
-	err = cl.send(packet)
-	if err != nil {
-		return fmt.Errorf(`%s: %w`, logp, err)
-	}
-
-	// Wait for server to response with CLOSE.
-	timer = time.NewTimer(defaultTimeout)
-	wait = true
+	// Wait for server to response with CLOSE or until 10 seconds without
+	// reponse.
 	for wait {
 		select {
 		case <-timer.C:
@@ -207,9 +203,13 @@ func (cl *Client) Close() (err error) {
 		}
 	}
 
+	if cl.conn == nil {
+		return nil
+	}
+
 	err = cl.conn.Close()
 	if err != nil {
-		err = fmt.Errorf(`%s: %w`, logp, err)
+		err = fmt.Errorf(`websocket: Close: %w`, err)
 	}
 	cl.conn = nil
 
@@ -221,15 +221,16 @@ func (cl *Client) Connect() (err error) {
 	cl.Lock()
 	defer cl.Unlock()
 
+	if cl.conn != nil {
+		_ = cl.conn.Close()
+		cl.conn = nil
+	}
+
 	err = cl.init()
 	if err != nil {
 		return fmt.Errorf("websocket: Connect: " + err.Error())
 	}
 
-	if cl.conn != nil {
-		cl.Quit()
-	}
-
 	err = cl.open()
 	if err != nil {
 		return fmt.Errorf("websocket: Connect: " + err.Error())