zogo/frame/frame.go

// Package frame 应用层帧协议：头 8 字节（版本1 + 命令1 + 长度2 + 流ID4，小端序）+ 数据，
// 提供帧与 net.Conn 间的读写封装。
//
// 迁移自 go-hua/common，迁移时修复两处问题：
//  1. SetData 超长时未返回错误（fmt.Errorf 结果未 return）；
//  2. ReadFrame 原实现用临时 bufio.Reader，帧头短读会错位、
//     一次 Peek 之外的缓冲字节会被丢弃（连续帧丢数据），现改用 io.ReadFull 精确读取。
package frame

import (
	"encoding/binary"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"net"
)

const FrameVersion byte = 1

// 帧头部
const (
	sizeOfVer    = 1
	sizeOfCmd    = 1
	sizeOfLength = 2
	sizeOfSid    = 4
	headerSize   = sizeOfVer + sizeOfCmd + sizeOfSid + sizeOfLength
)

const (
	frameBufSize = 0x10000
)

// Frame 应用层帧：头 8 字节（版本1 + 命令1 + 长度2 + 流ID4，小端序）+ 数据
type Frame struct {
	ver  byte
	cmd  byte
	sid  uint32
	data []byte
}

// NewFrame 创建帧，version 为协议版本，cmd 为命令字，sid 为流 ID
func NewFrame(version byte, cmd byte, sid uint32) Frame {
	return Frame{ver: version, cmd: cmd, sid: sid}
}

// ReadFrame 从连接读取一帧（阻塞直到完整帧读出；头/体不完整或连接中断返回错误）
func ReadFrame(conn net.Conn) (Frame, error) {
	f := NewFrame(0, 0, 0)
	hdr := RawHeader{}
	// 读取头（io.ReadFull 防止头部被 TCP 分段导致的短读）
	if _, err := io.ReadFull(conn, hdr[:]); err != nil {
		return f, err
	}

	f.ver = hdr.Version()
	f.cmd = hdr.Cmd()
	f.sid = hdr.StreamID()
	// 读取数据
	data := make([]byte, hdr.Length())
	if _, err := io.ReadFull(conn, data); err != nil {
		return f, err
	}
	f.data = data
	return f, nil
}

// WriteFrame 将帧写入连接
func WriteFrame(conn net.Conn, frame Frame) error {
	hdr := RawHeader{}
	hdr[0] = frame.ver
	hdr[1] = frame.cmd
	binary.LittleEndian.PutUint16(hdr[2:], uint16(len(frame.data)))
	binary.LittleEndian.PutUint32(hdr[4:], uint32(frame.sid))

	buf := hdr.Bytes()
	buf = append(buf, frame.data...)
	_, err := conn.Write(buf)
	if err != nil {
		return err
	}
	return nil
}

// Version 返回协议版本
func (f *Frame) Version() byte {
	return f.ver
}

// Cmd 返回命令字
func (f *Frame) Cmd() byte {
	return f.cmd
}

// Length 返回数据长度
func (f *Frame) Length() uint16 {
	return uint16(len(f.data))
}

// Data 返回数据字节集
func (f *Frame) Data() []byte {
	return f.data
}

// StreamID 返回流 ID
func (f *Frame) StreamID() uint32 {
	return f.sid
}

// SetData 设置数据（超出 65535 字节返回错误）
func (f *Frame) SetData(data []byte) error {
	if len(data) > frameBufSize-1 {
		return fmt.Errorf("超出最大长度:%d", frameBufSize-1)
	}
	f.data = data
	return nil
}

// Marshal 将 v 序列化为 JSON 后写入数据
func (f *Frame) Marshal(v any) error {
	buf, err := json.Marshal(v)
	if err != nil {
		return err
	}
	return f.SetData(buf)
}

// Unmarshal 将数据按 JSON 反序列化为 v
func (f *Frame) Unmarshal(v any) error {
	if f.data == nil || len(f.data) == 0 {
		return errors.New("没有数据")
	}
	return json.Unmarshal(f.data, v)
}

// String 返回帧的描述字符串
func (f *Frame) String() string {
	return fmt.Sprintf("Version:%d Cmd:%d StreamID:%d Length:%d",
		f.Version(), f.Cmd(), f.StreamID(), f.Length())
}

// RawHeader 原始帧头（8 字节数组）
type RawHeader [headerSize]byte

// Version 返回协议版本
func (h RawHeader) Version() byte {
	return h[0]
}

// Cmd 返回命令字
func (h RawHeader) Cmd() byte {
	return h[1]
}

// Length 返回数据长度
func (h RawHeader) Length() uint16 {
	return binary.LittleEndian.Uint16(h[2:])
}

// StreamID 返回流 ID
func (h RawHeader) StreamID() uint32 {
	return binary.LittleEndian.Uint32(h[4:])
}

// Bytes 返回帧头字节集
func (h RawHeader) Bytes() []byte {
	buf := make([]byte, headerSize)
	for i, _ := range buf {
		buf[i] = h[i]
	}
	return buf
}