280 lines
8.7 KiB
Go
280 lines
8.7 KiB
Go
package protocol
|
||
|
||
import (
|
||
"encoding/binary"
|
||
"encoding/json"
|
||
"errors"
|
||
"fmt"
|
||
"io"
|
||
"unicode/utf8"
|
||
)
|
||
|
||
// MaxFrameSize 用于限制单个帧的最大长度，
|
||
// 避免异常对端通过伪造超大长度值导致接收方无上限分配内存。
|
||
const MaxFrameSize = 8 * 1024 * 1024 // 先临时设置传输的视频帧不超过8MB
|
||
|
||
var (
|
||
ErrInvalidFrameLength = errors.New("protocol: invalid frame length") // 表示帧长度非法，例如长度为 0。
|
||
ErrFrameTooLarge = errors.New("protocol: frame too large") // 表示帧长度超过允许的上限。
|
||
ErrInvalidMessageType = errors.New("protocol: invalid message type") // 表示消息类型不是当前协议支持的类型。
|
||
ErrMissingFrom = errors.New("protocol: missing from") // 表示消息缺少发送方标识。
|
||
ErrMissingTo = errors.New("protocol: missing to") // 表示消息缺少接收方标识。
|
||
ErrMissingFileName = errors.New("protocol: missing file name") // 表示 file 消息缺少文件名。
|
||
ErrUnexpectedFileName = errors.New("protocol: unexpected file name") // 表示 text 消息错误地携带了文件名。
|
||
ErrInvalidTextBody = errors.New("protocol: invalid text body") // 表示 text 消息正文不是合法 UTF-8。
|
||
ErrUnexpectedBody = errors.New("protocol: unexpected body") // 表示某些控制消息不允许携带正文。
|
||
ErrInvalidRegisterTarget = errors.New("protocol: invalid register target") // 表示 register 消息没有发往 server。
|
||
ErrInvalidErrorSource = errors.New("protocol: invalid error source") // 表示 error 消息不是由 server 发出。
|
||
ErrInvalidHeaderLength = errors.New("protocol: invalid header length") // 表示 header 长度字段为 0、越界或无法完整切分。
|
||
ErrInvalidHeaderJSON = errors.New("protocol: invalid header json") // 表示 header JSON 无法解析，可能是格式错误或缺少必要字段。
|
||
ErrInvalidContentLength = errors.New("protocol: invalid content length") // 表示头部记录的正文长度与实际正文不一致。
|
||
)
|
||
|
||
// 应用层消息：[4字节 frameLength][4字节 headerLen][header JSON（下面自定义的Message头）][body bytes]
|
||
// 写了 tag：JSON 字段名是你指定的 type；不写 tag：JSON 字段名默认是 Go 字段名 Type
|
||
type messageHeader struct {
|
||
Type MessageType `json:"type"`
|
||
ID uint64 `json:"id"`
|
||
From string `json:"from"`
|
||
To string `json:"to"`
|
||
FileName string `json:"file_name,omitempty"`
|
||
ContentLength int `json:"content_length"`
|
||
}
|
||
|
||
// EncodeMessage 将逻辑消息编码为帧内字节格式：
|
||
// 1. 4 字节大端序 header 长度
|
||
// 2. header JSON
|
||
// 3. 原始 body 字节
|
||
func EncodeMessage(msg Message) ([]byte, error) {
|
||
if err := validateMessage(msg); err != nil {
|
||
return nil, err
|
||
}
|
||
|
||
header := messageHeader{
|
||
Type: msg.Type,
|
||
ID: msg.ID,
|
||
From: msg.From,
|
||
To: msg.To,
|
||
FileName: msg.FileName,
|
||
ContentLength: len(msg.Body),
|
||
}
|
||
|
||
headerPayload, err := json.Marshal(header)
|
||
if err != nil {
|
||
return nil, fmt.Errorf("protocol: encode header: %w", err)
|
||
}
|
||
// 创建一个新的字节切片来存储完整的帧内容，避免直接在 headerPayload 上修改导致数据混乱。
|
||
payload := make([]byte, 4+len(headerPayload)+len(msg.Body))
|
||
// 在 payload 前 4 字节写入 header 长度，后续内容依次是 header JSON（第五个字节开始） 和 body。
|
||
binary.BigEndian.PutUint32(payload[:4], uint32(len(headerPayload)))
|
||
copy(payload[4:], headerPayload)
|
||
copy(payload[4+len(headerPayload):], msg.Body)
|
||
|
||
//检查整个帧长度是否合法，避免上层调用者构造的消息过大导致发送失败。
|
||
if len(payload) > MaxFrameSize {
|
||
return nil, ErrFrameTooLarge
|
||
}
|
||
|
||
return payload, nil
|
||
}
|
||
|
||
// DecodeMessage 将帧内字节格式还原为 Message。
|
||
func DecodeMessage(data []byte) (Message, error) {
|
||
if len(data) > MaxFrameSize {
|
||
return Message{}, ErrFrameTooLarge
|
||
}
|
||
if len(data) < 4 {
|
||
return Message{}, ErrInvalidHeaderLength
|
||
}
|
||
|
||
headerLen := int(binary.BigEndian.Uint32(data[:4]))
|
||
if headerLen == 0 || headerLen > len(data)-4 {
|
||
return Message{}, ErrInvalidHeaderLength
|
||
}
|
||
|
||
headerPayload := data[4 : 4+headerLen]
|
||
body := data[4+headerLen:]
|
||
|
||
var header messageHeader
|
||
if err := json.Unmarshal(headerPayload, &header); err != nil {
|
||
return Message{}, fmt.Errorf("protocol: decode header: %w", errors.Join(ErrInvalidHeaderJSON, err))
|
||
}
|
||
|
||
if header.ContentLength < 0 || header.ContentLength != len(body) {
|
||
return Message{}, ErrInvalidContentLength
|
||
}
|
||
|
||
bodyCopy := make([]byte, len(body))
|
||
copy(bodyCopy, body)
|
||
|
||
msg := Message{
|
||
Type: header.Type,
|
||
ID: header.ID,
|
||
From: header.From,
|
||
To: header.To,
|
||
FileName: header.FileName,
|
||
Body: bodyCopy,
|
||
}
|
||
|
||
if err := validateMessage(msg); err != nil {
|
||
return Message{}, err
|
||
}
|
||
|
||
return msg, nil
|
||
}
|
||
|
||
// WriteFrame 向流中写入一个带长度前缀的帧。
|
||
// TCP帧格式如下：
|
||
// 1. 4 字节大端序长度
|
||
// 2. 后续 payload 内容
|
||
//
|
||
// TCP 是字节流协议，没有天然的消息边界。
|
||
// 增加显式长度前缀后，接收方就知道一条完整消息应该读取多少字节，
|
||
// 从而解决粘包和拆包问题。
|
||
func WriteFrame(w io.Writer, payload []byte) error {
|
||
size := len(payload)
|
||
//空帧
|
||
if size == 0 {
|
||
return ErrInvalidFrameLength
|
||
}
|
||
//帧过大
|
||
if size > MaxFrameSize {
|
||
return ErrFrameTooLarge
|
||
}
|
||
|
||
var header [4]byte
|
||
binary.BigEndian.PutUint32(header[:], uint32(size))
|
||
|
||
// 先写长度头，接收方才能根据长度一次性读取完整消息体。
|
||
if err := writeFull(w, header[:]); err != nil {
|
||
return err
|
||
}
|
||
|
||
return writeFull(w, payload)
|
||
}
|
||
|
||
// ReadFrame 从流中读取一个完整的长度前缀帧。
|
||
// 它会先读取固定 4 字节长度头，校验长度是否合法，
|
||
// 再使用 io.ReadFull 按长度读取完整消息体，
|
||
// 这样即使底层 TCP 发生分段读取，也不会把半条消息暴露给上层。
|
||
func ReadFrame(r io.Reader) ([]byte, error) {
|
||
var header [4]byte
|
||
if _, err := io.ReadFull(r, header[:]); err != nil {
|
||
return nil, err
|
||
}
|
||
|
||
size := binary.BigEndian.Uint32(header[:])
|
||
// 长度为 0 的帧被认为是非法输入，而不是合法的空消息。
|
||
if size == 0 {
|
||
return nil, ErrInvalidFrameLength
|
||
}
|
||
// 长度超过上限的帧会被拒绝，避免接收方无上限分配内存。
|
||
if size > MaxFrameSize {
|
||
return nil, ErrFrameTooLarge
|
||
}
|
||
|
||
payload := make([]byte, int(size))
|
||
if _, err := io.ReadFull(r, payload); err != nil {
|
||
return nil, err
|
||
}
|
||
|
||
return payload, nil
|
||
}
|
||
|
||
// WriteMessage 是给上层直接使用的完整发送路径：
|
||
// 把一条结构化消息完整编码并发送出去”的总入口。
|
||
// Message -> header+body -> 长度前缀帧 -> io.Writer。
|
||
func WriteMessage(w io.Writer, msg Message) error {
|
||
payload, err := EncodeMessage(msg)
|
||
if err != nil {
|
||
return fmt.Errorf("protocol: encode message: %w", err)
|
||
}
|
||
|
||
if err := WriteFrame(w, payload); err != nil {
|
||
return fmt.Errorf("protocol: write frame: %w", err)
|
||
}
|
||
|
||
return nil
|
||
}
|
||
|
||
// ReadMessage 是给上层直接使用的完整接收路径：
|
||
// io.Reader -> 长度前缀帧 -> header+body -> Message。
|
||
func ReadMessage(r io.Reader) (Message, error) {
|
||
payload, err := ReadFrame(r)
|
||
if err != nil {
|
||
return Message{}, fmt.Errorf("protocol: read frame: %w", err)
|
||
}
|
||
|
||
msg, err := DecodeMessage(payload)
|
||
if err != nil {
|
||
return Message{}, fmt.Errorf("protocol: decode message: %w", err)
|
||
}
|
||
|
||
return msg, nil
|
||
}
|
||
|
||
// validateMessage 检查 Message 传输的类型（只接受 text 和 file ）。
|
||
func validateMessage(msg Message) error {
|
||
if msg.From == "" {
|
||
return ErrMissingFrom
|
||
}
|
||
if msg.To == "" {
|
||
return ErrMissingTo
|
||
}
|
||
|
||
switch msg.Type {
|
||
case MessageTypeText:
|
||
if msg.FileName != "" {
|
||
return ErrUnexpectedFileName
|
||
}
|
||
if !utf8.Valid(msg.Body) {
|
||
return ErrInvalidTextBody
|
||
}
|
||
case MessageTypeFile:
|
||
if msg.FileName == "" {
|
||
return ErrMissingFileName
|
||
}
|
||
case MessageTypeRegister:
|
||
if msg.To != ServerPeerID {
|
||
return ErrInvalidRegisterTarget
|
||
}
|
||
if msg.FileName != "" {
|
||
return ErrUnexpectedFileName
|
||
}
|
||
if len(msg.Body) != 0 {
|
||
return ErrUnexpectedBody
|
||
}
|
||
case MessageTypeError:
|
||
if msg.From != ServerPeerID {
|
||
return ErrInvalidErrorSource
|
||
}
|
||
if msg.FileName != "" {
|
||
return ErrUnexpectedFileName
|
||
}
|
||
if !utf8.Valid(msg.Body) {
|
||
return ErrInvalidTextBody
|
||
}
|
||
default:
|
||
return ErrInvalidMessageType
|
||
}
|
||
|
||
return nil
|
||
}
|
||
|
||
// writeFull 会持续写入，直到所有字节都写完或者底层返回错误。
|
||
// 这样可以避免某些 Writer 发生部分写入时破坏帧格式。
|
||
func writeFull(w io.Writer, data []byte) error {
|
||
for len(data) > 0 {
|
||
n, err := w.Write(data)
|
||
if err != nil {
|
||
return err
|
||
}
|
||
if n == 0 {
|
||
return io.ErrShortWrite
|
||
}
|
||
data = data[n:]
|
||
}
|
||
|
||
return nil
|
||
}
|