带你写个自己的protoc生成工具

完整代码已上传至GitHub，在文章最下方获取喔~

• 你是否用过protobuf或gRPC？
• 你们公司项目的API有没有用到proto文件?

本文将带你一步一步写个类似protoc-gen-go-grpc的proto文件生成工具，从proto文件生成兼容Go标准库的HTTP框架代码。掌握它之后，你就能随心所欲的从proto文件生成gin、echo、net/http代码，甚至生成你自己的框架代码。

别担心，生成的内容不局限于Go语言，别的语言也没问题，甚至不是编程语言都可以！

你可以从proto文件生成任何它能描述的东西。😉

我们正在做什么?

在学习gRPC时，你执行了一段命令，就将proto文件变成了gRPC代码:

protoc -I api \
  --go_out=internal/genproto/$service \
  --go_opt=paths=source_relative \
  --go-grpc_out=internal/genproto/$service \
  --go-grpc_opt=paths=source_relative \
  $service.proto

这多亏了protoc-gen-go和protoc-gen-go-grpc这两个可执行文件。

我们在执行protoc -xxx_out=. -xxx_opt=.命令时，protoc会从你操作系统的$PATH目录下寻找protoc-gen-xxx这个可执行文件进行执行，并把后面的参数传给它。

你可以执行ls $GOPATH/bin | grep protoc命令来查看自己电脑上安装了哪些protoc生成工具:

本文会带你实现一个名叫protoc-gen-go-example的工具，在使用时，我们只需要执行以下命令，即可调用我们自己写的生成工具来生成代码:

protoc --go-example_out=. # 此处省略其他参数

好啦，看到这儿你也应该明白我们在做什么事情啦，我们直接进入正题吧！

站在巨人的肩膀上

如果你学过编译原理，你一定很清楚我们要做些什么(没学过的同学先别跑😢):

1.解析.proto文件，构建proto文件的AST(抽象语法树)

2.遍历AST，将其转换为想要生成的内容。

天哪，这要是从零实现，需要多大的工程量啊！更别提一些没学过编译原理的同学们了。我要是从零开始教，那能写一本书了...

幸运的是，我们有一些可以利用的东西！不需要我们自己去实现proto文件的Parser啦！

protocolbuffers/protobuf-go这个库(也就是protoc-gen-go)已经帮我们实现了工作量最大的parser部分。

这下我们可以继续一起愉快的玩耍了!

创建项目

我创建的项目叫protoc-gen-go-example，这就是我们最终生成的二进制文件名称。

我们在执行go install命令时，默认会以main.go的上一级目录名作为可执行文件的名称。

假设我们的main.go文件放在根目录下，我们执行go install github.com/bootun/protoc-gen-test-name后，就会在你的$GOPATH/bin目录下安装一个名为protoc-gen-test-name可执行文件。

你可能会觉得：“那这样我的项目名岂不是很丑😢”。如果你不想让项目名作为最终的文件名称，你可以参考protocolbuffers/protobuf-go的做法。 protobuf-go把main.go放在了项目的cmd/protoc-gen-go目录下，这样在执行go install时，生成的文件就不会与项目名一样了，但代价就是go install的路径也会变长:

go install google.golang.org/protobuf/cmd/protoc-gen-go

想了解更多可以去看看go的官方文档

执行以下命令来初始化项目:

mkdir protoc-gen-go-example
cd protoc-gen-go-example
touch main.go
go mod init github.com/bootun/protoc-gen-go-example

现在你的项目看起来像下面这样:

protoc-gen-go-example
├── main.go
└── go.mod

现在让我们来编辑main.go文件:

package main

import (
    "github.com/bootun/protoc-gen-go-example/parser"
    "google.golang.org/protobuf/compiler/protogen"
)

func main() {
    protogen.Options{}.Run(func(gen *protogen.Plugin) error {
        // 这个循环遍历所有要生成的proto文件
        for _, f := range gen.Files {
            if !f.Generate {
                // 如果该文件不需要生成，则跳过
                continue
            }
            // 如果需要生成，就把文件的相关信息传递给生成器
            if err := parser.GenerateFile(gen, f); err != nil {
                return err
            }
        }
        return nil
    })
}

还记得我之间说过的吗？我们在main.go中使用了protobuf-go中的组件，这样我们就不需要从零开始解析proto文件中的内容了。

protogen.Options{}.Run()的参数是一个回调函数，回调函数的gen参数里包括了所有已经解析好的信息。gen.Files表示所有proto文件的集合，我们需要遍历这些proto文件，来为它们生成代码。

我们把gen和file向下传递，以便下面的组件能够获得足够多的信息。现在parser.GenerateFile还在报错，我们来实现GenerateFile这个函数:

GenerateFile函数

GenerateFile函数还是比较清晰明了的:

func GenerateFile(gen *protogen.Plugin, file *protogen.File) error {
    // 如果这个proto文件里没写service
    // 我们就不需要为它生成代码
    if len(file.Services) == 0 {
        return nil
    }

    // 要生成的文件名称
    filename := file.GeneratedFilenamePrefix + ".example.pb.go"
    g := gen.NewGeneratedFile(filename, file.GoImportPath)

    return NewFile(g, file).Generate()
}

代码里有注释的部分我就不额外解释了，很容易理解。我们需要额外关注下NewGeneratedFile这个函数。

NewGeneratedFile使用给定的文件名和ImportPath创建一个新的生成文件实体，我们将它命名为g。那这ImportPath又是个啥东西呢？

看过我上一篇文章的小伙伴们应该比较清楚，ImportPath就是我们在proto文件中写的option go_package里的内容。protobuf-go帮我们做了许多处理，使得我们不需要过度关注像--xxx_opt=paths=source_relative等这种与代码生成逻辑无关的内容，感兴趣的话可以去看我的上篇文章——彻底搞清protobuf-go的文件生成位置。

有了g之后，我们只需要调用g.P("xxx")方法，即可在文件中写入对应的字符串。

看到这里你可能就明白了，我们只需要创建一套模板，将file参数中的信息套到这个模板上，然后传给g.P(模板字符串)就行了。没错，就是这么简单！

在NewGeneratedFile函数的最后，我们调用NewFile创建了一个文件结构，并执行了该结构体上的Generate方法，整个代码生成工作就完成了。让我们看看NewFile里做了什么工作吧。

NewFile函数

func NewFile(gen *protogen.GeneratedFile, protoFile *protogen.File) *File {
    f := &File{
        // 保存example.pb.go的文件实体
        // 以便后面操作
        gen: gen,
    }

    f.PackageName = string(protoFile.GoPackageName)

    for _, s := range protoFile.Services {
        f.ParseService(s)
    }

    return f
}

我在NewFile中创建了一个File结构体，这是我们自定义的一个结构体，用来表示一个proto文件的内容。

这个结构体不是必要的，甚至可能是多余的，它只是把protogen.File参数里的内容给转成了我们的内部表示，所有的信息protogen.File里都有，如果你想的话，你可以直接使用protogen.File+text/template来生成文件。这里我出于教学目的，希望你能更容易明白这个过程，同时也为了日后做些更骚的操作，就留下这个结构了。

proto文件的内部表示

刚刚我提到了File结构体，说它只是把protogen.File里的一部分信息复制出来，转为我们自己的内部结构体了，事实上除了File之外，还有几个同样表示proto信息的结构体，他们都是File结构的下属结构:

// 一个File表示一个proto文件的信息
type File struct {
    // File内部同时保存了example.pb.go的文件句柄
    // 方便我们直接调用gen.P向pb文件写入内容
    gen *protogen.GeneratedFile

    // 内嵌了一个FileDescription结构
    // 更多信息可以继续往下看
    FileDescription
}

// FileDescription 描述了一个解析过后的proto文件的信息
// 为我们后边的代码生成做准备
type FileDescription struct {
    // PackageName 代表我们生成后的example.pb.go文件的包名
    // 也就是go文件中的 package xxx
    PackageName string

    // Services 代表我们生成后的example.pb.go文件中的所有服务
    // 我们在proto文件中写的每个server都会转化为一个 Service 实体
    Services []*Service
}

type Service struct {
    // Service 的名称
    Name string

    // Service 里具有哪些方法
    Methods []*Method
}

type Method struct {
    // 方法名称
    Name string
    // 请求类型
    RequestType string
    // 响应类型
    ResponseType string
}

这些结构结合起来描述了一个简单的proto文件信息：

因为是教学的缘故，所以各种类型的信息都很简单，几乎都用字符串存储，只保留了最核心的内容。接下来，我们需要把信息从protogen.File里复制到我们自己的结构体里。

复制proto信息到内部表示中

还记得上面的NewFile函数吗？里面有这样一段代码:

for _, s := range protoFile.Services {
    f.ParseService(s)
}

这段代码遍历protoFile中所有的Service,并调用f.ParseService()方法来处理proto中的每个service:

func (f *File) ParseService(protoSvc *protogen.Service) {
    s := &Service{
        Name:    protoSvc.GoName,
        Methods: make([]*Method, 0, len(protoSvc.Methods)),
    }

    for _, m := range protoSvc.Methods {
        // 遍历并处理Service中的所有Method
        s.Methods = append(s.Methods, f.ParseMethod(m))
    }

    f.FileDescription.Services = append(f.FileDescription.Services, s)
}

func (f *File) ParseMethod(m *protogen.Method) *Method {
    return &Method{
        Name:         m.GoName,
        RequestType:  m.Input.GoIdent.GoName,
        ResponseType: m.Output.GoIdent.GoName,
    }
}

ParseService又会调用ParseMethod方法来遍历处理service中的每个method，我将它们两个的代码一起贴上来了，里面的逻辑很简单，就是从protogen的对应结构里找到我们需要的属性复制过来，解析工作就完成了。

现在，我们的File结构体被"填满了",里面保存了一个proto文件(比较粗略)的信息。接下来让我们来创建一套模板，这将是代码生成的最后一步。

模板代码

在给你代码之前，我想先明确一下，我在例子中生成的是“基于Go标准库net/http的框架代码”。当然，你可以生成gin或其他框架的代码，这全看你自己。但在写模板之前，我们要先想想，我们要生成什么样的代码?使用者又希望你能帮他做哪些事?

要知道，proto文件不是为gRPC而生的，除了gRPC, Transport层的框架多到数不清，gin/echo/chi等都算Transport层的框架。

因此，站在业务工程师的角度上，我希望能将关注点放在业务代码上，业务代码中不能包含任何传输层的细节，这样我就可以随时以很低的成本更换传输层的框架。

叠个甲: 这里的Transport层和传输层指的不是网络协议中的传输层，别喷。

所以站在使用者的角度上，我们可能会写出以下代码:

func main() {
    // 初始化Transport
    mux := http.NewServeMux()
    // 初始化业务依赖
    svc := UserService{
        store: make([]User, 0),
    }
    // 将业务Service注册到Transport框架中
    user_pb.RegisterUserServiceHTTPServeMux(mux, &svc)
    // 启动Transport框架
    if err := http.ListenAndServe(":8080", mux); err != nil {
        panic(err)
    }
}

// 业务Handler
type UserService struct {
    store []User
}

func (u *UserService) GetUser(ctx context.Context, req *user_pb.GetUserRequest) (resp *user_pb.GetUserResponse, err error) {
    // 这里写GetUser的业务代码
}

func (u *UserService) CreateUser(ctx context.Context, req *user_pb.CreateUserRequest) (resp *user_pb.CreateUserResponse, err error) {
    // 这里写CreateUser的业务代码
}

上面这段代码中，业务代码的GetUser和CreateUser中没有任何Transport层的内容，业务代码不知道上层使用的是HTTP还是gRPC，又或者是gin等其他框架。

那我们就按照这个格式，来抽象出一个接口，作为和业务之间的契约。

如果你用过gRPC，你会发现: gRPC也是这套“契约”，这意味着未来我们要从net/http迁移到gRPC时，业务代码不需要进行任何的改造，天然适配！

那为了能让上面那段业务代码能够正常运行，我们先来手写一遍框架代码，来“适配”上面的业务代码。

这有点类似TDD(Test-Driven Development)的味道，从使用者的角度上开始，来定义代码应该“长什么样”。

我们很容易就能写出下面的适配代码, 这将是我们模板的雏形:

// 这个接口就是业务和框架的“契约”
// 实现这个接口的结构都能够注册进我们的框架中
// 这个Service对应着proto文件的service
type ServiceNameService interface {
    // 这里就是service的方法列表，对应着proto文件中service的方法列表
    ServiceMethodName(ctx context.Context, req *MethodRequestName) (resp *MethodResponseName, err error)
}

// 业务代码通过下面这段代码将服务注册到我们的框架中
func RegisterServerNameServiceHTTPServeMux(mux *http.ServeMux, svc ServiceNameService) {
    // 这里用到了依赖注入的思想
    // 此时业务代码是依赖，通过接口的形式注入进来
    s := ServiceName{
        svc: svc,
    }
    // 将对应的方法绑定到相应的路由上
    mux.HandleFunc("/UserCode", s.Name)
}

// 框架service具体实现，里面通过接口保存了业务结构体
type ServiceName struct {
    svc ServiceNameService
}

// 每个service下都会有数个method
// 每个method也都对应着proto文件里service的method
// 这里用到了适配器(Adaptor)的设计思想
// 将业务代码(通过接口)与 net/http 转换，把它们“打通”
func (s *ServiceName) Name(rw http.ResponseWriter, r *http.Request) {
    // 我们在这个函数中要做的就是把HTTP请求中的内容解析出来
    // 尝试将其转换成业务需要的参数
    _ = r.ParseForm()
    var req MethodRequestName // 这个结构是protobuf生成的，和框架无关
    switch r.Method {
    // 出于教学目的，这里只支持了POST请求
    case http.MethodPost:
        if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
            rw.WriteHeader(http.StatusBadRequest)
            rw.Write([]byte(err.Error()))
            return
        }
    default:
        rw.WriteHeader(http.StatusMethodNotAllowed)
        rw.Write([]byte(`method not allowed`))
        return
    }
    // 到这里就顺利的把HTTP请求转为了业务所需要的Request类型了
    // 接下来我们把控制权交给业务代码吧
    resp, err := s.svc.ServiceMethodName(r.Context(), &req)
    if err != nil {
        rw.WriteHeader(http.StatusInternalServerError)
        rw.Write([]byte(err.Error()))
        return
    }
    // 将业务代码返回的Response类型再转为HTTP请求返回给客户端
    if err := json.NewEncoder(rw).Encode(resp); err != nil {
        rw.WriteHeader(http.StatusInternalServerError)
        rw.Write([]byte(err.Error()))
        return
    }
}

你可以看到，上面的适配代码其实挺简陋的，它只支持POST请求，甚至HTTP路由都是proto文件里method的名字。但这对于我们学习它的核心原理已经够用了。

即使是个玩具级别的demo，它依旧用到了许多设计模式。

PS: 如果这篇文章反响还不错的话，可能会考虑后续继续加点东西。这篇文章我从晚上8点开始写，写到这里已经凌晨1:15了😢

知道了我们的模板大概长什么样子后，剩下的就简单了，替换上面代码中的Name等各个部分，就得到了我们最终的模板代码:

package template

const HTTP = `// Code generated by github.com/bootun/protoc-gen-go-example. DO NOT EDIT.
package {{.PackageName}}

import (
    "context"
    "encoding/json"
    "net/http"
)

{{range $service := .Services}}
type {{$service.Name}}Service interface {
{{range $method := .Methods}}
    {{$method.Name}}(ctx context.Context, req *{{$method.RequestType}}) (resp *{{$method.ResponseType}}, err error){{end}}
}

type {{$service.Name}} struct {
    svc {{$service.Name}}Service
}

func Register{{$service.Name}}HTTPServeMux(mux *http.ServeMux, svc {{$service.Name}}Service) {
    s := {{$service.Name}}{
        svc: svc,
    }
    {{range $method := .Methods}}
    mux.HandleFunc("/{{$method.Name}}", s.{{$method.Name}}){{end}}
}

{{range $method := .Methods}}
func (s *{{$service.Name}}) {{$method.Name}}(rw http.ResponseWriter, r *http.Request) {
    _ = r.ParseForm()
    var req {{$method.RequestType}}
    switch r.Method {
    case http.MethodPost:
        if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
            rw.WriteHeader(http.StatusBadRequest)
            return
        }
    default:
        rw.WriteHeader(http.StatusMethodNotAllowed)
        return
    }
    resp, err := s.svc.{{$method.Name}}(r.Context(), &req)
    if err != nil {
        rw.WriteHeader(http.StatusInternalServerError)
        return
    }
    if err := json.NewEncoder(rw).Encode(resp); err != nil {
        rw.WriteHeader(http.StatusInternalServerError)
        return
    }
}
{{end}}

{{end}}
`

如果你看不懂上面的语法，你需要去看下go的text/template，或者你有其他的办法能拼凑渲染出这段字符串也可以。

我们只需要将我们内部结构中的数据“填充”到模板里，交给前文提到的g.P()进行打印就可以啦:

func (f *File) Generate() error {
    tmpl, err := template.New("example-template").Parse(example_tmpl.HTTP)
    if err != nil {
        return fmt.Errorf("failed to parse example template: %w", err)
    }
    buf := &bytes.Buffer{}
    if err := tmpl.Execute(buf, f.FileDescription); err != nil {
        return fmt.Errorf("failed to execute example template: %w", err)
    }
    f.gen.P(buf.String())
    return nil
}

至此，我们就完成了所有的代码编写。

我将完整代码上传到了GitHub上: github.com/bootun/prot…

你也可以使用以下命令来直接安装

go install github.com/bootun/protoc-gen-go-example@latest

然后使用--go-example_out来生成代码:

protoc -I ./api \
   --go_out=./user \
   --go_opt=paths=source_relative \
   --go-example_out=./user \
   --go-example_opt=paths=source_relative \
    api/user.proto

都看到最后了，点个关注呗~

写到这都凌晨1:42了😱，赶快发完睡了...

本文首发于微信公众号比特要塞，欢迎关注