protogen.go 35 KB

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