dsbqj_admin/pkg/util/proto_encode.go

//协议序列化
package util

import (
	"bytes"
	"encoding/binary"
	"errors"
	"fmt"
	"math"
	"reflect"
	"runtime/debug"
	"sync"
)

//协议转为二进制
func Pack(v interface{}) ([]byte, error) {
	e := newEncodeState()
	err := e.marshal(v)
	if err != nil {
		return nil, err
	}
	buf := append([]byte(nil), e.Bytes()...)
	encodeStatePool.Put(e)
	return buf, nil
}

var encodeStatePool sync.Pool

type encodeState struct {
	bytes.Buffer
}

//写入数字类型（v需要明确数字类型，不能为int，uint）
func (e *encodeState) WriteInt(v interface{}) {
	err := binary.Write(e, binary.BigEndian, v)
	if err != nil {
		panic(err)
	}
}

func (e *encodeState) reflectValue(v reflect.Value) {
	typeEncoder(v.Type())(e, v)
}

func (e *encodeState) marshal(v interface{}) (err error) {
	defer func() {
		if r := recover(); r != nil {
			if je, ok := r.(error); ok {
				err = je
			} else {
				err = errors.New("pack error:" + fmt.Sprint(r) + " debug:" + string(debug.Stack()))
			}
		}
	}()
	rv := reflect.ValueOf(v)
	if rv.Kind() == reflect.Ptr {
		rv = rv.Elem()
	}
	e.reflectValue(rv)
	return
}

func newEncodeState() *encodeState {
	if v := encodeStatePool.Get(); v != nil {
		e := v.(*encodeState)
		e.Reset()
		return e
	}
	return &encodeState{}
}

type protoError struct{ error }

type encoderFunc func(e *encodeState, v reflect.Value)

var encoderCache sync.Map // map[reflect.Type]encoderFunc

func typeEncoder(t reflect.Type) encoderFunc {
	if fi, ok := encoderCache.Load(t); ok {
		return fi.(encoderFunc)
	}
	var (
		wg sync.WaitGroup
		f  encoderFunc
	)
	wg.Add(1)
	fi, loaded := encoderCache.LoadOrStore(t, encoderFunc(func(e *encodeState, v reflect.Value) {
		wg.Wait()
		f(e, v)
	}))
	if loaded {
		return fi.(encoderFunc)
	}
	f = newTypeEncoder(t)
	wg.Done()
	encoderCache.Store(t, f)
	return f
}

func newTypeEncoder(t reflect.Type) encoderFunc {
	switch t.Kind() {
	case reflect.Bool:
		return boolEncoder
	case reflect.Int:
		return intEncoder
	case reflect.Uint:
		return uintEncoder
	case reflect.Int32, reflect.Int8, reflect.Int16, reflect.Int64, reflect.Uint8, reflect.Uint16, reflect.Uint32, reflect.Uint64:
		return intComEncoder
	case reflect.Float32:
		return float32Encoder
	case reflect.Float64:
		return float64Encoder
	case reflect.String:
		return stringEncoder
	case reflect.Struct:
		return newStructEncoder(t)
	case reflect.Slice:
		return newSliceEncoder(t)
	case reflect.Array:
		return newArrayEncoder(t)
	default:
		panic(errors.New("proto: unsupported type: " + t.String()))
	}
}

func boolEncoder(e *encodeState, v reflect.Value) {
	if v.Bool() {
		e.WriteByte(1)
	} else {
		e.WriteByte(0)
	}
}

//通用
func intComEncoder(e *encodeState, v reflect.Value) {
	e.WriteInt(v.Interface())
}

//4字节
func intEncoder(e *encodeState, v reflect.Value) {
	e.WriteInt(int32(v.Int()))
}

//4字节
func uintEncoder(e *encodeState, v reflect.Value) {
	e.WriteInt(uint32(v.Uint()))
}

//float32
func float32Encoder(e *encodeState, v reflect.Value) {
	e.WriteInt(math.Float32bits((float32)(v.Float())))
}

//float64
func float64Encoder(e *encodeState, v reflect.Value) {
	e.WriteInt(math.Float64bits((v.Float())))
}

//string
func stringEncoder(e *encodeState, v reflect.Value) {
	buf := []byte(v.String())
	l := len(buf)
	if l > 70000 {
		panic(errors.New(fmt.Sprint("proto: stringLen len > 70000 l:", l)))
	}
	e.WriteInt(uint32(l))
	e.Write(buf)
}

func encodeByteSlice(e *encodeState, v reflect.Value) {
	res := v.Bytes()
	l := len(res)
	if l >= 5000000 {
		panic(errors.New(fmt.Sprint("proto: []uint8 len > 5000000 l:", l)))
	}
	e.WriteInt(uint32(l))
	e.Write(res)
}

type sliceEncoder struct {
	arrayEnc encoderFunc
}

func (se sliceEncoder) encode(e *encodeState, v reflect.Value) {
	if v.IsNil() {
		e.WriteInt(uint32(0))
		return
	}
	se.arrayEnc(e, v)
}

func newSliceEncoder(t reflect.Type) encoderFunc {
	// Byte slices get special treatment; arrays don't.
	if t.Elem().Kind() == reflect.Uint8 {
		return encodeByteSlice
	}
	enc := sliceEncoder{newArrayEncoder(t)}
	return enc.encode
}

type arrayEncoder struct {
	elemEnc encoderFunc
}

func (ae arrayEncoder) encode(e *encodeState, v reflect.Value) {
	l := v.Len()
	if l > 5000000 {
		panic(errors.New("proto: not support arrLen> 5000000 type:" + v.Type().String()))
	}
	e.WriteInt(uint32(l))
	for i := 0; i < l; i++ {
		ae.elemEnc(e, v.Index(i))
	}
}

func newArrayEncoder(t reflect.Type) encoderFunc {
	enc := arrayEncoder{typeEncoder(t.Elem())}
	return enc.encode
}

type structEncoder struct {
	fields []field
}

func (se structEncoder) encode(e *encodeState, v reflect.Value) {
	for i := range se.fields {
		f := &se.fields[i]
		f.encoder(e, v.Field(i))
	}
}

func newStructEncoder(t reflect.Type) encoderFunc {
	return cachedTypeFields(t).encode
}

type field struct {
	encoder encoderFunc
}

func typeFields(t reflect.Type) structEncoder {
	var fields []field
	for i := 0; i < t.NumField(); i++ {
		sf := t.Field(i)
		st := sf.Type
		if st == t {
			panic("proto:not support recursive type:" + st.String())
		}
		if sf.PkgPath != "" {
			panic("proto:not support type:" + t.String() + ",name:" + sf.Name)
		}
		fields = append(fields, field{encoder: typeEncoder(st)})
	}
	return structEncoder{fields}
}

var fieldCache sync.Map // map[reflect.Type]structEncoder

func cachedTypeFields(t reflect.Type) structEncoder {
	if f, ok := fieldCache.Load(t); ok {
		return f.(structEncoder)
	}
	f, _ := fieldCache.LoadOrStore(t, typeFields(t))
	return f.(structEncoder)
}