github.com/renbou/grpcbridge@v0.0.2-0.20240416012907-bcbd8b12648a/proxy.go (about) 1 package grpcbridge 2 3 import ( 4 "context" 5 6 "github.com/renbou/grpcbridge/bridgelog" 7 "github.com/renbou/grpcbridge/grpcadapter" 8 "github.com/renbou/grpcbridge/routing" 9 "google.golang.org/grpc" 10 "google.golang.org/protobuf/proto" 11 ) 12 13 var _ grpc.StreamHandler = (*GRPCProxy)(nil).StreamHandler 14 15 // ProxyOption configures gRPC proxying options. 16 type ProxyOption interface { 17 applyProxy(*proxyOptions) 18 } 19 20 // GRPCProxy is a basic gRPC proxy implementation which uses a [routing.GRPCRouter] to route incoming gRPC requests 21 // to the proper target services. It should be registered with a [grpc.Server] as a [grpc.UnknownServiceHandler], 22 // which is precisely what the [GRPCProxy.AsServerOption] helper method returns. 23 // This way of routing is used instead of relying on the [grpc.ServiceRegistrar] interface because 24 // GRPCProxy supports dynamic routing, which is not possible with a classic gRPC service, since it expects 25 // all gRPC services to be registered before launch. 26 type GRPCProxy struct { 27 logger bridgelog.Logger 28 router routing.GRPCRouter 29 } 30 31 // NewGRPCProxy constructs a new [*GRPCProxy] with the given router and options. When no options are provided, sane defaults are used. 32 func NewGRPCProxy(router routing.GRPCRouter, opts ...ProxyOption) *GRPCProxy { 33 options := defaultProxyOptions() 34 35 for _, opt := range opts { 36 opt.applyProxy(&options) 37 } 38 39 return &GRPCProxy{ 40 logger: options.common.logger.WithComponent("grpcbridge.proxy"), 41 router: router, 42 } 43 } 44 45 // AsServerOption returns a [grpc.ServerOption] which can be used to register this proxy with 46 // a gRPC server as the handler to be used for all unknown services. 47 func (s *GRPCProxy) AsServerOption() grpc.ServerOption { 48 return grpc.UnknownServiceHandler(s.StreamHandler) 49 } 50 51 // StreamHandler allows proxying any incoming gRPC streams. 52 // It uses the router's RouteGRPC method with the incoming context to get a connection to the target 53 // and the description of the method. The actual forwarding is performed using the generic [grpcadapter.ForwardServerToClient] function. 54 func (s *GRPCProxy) StreamHandler(_ any, incoming grpc.ServerStream) error { 55 conn, route, err := s.router.RouteGRPC(incoming.Context()) 56 if err != nil { 57 return err 58 } 59 60 logger := s.logger.With("target", route.Target.Name, "method", route.Method.RPCName) 61 62 // TODO(renbou): timeouts for stream initiation and Recv/Sends 63 outgoing, err := conn.Stream(incoming.Context(), route.Method.RPCName) 64 if err != nil { 65 return err 66 } 67 68 // Always clean up the outgoing stream by explicitly closing it. 69 defer outgoing.Close() 70 71 logger.Debug("began proxying gRPC stream") 72 defer logger.Debug("ended proxying gRPC stream") 73 74 return grpcadapter.ForwardServerToClient(incoming.Context(), grpcadapter.ForwardS2C{ 75 Method: route.Method, 76 Incoming: grpcServerStream{incoming}, 77 Outgoing: outgoing, 78 }) 79 } 80 81 type grpcServerStream struct { 82 grpc.ServerStream 83 } 84 85 func (s grpcServerStream) Recv(_ context.Context, msg proto.Message) error { 86 return s.ServerStream.RecvMsg(msg) 87 } 88 89 func (s grpcServerStream) Send(_ context.Context, msg proto.Message) error { 90 return s.ServerStream.SendMsg(msg) 91 } 92 93 type proxyOptions struct { 94 common options 95 } 96 97 func defaultProxyOptions() proxyOptions { 98 return proxyOptions{ 99 common: defaultOptions(), 100 } 101 }