86 lines
2.0 KiB
Go
86 lines
2.0 KiB
Go
package codec
|
||
|
||
import (
|
||
"encoding/binary"
|
||
"errors"
|
||
"io"
|
||
|
||
"github.com/panjf2000/gnet/v2"
|
||
)
|
||
|
||
var ErrIncompletePacket = errors.New("incomplete packet")
|
||
|
||
// TomeeSocketCodec 协议格式:
|
||
//
|
||
// * 0 4
|
||
// * +-----------+
|
||
// * | body len |
|
||
// * +-----------+
|
||
// * | |
|
||
// * + +
|
||
// * | body bytes|
|
||
// * + +
|
||
// * | ... ... |
|
||
// * +-----------+
|
||
type TomeeSocketCodec struct{}
|
||
|
||
var _ SocketCodec = (*TomeeSocketCodec)(nil)
|
||
|
||
func NewTomeeSocketCodec() *TomeeSocketCodec {
|
||
return &TomeeSocketCodec{}
|
||
|
||
}
|
||
|
||
func (codec TomeeSocketCodec) Encode(buf []byte) ([]byte, error) {
|
||
bodyLen := len(buf)
|
||
data := make([]byte, 4+bodyLen)
|
||
|
||
// 写入4字节的包长度
|
||
binary.BigEndian.PutUint32(data[:4], uint32(bodyLen))
|
||
// 写入包体
|
||
copy(data[4:], buf)
|
||
|
||
return data, nil
|
||
}
|
||
|
||
func (codec TomeeSocketCodec) Decode(c gnet.Conn) ([]byte, error) {
|
||
const maxBodyLen = 10 * 1024 * 1024 // 业务最大包体长度，按需调整
|
||
|
||
// 1. 读取4字节长度头
|
||
lenBuf, err := c.Peek(4)
|
||
if err != nil {
|
||
if errors.Is(err, io.ErrShortBuffer) {
|
||
return nil, ErrIncompletePacket
|
||
}
|
||
return nil, err
|
||
}
|
||
|
||
// 2. 解析长度并校验上限
|
||
bodyLen := binary.BigEndian.Uint32(lenBuf)
|
||
if bodyLen > maxBodyLen {
|
||
return nil, errors.New("packet body exceeds max length")
|
||
}
|
||
totalLen := 4 + int(bodyLen) // 总包长=长度头(4)+包体
|
||
|
||
// 3. 检查整个包是否完整（关键修复：用totalLen判断）
|
||
fullBuf, err := c.Peek(totalLen)
|
||
if err != nil {
|
||
if errors.Is(err, io.ErrShortBuffer) {
|
||
return nil, ErrIncompletePacket
|
||
}
|
||
return nil, err
|
||
}
|
||
|
||
// 4. 提取包体（关键修复：跳过前4字节长度头）
|
||
body := make([]byte, bodyLen)
|
||
copy(body, fullBuf[4:]) // 从fullBuf[4:]开始拷贝，仅取包体
|
||
|
||
// 5. 丢弃已解析的完整数据
|
||
_, err = c.Discard(totalLen) // 原代码用_忽略err，建议处理（可选）
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
|
||
return body, nil
|
||
}
|