base-golang/pkg/websocket/codec/protobuf/protobuf.go

package protobuf

import (
	"errors"
	"fmt"

	"git.bvbej.com/bvbej/base-golang/pkg/websocket/codec"
	"git.bvbej.com/bvbej/base-golang/pkg/websocket/codec/protobuf/protocol"
	"google.golang.org/protobuf/proto"
)

type protobufCodec struct{}

func init() {
	codec.RegisterCodec("protobuf_codec", new(protobufCodec))
}

func (*protobufCodec) Marshal(router string, dataPtr any, retErr error) ([]byte, error) {
	if router == "" {
		return nil, fmt.Errorf("marshal: empty router")
	}
	if dataPtr == nil && retErr == nil {
		return nil, fmt.Errorf("marshal: empty data")
	}
	ack := &protocol.TransPack{
		Router: router,
	}
	if dataPtr != nil {
		pbMsg, ok := dataPtr.(proto.Message)
		if !ok {
			return nil, fmt.Errorf("marshal: dataptr only support proto.Message type. router:%s dt:%T ",
				router, dataPtr)
		}
		data, err := proto.Marshal(pbMsg)
		if err != nil {
			return nil, fmt.Errorf("marshal:protocol buffer marshal failed. router:%s dt:%T err:%v",
				router, dataPtr, err)
		}
		ack.Data = data
	} else {
		ack.Error = retErr.Error()
	}
	ackByte, err := proto.Marshal(ack)
	if err != nil {
		return nil, fmt.Errorf("marshal:protocol buffer marshal failed. router:%s dt:%T err:%v",
			router, ack, err)
	}
	return ackByte, nil
}

func (*protobufCodec) Unmarshal(msg []byte) (int, *codec.MsgPack, error) {
	var l = len(msg)
	req := &protocol.TransPack{}
	err := proto.Unmarshal(msg, req)
	if err != nil {
		return l, nil, errors.New("unmarshal split message id failed.")
	}
	var router = req.Router
	msgPack := &codec.MsgPack{Router: router}
	dt := codec.GetMessage(router)
	if dt == nil {
		return l, nil, fmt.Errorf("unmarshal message not registed. router:%s",
			router)
	}
	if req.Data != nil {
		err = proto.Unmarshal(req.Data, dt.(proto.Message))
		if err != nil {
			return l, nil, fmt.Errorf("unmarshal failed. router:%s", router)
		}
	}
	msgPack.DataPtr = dt
	if req.Error != "" {
		msgPack.Err = errors.New(req.Error)
	}
	return l, msgPack, nil
}

func (*protobufCodec) ToString(data any) string {
	pbMsg, ok := data.(proto.Message)
	if !ok {
		return fmt.Sprintf("invalid type %T", data)
	}
	marshal, _ := proto.Marshal(pbMsg)
	return string(marshal)
}