178 lines
4.3 KiB
Go
178 lines
4.3 KiB
Go
|
package hpds_net_framework
|
||
|
|
||
|
import (
|
||
|
"encoding/binary"
|
||
|
"fmt"
|
||
|
"git.hpds.cc/Component/logging"
|
||
|
"github.com/golang/protobuf/proto"
|
||
|
"go.uber.org/zap"
|
||
|
"reflect"
|
||
|
)
|
||
|
|
||
|
// PbProcessor one of Processor implement protoc --go_out=. *.proto
|
||
|
type PbProcessor struct {
|
||
|
bigEndian bool
|
||
|
enc Encryptor
|
||
|
msgTypes map[reflect.Type]int
|
||
|
handlers map[int]msgInfo
|
||
|
logger *logging.Logger
|
||
|
}
|
||
|
|
||
|
// NewPbProcessor return PB processor
|
||
|
func NewPbProcessor(logger *logging.Logger) *PbProcessor {
|
||
|
pb := PbProcessor{
|
||
|
msgTypes: make(map[reflect.Type]int),
|
||
|
handlers: make(map[int]msgInfo),
|
||
|
logger: logger,
|
||
|
}
|
||
|
return &pb
|
||
|
}
|
||
|
|
||
|
// OnReceivedPackage 收到完整数据包, 返回解包错误
|
||
|
func (pbf *PbProcessor) OnReceivedPackage(writer interface{}, body []byte) error {
|
||
|
// 解密
|
||
|
if pbf.enc != nil {
|
||
|
//log.Debug("before decode:: %v", body)
|
||
|
body = pbf.enc.Decode(body)
|
||
|
//log.Debug("after decode:: %v", body)
|
||
|
}
|
||
|
// 解码
|
||
|
var pack Protocol
|
||
|
if err := proto.UnmarshalMerge(body, &pack); err != nil {
|
||
|
pbf.logger.Error("Can't unmarshal pack body to protocolMsg",
|
||
|
zap.ByteString("body", body),
|
||
|
zap.Error(err),
|
||
|
)
|
||
|
return err
|
||
|
}
|
||
|
info, ok := pbf.handlers[int(pack.Id)]
|
||
|
if !ok {
|
||
|
pbf.logger.Error("Not register msg id",
|
||
|
zap.Uint32("pack Id", pack.Id),
|
||
|
)
|
||
|
// handler not found, not a dead err
|
||
|
return nil
|
||
|
}
|
||
|
msg := reflect.New(info.msgType.Elem()).Interface()
|
||
|
err := proto.UnmarshalMerge(pack.Content, msg.(proto.Message))
|
||
|
if err != nil {
|
||
|
pbf.logger.Error("Unmarshal Merge pack contents error",
|
||
|
zap.Uint32("pack id", pack.Id),
|
||
|
zap.Error(err),
|
||
|
)
|
||
|
return err
|
||
|
}
|
||
|
// 执行逻辑
|
||
|
execute(info, msg, writer, body, pack.Id)
|
||
|
return nil
|
||
|
}
|
||
|
|
||
|
// WrapMsg format the interface message to []byte
|
||
|
func (pbf *PbProcessor) WrapMsg(message interface{}) ([]byte, error) {
|
||
|
pbf.logger.Debug("Protobuf processor wrap for write",
|
||
|
zap.Any("reflect.TypeOf", reflect.TypeOf(message)),
|
||
|
)
|
||
|
data, err := proto.Marshal(message.(proto.Message))
|
||
|
if err != nil {
|
||
|
return nil, err
|
||
|
}
|
||
|
tp := reflect.TypeOf(message)
|
||
|
id, ok := pbf.msgTypes[tp]
|
||
|
if !ok {
|
||
|
return nil, fmt.Errorf("not register %v", tp)
|
||
|
}
|
||
|
protocol := Protocol{
|
||
|
Id: uint32(id),
|
||
|
Content: data,
|
||
|
}
|
||
|
data, err = proto.Marshal(&protocol)
|
||
|
if err != nil {
|
||
|
return nil, err
|
||
|
}
|
||
|
if pbf.enc != nil {
|
||
|
//log.Debug("before encode:: %v", data)
|
||
|
data = pbf.enc.Encode(data)
|
||
|
//log.Debug("after encode:: %v", data)
|
||
|
}
|
||
|
// head
|
||
|
head := make([]byte, 2)
|
||
|
if pbf.bigEndian {
|
||
|
binary.BigEndian.PutUint16(head, uint16(len(data)))
|
||
|
} else {
|
||
|
binary.LittleEndian.PutUint16(head, uint16(len(data)))
|
||
|
}
|
||
|
pkg := append(head, data...)
|
||
|
return pkg, nil
|
||
|
}
|
||
|
|
||
|
// WrapIdMsg format the interface message to []byte with id
|
||
|
func (pbf *PbProcessor) WrapIdMsg(id uint32, message interface{}) ([]byte, error) {
|
||
|
pbf.logger.Debug("Protobuf processor wrap for write",
|
||
|
zap.Any("reflect.TypeOf", reflect.TypeOf(message)),
|
||
|
)
|
||
|
data, err := proto.Marshal(message.(proto.Message))
|
||
|
if err != nil {
|
||
|
return nil, err
|
||
|
}
|
||
|
protocol := Protocol{
|
||
|
Id: id,
|
||
|
Content: data,
|
||
|
}
|
||
|
data, err = proto.Marshal(&protocol)
|
||
|
if err != nil {
|
||
|
return nil, err
|
||
|
}
|
||
|
if pbf.enc != nil {
|
||
|
//log.Debug("before encode:: %v", data)
|
||
|
data = pbf.enc.Encode(data)
|
||
|
//log.Debug("after encode:: %v", data)
|
||
|
}
|
||
|
// head
|
||
|
head := make([]byte, 2)
|
||
|
if pbf.bigEndian {
|
||
|
binary.BigEndian.PutUint16(head, uint16(len(data)))
|
||
|
} else {
|
||
|
binary.LittleEndian.PutUint16(head, uint16(len(data)))
|
||
|
}
|
||
|
pkg := append(head, data...)
|
||
|
return pkg, nil
|
||
|
}
|
||
|
|
||
|
// WrapMsgNoHeader without header length
|
||
|
func (pbf *PbProcessor) WrapMsgNoHeader(message interface{}) ([]byte, error) {
|
||
|
data, err := pbf.WrapMsg(message)
|
||
|
if err != nil {
|
||
|
return nil, err
|
||
|
}
|
||
|
return data[2:], nil
|
||
|
}
|
||
|
|
||
|
// RegisterHandler for logic
|
||
|
func (pbf *PbProcessor) RegisterHandler(id int, entity interface{}, handle func(args ...interface{})) {
|
||
|
if _, ok := pbf.handlers[id]; ok {
|
||
|
pbf.logger.Error("Already register handler", zap.Int("ID", id))
|
||
|
} else {
|
||
|
pbf.handlers[id] = msgInfo{
|
||
|
msgId: id,
|
||
|
msgType: reflect.TypeOf(entity),
|
||
|
msgCallback: handle,
|
||
|
}
|
||
|
pbf.msgTypes[reflect.TypeOf(entity)] = id
|
||
|
}
|
||
|
}
|
||
|
|
||
|
// SetBigEndian for order
|
||
|
func (pbf *PbProcessor) SetBigEndian() {
|
||
|
pbf.bigEndian = true
|
||
|
}
|
||
|
|
||
|
// GetBigEndian of the order
|
||
|
func (pbf *PbProcessor) GetBigEndian() bool {
|
||
|
return pbf.bigEndian
|
||
|
}
|
||
|
|
||
|
// SetEncryptor for processor
|
||
|
func (pbf *PbProcessor) SetEncryptor(enc Encryptor) {
|
||
|
pbf.enc = enc
|
||
|
}
|