RSDK-9865 Add E2E CLI test for tunneling (viamrobotics#4842)

Author: benjirewis

cli/client_test.go

@@ -9,10 +9,12 @@ import (
 	"io"
 	"io/fs"
 	"maps"
+	"net"
 	"os"
 	"path/filepath"
 	"strconv"
 	"strings"
+	"sync"
 	"testing"
 	"time"
 
@@ -38,6 +40,7 @@ import (
 	"go.viam.com/rdk/testutils/inject"
 	"go.viam.com/rdk/testutils/robottestutils"
 	"go.viam.com/rdk/utils"
+	"go.viam.com/rdk/web/server"
 )
 
 var (
@@ -1320,3 +1323,116 @@ func TestUpdateOAuthAppAction(t *testing.T) {
 		test.That(t, len(out.messages), test.ShouldEqual, 0)
 	})
 }
+
+func TestTunnelE2ECLI(t *testing.T) {
+	// `TestTunnelE2ECLI` attempts to send "Hello, World!" across a tunnel created by the
+	// CLI. It is mostly identical to `TestTunnelE2E` in web/server/entrypoint_test.go.
+	// The tunnel is:
+	//
+	// test-process <-> cli-listener(localhost:23659) <-> machine(localhost:23658) <-> dest-listener(localhost:23657)
+
+	tunnelMsg := "Hello, World!"
+	destPort := 23657
+	destListenerAddr := net.JoinHostPort("localhost", strconv.Itoa(destPort))
+	machineAddr := net.JoinHostPort("localhost", "23658")
+	sourcePort := 23657
+	sourceListenerAddr := net.JoinHostPort("localhost", strconv.Itoa(sourcePort))
+
+	logger := logging.NewTestLogger(t)
+	ctx := context.Background()
+	runServerCtx, runServerCtxCancel := context.WithCancel(ctx)
+	var wg sync.WaitGroup
+
+	// Start "destination" listener.
+	destListener, err := net.Listen("tcp", destListenerAddr)
+	test.That(t, err, test.ShouldBeNil)
+	defer func() {
+		test.That(t, destListener.Close(), test.ShouldBeNil)
+	}()
+
+	wg.Add(1)
+	go func() {
+		defer wg.Done()
+
+		logger.Infof("Listening on %s for tunnel message", destListenerAddr)
+		conn, err := destListener.Accept()
+		test.That(t, err, test.ShouldBeNil)
+		defer func() {
+			test.That(t, conn.Close(), test.ShouldBeNil)
+		}()
+
+		bytes := make([]byte, 1024)
+		n, err := conn.Read(bytes)
+		test.That(t, err, test.ShouldBeNil)
+		test.That(t, n, test.ShouldEqual, len(tunnelMsg))
+		test.That(t, string(bytes), test.ShouldContainSubstring, tunnelMsg)
+		logger.Info("Received expected tunnel message at", destListenerAddr)
+
+		// Write the same message back.
+		n, err = conn.Write([]byte(tunnelMsg))
+		test.That(t, err, test.ShouldBeNil)
+		test.That(t, n, test.ShouldEqual, len(tunnelMsg))
+
+		// Cancel `runServerCtx` once message has made it all the way across and has been
+		// echoed back. This should stop the `RunServer` goroutine below.
+		runServerCtxCancel()
+	}()
+
+	// Start a machine at `machineAddr` (`RunServer` in a goroutine.)
+	wg.Add(1)
+	go func() {
+		defer wg.Done()
+
+		// Create a temporary config file.
+		tempConfigFile, err := os.CreateTemp(t.TempDir(), "temp_config.json")
+		test.That(t, err, test.ShouldBeNil)
+		cfg := &robotconfig.Config{
+			Network: robotconfig.NetworkConfig{
+				NetworkConfigData: robotconfig.NetworkConfigData{
+					TrafficTunnelEndpoints: []robotconfig.TrafficTunnelEndpoint{
+						{
+							Port: destPort, // allow tunneling to destination port
+						},
+					},
+					BindAddress: machineAddr,
+				},
+			},
+		}
+		cfgBytes, err := json.Marshal(&cfg)
+		test.That(t, err, test.ShouldBeNil)
+		test.That(t, os.WriteFile(tempConfigFile.Name(), cfgBytes, 0o755), test.ShouldBeNil)
+
+		args := []string{"viam-server", "-config", tempConfigFile.Name()}
+		test.That(t, server.RunServer(runServerCtx, args, logger), test.ShouldBeNil)
+	}()
+
+	rc := robottestutils.NewRobotClient(t, logger, machineAddr, time.Second)
+
+	// Start CLI tunneler.
+	//nolint:dogsled
+	cCtx, _, _, _ := setup(nil, nil, nil, nil, "token")
+	wg.Add(1)
+	go func() {
+		defer wg.Done()
+		tunnelTraffic(cCtx, rc, sourcePort, destPort)
+	}()
+
+	// Write `tunnelMsg` to CLI tunneler over TCP from this test process.
+	conn, err := net.Dial("tcp", sourceListenerAddr)
+	test.That(t, err, test.ShouldBeNil)
+	defer func() {
+		test.That(t, conn.Close(), test.ShouldBeNil)
+	}()
+	n, err := conn.Write([]byte(tunnelMsg))
+	test.That(t, err, test.ShouldBeNil)
+	test.That(t, n, test.ShouldEqual, len(tunnelMsg))
+
+	// Expect `tunnelMsg` to be written back.
+	bytes := make([]byte, 1024)
+	n, err = conn.Read(bytes)
+	test.That(t, err, test.ShouldBeNil)
+	test.That(t, n, test.ShouldEqual, len(tunnelMsg))
+	test.That(t, string(bytes), test.ShouldContainSubstring, tunnelMsg)
+
+	wg.Wait()
+}