vendor/github.com/golang/protobuf/protoc-gen-go/grpc/grpc.go

   1 // Go support for Protocol Buffers - Google's data interchange format
   2 //
   3 // Copyright 2015 The Go Authors.  All rights reserved.
   4 // https://github.com/golang/protobuf
   5 //
   6 // Redistribution and use in source and binary forms, with or without
   7 // modification, are permitted provided that the following conditions are
   8 // met:
   9 //
  10 //     * Redistributions of source code must retain the above copyright
  11 // notice, this list of conditions and the following disclaimer.
  12 //     * Redistributions in binary form must reproduce the above
  13 // copyright notice, this list of conditions and the following disclaimer
  14 // in the documentation and/or other materials provided with the
  15 // distribution.
  16 //     * Neither the name of Google Inc. nor the names of its
  17 // contributors may be used to endorse or promote products derived from
  18 // this software without specific prior written permission.
  19 //
  20 // THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
  21 // "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
  22 // LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
  23 // A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
  24 // OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
  25 // SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
  26 // LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
  27 // DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
  28 // THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
  29 // (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
  30 // OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
  31
  32 // Package grpc outputs gRPC service descriptions in Go code.
  33 // It runs as a plugin for the Go protocol buffer compiler plugin.
  34 // It is linked in to protoc-gen-go.
  35 package grpc
  36
  37 import (
  38         "fmt"
  39         "path"
  40         "strconv"
  41         "strings"
  42
  43         pb "github.com/golang/protobuf/protoc-gen-go/descriptor"
  44         "github.com/golang/protobuf/protoc-gen-go/generator"
  45 )
  46
  47 // generatedCodeVersion indicates a version of the generated code.
  48 // It is incremented whenever an incompatibility between the generated code and
  49 // the grpc package is introduced; the generated code references
  50 // a constant, grpc.SupportPackageIsVersionN (where N is generatedCodeVersion).
  51 const generatedCodeVersion = 4
  52
  53 // Paths for packages used by code generated in this file,
  54 // relative to the import_prefix of the generator.Generator.
  55 const (
  56         contextPkgPath = "golang.org/x/net/context"
  57         grpcPkgPath    = "google.golang.org/grpc"
  58 )
  59
  60 func init() {
  61         generator.RegisterPlugin(new(grpc))
  62 }
  63
  64 // grpc is an implementation of the Go protocol buffer compiler's
  65 // plugin architecture.  It generates bindings for gRPC support.
  66 type grpc struct {
  67         gen *generator.Generator
  68 }
  69
  70 // Name returns the name of this plugin, "grpc".
  71 func (g *grpc) Name() string {
  72         return "grpc"
  73 }
  74
  75 // The names for packages imported in the generated code.
  76 // They may vary from the final path component of the import path
  77 // if the name is used by other packages.
  78 var (
  79         contextPkg string
  80         grpcPkg    string
  81 )
  82
  83 // Init initializes the plugin.
  84 func (g *grpc) Init(gen *generator.Generator) {
  85         g.gen = gen
  86         contextPkg = generator.RegisterUniquePackageName("context", nil)
  87         grpcPkg = generator.RegisterUniquePackageName("grpc", nil)
  88 }
  89
  90 // Given a type name defined in a .proto, return its object.
  91 // Also record that we're using it, to guarantee the associated import.
  92 func (g *grpc) objectNamed(name string) generator.Object {
  93         g.gen.RecordTypeUse(name)
  94         return g.gen.ObjectNamed(name)
  95 }
  96
  97 // Given a type name defined in a .proto, return its name as we will print it.
  98 func (g *grpc) typeName(str string) string {
  99         return g.gen.TypeName(g.objectNamed(str))
 100 }
 101
 102 // P forwards to g.gen.P.
 103 func (g *grpc) P(args ...interface{}) { g.gen.P(args...) }
 104
 105 // Generate generates code for the services in the given file.
 106 func (g *grpc) Generate(file *generator.FileDescriptor) {
 107         if len(file.FileDescriptorProto.Service) == 0 {
 108                 return
 109         }
 110
 111         g.P("// Reference imports to suppress errors if they are not otherwise used.")
 112         g.P("var _ ", contextPkg, ".Context")
 113         g.P("var _ ", grpcPkg, ".ClientConn")
 114         g.P()
 115
 116         // Assert version compatibility.
 117         g.P("// This is a compile-time assertion to ensure that this generated file")
 118         g.P("// is compatible with the grpc package it is being compiled against.")
 119         g.P("const _ = ", grpcPkg, ".SupportPackageIsVersion", generatedCodeVersion)
 120         g.P()
 121
 122         for i, service := range file.FileDescriptorProto.Service {
 123                 g.generateService(file, service, i)
 124         }
 125 }
 126
 127 // GenerateImports generates the import declaration for this file.
 128 func (g *grpc) GenerateImports(file *generator.FileDescriptor) {
 129         if len(file.FileDescriptorProto.Service) == 0 {
 130                 return
 131         }
 132         g.P("import (")
 133         g.P(contextPkg, " ", strconv.Quote(path.Join(g.gen.ImportPrefix, contextPkgPath)))
 134         g.P(grpcPkg, " ", strconv.Quote(path.Join(g.gen.ImportPrefix, grpcPkgPath)))
 135         g.P(")")
 136         g.P()
 137 }
 138
 139 // reservedClientName records whether a client name is reserved on the client side.
 140 var reservedClientName = map[string]bool{
 141 // TODO: do we need any in gRPC?
 142 }
 143
 144 func unexport(s string) string { return strings.ToLower(s[:1]) + s[1:] }
 145
 146 // generateService generates all the code for the named service.
 147 func (g *grpc) generateService(file *generator.FileDescriptor, service *pb.ServiceDescriptorProto, index int) {
 148         path := fmt.Sprintf("6,%d", index) // 6 means service.
 149
 150         origServName := service.GetName()
 151         fullServName := origServName
 152         if pkg := file.GetPackage(); pkg != "" {
 153                 fullServName = pkg + "." + fullServName
 154         }
 155         servName := generator.CamelCase(origServName)
 156
 157         g.P()
 158         g.P("// Client API for ", servName, " service")
 159         g.P()
 160
 161         // Client interface.
 162         g.P("type ", servName, "Client interface {")
 163         for i, method := range service.Method {
 164                 g.gen.PrintComments(fmt.Sprintf("%s,2,%d", path, i)) // 2 means method in a service.
 165                 g.P(g.generateClientSignature(servName, method))
 166         }
 167         g.P("}")
 168         g.P()
 169
 170         // Client structure.
 171         g.P("type ", unexport(servName), "Client struct {")
 172         g.P("cc *", grpcPkg, ".ClientConn")
 173         g.P("}")
 174         g.P()
 175
 176         // NewClient factory.
 177         g.P("func New", servName, "Client (cc *", grpcPkg, ".ClientConn) ", servName, "Client {")
 178         g.P("return &", unexport(servName), "Client{cc}")
 179         g.P("}")
 180         g.P()
 181
 182         var methodIndex, streamIndex int
 183         serviceDescVar := "_" + servName + "_serviceDesc"
 184         // Client method implementations.
 185         for _, method := range service.Method {
 186                 var descExpr string
 187                 if !method.GetServerStreaming() && !method.GetClientStreaming() {
 188                         // Unary RPC method
 189                         descExpr = fmt.Sprintf("&%s.Methods[%d]", serviceDescVar, methodIndex)
 190                         methodIndex++
 191                 } else {
 192                         // Streaming RPC method
 193                         descExpr = fmt.Sprintf("&%s.Streams[%d]", serviceDescVar, streamIndex)
 194                         streamIndex++
 195                 }
 196                 g.generateClientMethod(servName, fullServName, serviceDescVar, method, descExpr)
 197         }
 198
 199         g.P("// Server API for ", servName, " service")
 200         g.P()
 201
 202         // Server interface.
 203         serverType := servName + "Server"
 204         g.P("type ", serverType, " interface {")
 205         for i, method := range service.Method {
 206                 g.gen.PrintComments(fmt.Sprintf("%s,2,%d", path, i)) // 2 means method in a service.
 207                 g.P(g.generateServerSignature(servName, method))
 208         }
 209         g.P("}")
 210         g.P()
 211
 212         // Server registration.
 213         g.P("func Register", servName, "Server(s *", grpcPkg, ".Server, srv ", serverType, ") {")
 214         g.P("s.RegisterService(&", serviceDescVar, `, srv)`)
 215         g.P("}")
 216         g.P()
 217
 218         // Server handler implementations.
 219         var handlerNames []string
 220         for _, method := range service.Method {
 221                 hname := g.generateServerMethod(servName, fullServName, method)
 222                 handlerNames = append(handlerNames, hname)
 223         }
 224
 225         // Service descriptor.
 226         g.P("var ", serviceDescVar, " = ", grpcPkg, ".ServiceDesc {")
 227         g.P("ServiceName: ", strconv.Quote(fullServName), ",")
 228         g.P("HandlerType: (*", serverType, ")(nil),")
 229         g.P("Methods: []", grpcPkg, ".MethodDesc{")
 230         for i, method := range service.Method {
 231                 if method.GetServerStreaming() || method.GetClientStreaming() {
 232                         continue
 233                 }
 234                 g.P("{")
 235                 g.P("MethodName: ", strconv.Quote(method.GetName()), ",")
 236                 g.P("Handler: ", handlerNames[i], ",")
 237                 g.P("},")
 238         }
 239         g.P("},")
 240         g.P("Streams: []", grpcPkg, ".StreamDesc{")
 241         for i, method := range service.Method {
 242                 if !method.GetServerStreaming() && !method.GetClientStreaming() {
 243                         continue
 244                 }
 245                 g.P("{")
 246                 g.P("StreamName: ", strconv.Quote(method.GetName()), ",")
 247                 g.P("Handler: ", handlerNames[i], ",")
 248                 if method.GetServerStreaming() {
 249                         g.P("ServerStreams: true,")
 250                 }
 251                 if method.GetClientStreaming() {
 252                         g.P("ClientStreams: true,")
 253                 }
 254                 g.P("},")
 255         }
 256         g.P("},")
 257         g.P("Metadata: \"", file.GetName(), "\",")
 258         g.P("}")
 259         g.P()
 260 }
 261
 262 // generateClientSignature returns the client-side signature for a method.
 263 func (g *grpc) generateClientSignature(servName string, method *pb.MethodDescriptorProto) string {
 264         origMethName := method.GetName()
 265         methName := generator.CamelCase(origMethName)
 266         if reservedClientName[methName] {
 267                 methName += "_"
 268         }
 269         reqArg := ", in *" + g.typeName(method.GetInputType())
 270         if method.GetClientStreaming() {
 271                 reqArg = ""
 272         }
 273         respName := "*" + g.typeName(method.GetOutputType())
 274         if method.GetServerStreaming() || method.GetClientStreaming() {
 275                 respName = servName + "_" + generator.CamelCase(origMethName) + "Client"
 276         }
 277         return fmt.Sprintf("%s(ctx %s.Context%s, opts ...%s.CallOption) (%s, error)", methName, contextPkg, reqArg, grpcPkg, respName)
 278 }
 279
 280 func (g *grpc) generateClientMethod(servName, fullServName, serviceDescVar string, method *pb.MethodDescriptorProto, descExpr string) {
 281         sname := fmt.Sprintf("/%s/%s", fullServName, method.GetName())
 282         methName := generator.CamelCase(method.GetName())
 283         inType := g.typeName(method.GetInputType())
 284         outType := g.typeName(method.GetOutputType())
 285
 286         g.P("func (c *", unexport(servName), "Client) ", g.generateClientSignature(servName, method), "{")
 287         if !method.GetServerStreaming() && !method.GetClientStreaming() {
 288                 g.P("out := new(", outType, ")")
 289                 // TODO: Pass descExpr to Invoke.
 290                 g.P("err := ", grpcPkg, `.Invoke(ctx, "`, sname, `", in, out, c.cc, opts...)`)
 291                 g.P("if err != nil { return nil, err }")
 292                 g.P("return out, nil")
 293                 g.P("}")
 294                 g.P()
 295                 return
 296         }
 297         streamType := unexport(servName) + methName + "Client"
 298         g.P("stream, err := ", grpcPkg, ".NewClientStream(ctx, ", descExpr, `, c.cc, "`, sname, `", opts...)`)
 299         g.P("if err != nil { return nil, err }")
 300         g.P("x := &", streamType, "{stream}")
 301         if !method.GetClientStreaming() {
 302                 g.P("if err := x.ClientStream.SendMsg(in); err != nil { return nil, err }")
 303                 g.P("if err := x.ClientStream.CloseSend(); err != nil { return nil, err }")
 304         }
 305         g.P("return x, nil")
 306         g.P("}")
 307         g.P()
 308
 309         genSend := method.GetClientStreaming()
 310         genRecv := method.GetServerStreaming()
 311         genCloseAndRecv := !method.GetServerStreaming()
 312
 313         // Stream auxiliary types and methods.
 314         g.P("type ", servName, "_", methName, "Client interface {")
 315         if genSend {
 316                 g.P("Send(*", inType, ") error")
 317         }
 318         if genRecv {
 319                 g.P("Recv() (*", outType, ", error)")
 320         }
 321         if genCloseAndRecv {
 322                 g.P("CloseAndRecv() (*", outType, ", error)")
 323         }
 324         g.P(grpcPkg, ".ClientStream")
 325         g.P("}")
 326         g.P()
 327
 328         g.P("type ", streamType, " struct {")
 329         g.P(grpcPkg, ".ClientStream")
 330         g.P("}")
 331         g.P()
 332
 333         if genSend {
 334                 g.P("func (x *", streamType, ") Send(m *", inType, ") error {")
 335                 g.P("return x.ClientStream.SendMsg(m)")
 336                 g.P("}")
 337                 g.P()
 338         }
 339         if genRecv {
 340                 g.P("func (x *", streamType, ") Recv() (*", outType, ", error) {")
 341                 g.P("m := new(", outType, ")")
 342                 g.P("if err := x.ClientStream.RecvMsg(m); err != nil { return nil, err }")
 343                 g.P("return m, nil")
 344                 g.P("}")
 345                 g.P()
 346         }
 347         if genCloseAndRecv {
 348                 g.P("func (x *", streamType, ") CloseAndRecv() (*", outType, ", error) {")
 349                 g.P("if err := x.ClientStream.CloseSend(); err != nil { return nil, err }")
 350                 g.P("m := new(", outType, ")")
 351                 g.P("if err := x.ClientStream.RecvMsg(m); err != nil { return nil, err }")
 352                 g.P("return m, nil")
 353                 g.P("}")
 354                 g.P()
 355         }
 356 }
 357
 358 // generateServerSignature returns the server-side signature for a method.
 359 func (g *grpc) generateServerSignature(servName string, method *pb.MethodDescriptorProto) string {
 360         origMethName := method.GetName()
 361         methName := generator.CamelCase(origMethName)
 362         if reservedClientName[methName] {
 363                 methName += "_"
 364         }
 365
 366         var reqArgs []string
 367         ret := "error"
 368         if !method.GetServerStreaming() && !method.GetClientStreaming() {
 369                 reqArgs = append(reqArgs, contextPkg+".Context")
 370                 ret = "(*" + g.typeName(method.GetOutputType()) + ", error)"
 371         }
 372         if !method.GetClientStreaming() {
 373                 reqArgs = append(reqArgs, "*"+g.typeName(method.GetInputType()))
 374         }
 375         if method.GetServerStreaming() || method.GetClientStreaming() {
 376                 reqArgs = append(reqArgs, servName+"_"+generator.CamelCase(origMethName)+"Server")
 377         }
 378
 379         return methName + "(" + strings.Join(reqArgs, ", ") + ") " + ret
 380 }
 381
 382 func (g *grpc) generateServerMethod(servName, fullServName string, method *pb.MethodDescriptorProto) string {
 383         methName := generator.CamelCase(method.GetName())
 384         hname := fmt.Sprintf("_%s_%s_Handler", servName, methName)
 385         inType := g.typeName(method.GetInputType())
 386         outType := g.typeName(method.GetOutputType())
 387
 388         if !method.GetServerStreaming() && !method.GetClientStreaming() {
 389                 g.P("func ", hname, "(srv interface{}, ctx ", contextPkg, ".Context, dec func(interface{}) error, interceptor ", grpcPkg, ".UnaryServerInterceptor) (interface{}, error) {")
 390                 g.P("in := new(", inType, ")")
 391                 g.P("if err := dec(in); err != nil { return nil, err }")
 392                 g.P("if interceptor == nil { return srv.(", servName, "Server).", methName, "(ctx, in) }")
 393                 g.P("info := &", grpcPkg, ".UnaryServerInfo{")
 394                 g.P("Server: srv,")
 395                 g.P("FullMethod: ", strconv.Quote(fmt.Sprintf("/%s/%s", fullServName, methName)), ",")
 396                 g.P("}")
 397                 g.P("handler := func(ctx ", contextPkg, ".Context, req interface{}) (interface{}, error) {")
 398                 g.P("return srv.(", servName, "Server).", methName, "(ctx, req.(*", inType, "))")
 399                 g.P("}")
 400                 g.P("return interceptor(ctx, in, info, handler)")
 401                 g.P("}")
 402                 g.P()
 403                 return hname
 404         }
 405         streamType := unexport(servName) + methName + "Server"
 406         g.P("func ", hname, "(srv interface{}, stream ", grpcPkg, ".ServerStream) error {")
 407         if !method.GetClientStreaming() {
 408                 g.P("m := new(", inType, ")")
 409                 g.P("if err := stream.RecvMsg(m); err != nil { return err }")
 410                 g.P("return srv.(", servName, "Server).", methName, "(m, &", streamType, "{stream})")
 411         } else {
 412                 g.P("return srv.(", servName, "Server).", methName, "(&", streamType, "{stream})")
 413         }
 414         g.P("}")
 415         g.P()
 416
 417         genSend := method.GetServerStreaming()
 418         genSendAndClose := !method.GetServerStreaming()
 419         genRecv := method.GetClientStreaming()
 420
 421         // Stream auxiliary types and methods.
 422         g.P("type ", servName, "_", methName, "Server interface {")
 423         if genSend {
 424                 g.P("Send(*", outType, ") error")
 425         }
 426         if genSendAndClose {
 427                 g.P("SendAndClose(*", outType, ") error")
 428         }
 429         if genRecv {
 430                 g.P("Recv() (*", inType, ", error)")
 431         }
 432         g.P(grpcPkg, ".ServerStream")
 433         g.P("}")
 434         g.P()
 435
 436         g.P("type ", streamType, " struct {")
 437         g.P(grpcPkg, ".ServerStream")
 438         g.P("}")
 439         g.P()
 440
 441         if genSend {
 442                 g.P("func (x *", streamType, ") Send(m *", outType, ") error {")
 443                 g.P("return x.ServerStream.SendMsg(m)")
 444                 g.P("}")
 445                 g.P()
 446         }
 447         if genSendAndClose {
 448                 g.P("func (x *", streamType, ") SendAndClose(m *", outType, ") error {")
 449                 g.P("return x.ServerStream.SendMsg(m)")
 450                 g.P("}")
 451                 g.P()
 452         }
 453         if genRecv {
 454                 g.P("func (x *", streamType, ") Recv() (*", inType, ", error) {")
 455                 g.P("m := new(", inType, ")")
 456                 g.P("if err := x.ServerStream.RecvMsg(m); err != nil { return nil, err }")
 457                 g.P("return m, nil")
 458                 g.P("}")
 459                 g.P()
 460         }
 461
 462         return hname
 463 }