From 173d84f7e6fd6c956a9cdf01ee3288888e888c35 Mon Sep 17 00:00:00 2001
From: orgin <boyce.duan@gmail.com>
Date: Thu, 23 Jun 2022 10:07:50 +0800
Subject: [PATCH] =?UTF-8?q?=E4=BC=98=E5=8C=96websocket?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 network/ws_server.go | 16 +++++++++-------
 1 file changed, 9 insertions(+), 7 deletions(-)

diff --git a/network/ws_server.go b/network/ws_server.go
index 170a62f..2045b00 100644
--- a/network/ws_server.go
+++ b/network/ws_server.go
@@ -21,7 +21,7 @@ type WSServer struct {
 	NewAgent        func(*WSConn) Agent
 	ln              net.Listener
 	handler         *WSHandler
-	messageType 	int
+	messageType     int
 }
 
 type WSHandler struct {
@@ -33,10 +33,10 @@ type WSHandler struct {
 	conns           WebsocketConnSet
 	mutexConns      sync.Mutex
 	wg              sync.WaitGroup
-	messageType 	int
+	messageType     int
 }
 
-func (handler *WSHandler) SetMessageType(messageType int){
+func (handler *WSHandler) SetMessageType(messageType int) {
 	handler.messageType = messageType
 }
 
@@ -51,7 +51,9 @@ func (handler *WSHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 		return
 	}
 	conn.SetReadLimit(int64(handler.maxMsgLen))
-	handler.messageType = websocket.TextMessage
+	if handler.messageType == 0 {
+		handler.messageType = websocket.TextMessage
+	}
 
 	handler.wg.Add(1)
 	defer handler.wg.Done()
@@ -71,7 +73,7 @@ func (handler *WSHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 	handler.conns[conn] = struct{}{}
 	handler.mutexConns.Unlock()
 
-	wsConn := newWSConn(conn, handler.pendingWriteNum, handler.maxMsgLen,handler.messageType)
+	wsConn := newWSConn(conn, handler.pendingWriteNum, handler.maxMsgLen, handler.messageType)
 	agent := handler.newAgent(wsConn)
 	agent.Run()
 
@@ -83,9 +85,9 @@ func (handler *WSHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 	agent.OnClose()
 }
 
-func (server *WSServer) SetMessageType(messageType int){
+func (server *WSServer) SetMessageType(messageType int) {
 	server.messageType = messageType
-	if server.handler!= nil {
+	if server.handler != nil {
 		server.handler.SetMessageType(messageType)
 	}
 }