-
Notifications
You must be signed in to change notification settings - Fork 2.3k
/
registry.go
331 lines (291 loc) · 9.55 KB
/
registry.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
package descriptor
import (
"fmt"
"path"
"path/filepath"
"strings"
"github.com/golang/glog"
descriptor "github.com/golang/protobuf/protoc-gen-go/descriptor"
plugin "github.com/golang/protobuf/protoc-gen-go/plugin"
)
// Registry is a registry of information extracted from plugin.CodeGeneratorRequest.
type Registry struct {
// msgs is a mapping from fully-qualified message name to descriptor
msgs map[string]*Message
// enums is a mapping from fully-qualified enum name to descriptor
enums map[string]*Enum
// files is a mapping from file path to descriptor
files map[string]*File
// prefix is a prefix to be inserted to golang package paths generated from proto package names.
prefix string
// importPath is used as the package if no input files declare go_package. If it contains slashes, everything up to the rightmost slash is ignored.
importPath string
// pkgMap is a user-specified mapping from file path to proto package.
pkgMap map[string]string
// pkgAliases is a mapping from package aliases to package paths in go which are already taken.
pkgAliases map[string]string
// allowDeleteBody permits http delete methods to have a body
allowDeleteBody bool
}
// NewRegistry returns a new Registry.
func NewRegistry() *Registry {
return &Registry{
msgs: make(map[string]*Message),
enums: make(map[string]*Enum),
files: make(map[string]*File),
pkgMap: make(map[string]string),
pkgAliases: make(map[string]string),
}
}
// Load loads definitions of services, methods, messages, enumerations and fields from "req".
func (r *Registry) Load(req *plugin.CodeGeneratorRequest) error {
for _, file := range req.GetProtoFile() {
r.loadFile(file)
}
var targetPkg string
for _, name := range req.FileToGenerate {
target := r.files[name]
if target == nil {
return fmt.Errorf("no such file: %s", name)
}
name := packageIdentityName(target.FileDescriptorProto)
if targetPkg == "" {
targetPkg = name
} else {
if targetPkg != name {
return fmt.Errorf("inconsistent package names: %s %s", targetPkg, name)
}
}
if err := r.loadServices(target); err != nil {
return err
}
}
return nil
}
// loadFile loads messages, enumerations and fields from "file".
// It does not loads services and methods in "file". You need to call
// loadServices after loadFiles is called for all files to load services and methods.
func (r *Registry) loadFile(file *descriptor.FileDescriptorProto) {
pkg := GoPackage{
Path: r.goPackagePath(file),
Name: defaultGoPackageName(file),
}
if err := r.ReserveGoPackageAlias(pkg.Name, pkg.Path); err != nil {
for i := 0; ; i++ {
alias := fmt.Sprintf("%s_%d", pkg.Name, i)
if err := r.ReserveGoPackageAlias(alias, pkg.Path); err == nil {
pkg.Alias = alias
break
}
}
}
f := &File{
FileDescriptorProto: file,
GoPkg: pkg,
}
r.files[file.GetName()] = f
r.registerMsg(f, nil, file.GetMessageType())
r.registerEnum(f, nil, file.GetEnumType())
}
func (r *Registry) registerMsg(file *File, outerPath []string, msgs []*descriptor.DescriptorProto) {
for i, md := range msgs {
m := &Message{
File: file,
Outers: outerPath,
DescriptorProto: md,
Index: i,
}
for _, fd := range md.GetField() {
m.Fields = append(m.Fields, &Field{
Message: m,
FieldDescriptorProto: fd,
})
}
file.Messages = append(file.Messages, m)
r.msgs[m.FQMN()] = m
glog.V(1).Infof("register name: %s", m.FQMN())
var outers []string
outers = append(outers, outerPath...)
outers = append(outers, m.GetName())
r.registerMsg(file, outers, m.GetNestedType())
r.registerEnum(file, outers, m.GetEnumType())
}
}
func (r *Registry) registerEnum(file *File, outerPath []string, enums []*descriptor.EnumDescriptorProto) {
for i, ed := range enums {
e := &Enum{
File: file,
Outers: outerPath,
EnumDescriptorProto: ed,
Index: i,
}
file.Enums = append(file.Enums, e)
r.enums[e.FQEN()] = e
glog.V(1).Infof("register enum name: %s", e.FQEN())
}
}
// LookupMsg looks up a message type by "name".
// It tries to resolve "name" from "location" if "name" is a relative message name.
func (r *Registry) LookupMsg(location, name string) (*Message, error) {
glog.V(1).Infof("lookup %s from %s", name, location)
if strings.HasPrefix(name, ".") {
m, ok := r.msgs[name]
if !ok {
return nil, fmt.Errorf("no message found: %s", name)
}
return m, nil
}
if !strings.HasPrefix(location, ".") {
location = fmt.Sprintf(".%s", location)
}
components := strings.Split(location, ".")
for len(components) > 0 {
fqmn := strings.Join(append(components, name), ".")
if m, ok := r.msgs[fqmn]; ok {
return m, nil
}
components = components[:len(components)-1]
}
return nil, fmt.Errorf("no message found: %s", name)
}
// LookupEnum looks up a enum type by "name".
// It tries to resolve "name" from "location" if "name" is a relative enum name.
func (r *Registry) LookupEnum(location, name string) (*Enum, error) {
glog.V(1).Infof("lookup enum %s from %s", name, location)
if strings.HasPrefix(name, ".") {
e, ok := r.enums[name]
if !ok {
return nil, fmt.Errorf("no enum found: %s", name)
}
return e, nil
}
if !strings.HasPrefix(location, ".") {
location = fmt.Sprintf(".%s", location)
}
components := strings.Split(location, ".")
for len(components) > 0 {
fqen := strings.Join(append(components, name), ".")
if e, ok := r.enums[fqen]; ok {
return e, nil
}
components = components[:len(components)-1]
}
return nil, fmt.Errorf("no enum found: %s", name)
}
// LookupFile looks up a file by name.
func (r *Registry) LookupFile(name string) (*File, error) {
f, ok := r.files[name]
if !ok {
return nil, fmt.Errorf("no such file given: %s", name)
}
return f, nil
}
// AddPkgMap adds a mapping from a .proto file to proto package name.
func (r *Registry) AddPkgMap(file, protoPkg string) {
r.pkgMap[file] = protoPkg
}
// SetPrefix registers the prefix to be added to go package paths generated from proto package names.
func (r *Registry) SetPrefix(prefix string) {
r.prefix = prefix
}
// SetImportPath registers the importPath which is used as the package if no
// input files declare go_package. If it contains slashes, everything up to the
// rightmost slash is ignored.
func (r *Registry) SetImportPath(importPath string) {
r.importPath = importPath
}
// ReserveGoPackageAlias reserves the unique alias of go package.
// If succeeded, the alias will be never used for other packages in generated go files.
// If failed, the alias is already taken by another package, so you need to use another
// alias for the package in your go files.
func (r *Registry) ReserveGoPackageAlias(alias, pkgpath string) error {
if taken, ok := r.pkgAliases[alias]; ok {
if taken == pkgpath {
return nil
}
return fmt.Errorf("package name %s is already taken. Use another alias", alias)
}
r.pkgAliases[alias] = pkgpath
return nil
}
// goPackagePath returns the go package path which go files generated from "f" should have.
// It respects the mapping registered by AddPkgMap if exists. Or use go_package as import path
// if it includes a slash, Otherwide, it generates a path from the file name of "f".
func (r *Registry) goPackagePath(f *descriptor.FileDescriptorProto) string {
name := f.GetName()
if pkg, ok := r.pkgMap[name]; ok {
return path.Join(r.prefix, pkg)
}
gopkg := f.Options.GetGoPackage()
if len(gopkg) == 0 {
gopkg = r.importPath
}
idx := strings.LastIndex(gopkg, "/")
if idx >= 0 {
if sc := strings.LastIndex(gopkg, ";"); sc > 0 {
gopkg = gopkg[:sc+1-1]
}
return gopkg
}
return path.Join(r.prefix, path.Dir(name))
}
// GetAllFQMNs returns a list of all FQMNs
func (r *Registry) GetAllFQMNs() []string {
var keys []string
for k := range r.msgs {
keys = append(keys, k)
}
return keys
}
// GetAllFQENs returns a list of all FQENs
func (r *Registry) GetAllFQENs() []string {
var keys []string
for k := range r.enums {
keys = append(keys, k)
}
return keys
}
// SetAllowDeleteBody controls whether http delete methods may have a
// body or fail loading if encountered.
func (r *Registry) SetAllowDeleteBody(allow bool) {
r.allowDeleteBody = allow
}
// sanitizePackageName replaces unallowed character in package name
// with allowed character.
func sanitizePackageName(pkgName string) string {
pkgName = strings.Replace(pkgName, ".", "_", -1)
pkgName = strings.Replace(pkgName, "-", "_", -1)
return pkgName
}
// defaultGoPackageName returns the default go package name to be used for go files generated from "f".
// You might need to use an unique alias for the package when you import it. Use ReserveGoPackageAlias to get a unique alias.
func defaultGoPackageName(f *descriptor.FileDescriptorProto) string {
name := packageIdentityName(f)
return sanitizePackageName(name)
}
// packageIdentityName returns the identity of packages.
// protoc-gen-grpc-gateway rejects CodeGenerationRequests which contains more than one packages
// as protoc-gen-go does.
func packageIdentityName(f *descriptor.FileDescriptorProto) string {
if f.Options != nil && f.Options.GoPackage != nil {
gopkg := f.Options.GetGoPackage()
idx := strings.LastIndex(gopkg, "/")
if idx < 0 {
gopkg = gopkg[idx+1:]
}
gopkg = gopkg[idx+1:]
// package name is overrided with the string after the
// ';' character
sc := strings.IndexByte(gopkg, ';')
if sc < 0 {
return sanitizePackageName(gopkg)
}
return sanitizePackageName(gopkg[sc+1:])
}
if f.Package == nil {
base := filepath.Base(f.GetName())
ext := filepath.Ext(base)
return strings.TrimSuffix(base, ext)
}
return f.GetPackage()
}