在使用protoc的时候,可以通过指定不同的插件来生成不同的代码,它的参数统一是xx_out结尾的,制定了插件参数,就会到path下搜索protoc-gen-xx的插件。比如 protoc 通过 --foo_out 搜索插件 可执行文件 protoc-gen-foo, 也可使用参数 protoc --plugin=protoc-gen-foo=/path/to/protoc-gen-foo 指定插件位置。

protoc插件是一个独立的二进制程序,protoc进程通过fork生成子进程,并exec加载插件程序运行。父子进程间通过管道通信,并将管道的输入和输出重定向到标准输入和标准输出。

protoc进程将proto文件的信息封装为CodeGeneratorRequest传递给插件子进程,插件子进程将根据CodeGeneratorRequest中的信息,将要生成的代码数据封装为CodeGeneratorResponse对象传递给protoc进程。

插件进程从标准输入读取出CodeGeneratorRequest数据,将CodeGeneratorResponse数据写到标准输出。CodeGeneratorRequest和CodeGeneratorRequest两者也是使用proto定义的。于是一个protoc插件的开发可以简单分为三步:

  1. 从标准输入读取解析出CodeGeneratorRequest数据
  2. 利用读取的数据来生成对应的代码
  3. 将生成的结果封装为CodeGeneratorResponse写入标准输出

编写protoc插件的模板代码如下:

package main
import (
  "flag"
  "fmt"

  "google.golang.org/protobuf/compiler/protogen"
)
func main() {
  // 用于接收命令行参数
  var (
    flags        flag.FlagSet
    plugins      = flags.String("plugins", "", "list of plugins to enable (supported values: grpc)")
    importPrefix = flags.String("import_prefix", "", "prefix to prepend to import paths")
  )
  importRewriteFunc := func(importPath protogen.GoImportPath) protogen.GoImportPath {
    switch importPath {
    case "context", "fmt", "math":
      return importPath
    }
    if *importPrefix != "" {
      return protogen.GoImportPath(*importPrefix) + importPath
    }
    return importPath
  }
  protogen.Options{
    ParamFunc:         flags.Set,
    ImportRewriteFunc: importRewriteFunc,
  }.Run(func(gen *protogen.Plugin) error {
    // ...
    for _, f := range gen.Files {
      // 根据proto文件信息来生成新文件
      fmt.Println(plugins, f)
    }
    return nil
  })
}

其中

protogen.Options{
    ParamFunc:         flags.Set,
    ImportRewriteFunc: importRewriteFunc,
  }

Options有两个字段:

  1. ParamFunc:命令行中的插件参数会以--go_out=<param1>=<value1>,<param2>=<value2>:<output_directory>的形式输入并最总被解析为CodeGeneratorRequest字段,Run方法运行过程中会读取出键值对并调用ParamFunc函数,这样就可以将命令行参数绑定到flags对应的变量中了。
  2. ImportRewriteFunc:生成的新文件中的每个包导入的路径可以使用此函数进行重写

然后调用Run方法来进行相关代码的生成。下面我们实现一个简单的插件,实现解析消息体的字段名,并写文件。

package main

import (
  "strconv"
  "strings"

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

func main() {
  protogen.Options{}.Run(func(p *protogen.Plugin) error {
    // 遍历proto文件
    for _, f := range p.Files {
      fname := f.GeneratedFilenamePrefix + ".txt"
      // 后续使用t来写入新文件
      t := p.NewGeneratedFile(fname, f.GoImportPath)

      for _, msg := range f.Messages {
        builder := strings.Builder{}
        for _, field := range msg.Fields {
          builder.WriteString(field.Desc.TextName() + ": " + strconv.Itoa(field.Desc.Index()) + "\n")
        }
        t.Write([]byte(builder.String()))
      }
    }
    return nil
  })
}

可以定义一个文件来测试下

syntax = "proto3";

package api;

option go_package = "api/v1;v1";

message HelloRequest {
    string msg = 1;
}
 export PATH=$GOPATH/bin:$PATH
go build -o protoc-gen-test main.go
cp protoc-gen-test $GOPATH/bin
chmod +X $GOPATH/bin/proto-gen-test

然后就可以使用我们自定义的插件了

protoc --test_out=. test.proto

可以看到,生成了文件test.txt,内容是

msg: 0

至此一个简单的protoc插件开发完毕。