Moon's blog

write the code, change the world.

golang tcp拆包的正确姿势

最近在造一个叫im-go的服务，看名字也能猜出来，是一个基于Go的IM服务，因为不想引入任何的依赖库，所以是手写每个模块的。

之前看过Netty，于是也想做一个类似Netty Codec的，用于编码解码的模块, 方便地处理TCP粘包这种细节问题。

在网上做了一番搜索之后，发现排名靠前的实现，要么出乎意料地复杂，要么根本就是完全错误的，例如

出乎意料的复杂：

错误的：

分析一下这个错误的实现

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
func Decode(reader *bufio.Reader) (string, error) {
    lengthByte, _ := reader.Peek(4)
    lengthBuff := bytes.NewBuffer(lengthByte)
    var length int32
    err := binary.Read(lengthBuff, binary.LittleEndian, &length)
    if err != nil {
        return "", err
    }
    if int32(reader.Buffered()) < length+4 {
        return "", err
    }

    // 假设执行到了这里，那么已经成功读取了长度到length这个变量中
    pack := make([]byte, int(4+length))
    _, err = reader.Read(pack) //这里是不能保证就能完读到length长度的数据的!!
    if err != nil {
        return "", err
    }
    return string(pack[4:]), nil
}

我也受了它的误导，基于Peek()做了一个非常复杂的实现

正确的姿势

在翻了翻io和bufio这两个包之后，我找到了ReadFull

ReadFull，就是调用了ReadAtLeast

1
2
3
func ReadFull(r Reader, buf []byte) (n int, err error) {
    return ReadAtLeast(r, buf, len(buf))
}
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
func ReadAtLeast(r Reader, buf []byte, min int) (n int, err error) {
    if len(buf) < min {
        return 0, ErrShortBuffer
    }
    for n < min && err == nil {
        var nn int
        nn, err = r.Read(buf[n:])
        n += nn
    }
    if n >= min {
        err = nil
    } else if n > 0 && err == EOF {
        err = ErrUnexpectedEOF
    }
    return
}

标准库里的ReadAtLeast就非常优雅了，用n记录读取的总字节数，nn是每次读取到的字节数，一看就明白。

基于ReadFull的拆包代码

1
2
3
4
5
6
7
8
9
10
11
func (c *LenthCodec) Decode(conn net.Conn) (bodyBuf []byte, err error) {
    lengthBuf := make([]byte, 4)
    _, err = io.ReadFull(conn, lengthBuf)
    //check error
    length := binary.LittleEndian.Uint32(lengthBuf)
    
    bodyBuf = make([]byte, length)
    _, err = io.ReadFull(conn, bodyBuf)
    //check error
    return
}
golang造轮子