Multiplex socket to support both http and grpc requests

Author: janekbaraniewski

go.mod

@@ -51,6 +51,7 @@ require (
 	github.com/rhysd/go-github-selfupdate v1.2.3
 	github.com/sirupsen/logrus v1.9.3
 	github.com/skratchdot/open-golang v0.0.0-20200116055534-eef842397966
+	github.com/soheilhy/cmux v0.1.5
 	github.com/spf13/cobra v1.8.1
 	github.com/spf13/pflag v1.0.5
 	github.com/takama/daemon v1.0.0

go.sum

@@ -803,6 +803,7 @@ golang.org/x/net v0.0.0-20190404232315-eb5bcb51f2a3/go.mod h1:t9HGtf8HONx5eT2rtn
 golang.org/x/net v0.0.0-20190620200207-3b0461eec859/go.mod h1:z5CRVTTTmAJ677TzLLGU+0bjPO0LkuOLi4/5GtJWs/s=
 golang.org/x/net v0.0.0-20200226121028-0de0cce0169b/go.mod h1:z5CRVTTTmAJ677TzLLGU+0bjPO0LkuOLi4/5GtJWs/s=
 golang.org/x/net v0.0.0-20201021035429-f5854403a974/go.mod h1:sp8m0HH+o8qH0wwXwYZr8TS3Oi6o0r6Gce1SSxlDquU=
+golang.org/x/net v0.0.0-20201202161906-c7110b5ffcbb/go.mod h1:sp8m0HH+o8qH0wwXwYZr8TS3Oi6o0r6Gce1SSxlDquU=
 golang.org/x/net v0.0.0-20210226172049-e18ecbb05110/go.mod h1:m0MpNAwzfU5UDzcl9v0D8zg8gWTRqZa9RBIspLL5mdg=
 golang.org/x/net v0.0.0-20210316092652-d523dce5a7f4/go.mod h1:RBQZq4jEuRlivfhVLdyRGr576XBO4/greRjx4P4O3yc=
 golang.org/x/net v0.0.0-20210331212208-0fccb6fa2b5c/go.mod h1:p54w0d4576C0XHj96bSt6lcn1PtDYWL6XObtHCRCNQM=

pkg/daemon/workspace/network/network_proxy.go

@@ -5,28 +5,30 @@ import (
 	"errors"
 	"fmt"
 	"net"
+	"net/http"
 	"os"
-	"strconv"
+	"sync"
 	"time"
 
-	"github.com/loft-sh/devpod/pkg/ts"
 	"github.com/loft-sh/log"
 	"github.com/mwitkow/grpc-proxy/proxy"
+	"github.com/soheilhy/cmux"
 	"google.golang.org/grpc"
-	"google.golang.org/grpc/codes"
-	"google.golang.org/grpc/credentials/insecure"
-	"google.golang.org/grpc/metadata"
-	"google.golang.org/grpc/status"
 	"tailscale.com/tsnet"
 )
 
-// NetworkProxyService proxies gRPC requests based on metadata.
+// NetworkProxyService proxies gRPC and HTTP requests over DevPod network.
+// It coordinates the listener, cmux, and underlying servers.
 type NetworkProxyService struct {
-	listener   net.Listener
-	grpcServer *grpc.Server
-	tsServer   *tsnet.Server
-	log        log.Logger
-	socketPath string
+	mainListener net.Listener
+	grpcServer   *grpc.Server
+	httpServer   *http.Server
+	tsServer     *tsnet.Server
+	log          log.Logger
+	socketPath   string
+	mux          cmux.CMux
+	grpcDirector *GrpcDirector
+	httpProxy    *HttpProxyHandler
 }
 
 // NewNetworkProxyService creates a new instance listening on the given unix socket.
@@ -37,126 +39,171 @@ func NewNetworkProxyService(socketPath string, tsServer *tsnet.Server, log log.L
 		return nil, fmt.Errorf("failed to listen on socket %s: %w", socketPath, err)
 	}
 
-	if err := os.Chmod(socketPath, 0770); err != nil {
+	if err := os.Chmod(socketPath, 0777); err != nil {
 		l.Close()
 		return nil, fmt.Errorf("failed to set socket permissions on %s: %w", socketPath, err)
 	}
 
 	log.Infof("NetworkProxyService: network proxy listening on socket %s", socketPath)
+
+	grpcDirector := NewGrpcDirector(tsServer, log)
+	httpProxy := NewHttpProxyHandler(tsServer, log)
+
 	return &NetworkProxyService{
-		listener:   l,
-		tsServer:   tsServer,
-		log:        log,
-		socketPath: socketPath,
+		mainListener: l,
+		tsServer:     tsServer,
+		log:          log,
+		socketPath:   socketPath,
+		grpcDirector: grpcDirector,
+		httpProxy:    httpProxy,
 	}, nil
 }
 
 // Start runs the gRPC reverse proxy server.
 func (s *NetworkProxyService) Start(ctx context.Context) error {
-	director := func(ctx context.Context, fullMethodName string) (context.Context, *grpc.ClientConn, error) {
-		md, ok := metadata.FromIncomingContext(ctx)
-		if !ok {
-			s.log.Warnf("[NetworkProxyService] [gRPC] Director missing incoming metadata for call %q", fullMethodName)
-			return nil, nil, status.Errorf(codes.InvalidArgument, "missing metadata")
-		}
-		mdCopy := md.Copy()
-
-		targetHosts := mdCopy.Get("x-target-host")
-		targetPorts := mdCopy.Get("x-target-port")
-		proxyPorts := mdCopy.Get("x-proxy-port")
-		if len(targetHosts) == 0 || len(targetPorts) == 0 || len(proxyPorts) == 0 {
-			s.log.Errorf("[NetworkProxyService] [gRPC] Director missing x-target-host, x-proxy-port or x-target-port metadata for call %q", fullMethodName)
-			return nil, nil, status.Errorf(codes.InvalidArgument, "missing x-target-host, x-proxy-port or x-target-port metadata")
-		}
+	// Create connection multiplexer
+	s.mux = cmux.New(s.mainListener)
 
-		proxyPort, err := strconv.Atoi(proxyPorts[0])
-		if err != nil {
-			return nil, nil, err
-		}
-		targetAddr := ts.EnsureURL(targetHosts[0], proxyPort)
+	// Matchers
+	grpcL := s.mux.MatchWithWriters(cmux.HTTP2MatchHeaderFieldSendSettings("content-type", "application/grpc"))
+	httpL := s.mux.Match(cmux.Any())
 
-		s.log.Infof("[NetworkProxyService] [gRPC] Proxying call %q to target %s", fullMethodName, targetAddr)
+	// Servers
+	s.grpcServer = grpc.NewServer(
+		grpc.UnknownServiceHandler(proxy.TransparentHandler(s.grpcDirector.DirectorFunc)),
+	)
+	s.httpServer = &http.Server{
+		Handler: s.httpProxy,
+	}
 
-		// Create a custom dialer using the tsnet server.
-		tsDialer := func(ctx context.Context, addr string) (net.Conn, error) {
-			s.log.Debugf("[NetworkProxyService] [gRPC] Dialing target %s via tsnet", addr)
-			conn, err := s.tsServer.Dial(ctx, "tcp", addr)
-			if err != nil {
-				s.log.Errorf("[NetworkProxyService] [gRPC] Failed to dial target %s via tsnet: %v", addr, err)
-				return nil, err
-			}
-			return conn, nil
-		}
+	// Start servers
+	var runWg sync.WaitGroup
+	errChan := make(chan error, 3)
 
-		// Dial the target gRPC server (the second proxy) using the tsnet dialer.
-		conn, err := grpc.DialContext(ctx, targetAddr,
-			grpc.WithContextDialer(tsDialer),
-			grpc.WithTransportCredentials(insecure.NewCredentials()),
-			grpc.WithCodec(proxy.Codec()),
-		)
-		if err != nil {
-			s.log.Errorf("[NetworkProxyService] [gRPC] Failed to dial target backend %s: %v", targetAddr, err)
-			return nil, nil, status.Errorf(codes.Internal, "failed to dial target backend: %v", err)
+	runWg.Add(1)
+	go func() {
+		defer runWg.Done()
+		s.log.Debugf("NetworkProxyService: starting gRPC server...")
+		if err := s.grpcServer.Serve(grpcL); err != nil && !errors.Is(err, grpc.ErrServerStopped) && !errors.Is(err, cmux.ErrListenerClosed) {
+			s.log.Errorf("NetworkProxyService: gRPC server error: %v", err)
+			errChan <- fmt.Errorf("gRPC server error: %w", err)
+		} else {
+			s.log.Debugf("NetworkProxyService: gRPC server stopped.")
 		}
+	}()
 
-		outCtx := metadata.NewOutgoingContext(ctx, mdCopy)
-
-		return outCtx, conn, nil
-	}
-
-	// Create the gRPC server with the transparent proxy handler.
-	s.grpcServer = grpc.NewServer(
-		grpc.UnknownServiceHandler(proxy.TransparentHandler(director)),
-	)
+	runWg.Add(1)
+	go func() {
+		defer runWg.Done()
+		s.log.Debugf("NetworkProxyService: starting HTTP server...")
+		if err := s.httpServer.Serve(httpL); err != nil && !errors.Is(err, http.ErrServerClosed) && !errors.Is(err, cmux.ErrListenerClosed) {
+			s.log.Errorf("NetworkProxyService: HTTP server error: %v", err)
+			errChan <- fmt.Errorf("HTTP server error: %w", err)
+		} else {
+			s.log.Debugf("NetworkProxyService: HTTP server stopped.")
+		}
+	}()
 
-	s.log.Infof("NetworkProxyService: starting gRPC server on %s", s.socketPath)
+	runWg.Add(1)
 	go func() {
-		if err := s.grpcServer.Serve(s.listener); err != nil && !errors.Is(err, net.ErrClosed) {
-			s.log.Errorf("NetworkProxyService: gRPC server error: %v", err)
-		} else if errors.Is(err, net.ErrClosed) {
-			s.log.Infof("NetworkProxyService: gRPC server stopped gracefully.")
+		defer runWg.Done()
+		s.log.Infof("NetworkProxyService: starting server...")
+		err := s.mux.Serve()
+		if err != nil && !errors.Is(err, net.ErrClosed) && !errors.Is(err, cmux.ErrListenerClosed) {
+			s.log.Errorf("NetworkProxyService: server error: %v", err)
+			errChan <- fmt.Errorf("server error: %w", err)
+		} else {
+			s.log.Infof("NetworkProxyService: server stopped.")
 		}
 	}()
 
-	<-ctx.Done()
-	s.log.Infof("NetworkProxyService: context cancelled, shutting down proxy service")
+	s.log.Infof("NetworkProxyService: successfully started listeners on %s", s.socketPath)
+
+	var finalErr error
+	select {
+	case <-ctx.Done():
+		s.log.Infof("NetworkProxyService: context cancelled, shutting down proxy service")
+		finalErr = ctx.Err()
+	case err := <-errChan:
+		s.log.Errorf("NetworkProxyService: server error triggered shutdown: %v", err)
+		finalErr = err
+	}
+
 	s.Stop()
-	return ctx.Err()
+
+	s.log.Debugf("NetworkProxyService: Waiting for servers to exit...")
+	runWg.Wait()
+	s.log.Debugf("NetworkProxyService: All servers exited.")
+
+	return finalErr
 }
 
-// Stop gracefully shuts down the gRPC server and closes the listener.
 func (s *NetworkProxyService) Stop() {
 	s.log.Infof("NetworkProxyService: stopping proxy service...")
-	stopped := make(chan struct{})
+
+	shutdownCtx, cancel := context.WithTimeout(context.Background(), 15*time.Second)
+	defer cancel()
+
+	var shutdownWg sync.WaitGroup
+	shutdownWg.Add(2)
+
 	go func() {
+		defer shutdownWg.Done()
 		if s.grpcServer != nil {
 			s.grpcServer.GracefulStop()
+			s.log.Debugf("NetworkProxyService: gRPC server stopped.")
 		}
-		close(stopped)
 	}()
 
-	// Wait for graceful stop with a timeout
+	go func() {
+		defer shutdownWg.Done()
+		if s.httpServer != nil {
+			if err := s.httpServer.Shutdown(shutdownCtx); err != nil {
+				s.log.Warnf("NetworkProxyService: HTTP server shutdown error: %v", err)
+			} else {
+				s.log.Debugf("NetworkProxyService: HTTP server stopped.")
+			}
+		}
+	}()
+
+	s.log.Infof("NetworkProxyService: waiting for servers to stop...")
+
+	waitDone := make(chan struct{})
+	go func() {
+		defer close(waitDone)
+		shutdownWg.Wait()
+	}()
+
 	select {
-	case <-time.After(10 * time.Second):
-		s.log.Warnf("NetworkProxyService: Graceful stop timed out, forcing listener close.")
-	case <-stopped:
-		s.log.Infof("NetworkProxyService: gRPC server stopped.")
+	case <-waitDone:
+		s.log.Debugf("NetworkProxyService: All server shutdowns completed.")
+	case <-shutdownCtx.Done():
+		s.log.Warnf("NetworkProxyService: Graceful shutdown timed out after waiting for servers.")
 	}
 
-	if s.listener != nil {
-		if err := s.listener.Close(); err != nil {
-			if !errors.Is(err, net.ErrClosed) {
-				s.log.Errorf("NetworkProxyService: error closing listener: %v", err)
+	s.log.Debugf("NetworkProxyService: Listener and socket cleanup.")
+
+	if s.mainListener != nil {
+		s.log.Debugf("NetworkProxyService: Closing main listener...")
+		if err := s.mainListener.Close(); err != nil {
+			if !errors.Is(err, net.ErrClosed) && !errors.Is(err, cmux.ErrListenerClosed) {
+				s.log.Errorf("NetworkProxyService: Error closing main listener: %v", err)
+			} else {
+				s.log.Debugf("NetworkProxyService: Main listener closed.")
 			}
 		} else {
-			s.log.Infof("NetworkProxyService: Listener closed.")
+			s.log.Debugf("NetworkProxyService: Main listener closed successfully.")
 		}
+	} else {
+		s.log.Warnf("NetworkProxyService: Main listener was nil during stop.")
 	}
 
-	if s.socketPath != "" {
-		if err := os.Remove(s.socketPath); err != nil && !errors.Is(err, os.ErrNotExist) {
-			s.log.Warnf("NetworkProxyService: failed to remove socket file %s: %v", s.socketPath, err)
-		}
+	s.log.Debugf("NetworkProxyService: Removing socket file %s", s.socketPath)
+	if err := os.Remove(s.socketPath); err != nil && !errors.Is(err, os.ErrNotExist) {
+		s.log.Warnf("NetworkProxyService: Failed to remove socket file %s: %v", s.socketPath, err)
+	} else if err == nil {
+		s.log.Debugf("NetworkProxyService: Removed socket file %s", s.socketPath)
 	}
-	s.log.Infof("NetworkProxyService: proxy service stopped.")
+
+	s.log.Infof("NetworkProxyService: Proxy service stopped.")
 }