protogen.go 36 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182838485868788899091929394959697989910010110210310410510610710810911011111211311411511611711811912012112212312412512612712812913013113213313413513613713813914014114214314414514614714814915015115215315415515615715815916016116216316416516616716816917017117217317417517617717817918018118218318418518618718818919019119219319419519619719819920020120220320420520620720820921021121221321421521621721821922022122222322422522622722822923023123223323423523623723823924024124224324424524624724824925025125225325425525625725825926026126226326426526626726826927027127227327427527627727827928028128228328428528628728828929029129229329429529629729829930030130230330430530630730830931031131231331431531631731831932032132232332432532632732832933033133233333433533633733833934034134234334434534634734834935035135235335435535635735835936036136236336436536636736836937037137237337437537637737837938038138238338438538638738838939039139239339439539639739839940040140240340440540640740840941041141241341441541641741841942042142242342442542642742842943043143243343443543643743843944044144244344444544644744844945045145245345445545645745845946046146246346446546646746846947047147247347447547647747847948048148248348448548648748848949049149249349449549649749849950050150250350450550650750850951051151251351451551651751851952052152252352452552652752852953053153253353453553653753853954054154254354454554654754854955055155255355455555655755855956056156256356456556656756856957057157257357457557657757857958058158258358458558658758858959059159259359459559659759859960060160260360460560660760860961061161261361461561661761861962062162262362462562662762862963063163263363463563663763863964064164264364464564664764864965065165265365465565665765865966066166266366466566666766866967067167267367467567667767867968068168268368468568668768868969069169269369469569669769869970070170270370470570670770870971071171271371471571671771871972072172272372472572672772872973073173273373473573673773873974074174274374474574674774874975075175275375475575675775875976076176276376476576676776876977077177277377477577677777877978078178278378478578678778878979079179279379479579679779879980080180280380480580680780880981081181281381481581681781881982082182282382482582682782882983083183283383483583683783883984084184284384484584684784884985085185285385485585685785885986086186286386486586686786886987087187287387487587687787887988088188288388488588688788888989089189289389489589689789889990090190290390490590690790890991091191291391491591691791891992092192292392492592692792892993093193293393493593693793893994094194294394494594694794894995095195295395495595695795895996096196296396496596696796896997097197297397497597697797897998098198298398498598698798898999099199299399499599699799899910001001100210031004100510061007100810091010101110121013101410151016101710181019102010211022102310241025102610271028102910301031103210331034103510361037103810391040104110421043104410451046104710481049105010511052105310541055105610571058105910601061106210631064106510661067106810691070107110721073107410751076107710781079108010811082108310841085108610871088108910901091109210931094109510961097109810991100110111021103110411051106110711081109111011111112111311141115111611171118111911201121112211231124112511261127112811291130113111321133113411351136113711381139114011411142114311441145114611471148114911501151115211531154
  1. // Copyright 2018 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 protogen provides support for writing protoc plugins.
  5. //
  6. // Plugins for protoc, the Protocol Buffers Compiler, are programs which read
  7. // a CodeGeneratorRequest protocol buffer from standard input and write a
  8. // CodeGeneratorResponse protocol buffer to standard output. This package
  9. // provides support for writing plugins which generate Go code.
  10. package protogen
  11. import (
  12. "bufio"
  13. "bytes"
  14. "encoding/binary"
  15. "fmt"
  16. "go/ast"
  17. "go/parser"
  18. "go/printer"
  19. "go/token"
  20. "go/types"
  21. "io/ioutil"
  22. "os"
  23. "path"
  24. "path/filepath"
  25. "sort"
  26. "strconv"
  27. "strings"
  28. "github.com/golang/protobuf/proto"
  29. "github.com/golang/protobuf/v2/internal/descfield"
  30. "github.com/golang/protobuf/v2/internal/scalar"
  31. "github.com/golang/protobuf/v2/reflect/protodesc"
  32. "github.com/golang/protobuf/v2/reflect/protoreflect"
  33. "github.com/golang/protobuf/v2/reflect/protoregistry"
  34. descriptorpb "github.com/golang/protobuf/v2/types/descriptor"
  35. pluginpb "github.com/golang/protobuf/v2/types/plugin"
  36. )
  37. // Run executes a function as a protoc plugin.
  38. //
  39. // It reads a CodeGeneratorRequest message from os.Stdin, invokes the plugin
  40. // function, and writes a CodeGeneratorResponse message to os.Stdout.
  41. //
  42. // If a failure occurs while reading or writing, Run prints an error to
  43. // os.Stderr and calls os.Exit(1).
  44. //
  45. // Passing a nil options is equivalent to passing a zero-valued one.
  46. func Run(opts *Options, f func(*Plugin) error) {
  47. if err := run(opts, f); err != nil {
  48. fmt.Fprintf(os.Stderr, "%s: %v\n", filepath.Base(os.Args[0]), err)
  49. os.Exit(1)
  50. }
  51. }
  52. func run(opts *Options, f func(*Plugin) error) error {
  53. if len(os.Args) > 1 {
  54. return fmt.Errorf("unknown argument %q (this program should be run by protoc, not directly)", os.Args[1])
  55. }
  56. in, err := ioutil.ReadAll(os.Stdin)
  57. if err != nil {
  58. return err
  59. }
  60. req := &pluginpb.CodeGeneratorRequest{}
  61. if err := proto.Unmarshal(in, req); err != nil {
  62. return err
  63. }
  64. gen, err := New(req, opts)
  65. if err != nil {
  66. return err
  67. }
  68. if err := f(gen); err != nil {
  69. // Errors from the plugin function are reported by setting the
  70. // error field in the CodeGeneratorResponse.
  71. //
  72. // In contrast, errors that indicate a problem in protoc
  73. // itself (unparsable input, I/O errors, etc.) are reported
  74. // to stderr.
  75. gen.Error(err)
  76. }
  77. resp := gen.Response()
  78. out, err := proto.Marshal(resp)
  79. if err != nil {
  80. return err
  81. }
  82. if _, err := os.Stdout.Write(out); err != nil {
  83. return err
  84. }
  85. return nil
  86. }
  87. // A Plugin is a protoc plugin invocation.
  88. type Plugin struct {
  89. // Request is the CodeGeneratorRequest provided by protoc.
  90. Request *pluginpb.CodeGeneratorRequest
  91. // Files is the set of files to generate and everything they import.
  92. // Files appear in topological order, so each file appears before any
  93. // file that imports it.
  94. Files []*File
  95. filesByName map[string]*File
  96. fileReg *protoregistry.Files
  97. messagesByName map[protoreflect.FullName]*Message
  98. enumsByName map[protoreflect.FullName]*Enum
  99. annotateCode bool
  100. pathType pathType
  101. genFiles []*GeneratedFile
  102. opts *Options
  103. err error
  104. }
  105. // Options are optional parameters to New.
  106. type Options struct {
  107. // If ParamFunc is non-nil, it will be called with each unknown
  108. // generator parameter.
  109. //
  110. // Plugins for protoc can accept parameters from the command line,
  111. // passed in the --<lang>_out protoc, separated from the output
  112. // directory with a colon; e.g.,
  113. //
  114. // --go_out=<param1>=<value1>,<param2>=<value2>:<output_directory>
  115. //
  116. // Parameters passed in this fashion as a comma-separated list of
  117. // key=value pairs will be passed to the ParamFunc.
  118. //
  119. // The (flag.FlagSet).Set method matches this function signature,
  120. // so parameters can be converted into flags as in the following:
  121. //
  122. // var flags flag.FlagSet
  123. // value := flags.Bool("param", false, "")
  124. // opts := &protogen.Options{
  125. // ParamFunc: flags.Set,
  126. // }
  127. // protogen.Run(opts, func(p *protogen.Plugin) error {
  128. // if *value { ... }
  129. // })
  130. ParamFunc func(name, value string) error
  131. // ImportRewriteFunc is called with the import path of each package
  132. // imported by a generated file. It returns the import path to use
  133. // for this package.
  134. ImportRewriteFunc func(GoImportPath) GoImportPath
  135. }
  136. // New returns a new Plugin.
  137. //
  138. // Passing a nil Options is equivalent to passing a zero-valued one.
  139. func New(req *pluginpb.CodeGeneratorRequest, opts *Options) (*Plugin, error) {
  140. if opts == nil {
  141. opts = &Options{}
  142. }
  143. gen := &Plugin{
  144. Request: req,
  145. filesByName: make(map[string]*File),
  146. fileReg: protoregistry.NewFiles(),
  147. messagesByName: make(map[protoreflect.FullName]*Message),
  148. enumsByName: make(map[protoreflect.FullName]*Enum),
  149. opts: opts,
  150. }
  151. packageNames := make(map[string]GoPackageName) // filename -> package name
  152. importPaths := make(map[string]GoImportPath) // filename -> import path
  153. var packageImportPath GoImportPath
  154. for _, param := range strings.Split(req.GetParameter(), ",") {
  155. var value string
  156. if i := strings.Index(param, "="); i >= 0 {
  157. value = param[i+1:]
  158. param = param[0:i]
  159. }
  160. switch param {
  161. case "":
  162. // Ignore.
  163. case "import_path":
  164. packageImportPath = GoImportPath(value)
  165. case "paths":
  166. switch value {
  167. case "import":
  168. gen.pathType = pathTypeImport
  169. case "source_relative":
  170. gen.pathType = pathTypeSourceRelative
  171. default:
  172. return nil, fmt.Errorf(`unknown path type %q: want "import" or "source_relative"`, value)
  173. }
  174. case "annotate_code":
  175. switch value {
  176. case "true", "":
  177. gen.annotateCode = true
  178. case "false":
  179. default:
  180. return nil, fmt.Errorf(`bad value for parameter %q: want "true" or "false"`, param)
  181. }
  182. default:
  183. if param[0] == 'M' {
  184. importPaths[param[1:]] = GoImportPath(value)
  185. continue
  186. }
  187. if opts.ParamFunc != nil {
  188. if err := opts.ParamFunc(param, value); err != nil {
  189. return nil, err
  190. }
  191. }
  192. }
  193. }
  194. // Figure out the import path and package name for each file.
  195. //
  196. // The rules here are complicated and have grown organically over time.
  197. // Interactions between different ways of specifying package information
  198. // may be surprising.
  199. //
  200. // The recommended approach is to include a go_package option in every
  201. // .proto source file specifying the full import path of the Go package
  202. // associated with this file.
  203. //
  204. // option go_package = "github.com/golang/protobuf/ptypes/any";
  205. //
  206. // Build systems which want to exert full control over import paths may
  207. // specify M<filename>=<import_path> flags.
  208. //
  209. // Other approaches are not recommend.
  210. generatedFileNames := make(map[string]bool)
  211. for _, name := range gen.Request.FileToGenerate {
  212. generatedFileNames[name] = true
  213. }
  214. // We need to determine the import paths before the package names,
  215. // because the Go package name for a file is sometimes derived from
  216. // different file in the same package.
  217. packageNameForImportPath := make(map[GoImportPath]GoPackageName)
  218. for _, fdesc := range gen.Request.ProtoFile {
  219. filename := fdesc.GetName()
  220. packageName, importPath := goPackageOption(fdesc)
  221. switch {
  222. case importPaths[filename] != "":
  223. // Command line: M=foo.proto=quux/bar
  224. //
  225. // Explicit mapping of source file to import path.
  226. case generatedFileNames[filename] && packageImportPath != "":
  227. // Command line: import_path=quux/bar
  228. //
  229. // The import_path flag sets the import path for every file that
  230. // we generate code for.
  231. importPaths[filename] = packageImportPath
  232. case importPath != "":
  233. // Source file: option go_package = "quux/bar";
  234. //
  235. // The go_package option sets the import path. Most users should use this.
  236. importPaths[filename] = importPath
  237. default:
  238. // Source filename.
  239. //
  240. // Last resort when nothing else is available.
  241. importPaths[filename] = GoImportPath(path.Dir(filename))
  242. }
  243. if packageName != "" {
  244. packageNameForImportPath[importPaths[filename]] = packageName
  245. }
  246. }
  247. for _, fdesc := range gen.Request.ProtoFile {
  248. filename := fdesc.GetName()
  249. packageName, _ := goPackageOption(fdesc)
  250. defaultPackageName := packageNameForImportPath[importPaths[filename]]
  251. switch {
  252. case packageName != "":
  253. // Source file: option go_package = "quux/bar";
  254. packageNames[filename] = packageName
  255. case defaultPackageName != "":
  256. // A go_package option in another file in the same package.
  257. //
  258. // This is a poor choice in general, since every source file should
  259. // contain a go_package option. Supported mainly for historical
  260. // compatibility.
  261. packageNames[filename] = defaultPackageName
  262. case generatedFileNames[filename] && packageImportPath != "":
  263. // Command line: import_path=quux/bar
  264. packageNames[filename] = cleanPackageName(path.Base(string(packageImportPath)))
  265. case fdesc.GetPackage() != "":
  266. // Source file: package quux.bar;
  267. packageNames[filename] = cleanPackageName(fdesc.GetPackage())
  268. default:
  269. // Source filename.
  270. packageNames[filename] = cleanPackageName(baseName(filename))
  271. }
  272. }
  273. // Consistency check: Every file with the same Go import path should have
  274. // the same Go package name.
  275. packageFiles := make(map[GoImportPath][]string)
  276. for filename, importPath := range importPaths {
  277. if _, ok := packageNames[filename]; !ok {
  278. // Skip files mentioned in a M<file>=<import_path> parameter
  279. // but which do not appear in the CodeGeneratorRequest.
  280. continue
  281. }
  282. packageFiles[importPath] = append(packageFiles[importPath], filename)
  283. }
  284. for importPath, filenames := range packageFiles {
  285. for i := 1; i < len(filenames); i++ {
  286. if a, b := packageNames[filenames[0]], packageNames[filenames[i]]; a != b {
  287. return nil, fmt.Errorf("Go package %v has inconsistent names %v (%v) and %v (%v)",
  288. importPath, a, filenames[0], b, filenames[i])
  289. }
  290. }
  291. }
  292. for _, fdesc := range gen.Request.ProtoFile {
  293. filename := fdesc.GetName()
  294. if gen.filesByName[filename] != nil {
  295. return nil, fmt.Errorf("duplicate file name: %q", filename)
  296. }
  297. f, err := newFile(gen, fdesc, packageNames[filename], importPaths[filename])
  298. if err != nil {
  299. return nil, err
  300. }
  301. gen.Files = append(gen.Files, f)
  302. gen.filesByName[filename] = f
  303. }
  304. for _, filename := range gen.Request.FileToGenerate {
  305. f, ok := gen.FileByName(filename)
  306. if !ok {
  307. return nil, fmt.Errorf("no descriptor for generated file: %v", filename)
  308. }
  309. f.Generate = true
  310. }
  311. return gen, nil
  312. }
  313. // Error records an error in code generation. The generator will report the
  314. // error back to protoc and will not produce output.
  315. func (gen *Plugin) Error(err error) {
  316. if gen.err == nil {
  317. gen.err = err
  318. }
  319. }
  320. // Response returns the generator output.
  321. func (gen *Plugin) Response() *pluginpb.CodeGeneratorResponse {
  322. resp := &pluginpb.CodeGeneratorResponse{}
  323. if gen.err != nil {
  324. resp.Error = scalar.String(gen.err.Error())
  325. return resp
  326. }
  327. for _, g := range gen.genFiles {
  328. if g.skip {
  329. continue
  330. }
  331. content, err := g.Content()
  332. if err != nil {
  333. return &pluginpb.CodeGeneratorResponse{
  334. Error: scalar.String(err.Error()),
  335. }
  336. }
  337. resp.File = append(resp.File, &pluginpb.CodeGeneratorResponse_File{
  338. Name: scalar.String(g.filename),
  339. Content: scalar.String(string(content)),
  340. })
  341. if gen.annotateCode && strings.HasSuffix(g.filename, ".go") {
  342. meta, err := g.metaFile(content)
  343. if err != nil {
  344. return &pluginpb.CodeGeneratorResponse{
  345. Error: scalar.String(err.Error()),
  346. }
  347. }
  348. resp.File = append(resp.File, &pluginpb.CodeGeneratorResponse_File{
  349. Name: scalar.String(g.filename + ".meta"),
  350. Content: scalar.String(meta),
  351. })
  352. }
  353. }
  354. return resp
  355. }
  356. // FileByName returns the file with the given name.
  357. func (gen *Plugin) FileByName(name string) (f *File, ok bool) {
  358. f, ok = gen.filesByName[name]
  359. return f, ok
  360. }
  361. // A File describes a .proto source file.
  362. type File struct {
  363. Desc protoreflect.FileDescriptor
  364. Proto *descriptorpb.FileDescriptorProto
  365. GoDescriptorIdent GoIdent // name of Go variable for the file descriptor
  366. GoPackageName GoPackageName // name of this file's Go package
  367. GoImportPath GoImportPath // import path of this file's Go package
  368. Messages []*Message // top-level message declarations
  369. Enums []*Enum // top-level enum declarations
  370. Extensions []*Extension // top-level extension declarations
  371. Services []*Service // top-level service declarations
  372. Generate bool // true if we should generate code for this file
  373. // GeneratedFilenamePrefix is used to construct filenames for generated
  374. // files associated with this source file.
  375. //
  376. // For example, the source file "dir/foo.proto" might have a filename prefix
  377. // of "dir/foo". Appending ".pb.go" produces an output file of "dir/foo.pb.go".
  378. GeneratedFilenamePrefix string
  379. sourceInfo map[pathKey][]*descriptorpb.SourceCodeInfo_Location
  380. }
  381. func newFile(gen *Plugin, p *descriptorpb.FileDescriptorProto, packageName GoPackageName, importPath GoImportPath) (*File, error) {
  382. desc, err := protodesc.NewFile(p, gen.fileReg)
  383. if err != nil {
  384. return nil, fmt.Errorf("invalid FileDescriptorProto %q: %v", p.GetName(), err)
  385. }
  386. if err := gen.fileReg.Register(desc); err != nil {
  387. return nil, fmt.Errorf("cannot register descriptor %q: %v", p.GetName(), err)
  388. }
  389. f := &File{
  390. Desc: desc,
  391. Proto: p,
  392. GoPackageName: packageName,
  393. GoImportPath: importPath,
  394. sourceInfo: make(map[pathKey][]*descriptorpb.SourceCodeInfo_Location),
  395. }
  396. // Determine the prefix for generated Go files.
  397. prefix := p.GetName()
  398. if ext := path.Ext(prefix); ext == ".proto" || ext == ".protodevel" {
  399. prefix = prefix[:len(prefix)-len(ext)]
  400. }
  401. if gen.pathType == pathTypeImport {
  402. // If paths=import (the default) and the file contains a go_package option
  403. // with a full import path, the output filename is derived from the Go import
  404. // path.
  405. //
  406. // Pass the paths=source_relative flag to always derive the output filename
  407. // from the input filename instead.
  408. if _, importPath := goPackageOption(p); importPath != "" {
  409. prefix = path.Join(string(importPath), path.Base(prefix))
  410. }
  411. }
  412. f.GoDescriptorIdent = GoIdent{
  413. GoName: "File_" + cleanGoName(p.GetName()),
  414. GoImportPath: f.GoImportPath,
  415. }
  416. f.GeneratedFilenamePrefix = prefix
  417. for _, loc := range p.GetSourceCodeInfo().GetLocation() {
  418. key := newPathKey(loc.Path)
  419. f.sourceInfo[key] = append(f.sourceInfo[key], loc)
  420. }
  421. for i, mdescs := 0, desc.Messages(); i < mdescs.Len(); i++ {
  422. f.Messages = append(f.Messages, newMessage(gen, f, nil, mdescs.Get(i)))
  423. }
  424. for i, edescs := 0, desc.Enums(); i < edescs.Len(); i++ {
  425. f.Enums = append(f.Enums, newEnum(gen, f, nil, edescs.Get(i)))
  426. }
  427. for i, extdescs := 0, desc.Extensions(); i < extdescs.Len(); i++ {
  428. f.Extensions = append(f.Extensions, newField(gen, f, nil, extdescs.Get(i)))
  429. }
  430. for i, sdescs := 0, desc.Services(); i < sdescs.Len(); i++ {
  431. f.Services = append(f.Services, newService(gen, f, sdescs.Get(i)))
  432. }
  433. for _, message := range f.Messages {
  434. if err := message.init(gen); err != nil {
  435. return nil, err
  436. }
  437. }
  438. for _, extension := range f.Extensions {
  439. if err := extension.init(gen); err != nil {
  440. return nil, err
  441. }
  442. }
  443. for _, service := range f.Services {
  444. for _, method := range service.Methods {
  445. if err := method.init(gen); err != nil {
  446. return nil, err
  447. }
  448. }
  449. }
  450. return f, nil
  451. }
  452. func (f *File) location(path ...int32) Location {
  453. return Location{
  454. SourceFile: f.Desc.Path(),
  455. Path: path,
  456. }
  457. }
  458. // goPackageOption interprets a file's go_package option.
  459. // If there is no go_package, it returns ("", "").
  460. // If there's a simple name, it returns (pkg, "").
  461. // If the option implies an import path, it returns (pkg, impPath).
  462. func goPackageOption(d *descriptorpb.FileDescriptorProto) (pkg GoPackageName, impPath GoImportPath) {
  463. opt := d.GetOptions().GetGoPackage()
  464. if opt == "" {
  465. return "", ""
  466. }
  467. // A semicolon-delimited suffix delimits the import path and package name.
  468. if i := strings.Index(opt, ";"); i >= 0 {
  469. return cleanPackageName(opt[i+1:]), GoImportPath(opt[:i])
  470. }
  471. // The presence of a slash implies there's an import path.
  472. if i := strings.LastIndex(opt, "/"); i >= 0 {
  473. return cleanPackageName(opt[i+1:]), GoImportPath(opt)
  474. }
  475. return cleanPackageName(opt), ""
  476. }
  477. // A Message describes a message.
  478. type Message struct {
  479. Desc protoreflect.MessageDescriptor
  480. GoIdent GoIdent // name of the generated Go type
  481. Fields []*Field // message field declarations
  482. Oneofs []*Oneof // oneof declarations
  483. Messages []*Message // nested message declarations
  484. Enums []*Enum // nested enum declarations
  485. Extensions []*Extension // nested extension declarations
  486. Location Location // location of this message
  487. }
  488. func newMessage(gen *Plugin, f *File, parent *Message, desc protoreflect.MessageDescriptor) *Message {
  489. var loc Location
  490. if parent != nil {
  491. loc = parent.Location.appendPath(descfield.DescriptorProto_NestedType, int32(desc.Index()))
  492. } else {
  493. loc = f.location(descfield.FileDescriptorProto_MessageType, int32(desc.Index()))
  494. }
  495. message := &Message{
  496. Desc: desc,
  497. GoIdent: newGoIdent(f, desc),
  498. Location: loc,
  499. }
  500. gen.messagesByName[desc.FullName()] = message
  501. for i, mdescs := 0, desc.Messages(); i < mdescs.Len(); i++ {
  502. message.Messages = append(message.Messages, newMessage(gen, f, message, mdescs.Get(i)))
  503. }
  504. for i, edescs := 0, desc.Enums(); i < edescs.Len(); i++ {
  505. message.Enums = append(message.Enums, newEnum(gen, f, message, edescs.Get(i)))
  506. }
  507. for i, odescs := 0, desc.Oneofs(); i < odescs.Len(); i++ {
  508. message.Oneofs = append(message.Oneofs, newOneof(gen, f, message, odescs.Get(i)))
  509. }
  510. for i, fdescs := 0, desc.Fields(); i < fdescs.Len(); i++ {
  511. message.Fields = append(message.Fields, newField(gen, f, message, fdescs.Get(i)))
  512. }
  513. for i, extdescs := 0, desc.Extensions(); i < extdescs.Len(); i++ {
  514. message.Extensions = append(message.Extensions, newField(gen, f, message, extdescs.Get(i)))
  515. }
  516. // Field name conflict resolution.
  517. //
  518. // We assume well-known method names that may be attached to a generated
  519. // message type, as well as a 'Get*' method for each field. For each
  520. // field in turn, we add _s to its name until there are no conflicts.
  521. //
  522. // Any change to the following set of method names is a potential
  523. // incompatible API change because it may change generated field names.
  524. //
  525. // TODO: If we ever support a 'go_name' option to set the Go name of a
  526. // field, we should consider dropping this entirely. The conflict
  527. // resolution algorithm is subtle and surprising (changing the order
  528. // in which fields appear in the .proto source file can change the
  529. // names of fields in generated code), and does not adapt well to
  530. // adding new per-field methods such as setters.
  531. usedNames := map[string]bool{
  532. "Reset": true,
  533. "String": true,
  534. "ProtoMessage": true,
  535. "Marshal": true,
  536. "Unmarshal": true,
  537. "ExtensionRangeArray": true,
  538. "ExtensionMap": true,
  539. "Descriptor": true,
  540. }
  541. makeNameUnique := func(name string, hasGetter bool) string {
  542. for usedNames[name] || (hasGetter && usedNames["Get"+name]) {
  543. name += "_"
  544. }
  545. usedNames[name] = true
  546. usedNames["Get"+name] = hasGetter
  547. return name
  548. }
  549. seenOneofs := make(map[int]bool)
  550. for _, field := range message.Fields {
  551. field.GoName = makeNameUnique(field.GoName, true)
  552. if field.OneofType != nil {
  553. if !seenOneofs[field.OneofType.Desc.Index()] {
  554. // If this is a field in a oneof that we haven't seen before,
  555. // make the name for that oneof unique as well.
  556. field.OneofType.GoName = makeNameUnique(field.OneofType.GoName, false)
  557. seenOneofs[field.OneofType.Desc.Index()] = true
  558. }
  559. }
  560. }
  561. return message
  562. }
  563. func (message *Message) init(gen *Plugin) error {
  564. for _, child := range message.Messages {
  565. if err := child.init(gen); err != nil {
  566. return err
  567. }
  568. }
  569. for _, field := range message.Fields {
  570. if err := field.init(gen); err != nil {
  571. return err
  572. }
  573. }
  574. for _, oneof := range message.Oneofs {
  575. oneof.init(gen, message)
  576. }
  577. for _, extension := range message.Extensions {
  578. if err := extension.init(gen); err != nil {
  579. return err
  580. }
  581. }
  582. return nil
  583. }
  584. // A Field describes a message field.
  585. type Field struct {
  586. Desc protoreflect.FieldDescriptor
  587. // GoName is the base name of this field's Go field and methods.
  588. // For code generated by protoc-gen-go, this means a field named
  589. // '{{GoName}}' and a getter method named 'Get{{GoName}}'.
  590. GoName string
  591. ParentMessage *Message // message in which this field is defined; nil if top-level extension
  592. ExtendedType *Message // extended message for extension fields; nil otherwise
  593. MessageType *Message // type for message or group fields; nil otherwise
  594. EnumType *Enum // type for enum fields; nil otherwise
  595. OneofType *Oneof // containing oneof; nil if not part of a oneof
  596. Location Location // location of this field
  597. }
  598. func newField(gen *Plugin, f *File, message *Message, desc protoreflect.FieldDescriptor) *Field {
  599. var loc Location
  600. switch {
  601. case desc.ExtendedType() != nil && message == nil:
  602. loc = f.location(descfield.FileDescriptorProto_Extension, int32(desc.Index()))
  603. case desc.ExtendedType() != nil && message != nil:
  604. loc = message.Location.appendPath(descfield.DescriptorProto_Extension, int32(desc.Index()))
  605. default:
  606. loc = message.Location.appendPath(descfield.DescriptorProto_Field, int32(desc.Index()))
  607. }
  608. field := &Field{
  609. Desc: desc,
  610. GoName: camelCase(string(desc.Name())),
  611. ParentMessage: message,
  612. Location: loc,
  613. }
  614. if desc.OneofType() != nil {
  615. field.OneofType = message.Oneofs[desc.OneofType().Index()]
  616. }
  617. return field
  618. }
  619. // Extension is an alias of Field for documentation.
  620. type Extension = Field
  621. func (field *Field) init(gen *Plugin) error {
  622. desc := field.Desc
  623. switch desc.Kind() {
  624. case protoreflect.MessageKind, protoreflect.GroupKind:
  625. mname := desc.MessageType().FullName()
  626. message, ok := gen.messagesByName[mname]
  627. if !ok {
  628. return fmt.Errorf("field %v: no descriptor for type %v", desc.FullName(), mname)
  629. }
  630. field.MessageType = message
  631. case protoreflect.EnumKind:
  632. ename := field.Desc.EnumType().FullName()
  633. enum, ok := gen.enumsByName[ename]
  634. if !ok {
  635. return fmt.Errorf("field %v: no descriptor for enum %v", desc.FullName(), ename)
  636. }
  637. field.EnumType = enum
  638. }
  639. if desc.ExtendedType() != nil {
  640. mname := desc.ExtendedType().FullName()
  641. message, ok := gen.messagesByName[mname]
  642. if !ok {
  643. return fmt.Errorf("field %v: no descriptor for type %v", desc.FullName(), mname)
  644. }
  645. field.ExtendedType = message
  646. }
  647. return nil
  648. }
  649. // A Oneof describes a oneof field.
  650. type Oneof struct {
  651. Desc protoreflect.OneofDescriptor
  652. GoName string // Go field name of this oneof
  653. ParentMessage *Message // message in which this oneof occurs
  654. Fields []*Field // fields that are part of this oneof
  655. Location Location // location of this oneof
  656. }
  657. func newOneof(gen *Plugin, f *File, message *Message, desc protoreflect.OneofDescriptor) *Oneof {
  658. return &Oneof{
  659. Desc: desc,
  660. ParentMessage: message,
  661. GoName: camelCase(string(desc.Name())),
  662. Location: message.Location.appendPath(descfield.DescriptorProto_OneofDecl, int32(desc.Index())),
  663. }
  664. }
  665. func (oneof *Oneof) init(gen *Plugin, parent *Message) {
  666. for i, fdescs := 0, oneof.Desc.Fields(); i < fdescs.Len(); i++ {
  667. oneof.Fields = append(oneof.Fields, parent.Fields[fdescs.Get(i).Index()])
  668. }
  669. }
  670. // An Enum describes an enum.
  671. type Enum struct {
  672. Desc protoreflect.EnumDescriptor
  673. GoIdent GoIdent // name of the generated Go type
  674. Values []*EnumValue // enum values
  675. Location Location // location of this enum
  676. }
  677. func newEnum(gen *Plugin, f *File, parent *Message, desc protoreflect.EnumDescriptor) *Enum {
  678. var loc Location
  679. if parent != nil {
  680. loc = parent.Location.appendPath(descfield.DescriptorProto_EnumType, int32(desc.Index()))
  681. } else {
  682. loc = f.location(descfield.FileDescriptorProto_EnumType, int32(desc.Index()))
  683. }
  684. enum := &Enum{
  685. Desc: desc,
  686. GoIdent: newGoIdent(f, desc),
  687. Location: loc,
  688. }
  689. gen.enumsByName[desc.FullName()] = enum
  690. for i, evdescs := 0, enum.Desc.Values(); i < evdescs.Len(); i++ {
  691. enum.Values = append(enum.Values, newEnumValue(gen, f, parent, enum, evdescs.Get(i)))
  692. }
  693. return enum
  694. }
  695. // An EnumValue describes an enum value.
  696. type EnumValue struct {
  697. Desc protoreflect.EnumValueDescriptor
  698. GoIdent GoIdent // name of the generated Go type
  699. Location Location // location of this enum value
  700. }
  701. func newEnumValue(gen *Plugin, f *File, message *Message, enum *Enum, desc protoreflect.EnumValueDescriptor) *EnumValue {
  702. // A top-level enum value's name is: EnumName_ValueName
  703. // An enum value contained in a message is: MessageName_ValueName
  704. //
  705. // Enum value names are not camelcased.
  706. parentIdent := enum.GoIdent
  707. if message != nil {
  708. parentIdent = message.GoIdent
  709. }
  710. name := parentIdent.GoName + "_" + string(desc.Name())
  711. return &EnumValue{
  712. Desc: desc,
  713. GoIdent: f.GoImportPath.Ident(name),
  714. Location: enum.Location.appendPath(descfield.EnumDescriptorProto_Value, int32(desc.Index())),
  715. }
  716. }
  717. // A Service describes a service.
  718. type Service struct {
  719. Desc protoreflect.ServiceDescriptor
  720. GoName string
  721. Location Location // location of this service
  722. Methods []*Method // service method definitions
  723. }
  724. func newService(gen *Plugin, f *File, desc protoreflect.ServiceDescriptor) *Service {
  725. service := &Service{
  726. Desc: desc,
  727. GoName: camelCase(string(desc.Name())),
  728. Location: f.location(descfield.FileDescriptorProto_Service, int32(desc.Index())),
  729. }
  730. for i, mdescs := 0, desc.Methods(); i < mdescs.Len(); i++ {
  731. service.Methods = append(service.Methods, newMethod(gen, f, service, mdescs.Get(i)))
  732. }
  733. return service
  734. }
  735. // A Method describes a method in a service.
  736. type Method struct {
  737. Desc protoreflect.MethodDescriptor
  738. GoName string
  739. ParentService *Service
  740. Location Location // location of this method
  741. InputType *Message
  742. OutputType *Message
  743. }
  744. func newMethod(gen *Plugin, f *File, service *Service, desc protoreflect.MethodDescriptor) *Method {
  745. method := &Method{
  746. Desc: desc,
  747. GoName: camelCase(string(desc.Name())),
  748. ParentService: service,
  749. Location: service.Location.appendPath(descfield.ServiceDescriptorProto_Method, int32(desc.Index())),
  750. }
  751. return method
  752. }
  753. func (method *Method) init(gen *Plugin) error {
  754. desc := method.Desc
  755. inName := desc.InputType().FullName()
  756. in, ok := gen.messagesByName[inName]
  757. if !ok {
  758. return fmt.Errorf("method %v: no descriptor for type %v", desc.FullName(), inName)
  759. }
  760. method.InputType = in
  761. outName := desc.OutputType().FullName()
  762. out, ok := gen.messagesByName[outName]
  763. if !ok {
  764. return fmt.Errorf("method %v: no descriptor for type %v", desc.FullName(), outName)
  765. }
  766. method.OutputType = out
  767. return nil
  768. }
  769. // A GeneratedFile is a generated file.
  770. type GeneratedFile struct {
  771. gen *Plugin
  772. skip bool
  773. filename string
  774. goImportPath GoImportPath
  775. buf bytes.Buffer
  776. packageNames map[GoImportPath]GoPackageName
  777. usedPackageNames map[GoPackageName]bool
  778. manualImports map[GoImportPath]bool
  779. annotations map[string][]Location
  780. }
  781. // NewGeneratedFile creates a new generated file with the given filename
  782. // and import path.
  783. func (gen *Plugin) NewGeneratedFile(filename string, goImportPath GoImportPath) *GeneratedFile {
  784. g := &GeneratedFile{
  785. gen: gen,
  786. filename: filename,
  787. goImportPath: goImportPath,
  788. packageNames: make(map[GoImportPath]GoPackageName),
  789. usedPackageNames: make(map[GoPackageName]bool),
  790. manualImports: make(map[GoImportPath]bool),
  791. annotations: make(map[string][]Location),
  792. }
  793. // All predeclared identifiers in Go are already used.
  794. for _, s := range types.Universe.Names() {
  795. g.usedPackageNames[GoPackageName(s)] = true
  796. }
  797. gen.genFiles = append(gen.genFiles, g)
  798. return g
  799. }
  800. // P prints a line to the generated output. It converts each parameter to a
  801. // string following the same rules as fmt.Print. It never inserts spaces
  802. // between parameters.
  803. func (g *GeneratedFile) P(v ...interface{}) {
  804. for _, x := range v {
  805. switch x := x.(type) {
  806. case GoIdent:
  807. fmt.Fprint(&g.buf, g.QualifiedGoIdent(x))
  808. default:
  809. fmt.Fprint(&g.buf, x)
  810. }
  811. }
  812. fmt.Fprintln(&g.buf)
  813. }
  814. // PrintLeadingComments writes the comment appearing before a location in
  815. // the .proto source to the generated file.
  816. //
  817. // It returns true if a comment was present at the location.
  818. func (g *GeneratedFile) PrintLeadingComments(loc Location) (hasComment bool) {
  819. f := g.gen.filesByName[loc.SourceFile]
  820. if f == nil {
  821. return false
  822. }
  823. for _, infoLoc := range f.sourceInfo[newPathKey(loc.Path)] {
  824. if infoLoc.LeadingComments == nil {
  825. continue
  826. }
  827. for _, line := range strings.Split(strings.TrimSuffix(infoLoc.GetLeadingComments(), "\n"), "\n") {
  828. g.buf.WriteString("//")
  829. g.buf.WriteString(line)
  830. g.buf.WriteString("\n")
  831. }
  832. return true
  833. }
  834. return false
  835. }
  836. // QualifiedGoIdent returns the string to use for a Go identifier.
  837. //
  838. // If the identifier is from a different Go package than the generated file,
  839. // the returned name will be qualified (package.name) and an import statement
  840. // for the identifier's package will be included in the file.
  841. func (g *GeneratedFile) QualifiedGoIdent(ident GoIdent) string {
  842. if ident.GoImportPath == g.goImportPath {
  843. return ident.GoName
  844. }
  845. if packageName, ok := g.packageNames[ident.GoImportPath]; ok {
  846. return string(packageName) + "." + ident.GoName
  847. }
  848. packageName := cleanPackageName(baseName(string(ident.GoImportPath)))
  849. for i, orig := 1, packageName; g.usedPackageNames[packageName]; i++ {
  850. packageName = orig + GoPackageName(strconv.Itoa(i))
  851. }
  852. g.packageNames[ident.GoImportPath] = packageName
  853. g.usedPackageNames[packageName] = true
  854. return string(packageName) + "." + ident.GoName
  855. }
  856. // Import ensures a package is imported by the generated file.
  857. //
  858. // Packages referenced by QualifiedGoIdent are automatically imported.
  859. // Explicitly importing a package with Import is generally only necessary
  860. // when the import will be blank (import _ "package").
  861. func (g *GeneratedFile) Import(importPath GoImportPath) {
  862. g.manualImports[importPath] = true
  863. }
  864. // Write implements io.Writer.
  865. func (g *GeneratedFile) Write(p []byte) (n int, err error) {
  866. return g.buf.Write(p)
  867. }
  868. // Skip removes the generated file from the plugin output.
  869. func (g *GeneratedFile) Skip() {
  870. g.skip = true
  871. }
  872. // Annotate associates a symbol in a generated Go file with a location in a
  873. // source .proto file.
  874. //
  875. // The symbol may refer to a type, constant, variable, function, method, or
  876. // struct field. The "T.sel" syntax is used to identify the method or field
  877. // 'sel' on type 'T'.
  878. func (g *GeneratedFile) Annotate(symbol string, loc Location) {
  879. g.annotations[symbol] = append(g.annotations[symbol], loc)
  880. }
  881. // Content returns the contents of the generated file.
  882. func (g *GeneratedFile) Content() ([]byte, error) {
  883. if !strings.HasSuffix(g.filename, ".go") {
  884. return g.buf.Bytes(), nil
  885. }
  886. // Reformat generated code.
  887. original := g.buf.Bytes()
  888. fset := token.NewFileSet()
  889. file, err := parser.ParseFile(fset, "", original, parser.ParseComments)
  890. if err != nil {
  891. // Print out the bad code with line numbers.
  892. // This should never happen in practice, but it can while changing generated code
  893. // so consider this a debugging aid.
  894. var src bytes.Buffer
  895. s := bufio.NewScanner(bytes.NewReader(original))
  896. for line := 1; s.Scan(); line++ {
  897. fmt.Fprintf(&src, "%5d\t%s\n", line, s.Bytes())
  898. }
  899. return nil, fmt.Errorf("%v: unparsable Go source: %v\n%v", g.filename, err, src.String())
  900. }
  901. // Collect a sorted list of all imports.
  902. var importPaths [][2]string
  903. rewriteImport := func(importPath string) string {
  904. if f := g.gen.opts.ImportRewriteFunc; f != nil {
  905. return string(f(GoImportPath(importPath)))
  906. }
  907. return importPath
  908. }
  909. for importPath := range g.packageNames {
  910. pkgName := string(g.packageNames[GoImportPath(importPath)])
  911. pkgPath := rewriteImport(string(importPath))
  912. importPaths = append(importPaths, [2]string{pkgName, pkgPath})
  913. }
  914. for importPath := range g.manualImports {
  915. if _, ok := g.packageNames[importPath]; !ok {
  916. pkgPath := rewriteImport(string(importPath))
  917. importPaths = append(importPaths, [2]string{"_", pkgPath})
  918. }
  919. }
  920. sort.Slice(importPaths, func(i, j int) bool {
  921. return importPaths[i][1] < importPaths[j][1]
  922. })
  923. // Modify the AST to include a new import block.
  924. if len(importPaths) > 0 {
  925. // Insert block after package statement or
  926. // possible comment attached to the end of the package statement.
  927. pos := file.Package
  928. tokFile := fset.File(file.Package)
  929. pkgLine := tokFile.Line(file.Package)
  930. for _, c := range file.Comments {
  931. if tokFile.Line(c.Pos()) > pkgLine {
  932. break
  933. }
  934. pos = c.End()
  935. }
  936. // Construct the import block.
  937. impDecl := &ast.GenDecl{
  938. Tok: token.IMPORT,
  939. TokPos: pos,
  940. Lparen: pos,
  941. Rparen: pos,
  942. }
  943. for _, importPath := range importPaths {
  944. impDecl.Specs = append(impDecl.Specs, &ast.ImportSpec{
  945. Name: &ast.Ident{
  946. Name: importPath[0],
  947. NamePos: pos,
  948. },
  949. Path: &ast.BasicLit{
  950. Kind: token.STRING,
  951. Value: strconv.Quote(importPath[1]),
  952. ValuePos: pos,
  953. },
  954. EndPos: pos,
  955. })
  956. }
  957. file.Decls = append([]ast.Decl{impDecl}, file.Decls...)
  958. }
  959. var out bytes.Buffer
  960. if err = (&printer.Config{Mode: printer.TabIndent | printer.UseSpaces, Tabwidth: 8}).Fprint(&out, fset, file); err != nil {
  961. return nil, fmt.Errorf("%v: can not reformat Go source: %v", g.filename, err)
  962. }
  963. return out.Bytes(), nil
  964. }
  965. // metaFile returns the contents of the file's metadata file, which is a
  966. // text formatted string of the google.protobuf.GeneratedCodeInfo.
  967. func (g *GeneratedFile) metaFile(content []byte) (string, error) {
  968. fset := token.NewFileSet()
  969. astFile, err := parser.ParseFile(fset, "", content, 0)
  970. if err != nil {
  971. return "", err
  972. }
  973. info := &descriptorpb.GeneratedCodeInfo{}
  974. seenAnnotations := make(map[string]bool)
  975. annotate := func(s string, ident *ast.Ident) {
  976. seenAnnotations[s] = true
  977. for _, loc := range g.annotations[s] {
  978. info.Annotation = append(info.Annotation, &descriptorpb.GeneratedCodeInfo_Annotation{
  979. SourceFile: scalar.String(loc.SourceFile),
  980. Path: loc.Path,
  981. Begin: scalar.Int32(int32(fset.Position(ident.Pos()).Offset)),
  982. End: scalar.Int32(int32(fset.Position(ident.End()).Offset)),
  983. })
  984. }
  985. }
  986. for _, decl := range astFile.Decls {
  987. switch decl := decl.(type) {
  988. case *ast.GenDecl:
  989. for _, spec := range decl.Specs {
  990. switch spec := spec.(type) {
  991. case *ast.TypeSpec:
  992. annotate(spec.Name.Name, spec.Name)
  993. switch st := spec.Type.(type) {
  994. case *ast.StructType:
  995. for _, field := range st.Fields.List {
  996. for _, name := range field.Names {
  997. annotate(spec.Name.Name+"."+name.Name, name)
  998. }
  999. }
  1000. case *ast.InterfaceType:
  1001. for _, field := range st.Methods.List {
  1002. for _, name := range field.Names {
  1003. annotate(spec.Name.Name+"."+name.Name, name)
  1004. }
  1005. }
  1006. }
  1007. case *ast.ValueSpec:
  1008. for _, name := range spec.Names {
  1009. annotate(name.Name, name)
  1010. }
  1011. }
  1012. }
  1013. case *ast.FuncDecl:
  1014. if decl.Recv == nil {
  1015. annotate(decl.Name.Name, decl.Name)
  1016. } else {
  1017. recv := decl.Recv.List[0].Type
  1018. if s, ok := recv.(*ast.StarExpr); ok {
  1019. recv = s.X
  1020. }
  1021. if id, ok := recv.(*ast.Ident); ok {
  1022. annotate(id.Name+"."+decl.Name.Name, decl.Name)
  1023. }
  1024. }
  1025. }
  1026. }
  1027. for a := range g.annotations {
  1028. if !seenAnnotations[a] {
  1029. return "", fmt.Errorf("%v: no symbol matching annotation %q", g.filename, a)
  1030. }
  1031. }
  1032. return strings.TrimSpace(proto.CompactTextString(info)), nil
  1033. }
  1034. type pathType int
  1035. const (
  1036. pathTypeImport pathType = iota
  1037. pathTypeSourceRelative
  1038. )
  1039. // A Location is a location in a .proto source file.
  1040. //
  1041. // See the google.protobuf.SourceCodeInfo documentation in descriptor.proto
  1042. // for details.
  1043. type Location struct {
  1044. SourceFile string
  1045. Path []int32
  1046. }
  1047. // appendPath add elements to a Location's path, returning a new Location.
  1048. func (loc Location) appendPath(a ...int32) Location {
  1049. var n []int32
  1050. n = append(n, loc.Path...)
  1051. n = append(n, a...)
  1052. return Location{
  1053. SourceFile: loc.SourceFile,
  1054. Path: n,
  1055. }
  1056. }
  1057. // A pathKey is a representation of a location path suitable for use as a map key.
  1058. type pathKey struct {
  1059. s string
  1060. }
  1061. // newPathKey converts a location path to a pathKey.
  1062. func newPathKey(path []int32) pathKey {
  1063. buf := make([]byte, 4*len(path))
  1064. for i, x := range path {
  1065. binary.LittleEndian.PutUint32(buf[i*4:], uint32(x))
  1066. }
  1067. return pathKey{string(buf)}
  1068. }