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