package unpacker import ( unpackerpkg "git.noahlan.cn/noahlan/nnet/v2/pkg/unpacker" ) // frameHeaderUnpacker 帧头拆包器实现 type frameHeaderUnpacker struct { headerLength int getLength func(header []byte) int buffer []byte maxBufferSize int } // NewFrameHeaderUnpacker 创建帧头拆包器 func NewFrameHeaderUnpacker(config unpackerpkg.FrameHeaderUnpacker) unpackerpkg.Unpacker { headerLength := config.HeaderLength if headerLength <= 0 { headerLength = 4 // 默认4字节 } getLength := config.GetLength if getLength == nil { // 默认实现:假设前4字节是长度 getLength = func(header []byte) int { if len(header) < 4 { return 0 } return int(header[0])<<24 | int(header[1])<<16 | int(header[2])<<8 | int(header[3]) } } maxBufferSize := config.MaxBufferSize if maxBufferSize <= 0 { maxBufferSize = unpackerpkg.DefaultMaxBufferSize } return &frameHeaderUnpacker{ headerLength: headerLength, getLength: getLength, buffer: make([]byte, 0, 4096), // 预分配初始容量 maxBufferSize: maxBufferSize, } } // Unpack 拆包 func (u *frameHeaderUnpacker) Unpack(data []byte) ([][]byte, []byte, int, error) { // 检查buffer大小限制 newSize := len(u.buffer) + len(data) if newSize > u.maxBufferSize { return nil, nil, 0, unpackerpkg.NewErrorf("unpacker buffer size exceeded: %d > %d", newSize, u.maxBufferSize) } // 优化:如果容量不足,预分配更大的容量(零拷贝优化) if cap(u.buffer) < newSize { newCap := cap(u.buffer) * 2 if newCap < newSize { newCap = newSize } if newCap > u.maxBufferSize { newCap = u.maxBufferSize } // 如果现有 buffer 为空,直接分配新 buffer(避免不必要的复制) if len(u.buffer) == 0 { u.buffer = make([]byte, 0, newCap) } else { newBuffer := make([]byte, len(u.buffer), newCap) copy(newBuffer, u.buffer) u.buffer = newBuffer } } u.buffer = append(u.buffer, data...) var messages [][]byte for { if len(u.buffer) < u.headerLength { // 数据不足,等待更多数据 break } // 读取帧头 header := u.buffer[:u.headerLength] messageLength := u.getLength(header) if messageLength <= 0 { // 无效长度,跳过 u.buffer = u.buffer[1:] continue } // 验证长度字段的合理性(防止恶意数据) if messageLength > u.maxBufferSize { return nil, nil, 0, unpackerpkg.NewErrorf("invalid message length: %d > %d", messageLength, u.maxBufferSize) } totalLength := u.headerLength + messageLength if totalLength > u.maxBufferSize { return nil, nil, 0, unpackerpkg.NewErrorf("message too large: %d > %d", totalLength, u.maxBufferSize) } if len(u.buffer) < totalLength { // 数据不足,等待更多数据 break } // 提取消息 message := make([]byte, totalLength) copy(message, u.buffer[:totalLength]) messages = append(messages, message) // 移除已处理的数据(优化:使用切片操作,避免复制) u.buffer = u.buffer[totalLength:] // 如果 buffer 太大但剩余数据很少,压缩 buffer(减少内存占用) // 注意:压缩不会改变buffer的长度,只改变容量 if len(u.buffer) < cap(u.buffer)/4 && cap(u.buffer) > 4096 { compressed := make([]byte, len(u.buffer), cap(u.buffer)/2) copy(compressed, u.buffer) u.buffer = compressed } } // 输入字节已经被复制进连接级buffer;调用方应从底层读缓冲中丢弃本次输入,避免重复处理。 return messages, u.buffer, len(data), nil } // Pack 打包 func (u *frameHeaderUnpacker) Pack(data []byte) ([]byte, error) { messageLength := len(data) header := make([]byte, u.headerLength) // 写入长度到帧头 if u.headerLength >= 4 { header[0] = byte(messageLength >> 24) header[1] = byte(messageLength >> 16) header[2] = byte(messageLength >> 8) header[3] = byte(messageLength) } result := make([]byte, u.headerLength+messageLength) copy(result, header) copy(result[u.headerLength:], data) return result, nil }