func (g *authenticatedWrapperGen) genServerStreamingMethod(s *descriptor.ServiceDescriptorProto, m *descriptor.MethodDescriptorProto) {
g.gen.P(sigPrefix(s, m) + "r *" + getInputTypeName(m) + ", stream " + s.GetName() + "_" + m.GetName() + "Server) error {")
authIntf, err := proto.GetExtension(m.Options, plugin.E_TlsAuthorization)
if err != nil {
g.gen.P(`
panic("no authorization information in protobuf")`)
g.gen.P(`}`)
return
}
auth := authIntf.(*plugin.TLSAuthorization)
if auth.Insecure != nil && *auth.Insecure {
if len(auth.Roles) != 0 {
panic("Roles and Insecure cannot both be specified")
}
g.gen.P(`
return p.local.` + m.GetName() + `(r, stream)`)
g.gen.P(`}`)
return
}
g.gen.P(`
if err := p.authorize(stream.Context(),` + genRoles(auth) + `); err != nil {
return err
}
return p.local.` + m.GetName() + `(r, stream)`)
g.gen.P("}")
}
func (p *html) getInputType(method *descriptor.MethodDescriptorProto) *descriptor.DescriptorProto {
fileDescriptorSet := p.AllFiles()
inputs := strings.Split(method.GetInputType(), ".")
packageName := inputs[1]
messageName := inputs[2]
msg := fileDescriptorSet.GetMessage(packageName, messageName)
if msg == nil {
p.Fail("could not find message ", method.GetInputType())
}
return msg
}
func (g *authenticatedWrapperGen) genAuthenticatedMethod(s *descriptor.ServiceDescriptorProto, m *descriptor.MethodDescriptorProto) {
g.gen.P()
switch {
case m.GetClientStreaming():
g.genClientServerStreamingMethod(s, m)
case m.GetServerStreaming():
g.genServerStreamingMethod(s, m)
default:
g.genSimpleMethod(s, m)
}
g.gen.P()
}
func (g *raftProxyGen) genSimpleMethod(s *descriptor.ServiceDescriptorProto, m *descriptor.MethodDescriptorProto) {
g.gen.P(sigPrefix(s, m) + "ctx context.Context, r *" + getInputTypeName(m) + ") (*" + getOutputTypeName(m) + ", error) {")
g.gen.P(`
conn, err := p.connSelector.LeaderConn(ctx)
if err != nil {
if err == raftselector.ErrIsLeader {
return p.local.` + m.GetName() + `(ctx, r)
}
return nil, err
}
modCtx, err := p.runCtxMods(ctx)
if err != nil {
return nil, err
}`)
g.gen.P(`
resp, err := New` + s.GetName() + `Client(conn).` + m.GetName() + `(modCtx, r)
if err != nil {
if !strings.Contains(err.Error(), "is closing") && !strings.Contains(err.Error(), "the connection is unavailable") && !strings.Contains(err.Error(), "connection error") {
return resp, err
}
conn, err := p.pollNewLeaderConn(ctx)
if err != nil {
if err == raftselector.ErrIsLeader {
return p.local.` + m.GetName() + `(ctx, r)
}
return nil, err
}
return New` + s.GetName() + `Client(conn).` + m.GetName() + `(modCtx, r)
}`)
g.gen.P("return resp, err")
g.gen.P("}")
}
func (g *raftProxyGen) genSimpleMethod(s *descriptor.ServiceDescriptorProto, m *descriptor.MethodDescriptorProto) {
g.gen.P(sigPrefix(s, m) + "ctx context.Context, r *" + getInputTypeName(m) + ") (*" + getOutputTypeName(m) + ", error) {")
g.gen.P(`
conn, err := p.connSelector.LeaderConn(ctx)
if err != nil {
if err == raftselector.ErrIsLeader {
return p.local.` + m.GetName() + `(ctx, r)
}
return nil, err
}
ctx, err = p.runCtxMods(ctx)
if err != nil {
return nil, err
}`)
g.gen.P("return New" + s.GetName() + "Client(conn)." + m.GetName() + "(ctx, r)")
g.gen.P("}")
}
func (p *html) generateFormFunc(servName string, method *descriptor.MethodDescriptorProto) {
inputs := strings.Split(method.GetInputType(), ".")
packageName := inputs[1]
messageName := inputs[2]
s := `<div class="container"><div class="jumbotron">
<h3>` + servName + `: ` + method.GetName() + `</h3>
` + form.Create(method.GetName(), packageName, messageName, p.Generator) + `
</div>`
p.P(`var Form`, servName, "_", method.GetName(), " string = `", s, "`")
}
func (g *raftProxyGen) genProxyMethod(s *descriptor.ServiceDescriptorProto, m *descriptor.MethodDescriptorProto) {
g.gen.P()
switch {
case m.GetServerStreaming() && m.GetClientStreaming():
g.genClientServerStreamingMethod(s, m)
case m.GetServerStreaming():
g.genServerStreamingMethod(s, m)
case m.GetClientStreaming():
g.genClientStreamingMethod(s, m)
default:
g.genSimpleMethod(s, m)
}
g.gen.P()
}
func (g *raftProxyGen) genServerStreamingMethod(s *descriptor.ServiceDescriptorProto, m *descriptor.MethodDescriptorProto) {
g.gen.P(sigPrefix(s, m) + "r *" + getInputTypeName(m) + ", stream " + s.GetName() + "_" + m.GetName() + "Server) error {")
g.gen.P(`
ctx := stream.Context()
conn, err := p.connSelector.LeaderConn(ctx)
if err != nil {
if err == raftselector.ErrIsLeader {
return p.local.` + m.GetName() + `(r, stream)
}
return err
}
ctx, err = p.runCtxMods(ctx)
if err != nil {
return err
}`)
g.gen.P("clientStream, err := New" + s.GetName() + "Client(conn)." + m.GetName() + "(ctx, r)")
g.gen.P(`
if err != nil {
return err
}`)
g.gen.P(`
for {
msg, err := clientStream.Recv()
if err == io.EOF {
break
}
if err != nil {
return err
}
if err := stream.Send(msg); err != nil {
return err
}
}
return nil`)
g.gen.P("}")
}
func (g *raftProxyGen) genClientStreamingMethod(s *descriptor.ServiceDescriptorProto, m *descriptor.MethodDescriptorProto) {
streamType := s.GetName() + "_" + m.GetName() + "Server"
// Generate stream wrapper that returns a modified context
g.genStreamWrapper(streamType)
g.gen.P(sigPrefix(s, m) + "stream " + streamType + `) error {
ctx := stream.Context()
conn, err := p.connSelector.LeaderConn(ctx)
if err != nil {
if err == raftselector.ErrIsLeader {
ctx, err = p.runCtxMods(ctx, p.localCtxMods)
if err != nil {
return err
}
streamWrapper := ` + streamType + `Wrapper{
` + streamType + `: stream,
ctx: ctx,
}
return p.local.` + m.GetName() + `(streamWrapper)
}
return err
}
ctx, err = p.runCtxMods(ctx, p.remoteCtxMods)
if err != nil {
return err
}`)
g.gen.P("clientStream, err := New" + s.GetName() + "Client(conn)." + m.GetName() + "(ctx)")
g.gen.P(`
if err != nil {
return err
}`)
g.gen.P(`
for {
msg, err := stream.Recv()
if err == io.EOF {
break
}
if err != nil {
return err
}
if err := clientStream.Send(msg); err != nil {
return err
}
}
reply, err := clientStream.CloseAndRecv()
if err != nil {
return err
}
return stream.SendAndClose(reply)`)
g.gen.P("}")
}
func (g *raftProxyGen) genClientServerStreamingMethod(s *descriptor.ServiceDescriptorProto, m *descriptor.MethodDescriptorProto) {
g.gen.P(sigPrefix(s, m) + "stream " + s.GetName() + "_" + m.GetName() + "Server) error {")
g.gen.P(`
ctx := stream.Context()
conn, err := p.connSelector.LeaderConn(ctx)
if err != nil {
if err == raftselector.ErrIsLeader {
return p.local.` + m.GetName() + `(stream)
}
return err
}
ctx, err = p.runCtxMods(ctx)
if err != nil {
return err
}`)
g.gen.P("clientStream, err := New" + s.GetName() + "Client(conn)." + m.GetName() + "(ctx)")
g.gen.P(`
if err != nil {
return err
}`)
g.gen.P(`errc := make(chan error, 1)
go func() {
msg, err := stream.Recv()
if err == io.EOF {
close(errc)
return
}
if err != nil {
errc <- err
return
}
if err := clientStream.Send(msg); err != nil {
errc <- err
return
}
}()`)
g.gen.P(`
for {
msg, err := clientStream.Recv()
if err == io.EOF {
break
}
if err != nil {
return err
}
if err := stream.Send(msg); err != nil {
return err
}
}
clientStream.CloseSend()
return <-errc`)
g.gen.P("}")
}
func getOutputTypeName(m *descriptor.MethodDescriptorProto) string {
parts := strings.Split(m.GetOutputType(), ".")
return parts[len(parts)-1]
}
func (g *grpc) generateServerMethod(servName, fullServName string, method *pb.MethodDescriptorProto) string {
methName := generator.CamelCase(method.GetName())
hname := fmt.Sprintf("_%s_%s_Handler", servName, methName)
inType := g.typeName(method.GetInputType())
outType := g.typeName(method.GetOutputType())
if !method.GetServerStreaming() && !method.GetClientStreaming() {
g.P("func ", hname, "(srv interface{}, ctx ", contextPkg, ".Context, dec func(interface{}) error, interceptor ", grpcPkg, ".UnaryServerInterceptor) (interface{}, error) {")
g.P("in := new(", inType, ")")
g.P("if err := dec(in); err != nil { return nil, err }")
g.P("if interceptor == nil { return srv.(", servName, "Server).", methName, "(ctx, in) }")
g.P("info := &", grpcPkg, ".UnaryServerInfo{")
g.P("Server: srv,")
g.P("FullMethod: ", strconv.Quote(fmt.Sprintf("/%s/%s", fullServName, methName)), ",")
g.P("}")
g.P("handler := func(ctx ", contextPkg, ".Context, req interface{}) (interface{}, error) {")
g.P("return srv.(", servName, "Server).", methName, "(ctx, req.(*", inType, "))")
g.P("}")
g.P("return interceptor(ctx, in, info, handler)")
g.P("}")
g.P()
return hname
}
streamType := unexport(servName) + methName + "Server"
g.P("func ", hname, "(srv interface{}, stream ", grpcPkg, ".ServerStream) error {")
if !method.GetClientStreaming() {
g.P("m := new(", inType, ")")
g.P("if err := stream.RecvMsg(m); err != nil { return err }")
g.P("return srv.(", servName, "Server).", methName, "(m, &", streamType, "{stream})")
} else {
g.P("return srv.(", servName, "Server).", methName, "(&", streamType, "{stream})")
}
g.P("}")
g.P()
genSend := method.GetServerStreaming()
genSendAndClose := !method.GetServerStreaming()
genRecv := method.GetClientStreaming()
// Stream auxiliary types and methods.
g.P("type ", servName, "_", methName, "Server interface {")
if genSend {
g.P("Send(*", outType, ") error")
}
if genSendAndClose {
g.P("SendAndClose(*", outType, ") error")
}
if genRecv {
g.P("Recv() (*", inType, ", error)")
}
g.P(grpcPkg, ".ServerStream")
g.P("}")
g.P()
g.P("type ", streamType, " struct {")
g.P(grpcPkg, ".ServerStream")
g.P("}")
g.P()
if genSend {
g.P("func (x *", streamType, ") Send(m *", outType, ") error {")
g.P("return x.ServerStream.SendMsg(m)")
g.P("}")
g.P()
}
if genSendAndClose {
g.P("func (x *", streamType, ") SendAndClose(m *", outType, ") error {")
g.P("return x.ServerStream.SendMsg(m)")
g.P("}")
g.P()
}
if genRecv {
g.P("func (x *", streamType, ") Recv() (*", inType, ", error) {")
g.P("m := new(", inType, ")")
g.P("if err := x.ServerStream.RecvMsg(m); err != nil { return nil, err }")
g.P("return m, nil")
g.P("}")
g.P()
}
return hname
}
// generateServerSignature returns the server-side signature for a method.
func (g *grpc) generateServerSignature(servName string, method *pb.MethodDescriptorProto) string {
origMethName := method.GetName()
methName := generator.CamelCase(origMethName)
if reservedClientName[methName] {
methName += "_"
}
var reqArgs []string
ret := "error"
if !method.GetServerStreaming() && !method.GetClientStreaming() {
reqArgs = append(reqArgs, contextPkg+".Context")
ret = "(*" + g.typeName(method.GetOutputType()) + ", error)"
}
if !method.GetClientStreaming() {
reqArgs = append(reqArgs, "*"+g.typeName(method.GetInputType()))
}
if method.GetServerStreaming() || method.GetClientStreaming() {
reqArgs = append(reqArgs, servName+"_"+generator.CamelCase(origMethName)+"Server")
}
return methName + "(" + strings.Join(reqArgs, ", ") + ") " + ret
}
func (g *grpc) generateClientMethod(servName, fullServName, serviceDescVar string, method *pb.MethodDescriptorProto, descExpr string) {
sname := fmt.Sprintf("/%s/%s", fullServName, method.GetName())
methName := generator.CamelCase(method.GetName())
inType := g.typeName(method.GetInputType())
outType := g.typeName(method.GetOutputType())
g.P("func (c *", unexport(servName), "Client) ", g.generateClientSignature(servName, method), "{")
if !method.GetServerStreaming() && !method.GetClientStreaming() {
g.P("out := new(", outType, ")")
// TODO: Pass descExpr to Invoke.
g.P("err := ", grpcPkg, `.Invoke(ctx, "`, sname, `", in, out, c.cc, opts...)`)
g.P("if err != nil { return nil, err }")
g.P("return out, nil")
g.P("}")
g.P()
return
}
streamType := unexport(servName) + methName + "Client"
g.P("stream, err := ", grpcPkg, ".NewClientStream(ctx, ", descExpr, `, c.cc, "`, sname, `", opts...)`)
g.P("if err != nil { return nil, err }")
g.P("x := &", streamType, "{stream}")
if !method.GetClientStreaming() {
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.GetClientStreaming()
genRecv := method.GetServerStreaming()
genCloseAndRecv := !method.GetServerStreaming()
// Stream auxiliary types and methods.
g.P("type ", servName, "_", methName, "Client interface {")
if genSend {
g.P("Send(*", inType, ") error")
}
if genRecv {
g.P("Recv() (*", outType, ", error)")
}
if genCloseAndRecv {
g.P("CloseAndRecv() (*", outType, ", error)")
}
g.P(grpcPkg, ".ClientStream")
g.P("}")
g.P()
g.P("type ", streamType, " struct {")
g.P(grpcPkg, ".ClientStream")
g.P("}")
g.P()
if genSend {
g.P("func (x *", streamType, ") Send(m *", inType, ") error {")
g.P("return x.ClientStream.SendMsg(m)")
g.P("}")
g.P()
}
if genRecv {
g.P("func (x *", streamType, ") Recv() (*", outType, ", error) {")
g.P("m := new(", outType, ")")
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() (*", outType, ", error) {")
g.P("if err := x.ClientStream.CloseSend(); err != nil { return nil, err }")
g.P("m := new(", outType, ")")
g.P("if err := x.ClientStream.RecvMsg(m); err != nil { return nil, err }")
g.P("return m, nil")
g.P("}")
g.P()
}
}
// generateClientSignature returns the client-side signature for a method.
func (g *grpc) generateClientSignature(servName string, method *pb.MethodDescriptorProto) string {
origMethName := method.GetName()
methName := generator.CamelCase(origMethName)
if reservedClientName[methName] {
methName += "_"
}
reqArg := ", in *" + g.typeName(method.GetInputType())
if method.GetClientStreaming() {
reqArg = ""
}
respName := "*" + g.typeName(method.GetOutputType())
if method.GetServerStreaming() || method.GetClientStreaming() {
respName = servName + "_" + generator.CamelCase(origMethName) + "Client"
}
return fmt.Sprintf("%s(ctx %s.Context%s, opts ...%s.CallOption) (%s, error)", methName, contextPkg, reqArg, grpcPkg, respName)
}
func (g *grpc) generateServerMethod(servName string, method *pb.MethodDescriptorProto) string {
methName := generator.CamelCase(method.GetName())
hname := fmt.Sprintf("_%s_%s_Handler", servName, methName)
inType := g.typeName(method.GetInputType())
outType := g.typeName(method.GetOutputType())
if !method.GetServerStreaming() && !method.GetClientStreaming() {
g.P("func ", hname, "(srv interface{}, ctx ", contextPkg, ".Context, codec ", grpcPkg, ".Codec, buf []byte) (interface{}, error) {")
g.P("in := new(", inType, ")")
g.P("if err := codec.Unmarshal(buf, in); err != nil { return nil, err }")
g.P("out, err := srv.(", servName, "Server).", methName, "(ctx, in)")
g.P("if err != nil { return nil, err }")
g.P("return out, nil")
g.P("}")
g.P()
return hname
}
streamType := unexport(servName) + methName + "Server"
g.P("func ", hname, "(srv interface{}, stream ", grpcPkg, ".ServerStream) error {")
if !method.GetClientStreaming() {
g.P("m := new(", inType, ")")
g.P("if err := stream.RecvMsg(m); err != nil { return err }")
g.P("return srv.(", servName, "Server).", methName, "(m, &", streamType, "{stream})")
} else {
g.P("return srv.(", servName, "Server).", methName, "(&", streamType, "{stream})")
}
g.P("}")
g.P()
genSend := method.GetServerStreaming()
genSendAndClose := !method.GetServerStreaming()
genRecv := method.GetClientStreaming()
// Stream auxiliary types and methods.
g.P("type ", servName, "_", methName, "Server interface {")
if genSend {
g.P("Send(*", outType, ") error")
}
if genSendAndClose {
g.P("SendAndClose(*", outType, ") error")
}
if genRecv {
g.P("Recv() (*", inType, ", error)")
}
g.P(grpcPkg, ".ServerStream")
g.P("}")
g.P()
g.P("type ", streamType, " struct {")
g.P(grpcPkg, ".ServerStream")
g.P("}")
g.P()
if genSend {
g.P("func (x *", streamType, ") Send(m *", outType, ") error {")
g.P("return x.ServerStream.SendMsg(m)")
g.P("}")
g.P()
}
if genSendAndClose {
g.P("func (x *", streamType, ") SendAndClose(m *", outType, ") error {")
g.P("return x.ServerStream.SendMsg(m)")
g.P("}")
g.P()
}
if genRecv {
g.P("func (x *", streamType, ") Recv() (*", inType, ", error) {")
g.P("m := new(", inType, ")")
g.P("if err := x.ServerStream.RecvMsg(m); err != nil { return nil, err }")
g.P("return m, nil")
g.P("}")
g.P()
}
return hname
}
func sigPrefix(s *descriptor.ServiceDescriptorProto, m *descriptor.MethodDescriptorProto) string {
return "func (p *" + serviceTypeName(s) + ") " + m.GetName() + "("
}