fix closing the connection when OpClose receive, add test to this case

Author: exelban

channel.go

@@ -25,6 +25,7 @@ func newChannel(id string) *Channel {
 			select {
 			case conn := <-c.delConn:
 				c.mu.Lock()
+				_ = conn.Close()
 				delete(c.connections, conn)
 				c.mu.Unlock()
 			}

channel_test.go

@@ -81,9 +81,9 @@ func TestChannel_Emit(t *testing.T) {
 
 func TestChannel_Remove(t *testing.T) {
 	ts, wsServer := wsServer()
-	defer ts.Close()
 	defer func() {
 		require.NoError(t, wsServer.Shutdown())
+		ts.Close()
 	}()
 
 	ch := wsServer.NewChannel("test-channel-add")

conn.go

@@ -92,6 +92,10 @@ func (c *Conn) Send(data interface{}) error {
 
 // Close closing websocket connection.
 func (c *Conn) Close() error {
+	if c.conn == nil {
+		return nil
+	}
+
 	c.done <- true
 
 	err := c.conn.Close()

websocket.go

@@ -207,8 +207,9 @@ func (s *Server) Handler(w http.ResponseWriter, r *http.Request) {
 	for {
 		header, _ := ws.ReadHeader(conn)
 		if err = ws.CheckHeader(header, state); err != nil {
+			log.Printf("drop connection: CheckHeader: %v", err)
 			s.dropConn(connection)
-			return
+			break
 		}
 
 		cipherReader.Reset(io.LimitReader(conn, header.Length), header.Mask)
@@ -261,8 +262,9 @@ func (s *Server) Handler(w http.ResponseWriter, r *http.Request) {
 		}
 
 		if err != nil || header.OpCode == ws.OpClose {
+			log.Printf("drop connection: %v or OpClose", err)
 			s.dropConn(connection)
-			return
+			break
 		}
 
 		header.Masked = false

websocket_test.go

@@ -9,6 +9,7 @@ import (
 	"github.com/gobwas/ws"
 	"github.com/gobwas/ws/wsutil"
 	"github.com/stretchr/testify/require"
+	"log"
 	"math/rand"
 	"net"
 	"net/http"
@@ -529,6 +530,84 @@ func TestServerProcessMessage(t *testing.T) {
 	}
 }
 
+func TestServer_ConnectionClose(t *testing.T) {
+	ts, wsServer := wsServer()
+	defer func() {
+		require.NoError(t, wsServer.Shutdown())
+		ts.Close()
+	}()
+
+	ch := wsServer.NewChannel("test-channel-add")
+	msg := Message{
+		Name: "test",
+		Data: []byte("Hello World"),
+	}
+	messageBytes, err := json.Marshal(msg)
+	require.NoError(t, err)
+
+	ticker := time.NewTicker(time.Millisecond * 1)
+	done := make(chan bool, 1)
+
+	wsServer.OnConnect(func(c *Conn) {
+		ch.Add(c)
+		require.Equal(t, 1, ch.Count(), "channel must contain only 1 connection")
+		log.Print("Connected")
+	})
+	wsServer.OnDisconnect(func(c *Conn) {
+		log.Print("Disconnected")
+	})
+	wsServer.On("test", func(c *Conn, msg *Message) {
+		log.Printf("message: %s", msg.Name)
+	})
+
+	u := url.URL{Scheme: "ws", Host: strings.Replace(ts.URL, "http://", "", 1), Path: "/ws"}
+	c, _, _, err := ws.Dial(context.Background(), u.String())
+	require.NoError(t, err)
+	defer func() {
+		require.NoError(t, c.Close())
+	}()
+
+	go func() {
+		for {
+			select {
+			case <-ticker.C:
+				err = ws.WriteHeader(c, ws.Header{
+					Fin:    true,
+					OpCode: ws.OpText,
+					Masked: true,
+					Length: int64(len(messageBytes)),
+				})
+				require.NoError(t, err)
+
+				n, err := c.Write(messageBytes)
+				require.NoError(t, err)
+				require.Equal(t, len(messageBytes), n)
+			case <-done:
+				ticker.Stop()
+				return
+			}
+		}
+	}()
+
+	time.Sleep(time.Millisecond * 50)
+
+	done <- true
+
+	time.Sleep(time.Millisecond * 5)
+
+	err = ws.WriteHeader(c, ws.Header{
+		Fin:    true,
+		OpCode: ws.OpClose,
+		Masked: true,
+		Length: 0,
+	})
+	require.NoError(t, err)
+
+	time.Sleep(time.Millisecond * 5)
+
+	require.Equal(t, 0, ch.Count())
+}
+
 func wsServer() (*httptest.Server, *Server) {
 	wsServer := Start(context.Background())
 	r := chi.NewRouter()