bl/logic/service/player/wscodec.go

package player

import (
	"bytes"
	"encoding/binary"
	"errors"
	"io"

	"github.com/gobwas/ws"
	"github.com/gobwas/ws/wsutil"
	"github.com/panjf2000/gnet/v2"
	"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 消息缓存
	bufferedInbound int          // 已镜像到 Buf 中的 inbound 字节数
}

type wsMessageBuf struct {
	curHeader *ws.Header
	cachedBuf bytes.Buffer
}

type UpgradeState uint8

const (
	UpgradeNeedMoreData UpgradeState = iota
	UpgradeUseTCP
	UpgradeUseWS
)

type readWrite struct {
	io.Reader
	io.Writer
}

func (w *WsCodec) Upgrade(c gnet.Conn) (state UpgradeState, action gnet.Action) {
	if w.Upgraded {
		state = UpgradeUseWS
		return
	}
	if w.Tcp {
		state = UpgradeUseTCP
		return
	}

	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) {
			state = UpgradeNeedMoreData
			return
		}
		w.Buf.Next(skipN)
		logging.Errorf("conn[%v] [err=%v]", c.RemoteAddr().String(), err.Error())
		action = gnet.Close
		return
	}

	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()
	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
	}
	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) {
	messages, err := w.readWsMessages()
	if err != nil {
		logging.Errorf("Error reading message! %v", err)
		return nil, err
	}
	if len(messages) <= 0 {
		return
	}
	for _, message := range messages {
		if message.OpCode.IsControl() {
			err = wsutil.HandleClientControlMessage(c, message)
			if err != nil {
				return
			}
			continue
		}
		if message.OpCode == ws.OpText || message.OpCode == ws.OpBinary {
			outs = append(outs, message)
		}
	}
	return
}

func (w *WsCodec) readWsMessages() (messages []wsutil.Message, err error) {
	msgBuf := &w.wsMsgBuf
	in := &w.Buf
	for {
		if msgBuf.curHeader == nil {
			if in.Len() < ws.MinHeaderSize {
				return
			}
			var head ws.Header
			if in.Len() >= ws.MaxHeaderSize {
				head, err = ws.ReadHeader(in)
				if err != nil {
					return messages, err
				}
			} 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) {
						return messages, nil
					}
					in.Next(skipN)
					return nil, err
				}
				in.Next(skipN)
			}

			msgBuf.curHeader = &head
			err = ws.WriteHeader(&msgBuf.cachedBuf, head)
			if err != nil {
				return nil, err
			}
		}

		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 {
			messages, err = wsutil.ReadClientMessage(&msgBuf.cachedBuf, messages)
			if err != nil {
				return nil, err
			}
			msgBuf.cachedBuf.Reset()
		} else {
			logging.Infof("The data is split into multiple frames")
		}
		msgBuf.curHeader = nil
	}
}