Go1.7.3一個收發消息的小接口

package transport

import (
    "encoding/binary"
    "errors"
    "io"
    "math"
    "sync"
)

const (
    sizeLen = 5 //設置消息多少字節大小,1<<4*8-1長度,最後一個字節用來做消息分割的,防止惡意鏈接
)

var (
    ErrLength = errors.New("message too large")
    ErrHeader = errors.New("message head error")
)

//添加數據頭併發送
func SedMsg(w io.Writer, p []byte) error {
    msglen := len(p)
    if msglen > math.MaxUint32 {
        return ErrLength
    }
    var buf = make([]byte, sizeLen+msglen)
    binary.BigEndian.PutUint32(buf, uint32(msglen))
    copy(buf[sizeLen:], p)
    _, err := w.Write(buf)
    return err
}

//從數據流中根據頭長度讀取完整的消息
func RcvMsg(r io.Reader, lock sync.Locker) ([]byte, error) {
    lock.Lock()
    defer lock.Unlock()

    var header [sizeLen]byte
    if _, err := io.ReadFull(r, header[:]); err != nil {
        return nil, err
    }

    if header[sizeLen-1] == 0 {
        length := binary.BigEndian.Uint32(header[:])
        if length == 0 {
            return nil, nil
        }
        msg := make([]byte, int(length))
        if _, err := io.ReadFull(r, msg); err != nil {
            if err == io.EOF {
                err = io.ErrUnexpectedEOF
            }
            msg = msg[:0]
            return nil, err
        }
        return msg, nil
    }

    return nil, ErrHeader
}
發表評論
所有評論
還沒有人評論,想成為第一個評論的人麼? 請在上方評論欄輸入並且點擊發布.
相關文章