registry.go 11 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323
  1. // Copyright 2019 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. "bytes"
  7. "compress/gzip"
  8. "fmt"
  9. "io/ioutil"
  10. "reflect"
  11. "strings"
  12. "sync"
  13. "google.golang.org/protobuf/reflect/protoreflect"
  14. "google.golang.org/protobuf/reflect/protoregistry"
  15. "google.golang.org/protobuf/runtime/protoimpl"
  16. )
  17. // filePath is the path to the proto source file.
  18. type filePath = string // e.g., "google/protobuf/descriptor.proto"
  19. // fileDescGZIP is the compressed contents of the encoded FileDescriptorProto.
  20. type fileDescGZIP = []byte
  21. var fileCache sync.Map // map[filePath]fileDescGZIP
  22. // RegisterFile is called from generated code to register the compressed
  23. // FileDescriptorProto with the file path for a proto source file.
  24. //
  25. // Deprecated: Use protoregistry.GlobalFiles.RegisterFile instead.
  26. func RegisterFile(s filePath, d fileDescGZIP) {
  27. // Decompress the descriptor.
  28. zr, err := gzip.NewReader(bytes.NewReader(d))
  29. if err != nil {
  30. panic(fmt.Sprintf("proto: invalid compressed file descriptor: %v", err))
  31. }
  32. b, err := ioutil.ReadAll(zr)
  33. if err != nil {
  34. panic(fmt.Sprintf("proto: invalid compressed file descriptor: %v", err))
  35. }
  36. // Construct a protoreflect.FileDescriptor from the raw descriptor.
  37. // Note that DescBuilder.Build automatically registers the constructed
  38. // file descriptor with the v2 registry.
  39. protoimpl.DescBuilder{RawDescriptor: b}.Build()
  40. // Locally cache the raw descriptor form for the file.
  41. fileCache.Store(s, d)
  42. }
  43. // FileDescriptor returns the compressed FileDescriptorProto given the file path
  44. // for a proto source file. It returns nil if not found.
  45. //
  46. // Deprecated: Use protoregistry.GlobalFiles.FindFileByPath instead.
  47. func FileDescriptor(s filePath) fileDescGZIP {
  48. if v, ok := fileCache.Load(s); ok {
  49. return v.(fileDescGZIP)
  50. }
  51. // Find the descriptor in the v2 registry.
  52. var b []byte
  53. if fd, _ := protoregistry.GlobalFiles.FindFileByPath(s); fd != nil {
  54. if fd, ok := fd.(interface{ ProtoLegacyRawDesc() []byte }); ok {
  55. b = fd.ProtoLegacyRawDesc()
  56. } else {
  57. // TODO: Use protodesc.ToFileDescriptorProto to construct
  58. // a descriptorpb.FileDescriptorProto and marshal it.
  59. // However, doing so causes the proto package to have a dependency
  60. // on descriptorpb, leading to cyclic dependency issues.
  61. }
  62. }
  63. // Locally cache the raw descriptor form for the file.
  64. if len(b) > 0 {
  65. v, _ := fileCache.LoadOrStore(s, protoimpl.X.CompressGZIP(b))
  66. return v.(fileDescGZIP)
  67. }
  68. return nil
  69. }
  70. // enumName is the name of an enum. For historical reasons, the enum name is
  71. // neither the full Go name nor the full protobuf name of the enum.
  72. // The name is the dot-separated combination of just the proto package that the
  73. // enum is declared within followed by the Go type name of the generated enum.
  74. type enumName = string // e.g., "my.proto.package.GoMessage_GoEnum"
  75. // enumsByName maps enum values by name to their numeric counterpart.
  76. type enumsByName = map[string]int32
  77. // enumsByNumber maps enum values by number to their name counterpart.
  78. type enumsByNumber = map[int32]string
  79. var enumCache sync.Map // map[enumName]enumsByName
  80. var numFilesCache sync.Map // map[protoreflect.FullName]int
  81. // RegisterEnum is called from the generated code to register the mapping of
  82. // enum value names to enum numbers for the enum identified by s.
  83. //
  84. // Deprecated: Use protoregistry.GlobalTypes.RegisterEnum instead.
  85. func RegisterEnum(s enumName, _ enumsByNumber, m enumsByName) {
  86. if _, ok := enumCache.Load(s); ok {
  87. panic("proto: duplicate enum registered: " + s)
  88. }
  89. enumCache.Store(s, m)
  90. // This does not forward registration to the v2 registry since this API
  91. // lacks sufficient information to construct a complete v2 enum descriptor.
  92. }
  93. // EnumValueMap returns the mapping from enum value names to enum numbers for
  94. // the enum of the given name. It returns nil if not found.
  95. //
  96. // Deprecated: Use protoregistry.GlobalTypes.FindEnumByName instead.
  97. func EnumValueMap(s enumName) enumsByName {
  98. if v, ok := enumCache.Load(s); ok {
  99. return v.(enumsByName)
  100. }
  101. // Check whether the cache is stale. If the number of files in the current
  102. // package differs, then it means that some enums may have been recently
  103. // registered upstream that we do not know about.
  104. var protoPkg protoreflect.FullName
  105. if i := strings.LastIndexByte(s, '.'); i >= 0 {
  106. protoPkg = protoreflect.FullName(s[:i])
  107. }
  108. v, _ := numFilesCache.Load(protoPkg)
  109. numFiles, _ := v.(int)
  110. if protoregistry.GlobalFiles.NumFilesByPackage(protoPkg) == numFiles {
  111. return nil // cache is up-to-date; was not found earlier
  112. }
  113. // Update the enum cache for all enums declared in the given proto package.
  114. numFiles = 0
  115. protoregistry.GlobalFiles.RangeFilesByPackage(protoPkg, func(fd protoreflect.FileDescriptor) bool {
  116. walkEnums(fd, func(ed protoreflect.EnumDescriptor) {
  117. name := protoimpl.X.LegacyEnumName(ed)
  118. if _, ok := enumCache.Load(name); !ok {
  119. m := make(enumsByName)
  120. evs := ed.Values()
  121. for i := evs.Len() - 1; i >= 0; i-- {
  122. ev := evs.Get(i)
  123. m[string(ev.Name())] = int32(ev.Number())
  124. }
  125. enumCache.LoadOrStore(name, m)
  126. }
  127. })
  128. numFiles++
  129. return true
  130. })
  131. numFilesCache.Store(protoPkg, numFiles)
  132. // Check cache again for enum map.
  133. if v, ok := enumCache.Load(s); ok {
  134. return v.(enumsByName)
  135. }
  136. return nil
  137. }
  138. // walkEnums recursively walks all enums declared in d.
  139. func walkEnums(d interface {
  140. Enums() protoreflect.EnumDescriptors
  141. Messages() protoreflect.MessageDescriptors
  142. }, f func(protoreflect.EnumDescriptor)) {
  143. eds := d.Enums()
  144. for i := eds.Len() - 1; i >= 0; i-- {
  145. f(eds.Get(i))
  146. }
  147. mds := d.Messages()
  148. for i := mds.Len() - 1; i >= 0; i-- {
  149. walkEnums(mds.Get(i), f)
  150. }
  151. }
  152. // messageName is the full name of protobuf message.
  153. type messageName = string
  154. var messageTypeCache sync.Map // map[messageName]reflect.Type
  155. // RegisterType is called from generated code to register the message Go type
  156. // for a message of the given name.
  157. //
  158. // Deprecated: Use protoregistry.GlobalTypes.RegisterMessage instead.
  159. func RegisterType(m Message, s messageName) {
  160. mt := protoimpl.X.LegacyMessageTypeOf(m, protoreflect.FullName(s))
  161. if err := protoregistry.GlobalTypes.RegisterMessage(mt); err != nil {
  162. panic(err)
  163. }
  164. messageTypeCache.Store(s, reflect.TypeOf(m))
  165. }
  166. // RegisterMapType is called from generated code to register the Go map type
  167. // for a protobuf message representing a map entry.
  168. //
  169. // Deprecated: Do not use.
  170. func RegisterMapType(m interface{}, s messageName) {
  171. t := reflect.TypeOf(m)
  172. if t.Kind() != reflect.Map {
  173. panic(fmt.Sprintf("invalid map kind: %v", t))
  174. }
  175. if _, ok := messageTypeCache.Load(s); ok {
  176. panic(fmt.Errorf("proto: duplicate proto message registered: %s", s))
  177. }
  178. messageTypeCache.Store(s, t)
  179. }
  180. // MessageType returns the message type for a named message.
  181. // It returns nil if not found.
  182. //
  183. // Deprecated: Use protoregistry.GlobalTypes.FindMessageByName instead.
  184. func MessageType(s messageName) reflect.Type {
  185. if v, ok := messageTypeCache.Load(s); ok {
  186. return v.(reflect.Type)
  187. }
  188. // Derive the message type from the v2 registry.
  189. var t reflect.Type
  190. if mt, _ := protoregistry.GlobalTypes.FindMessageByName(protoreflect.FullName(s)); mt != nil {
  191. t = messageGoType(mt)
  192. }
  193. // If we could not get a concrete type, it is possible that it is a
  194. // pseudo-message for a map entry.
  195. if t == nil {
  196. d, _ := protoregistry.GlobalFiles.FindDescriptorByName(protoreflect.FullName(s))
  197. if md, _ := d.(protoreflect.MessageDescriptor); md != nil && md.IsMapEntry() {
  198. kt := goTypeForField(md.Fields().ByNumber(1))
  199. vt := goTypeForField(md.Fields().ByNumber(2))
  200. t = reflect.MapOf(kt, vt)
  201. }
  202. }
  203. // Locally cache the message type for the given name.
  204. if t != nil {
  205. v, _ := messageTypeCache.LoadOrStore(s, t)
  206. return v.(reflect.Type)
  207. }
  208. return nil
  209. }
  210. func goTypeForField(fd protoreflect.FieldDescriptor) reflect.Type {
  211. switch k := fd.Kind(); k {
  212. case protoreflect.EnumKind:
  213. if et, _ := protoregistry.GlobalTypes.FindEnumByName(fd.Enum().FullName()); et != nil {
  214. return enumGoType(et)
  215. }
  216. return reflect.TypeOf(protoreflect.EnumNumber(0))
  217. case protoreflect.MessageKind, protoreflect.GroupKind:
  218. if mt, _ := protoregistry.GlobalTypes.FindMessageByName(fd.Message().FullName()); mt != nil {
  219. return messageGoType(mt)
  220. }
  221. return reflect.TypeOf((*protoreflect.Message)(nil)).Elem()
  222. default:
  223. return reflect.TypeOf(fd.Default().Interface())
  224. }
  225. }
  226. func enumGoType(et protoreflect.EnumType) reflect.Type {
  227. return reflect.TypeOf(et.New(0))
  228. }
  229. func messageGoType(mt protoreflect.MessageType) reflect.Type {
  230. return reflect.TypeOf(MessageV1(mt.Zero().Interface()))
  231. }
  232. // MessageName returns the full protobuf name for the given message type.
  233. //
  234. // Deprecated: Use protoreflect.MessageDescriptor.FullName instead.
  235. func MessageName(m Message) messageName {
  236. if m == nil {
  237. return ""
  238. }
  239. if m, ok := m.(interface{ XXX_MessageName() messageName }); ok {
  240. return m.XXX_MessageName()
  241. }
  242. return messageName(protoimpl.X.MessageDescriptorOf(m).FullName())
  243. }
  244. // RegisterExtension is called from the generated code to register
  245. // the extension descriptor.
  246. //
  247. // Deprecated: Use protoregistry.GlobalTypes.RegisterExtension instead.
  248. func RegisterExtension(d *ExtensionDesc) {
  249. if err := protoregistry.GlobalTypes.RegisterExtension(d); err != nil {
  250. panic(err)
  251. }
  252. }
  253. type extensionsByNumber = map[int32]*ExtensionDesc
  254. var extensionCache sync.Map // map[messageName]extensionsByNumber
  255. // RegisteredExtensions returns a map of the registered extensions for the
  256. // provided protobuf message, indexed by the extension field number.
  257. //
  258. // Deprecated: Use protoregistry.GlobalTypes.RangeExtensionsByMessage instead.
  259. func RegisteredExtensions(m Message) extensionsByNumber {
  260. // Check whether the cache is stale. If the number of extensions for
  261. // the given message differs, then it means that some extensions were
  262. // recently registered upstream that we do not know about.
  263. s := MessageName(m)
  264. v, _ := extensionCache.Load(s)
  265. xs, _ := v.(extensionsByNumber)
  266. if protoregistry.GlobalTypes.NumExtensionsByMessage(protoreflect.FullName(s)) == len(xs) {
  267. return xs // cache is up-to-date
  268. }
  269. // Cache is stale, re-compute the extensions map.
  270. xs = make(extensionsByNumber)
  271. protoregistry.GlobalTypes.RangeExtensionsByMessage(protoreflect.FullName(s), func(xt protoreflect.ExtensionType) bool {
  272. if xd, ok := xt.(*ExtensionDesc); ok {
  273. xs[int32(xt.TypeDescriptor().Number())] = xd
  274. } else {
  275. // TODO: This implies that the protoreflect.ExtensionType is a
  276. // custom type not generated by protoc-gen-go. We could try and
  277. // convert the type to an ExtensionDesc.
  278. }
  279. return true
  280. })
  281. extensionCache.Store(s, xs)
  282. return xs
  283. }