websocket: add method to gracefully Close the client connection

The Close method send the control CLOSE frame to server and wait for
the CLOSE response.

This changes unexport the SendClose method.

2 files changed, 59 insertions, 13 deletions

diff --git a/lib/websocket/client.go b/lib/websocket/client.go
index dd8cd666..abf38ab6 100644
--- a/lib/websocket/client.go
+++ b/lib/websocket/client.go
@@ -151,6 +151,49 @@ type Client struct {
 	allowRsv1 bool
 	allowRsv2 bool
 	allowRsv3 bool
+
+	// gracefulClose is a channel to gracefully close connection by
+	// client.
+	gracefulClose chan bool
+}
+
+//
+// Close gracefully close the client connection.
+//
+func (cl *Client) Close() (err error) {
+	packet := NewFrameClose(true, StatusNormal, nil)
+
+	cl.gracefulClose = make(chan bool, 1)
+
+	err = cl.send(packet)
+	if err != nil {
+		return fmt.Errorf("websocket: Close: %w", err)
+	}
+
+	// Wait for server to response with CLOSE.
+	timer := time.NewTimer(defaultTimeout)
+loop:
+	for {
+		select {
+		case <-timer.C:
+			// We did not receive server CLOSE frame in timely
+			// manner.
+			break loop
+		case <-cl.gracefulClose:
+			timer.Stop()
+			break loop
+		}
+	}
+
+	cl.Lock()
+	err = cl.conn.Close()
+	if err != nil {
+		err = fmt.Errorf("websocket: Close: %w", err)
+	}
+	cl.conn = nil
+	cl.Unlock()
+
+	return err
 }
 
 //
@@ -528,7 +571,12 @@ func (cl *Client) handleFrame(frame *Frame) (isClosing bool) {
 		cl.handleBadRequest()
 		return true
 	case OpcodeClose:
-		cl.handleClose(cl, frame)
+		// Check if we are requesting the close.
+		if cl.gracefulClose != nil {
+			cl.gracefulClose <- true
+		} else {
+			cl.handleClose(cl, frame)
+		}
 		return true
 	case OpcodePing:
 		_ = cl.handlePing(cl, frame)
@@ -622,11 +670,9 @@ func (cl *Client) SendBin(payload []byte) error {
 }
 
 //
-// SendClose send the control CLOSE frame to server.
-// If waitResponse is true, client will wait for CLOSE response from server
-// before closing the connection.
+// sendClose send the control CLOSE frame to server.
 //
-func (cl *Client) SendClose(status CloseCode, payload []byte) (err error) {
+func (cl *Client) sendClose(status CloseCode, payload []byte) (err error) {
 	packet := NewFrameClose(true, status, payload)
 	return cl.send(packet)
 }
diff --git a/lib/websocket/client_test.go b/lib/websocket/client_test.go
index ec019c81..9aae9cfc 100644
--- a/lib/websocket/client_test.go
+++ b/lib/websocket/client_test.go
@@ -56,7 +56,7 @@ func TestConnect(t *testing.T) {
 			continue
 		}
 
-		client.SendClose(StatusNormal, nil)
+		client.sendClose(StatusNormal, nil)
 	}
 }
 
@@ -138,7 +138,7 @@ func TestClientPing(t *testing.T) {
 				got.payload = got.payload[2:]
 			}
 
-			cl.SendClose(got.closeCode, got.payload)
+			cl.sendClose(got.closeCode, got.payload)
 			cl.Quit()
 			wg.Done()
 			return nil
@@ -268,7 +268,7 @@ func TestClientText(t *testing.T) {
 		testClient.handleClose = func(cl *Client, got *Frame) error {
 			exp := c.expClose
 			test.Assert(t, "close", exp, got, true)
-			cl.SendClose(got.closeCode, got.payload)
+			cl.sendClose(got.closeCode, got.payload)
 			cl.Quit()
 			wg.Done()
 			return nil
@@ -399,7 +399,7 @@ func TestClientFragmentation(t *testing.T) {
 		testClient.handleClose = func(cl *Client, got *Frame) error {
 			exp := c.expClose
 			test.Assert(t, "close", exp, got, true)
-			cl.SendClose(got.closeCode, got.payload)
+			cl.sendClose(got.closeCode, got.payload)
 			cl.Quit()
 			wg.Done()
 			return nil
@@ -628,7 +628,7 @@ func TestClientSendPing(t *testing.T) {
 	}
 }
 
-func TestClientSendClose(t *testing.T) {
+func TestClient_sendClose(t *testing.T) {
 	if _testServer == nil {
 		runTestServer()
 	}
@@ -642,7 +642,7 @@ func TestClientSendClose(t *testing.T) {
 
 	err := testClient.Connect()
 	if err != nil {
-		t.Fatal("TestClientSendClose: Connect: " + err.Error())
+		t.Fatal("TestClient_sendClose: Connect: " + err.Error())
 	}
 
 	testClient.handleClose = func(cl *Client, got *Frame) error {
@@ -661,9 +661,9 @@ func TestClientSendClose(t *testing.T) {
 	}
 
 	wg.Add(1)
-	err = testClient.SendClose(StatusNormal, []byte("normal"))
+	err = testClient.sendClose(StatusNormal, []byte("normal"))
 	if err != nil {
-		t.Fatal("TestClientSendClose: " + err.Error())
+		t.Fatal("TestClient_sendClose: " + err.Error())
 	}
 
 	wg.Wait()