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

common/socket/ServerEvent.go

@@ -1,24 +1,29 @@
 package socket
 
 import (
+	"blazing/common/socket/codec"
+	"blazing/cool"
+	"blazing/logic/service/player"
+	"blazing/modules/config/service"
+	"bytes"
 	"context"
 	"encoding/binary"
 	"errors"
-	"io"
 	"log"
 	"os"
 	"sync/atomic"
 	"time"
 
-	"blazing/cool"
-	"blazing/logic/service/player"
-	"blazing/modules/config/service"
-
 	"github.com/gogf/gf/v2/frame/g"
 	"github.com/gogf/gf/v2/os/gtime"
 	"github.com/panjf2000/gnet/v2"
 )
 
+const (
+	minPacketLen = 17
+	maxPacketLen = 10 * 1024
+)
+
 func (s *Server) Boot(serverid, port uint32) error {
 	//	go s.bootws()
 	s.serverid = serverid
@@ -53,36 +58,19 @@ func (s *Server) Stop() error {
 func (s *Server) OnClose(c gnet.Conn, err error) (action gnet.Action) {
 	defer func() {
 		if err := recover(); err != nil { // 恢复 panic，err 为 panic 错误值
-			// 1. 打印错误信息
 			if t, ok := c.Context().(*player.ClientData); ok {
-				if t.Player != nil {
-					if t.Player.Info != nil {
-						cool.Logger.Error(context.TODO(), "OnClose 错误：", cool.Config.ServerInfo.OnlineID, t.Player.Info.UserID, err)
-						t.Player.Service.Info.Save(*t.Player.Info)
-					}
-
+				if t.Player != nil && t.Player.Info != nil {
+					cool.Logger.Error(context.TODO(), "OnClose 错误：", cool.Config.ServerInfo.OnlineID, t.Player.Info.UserID, err)
+					t.Player.Service.Info.Save(*t.Player.Info)
 				}
-
 			} else {
 				cool.Logger.Error(context.TODO(), "OnClose 错误：", cool.Config.ServerInfo.OnlineID, err)
-
 			}
-
 		}
 	}()
-	// 识别 RST 导致的连接中断（错误信息含 "connection reset"）
-	// if err != nil && (strings.Contains(err.Error(), "connection reset") || strings.Contains(err.Error(), "reset by peer")) {
-	// 	remoteIP := c.RemoteAddr().(*net.TCPAddr).IP.String()
 
-	// 	log.Printf("RST 攻击检测: 来源 %s, 累计攻击次数 %d", remoteIP)
-
-	// 	// 防护逻辑：临时封禁异常 IP（可扩展为 IP 黑名单）
-	// 	//	go s.tempBlockIP(remoteIP, 5*time.Minute)
-	// }
-	//fmt.Println(err, c.RemoteAddr().String(), "断开连接")
 	atomic.AddInt64(&cool.Connected, -1)
 
-	//logging.Infof("conn[%v] disconnected", c.RemoteAddr().String())
 	v, _ := c.Context().(*player.ClientData)
 	if v != nil {
 		v.Close()
@@ -90,23 +78,20 @@ func (s *Server) OnClose(c gnet.Conn, err error) (action gnet.Action) {
 			v.Player.Save() //保存玩家数据
 		}
 	}
-
-	//}
-	//关闭连接
 	return
 }
+
 func (s *Server) OnTick() (delay time.Duration, action gnet.Action) {
 	g.Log().Async().Info(context.Background(), gtime.Now().ISO8601(), "服务器ID", cool.Config.ServerInfo.OnlineID, "链接数", atomic.LoadInt64(&cool.Connected))
 	if s.quit && atomic.LoadInt64(&cool.Connected) == 0 {
-		//执行正常退出逻辑
 		os.Exit(0)
 	}
 	return 30 * time.Second, gnet.None
 }
+
 func (s *Server) OnBoot(eng gnet.Engine) gnet.Action {
 	s.eng = eng
-
-	service.NewServerService().SetServerID(s.serverid, s.port) //设置当前服务器端口
+	service.NewServerService().SetServerID(s.serverid, s.port)
 	return gnet.None
 }
 
@@ -114,59 +99,52 @@ func (s *Server) OnOpen(conn gnet.Conn) (out []byte, action gnet.Action) {
 	if s.network != "tcp" {
 		return nil, gnet.Close
 	}
-
 	if conn.Context() == nil {
-		conn.SetContext(player.NewClientData(conn)) //注入data
+		conn.SetContext(player.NewClientData(conn))
 	}
-
 	atomic.AddInt64(&cool.Connected, 1)
-
 	return nil, gnet.None
 }
 
 func (s *Server) OnTraffic(c gnet.Conn) (action gnet.Action) {
 	defer func() {
-		if err := recover(); err != nil { // 恢复 panic，err 为 panic 错误值
-			// 1. 打印错误信息
+		if err := recover(); err != nil {
 			if t, ok := c.Context().(*player.ClientData); ok {
-				if t.Player != nil {
-					if t.Player.Info != nil {
-						cool.Logger.Error(context.TODO(), "OnTraffic 错误：", cool.Config.ServerInfo.OnlineID, t.Player.Info.UserID, err)
-						t.Player.Service.Info.Save(*t.Player.Info)
-
-					}
-
+				if t.Player != nil && t.Player.Info != nil {
+					cool.Logger.Error(context.TODO(), "OnTraffic 错误：", cool.Config.ServerInfo.OnlineID, t.Player.Info.UserID, err)
+					t.Player.Service.Info.Save(*t.Player.Info)
 				}
-
 			}
-
 		}
 	}()
 
 	ws := c.Context().(*player.ClientData).Wsmsg
-	if ws.Tcp { //升级失败时候防止缓冲区溢出
+	if ws.Tcp {
 		return s.handleTCP(c)
-
 	}
 
-	tt, len1 := ws.ReadBufferBytes(c)
-	if tt == gnet.Close {
-
+	readAction, inboundLen := ws.ReadBufferBytes(c)
+	if readAction == gnet.Close {
 		return gnet.Close
 	}
 
-	ok, action := ws.Upgrade(c)
-	if action != gnet.None { //连接断开
+	state, action := ws.Upgrade(c)
+	if action != gnet.None {
 		return action
 	}
-	if !ok { //升级失败,说明是tcp连接
-		ws.Tcp = true
-
-		return s.handleTCP(c)
-
+	if state == player.UpgradeNeedMoreData {
+		return gnet.None
+	}
+	if state == player.UpgradeUseTCP {
+		return s.handleTCP(c)
+	}
+
+	if inboundLen > 0 {
+		if _, err := c.Discard(inboundLen); err != nil {
+			return gnet.Close
+		}
+		ws.ResetInboundMirror()
 	}
-	//	fmt.Println(ws.Buf.Bytes())
-	c.Discard(len1)
 
 	messages, err := ws.Decode(c)
 	if err != nil {
@@ -177,91 +155,93 @@ func (s *Server) OnTraffic(c gnet.Conn) (action gnet.Action) {
 	}
 
 	for _, msg := range messages {
-
-		s.onevent(c, msg.Payload)
-		//t.OnEvent(msg.Payload)
+		if !s.onevent(c, msg.Payload) {
+			return gnet.Close
+		}
 	}
-
 	return gnet.None
 }
 
-const maxBodyLen = 10 * 1024 // 业务最大包体长度，按需调整
 func (s *Server) handleTCP(conn gnet.Conn) (action gnet.Action) {
+	client := conn.Context().(*player.ClientData)
+	if s.discorse && !client.IsCrossDomainChecked() {
+		handled, ready, action := handle(conn)
+		if action != gnet.None {
+			return action
+		}
+		if !ready {
+			return gnet.None
+		}
+		if handled {
+			client.MarkCrossDomainChecked()
+			return gnet.None
+		}
+		client.MarkCrossDomainChecked()
+	}
 
-	conn.Context().(*player.ClientData).IsCrossDomain.Do(func() { //跨域检测
-		handle(conn)
-	})
-
-	//	handle(c)
-	// 先读取4字节的包长度
-	lenBuf, err := conn.Peek(4)
-
+	body, err := s.codec.Decode(conn)
 	if err != nil {
-		if errors.Is(err, io.ErrShortBuffer) {
-			return
+		if errors.Is(err, codec.ErrIncompletePacket) {
+			return gnet.None
 		}
 		return gnet.Close
 	}
-
-	bodyLen := binary.BigEndian.Uint32(lenBuf)
-
-	if bodyLen > maxBodyLen {
+	if !s.onevent(conn, body) {
 		return gnet.Close
 	}
-
-	if conn.InboundBuffered() < int(bodyLen) {
-		return
-	}
-	// 提取包体
-	body, err := conn.Next(int(bodyLen))
-	if err != nil {
-		if errors.Is(err, io.ErrShortBuffer) {
-			return
-		}
-		return gnet.Close
-	}
-
-	s.onevent(conn, body)
-
 	if conn.InboundBuffered() > 0 {
-		if err := conn.Wake(nil); err != nil { // wake up the connection manually to avoid missing the leftover data
-
+		if err := conn.Wake(nil); err != nil {
 			return gnet.Close
 		}
 	}
 	return action
-
 }
 
-// CROSS_DOMAIN 定义跨域策略文件内容
 const CROSS_DOMAIN = "<?xml version=\"1.0\"?><!DOCTYPE cross-domain-policy><cross-domain-policy><allow-access-from domain=\"*\" to-ports=\"*\" /></cross-domain-policy>\x00"
-
-// TEXT 定义跨域请求的文本格式
 const TEXT = "<policy-file-request/>\x00"
 
-func handle(c gnet.Conn) {
+func handle(c gnet.Conn) (handled bool, ready bool, action gnet.Action) {
+	probeLen := c.InboundBuffered()
+	if probeLen == 0 {
+		return false, false, gnet.None
+	}
+	if probeLen > len(TEXT) {
+		probeLen = len(TEXT)
+	}
 
-	// 读取数据并检查是否为跨域请求
-	data, err := c.Peek(len(TEXT))
+	data, err := c.Peek(probeLen)
 	if err != nil {
 		log.Printf("Error reading cross-domain request: %v", err)
-		return
+		return false, false, gnet.Close
 	}
-
-	if string(data) == TEXT { //判断是否是跨域请求
-		//log.Printf("Received cross-domain request from %s", c.RemoteAddr())
-		// 处理跨域请求
-		c.Write([]byte(CROSS_DOMAIN))
-		c.Discard(len(TEXT))
-
-		return
+	if !bytes.Equal(data, []byte(TEXT[:probeLen])) {
+		return false, true, gnet.None
 	}
-
-	//return
+	if probeLen < len(TEXT) {
+		return false, false, gnet.None
+	}
+	if _, err := c.Write([]byte(CROSS_DOMAIN)); err != nil {
+		return false, true, gnet.Close
+	}
+	if _, err := c.Discard(len(TEXT)); err != nil {
+		return false, true, gnet.Close
+	}
+	return true, true, gnet.None
 }
 
-func (s *Server) onevent(c gnet.Conn, v []byte) {
+func (s *Server) onevent(c gnet.Conn, v []byte) bool {
+	if !isValidPacket(v) {
+		return false
+	}
 	if t, ok := c.Context().(*player.ClientData); ok {
 		t.PushEvent(v, s.workerPool.Submit)
 	}
+	return true
+}
+
+func isValidPacket(v []byte) bool {
+	if len(v) < minPacketLen || len(v) > maxPacketLen {
+		return false
+	}
+	return binary.BigEndian.Uint32(v[0:4]) == uint32(len(v))
 }