extensions.go 11 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356
  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. import (
  6. "errors"
  7. "fmt"
  8. "reflect"
  9. "google.golang.org/protobuf/encoding/protowire"
  10. "google.golang.org/protobuf/proto"
  11. "google.golang.org/protobuf/reflect/protoreflect"
  12. "google.golang.org/protobuf/reflect/protoregistry"
  13. "google.golang.org/protobuf/runtime/protoiface"
  14. "google.golang.org/protobuf/runtime/protoimpl"
  15. )
  16. type (
  17. // ExtensionDesc represents an extension descriptor and
  18. // is used to interact with an extension field in a message.
  19. //
  20. // Variables of this type are generated in code by protoc-gen-go.
  21. ExtensionDesc = protoimpl.ExtensionInfo
  22. // ExtensionRange represents a range of message extensions.
  23. // Used in code generated by protoc-gen-go.
  24. ExtensionRange = protoiface.ExtensionRangeV1
  25. // Deprecated: Do not use; this is an internal type.
  26. Extension = protoimpl.ExtensionFieldV1
  27. // Deprecated: Do not use; this is an internal type.
  28. XXX_InternalExtensions = protoimpl.ExtensionFields
  29. )
  30. // ErrMissingExtension reports whether the extension was not present.
  31. var ErrMissingExtension = errors.New("proto: missing extension")
  32. var errNotExtendable = errors.New("proto: not an extendable proto.Message")
  33. // HasExtension reports whether the extension field is present in m
  34. // either as an explicitly populated field or as an unknown field.
  35. func HasExtension(m Message, xt *ExtensionDesc) (has bool) {
  36. mr := MessageReflect(m)
  37. if mr == nil || !mr.IsValid() {
  38. return false
  39. }
  40. // Check whether any populated known field matches the field number.
  41. xtd := xt.TypeDescriptor()
  42. if isValidExtension(mr.Descriptor(), xtd) {
  43. has = mr.Has(xtd)
  44. } else {
  45. mr.Range(func(fd protoreflect.FieldDescriptor, _ protoreflect.Value) bool {
  46. has = int32(fd.Number()) == xt.Field
  47. return !has
  48. })
  49. }
  50. // Check whether any unknown field matches the field number.
  51. for b := mr.GetUnknown(); !has && len(b) > 0; {
  52. num, _, n := protowire.ConsumeField(b)
  53. has = int32(num) == xt.Field
  54. b = b[n:]
  55. }
  56. return has
  57. }
  58. // ClearExtension removes the extension field from m
  59. // either as an explicitly populated field or as an unknown field.
  60. func ClearExtension(m Message, xt *ExtensionDesc) {
  61. mr := MessageReflect(m)
  62. if mr == nil || !mr.IsValid() {
  63. return
  64. }
  65. xtd := xt.TypeDescriptor()
  66. if isValidExtension(mr.Descriptor(), xtd) {
  67. mr.Clear(xtd)
  68. } else {
  69. mr.Range(func(fd protoreflect.FieldDescriptor, _ protoreflect.Value) bool {
  70. if int32(fd.Number()) == xt.Field {
  71. mr.Clear(fd)
  72. return false
  73. }
  74. return true
  75. })
  76. }
  77. clearUnknown(mr, fieldNum(xt.Field))
  78. }
  79. // ClearAllExtensions clears all extensions from m.
  80. // This includes populated fields and unknown fields in the extension range.
  81. func ClearAllExtensions(m Message) {
  82. mr := MessageReflect(m)
  83. if mr == nil || !mr.IsValid() {
  84. return
  85. }
  86. mr.Range(func(fd protoreflect.FieldDescriptor, _ protoreflect.Value) bool {
  87. if fd.IsExtension() {
  88. mr.Clear(fd)
  89. }
  90. return true
  91. })
  92. clearUnknown(mr, mr.Descriptor().ExtensionRanges())
  93. }
  94. // GetExtension retrieves a proto2 extended field from m.
  95. //
  96. // If the descriptor is type complete (i.e., ExtensionDesc.ExtensionType is non-nil),
  97. // then GetExtension parses the encoded field and returns a Go value of the specified type.
  98. // If the field is not present, then the default value is returned (if one is specified),
  99. // otherwise ErrMissingExtension is reported.
  100. //
  101. // If the descriptor is type incomplete (i.e., ExtensionDesc.ExtensionType is nil),
  102. // then GetExtension returns the raw encoded bytes for the extension field.
  103. func GetExtension(m Message, xt *ExtensionDesc) (interface{}, error) {
  104. mr := MessageReflect(m)
  105. if mr == nil || !mr.IsValid() || mr.Descriptor().ExtensionRanges().Len() == 0 {
  106. return nil, errNotExtendable
  107. }
  108. // Retrieve the unknown fields for this extension field.
  109. var bo protoreflect.RawFields
  110. for bi := mr.GetUnknown(); len(bi) > 0; {
  111. num, _, n := protowire.ConsumeField(bi)
  112. if int32(num) == xt.Field {
  113. bo = append(bo, bi[:n]...)
  114. }
  115. bi = bi[n:]
  116. }
  117. // For type incomplete descriptors, only retrieve the unknown fields.
  118. if xt.ExtensionType == nil {
  119. return []byte(bo), nil
  120. }
  121. // If the extension field only exists as unknown fields, unmarshal it.
  122. // This is rarely done since proto.Unmarshal eagerly unmarshals extensions.
  123. xtd := xt.TypeDescriptor()
  124. if !isValidExtension(mr.Descriptor(), xtd) {
  125. return nil, fmt.Errorf("proto: bad extended type; %T does not extend %T", xt.ExtendedType, m)
  126. }
  127. if !mr.Has(xtd) && len(bo) > 0 {
  128. m2 := mr.New()
  129. if err := (proto.UnmarshalOptions{
  130. Resolver: extensionResolver{xt},
  131. }.Unmarshal(bo, m2.Interface())); err != nil {
  132. return nil, err
  133. }
  134. if m2.Has(xtd) {
  135. mr.Set(xtd, m2.Get(xtd))
  136. clearUnknown(mr, fieldNum(xt.Field))
  137. }
  138. }
  139. // Check whether the message has the extension field set or a default.
  140. var pv protoreflect.Value
  141. switch {
  142. case mr.Has(xtd):
  143. pv = mr.Get(xtd)
  144. case xtd.HasDefault():
  145. pv = xtd.Default()
  146. default:
  147. return nil, ErrMissingExtension
  148. }
  149. v := xt.InterfaceOf(pv)
  150. rv := reflect.ValueOf(v)
  151. if isScalarKind(rv.Kind()) {
  152. rv2 := reflect.New(rv.Type())
  153. rv2.Elem().Set(rv)
  154. v = rv2.Interface()
  155. }
  156. return v, nil
  157. }
  158. // extensionResolver is a custom extension resolver that stores a single
  159. // extension type that takes precedence over the global registry.
  160. type extensionResolver struct{ xt protoreflect.ExtensionType }
  161. func (r extensionResolver) FindExtensionByName(field protoreflect.FullName) (protoreflect.ExtensionType, error) {
  162. if xtd := r.xt.TypeDescriptor(); xtd.FullName() == field {
  163. return r.xt, nil
  164. }
  165. return protoregistry.GlobalTypes.FindExtensionByName(field)
  166. }
  167. func (r extensionResolver) FindExtensionByNumber(message protoreflect.FullName, field protoreflect.FieldNumber) (protoreflect.ExtensionType, error) {
  168. if xtd := r.xt.TypeDescriptor(); xtd.ContainingMessage().FullName() == message && xtd.Number() == field {
  169. return r.xt, nil
  170. }
  171. return protoregistry.GlobalTypes.FindExtensionByNumber(message, field)
  172. }
  173. // GetExtensions returns a list of the extensions values present in m,
  174. // corresponding with the provided list of extension descriptors, xts.
  175. // If an extension is missing in m, the corresponding value is nil.
  176. func GetExtensions(m Message, xts []*ExtensionDesc) ([]interface{}, error) {
  177. mr := MessageReflect(m)
  178. if mr == nil || !mr.IsValid() {
  179. return nil, errNotExtendable
  180. }
  181. vs := make([]interface{}, len(xts))
  182. for i, xt := range xts {
  183. v, err := GetExtension(m, xt)
  184. if err != nil {
  185. if err == ErrMissingExtension {
  186. continue
  187. }
  188. return vs, err
  189. }
  190. vs[i] = v
  191. }
  192. return vs, nil
  193. }
  194. // SetExtension sets an extension field in m to the provided value.
  195. func SetExtension(m Message, xt *ExtensionDesc, v interface{}) error {
  196. mr := MessageReflect(m)
  197. if mr == nil || !mr.IsValid() || mr.Descriptor().ExtensionRanges().Len() == 0 {
  198. return errNotExtendable
  199. }
  200. rv := reflect.ValueOf(v)
  201. if reflect.TypeOf(v) != reflect.TypeOf(xt.ExtensionType) {
  202. return fmt.Errorf("proto: bad extension value type. got: %T, want: %T", v, xt.ExtensionType)
  203. }
  204. if rv.Kind() == reflect.Ptr {
  205. if rv.IsNil() {
  206. return fmt.Errorf("proto: SetExtension called with nil value of type %T", v)
  207. }
  208. if isScalarKind(rv.Elem().Kind()) {
  209. v = rv.Elem().Interface()
  210. }
  211. }
  212. xtd := xt.TypeDescriptor()
  213. if !isValidExtension(mr.Descriptor(), xtd) {
  214. return fmt.Errorf("proto: bad extended type; %T does not extend %T", xt.ExtendedType, m)
  215. }
  216. mr.Set(xtd, xt.ValueOf(v))
  217. clearUnknown(mr, fieldNum(xt.Field))
  218. return nil
  219. }
  220. // SetRawExtension inserts b into the unknown fields of m.
  221. //
  222. // Deprecated: Use Message.ProtoReflect.SetUnknown instead.
  223. func SetRawExtension(m Message, fnum int32, b []byte) {
  224. mr := MessageReflect(m)
  225. if mr == nil || !mr.IsValid() {
  226. return
  227. }
  228. // Verify that the raw field is valid.
  229. for b0 := b; len(b0) > 0; {
  230. num, _, n := protowire.ConsumeField(b0)
  231. if int32(num) != fnum {
  232. panic(fmt.Sprintf("mismatching field number: got %d, want %d", num, fnum))
  233. }
  234. b0 = b0[n:]
  235. }
  236. ClearExtension(m, &ExtensionDesc{Field: fnum})
  237. mr.SetUnknown(append(mr.GetUnknown(), b...))
  238. }
  239. // ExtensionDescs returns a list of extension descriptors found in m,
  240. // containing descriptors for both populated extension fields in m and
  241. // also unknown fields of m that are in the extension range.
  242. // For the later case, an type incomplete descriptor is provided where only
  243. // the ExtensionDesc.Field field is populated.
  244. // The order of the extension descriptors is undefined.
  245. func ExtensionDescs(m Message) ([]*ExtensionDesc, error) {
  246. mr := MessageReflect(m)
  247. if mr == nil || !mr.IsValid() || mr.Descriptor().ExtensionRanges().Len() == 0 {
  248. return nil, errNotExtendable
  249. }
  250. // Collect a set of known extension descriptors.
  251. extDescs := make(map[protoreflect.FieldNumber]*ExtensionDesc)
  252. mr.Range(func(fd protoreflect.FieldDescriptor, v protoreflect.Value) bool {
  253. if fd.IsExtension() {
  254. xt := fd.(protoreflect.ExtensionTypeDescriptor)
  255. if xd, ok := xt.Type().(*ExtensionDesc); ok {
  256. extDescs[fd.Number()] = xd
  257. }
  258. }
  259. return true
  260. })
  261. // Collect a set of unknown extension descriptors.
  262. extRanges := mr.Descriptor().ExtensionRanges()
  263. for b := mr.GetUnknown(); len(b) > 0; {
  264. num, _, n := protowire.ConsumeField(b)
  265. if extRanges.Has(num) && extDescs[num] == nil {
  266. extDescs[num] = nil
  267. }
  268. b = b[n:]
  269. }
  270. // Transpose the set of descriptors into a list.
  271. var xts []*ExtensionDesc
  272. for num, xt := range extDescs {
  273. if xt == nil {
  274. xt = &ExtensionDesc{Field: int32(num)}
  275. }
  276. xts = append(xts, xt)
  277. }
  278. return xts, nil
  279. }
  280. // isValidExtension reports whether xtd is a valid extension descriptor for md.
  281. func isValidExtension(md protoreflect.MessageDescriptor, xtd protoreflect.ExtensionTypeDescriptor) bool {
  282. return xtd.ContainingMessage() == md && md.ExtensionRanges().Has(xtd.Number())
  283. }
  284. // isScalarKind reports whether k is a protobuf scalar kind (except bytes).
  285. // This function exists for historical reasons since the representation of
  286. // scalars differs between v1 and v2, where v1 uses *T and v2 uses T.
  287. func isScalarKind(k reflect.Kind) bool {
  288. switch k {
  289. case reflect.Bool, reflect.Int32, reflect.Int64, reflect.Uint32, reflect.Uint64, reflect.Float32, reflect.Float64, reflect.String:
  290. return true
  291. default:
  292. return false
  293. }
  294. }
  295. // clearUnknown removes unknown fields from m where remover.Has reports true.
  296. func clearUnknown(m protoreflect.Message, remover interface {
  297. Has(protoreflect.FieldNumber) bool
  298. }) {
  299. var bo protoreflect.RawFields
  300. for bi := m.GetUnknown(); len(bi) > 0; {
  301. num, _, n := protowire.ConsumeField(bi)
  302. if !remover.Has(num) {
  303. bo = append(bo, bi[:n]...)
  304. }
  305. bi = bi[n:]
  306. }
  307. if bi := m.GetUnknown(); len(bi) != len(bo) {
  308. m.SetUnknown(bo)
  309. }
  310. }
  311. type fieldNum protoreflect.FieldNumber
  312. func (n1 fieldNum) Has(n2 protoreflect.FieldNumber) bool {
  313. return protoreflect.FieldNumber(n1) == n2
  314. }