Split workspace network server into separate services

Author: janekbaraniewski

cmd/agent/git_credentials.go

@@ -14,6 +14,7 @@ import (
 
 	"github.com/loft-sh/devpod/cmd/flags"
 	workspaced "github.com/loft-sh/devpod/pkg/daemon/workspace"
+	"github.com/loft-sh/devpod/pkg/daemon/workspace/network"
 	"github.com/loft-sh/devpod/pkg/gitcredentials"
 	devpodhttp "github.com/loft-sh/devpod/pkg/http"
 	"github.com/loft-sh/log"
@@ -78,15 +79,15 @@ func (cmd *GitCredentialsCmd) Run(ctx context.Context, args []string, log log.Lo
 }
 
 func getCredentialsFromWorkspaceServer(credentials *gitcredentials.GitCredentials) *gitcredentials.GitCredentials {
-	if _, err := os.Stat(filepath.Join(workspaced.RootDir, workspaced.RunnerProxySocket)); err != nil {
+	if _, err := os.Stat(filepath.Join(workspaced.RootDir, network.RunnerProxySocket)); err != nil {
 		// workspace server is not running
 		return nil
 	}
 
 	httpClient := &http.Client{
 		Transport: &http.Transport{
 			DialContext: func(_ context.Context, _, _ string) (net.Conn, error) {
-				return net.Dial("unix", filepath.Join(workspaced.RootDir, workspaced.RunnerProxySocket))
+				return net.Dial("unix", filepath.Join(workspaced.RootDir, network.RunnerProxySocket))
 			},
 		},
 	}

pkg/credentials/server.go

@@ -16,6 +16,7 @@ import (
 	"github.com/loft-sh/devpod/pkg/agent/tunnel"
 	locald "github.com/loft-sh/devpod/pkg/daemon/platform"
 	workspaced "github.com/loft-sh/devpod/pkg/daemon/workspace"
+	network "github.com/loft-sh/devpod/pkg/daemon/workspace/network"
 	devpodlog "github.com/loft-sh/devpod/pkg/log"
 	"github.com/loft-sh/devpod/pkg/ts"
 	"github.com/loft-sh/log"
@@ -159,7 +160,7 @@ func handleGitCredentialsOverTSNet(ctx context.Context, writer http.ResponseWrit
 
 	log.Infof("Received git credentials post data: %s", string(bodyBytes))
 	// Set up HTTP transport that uses our network socket.
-	socketPath := filepath.Join(workspaced.RootDir, workspaced.TSNetProxySocket)
+	socketPath := filepath.Join(workspaced.RootDir, network.TSNetProxySocket)
 	transport := &http.Transport{
 		Dial: func(network, addr string) (net.Conn, error) {
 			return net.Dial("unix", socketPath)

pkg/daemon/workspace/network.go

@@ -7,6 +7,7 @@ import (
 
 	"sync"
 
+	"github.com/loft-sh/devpod/pkg/daemon/workspace/network"
 	"github.com/loft-sh/devpod/pkg/platform/client"
 	"github.com/loft-sh/devpod/pkg/ts"
 	"github.com/loft-sh/log"
@@ -30,7 +31,7 @@ func RunNetworkServer(ctx context.Context, d *Daemon, errChan chan<- error, wg *
 		errChan <- fmt.Errorf("failed to refresh client: %w", err)
 		return
 	}
-	tsServer := NewWorkspaceServer(&WorkspaceServerConfig{
+	networkServer := network.NewWorkspaceServer(&network.WorkspaceServerConfig{
 		AccessKey:     d.Config.Platform.AccessKey,
 		PlatformHost:  ts.RemoveProtocol(d.Config.Platform.PlatformHost),
 		WorkspaceHost: d.Config.Platform.WorkspaceHost,
@@ -40,7 +41,7 @@ func RunNetworkServer(ctx context.Context, d *Daemon, errChan chan<- error, wg *
 			logger.Infof(format, args...)
 		},
 	}, logger)
-	if err := tsServer.Start(ctx); err != nil {
+	if err := networkServer.Start(ctx); err != nil {
 		errChan <- fmt.Errorf("network server: %w", err)
 	}
 }

pkg/daemon/workspace/network/connection_tracker.go

@@ -0,0 +1,27 @@
+package network
+
+import "sync"
+
+// ConnTracker is a simple connection counter used by several services.
+type ConnTracker struct {
+	mu    sync.Mutex
+	count int
+}
+
+func (c *ConnTracker) Add() {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+	c.count++
+}
+
+func (c *ConnTracker) Remove() {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+	c.count--
+}
+
+func (c *ConnTracker) Count() int {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+	return c.count
+}

pkg/daemon/workspace/network/heartbeat.go

@@ -0,0 +1,82 @@
+package network
+
+import (
+	"context"
+	"fmt"
+	"net/http"
+	"time"
+
+	"github.com/loft-sh/log"
+	"tailscale.com/client/tailscale"
+	"tailscale.com/tsnet"
+)
+
+// HeartbeatService sends periodic heartbeats when there are active connections.
+type HeartbeatService struct {
+	tsServer      *tsnet.Server
+	lc            *tailscale.LocalClient
+	config        *WorkspaceServerConfig
+	projectName   string
+	workspaceName string
+	log           log.Logger
+	tracker       *ConnTracker
+}
+
+// NewHeartbeatService creates a new HeartbeatService.
+func NewHeartbeatService(config *WorkspaceServerConfig, tsServer *tsnet.Server, lc *tailscale.LocalClient, projectName, workspaceName string, tracker *ConnTracker, log log.Logger) *HeartbeatService {
+	return &HeartbeatService{
+		tsServer:      tsServer,
+		lc:            lc,
+		config:        config,
+		projectName:   projectName,
+		workspaceName: workspaceName,
+		log:           log,
+		tracker:       tracker,
+	}
+}
+
+// Start begins the heartbeat loop.
+func (s *HeartbeatService) Start(ctx context.Context) {
+	transport := &http.Transport{DialContext: s.tsServer.Dial}
+	client := &http.Client{Transport: transport, Timeout: 10 * time.Second}
+	ticker := time.NewTicker(10 * time.Second)
+	defer ticker.Stop()
+	for {
+		select {
+		case <-ctx.Done():
+			return
+		case <-ticker.C:
+			if s.tracker.Count() > 0 {
+				if err := s.sendHeartbeat(ctx, client); err != nil {
+					s.log.Errorf("HeartbeatService: failed to send heartbeat: %v", err)
+				}
+			}
+		}
+	}
+}
+
+func (s *HeartbeatService) sendHeartbeat(ctx context.Context, client *http.Client) error {
+	hbCtx, cancel := context.WithTimeout(ctx, 10*time.Second)
+	defer cancel()
+	discoveredRunner, err := discoverRunner(hbCtx, s.lc, s.log)
+	if err != nil {
+		return fmt.Errorf("failed to discover runner: %w", err)
+	}
+	heartbeatURL := fmt.Sprintf("http://%s.ts.loft/devpod/%s/%s/heartbeat", discoveredRunner, s.projectName, s.workspaceName)
+	s.log.Infof("HeartbeatService: sending heartbeat to %s, active connections: %d", heartbeatURL, s.tracker.Count())
+	req, err := http.NewRequestWithContext(hbCtx, "GET", heartbeatURL, nil)
+	if err != nil {
+		return fmt.Errorf("failed to create request for %s: %w", heartbeatURL, err)
+	}
+	req.Header.Set("Authorization", "Bearer "+s.config.AccessKey)
+	resp, err := client.Do(req)
+	if err != nil {
+		return fmt.Errorf("request to %s failed: %w", heartbeatURL, err)
+	}
+	defer resp.Body.Close()
+	if resp.StatusCode != http.StatusOK {
+		return fmt.Errorf("received response from %s - Status: %d", heartbeatURL, resp.StatusCode)
+	}
+	s.log.Infof("HeartbeatService: received response from %s - Status: %d", heartbeatURL, resp.StatusCode)
+	return nil
+}

pkg/daemon/workspace/network/netmap.go

@@ -0,0 +1,51 @@
+package network
+
+import (
+	"context"
+	"encoding/json"
+	"os"
+	"path/filepath"
+	"time"
+
+	"github.com/loft-sh/devpod/pkg/ts"
+	"github.com/loft-sh/log"
+	"tailscale.com/client/tailscale"
+	"tailscale.com/types/netmap"
+)
+
+// NetmapWatcherService watches the Tailscale netmap and writes it to a file.
+type NetmapWatcherService struct {
+	rootDir string
+	lc      *tailscale.LocalClient
+	log     log.Logger
+}
+
+// NewNetmapWatcherService creates a new NetmapWatcherService.
+func NewNetmapWatcherService(rootDir string, lc *tailscale.LocalClient, log log.Logger) *NetmapWatcherService {
+	return &NetmapWatcherService{
+		rootDir: rootDir,
+		lc:      lc,
+		log:     log,
+	}
+}
+
+// Start begins watching the netmap.
+func (s *NetmapWatcherService) Start(ctx context.Context) {
+	go func() {
+		lastUpdate := time.Now()
+		if err := ts.WatchNetmap(ctx, s.lc, func(netMap *netmap.NetworkMap) {
+			if time.Since(lastUpdate) < netMapCooldown {
+				return
+			}
+			lastUpdate = time.Now()
+			nm, err := json.Marshal(netMap)
+			if err != nil {
+				s.log.Errorf("NetmapWatcherService: failed to marshal netmap: %v", err)
+			} else {
+				_ = os.WriteFile(filepath.Join(s.rootDir, "netmap.json"), nm, 0644)
+			}
+		}); err != nil {
+			s.log.Errorf("NetmapWatcherService: failed to watch netmap: %v", err)
+		}
+	}()
+}

pkg/daemon/workspace/network/network_proxy.go

@@ -0,0 +1,100 @@
+package network
+
+import (
+	"bufio"
+	"context"
+	"fmt"
+	"io"
+	"net"
+	"net/http"
+	"os"
+	"strings"
+
+	"github.com/loft-sh/log"
+	"tailscale.com/tsnet"
+)
+
+// NetworkProxyService listens on a Unix socket and proxies requests to TSNet peers.
+type NetworkProxyService struct {
+	listener net.Listener
+	tsServer *tsnet.Server
+	log      log.Logger
+}
+
+// NewNetworkProxyService creates a new NetworkProxyService.
+func NewNetworkProxyService(socketPath string, tsServer *tsnet.Server, log log.Logger) (*NetworkProxyService, error) {
+	_ = os.Remove(socketPath)
+	l, err := net.Listen("unix", socketPath)
+	if err != nil {
+		return nil, fmt.Errorf("failed to listen on socket %s: %w", socketPath, err)
+	}
+	if err := os.Chmod(socketPath, 0777); err != nil {
+		log.Errorf("failed to set socket permissions on %s: %v", socketPath, err)
+	}
+	log.Infof("TSProxyServer: network proxy listening on socket %s", socketPath)
+	return &NetworkProxyService{
+		listener: l,
+		tsServer: tsServer,
+		log:      log,
+	}, nil
+}
+
+// Start begins accepting connections on the TS proxy socket.
+func (s *NetworkProxyService) Start(ctx context.Context) {
+	go func() {
+		defer s.listener.Close()
+		for {
+			conn, err := s.listener.Accept()
+			if err != nil {
+				select {
+				case <-ctx.Done():
+					s.log.Infof("TSProxyServer: listener shutting down (context cancelled)")
+					return
+				default:
+					s.log.Errorf("TSProxyServer: error accepting connection: %v", err)
+					continue
+				}
+			}
+			go s.handleConnection(ctx, conn)
+		}
+	}()
+}
+
+func (s *NetworkProxyService) handleConnection(ctx context.Context, conn net.Conn) {
+	defer conn.Close()
+	reader := bufio.NewReader(conn)
+	req, err := http.ReadRequest(reader)
+	if err != nil {
+		s.log.Errorf("TSProxyServer: failed to read HTTP request: %v", err)
+		return
+	}
+	target := req.Host
+	if target == "" {
+		s.log.Errorf("TSProxyServer: HTTP request does not contain a Host header")
+		return
+	}
+	if !strings.Contains(target, ":") {
+		target = target + ":80"
+	}
+	s.log.Infof("TSProxyServer: proxying request to target %s", target)
+	tsConn, err := s.tsServer.Dial(ctx, "tcp", target)
+	if err != nil {
+		s.log.Errorf("TSProxyServer: error dialing target %s: %v", target, err)
+		return
+	}
+	defer tsConn.Close()
+	if err := req.Write(tsConn); err != nil {
+		s.log.Errorf("TSProxyServer: error forwarding request to target %s: %v", target, err)
+		return
+	}
+	if _, err := io.Copy(conn, tsConn); err != nil {
+		s.log.Errorf("TSProxyServer: error forwarding response from target: %v", err)
+	}
+}
+
+// Stop stops the TSProxyServer by closing its listener.
+func (s *NetworkProxyService) Stop() {
+	if s.listener != nil {
+		s.listener.Close()
+	}
+}