@@ -4,123 +4,128 @@ import (
44 "context"
55 "fmt"
66 "net"
7- "net/http"
8- "net/http/httputil"
9- "net/url"
107 "time"
118
129 "github.com/loft-sh/log"
10+ "github.com/mwitkow/grpc-proxy/proxy"
11+ "google.golang.org/grpc"
12+ "google.golang.org/grpc/codes"
13+ "google.golang.org/grpc/credentials/insecure"
14+ "google.golang.org/grpc/metadata"
15+ "google.golang.org/grpc/status"
1316 "tailscale.com/tsnet"
1417)
1518
1619const (
17- // Listen on this port via tsnet.
18- LocalCredentialsServerPort = 9999 // FIXME - use random prot
19- // Target server: local gRPC server running on port 5555.
20- TargetServer = "http://localhost:5555" // FIXME - get port from request
20+ LocalCredentialsServerPort int = 9999
21+ DefaultTargetHost string = "localhost"
2122)
2223
23- // LocalCredentialsServerProxy acts as a reverse proxy that blindly forwards
24- // all incoming traffic to the local gRPC server on port 5555.
2524type LocalCredentialsServerProxy struct {
26- log log.Logger
27- tsServer * tsnet.Server
28-
29- ln net.Listener
30- srv * http.Server
25+ log log.Logger
26+ tsServer * tsnet.Server
27+ grpcServer * grpc.Server
28+ ln net.Listener
3129}
3230
33- // NewLocalCredentialsServerProxy initializes a new LocalCredentialsServerProxy.
34- func NewLocalCredentialsServerProxy (tsServer * tsnet.Server , log log.Logger ) (* LocalCredentialsServerProxy , error ) {
31+ func NewLocalCredentialsServerProxy (tsServer * tsnet.Server , logger log.Logger ) (* LocalCredentialsServerProxy , error ) {
32+ logger .Infof ("NewLocalCredentialsServerProxy: initializing local reverse proxy" )
33+ if tsServer == nil {
34+ return nil , fmt .Errorf ("tsnet.Server cannot be nil" )
35+ }
3536 return & LocalCredentialsServerProxy {
36- log : log ,
37+ log : logger ,
3738 tsServer : tsServer ,
3839 }, nil
3940}
4041
41- // Listen creates the tsnet listener and HTTP server,
42- // and registers a catch-all handler that acts as the reverse proxy.
4342func (s * LocalCredentialsServerProxy ) Listen (ctx context.Context ) error {
44- s .log .Info ( " Starting reverse proxy for local gRPC server" )
43+ s .log .Infof ( "LocalCredentialsServerProxy: Starting reverse proxy on tsnet port %d" , LocalCredentialsServerPort )
4544
46- // Create a tsnet listener.
47- ln , err := s .tsServer .Listen ("tcp" , fmt . Sprintf ( ":%d" , LocalCredentialsServerPort ) )
45+ listenAddr := fmt . Sprintf ( ":%d" , LocalCredentialsServerPort )
46+ ln , err := s .tsServer .Listen ("tcp" , listenAddr )
4847 if err != nil {
49- s .log .Infof ( " Failed to listen on tsnet port %d : %v" , LocalCredentialsServerPort , err )
50- return fmt .Errorf ("failed to listen on tsnet port %d : %w" , LocalCredentialsServerPort , err )
48+ s .log .Errorf ( "LocalCredentialsServerProxy: Failed to listen on tsnet %s : %v" , listenAddr , err )
49+ return fmt .Errorf ("failed to listen on tsnet %s : %w" , listenAddr , err )
5150 }
5251 s .ln = ln
5352
54- mux := http .NewServeMux ()
55- mux .HandleFunc ("/" , s .handleReverseProxy )
53+ s .log .Infof ("LocalCredentialsServerProxy: tsnet listener started on %s" , ln .Addr ().String ())
5654
57- // Create the HTTP server.
58- s .srv = & http.Server {
59- Handler : mux ,
60- }
55+ director := func (ctx context.Context , fullMethodName string ) (context.Context , * grpc.ClientConn , error ) {
56+ md , ok := metadata .FromIncomingContext (ctx )
57+ if ! ok {
58+ return nil , nil , status .Errorf (codes .InvalidArgument , "missing metadata" )
59+ }
6160
62- go func () {
63- <- ctx .Done ()
64- s .log .Info ("Context canceled, shutting down reverse proxy" )
65- shutdownCtx , cancel := context .WithTimeout (context .Background (), 5 * time .Second )
66- defer cancel ()
67- if err := s .srv .Shutdown (shutdownCtx ); err != nil {
68- s .log .Errorf ("Error shutting down reverse proxy: %v" , err )
61+ // Get the target port from metadata. Host is always localhost.
62+ targetPorts := md .Get ("x-target-port" )
63+ if len (targetPorts ) == 0 {
64+ s .log .Error ("LocalCredentialsServerProxy: Director missing x-target-port metadata" )
65+ return nil , nil , status .Errorf (codes .InvalidArgument , "missing x-target-port metadata" )
6966 }
70- }()
67+ // targetPort := targetPorts[0]
68+ targetPort := "4795" // FIXME
7169
72- s .log .Infof ("Reverse proxy listening on tsnet port %d" , LocalCredentialsServerPort )
73- err = s .srv .Serve (ln )
74- if err != nil && err != http .ErrServerClosed {
75- s .log .Errorf ("Reverse proxy error: %v" , err )
76- return err
77- }
70+ targetAddr := net .JoinHostPort (DefaultTargetHost , targetPort )
7871
79- return nil
80- }
72+ s .log .Infof ("[LocalCredentialsServerProxy] [gRPC] Proxying call %q to target %s" , fullMethodName , targetAddr )
8173
82- // handleReverseProxy forwards every request to the target gRPC server.
83- func (s * LocalCredentialsServerProxy ) handleReverseProxy (w http.ResponseWriter , r * http.Request ) {
84- s .log .Infof ("Forwarding request %s %s to target server" , r .Method , r .URL .String ())
74+ conn , err := grpc .DialContext (ctx , targetAddr ,
75+ grpc .WithTransportCredentials (insecure .NewCredentials ()),
76+ grpc .WithCodec (proxy .Codec ()), // Use proxy codec for transparency
77+ )
78+ if err != nil {
79+ s .log .Errorf ("[LocalCredentialsServerProxy] [gRPC] Failed to dial local target backend %s: %v" , targetAddr , err )
80+ return nil , nil , status .Errorf (codes .Internal , "failed to dial local target backend: %v" , err )
81+ }
8582
86- // Parse the target URL.
87- targetURL , err := url .Parse (TargetServer )
88- if err != nil {
89- s .log .Errorf ("Error parsing target URL %s: %v" , TargetServer , err )
90- http .Error (w , "Bad Gateway" , http .StatusBadGateway )
91- return
83+ return ctx , conn , nil
9284 }
9385
94- // Create the reverse proxy.
95- proxy := httputil .NewSingleHostReverseProxy (targetURL )
86+ // Create the gRPC server using the transparent handler.
87+ // It will forward any unknown service call based on the director logic.
88+ s .grpcServer = grpc .NewServer (
89+ grpc .UnknownServiceHandler (proxy .TransparentHandler (director )),
90+ )
9691
97- // Customize the director to forward the Host header to the target.
98- originalDirector := proxy .Director
99- proxy .Director = func (req * http.Request ) {
100- originalDirector (req )
101- req .Host = targetURL .Host
102- }
92+ s .log .Infof ("LocalCredentialsServerProxy: gRPC reverse proxy configured, starting server on %s" , ln .Addr ().String ())
10393
104- // Use an error handler to log any errors that occur.
105- proxy .ErrorHandler = func (w http.ResponseWriter , r * http.Request , err error ) {
106- s .log .Errorf ("Reverse proxy error: %v" , err )
107- http .Error (w , "Bad Gateway" , http .StatusBadGateway )
94+ if err := s .grpcServer .Serve (s .ln ); err != nil {
95+ if err .Error () != "grpc: the server has been stopped" {
96+ s .log .Errorf ("LocalCredentialsServerProxy: failed to serve: %v" , err )
97+ return fmt .Errorf ("gRPC server error: %w" , err )
98+ } else {
99+ s .log .Infof ("LocalCredentialsServerProxy: gRPC server stopped gracefully." )
100+ }
108101 }
109-
110- // Forward the request.
111- proxy .ServeHTTP (w , r )
102+ return nil
112103}
113104
114- // Close gracefully shuts down the reverse proxy.
115- func (s * LocalCredentialsServerProxy ) Close () error {
116- s .log .Info ("Closing reverse proxy" )
117- if s .srv != nil {
118- shutdownCtx , cancel := context .WithTimeout (context .Background (), 5 * time .Second )
119- defer cancel ()
120- if err := s .srv .Shutdown (shutdownCtx ); err != nil {
121- s .log .Errorf ("Error during reverse proxy shutdown: %v" , err )
122- return err
105+ func (s * LocalCredentialsServerProxy ) Stop () {
106+ s .log .Info ("LocalCredentialsServerProxy: Stopping reverse proxy..." )
107+ if s .grpcServer != nil {
108+ stopped := make (chan struct {})
109+ go func () {
110+ s .grpcServer .GracefulStop ()
111+ close (stopped )
112+ }()
113+
114+ select {
115+ case <- time .After (10 * time .Second ):
116+ s .log .Warnf ("LocalCredentialsServerProxy: Graceful shutdown timed out after 10 seconds, forcing stop." )
117+ s .grpcServer .Stop ()
118+ case <- stopped :
119+ s .log .Infof ("LocalCredentialsServerProxy: gRPC server stopped gracefully." )
123120 }
124121 }
125- return nil
122+
123+ if s .ln != nil {
124+ if err := s .ln .Close (); err != nil {
125+ s .log .Errorf ("LocalCredentialsServerProxy: Error closing listener: %v" , err )
126+ } else {
127+ s .log .Infof ("LocalCredentialsServerProxy: Listener closed." )
128+ }
129+ }
130+ s .log .Info ("LocalCredentialsServerProxy: Reverse proxy stopped." )
126131}
0 commit comments