Change local credentials server into credentials proxy

Author: janekbaraniewski

pkg/daemon/local/credentials_proxy.go

@@ -0,0 +1,126 @@
+package local
+
+import (
+	"context"
+	"fmt"
+	"net"
+	"net/http"
+	"net/http/httputil"
+	"net/url"
+	"time"
+
+	"github.com/loft-sh/log"
+	"tailscale.com/tsnet"
+)
+
+const (
+	// Listen on this port via tsnet.
+	LocalCredentialsServerPort = 9999 // FIXME - use random prot
+	// Target server: local gRPC server running on port 5555.
+	TargetServer = "http://localhost:5555" // FIXME - get port from request
+)
+
+// LocalCredentialsServerProxy acts as a reverse proxy that blindly forwards
+// all incoming traffic to the local gRPC server on port 5555.
+type LocalCredentialsServerProxy struct {
+	log      log.Logger
+	tsServer *tsnet.Server
+
+	ln  net.Listener
+	srv *http.Server
+}
+
+// NewLocalCredentialsServerProxy initializes a new LocalCredentialsServerProxy.
+func NewLocalCredentialsServerProxy(tsServer *tsnet.Server, log log.Logger) (*LocalCredentialsServerProxy, error) {
+	return &LocalCredentialsServerProxy{
+		log:      log,
+		tsServer: tsServer,
+	}, nil
+}
+
+// Listen creates the tsnet listener and HTTP server,
+// and registers a catch-all handler that acts as the reverse proxy.
+func (s *LocalCredentialsServerProxy) Listen(ctx context.Context) error {
+	s.log.Info("Starting reverse proxy for local gRPC server")
+
+	// Create a tsnet listener.
+	ln, err := s.tsServer.Listen("tcp", fmt.Sprintf(":%d", LocalCredentialsServerPort))
+	if err != nil {
+		s.log.Infof("Failed to listen on tsnet port %d: %v", LocalCredentialsServerPort, err)
+		return fmt.Errorf("failed to listen on tsnet port %d: %w", LocalCredentialsServerPort, err)
+	}
+	s.ln = ln
+
+	mux := http.NewServeMux()
+	mux.HandleFunc("/", s.handleReverseProxy)
+
+	// Create the HTTP server.
+	s.srv = &http.Server{
+		Handler: mux,
+	}
+
+	go func() {
+		<-ctx.Done()
+		s.log.Info("Context canceled, shutting down reverse proxy")
+		shutdownCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+		defer cancel()
+		if err := s.srv.Shutdown(shutdownCtx); err != nil {
+			s.log.Errorf("Error shutting down reverse proxy: %v", err)
+		}
+	}()
+
+	s.log.Infof("Reverse proxy listening on tsnet port %d", LocalCredentialsServerPort)
+	err = s.srv.Serve(ln)
+	if err != nil && err != http.ErrServerClosed {
+		s.log.Errorf("Reverse proxy error: %v", err)
+		return err
+	}
+
+	return nil
+}
+
+// handleReverseProxy forwards every request to the target gRPC server.
+func (s *LocalCredentialsServerProxy) handleReverseProxy(w http.ResponseWriter, r *http.Request) {
+	s.log.Infof("Forwarding request %s %s to target server", r.Method, r.URL.String())
+
+	// Parse the target URL.
+	targetURL, err := url.Parse(TargetServer)
+	if err != nil {
+		s.log.Errorf("Error parsing target URL %s: %v", TargetServer, err)
+		http.Error(w, "Bad Gateway", http.StatusBadGateway)
+		return
+	}
+
+	// Create the reverse proxy.
+	proxy := httputil.NewSingleHostReverseProxy(targetURL)
+
+	// Customize the director to forward the Host header to the target.
+	originalDirector := proxy.Director
+	proxy.Director = func(req *http.Request) {
+		originalDirector(req)
+		req.Host = targetURL.Host
+	}
+
+	// Use an error handler to log any errors that occur.
+	proxy.ErrorHandler = func(w http.ResponseWriter, r *http.Request, err error) {
+		s.log.Errorf("Reverse proxy error: %v", err)
+		http.Error(w, "Bad Gateway", http.StatusBadGateway)
+	}
+
+	// Forward the request.
+	proxy.ServeHTTP(w, r)
+}
+
+// Close gracefully shuts down the reverse proxy.
+func (s *LocalCredentialsServerProxy) Close() error {
+	s.log.Info("Closing reverse proxy")
+	if s.srv != nil {
+		shutdownCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+		defer cancel()
+		if err := s.srv.Shutdown(shutdownCtx); err != nil {
+			s.log.Errorf("Error during reverse proxy shutdown: %v", err)
+			return err
+		}
+	}
+	return nil
+}

pkg/daemon/local/credentials_server.go

@@ -23,12 +23,12 @@ import (
 )
 
 type Daemon struct {
-	socketListener    net.Listener
-	tsServer          *tsnet.Server
-	localServer       *localServer
-	credentialsServer *LocalCredentialsServer
-	rootDir           string
-	log               log.Logger
+	socketListener         net.Listener
+	tsServer               *tsnet.Server
+	localServer            *localServer
+	credentialsServerProxy *LocalCredentialsServerProxy
+	rootDir                string
+	log                    log.Logger
 }
 
 type InitConfig struct {
@@ -63,18 +63,18 @@ func Init(ctx context.Context, config InitConfig) (*Daemon, error) {
 		return nil, fmt.Errorf("create local server: %w", err)
 	}
 
-	credentialsServer, err := NewLocalCredentialsServer(tsServer, log)
+	credentialsProxy, err := NewLocalCredentialsServerProxy(tsServer, log)
 	if err != nil {
 		return nil, fmt.Errorf("create local credentials server: %w", err)
 	}
 
 	return &Daemon{
-		socketListener:    socketListener,
-		tsServer:          tsServer,
-		localServer:       localServer,
-		credentialsServer: credentialsServer,
-		rootDir:           config.RootDir,
-		log:               log,
+		socketListener:         socketListener,
+		tsServer:               tsServer,
+		localServer:            localServer,
+		credentialsServerProxy: credentialsProxy,
+		rootDir:                config.RootDir,
+		log:                    log,
 	}, nil
 }
 func (d *Daemon) Start(ctx context.Context) error {
@@ -94,7 +94,7 @@ func (d *Daemon) Start(ctx context.Context) error {
 	}()
 	go func() {
 		d.log.Info("Start credentials server")
-		errChan <- d.credentialsServer.Listen(ctx)
+		errChan <- d.credentialsServerProxy.Listen(ctx)
 	}()
 	defer func() {
 		d.log.Info("Cleaning up daemon resources")