golang源码分析：自定义proto插件

在使用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插件的开发可以简单分为三步：

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

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

代码语言：javascript复制

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
  })
}

其中

代码语言：javascript复制

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

Options有两个字段：

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

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

代码语言：javascript复制

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
  })
}

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

代码语言：javascript复制

syntax = "proto3";

package api;

option go_package = "api/v1;v1";

message HelloRequest {
    string msg = 1;
}

代码语言：javascript复制

 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

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

代码语言：javascript复制

protoc --test_out=. test.proto

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

代码语言：javascript复制

msg: 0

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

directory exec path 插件源码分析

0 人点赞