message_set.go 3.9 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150
  1. // Copyright 2010 The Go Authors. All rights reserved.
  2. // Use of this source code is governed by a BSD-style
  3. // license that can be found in the LICENSE file.
  4. package proto
  5. /*
  6. * Support for message sets.
  7. */
  8. import (
  9. "errors"
  10. "github.com/golang/protobuf/v2/reflect/protoreflect"
  11. )
  12. // errNoMessageTypeID occurs when a protocol buffer does not have a message type ID.
  13. // A message type ID is required for storing a protocol buffer in a message set.
  14. var errNoMessageTypeID = errors.New("proto does not have a message type ID")
  15. // The first two types (_MessageSet_Item and messageSet)
  16. // model what the protocol compiler produces for the following protocol message:
  17. // message MessageSet {
  18. // repeated group Item = 1 {
  19. // required int32 type_id = 2;
  20. // required string message = 3;
  21. // };
  22. // }
  23. // That is the MessageSet wire format. We can't use a proto to generate these
  24. // because that would introduce a circular dependency between it and this package.
  25. type _MessageSet_Item struct {
  26. TypeId *int32 `protobuf:"varint,2,req,name=type_id"`
  27. Message []byte `protobuf:"bytes,3,req,name=message"`
  28. }
  29. type messageSet struct {
  30. Item []*_MessageSet_Item `protobuf:"group,1,rep"`
  31. XXX_unrecognized []byte
  32. // TODO: caching?
  33. }
  34. // Make sure messageSet is a Message.
  35. var _ Message = (*messageSet)(nil)
  36. // messageTypeIder is an interface satisfied by a protocol buffer type
  37. // that may be stored in a MessageSet.
  38. type messageTypeIder interface {
  39. MessageTypeId() int32
  40. }
  41. func (ms *messageSet) find(pb Message) *_MessageSet_Item {
  42. mti, ok := pb.(messageTypeIder)
  43. if !ok {
  44. return nil
  45. }
  46. id := mti.MessageTypeId()
  47. for _, item := range ms.Item {
  48. if *item.TypeId == id {
  49. return item
  50. }
  51. }
  52. return nil
  53. }
  54. func (ms *messageSet) Has(pb Message) bool {
  55. return ms.find(pb) != nil
  56. }
  57. func (ms *messageSet) Unmarshal(pb Message) error {
  58. if item := ms.find(pb); item != nil {
  59. return Unmarshal(item.Message, pb)
  60. }
  61. if _, ok := pb.(messageTypeIder); !ok {
  62. return errNoMessageTypeID
  63. }
  64. return nil // TODO: return error instead?
  65. }
  66. func (ms *messageSet) Marshal(pb Message) error {
  67. msg, err := Marshal(pb)
  68. if err != nil {
  69. return err
  70. }
  71. if item := ms.find(pb); item != nil {
  72. // reuse existing item
  73. item.Message = msg
  74. return nil
  75. }
  76. mti, ok := pb.(messageTypeIder)
  77. if !ok {
  78. return errNoMessageTypeID
  79. }
  80. mtid := mti.MessageTypeId()
  81. ms.Item = append(ms.Item, &_MessageSet_Item{
  82. TypeId: &mtid,
  83. Message: msg,
  84. })
  85. return nil
  86. }
  87. func (ms *messageSet) Reset() { *ms = messageSet{} }
  88. func (ms *messageSet) String() string { return CompactTextString(ms) }
  89. func (*messageSet) ProtoMessage() {}
  90. // Support for the message_set_wire_format message option.
  91. func skipVarint(buf []byte) []byte {
  92. i := 0
  93. for ; buf[i]&0x80 != 0; i++ {
  94. }
  95. return buf[i+1:]
  96. }
  97. // unmarshalMessageSet decodes the extension map encoded in buf in the message set wire format.
  98. // It is called by Unmarshal methods on protocol buffer messages with the message_set_wire_format option.
  99. func unmarshalMessageSet(buf []byte, exts interface{}) error {
  100. m := extensionFieldsOf(exts)
  101. ms := new(messageSet)
  102. if err := Unmarshal(buf, ms); err != nil {
  103. return err
  104. }
  105. for _, item := range ms.Item {
  106. id := protoreflect.FieldNumber(*item.TypeId)
  107. msg := item.Message
  108. // Restore wire type and field number varint, plus length varint.
  109. // Be careful to preserve duplicate items.
  110. b := EncodeVarint(uint64(id)<<3 | WireBytes)
  111. if m.Has(id) {
  112. ext := m.Get(id)
  113. // Existing data; rip off the tag and length varint
  114. // so we join the new data correctly.
  115. // We can assume that ext.Raw is set because we are unmarshaling.
  116. o := ext.Raw[len(b):] // skip wire type and field number
  117. _, n := DecodeVarint(o) // calculate length of length varint
  118. o = o[n:] // skip length varint
  119. msg = append(o, msg...) // join old data and new data
  120. }
  121. b = append(b, EncodeVarint(uint64(len(msg)))...)
  122. b = append(b, msg...)
  123. m.Set(id, Extension{Raw: b})
  124. }
  125. return nil
  126. }