白菜君の技术库

【Golang】- protobuf插件扩展开发

前言

最近，项目需要用到protobuf来定义消息，但是我们需要一个更灵活的代码片段，如何通过proto文件来创建自定义的代码呢？
可以通过proto的plugin对方式来自己是一个proto-gen。

在网上看了一些教程，发现有一些教程已经过时了，而且过于片面，没有把整套思想很好的说明。并且也有一些功能点并没有完全实现。
这里总结一下相关的内容，并且说一下最近实现的一个插件。

对于已经了解大概proto的人来说，相对简单，但是如果是自定义option呢？你又了解吗？

需求

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
// test.proto

syntax = "proto3";
package pb;
option go_package = "/;pb";

import "unknow.proto";

service ApiIMService {
    rpc RegisterDevice (ApiIMRegisterDeviceReq) returns (ApiIMRegisterDeviceResp) {
        option (unknow.api.http) = {
            method: "post"
            url: "/v1/im/register_device"
        };
    }
}

// 设备类型
enum ApiIMDeviceType {
    API_IM_UNKNOWN = 0;
    API_IM_Android = 1;
    API_IM_IOS = 2;
    API_IM_Window = 3;
    API_IM_MacOS = 4;
    API_IM_Web = 5;
}

message ApiIMRegisterDeviceReq {
    ApiIMDeviceType type = 1; // 设备类型
}

message ApiIMRegisterDeviceResp {
    int64 device_id = 1; // 设备id
}

我们看到这一个proto文件，常规的我就不说了。主要是看到import "unknow.proto", option (unknow.api.http)

可以看到，我这里引入一个unknow.proto的文件。

我希望最终生成一个api.unknow.go的文件。里面的内容期待如下：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
// Copyright (c) 2021, whiteCcinn Inc.
// Code generated by protoc-gen-unknow. DO NOT EDIT.
// source: test.proto

package pb

type Api struct {
	Name   string
	Method string
	Url    string
}

func newApi(name, method, url string) *Api {
	return &Api{
		name, method, url,
	}
}

var (
	RegisterDeviceApi = newApi("RegisterDevice", "post", "/v1/im/register_device")
)

生成的文件，我可以在项目通过pb.RegisterDeviceApi拿到在proto定义好的API内容。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
// unknow.proto

syntax = "proto3";

package unknow.api;

option go_package = "/;pb";

import "google/protobuf/descriptor.proto";

extend google.protobuf.MethodOptions {
    // See `HttpRule`.
    HttpRule http = 72295728;
}

message HttpRule {
    string url = 1;
    string method = 2;
}

如果有了解过google/api/annotations.proto 和 google/api/http.proto的人应该不会陌生，当你需要用到grpc-gateway或者proto-gen-swaager的时候，都会有介绍到option(goggle.api.http)的用法。

这里我们看到unknow.proto的结构体，这就是一个简化版本。用于实现自定义的option用的。

unknow.proto

对于基础的语法规则来说，我们来看剖析一下unknow.proto，常规的就不说了。

1
2
3
4
5
6
7
8
9
10
11
package unknow.api;

extend google.protobuf.MethodOptions {
    // See `HttpRule`.
    HttpRule http = 72295728;
}

message HttpRule {
    string url = 1;
    string method = 2;
}

对于extend的用法，我这里列一下官方的链接。

我们看到这里，extend google.protobuf.MethodOptions。

代表着，我这里自定义个作用于Method的option。所以在真正的proto文件中，我就可以使用unknow.api.option的标签。也就是option (unknow.api.http)。

接着我们看到，这个option我们定义了一个元素，类型是Message HttpRule，命名为http，并且给它定义一个唯一的Number。

接着我们看到HttpRule，内部存在2个元素，一个是string url 和 string method，这意味着我们可以使用独立行写法：option (unknow.api.http).url = "/v1/im/register_device"，然后再下一行使用 option (unknow.api.http).method = "post"，一个完整的例如如下：

刚才说到，这是一个method的option，也就是我们这里定义的rpc下的 option。

1
2
3
4
5
6
service ApiIMService {
    rpc RegisterDevice (ApiIMRegisterDeviceReq) returns (ApiIMRegisterDeviceResp) {
        option (unknow.api.http).method = "post"
        option (unknow.api.http).url = "/v1/im/register_device"
    }
}

除了这个写法，我更推荐如下更简洁的写法，用map的写法：

1
2
3
4
5
6
7
8
service ApiIMService {
    rpc RegisterDevice (ApiIMRegisterDeviceReq) returns (ApiIMRegisterDeviceResp) {
        option (unknow.api.http) = {
            method: "post"
            url: "/v1/im/register_device"
        };
    }
}

这样子，我们就看懂了test.proto的内容了对吧。连贯起来，那么我们的这个unknow.proto就是用于实现option(unknow.api.http)。而option(unknow.api.http)的使用则在test.proto进行采用。

好了，有了定义的文件，那么接下来就是我们的重点了，如何编写一个实现自定义代码的protobuf插件扩展

扩展实现

protobuf解析流程图

protobuf解析旧版的流程图，便于我们理解。新版的后续我抽空再画画

不科学的例子

第一个例子，以golang旧版proto-gen-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
package main

import (
	"io/ioutil"
	"os"

	"github.com/golang/protobuf/proto"
	"github.com/golang/protobuf/protoc-gen-go/generator"
)

func main() {
	// Begin by allocating a generator. The request and response structures are stored there
	// so we can do error handling easily - the response structure contains the field to
	// report failure.
	g := generator.New()

	data, err := ioutil.ReadAll(os.Stdin)
	if err != nil {
		g.Error(err, "reading input")
	}

	if err := proto.Unmarshal(data, g.Request); err != nil {
		g.Error(err, "parsing input proto")
	}

	if len(g.Request.FileToGenerate) == 0 {
		g.Fail("no files to generate")
	}

	g.CommandLineParameters(g.Request.GetParameter())

	// Create a wrapped version of the Descriptors and EnumDescriptors that
	// point to the file that defines them.
	g.WrapTypes()

	g.SetPackageNames()
	g.BuildTypeNameMap()

	g.GenerateAllFiles()

	// Send back the results.
	data, err = proto.Marshal(g.Response)
	if err != nil {
		g.Error(err, "failed to marshal output proto")
	}
	_, err = os.Stdout.Write(data)
	if err != nil {
		g.Error(err, "failed to write output proto")
	}
}

我发现现在网上很多教程都会以旧版的为主，并且都是沿用参考了proto-gen-go的这个旧版的写法，所以导致，我们在学习写的时候，会出现一些问题。并且其实你用了旧版的这个写法，当你在用protoc去编译的情况下，protoc也会友好的提示你，该API已经被废弃，将在未来的版本下移除，请使用新的写法。虽然你还是能编译通过，但是你不敢保证未来哪一个版本就直接不向后兼容了。

第二个例子。

1
2
3
4
5
6
7
8
9
10
input, _ := ioutil.ReadAll(os.Stdin)
var req pluginpb.CodeGeneratorRequest
proto.Unmarshal(input, &req)

// 使用默认选项初始化我们的插件
opts := protogen.Options{}
plugin, err := opts.New(&req)
if err != nil {
   panic(err)
...

这种方式就是从头到尾都自己写，把整个pipeline的流程都自己处理。但是大可不必，因为其实有很多流程化的东西，在新版的genpro下已经封装成了一个Run传递回调函数即可。

正确的例子

如果真的要参考的话，推荐参考最新版本的proto-gen-go

首先，我们需要知道一点，我们在采用protoc对proto文件进行编译的时候，经常是执行类似如下代码：

1
2
protoc -I.:${GOGO_PROTOBUF} \
--gofast_out=plugins=grpc:./go-pb


由于这里我用的是gogoprotobuf，所以你会看到我这里的是--gofast_out=plugins=grpc:./go-pb，如果你是protobuf的官方的proto-gen的话，那么你这里应该是--go_out=plugins=grpc:./go-pb

这里我们需要注意的是什么呢，就是插件二进制文件名，这是有一定规则的，这是我之前在自定义git凭据文章中说明的一样，这二进制文件需要在你的环境变量中，否则你就需要通过-I来指定文件路径。然后命名规则为proto-gen-xxx，而这个xxx就是你的自定义的名字。在本例子中，我的扩展名字就是proto-gen-unknow。

因此，最终如果我要执行的话，那么就是执行如下命令：

1
2
protoc -I.:${GOGO_PROTOBUF} \
--unknow_out=./go-pb

其实如果有接触过thrift 或者 Rust的元编程 甚至是 Python的 lark-parser自定义抽象语法树，或者其他经由AST抽象语法树写代码的同学应该都知道，这其实抽象出来就是一个AST的解析处理而已。所以我首先需要了解他的部署。

一个简陋的AST定义如下（哈哈，略显简陋，但是了解AST是啥的应该多少能看懂一些）：

1
2
3
4
5
6
7
8
9
10
FileDescriptor -> ServiceDescriptor
					-> ServiceOptionDescriptor
					-> MethodDescriptor
						-> MethodOptionDescriptor
					    -> [MessageDescriptor]
               -> MessageDescriptor
					-> MessageOptionDescriptor
					-> FieldDescriptor
						-> FieldOptionDescriptor
			   -> FieldOptionDescriptor

知道怎么执行了，和AST, 我们就来看看怎么编写代码。

1
2
3
4
5
6
7
8
9
10
11
12
13
➜  protoc-gen-unknow git:(main) tree
.
├── README.md
├── go.mod
├── go.sum
├── internal
│   └── unknow.go
├── main.go
├── out
│   └── unknow.pb.go
└── proto
    ├── descriptor.proto
    └── unknow.proto

可以看到，这是我的一个代码结构目录

1
2
3
4
5
6
7
8
9
10
11
package main

import (
	"internal"
	"google.golang.org/protobuf/compiler/protogen"
)

func main() {
	u := internal.Unknow{}
	protogen.Options{}.Run(u.Generate)
}

这里，我们借助protobuf 编译包下的 protogen 工具来解析 proto 文件。我们接下来看一下 internal 包下具体的写法。

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
package internal

import (
	pb "out"
	"google.golang.org/protobuf/compiler/protogen"
	"google.golang.org/protobuf/proto"
	"google.golang.org/protobuf/types/descriptorpb"
)

type Unknow struct {
}

func (u Unknow) Generate(plugin *protogen.Plugin) error {
	if len(plugin.Files) < 1 {
		return nil
	}
	// 指定生成文件的文件名
	filename := "/api.unknow.go"

	// 创建一个文件生成器对象
	g := plugin.NewGeneratedFile(filename, plugin.Files[len(plugin.Files)-1].GoImportPath)

	// 调用g.P就是往文件开始写入自己期待的代码
	g.P(`// Copyright (c) 2021, whiteCcinn Inc.`)
	g.P("// Code generated by protoc-gen-unknow. DO NOT EDIT.")
	g.P("// source: all MethodOptions(unknow.api.http) in proto file")
	g.P()
	g.P("package ", plugin.Files[len(plugin.Files)-1].GoPackageName)

	g.P(`
type Api struct {
	Name   string
	Method string
	Url    string
}

func newApi(name, method, url string) *Api {
	return &Api{
		name, method, url,
	}
}
`)

	g.P("var (")

	// 通过plugin.Fiels，我们可以拿到所有的输入的proto文件
	// 如果我们需要对这个文件生成代码的话，那么就进入到generateFile()逻辑
	// 并且把g和f一起传递过去
	for _, f := range plugin.Files {
		if f.Generate {
			if _, err := u.generateFile(g, f); err != nil {
				return err
			}
		}
	}

	g.P(")")

	return nil
}

func (u Unknow) generateFile(g *protogen.GeneratedFile, file *protogen.File) (*protogen.GeneratedFile, error) {
	// 这一段代码仅仅只是为了忽略包含proto文件中包含了streamClient和streamServer的代码
	isGenerated := false
	for _, srv := range file.Services {
		for _, method := range srv.Methods {
			if method.Desc.IsStreamingClient() || method.Desc.IsStreamingServer() {
				continue
			}
			isGenerated = true
		}
	}

	if !isGenerated {
		return nil, nil
	}

	// file 的下一层级就是 services 层级
	for _, srv := range file.Services {
		if err := u.genService(g, srv); err != nil {
			return nil, err
		}
	}

	return g, nil
}

func (u Unknow) genService(g *protogen.GeneratedFile, srv *protogen.Service) error {
	// service 内部有很多 rpc 关键字的方法
	for _, method := range srv.Methods {
		if method.Desc.IsStreamingClient() || method.Desc.IsStreamingServer() {
			continue
		}

		// 由于我们自定义的是就是MethodOptions，所以就来到了这里来进行判断
		if err := u.genMethodHTTPRule(g, method); err != nil {
			return err
		}
	}

	return nil
}

func (u Unknow) genMethodHTTPRule(g *protogen.GeneratedFile, method *protogen.Method) error {
	// 因为我们通过method.Desc.Options() 拿到的数据类型是`interface{}` 类型
	// 所以这里我们需要对Options，明确指定转换为 *descriptorpb.MethodOptions 类型
	// 这样子就能拿到我们的MethodOption对象
	options, ok := method.Desc.Options().(*descriptorpb.MethodOptions)
	if !ok {
		return nil
	}

    // PS：重点
	// 这里我们看到我们借助了一个非protogen下的包的内容
	// 原因就是，protobuf编译器会把自定义的Option全部指定为Extension，由于并非内置的属性和值
	// protobuf官方是没办法拿到和你对应的可读的内容的，只能通过拿到经过序列化之后的数据。
	// 因此，我们这里通过 proto.GetExtension的方法，把刚才unknow.proto单独编译好的 unknow.pb.proto 文件下的 pb. E_HTTP 加载进来，指定了我需要在自定义扩展的MethodOptions中，拿到该Http下里面的value
	// 也因此，我们可以再经过一次类型转换，就可以拿到了具体的httpRule
	httpRule, ok := proto.GetExtension(options, pb.E_Http).(*pb.HttpRule)
	if !ok {
		return nil
	}

	// 接下来，我们就可以通过GetXxx的方式，来获取我们设置在其Message内部filed
	m := httpRule.GetMethod()
	url := httpRule.GetUrl()

	if len(m) == 0 && len(url) == 0 {
		return nil
	}

	g.P("     ", method.GoName, "Api = ", "newApi(\"", method.GoName, "\", \"", m, "\", \"", url, "\")")

	return nil
}

请跟着说明中注释一步步查看详解

查看这段代码逻辑，也是非常简单，因为没有特别复杂的逻辑，尽量不要跳过，因为里面涉及到如何读取自定义的Option的问题。代码大致定位在 proto.GetExtension方法附近。

因为这里的demo生成的代码比较简单。最终，我们生成的代码就是：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
// Copyright (c) 2021, whiteCcinn Inc.
// Code generated by protoc-gen-unknow. DO NOT EDIT.
// source: source: all MethodOptions(unknow.api.http) in proto file

package pb

type Api struct {
	Name   string
	Method string
	Url    string
}

func newApi(name, method, url string) *Api {
	return &Api{
		name, method, url,
	}
}

var (
	RegisterDeviceApi = newApi("RegisterDevice", "post", "/v1/im/register_device")
)

放一下完整的测试命令:

1
go install . && protoc --proto_path proto/ -I=. test.proto  test2.proto --unknow_out=./out --go_out=./out

最后生成的文件目录结构如下：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
➜  protoc-gen-unknow git:(main) tree
.
├── README.md
├── go.mod
├── go.sum
├── internal
│   └── unknow.go
├── main.go
├── out
│   ├── api.unknow.go
│   ├── test.pb.go
│   ├── test2.pb.go
│   └── unknow.pb.go
└── proto
    ├── descriptor.proto
    ├── test.proto
    ├── test2.proto
    └── unknow.protos

最后推荐几个扩展库写得不错的扩展插件:

Golang Golang