add test case for multiple socket forwards

mafredri · mafredri · commit 8da5d2c4d5ef · 2024-01-19T14:23:07.000+02:00
diff --git a/cli/ssh_test.go b/cli/ssh_test.go
@@ -883,6 +883,104 @@ func TestSSH(t *testing.T) {
 		require.NoError(t, err)
 	})
 
+	// Test that we can forward a local unix socket to a remote unix socket and
+	// that new SSH sessions take over the socket without closing active socket
+	// connections.
+	t.Run("RemoteForwardMultipleUnixSockets", func(t *testing.T) {
+		if runtime.GOOS == "windows" {
+			t.Skip("Test not supported on windows")
+		}
+
+		t.Parallel()
+
+		client, workspace, agentToken := setupWorkspaceForAgent(t)
+
+		_ = agenttest.New(t, client.URL, agentToken)
+		coderdtest.AwaitWorkspaceAgents(t, client, workspace.ID)
+
+		// Wait super long so this doesn't flake on -race test.
+		ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitSuperLong)
+		defer cancel()
+
+		tmpdir := tempDirUnixSocket(t)
+
+		type testSocket struct {
+			local  string
+			remote string
+		}
+
+		args := []string{"ssh", workspace.Name}
+		var sockets []testSocket
+		for i := 0; i < 2; i++ {
+			localSock := filepath.Join(tmpdir, fmt.Sprintf("local-%d.sock", i))
+			remoteSock := filepath.Join(tmpdir, fmt.Sprintf("remote-%d.sock", i))
+			sockets = append(sockets, testSocket{
+				local:  localSock,
+				remote: remoteSock,
+			})
+			args = append(args, "--remote-forward", fmt.Sprintf("%s:%s", remoteSock, localSock))
+		}
+
+		inv, root := clitest.New(t, args...)
+		clitest.SetupConfig(t, client, root)
+		pty := ptytest.New(t).Attach(inv)
+		inv.Stderr = pty.Output()
+
+		clitest.Start(t, inv.WithContext(ctx))
+
+		// Since something was output, it should be safe to write input.
+		// This could show a prompt or "running startup scripts", so it's
+		// not indicative of the SSH connection being ready.
+		_ = pty.Peek(ctx, 1)
+
+		// Ensure the SSH connection is ready by testing the shell
+		// input/output.
+		pty.WriteLine("echo ping' 'pong")
+		pty.ExpectMatchContext(ctx, "ping pong")
+
+		for i, sock := range sockets {
+			i := i
+			// Start the listener on the "local machine".
+			l, err := net.Listen("unix", sock.local)
+			require.NoError(t, err)
+			defer l.Close() //nolint:revive // Defer is fine in this loop, we only run it twice.
+			testutil.Go(t, func() {
+				for {
+					fd, err := l.Accept()
+					if err != nil {
+						if !errors.Is(err, net.ErrClosed) {
+							assert.NoError(t, err, "listener accept failed", i)
+						}
+						return
+					}
+
+					testutil.Go(t, func() {
+						defer fd.Close()
+						agentssh.Bicopy(ctx, fd, fd)
+					})
+				}
+			})
+
+			// Dial the forwarded socket on the "remote machine".
+			d := &net.Dialer{}
+			fd, err := d.DialContext(ctx, "unix", sock.remote)
+			require.NoError(t, err, i)
+			defer fd.Close() //nolint:revive // Defer is fine in this loop, we only run it twice.
+
+			// Ping / pong to ensure the socket is working.
+			_, err = fd.Write([]byte("hello world"))
+			require.NoError(t, err, i)
+
+			buf := make([]byte, 11)
+			_, err = fd.Read(buf)
+			require.NoError(t, err, i)
+			require.Equal(t, "hello world", string(buf), i)
+		}
+
+		// And we're done.
+		pty.WriteLine("exit")
+	})
+
 	t.Run("FileLogging", func(t *testing.T) {
 		t.Parallel()
 
