如何自定义 protoc 插件

 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

syntax = "proto3";

package hello.service.v1;
option go_package = "api/hello/service/v1;v1";

// 下载 `github.com/googleapis/googleapis` 至`GOPATH`, 生成 http 代码需要。
import "google/api/annotations.proto";

service Hello {
    rpc Add(AddRequest) returns (AddResponse) {
        option (google.api.http) = {
            post: "/api/hello/service/v1/add"
            body: "*"
        };
    }

    rpc Get(GetRequest) returns (GetResponse) {
        option (google.api.http) = {
            get: "/api/hello/service/v1/get"
        };
    }
}

message AddRequest {
    uint32 id = 1;
    string name = 2;
}

message AddResponse {
    uint32 id = 1;
    string name = 2;
}

message GetRequest {
    uint32 id = 1;
}

message GetResponse {
    uint32 id = 1;
    string name = 2;
    float score = 3;
    bytes bs = 4;
    map<string, string> m = 5;
}

1
2
3

# --go-http 为我自己的插件
# 其中参数是 key=v,key2=v2 方式传，最后冒号后面写输出目录
protoc -I$GOPATH/src/github.com/googleapis/googleapis --proto_path=$GOPATH/src:. --go_out=. --go-http_out=router=gin:. --micro_out=. test.proto

 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

// Code generated by protoc-gen-go-http. DO NOT EDIT.
// versions:
// protoc-gen-go-http v0.0.9

package v1

import (
    context "context"
    gin "github.com/gin-gonic/gin"
)

// This is a compile-time assertion to ensure that this generated file
// is compatible with the galaxy package it is being compiled against.
var _ context.Context
const _ = gin.Version

type HelloHTTPHandler interface {
    Add(context.Context, *AddRequest, *AddResponse) error
    Get(context.Context, *GetRequest, *GetResponse) error
}

// RegisterHelloHTTPHandler define http router handle by gin.
func RegisterHelloHTTPHandler(g *gin.RouterGroup, srv HelloHTTPHandler) {
    g.POST("/api/hello/service/v1/add", _Hello_Add0_HTTP_Handler(srv))
    g.GET("/api/hello/service/v1/get", _Hello_Get0_HTTP_Handler(srv))
}

func _Hello_Add0_HTTP_Handler(srv HelloHTTPHandler) func(c *gin.Context) {
    return func(c *gin.Context) {
        var (
            in  AddRequest
            out AddResponse
        )

        if err := c.ShouldBind(&in); err != nil {
            c.AbortWithStatusJSON(400, gin.H{"err": err.Error()})
            return
        }

        err := srv.Add(context.Background(), &in, &out)
        if err != nil {
            c.AbortWithStatusJSON(500, gin.H{"err": err.Error()})
            return
        }

        c.JSON(200, &out)
    }
}

func _Hello_Get0_HTTP_Handler(srv HelloHTTPHandler) func(c *gin.Context) {
    return func(c *gin.Context) {
        var (
            in  GetRequest
            out GetResponse
        )

        if err := c.ShouldBind(&in); err != nil {
            c.AbortWithStatusJSON(400, gin.H{"err": err.Error()})
            return
        }

        err := srv.Get(context.Background(), &in, &out)
        if err != nil {
            c.AbortWithStatusJSON(500, gin.H{"err": err.Error()})
            return
        }

        c.JSON(200, &out)
    }
}

1
2
3
4
5
6

type HelloHandler interface {
    Add(context.Context, *AddRequest, *AddResponse) error
    Get(context.Context, *GetRequest, *GetResponse) error
}

func RegisterHelloHandler(s server.Server, hdlr HelloHandler, opts ...server.HandlerOption) error {}

1
2
3
4
5
6
7
8

// 它实现了 HelloHandler
type implHello struct{}

RegisterHelloHandler(micro.Server, &implHello)
g := gin.New()
// implHello 实现HelloHandler 那就是实现了HelloHTTPHandler
RegisterHelloHTTPHandler(g.Group("/"), &implHello)

 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

package main

import (
    "flag"

    "google.golang.org/protobuf/compiler/protogen"
    "google.golang.org/protobuf/types/pluginpb"
)

// protoc-gen-go-http 工具版本
// 与 GalaxyMicroVersion 保持一致
const version = "v0.0.12"

func main() {
    // 1. 传参定义
    // 即 插件是支持自定义参数的，这样我们可以更加灵活，针对不同的场景生成不同的代码
    var flags flag.FlagSet
    // 是否忽略没有指定 google.api 的方法
    omitempty := flags.Bool("omitempty", true, "omit if google.api is empty")
    // 我这里同时支持了 gin 和 iris 可以通过参数指定生成
    routerEngine := flags.String("router", "gin", "http router engine, choose between gin and iris")
    // 是否生校验代码块
    // 发现了一个很有用的插件 github.com/envoyproxy/protoc-gen-validate
    // 可以在 pb 的 message 中设置参数规则，然后会生成一个 validate.go 的文件 针对每个 message 生成一个 Validate() 方法
    // 我在每个 handler 处理业务前做了一次参数校验判断，通过这个 flag 控制是否生成这段校验代码
    genValidateCode := flags.Bool("validate", false, "add validate request params in handler")
    // 生成代码时参数 这么传：--go-http_out=router=iris,validate=true:.

    gp := &GenParam{
        Omitempty:       omitempty,
        RouterEngine:    routerEngine,
        GenValidateCode: genValidateCode,
    }
    // 这里就是入口，指定 option 后执行 Run 方法 ，我们的主逻辑就是在 Run 方法
    protogen.Options{
        ParamFunc: flags.Set,
    }.Run(func(gen *protogen.Plugin) error {
        gen.SupportedFeatures = uint64(pluginpb.CodeGeneratorResponse_FEATURE_PROTO3_OPTIONAL)
        for _, f := range gen.Files {
            if !f.Generate {
                continue
            }
            // 这里是我们的生成代码方法
            generateFile(gen, f, gp)
        }
        return nil
    })
}

type GenParam struct {
    Omitempty       *bool
    RouterEngine    *string
    GenValidateCode *bool
}

  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

import (
    "fmt"
    "strings"

    "google.golang.org/genproto/googleapis/api/annotations"
    "google.golang.org/protobuf/compiler/protogen"
    "google.golang.org/protobuf/proto"
    "google.golang.org/protobuf/types/descriptorpb"
)

const (
    contextPackage = protogen.GoImportPath("context")
    ginPackage     = protogen.GoImportPath("github.com/gin-gonic/gin")
    irisPackage    = protogen.GoImportPath("github.com/kataras/iris/v12")
)

var methodSets = make(map[string]int)

// generateFile generates a _http.pb.go file containing gin/iris handler.
func generateFile(gen *protogen.Plugin, file *protogen.File, gp *GenParam) *protogen.GeneratedFile {
    if len(file.Services) == 0 || (*gp.Omitempty && !hasHTTPRule(file.Services)) {
        return nil
    }
    // 这里我们可以自定义文件名
    filename := file.GeneratedFilenamePrefix + ".pb.http.go"
    g := gen.NewGeneratedFile(filename, file.GoImportPath)
    // 写入一些警告之类的 告诉用户不要修改
    g.P("// Code generated by protoc-gen-go-http. DO NOT EDIT.")
    g.P("// versions:")
    g.P(fmt.Sprintf("// protoc-gen-go-http %s", version))
    g.P()
    g.P("package ", file.GoPackageName)
    g.P()
    generateFileContent(gen, file, g, gp)
    return g
}

// generateFileContent generates the _http.pb.go file content, excluding the package statement.
func generateFileContent(gen *protogen.Plugin, file *protogen.File, g *protogen.GeneratedFile, gp *GenParam) {
    if len(file.Services) == 0 {
        return
    }
    // import
    // 这里有个插曲：其实 import 相关的代码我们这么不需要特殊指定，protogen 包会帮我们处理，
    // 但是import 的 path 前的别名默认取 path 最后一个 `/` 之后的字符，
    // 比如：github.com/kataras/iris/v12 被处理成 v12 "github.com/kataras/iris/v12"
    // 这个我不太愿意接受 所以自己写入 import
    g.P("// This imports are custom by galaxy micro framework.")
    g.P("import (")
    switch *gp.RouterEngine {
    case "gin":
        g.P("gin", " ", ginPackage)
    case "iris":
        g.P("iris", " ", irisPackage)
    }
    g.P(")")
    
    // 注： 我们难免有一些 _ "my/package" 这种需求，这其实不用自己写 直接调 g.Import("my/package") 就可以

    // 这里定义一堆变量是为了程序编译的时候确保这些包是正确的，如果包不存在或者这些定义的包变量不存在都会编译失败
    g.P("// This is a compile-time assertion to ensure that this generated file")
    g.P("// is compatible with the galaxy package it is being compiled against.")
    // 只要调用这个 Ident 方法 就会自动写入到 import 中 ，所以如果对 import 的包名没有特殊要求，那就直接使用 Ident
    g.P("var _ ", contextPackage.Ident("Context"))

    // 像我自己自定义 import 的包就不要使用 Ident 方法，否则生成的代码文件里有两个同一个包的引入导致语法错误
    switch *gp.RouterEngine {
    case "gin":
        g.P("const _ = ", "gin.", "Version")
    case "iris":
        g.P("const _ = ", "iris.", "Version")
    }
    g.P()

    // 到这里我们就把包名 import 和变量写入成功了，剩下的就是针对 rpc service 生成对应的 handler
    for _, service := range file.Services {
        genService(gen, file, g, service, gp)
    }
}


// rpc service 信息
type serviceDesc struct {
    ServiceType string // Greeter
    ServiceName string // helloworld.Greeter
    Metadata    string // api/helloworld/helloworld.proto
    GenValidate bool
    Methods     []*methodDesc
    MethodSets  map[string]*methodDesc
}

// rpc 方法信息
type methodDesc struct {
    // method
    Name    string
    Num     int
    Request string
    Reply   string
    // http_rule
    Path            string
    Method          string
    CamelCaseMethod string
    HasVars         bool
    HasBody         bool
    Body            string
    ResponseBody    string
}

// 生成 service 相关代码
func genService(gen *protogen.Plugin, file *protogen.File, g *protogen.GeneratedFile, service *protogen.Service, gp *GenParam) {
    if service.Desc.Options().(*descriptorpb.ServiceOptions).GetDeprecated() {
        g.P("//")
        g.P(deprecationComment)
    }

    // HTTP Server.
    // 服务的主要变量，比如服务名 服务类型等
    sd := &serviceDesc{
        ServiceType: service.GoName,
        ServiceName: string(service.Desc.FullName()),
        Metadata:    file.Desc.Path(),
        GenValidate: *gp.GenValidateCode,
    }
    // 开始遍历服务的方法
    for _, method := range service.Methods {
        // 不处理
        if method.Desc.IsStreamingClient() || method.Desc.IsStreamingServer() {
            continue
        }
        // annotations 这个就是我们在 rpc 方法里 option 里定义的 http 路由
        rule, ok := proto.GetExtension(method.Desc.Options(), annotations.E_Http).(*annotations.HttpRule)
        if rule != nil && ok {
            for _, bind := range rule.AdditionalBindings {
                // 拿到 option里定义的路由， http method等信息
                sd.Methods = append(sd.Methods, buildHTTPRule(g, method, bind))
            }
            sd.Methods = append(sd.Methods, buildHTTPRule(g, method, rule))
        } else if !*gp.Omitempty {
            path := fmt.Sprintf("/%s/%s", service.Desc.FullName(), method.Desc.Name())
            sd.Methods = append(sd.Methods, buildMethodDesc(g, method, "POST", path))
        }
    }

    // 拿到了 n 个 rpc 方法，开始生成了
    if len(sd.Methods) != 0 {
        // 渲染
        g.P(sd.execute(*gp.RouterEngine))
    }
}

// 检查是否有 http 规则 即 
// option (google.api.http) = {
//      get: "/user/query"
//    };
func hasHTTPRule(services []*protogen.Service) bool {
    for _, service := range services {
        for _, method := range service.Methods {
            if method.Desc.IsStreamingClient() || method.Desc.IsStreamingServer() {
                continue
            }
            rule, ok := proto.GetExtension(method.Desc.Options(), annotations.E_Http).(*annotations.HttpRule)
            if rule != nil && ok {
                return true
            }
        }
    }
    return false
}

// 解析 http 规则，读取内容
func buildHTTPRule(g *protogen.GeneratedFile, m *protogen.Method, rule *annotations.HttpRule) *methodDesc {
    var (
        path         string
        method       string
        body         string
        responseBody string
    )
    // 读取 路由和方法
    switch pattern := rule.Pattern.(type) {
    case *annotations.HttpRule_Get:
        path = pattern.Get
        method = "GET"
    case *annotations.HttpRule_Put:
        path = pattern.Put
        method = "PUT"
    case *annotations.HttpRule_Post:
        path = pattern.Post
        method = "POST"
    case *annotations.HttpRule_Delete:
        path = pattern.Delete
        method = "DELETE"
    case *annotations.HttpRule_Patch:
        path = pattern.Patch
        method = "PATCH"
    case *annotations.HttpRule_Custom:
        path = pattern.Custom.Path
        method = pattern.Custom.Kind
    }
    body = rule.Body
    responseBody = rule.ResponseBody
    md := buildMethodDesc(g, m, method, path)
    if method == "GET" {
        md.HasBody = false
    } else if body == "*" {
        md.HasBody = true
        md.Body = ""
    } else if body != "" {
        md.HasBody = true
        md.Body = "." + camelCaseVars(body)
    } else {
        md.HasBody = false
    }
    if responseBody == "*" {
        md.ResponseBody = ""
    } else if responseBody != "" {
        md.ResponseBody = "." + camelCaseVars(responseBody)
    }
    return md
}

// 构建 每个方法的基础信息
// 到这里我们拿到了 我们需要生成一个 handler 的所有信息
// 名称，输入，输出，方法类型，路由
func buildMethodDesc(g *protogen.GeneratedFile, m *protogen.Method, method, path string) *methodDesc {
    defer func() { methodSets[m.GoName]++ }()
    return &methodDesc{
        Name:            m.GoName,
        Num:             methodSets[m.GoName],
        Request:         g.QualifiedGoIdent(m.Input.GoIdent), // rpc 方法中的 request
        Reply:           g.QualifiedGoIdent(m.Output.GoIdent), // rpc 方法中的 response
        Path:            path,
        Method:          method,
        CamelCaseMethod: camelCase(strings.ToLower(method)),
        HasVars:         len(buildPathVars(m, path)) > 0,
    }
}

// 处理 路由中 /api/user/{name} 这种情况
func buildPathVars(method *protogen.Method, path string) (res []string) {
    for _, v := range strings.Split(path, "/") {
        if strings.HasPrefix(v, "{") && strings.HasSuffix(v, "}") {
            name := strings.TrimRight(strings.TrimLeft(v, "{"), "}")
            res = append(res, name)
        }
    }
    return
}

 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

// execute 方法实现也其实不复杂，总起来就是 go 的 temple 包的使用
// 提前写好模板文件，然后拿到所有需要的变量，进行模板渲染，写入文件
func (s *serviceDesc) execute(routerEngine string) string {
    var (
        name = routerEngine
        tmp  string
    )
    switch routerEngine {
    case "gin":
        tmp = ginTemplate
    case "iris":
        tmp = irisTemplate
    default:
        panic("unknown http engine")
    }
    s.MethodSets = make(map[string]*methodDesc)
    for _, m := range s.Methods {
        s.MethodSets[m.Name] = m
    }
    buf := new(bytes.Buffer)
    tmpl, err := template.New(name).Parse(strings.TrimSpace(tmp))
    if err != nil {
        panic(err)
    }
    if err = tmpl.Execute(buf, s); err != nil {
        panic(err)
    }

    return strings.Trim(buf.String(), "\r\n")
}

 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

var ginTemplate = `
{{$svrType := .ServiceType}}
{{$svrName := .ServiceName}}
{{$validate := .GenValidate}}
// 这里定义 handler interface
type {{.ServiceType}}HTTPHandler interface {
{{- range .MethodSets}}
    {{.Name}}(context.Context, *{{.Request}}, *{{.Reply}}) error
{{- end}}
}

// Register{{.ServiceType}}HTTPHandler define http router handle by gin. 
// 注册路由 handler
func Register{{.ServiceType}}HTTPHandler(g *gin.RouterGroup, srv {{.ServiceType}}HTTPHandler) {
    {{- range .Methods}}
    g.{{.Method}}("{{.Path}}", _{{$svrType}}_{{.Name}}{{.Num}}_HTTP_Handler(srv))
    {{- end}}
}

// 定义 handler
// 遍历之前解析到所有 rpc 方法信息
{{range .Methods}}
func _{{$svrType}}_{{.Name}}{{.Num}}_HTTP_Handler(srv {{$svrType}}HTTPHandler) func(c *gin.Context) {
    return func(c *gin.Context) {
        var (
            in  = new({{.Request}})
            out = new({{.Reply}})
            ctx = middleware.GetContextFromGinCtx(c)
        )

        if err := c.ShouldBind(in{{.Body}}); err != nil {
            c.AbortWithStatusJSON(400, gin.H{"err": err.Error()})
            return
        }
        
        // 这里就是最开始提到的判断是否启用 validate
        // 其中这个 api.Validator 接口只有一个方法 Validate() error
        // 所以需要在一个统一的地方定义好引入使用，建议不要在生成的时候写入，因为这个是通用的 interface{}
        {{if $validate -}} 
        // check param
        if v, ok := interface{}(in).(api.Validator);ok {
            if err := v.Validate();err != nil {
                c.AbortWithStatusJSON(400, gin.H{"err": err.Error()})
                return
            }
        }
        {{end -}}

        // 执行方法
        err := srv.{{.Name}}(ctx, in, out)
        if err != nil {
            c.AbortWithStatusJSON(500, gin.H{"err": err.Error()})
            return
        }
        
        c.JSON(200, out)
    }
}
{{end}}
`