package main import ( "fmt" protogen2 "git.diulo.com/mogfee/kit/protogen" "google.golang.org/protobuf/compiler/protogen" "google.golang.org/protobuf/types/descriptorpb" "strings" ) func main() { u := &Kit{ imports: map[string]string{}, } protogen.Options{}.Run(u.Generate) } type Kit struct { imports map[string]string } func (u *Kit) addImports(imp string) { u.imports[imp] = imp } func (u *Kit) packageName(f *protogen.File) string { a := f.Desc.Options().(*descriptorpb.FileOptions) p := strings.Split(*a.GoPackage, ";") if len(p) == 2 { return p[1] } return fmt.Sprintf("%s", f.Desc.Name()) } func (u *Kit) Generate(plugin *protogen.Plugin) error { if len(plugin.Files) < 1 { return nil } u.addImports("context") u.addImports("git.diulo.com/mogfee/kit/transport/http") for _, f := range plugin.Files { if len(f.Services) == 0 { continue } fname := f.GeneratedFilenamePrefix + "_http.pb.go" t := plugin.NewGeneratedFile(fname, f.GoImportPath) t.P("package ", u.packageName(f)) t.P("import (") for _, v := range u.imports { t.P(`"` + v + `"`) } t.P(")") for _, s := range f.Services { t.P(`type `, s.GoName, `HTTPServer interface {`) for _, m := range s.Methods { t.P(m.GoName, `(context.Context, *`, m.Input.GoIdent.GoName, `) (*`, m.Output.GoIdent.GoName, `,error)`) } t.P(`}`) } for _, s := range f.Services { serverName := s.GoName t.P(`func Register`, serverName, `HTTPServer(s *http.Server,srv `, serverName, `Server) {`) t.P(`r:=s.Route("/")`) for _, m := range s.Methods { method, path := protogen2.GetProtoMethod(m) if method == "" { continue } t.P(`r.`, method, `("`, path, `",_`, s.GoName, `_`, m.GoName, `0_HTTP_Handler(srv))`) } t.P(`}`) } for _, s := range f.Services { for _, m := range s.Methods { method, _ := protogen2.GetProtoMethod(m) if method == "" { continue } u.genGet(f, s, t, m) } } } return nil } func (u *Kit) genGet(f *protogen.File, s *protogen.Service, t *protogen.GeneratedFile, m *protogen.Method) { method, path := protogen2.GetProtoMethod(m) if method == "" { return } t.P(`func _`, s.GoName, `_`, m.GoName, `0_HTTP_Handler(srv `, s.GoName, `HTTPServer) func(ctx http.Context) error { return func(ctx http.Context) error { var in `, m.Input.GoIdent.GoName) t.P("//", protogen2.GetAuthKey(m)) if method == protogen2.METHOD_GET { t.P(`if err := ctx.BindQuery(&in); err != nil { return err }`) } else if method == protogen2.METHOD_POST { t.P(`if err := ctx.Bind(&in); err != nil { return err }`) } if strings.LastIndexByte(path, '{') != -1 { t.P(`if err := ctx.BindVars(&in); err != nil { return err }`) } t.P(`http.SetOperation(ctx, "/`, f.Desc.Package(), `.`, s.Desc.Name(), `/`, m.Desc.Name(), `") h := ctx.Middleware(func(ctx context.Context, req interface{}) (interface{}, error) { return srv.`, m.GoName, `(ctx, req.(*`, m.Input.GoIdent.GoName, `)) }) out, err := h(ctx, &in) if err != nil { return err } reply := out.(*`, m.Output.GoIdent.GoName, `) return ctx.Result(200, reply) } }`) }