fix: use fake local network for port-forward tests (#11119)

spikecurtis · web-flow · commit 50575e1a9a6a · 2023-12-11T14:51:56.000+04:00
Fixes #10979 Testing code that listens on a specific port has created a long battle with flakes. Previous attempts to deal with this include opening a listener on a port chosen by the OS, then closing the listener, noting the port and starting the test with that port. This still flakes, notably in macOS which has a proclivity to reuse ports quickly. Instead of fighting with the chaos that is an OS networking stack, this PR fakes the host networking in tests. I've taken a small step here, only faking out the Listen() calls that port-forward makes, but I think over time we should be transitioning all networking the CLI does to an abstract interface so we can fake it. This allows us to run in parallel without flakes and presents an opportunity to test error paths as well.
diff --git a/cli/clibase/cmd.go b/cli/clibase/cmd.go
@@ -189,6 +189,7 @@ type Invocation struct {
 	Stderr  io.Writer
 	Stdin   io.Reader
 	Logger  slog.Logger
+	Net     Net
 
 	// testing
 	signalNotifyContext func(parent context.Context, signals ...os.Signal) (ctx context.Context, stop context.CancelFunc)
@@ -203,6 +204,7 @@ func (inv *Invocation) WithOS() *Invocation {
 		i.Stdin = os.Stdin
 		i.Args = os.Args[1:]
 		i.Environ = ParseEnviron(os.Environ(), "")
+		i.Net = osNet{}
 	})
 }
 
diff --git a/cli/clibase/net.go b/cli/clibase/net.go
@@ -0,0 +1,50 @@
+package clibase
+
+import (
+	"net"
+	"strconv"
+
+	"github.com/pion/udp"
+	"golang.org/x/xerrors"
+)
+
+// Net abstracts CLI commands interacting with the operating system networking.
+//
+// At present, it covers opening local listening sockets, since doing this
+// in testing is a challenge without flakes, since it's hard to pick a port we
+// know a priori will be free.
+type Net interface {
+	// Listen has the same semantics as `net.Listen` but also supports `udp`
+	Listen(network, address string) (net.Listener, error)
+}
+
+// osNet is an implementation that call the real OS for networking.
+type osNet struct{}
+
+func (osNet) Listen(network, address string) (net.Listener, error) {
+	switch network {
+	case "tcp", "tcp4", "tcp6", "unix", "unixpacket":
+		return net.Listen(network, address)
+	case "udp":
+		host, port, err := net.SplitHostPort(address)
+		if err != nil {
+			return nil, xerrors.Errorf("split %q: %w", address, err)
+		}
+
+		var portInt int
+		portInt, err = strconv.Atoi(port)
+		if err != nil {
+			return nil, xerrors.Errorf("parse port %v from %q as int: %w", port, address, err)
+		}
+
+		// Use pion here so that we get a stream-style net.Conn listener, instead
+		// of a packet-oriented connection that can read and write to multiple
+		// addresses.
+		return udp.Listen(network, &net.UDPAddr{
+			IP:   net.ParseIP(host),
+			Port: portInt,
+		})
+	default:
+		return nil, xerrors.Errorf("unknown listen network %q", network)
+	}
+}
diff --git a/cli/portforward.go b/cli/portforward.go
@@ -12,7 +12,6 @@ import (
 	"sync"
 	"syscall"
 
-	"github.com/pion/udp"
 	"golang.org/x/xerrors"
 
 	"cdr.dev/slog"
@@ -121,6 +120,7 @@ func (r *RootCmd) portForward() *clibase.Cmd {
 				wg                = new(sync.WaitGroup)
 				listeners         = make([]net.Listener, len(specs))
 				closeAllListeners = func() {
+					logger.Debug(ctx, "closing all listeners")
 					for _, l := range listeners {
 						if l == nil {
 							continue
@@ -134,6 +134,7 @@ func (r *RootCmd) portForward() *clibase.Cmd {
 			for i, spec := range specs {
 				l, err := listenAndPortForward(ctx, inv, conn, wg, spec, logger)
 				if err != nil {
+					logger.Error(ctx, "failed to listen", slog.F("spec", spec), slog.Error(err))
 					return err
 				}
 				listeners[i] = l
@@ -151,8 +152,10 @@ func (r *RootCmd) portForward() *clibase.Cmd {
 
 				select {
 				case <-ctx.Done():
+					logger.Debug(ctx, "command context expired waiting for signal", slog.Error(ctx.Err()))
 					closeErr = ctx.Err()
-				case <-sigs:
+				case sig := <-sigs:
+					logger.Debug(ctx, "received signal", slog.F("signal", sig))
 					_, _ = fmt.Fprintln(inv.Stderr, "\nReceived signal, closing all listeners and active connections")
 				}
 
@@ -161,6 +164,7 @@ func (r *RootCmd) portForward() *clibase.Cmd {
 			}()
 
 			conn.AwaitReachable(ctx)
+			logger.Debug(ctx, "read to accept connections to forward")
 			_, _ = fmt.Fprintln(inv.Stderr, "Ready!")
 			wg.Wait()
 			return closeErr
@@ -198,33 +202,7 @@ func listenAndPortForward(
 	logger = logger.With(slog.F("network", spec.listenNetwork), slog.F("address", spec.listenAddress))
 	_, _ = fmt.Fprintf(inv.Stderr, "Forwarding '%v://%v' locally to '%v://%v' in the workspace\n", spec.listenNetwork, spec.listenAddress, spec.dialNetwork, spec.dialAddress)
 
-	var (
-		l   net.Listener
-		err error
-	)
-	switch spec.listenNetwork {
-	case "tcp":
-		l, err = net.Listen(spec.listenNetwork, spec.listenAddress)
-	case "udp":
-		var host, port string
-		host, port, err = net.SplitHostPort(spec.listenAddress)
-		if err != nil {
-			return nil, xerrors.Errorf("split %q: %w", spec.listenAddress, err)
-		}
-
-		var portInt int
-		portInt, err = strconv.Atoi(port)
-		if err != nil {
-			return nil, xerrors.Errorf("parse port %v from %q as int: %w", port, spec.listenAddress, err)
-		}
-
-		l, err = udp.Listen(spec.listenNetwork, &net.UDPAddr{
-			IP:   net.ParseIP(host),
-			Port: portInt,
-		})
-	default:
-		return nil, xerrors.Errorf("unknown listen network %q", spec.listenNetwork)
-	}
+	l, err := inv.Net.Listen(spec.listenNetwork, spec.listenAddress)
 	if err != nil {
 		return nil, xerrors.Errorf("listen '%v://%v': %w", spec.listenNetwork, spec.listenAddress, err)
 	}
diff --git a/cli/portforward_test.go b/cli/portforward_test.go
