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  }