Golang实现自定义TCP协议

scanner.go

package main

import (
	"bufio"
	"bytes"
	"encoding/binary"
	"fmt"
	"io"
)

type Package struct {
	Version [2]int8
	Length  int16
	Data    []byte
}

func (p *Package) String() string {
	return fmt.Sprintf("Version:%d Length:%d Data:%s", p.Version, p.Length, p.Data)
}

// 打包
func (p *Package) Pack(writer io.Writer) {
	binary.Write(writer, binary.BigEndian, p.Version)
	binary.Write(writer, binary.BigEndian, p.Length)
	binary.Write(writer, binary.BigEndian, p.Data)
}

// 解包
func (p *Package) Unpack(r io.Reader) {
	binary.Read(r, binary.BigEndian, &p.Version)
	binary.Read(r, binary.BigEndian, &p.Length)
	if p.Length > 0 {
		p.Data = make([]byte, p.Length)
	}
	binary.Read(r, binary.BigEndian, &p.Data)
}

func main() {
	pkg := &Package{
		Version: [2]int8{'V', 1},
		Data:    []byte("ABCDEFGHIJK"),
	}
	pkg.Length = int16(len(pkg.Data))
	buf := new(bytes.Buffer)
	pkg.Pack(buf)
	pkg.Pack(buf)
	pkg.Pack(buf)
	// scanner
	scanner := bufio.NewScanner(buf)

	scanner.Split(func(data []byte, atEOF bool) (advance int, token []byte, err error) {
		if !atEOF && data[0] == 'V' {
			if len(data) > 4 { // 2字节版本号+2字节数据长度,最少4字节
				length := int16(0)                                                 // 数据包长度
				binary.Read(bytes.NewReader(data[2:4]), binary.BigEndian, &length) // 第3-4字节为数据包长度
				if int(length)+4 <= len(data) { // 数据已经读完了，但是产生了粘包现象
					return int(length) + 4, data[:int(length)+4], nil
				}
			}
		}
		return
	})

	for scanner.Scan() {
		pack := new(Package)
		pack.Unpack(bytes.NewReader(scanner.Bytes()))
		fmt.Println(pkg)
	}
	if err := scanner.Err(); err != nil {
		fmt.Printf("error:%s", err)
	}
}