package gengateway import ( "errors" "fmt" "go/format" "path" "path/filepath" "strings" "github.com/golang/glog" "github.com/golang/protobuf/proto" plugin "github.com/golang/protobuf/protoc-gen-go/plugin" "github.com/grpc-ecosystem/grpc-gateway/protoc-gen-grpc-gateway/descriptor" gen "github.com/grpc-ecosystem/grpc-gateway/protoc-gen-grpc-gateway/generator" ) var ( errNoTargetService = errors.New("no target service defined in the file") ) type pathType int const ( pathTypeImport pathType = iota pathTypeSourceRelative ) type generator struct { reg *descriptor.Registry baseImports []descriptor.GoPackage useRequestContext bool registerFuncSuffix string pathType pathType allowPatchFeature bool } // New returns a new generator which generates grpc gateway files. func New(reg *descriptor.Registry, useRequestContext bool, registerFuncSuffix, pathTypeString string, allowPatchFeature bool) gen.Generator { var imports []descriptor.GoPackage for _, pkgpath := range []string{ "context", "io", "net/http", "github.com/grpc-ecosystem/grpc-gateway/runtime", "github.com/grpc-ecosystem/grpc-gateway/utilities", "github.com/golang/protobuf/proto", "google.golang.org/grpc", "google.golang.org/grpc/codes", "google.golang.org/grpc/grpclog", "google.golang.org/grpc/status", } { pkg := descriptor.GoPackage{ Path: pkgpath, Name: path.Base(pkgpath), } if err := reg.ReserveGoPackageAlias(pkg.Name, pkg.Path); err != nil { for i := 0; ; i++ { alias := fmt.Sprintf("%s_%d", pkg.Name, i) if err := reg.ReserveGoPackageAlias(alias, pkg.Path); err != nil { continue } pkg.Alias = alias break } } imports = append(imports, pkg) } var pathType pathType switch pathTypeString { case "", "import": // paths=import is default case "source_relative": pathType = pathTypeSourceRelative default: glog.Fatalf(`Unknown path type %q: want "import" or "source_relative".`, pathTypeString) } return &generator{ reg: reg, baseImports: imports, useRequestContext: useRequestContext, registerFuncSuffix: registerFuncSuffix, pathType: pathType, allowPatchFeature: allowPatchFeature, } } func (g *generator) Generate(targets []*descriptor.File) ([]*plugin.CodeGeneratorResponse_File, error) { var files []*plugin.CodeGeneratorResponse_File for _, file := range targets { glog.V(1).Infof("Processing %s", file.GetName()) code, err := g.generate(file) if err == errNoTargetService { glog.V(1).Infof("%s: %v", file.GetName(), err) continue } if err != nil { return nil, err } formatted, err := format.Source([]byte(code)) if err != nil { glog.Errorf("%v: %s", err, code) return nil, err } name := file.GetName() if g.pathType == pathTypeImport && file.GoPkg.Path != "" { name = fmt.Sprintf("%s/%s", file.GoPkg.Path, filepath.Base(name)) } ext := filepath.Ext(name) base := strings.TrimSuffix(name, ext) output := fmt.Sprintf("%s.pb.gw.go", base) files = append(files, &plugin.CodeGeneratorResponse_File{ Name: proto.String(output), Content: proto.String(string(formatted)), }) glog.V(1).Infof("Will emit %s", output) } return files, nil } func (g *generator) generate(file *descriptor.File) (string, error) { pkgSeen := make(map[string]bool) var imports []descriptor.GoPackage for _, pkg := range g.baseImports { pkgSeen[pkg.Path] = true imports = append(imports, pkg) } for _, svc := range file.Services { for _, m := range svc.Methods { imports = append(imports, g.addEnumPathParamImports(file, m, pkgSeen)...) pkg := m.RequestType.File.GoPkg if len(m.Bindings) == 0 || pkg == file.GoPkg || pkgSeen[pkg.Path] { continue } pkgSeen[pkg.Path] = true imports = append(imports, pkg) } } params := param{ File: file, Imports: imports, UseRequestContext: g.useRequestContext, RegisterFuncSuffix: g.registerFuncSuffix, AllowPatchFeature: g.allowPatchFeature, } return applyTemplate(params, g.reg) } // addEnumPathParamImports handles adding import of enum path parameter go packages func (g *generator) addEnumPathParamImports(file *descriptor.File, m *descriptor.Method, pkgSeen map[string]bool) []descriptor.GoPackage { var imports []descriptor.GoPackage for _, b := range m.Bindings { for _, p := range b.PathParams { e, err := g.reg.LookupEnum("", p.Target.GetTypeName()) if err != nil { continue } pkg := e.File.GoPkg if pkg == file.GoPkg || pkgSeen[pkg.Path] { continue } pkgSeen[pkg.Path] = true imports = append(imports, pkg) } } return imports }