protobuf.go 2.1 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384
  1. package protobuf
  2. import (
  3. "errors"
  4. "fmt"
  5. "git.bvbej.com/bvbej/base-golang/pkg/websocket/codec"
  6. "git.bvbej.com/bvbej/base-golang/pkg/websocket/codec/protobuf/protocol"
  7. "github.com/golang/protobuf/proto"
  8. )
  9. type protobufCodec struct{}
  10. func init() {
  11. codec.RegisterCodec("protobuf_codec", new(protobufCodec))
  12. }
  13. func (*protobufCodec) Marshal(router string, dataPtr any, retErr error) ([]byte, error) {
  14. if router == "" {
  15. return nil, fmt.Errorf("marshal: empty router")
  16. }
  17. if dataPtr == nil && retErr == nil {
  18. return nil, fmt.Errorf("marshal: empty data")
  19. }
  20. ack := &protocol.TransPack{
  21. Router: router,
  22. }
  23. if dataPtr != nil {
  24. pbMsg, ok := dataPtr.(proto.Message)
  25. if !ok {
  26. return nil, fmt.Errorf("marshal: dataptr only support proto.Message type. router:%s dt:%T ",
  27. router, dataPtr)
  28. }
  29. data, err := proto.Marshal(pbMsg)
  30. if err != nil {
  31. return nil, fmt.Errorf("marshal:protocol buffer marshal failed. router:%s dt:%T err:%v",
  32. router, dataPtr, err)
  33. }
  34. ack.Data = data
  35. } else {
  36. ack.Error = retErr.Error()
  37. }
  38. ackByte, err := proto.Marshal(ack)
  39. if err != nil {
  40. return nil, fmt.Errorf("marshal:protocol buffer marshal failed. router:%s dt:%T err:%v",
  41. router, ack, err)
  42. }
  43. return ackByte, nil
  44. }
  45. func (*protobufCodec) Unmarshal(msg []byte) (int, *codec.MsgPack, error) {
  46. var l = len(msg)
  47. req := &protocol.TransPack{}
  48. err := proto.Unmarshal(msg, req)
  49. if err != nil {
  50. return l, nil, errors.New("unmarshal split message id failed.")
  51. }
  52. var router = req.Router
  53. msgPack := &codec.MsgPack{Router: router}
  54. dt := codec.GetMessage(router)
  55. if dt == nil {
  56. return l, nil, fmt.Errorf("unmarshal message not registed. router:%s",
  57. router)
  58. }
  59. if req.Data != nil {
  60. err = proto.Unmarshal(req.Data, dt.(proto.Message))
  61. if err != nil {
  62. return l, nil, fmt.Errorf("unmarshal failed. router:%s", router)
  63. }
  64. }
  65. msgPack.DataPtr = dt
  66. if req.Error != "" {
  67. msgPack.Err = errors.New(req.Error)
  68. }
  69. return l, msgPack, nil
  70. }
  71. func (*protobufCodec) ToString(data any) string {
  72. pbMsg, ok := data.(proto.Message)
  73. if !ok {
  74. return fmt.Sprintf("invalid type %T", data)
  75. }
  76. return pbMsg.String()
  77. }