0 Star 1 Fork 0

蒋佳李 / leaf

加入 Gitee
与超过 1200万 开发者一起发现、参与优秀开源项目,私有仓库也完全免费 :)
免费加入
克隆/下载
tcp_msg.go 3.13 KB
一键复制 编辑 原始数据 按行查看 历史
蒋佳李 提交于 2023-03-01 10:43 . 添加判断是否加密传输数据
package network
import (
"encoding/binary"
"errors"
"gitee.com/jiangjiali/leaf/conf"
"gitee.com/jiangjiali/leaf/sdk/sm4"
"io"
"math"
)
// MsgParser
// --------------
// | len | data |
// --------------
type MsgParser struct {
lenMsgLen int
minMsgLen uint32
maxMsgLen uint32
littleEndian bool
}
func NewMsgParser() *MsgParser {
p := new(MsgParser)
p.lenMsgLen = 2
p.minMsgLen = 1
p.maxMsgLen = 4096
p.littleEndian = false
return p
}
// SetMsgLen It's dangerous to call the method on reading or writing
func (p *MsgParser) SetMsgLen(lenMsgLen int, minMsgLen uint32, maxMsgLen uint32) {
if lenMsgLen == 1 || lenMsgLen == 2 || lenMsgLen == 4 {
p.lenMsgLen = lenMsgLen
}
if minMsgLen != 0 {
p.minMsgLen = minMsgLen
}
if maxMsgLen != 0 {
p.maxMsgLen = maxMsgLen
}
var max uint32
switch p.lenMsgLen {
case 1:
max = math.MaxUint8
case 2:
max = math.MaxUint16
case 4:
max = math.MaxUint32
}
if p.minMsgLen > max {
p.minMsgLen = max
}
if p.maxMsgLen > max {
p.maxMsgLen = max
}
}
// SetByteOrder It's dangerous to call the method on reading or writing
func (p *MsgParser) SetByteOrder(littleEndian bool) {
p.littleEndian = littleEndian
}
// goroutine safe
func (p *MsgParser) Read(conn *TCPConn) ([]byte, error) {
var b [4]byte
bufMsgLen := b[:p.lenMsgLen]
// read len
if _, err := io.ReadFull(conn, bufMsgLen); err != nil {
return nil, err
}
// parse len
var msgLen uint32
switch p.lenMsgLen {
case 1:
msgLen = uint32(bufMsgLen[0])
case 2:
if p.littleEndian {
msgLen = uint32(binary.LittleEndian.Uint16(bufMsgLen))
} else {
msgLen = uint32(binary.BigEndian.Uint16(bufMsgLen))
}
case 4:
if p.littleEndian {
msgLen = binary.LittleEndian.Uint32(bufMsgLen)
} else {
msgLen = binary.BigEndian.Uint32(bufMsgLen)
}
}
// check len
if msgLen > p.maxMsgLen {
return nil, errors.New("message too long")
} else if msgLen < p.minMsgLen {
return nil, errors.New("message too short")
}
// data
msgData := make([]byte, msgLen)
if _, err := io.ReadFull(conn, msgData); err != nil {
return nil, err
}
return msgData, nil
}
// goroutine safe
func (p *MsgParser) Write(conn *TCPConn, args ...[]byte) error {
// get len
var msgLen uint32
for i := 0; i < len(args); i++ {
msgLen += uint32(len(args[i]))
}
// check len
if msgLen > p.maxMsgLen {
return errors.New("message too long")
} else if msgLen < p.minMsgLen {
return errors.New("message too short")
}
msg := make([]byte, uint32(p.lenMsgLen)+msgLen)
// write len
switch p.lenMsgLen {
case 1:
msg[0] = byte(msgLen)
case 2:
if p.littleEndian {
binary.LittleEndian.PutUint16(msg, uint16(msgLen))
} else {
binary.BigEndian.PutUint16(msg, uint16(msgLen))
}
case 4:
if p.littleEndian {
binary.LittleEndian.PutUint32(msg, msgLen)
} else {
binary.BigEndian.PutUint32(msg, msgLen)
}
}
// write data
l := p.lenMsgLen
for i := 0; i < len(args); i++ {
copy(msg[l:], args[i])
l += len(args[i])
}
if conf.IsEncode {
// 加密发送的信息
cbcMsg, err := sm4.Cbc([]byte(conf.SM4Key), msg, true)
if err != nil {
return err
}
conn.Write(cbcMsg)
return nil
}
conn.Write(msg)
return nil
}
Go
1
https://gitee.com/jiangjiali/leaf.git
git@gitee.com:jiangjiali/leaf.git
jiangjiali
leaf
leaf
v1.1.44

搜索帮助