修复websocket

util/ws/websocket.go

@@ -3,6 +3,7 @@ package ws
 import (
 	"context"
 	"github.com/gorilla/websocket"
+	"log"
 )
 
 type Packet struct {
@@ -27,16 +28,21 @@ func Binary(buf []byte) *Packet {
 type Conn struct {
 	WriteChan chan *Packet
 	ReadChan  chan []byte
-	ctx       context.Context
+	Context   context.Context
+	cancel    context.CancelFunc
 	conn      *websocket.Conn
 	close     bool
 }
 
 func Wrap(conn *websocket.Conn) *Conn {
+	ctx, cancel := context.WithCancel(context.Background())
 	return &Conn{
 		WriteChan: make(chan *Packet),
 		ReadChan:  make(chan []byte),
 		conn:      conn,
+		Context:   ctx,
+		cancel:    cancel,
+		close:     false,
 	}
 }
 
@@ -44,6 +50,9 @@ func (conn *Conn) ReadLoop() {
 	for !conn.close {
 		_, message, err := conn.conn.ReadMessage()
 		if err != nil {
+			if !websocket.IsCloseError(err, websocket.CloseNormalClosure, websocket.CloseAbnormalClosure) {
+				log.Println("websocket read error:", err)
+			}
 			conn.Close()
 			break
 		}
@@ -63,16 +72,14 @@ func (conn *Conn) WriteLoop() {
 			if err != nil {
 				break
 			}
-		default:
-			if conn.close {
-				return
-			}
+		case <-conn.Context.Done():
+			return
 		}
 	}
 }
 
 func (conn *Conn) Close() error {
-	conn.ReadChan <- nil
+	conn.cancel()
 	conn.close = true
 	return conn.conn.Close()
 }