405 lines
9.1 KiB
Go
405 lines
9.1 KiB
Go
package algos
|
|
|
|
import (
|
|
"fmt"
|
|
)
|
|
|
|
// ProtoFieldType 协议字段类型
|
|
type ProtoFieldType int
|
|
|
|
const (
|
|
ProtoFieldTypeVARINT ProtoFieldType = 0
|
|
ProtoFieldTypeINT64 ProtoFieldType = 1
|
|
ProtoFieldTypeSTRING ProtoFieldType = 2
|
|
ProtoFieldTypeGROUPSTART ProtoFieldType = 3
|
|
ProtoFieldTypeGROUPEND ProtoFieldType = 4
|
|
ProtoFieldTypeINT32 ProtoFieldType = 5
|
|
ProtoFieldTypeERROR1 ProtoFieldType = 6
|
|
ProtoFieldTypeERROR2 ProtoFieldType = 7
|
|
)
|
|
|
|
// ProtoField 协议字段
|
|
type ProtoField struct {
|
|
Idx int
|
|
Type ProtoFieldType
|
|
Val interface{}
|
|
}
|
|
|
|
// NewProtoField 创建新的协议字段
|
|
func NewProtoField(idx int, fieldType ProtoFieldType, val interface{}) *ProtoField {
|
|
return &ProtoField{
|
|
Idx: idx,
|
|
Type: fieldType,
|
|
Val: val,
|
|
}
|
|
}
|
|
|
|
// IsAsciiStr 检查是否为ASCII字符串
|
|
func (pf *ProtoField) IsAsciiStr() bool {
|
|
if data, ok := pf.Val.([]byte); ok {
|
|
for _, b := range data {
|
|
if b < 0x20 || b > 0x7E {
|
|
return false
|
|
}
|
|
}
|
|
return true
|
|
}
|
|
return false
|
|
}
|
|
|
|
// ProtoReader 协议读取器
|
|
type ProtoReader struct {
|
|
data []byte
|
|
pos int
|
|
}
|
|
|
|
// NewProtoReader 创建新的协议读取器
|
|
func NewProtoReader(data []byte) *ProtoReader {
|
|
return &ProtoReader{
|
|
data: data,
|
|
pos: 0,
|
|
}
|
|
}
|
|
|
|
// Seek 设置位置
|
|
func (pr *ProtoReader) Seek(pos int) {
|
|
pr.pos = pos
|
|
}
|
|
|
|
// IsRemain 检查是否还有剩余数据
|
|
func (pr *ProtoReader) IsRemain(length int) bool {
|
|
return pr.pos+length <= len(pr.data)
|
|
}
|
|
|
|
// Read0 读取一个字节
|
|
func (pr *ProtoReader) Read0() byte {
|
|
if !pr.IsRemain(1) {
|
|
panic("buffer overflow")
|
|
}
|
|
ret := pr.data[pr.pos]
|
|
pr.pos++
|
|
return ret & 0xFF
|
|
}
|
|
|
|
// Read 读取指定长度的字节
|
|
func (pr *ProtoReader) Read(length int) []byte {
|
|
if !pr.IsRemain(length) {
|
|
panic("buffer overflow")
|
|
}
|
|
ret := pr.data[pr.pos : pr.pos+length]
|
|
pr.pos += length
|
|
return ret
|
|
}
|
|
|
|
// ReadInt32 读取32位整数
|
|
func (pr *ProtoReader) ReadInt32() uint32 {
|
|
bytes := pr.Read(4)
|
|
return uint32(bytes[0]) | uint32(bytes[1])<<8 | uint32(bytes[2])<<16 | uint32(bytes[3])<<24
|
|
}
|
|
|
|
// ReadInt64 读取64位整数
|
|
func (pr *ProtoReader) ReadInt64() uint64 {
|
|
bytes := pr.Read(8)
|
|
return uint64(bytes[0]) | uint64(bytes[1])<<8 | uint64(bytes[2])<<16 | uint64(bytes[3])<<24 |
|
|
uint64(bytes[4])<<32 | uint64(bytes[5])<<40 | uint64(bytes[6])<<48 | uint64(bytes[7])<<56
|
|
}
|
|
|
|
// ReadVarint 读取可变长度整数
|
|
func (pr *ProtoReader) ReadVarint() uint64 {
|
|
vint := uint64(0)
|
|
n := 0
|
|
for {
|
|
b := pr.Read0()
|
|
vint |= uint64(b&0x7F) << (7 * n)
|
|
if b < 0x80 {
|
|
break
|
|
}
|
|
n++
|
|
}
|
|
return vint
|
|
}
|
|
|
|
// ReadString 读取字符串
|
|
func (pr *ProtoReader) ReadString() []byte {
|
|
length := pr.ReadVarint()
|
|
return pr.Read(int(length))
|
|
}
|
|
|
|
// ProtoWriter 协议写入器
|
|
type ProtoWriter struct {
|
|
data []byte
|
|
}
|
|
|
|
// NewProtoWriter 创建新的协议写入器
|
|
func NewProtoWriter() *ProtoWriter {
|
|
return &ProtoWriter{
|
|
data: make([]byte, 0),
|
|
}
|
|
}
|
|
|
|
// Write0 写入一个字节
|
|
func (pw *ProtoWriter) Write0(b byte) {
|
|
pw.data = append(pw.data, b&0xFF)
|
|
}
|
|
|
|
// Write 写入字节数组
|
|
func (pw *ProtoWriter) Write(bytes []byte) {
|
|
pw.data = append(pw.data, bytes...)
|
|
}
|
|
|
|
// WriteInt32 写入32位整数
|
|
func (pw *ProtoWriter) WriteInt32(val uint32) {
|
|
bs := []byte{
|
|
byte(val & 0xFF),
|
|
byte((val >> 8) & 0xFF),
|
|
byte((val >> 16) & 0xFF),
|
|
byte((val >> 24) & 0xFF),
|
|
}
|
|
pw.Write(bs)
|
|
}
|
|
|
|
// WriteInt64 写入64位整数
|
|
func (pw *ProtoWriter) WriteInt64(val uint64) {
|
|
bs := []byte{
|
|
byte(val & 0xFF),
|
|
byte((val >> 8) & 0xFF),
|
|
byte((val >> 16) & 0xFF),
|
|
byte((val >> 24) & 0xFF),
|
|
byte((val >> 32) & 0xFF),
|
|
byte((val >> 40) & 0xFF),
|
|
byte((val >> 48) & 0xFF),
|
|
byte((val >> 56) & 0xFF),
|
|
}
|
|
pw.Write(bs)
|
|
}
|
|
|
|
// WriteVarint 写入可变长度整数
|
|
func (pw *ProtoWriter) WriteVarint(vint uint64) {
|
|
vint = vint & 0xFFFFFFFF
|
|
for vint > 0x80 {
|
|
pw.Write0(byte((vint & 0x7F) | 0x80))
|
|
vint >>= 7
|
|
}
|
|
pw.Write0(byte(vint & 0x7F))
|
|
}
|
|
|
|
// WriteString 写入字符串
|
|
func (pw *ProtoWriter) WriteString(bytes []byte) {
|
|
pw.WriteVarint(uint64(len(bytes)))
|
|
pw.Write(bytes)
|
|
}
|
|
|
|
// ToBytes 转换为字节数组
|
|
func (pw *ProtoWriter) ToBytes() []byte {
|
|
return pw.data
|
|
}
|
|
|
|
// ProtoBuf 协议缓冲区
|
|
type ProtoBuf struct {
|
|
fields []*ProtoField
|
|
}
|
|
|
|
// NewProtoBuf 创建新的协议缓冲区
|
|
func NewProtoBuf(data interface{}) *ProtoBuf {
|
|
pb := &ProtoBuf{
|
|
fields: make([]*ProtoField, 0),
|
|
}
|
|
|
|
if data != nil {
|
|
switch v := data.(type) {
|
|
case []byte:
|
|
if len(v) > 0 {
|
|
pb.parseBuf(v)
|
|
}
|
|
case map[int]interface{}:
|
|
if len(v) > 0 {
|
|
pb.parseDict(v)
|
|
}
|
|
default:
|
|
panic(fmt.Sprintf("unsupported type(%T) to protobuf", data))
|
|
}
|
|
}
|
|
|
|
return pb
|
|
}
|
|
|
|
// parseBuf 解析字节数组
|
|
func (pb *ProtoBuf) parseBuf(bytes []byte) {
|
|
reader := NewProtoReader(bytes)
|
|
for reader.IsRemain(1) {
|
|
key := reader.ReadVarint()
|
|
fieldType := ProtoFieldType(key & 0x7)
|
|
fieldIdx := int(key >> 3)
|
|
|
|
if fieldIdx == 0 {
|
|
break
|
|
}
|
|
|
|
switch fieldType {
|
|
case ProtoFieldTypeINT32:
|
|
pb.Put(NewProtoField(fieldIdx, fieldType, reader.ReadInt32()))
|
|
case ProtoFieldTypeINT64:
|
|
pb.Put(NewProtoField(fieldIdx, fieldType, reader.ReadInt64()))
|
|
case ProtoFieldTypeVARINT:
|
|
pb.Put(NewProtoField(fieldIdx, fieldType, reader.ReadVarint()))
|
|
case ProtoFieldTypeSTRING:
|
|
pb.Put(NewProtoField(fieldIdx, fieldType, reader.ReadString()))
|
|
default:
|
|
panic(fmt.Sprintf("parse protobuf error, unexpected field type: %d", fieldType))
|
|
}
|
|
}
|
|
}
|
|
|
|
// ToBuf 转换为字节数组
|
|
func (pb *ProtoBuf) ToBuf() []byte {
|
|
writer := NewProtoWriter()
|
|
for _, field := range pb.fields {
|
|
key := uint64((field.Idx << 3) | int(field.Type&7))
|
|
writer.WriteVarint(key)
|
|
|
|
switch field.Type {
|
|
case ProtoFieldTypeINT32:
|
|
writer.WriteInt32(field.Val.(uint32))
|
|
case ProtoFieldTypeINT64:
|
|
writer.WriteInt64(field.Val.(uint64))
|
|
case ProtoFieldTypeVARINT:
|
|
writer.WriteVarint(field.Val.(uint64))
|
|
case ProtoFieldTypeSTRING:
|
|
writer.WriteString(field.Val.([]byte))
|
|
default:
|
|
panic(fmt.Sprintf("encode to protobuf error, unexpected field type: %d", field.Type))
|
|
}
|
|
}
|
|
return writer.ToBytes()
|
|
}
|
|
|
|
// Get 获取字段
|
|
func (pb *ProtoBuf) Get(idx int) *ProtoField {
|
|
for _, field := range pb.fields {
|
|
if field.Idx == idx {
|
|
return field
|
|
}
|
|
}
|
|
return nil
|
|
}
|
|
|
|
// GetList 获取字段列表
|
|
func (pb *ProtoBuf) GetList(idx int) []*ProtoField {
|
|
result := make([]*ProtoField, 0)
|
|
for _, field := range pb.fields {
|
|
if field.Idx == idx {
|
|
result = append(result, field)
|
|
}
|
|
}
|
|
return result
|
|
}
|
|
|
|
// GetInt 获取整数
|
|
func (pb *ProtoBuf) GetInt(idx int) uint64 {
|
|
pf := pb.Get(idx)
|
|
if pf == nil {
|
|
return 0
|
|
}
|
|
|
|
switch pf.Type {
|
|
case ProtoFieldTypeINT32:
|
|
return uint64(pf.Val.(uint32))
|
|
case ProtoFieldTypeINT64:
|
|
return pf.Val.(uint64)
|
|
case ProtoFieldTypeVARINT:
|
|
return pf.Val.(uint64)
|
|
default:
|
|
panic(fmt.Sprintf("getInt(%d) -> %d", idx, pf.Type))
|
|
}
|
|
}
|
|
|
|
// GetBytes 获取字节数组
|
|
func (pb *ProtoBuf) GetBytes(idx int) []byte {
|
|
pf := pb.Get(idx)
|
|
if pf == nil {
|
|
return nil
|
|
}
|
|
|
|
if pf.Type == ProtoFieldTypeSTRING {
|
|
return pf.Val.([]byte)
|
|
}
|
|
|
|
panic(fmt.Sprintf("getBytes(%d) -> %d", idx, pf.Type))
|
|
}
|
|
|
|
// GetUTF8 获取UTF8字符串
|
|
func (pb *ProtoBuf) GetUTF8(idx int) string {
|
|
bs := pb.GetBytes(idx)
|
|
if bs == nil {
|
|
return ""
|
|
}
|
|
return string(bs)
|
|
}
|
|
|
|
// GetProtoBuf 获取嵌套的ProtoBuf
|
|
func (pb *ProtoBuf) GetProtoBuf(idx int) *ProtoBuf {
|
|
bs := pb.GetBytes(idx)
|
|
if bs == nil {
|
|
return nil
|
|
}
|
|
return NewProtoBuf(bs)
|
|
}
|
|
|
|
// Put 添加字段
|
|
func (pb *ProtoBuf) Put(field *ProtoField) {
|
|
pb.fields = append(pb.fields, field)
|
|
}
|
|
|
|
// PutInt32 添加32位整数
|
|
func (pb *ProtoBuf) PutInt32(idx int, val uint32) {
|
|
pb.Put(NewProtoField(idx, ProtoFieldTypeINT32, val))
|
|
}
|
|
|
|
// PutInt64 添加64位整数
|
|
func (pb *ProtoBuf) PutInt64(idx int, val uint64) {
|
|
pb.Put(NewProtoField(idx, ProtoFieldTypeINT64, val))
|
|
}
|
|
|
|
// PutVarint 添加可变长度整数
|
|
func (pb *ProtoBuf) PutVarint(idx int, val uint64) {
|
|
pb.Put(NewProtoField(idx, ProtoFieldTypeVARINT, val))
|
|
}
|
|
|
|
// PutBytes 添加字节数组
|
|
func (pb *ProtoBuf) PutBytes(idx int, data []byte) {
|
|
pb.Put(NewProtoField(idx, ProtoFieldTypeSTRING, data))
|
|
}
|
|
|
|
// PutUTF8 添加UTF8字符串
|
|
func (pb *ProtoBuf) PutUTF8(idx int, data string) {
|
|
pb.Put(NewProtoField(idx, ProtoFieldTypeSTRING, []byte(data)))
|
|
}
|
|
|
|
// PutProtoBuf 添加嵌套的ProtoBuf
|
|
func (pb *ProtoBuf) PutProtoBuf(idx int, data *ProtoBuf) {
|
|
pb.Put(NewProtoField(idx, ProtoFieldTypeSTRING, data.ToBuf()))
|
|
}
|
|
|
|
// parseDict 解析字典
|
|
func (pb *ProtoBuf) parseDict(data map[int]interface{}) {
|
|
for k, v := range data {
|
|
switch val := v.(type) {
|
|
case int:
|
|
pb.PutVarint(k, uint64(val))
|
|
case uint64:
|
|
pb.PutVarint(k, val)
|
|
case int64:
|
|
pb.PutVarint(k, uint64(val))
|
|
case string:
|
|
pb.PutUTF8(k, val)
|
|
case []byte:
|
|
pb.PutBytes(k, val)
|
|
case map[int]interface{}:
|
|
pb.PutProtoBuf(k, NewProtoBuf(val))
|
|
default:
|
|
panic(fmt.Sprintf("unsupported type(%T) to protobuf", v))
|
|
}
|
|
}
|
|
}
|
|
|