Damien Neil | 220c202 | 2018-08-15 11:24:18 -0700 | [diff] [blame] | 1 | // Copyright 2018 The Go Authors. All rights reserved. |
| 2 | // Use of this source code is governed by a BSD-style |
| 3 | // license that can be found in the LICENSE file. |
| 4 | |
| 5 | // The protoc-gen-go binary is a protoc plugin to generate a Go protocol |
| 6 | // buffer package. |
| 7 | package main |
| 8 | |
| 9 | import ( |
Damien Neil | 7779e05 | 2018-09-07 14:14:06 -0700 | [diff] [blame] | 10 | "bytes" |
| 11 | "compress/gzip" |
| 12 | "crypto/sha256" |
| 13 | "encoding/hex" |
| 14 | "fmt" |
| 15 | "strconv" |
Damien Neil | cab8dfe | 2018-09-06 14:51:28 -0700 | [diff] [blame] | 16 | "strings" |
Damien Neil | 7779e05 | 2018-09-07 14:14:06 -0700 | [diff] [blame] | 17 | |
| 18 | "github.com/golang/protobuf/proto" |
| 19 | descpb "github.com/golang/protobuf/protoc-gen-go/descriptor" |
Damien Neil | 220c202 | 2018-08-15 11:24:18 -0700 | [diff] [blame] | 20 | "google.golang.org/proto/protogen" |
Damien Neil | 46abb57 | 2018-09-07 12:45:37 -0700 | [diff] [blame^] | 21 | "google.golang.org/proto/reflect/protoreflect" |
Damien Neil | 220c202 | 2018-08-15 11:24:18 -0700 | [diff] [blame] | 22 | ) |
| 23 | |
Damien Neil | 46abb57 | 2018-09-07 12:45:37 -0700 | [diff] [blame^] | 24 | const protoPackage = "github.com/golang/protobuf/proto" |
| 25 | |
Damien Neil | 220c202 | 2018-08-15 11:24:18 -0700 | [diff] [blame] | 26 | func main() { |
| 27 | protogen.Run(func(gen *protogen.Plugin) error { |
| 28 | for _, f := range gen.Files { |
| 29 | if !f.Generate { |
| 30 | continue |
| 31 | } |
| 32 | genFile(gen, f) |
| 33 | } |
| 34 | return nil |
| 35 | }) |
| 36 | } |
| 37 | |
Damien Neil | cab8dfe | 2018-09-06 14:51:28 -0700 | [diff] [blame] | 38 | type File struct { |
| 39 | *protogen.File |
Damien Neil | 46abb57 | 2018-09-07 12:45:37 -0700 | [diff] [blame^] | 40 | locationMap map[string][]*descpb.SourceCodeInfo_Location |
| 41 | descriptorVar string // var containing the gzipped FileDescriptorProto |
| 42 | init []string |
Damien Neil | cab8dfe | 2018-09-06 14:51:28 -0700 | [diff] [blame] | 43 | } |
| 44 | |
| 45 | func genFile(gen *protogen.Plugin, file *protogen.File) { |
| 46 | f := &File{ |
| 47 | File: file, |
| 48 | locationMap: make(map[string][]*descpb.SourceCodeInfo_Location), |
| 49 | } |
| 50 | for _, loc := range file.Proto.GetSourceCodeInfo().GetLocation() { |
| 51 | key := pathKey(loc.Path) |
| 52 | f.locationMap[key] = append(f.locationMap[key], loc) |
| 53 | } |
| 54 | |
Damien Neil | 46abb57 | 2018-09-07 12:45:37 -0700 | [diff] [blame^] | 55 | // Determine the name of the var holding the file descriptor: |
| 56 | // |
| 57 | // fileDescriptor_<hash of filename> |
| 58 | filenameHash := sha256.Sum256([]byte(f.Desc.Path())) |
| 59 | f.descriptorVar = fmt.Sprintf("fileDescriptor_%s", hex.EncodeToString(filenameHash[:8])) |
| 60 | |
Damien Neil | 082ce92 | 2018-09-06 10:23:53 -0700 | [diff] [blame] | 61 | g := gen.NewGeneratedFile(f.GeneratedFilenamePrefix+".pb.go", f.GoImportPath) |
Damien Neil | 220c202 | 2018-08-15 11:24:18 -0700 | [diff] [blame] | 62 | g.P("// Code generated by protoc-gen-go. DO NOT EDIT.") |
Damien Neil | abc6fc1 | 2018-08-23 14:39:30 -0700 | [diff] [blame] | 63 | g.P("// source: ", f.Desc.Path()) |
Damien Neil | 220c202 | 2018-08-15 11:24:18 -0700 | [diff] [blame] | 64 | g.P() |
Damien Neil | cab8dfe | 2018-09-06 14:51:28 -0700 | [diff] [blame] | 65 | const filePackageField = 2 // FileDescriptorProto.package |
| 66 | genComment(g, f, []int32{filePackageField}) |
| 67 | g.P() |
Damien Neil | 082ce92 | 2018-09-06 10:23:53 -0700 | [diff] [blame] | 68 | g.P("package ", f.GoPackageName) |
Damien Neil | c7d07d9 | 2018-08-22 13:46:02 -0700 | [diff] [blame] | 69 | g.P() |
| 70 | |
Damien Neil | 46abb57 | 2018-09-07 12:45:37 -0700 | [diff] [blame^] | 71 | for _, enum := range f.Enums { |
| 72 | genEnum(gen, g, f, enum) |
| 73 | } |
Damien Neil | cab8dfe | 2018-09-06 14:51:28 -0700 | [diff] [blame] | 74 | for _, message := range f.Messages { |
| 75 | genMessage(gen, g, f, message) |
Damien Neil | c7d07d9 | 2018-08-22 13:46:02 -0700 | [diff] [blame] | 76 | } |
Damien Neil | 220c202 | 2018-08-15 11:24:18 -0700 | [diff] [blame] | 77 | |
Damien Neil | 46abb57 | 2018-09-07 12:45:37 -0700 | [diff] [blame^] | 78 | if len(f.init) != 0 { |
| 79 | g.P("func init() {") |
| 80 | for _, s := range f.init { |
| 81 | g.P(s) |
| 82 | } |
| 83 | g.P("}") |
| 84 | g.P() |
| 85 | } |
| 86 | |
Damien Neil | 7779e05 | 2018-09-07 14:14:06 -0700 | [diff] [blame] | 87 | genFileDescriptor(gen, g, f) |
| 88 | } |
| 89 | |
Damien Neil | cab8dfe | 2018-09-06 14:51:28 -0700 | [diff] [blame] | 90 | func genFileDescriptor(gen *protogen.Plugin, g *protogen.GeneratedFile, f *File) { |
Damien Neil | 7779e05 | 2018-09-07 14:14:06 -0700 | [diff] [blame] | 91 | // Trim the source_code_info from the descriptor. |
| 92 | // Marshal and gzip it. |
| 93 | descProto := proto.Clone(f.Proto).(*descpb.FileDescriptorProto) |
| 94 | descProto.SourceCodeInfo = nil |
| 95 | b, err := proto.Marshal(descProto) |
| 96 | if err != nil { |
| 97 | gen.Error(err) |
| 98 | return |
| 99 | } |
| 100 | var buf bytes.Buffer |
| 101 | w, _ := gzip.NewWriterLevel(&buf, gzip.BestCompression) |
| 102 | w.Write(b) |
| 103 | w.Close() |
| 104 | b = buf.Bytes() |
| 105 | |
Damien Neil | 46abb57 | 2018-09-07 12:45:37 -0700 | [diff] [blame^] | 106 | g.P("func init() { proto.RegisterFile(", strconv.Quote(f.Desc.Path()), ", ", f.descriptorVar, ") }") |
Damien Neil | 7779e05 | 2018-09-07 14:14:06 -0700 | [diff] [blame] | 107 | g.P() |
Damien Neil | 46abb57 | 2018-09-07 12:45:37 -0700 | [diff] [blame^] | 108 | g.P("var ", f.descriptorVar, " = []byte{") |
Damien Neil | 7779e05 | 2018-09-07 14:14:06 -0700 | [diff] [blame] | 109 | g.P("// ", len(b), " bytes of a gzipped FileDescriptorProto") |
| 110 | for len(b) > 0 { |
| 111 | n := 16 |
| 112 | if n > len(b) { |
| 113 | n = len(b) |
| 114 | } |
| 115 | |
| 116 | s := "" |
| 117 | for _, c := range b[:n] { |
| 118 | s += fmt.Sprintf("0x%02x,", c) |
| 119 | } |
| 120 | g.P(s) |
| 121 | |
| 122 | b = b[n:] |
| 123 | } |
| 124 | g.P("}") |
| 125 | g.P() |
Damien Neil | 220c202 | 2018-08-15 11:24:18 -0700 | [diff] [blame] | 126 | } |
Damien Neil | c7d07d9 | 2018-08-22 13:46:02 -0700 | [diff] [blame] | 127 | |
Damien Neil | 46abb57 | 2018-09-07 12:45:37 -0700 | [diff] [blame^] | 128 | func genEnum(gen *protogen.Plugin, g *protogen.GeneratedFile, f *File, enum *protogen.Enum) { |
| 129 | genComment(g, f, enum.Path) |
| 130 | // TODO: deprecation |
| 131 | g.P("type ", enum.GoIdent, " int32") |
| 132 | g.P("const (") |
| 133 | for _, value := range enum.Values { |
| 134 | genComment(g, f, value.Path) |
| 135 | // TODO: deprecation |
| 136 | g.P(value.GoIdent, " ", enum.GoIdent, " = ", value.Desc.Number()) |
| 137 | } |
| 138 | g.P(")") |
| 139 | g.P() |
| 140 | nameMap := enum.GoIdent.GoName + "_name" |
| 141 | g.P("var ", nameMap, " = map[int32]string{") |
| 142 | generated := make(map[protoreflect.EnumNumber]bool) |
| 143 | for _, value := range enum.Values { |
| 144 | duplicate := "" |
| 145 | if _, present := generated[value.Desc.Number()]; present { |
| 146 | duplicate = "// Duplicate value: " |
| 147 | } |
| 148 | g.P(duplicate, value.Desc.Number(), ": ", strconv.Quote(string(value.Desc.Name())), ",") |
| 149 | generated[value.Desc.Number()] = true |
| 150 | } |
| 151 | g.P("}") |
| 152 | g.P() |
| 153 | valueMap := enum.GoIdent.GoName + "_value" |
| 154 | g.P("var ", valueMap, " = map[string]int32{") |
| 155 | for _, value := range enum.Values { |
| 156 | g.P(strconv.Quote(string(value.Desc.Name())), ": ", value.Desc.Number(), ",") |
| 157 | } |
| 158 | g.P("}") |
| 159 | g.P() |
| 160 | if enum.Desc.Syntax() != protoreflect.Proto3 { |
| 161 | g.P("func (x ", enum.GoIdent, ") Enum() *", enum.GoIdent, " {") |
| 162 | g.P("p := new(", enum.GoIdent, ")") |
| 163 | g.P("*p = x") |
| 164 | g.P("return p") |
| 165 | g.P("}") |
| 166 | g.P() |
| 167 | } |
| 168 | g.P("func (x ", enum.GoIdent, ") String() string {") |
| 169 | g.P("return ", protogen.GoIdent{GoImportPath: protoPackage, GoName: "EnumName"}, "(", enum.GoIdent, "_name, int32(x))") |
| 170 | g.P("}") |
| 171 | g.P() |
| 172 | |
| 173 | if enum.Desc.Syntax() != protoreflect.Proto3 { |
| 174 | g.P("func (x *", enum.GoIdent, ") UnmarshalJSON(data []byte) error {") |
| 175 | g.P("value, err := ", protogen.GoIdent{GoImportPath: protoPackage, GoName: "UnmarshalJSONEnum"}, "(", enum.GoIdent, `_value, data, "`, enum.GoIdent, `")`) |
| 176 | g.P("if err != nil {") |
| 177 | g.P("return err") |
| 178 | g.P("}") |
| 179 | g.P("*x = ", enum.GoIdent, "(value)") |
| 180 | g.P("return nil") |
| 181 | g.P("}") |
| 182 | g.P() |
| 183 | } |
| 184 | |
| 185 | var indexes []string |
| 186 | for i := 1; i < len(enum.Path); i += 2 { |
| 187 | indexes = append(indexes, strconv.Itoa(int(enum.Path[i]))) |
| 188 | } |
| 189 | g.P("func (", enum.GoIdent, ") EnumDescriptor() ([]byte, []int) {") |
| 190 | g.P("return ", f.descriptorVar, ", []int{", strings.Join(indexes, ","), "}") |
| 191 | g.P("}") |
| 192 | g.P() |
| 193 | |
| 194 | genWellKnownType(g, enum.GoIdent, enum.Desc) |
| 195 | |
| 196 | // The name registered is, confusingly, <proto_package>.<go_ident>. |
| 197 | // This probably should have been the full name of the proto enum |
| 198 | // type instead, but changing it at this point would require thought. |
| 199 | regName := string(f.Desc.Package()) + "." + enum.GoIdent.GoName |
| 200 | f.init = append(f.init, fmt.Sprintf("%s(%q, %s, %s)", |
| 201 | g.QualifiedGoIdent(protogen.GoIdent{ |
| 202 | GoImportPath: protoPackage, |
| 203 | GoName: "RegisterEnum", |
| 204 | }), |
| 205 | regName, nameMap, valueMap, |
| 206 | )) |
| 207 | } |
| 208 | |
Damien Neil | cab8dfe | 2018-09-06 14:51:28 -0700 | [diff] [blame] | 209 | func genMessage(gen *protogen.Plugin, g *protogen.GeneratedFile, f *File, message *protogen.Message) { |
Damien Neil | 46abb57 | 2018-09-07 12:45:37 -0700 | [diff] [blame^] | 210 | for _, enum := range message.Enums { |
| 211 | genEnum(gen, g, f, enum) |
| 212 | } |
| 213 | |
Damien Neil | cab8dfe | 2018-09-06 14:51:28 -0700 | [diff] [blame] | 214 | genComment(g, f, message.Path) |
| 215 | g.P("type ", message.GoIdent, " struct {") |
Damien Neil | c7d07d9 | 2018-08-22 13:46:02 -0700 | [diff] [blame] | 216 | g.P("}") |
| 217 | g.P() |
| 218 | |
Damien Neil | cab8dfe | 2018-09-06 14:51:28 -0700 | [diff] [blame] | 219 | for _, nested := range message.Messages { |
| 220 | genMessage(gen, g, f, nested) |
Damien Neil | c7d07d9 | 2018-08-22 13:46:02 -0700 | [diff] [blame] | 221 | } |
| 222 | } |
Damien Neil | cab8dfe | 2018-09-06 14:51:28 -0700 | [diff] [blame] | 223 | |
| 224 | func genComment(g *protogen.GeneratedFile, f *File, path []int32) { |
| 225 | for _, loc := range f.locationMap[pathKey(path)] { |
| 226 | if loc.LeadingComments == nil { |
| 227 | continue |
| 228 | } |
| 229 | for _, line := range strings.Split(strings.TrimSuffix(loc.GetLeadingComments(), "\n"), "\n") { |
| 230 | g.P("//", line) |
| 231 | } |
| 232 | return |
| 233 | } |
| 234 | } |
| 235 | |
| 236 | // pathKey converts a location path to a string suitable for use as a map key. |
| 237 | func pathKey(path []int32) string { |
| 238 | var buf []byte |
| 239 | for i, x := range path { |
| 240 | if i != 0 { |
| 241 | buf = append(buf, ',') |
| 242 | } |
| 243 | buf = strconv.AppendInt(buf, int64(x), 10) |
| 244 | } |
| 245 | return string(buf) |
| 246 | } |
Damien Neil | 46abb57 | 2018-09-07 12:45:37 -0700 | [diff] [blame^] | 247 | |
| 248 | func genWellKnownType(g *protogen.GeneratedFile, ident protogen.GoIdent, desc protoreflect.Descriptor) { |
| 249 | if wellKnownTypes[desc.FullName()] { |
| 250 | g.P("func (", ident, `) XXX_WellKnownType() string { return "`, desc.Name(), `" }`) |
| 251 | g.P() |
| 252 | } |
| 253 | } |
| 254 | |
| 255 | // Names of messages and enums for which we will generate XXX_WellKnownType methods. |
| 256 | var wellKnownTypes = map[protoreflect.FullName]bool{ |
| 257 | "google.protobuf.NullValue": true, |
| 258 | } |