/* * 版权所有 (c) 上海元泓软件科技有限公司 2022. * 严禁通过任何媒介未经授权复制本文件. * * 作者:mic * Email:funui@outlook.com */ package YHProto import ( "bytes" "encoding/binary" "encoding/json" "errors" "fmt" "github.com/gobwas/ws/wsutil" "github.com/panjf2000/gnet/v2" "github.com/zeromicro/go-zero/core/logx" "time" "unsafe" ) //////////////////////////////////////////// // echo command //////////////////////////////////////////// type echoCommand uint32 var echoCommandStrings = [...]string{ "heartbeat", "echo", } func (c echoCommand) String() string { return echoCommandStrings[c] } var ( echoPkgMagic uint32 = 0x20160905 maxEchoStringLen uint32 = 0xffff echoPkgKey string = "123456" ) var ( ErrNotEnoughStream = errors.New("packet stream is not enough") ErrTooLargePackage = errors.New("package length is exceed the echo package's legal maximum length.") ErrIllegalMagic = errors.New("package magic is not right.") ) var echoPkgHeaderLen int func init() { echoPkgHeaderLen = (int)((uint)(unsafe.Sizeof(EchoPkgHeader{}))) } func SetEchoPkgMagic(magic uint32) { echoPkgMagic = magic } func SetEchoPkgKey(key string) { echoPkgKey = key } type ( // EchoPkgHeader // 长度20字节.SimpleCodec Protocol format: // // * 0 4 6 8 12 16 20 // * +----------------+----------+----------+----------------+----------------+----------------+ // * | Magic | Sequence | Command | UT | ChkSum | Len | // * +----------------+----------+----------+----------------+----------------+----------------+ // * | | // * + + // * | body bytes | // * + + // * | ... ... | // * +-----------------------------------------------------------------------------------------+ EchoPkgHeader struct { Magic uint32 //包头,协议版本号 Sequence uint16 // 请求序列 Command uint16 // 命令编号 UT uint32 // 客户端本地时间戳(秒) ChkSum int32 // 检验和 Len uint32 // 数据包长度 } EchoPackage struct { H EchoPkgHeader B json.RawMessage } ) func (p EchoPkgHeader) String() string { return fmt.Sprintf("请求ID:%d, 命令ID:%s, 数据包长度:%d", p.Sequence, (echoCommand(p.Command)).String(), p.Len) } // Encode 编码 func (codec *EchoPackage) Encode() ([]byte, error) { codec.H.Magic = echoPkgMagic codec.H.UT = uint32(time.Now().Unix()) buf := &bytes.Buffer{} codec.H.Len = uint32(len(codec.B)) err := binary.Write(buf, binary.LittleEndian, codec.H) if err != nil { return nil, err } buf.Write(codec.B) return buf.Bytes(), nil } // Decode 解码buffer func (codec *EchoPackage) DecodeBuffer(msg []byte) (err error) { if len(msg) < echoPkgHeaderLen { return ErrNotEnoughStream } hBuf := msg[:echoPkgHeaderLen] codec.H.Magic = binary.LittleEndian.Uint32(hBuf[0:4]) if codec.H.Magic != echoPkgMagic { logx.Errorf("@p.H.Magic{%x}, right magic{%x}", codec.H.Magic, echoPkgMagic) return ErrIllegalMagic } codec.H.Sequence = binary.LittleEndian.Uint16(hBuf[4:6]) codec.H.Command = binary.LittleEndian.Uint16(hBuf[6:8]) codec.H.UT = binary.LittleEndian.Uint32(hBuf[8:12]) codec.H.ChkSum = int32(binary.LittleEndian.Uint32(hBuf[12:16])) codec.H.Len = binary.LittleEndian.Uint32(hBuf[16:20]) // 防止恶意客户端把这个字段设置过大导致服务端死等或者服务端在准备对应的缓冲区时内存崩溃 if maxEchoStringLen < codec.H.Len { return ErrTooLargePackage } if len(msg)-echoPkgHeaderLen < int(codec.H.Len) { return ErrNotEnoughStream } bBuf := msg[echoPkgHeaderLen:codec.H.Len] codec.B = make([]byte, codec.H.Len) copy(codec.B, bBuf) return nil } // Decode 解码gnet func (codec *EchoPackage) DecodeGnet(c gnet.Conn) (err error) { if c.InboundBuffered() < echoPkgHeaderLen { return ErrNotEnoughStream } hBuf, _ := c.Peek(echoPkgHeaderLen) codec.H.Magic = binary.LittleEndian.Uint32(hBuf[0:4]) if codec.H.Magic != echoPkgMagic { logx.Errorf("@p.H.Magic{%x}, right magic{%x}", codec.H.Magic, echoPkgMagic) return ErrIllegalMagic } codec.H.Sequence = binary.LittleEndian.Uint16(hBuf[4:6]) codec.H.Command = binary.LittleEndian.Uint16(hBuf[6:8]) codec.H.UT = binary.LittleEndian.Uint32(hBuf[8:12]) codec.H.ChkSum = int32(binary.LittleEndian.Uint32(hBuf[12:16])) codec.H.Len = binary.LittleEndian.Uint32(hBuf[16:20]) // 防止恶意客户端把这个字段设置过大导致服务端死等或者服务端在准备对应的缓冲区时内存崩溃 if maxEchoStringLen < codec.H.Len { return ErrTooLargePackage } msgLen := echoPkgHeaderLen + int(codec.H.Len) if c.InboundBuffered() < msgLen { return ErrNotEnoughStream } _, _ = c.Discard(echoPkgHeaderLen) bBuf, _ := c.Peek(int(codec.H.Len)) codec.B = make([]byte, codec.H.Len) copy(codec.B, bBuf) _, _ = c.Discard(int(codec.H.Len)) return nil } // Decode 解码wsutil func (codec *EchoPackage) DecodeWsMessage(msg wsutil.Message) (err error) { if len(msg.Payload) < echoPkgHeaderLen { return ErrNotEnoughStream } hBuf := msg.Payload[:echoPkgHeaderLen] codec.H.Magic = binary.LittleEndian.Uint32(hBuf[0:4]) if codec.H.Magic != echoPkgMagic { logx.Errorf("@p.H.Magic{%x}, right magic{%x}", codec.H.Magic, echoPkgMagic) return ErrIllegalMagic } codec.H.Sequence = binary.LittleEndian.Uint16(hBuf[4:6]) codec.H.Command = binary.LittleEndian.Uint16(hBuf[6:8]) codec.H.UT = binary.LittleEndian.Uint32(hBuf[8:12]) codec.H.ChkSum = int32(binary.LittleEndian.Uint32(hBuf[12:16])) codec.H.Len = binary.LittleEndian.Uint32(hBuf[16:20]) // 防止恶意客户端把这个字段设置过大导致服务端死等或者服务端在准备对应的缓冲区时内存崩溃 if maxEchoStringLen < codec.H.Len { return ErrTooLargePackage } msgLen := echoPkgHeaderLen + int(codec.H.Len) if len(msg.Payload)-echoPkgHeaderLen < msgLen { return ErrNotEnoughStream } bBuf := msg.Payload[echoPkgHeaderLen:codec.H.Len] codec.B = make([]byte, codec.H.Len) copy(codec.B, bBuf) return nil }