mirror of
				https://gitea.com/Lydanne/buildx.git
				synced 2025-11-04 10:03:42 +08:00 
			
		
		
		
	
		
			
				
	
	
		
			399 lines
		
	
	
		
			13 KiB
		
	
	
	
		
			Go
		
	
	
	
	
	
			
		
		
	
	
			399 lines
		
	
	
		
			13 KiB
		
	
	
	
		
			Go
		
	
	
	
	
	
// Copyright 2018 The Go Authors. All rights reserved.
 | 
						|
// Use of this source code is governed by a BSD-style
 | 
						|
// license that can be found in the LICENSE file.
 | 
						|
 | 
						|
// Package gengogrpc contains the gRPC code generator.
 | 
						|
package gengogrpc
 | 
						|
 | 
						|
import (
 | 
						|
	"fmt"
 | 
						|
	"strconv"
 | 
						|
	"strings"
 | 
						|
 | 
						|
	"google.golang.org/protobuf/compiler/protogen"
 | 
						|
 | 
						|
	"google.golang.org/protobuf/types/descriptorpb"
 | 
						|
)
 | 
						|
 | 
						|
const (
 | 
						|
	contextPackage = protogen.GoImportPath("context")
 | 
						|
	grpcPackage    = protogen.GoImportPath("google.golang.org/grpc")
 | 
						|
	codesPackage   = protogen.GoImportPath("google.golang.org/grpc/codes")
 | 
						|
	statusPackage  = protogen.GoImportPath("google.golang.org/grpc/status")
 | 
						|
)
 | 
						|
 | 
						|
// GenerateFile generates a _grpc.pb.go file containing gRPC service definitions.
 | 
						|
func GenerateFile(gen *protogen.Plugin, file *protogen.File) *protogen.GeneratedFile {
 | 
						|
	if len(file.Services) == 0 {
 | 
						|
		return nil
 | 
						|
	}
 | 
						|
	filename := file.GeneratedFilenamePrefix + "_grpc.pb.go"
 | 
						|
	g := gen.NewGeneratedFile(filename, file.GoImportPath)
 | 
						|
	g.P("// Code generated by protoc-gen-go-grpc. DO NOT EDIT.")
 | 
						|
	g.P()
 | 
						|
	g.P("package ", file.GoPackageName)
 | 
						|
	g.P()
 | 
						|
	GenerateFileContent(gen, file, g)
 | 
						|
	return g
 | 
						|
}
 | 
						|
 | 
						|
// GenerateFileContent generates the gRPC service definitions, excluding the package statement.
 | 
						|
func GenerateFileContent(gen *protogen.Plugin, file *protogen.File, g *protogen.GeneratedFile) {
 | 
						|
	if len(file.Services) == 0 {
 | 
						|
		return
 | 
						|
	}
 | 
						|
 | 
						|
	// TODO: Remove this. We don't need to include these references any more.
 | 
						|
	g.P("// Reference imports to suppress errors if they are not otherwise used.")
 | 
						|
	g.P("var _ ", contextPackage.Ident("Context"))
 | 
						|
	g.P("var _ ", grpcPackage.Ident("ClientConnInterface"))
 | 
						|
	g.P()
 | 
						|
 | 
						|
	g.P("// This is a compile-time assertion to ensure that this generated file")
 | 
						|
	g.P("// is compatible with the grpc package it is being compiled against.")
 | 
						|
	g.P("const _ = ", grpcPackage.Ident("SupportPackageIsVersion6"))
 | 
						|
	g.P()
 | 
						|
	for _, service := range file.Services {
 | 
						|
		genService(gen, file, g, service)
 | 
						|
	}
 | 
						|
}
 | 
						|
 | 
						|
func genService(gen *protogen.Plugin, file *protogen.File, g *protogen.GeneratedFile, service *protogen.Service) {
 | 
						|
	clientName := service.GoName + "Client"
 | 
						|
 | 
						|
	g.P("// ", clientName, " is the client API for ", service.GoName, " service.")
 | 
						|
	g.P("//")
 | 
						|
	g.P("// For semantics around ctx use and closing/ending streaming RPCs, please refer to https://godoc.org/google.golang.org/grpc#ClientConn.NewStream.")
 | 
						|
 | 
						|
	// Client interface.
 | 
						|
	if service.Desc.Options().(*descriptorpb.ServiceOptions).GetDeprecated() {
 | 
						|
		g.P("//")
 | 
						|
		g.P(deprecationComment)
 | 
						|
	}
 | 
						|
	g.Annotate(clientName, service.Location)
 | 
						|
	g.P("type ", clientName, " interface {")
 | 
						|
	for _, method := range service.Methods {
 | 
						|
		g.Annotate(clientName+"."+method.GoName, method.Location)
 | 
						|
		if method.Desc.Options().(*descriptorpb.MethodOptions).GetDeprecated() {
 | 
						|
			g.P(deprecationComment)
 | 
						|
		}
 | 
						|
		g.P(method.Comments.Leading,
 | 
						|
			clientSignature(g, method))
 | 
						|
	}
 | 
						|
	g.P("}")
 | 
						|
	g.P()
 | 
						|
 | 
						|
	// Client structure.
 | 
						|
	g.P("type ", unexport(clientName), " struct {")
 | 
						|
	g.P("cc ", grpcPackage.Ident("ClientConnInterface"))
 | 
						|
	g.P("}")
 | 
						|
	g.P()
 | 
						|
 | 
						|
	// NewClient factory.
 | 
						|
	if service.Desc.Options().(*descriptorpb.ServiceOptions).GetDeprecated() {
 | 
						|
		g.P(deprecationComment)
 | 
						|
	}
 | 
						|
	g.P("func New", clientName, " (cc ", grpcPackage.Ident("ClientConnInterface"), ") ", clientName, " {")
 | 
						|
	g.P("return &", unexport(clientName), "{cc}")
 | 
						|
	g.P("}")
 | 
						|
	g.P()
 | 
						|
 | 
						|
	var methodIndex, streamIndex int
 | 
						|
	// Client method implementations.
 | 
						|
	for _, method := range service.Methods {
 | 
						|
		if !method.Desc.IsStreamingServer() && !method.Desc.IsStreamingClient() {
 | 
						|
			// Unary RPC method
 | 
						|
			genClientMethod(gen, file, g, method, methodIndex)
 | 
						|
			methodIndex++
 | 
						|
		} else {
 | 
						|
			// Streaming RPC method
 | 
						|
			genClientMethod(gen, file, g, method, streamIndex)
 | 
						|
			streamIndex++
 | 
						|
		}
 | 
						|
	}
 | 
						|
 | 
						|
	// Server interface.
 | 
						|
	serverType := service.GoName + "Server"
 | 
						|
	g.P("// ", serverType, " is the server API for ", service.GoName, " service.")
 | 
						|
	if service.Desc.Options().(*descriptorpb.ServiceOptions).GetDeprecated() {
 | 
						|
		g.P("//")
 | 
						|
		g.P(deprecationComment)
 | 
						|
	}
 | 
						|
	g.Annotate(serverType, service.Location)
 | 
						|
	g.P("type ", serverType, " interface {")
 | 
						|
	for _, method := range service.Methods {
 | 
						|
		g.Annotate(serverType+"."+method.GoName, method.Location)
 | 
						|
		if method.Desc.Options().(*descriptorpb.MethodOptions).GetDeprecated() {
 | 
						|
			g.P(deprecationComment)
 | 
						|
		}
 | 
						|
		g.P(method.Comments.Leading,
 | 
						|
			serverSignature(g, method))
 | 
						|
	}
 | 
						|
	g.P("}")
 | 
						|
	g.P()
 | 
						|
 | 
						|
	// Server Unimplemented struct for forward compatibility.
 | 
						|
	g.P("// Unimplemented", serverType, " can be embedded to have forward compatible implementations.")
 | 
						|
	g.P("type Unimplemented", serverType, " struct {")
 | 
						|
	g.P("}")
 | 
						|
	g.P()
 | 
						|
	for _, method := range service.Methods {
 | 
						|
		nilArg := ""
 | 
						|
		if !method.Desc.IsStreamingClient() && !method.Desc.IsStreamingServer() {
 | 
						|
			nilArg = "nil,"
 | 
						|
		}
 | 
						|
		g.P("func (*Unimplemented", serverType, ") ", serverSignature(g, method), "{")
 | 
						|
		g.P("return ", nilArg, statusPackage.Ident("Errorf"), "(", codesPackage.Ident("Unimplemented"), `, "method `, method.GoName, ` not implemented")`)
 | 
						|
		g.P("}")
 | 
						|
	}
 | 
						|
	g.P()
 | 
						|
 | 
						|
	// Server registration.
 | 
						|
	if service.Desc.Options().(*descriptorpb.ServiceOptions).GetDeprecated() {
 | 
						|
		g.P(deprecationComment)
 | 
						|
	}
 | 
						|
	serviceDescVar := "_" + service.GoName + "_serviceDesc"
 | 
						|
	g.P("func Register", service.GoName, "Server(s *", grpcPackage.Ident("Server"), ", srv ", serverType, ") {")
 | 
						|
	g.P("s.RegisterService(&", serviceDescVar, `, srv)`)
 | 
						|
	g.P("}")
 | 
						|
	g.P()
 | 
						|
 | 
						|
	// Server handler implementations.
 | 
						|
	var handlerNames []string
 | 
						|
	for _, method := range service.Methods {
 | 
						|
		hname := genServerMethod(gen, file, g, method)
 | 
						|
		handlerNames = append(handlerNames, hname)
 | 
						|
	}
 | 
						|
 | 
						|
	// Service descriptor.
 | 
						|
	g.P("var ", serviceDescVar, " = ", grpcPackage.Ident("ServiceDesc"), " {")
 | 
						|
	g.P("ServiceName: ", strconv.Quote(string(service.Desc.FullName())), ",")
 | 
						|
	g.P("HandlerType: (*", serverType, ")(nil),")
 | 
						|
	g.P("Methods: []", grpcPackage.Ident("MethodDesc"), "{")
 | 
						|
	for i, method := range service.Methods {
 | 
						|
		if method.Desc.IsStreamingClient() || method.Desc.IsStreamingServer() {
 | 
						|
			continue
 | 
						|
		}
 | 
						|
		g.P("{")
 | 
						|
		g.P("MethodName: ", strconv.Quote(string(method.Desc.Name())), ",")
 | 
						|
		g.P("Handler: ", handlerNames[i], ",")
 | 
						|
		g.P("},")
 | 
						|
	}
 | 
						|
	g.P("},")
 | 
						|
	g.P("Streams: []", grpcPackage.Ident("StreamDesc"), "{")
 | 
						|
	for i, method := range service.Methods {
 | 
						|
		if !method.Desc.IsStreamingClient() && !method.Desc.IsStreamingServer() {
 | 
						|
			continue
 | 
						|
		}
 | 
						|
		g.P("{")
 | 
						|
		g.P("StreamName: ", strconv.Quote(string(method.Desc.Name())), ",")
 | 
						|
		g.P("Handler: ", handlerNames[i], ",")
 | 
						|
		if method.Desc.IsStreamingServer() {
 | 
						|
			g.P("ServerStreams: true,")
 | 
						|
		}
 | 
						|
		if method.Desc.IsStreamingClient() {
 | 
						|
			g.P("ClientStreams: true,")
 | 
						|
		}
 | 
						|
		g.P("},")
 | 
						|
	}
 | 
						|
	g.P("},")
 | 
						|
	g.P("Metadata: \"", file.Desc.Path(), "\",")
 | 
						|
	g.P("}")
 | 
						|
	g.P()
 | 
						|
}
 | 
						|
 | 
						|
func clientSignature(g *protogen.GeneratedFile, method *protogen.Method) string {
 | 
						|
	s := method.GoName + "(ctx " + g.QualifiedGoIdent(contextPackage.Ident("Context"))
 | 
						|
	if !method.Desc.IsStreamingClient() {
 | 
						|
		s += ", in *" + g.QualifiedGoIdent(method.Input.GoIdent)
 | 
						|
	}
 | 
						|
	s += ", opts ..." + g.QualifiedGoIdent(grpcPackage.Ident("CallOption")) + ") ("
 | 
						|
	if !method.Desc.IsStreamingClient() && !method.Desc.IsStreamingServer() {
 | 
						|
		s += "*" + g.QualifiedGoIdent(method.Output.GoIdent)
 | 
						|
	} else {
 | 
						|
		s += method.Parent.GoName + "_" + method.GoName + "Client"
 | 
						|
	}
 | 
						|
	s += ", error)"
 | 
						|
	return s
 | 
						|
}
 | 
						|
 | 
						|
func genClientMethod(gen *protogen.Plugin, file *protogen.File, g *protogen.GeneratedFile, method *protogen.Method, index int) {
 | 
						|
	service := method.Parent
 | 
						|
	sname := fmt.Sprintf("/%s/%s", service.Desc.FullName(), method.Desc.Name())
 | 
						|
 | 
						|
	if method.Desc.Options().(*descriptorpb.MethodOptions).GetDeprecated() {
 | 
						|
		g.P(deprecationComment)
 | 
						|
	}
 | 
						|
	g.P("func (c *", unexport(service.GoName), "Client) ", clientSignature(g, method), "{")
 | 
						|
	if !method.Desc.IsStreamingServer() && !method.Desc.IsStreamingClient() {
 | 
						|
		g.P("out := new(", method.Output.GoIdent, ")")
 | 
						|
		g.P(`err := c.cc.Invoke(ctx, "`, sname, `", in, out, opts...)`)
 | 
						|
		g.P("if err != nil { return nil, err }")
 | 
						|
		g.P("return out, nil")
 | 
						|
		g.P("}")
 | 
						|
		g.P()
 | 
						|
		return
 | 
						|
	}
 | 
						|
	streamType := unexport(service.GoName) + method.GoName + "Client"
 | 
						|
	serviceDescVar := "_" + service.GoName + "_serviceDesc"
 | 
						|
	g.P("stream, err := c.cc.NewStream(ctx, &", serviceDescVar, ".Streams[", index, `], "`, sname, `", opts...)`)
 | 
						|
	g.P("if err != nil { return nil, err }")
 | 
						|
	g.P("x := &", streamType, "{stream}")
 | 
						|
	if !method.Desc.IsStreamingClient() {
 | 
						|
		g.P("if err := x.ClientStream.SendMsg(in); err != nil { return nil, err }")
 | 
						|
		g.P("if err := x.ClientStream.CloseSend(); err != nil { return nil, err }")
 | 
						|
	}
 | 
						|
	g.P("return x, nil")
 | 
						|
	g.P("}")
 | 
						|
	g.P()
 | 
						|
 | 
						|
	genSend := method.Desc.IsStreamingClient()
 | 
						|
	genRecv := method.Desc.IsStreamingServer()
 | 
						|
	genCloseAndRecv := !method.Desc.IsStreamingServer()
 | 
						|
 | 
						|
	// Stream auxiliary types and methods.
 | 
						|
	g.P("type ", service.GoName, "_", method.GoName, "Client interface {")
 | 
						|
	if genSend {
 | 
						|
		g.P("Send(*", method.Input.GoIdent, ") error")
 | 
						|
	}
 | 
						|
	if genRecv {
 | 
						|
		g.P("Recv() (*", method.Output.GoIdent, ", error)")
 | 
						|
	}
 | 
						|
	if genCloseAndRecv {
 | 
						|
		g.P("CloseAndRecv() (*", method.Output.GoIdent, ", error)")
 | 
						|
	}
 | 
						|
	g.P(grpcPackage.Ident("ClientStream"))
 | 
						|
	g.P("}")
 | 
						|
	g.P()
 | 
						|
 | 
						|
	g.P("type ", streamType, " struct {")
 | 
						|
	g.P(grpcPackage.Ident("ClientStream"))
 | 
						|
	g.P("}")
 | 
						|
	g.P()
 | 
						|
 | 
						|
	if genSend {
 | 
						|
		g.P("func (x *", streamType, ") Send(m *", method.Input.GoIdent, ") error {")
 | 
						|
		g.P("return x.ClientStream.SendMsg(m)")
 | 
						|
		g.P("}")
 | 
						|
		g.P()
 | 
						|
	}
 | 
						|
	if genRecv {
 | 
						|
		g.P("func (x *", streamType, ") Recv() (*", method.Output.GoIdent, ", error) {")
 | 
						|
		g.P("m := new(", method.Output.GoIdent, ")")
 | 
						|
		g.P("if err := x.ClientStream.RecvMsg(m); err != nil { return nil, err }")
 | 
						|
		g.P("return m, nil")
 | 
						|
		g.P("}")
 | 
						|
		g.P()
 | 
						|
	}
 | 
						|
	if genCloseAndRecv {
 | 
						|
		g.P("func (x *", streamType, ") CloseAndRecv() (*", method.Output.GoIdent, ", error) {")
 | 
						|
		g.P("if err := x.ClientStream.CloseSend(); err != nil { return nil, err }")
 | 
						|
		g.P("m := new(", method.Output.GoIdent, ")")
 | 
						|
		g.P("if err := x.ClientStream.RecvMsg(m); err != nil { return nil, err }")
 | 
						|
		g.P("return m, nil")
 | 
						|
		g.P("}")
 | 
						|
		g.P()
 | 
						|
	}
 | 
						|
}
 | 
						|
 | 
						|
func serverSignature(g *protogen.GeneratedFile, method *protogen.Method) string {
 | 
						|
	var reqArgs []string
 | 
						|
	ret := "error"
 | 
						|
	if !method.Desc.IsStreamingClient() && !method.Desc.IsStreamingServer() {
 | 
						|
		reqArgs = append(reqArgs, g.QualifiedGoIdent(contextPackage.Ident("Context")))
 | 
						|
		ret = "(*" + g.QualifiedGoIdent(method.Output.GoIdent) + ", error)"
 | 
						|
	}
 | 
						|
	if !method.Desc.IsStreamingClient() {
 | 
						|
		reqArgs = append(reqArgs, "*"+g.QualifiedGoIdent(method.Input.GoIdent))
 | 
						|
	}
 | 
						|
	if method.Desc.IsStreamingClient() || method.Desc.IsStreamingServer() {
 | 
						|
		reqArgs = append(reqArgs, method.Parent.GoName+"_"+method.GoName+"Server")
 | 
						|
	}
 | 
						|
	return method.GoName + "(" + strings.Join(reqArgs, ", ") + ") " + ret
 | 
						|
}
 | 
						|
 | 
						|
func genServerMethod(gen *protogen.Plugin, file *protogen.File, g *protogen.GeneratedFile, method *protogen.Method) string {
 | 
						|
	service := method.Parent
 | 
						|
	hname := fmt.Sprintf("_%s_%s_Handler", service.GoName, method.GoName)
 | 
						|
 | 
						|
	if !method.Desc.IsStreamingClient() && !method.Desc.IsStreamingServer() {
 | 
						|
		g.P("func ", hname, "(srv interface{}, ctx ", contextPackage.Ident("Context"), ", dec func(interface{}) error, interceptor ", grpcPackage.Ident("UnaryServerInterceptor"), ") (interface{}, error) {")
 | 
						|
		g.P("in := new(", method.Input.GoIdent, ")")
 | 
						|
		g.P("if err := dec(in); err != nil { return nil, err }")
 | 
						|
		g.P("if interceptor == nil { return srv.(", service.GoName, "Server).", method.GoName, "(ctx, in) }")
 | 
						|
		g.P("info := &", grpcPackage.Ident("UnaryServerInfo"), "{")
 | 
						|
		g.P("Server: srv,")
 | 
						|
		g.P("FullMethod: ", strconv.Quote(fmt.Sprintf("/%s/%s", service.Desc.FullName(), method.GoName)), ",")
 | 
						|
		g.P("}")
 | 
						|
		g.P("handler := func(ctx ", contextPackage.Ident("Context"), ", req interface{}) (interface{}, error) {")
 | 
						|
		g.P("return srv.(", service.GoName, "Server).", method.GoName, "(ctx, req.(*", method.Input.GoIdent, "))")
 | 
						|
		g.P("}")
 | 
						|
		g.P("return interceptor(ctx, in, info, handler)")
 | 
						|
		g.P("}")
 | 
						|
		g.P()
 | 
						|
		return hname
 | 
						|
	}
 | 
						|
	streamType := unexport(service.GoName) + method.GoName + "Server"
 | 
						|
	g.P("func ", hname, "(srv interface{}, stream ", grpcPackage.Ident("ServerStream"), ") error {")
 | 
						|
	if !method.Desc.IsStreamingClient() {
 | 
						|
		g.P("m := new(", method.Input.GoIdent, ")")
 | 
						|
		g.P("if err := stream.RecvMsg(m); err != nil { return err }")
 | 
						|
		g.P("return srv.(", service.GoName, "Server).", method.GoName, "(m, &", streamType, "{stream})")
 | 
						|
	} else {
 | 
						|
		g.P("return srv.(", service.GoName, "Server).", method.GoName, "(&", streamType, "{stream})")
 | 
						|
	}
 | 
						|
	g.P("}")
 | 
						|
	g.P()
 | 
						|
 | 
						|
	genSend := method.Desc.IsStreamingServer()
 | 
						|
	genSendAndClose := !method.Desc.IsStreamingServer()
 | 
						|
	genRecv := method.Desc.IsStreamingClient()
 | 
						|
 | 
						|
	// Stream auxiliary types and methods.
 | 
						|
	g.P("type ", service.GoName, "_", method.GoName, "Server interface {")
 | 
						|
	if genSend {
 | 
						|
		g.P("Send(*", method.Output.GoIdent, ") error")
 | 
						|
	}
 | 
						|
	if genSendAndClose {
 | 
						|
		g.P("SendAndClose(*", method.Output.GoIdent, ") error")
 | 
						|
	}
 | 
						|
	if genRecv {
 | 
						|
		g.P("Recv() (*", method.Input.GoIdent, ", error)")
 | 
						|
	}
 | 
						|
	g.P(grpcPackage.Ident("ServerStream"))
 | 
						|
	g.P("}")
 | 
						|
	g.P()
 | 
						|
 | 
						|
	g.P("type ", streamType, " struct {")
 | 
						|
	g.P(grpcPackage.Ident("ServerStream"))
 | 
						|
	g.P("}")
 | 
						|
	g.P()
 | 
						|
 | 
						|
	if genSend {
 | 
						|
		g.P("func (x *", streamType, ") Send(m *", method.Output.GoIdent, ") error {")
 | 
						|
		g.P("return x.ServerStream.SendMsg(m)")
 | 
						|
		g.P("}")
 | 
						|
		g.P()
 | 
						|
	}
 | 
						|
	if genSendAndClose {
 | 
						|
		g.P("func (x *", streamType, ") SendAndClose(m *", method.Output.GoIdent, ") error {")
 | 
						|
		g.P("return x.ServerStream.SendMsg(m)")
 | 
						|
		g.P("}")
 | 
						|
		g.P()
 | 
						|
	}
 | 
						|
	if genRecv {
 | 
						|
		g.P("func (x *", streamType, ") Recv() (*", method.Input.GoIdent, ", error) {")
 | 
						|
		g.P("m := new(", method.Input.GoIdent, ")")
 | 
						|
		g.P("if err := x.ServerStream.RecvMsg(m); err != nil { return nil, err }")
 | 
						|
		g.P("return m, nil")
 | 
						|
		g.P("}")
 | 
						|
		g.P()
 | 
						|
	}
 | 
						|
 | 
						|
	return hname
 | 
						|
}
 | 
						|
 | 
						|
const deprecationComment = "// Deprecated: Do not use."
 | 
						|
 | 
						|
func unexport(s string) string { return strings.ToLower(s[:1]) + s[1:] }
 |