CodeGeneratorRequest ジェネレータ

protocのプラグインを作るときに、protocから呼び出される形で動かすとプロセスをデバッガにアタッチさせるのが難しいので単体で動かしたい。
単体で動かす場合stdinから CodeGeneratorRequest をバイナリで入力する必要がある。

CodeGeneratorRequest をバイナリファイルとしてとりだすprotoc pluginを作った。

//
// CodeGeneratorRequest binary generator
// go build  -o protoc-gen-echo  main.go
// protoc --plugin=./protoc-gen-echo --echo_out=. test.pb

package main

import (
    "io"
    "io/ioutil"
    "log"
    "os"

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

func parseReq(r io.Reader) (*plugin.CodeGeneratorRequest, error) {
    buf, err := ioutil.ReadAll(r)
    if err != nil {
        return nil, err
    }

    ioutil.WriteFile("CodeGeneratorRequest.bin.pb", buf, 0644)

    var req plugin.CodeGeneratorRequest
    if err = proto.Unmarshal(buf, &req); err != nil {
        return nil, err
    }
    return &req, nil
}

func processReq(req *plugin.CodeGeneratorRequest) *plugin.CodeGeneratorResponse {
    files := make(map[string]*descriptor.FileDescriptorProto)
    for _, f := range req.ProtoFile {
        files[f.GetName()] = f
    }

    var resp plugin.CodeGeneratorResponse
    for _, fname := range req.FileToGenerate {
        f := files[fname]
        out := fname + ".dump"
        resp.File = append(resp.File, &plugin.CodeGeneratorResponse_File{
            Name:    proto.String(out),
            Content: proto.String(proto.MarshalTextString(f)),
        })
    }
    return &resp
}

func emitResp(resp *plugin.CodeGeneratorResponse) error {
    buf, err := proto.Marshal(resp)
    if err != nil {
        return err
    }
    _, err = os.Stdout.Write(buf)
    return err
}

func run() error {
    req, err := parseReq(os.Stdin)
    if err != nil {
        return err
    }

    resp := processReq(req)

    return emitResp(resp)
}

func main() {
    if err := run(); err != nil {
        log.Fatalln(err)
    }
}