fix: 优化TCP/WebSocket协议检测与处理逻辑

logic/service/player/pack.go

@@ -23,6 +23,11 @@ import (
 	"github.com/panjf2000/gnet/v2"
 )
 
+const (
+	minPacketLen = 17
+	maxPacketLen = 10 * 1024
+)
+
 // getUnderlyingValue 递归解析reflect.Value，解包指针、interface{}到底层具体类型
 func getUnderlyingValue(val reflect.Value) (reflect.Value, error) {
 	for {
@@ -106,6 +111,12 @@ func (h *ClientData) PushEvent(v []byte, submit func(task func()) error) {
 	if h == nil || h.IsClosed() {
 		return
 	}
+	if len(v) < minPacketLen || len(v) > maxPacketLen {
+		return
+	}
+	if binary.BigEndian.Uint32(v[0:4]) != uint32(len(v)) {
+		return
+	}
 
 	var header common.TomeeHeader
 	header.Len = binary.BigEndian.Uint32(v[0:4])
@@ -249,13 +260,13 @@ func (h *ClientData) OnEvent(data common.TomeeHeader) {
 }
 
 type ClientData struct {
-	IsCrossDomain sync.Once //是否跨域过
-	Player        *Player   //客户实体
-	ERROR_CONNUT  int
-	Wsmsg         *WsCodec
-	Conn          gnet.Conn
-	LF            *lockfree.Lockfree[common.TomeeHeader]
-	closed        int32
+	Player             *Player //客户实体
+	ERROR_CONNUT       int
+	Wsmsg              *WsCodec
+	Conn               gnet.Conn
+	LF                 *lockfree.Lockfree[common.TomeeHeader]
+	closed             int32
+	crossDomainChecked uint32
 }
 
 func (p *ClientData) IsClosed() bool {
@@ -271,6 +282,14 @@ func (p *ClientData) Close() {
 	}
 }
 
+func (p *ClientData) IsCrossDomainChecked() bool {
+	return atomic.LoadUint32(&p.crossDomainChecked) == 1
+}
+
+func (p *ClientData) MarkCrossDomainChecked() {
+	atomic.StoreUint32(&p.crossDomainChecked, 1)
+}
+
 func (p *ClientData) GetPlayer(userid uint32) *Player { //TODO 这里待优化,可能存在内存泄漏问题
 	if p.Player == nil {
 		p.Player = NewPlayer(p.Conn)

logic/service/player/wscodec.go

@@ -2,6 +2,7 @@ package player
 
 import (
 	"bytes"
+	"encoding/binary"
 	"errors"
 	"io"
 
@@ -11,12 +12,18 @@ import (
 	"github.com/panjf2000/gnet/v2/pkg/logging"
 )
 
+const (
+	minTCPPacketLen = 17
+	maxTCPPacketLen = 10 * 1024
+	tomeeVersion    = 49
+)
+
 type WsCodec struct {
-	Tcp      bool
-	Upgraded bool         // 链接是否升级
-	Buf      bytes.Buffer // 从实际socket中读取到的数据缓存
-	wsMsgBuf wsMessageBuf // ws 消息缓存
-	//Isinitws bool
+	Tcp             bool
+	Upgraded        bool         // 链接是否升级
+	Buf             bytes.Buffer // 从实际socket中读取到的数据缓存
+	wsMsgBuf        wsMessageBuf // ws 消息缓存
+	bufferedInbound int          // 已镜像到 Buf 中的 inbound 字节数
 }
 
 type wsMessageBuf struct {
@@ -24,92 +31,115 @@ type wsMessageBuf struct {
 	cachedBuf bytes.Buffer
 }
 
+type UpgradeState uint8
+
+const (
+	UpgradeNeedMoreData UpgradeState = iota
+	UpgradeUseTCP
+	UpgradeUseWS
+)
+
 type readWrite struct {
 	io.Reader
 	io.Writer
 }
 
-func CompareLeftBytes(array1, array2 []byte, leftBytesCount int) bool {
-	// 检查切片长度是否足够比较左边的字节
-	if len(array1) < leftBytesCount || len(array2) < leftBytesCount {
-		return false
-	}
-
-	// 提取左边的字节切片
-	left1 := array1[:leftBytesCount]
-	left2 := array2[:leftBytesCount]
-
-	// 比较左边的字节切片
-	for i := 0; i < leftBytesCount; i++ {
-		if left1[i] != left2[i] {
-			return false
-		}
-	}
-
-	return true
-}
-func (w *WsCodec) Upgrade(c gnet.Conn) (ok bool, action gnet.Action) {
+func (w *WsCodec) Upgrade(c gnet.Conn) (state UpgradeState, action gnet.Action) {
 	if w.Upgraded {
-		ok = true
+		state = UpgradeUseWS
 		return
 	}
-
 	if w.Tcp {
-		ok = false
+		state = UpgradeUseTCP
 		return
 	}
-	buf := &w.Buf
-	if CompareLeftBytes(buf.Bytes(), []byte{0, 0}, 2) {
-		w.Tcp = true
-		return
-	}
-	tmpReader := bytes.NewReader(buf.Bytes())
-	oldLen := tmpReader.Len()
-	//logging.Infof("do Upgrade")
 
+	buf := w.Buf.Bytes()
+	if looksLikeTCPPacket(buf) {
+		w.SwitchToTCP()
+		state = UpgradeUseTCP
+		return
+	}
+	if len(buf) == 0 {
+		state = UpgradeNeedMoreData
+		return
+	}
+
+	tmpReader := bytes.NewReader(buf)
+	oldLen := tmpReader.Len()
 	hs, err := ws.Upgrade(readWrite{tmpReader, c})
 	skipN := oldLen - tmpReader.Len()
 	if err != nil {
-		if err == io.EOF || errors.Is(err, io.ErrUnexpectedEOF) { //数据不完整，不跳过 buf 中的 skipN 字节（此时 buf 中存放的仅是部分 "handshake data" bytes），下次再尝试读取
+		if err == io.EOF || errors.Is(err, io.ErrUnexpectedEOF) {
+			state = UpgradeNeedMoreData
 			return
 		}
-		buf.Next(skipN)
+		w.Buf.Next(skipN)
 		logging.Errorf("conn[%v] [err=%v]", c.RemoteAddr().String(), err.Error())
 		action = gnet.Close
-		//ok = true
-		//w.Tcp = true
 		return
 	}
-	buf.Next(skipN)
-	logging.Infof("conn[%v] upgrade websocket protocol! Handshake: %v", c.RemoteAddr().String(), hs)
 
-	ok = true
+	w.Buf.Next(skipN)
+	logging.Infof("conn[%v] upgrade websocket protocol! Handshake: %v", c.RemoteAddr().String(), hs)
 	w.Upgraded = true
+	state = UpgradeUseWS
 	return
 }
+
+func looksLikeTCPPacket(buf []byte) bool {
+	if len(buf) < 4 {
+		return false
+	}
+	packetLen := binary.BigEndian.Uint32(buf[:4])
+	if packetLen < minTCPPacketLen || packetLen > maxTCPPacketLen {
+		return false
+	}
+	if len(buf) >= 5 && buf[4] != tomeeVersion {
+		return false
+	}
+	return true
+}
+
 func (w *WsCodec) ReadBufferBytes(c gnet.Conn) (gnet.Action, int) {
 	size := c.InboundBuffered()
-	//buf := make([]byte, size)
+	if size < w.bufferedInbound {
+		w.bufferedInbound = 0
+	}
+	if size == w.bufferedInbound {
+		return gnet.None, size
+	}
+
 	read, err := c.Peek(size)
 	if err != nil {
 		logging.Errorf("read err! %v", err)
 		return gnet.Close, 0
 	}
-	// if read < size {
-	// 	logging.Errorf("read bytes len err! size: %d read: %d", size, read)
-	// 	return gnet.Close
-	// }
-	w.Buf.Write(read)
+	w.Buf.Write(read[w.bufferedInbound:])
+	w.bufferedInbound = size
 	return gnet.None, size
 }
+
+func (w *WsCodec) ResetInboundMirror() {
+	w.bufferedInbound = 0
+}
+
+func (w *WsCodec) SwitchToTCP() {
+	w.Tcp = true
+	w.Upgraded = false
+	w.bufferedInbound = 0
+	w.Buf.Reset()
+	w.wsMsgBuf.curHeader = nil
+	w.wsMsgBuf.cachedBuf.Reset()
+}
+
 func (w *WsCodec) Decode(c gnet.Conn) (outs []wsutil.Message, err error) {
-	//	fmt.Println("do Decode")
 	messages, err := w.readWsMessages()
 	if err != nil {
 		logging.Errorf("Error reading message! %v", err)
 		return nil, err
 	}
-	if len(messages) <= 0 { //没有读到完整数据 不处理
+	if len(messages) <= 0 {
 		return
 	}
 	for _, message := range messages {
@@ -131,9 +161,8 @@ func (w *WsCodec) readWsMessages() (messages []wsutil.Message, err error) {
 	msgBuf := &w.wsMsgBuf
 	in := &w.Buf
 	for {
-		// 从 in 中读出 header，并将 header bytes 写入 msgBuf.cachedBuf
 		if msgBuf.curHeader == nil {
-			if in.Len() < ws.MinHeaderSize { //头长度至少是2
+			if in.Len() < ws.MinHeaderSize {
 				return
 			}
 			var head ws.Header
@@ -142,13 +171,13 @@ func (w *WsCodec) readWsMessages() (messages []wsutil.Message, err error) {
 				if err != nil {
 					return messages, err
 				}
-			} else { //有可能不完整，构建新的 reader 读取 head，读取成功才实际对 in 进行读操作
+			} else {
 				tmpReader := bytes.NewReader(in.Bytes())
 				oldLen := tmpReader.Len()
 				head, err = ws.ReadHeader(tmpReader)
 				skipN := oldLen - tmpReader.Len()
 				if err != nil {
-					if err == io.EOF || errors.Is(err, io.ErrUnexpectedEOF) { //数据不完整
+					if err == io.EOF || errors.Is(err, io.ErrUnexpectedEOF) {
 						return messages, nil
 					}
 					in.Next(skipN)
@@ -163,21 +192,19 @@ func (w *WsCodec) readWsMessages() (messages []wsutil.Message, err error) {
 				return nil, err
 			}
 		}
-		dataLen := (int)(msgBuf.curHeader.Length)
-		// 从 in 中读出 data，并将 data bytes 写入 msgBuf.cachedBuf
-		if dataLen > 0 {
-			if in.Len() < dataLen { //数据不完整
 
+		dataLen := int(msgBuf.curHeader.Length)
+		if dataLen > 0 {
+			if in.Len() < dataLen {
 				logging.Infof("incomplete data")
 				return
 			}
-
 			_, err = io.CopyN(&msgBuf.cachedBuf, in, int64(dataLen))
 			if err != nil {
 				return
 			}
 		}
-		if msgBuf.curHeader.Fin { //当前 header 已经是一个完整消息
+		if msgBuf.curHeader.Fin {
 			messages, err = wsutil.ReadClientMessage(&msgBuf.cachedBuf, messages)
 			if err != nil {
 				return nil, err