Merge pull request #118 from Jardaliao/fix-closing-connection-returns…

…-invalid-code

fix: closing websocket connection returns invalid status code

client_test.go

@@ -29,6 +29,7 @@ import (
 )
 
 import (
+	"github.com/gorilla/websocket"
 	"github.com/stretchr/testify/assert"
 )
 
@@ -352,6 +353,26 @@ func TestNewWSClient(t *testing.T) {
 	err = conn.writePing()
 	assert.Nil(t, err)
 
+	done := make(chan int)
+	conn.conn.SetCloseHandler(func(code int, text string) error {
+		defer func() {
+			done <- code
+			close(done)
+		}()
+		message := websocket.FormatCloseMessage(code, "")
+		conn.conn.WriteControl(websocket.CloseMessage, message, time.Now().Add(1e9))
+		return nil
+	})
+	serverSession := serverMsgHandler.array[0]
+	serverSession.Close()
+	select {
+	case code := <-done:
+		assert.True(t, code == websocket.CloseNormalClosure)
+	case <-time.After(5e9):
+		assert.True(t, false)
+	}
+	assert.True(t, serverSession.IsClosed())
+
 	ss.SetReader(nil)
 	assert.Nil(t, ss.(*session).reader)
 	ss.SetWriter(nil)

connection.go

@@ -628,7 +628,7 @@ func (w *gettyWSConn) writePong(message []byte) error {
 // close websocket connection
 func (w *gettyWSConn) CloseConn(waitSec int) {
 	w.updateWriteDeadline()
-	w.conn.WriteMessage(websocket.CloseMessage, []byte("bye-bye!!!"))
+	w.conn.WriteMessage(websocket.CloseMessage, websocket.FormatCloseMessage(websocket.CloseNormalClosure, "bye-bye!!!"))
 	conn := w.conn.UnderlyingConn()
 	if tcpConn, ok := conn.(*net.TCPConn); ok {
 		tcpConn.SetLinger(waitSec)