add tests and check for socket file

monika-canva · monika-canva · commit 819fd9af9501 · 2023-09-29T12:09:36.000+08:00
diff --git a/cli/remoteforward.go b/cli/remoteforward.go
@@ -5,6 +5,7 @@ import (
 	"fmt"
 	"io"
 	"net"
+	"os"
 	"regexp"
 	"strconv"
 
@@ -70,6 +71,15 @@ func parseRemoteForwardUnixSocket(matches []string) (net.Addr, net.Addr, error)
 	remoteSocket := matches[1]
 	localSocket := matches[2]
 
+	fileInfo, err := os.Stat(localSocket)
+	if err != nil {
+		return nil, nil, err
+	}
+
+	if fileInfo.Mode()&os.ModeSocket == 0 {
+		return nil, nil, xerrors.New("File is not a Unix domain socket file")
+	}
+
 	remoteAddr := &net.UnixAddr{
 		Name: remoteSocket,
 		Net:  "unix",
diff --git a/cli/ssh_test.go b/cli/ssh_test.go
@@ -428,6 +428,54 @@ func TestSSH(t *testing.T) {
 		<-cmdDone
 	})
 
+	t.Run("RemoteForwardUnixSocket", func(t *testing.T) {
+		if runtime.GOOS == "windows" {
+			t.Skip("Test not supported on windows")
+		}
+
+		t.Parallel()
+
+		client, workspace, agentToken := setupWorkspaceForAgent(t, nil)
+
+		_ = agenttest.New(t, client.URL, agentToken)
+		coderdtest.AwaitWorkspaceAgents(t, client, workspace.ID)
+
+		ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong)
+		defer cancel()
+
+		tmpdir := tempDirUnixSocket(t)
+		agentSock := filepath.Join(tmpdir, "agent.sock")
+		l, err := net.Listen("unix", agentSock)
+		require.NoError(t, err)
+		defer l.Close()
+
+		inv, root := clitest.New(t,
+			"ssh",
+			workspace.Name,
+			"--remote-forward",
+			agentSock+":/tmp/test.sock",
+		)
+		clitest.SetupConfig(t, client, root)
+		pty := ptytest.New(t).Attach(inv)
+		inv.Stderr = pty.Output()
+		cmdDone := tGo(t, func() {
+			err := inv.WithContext(ctx).Run()
+			assert.NoError(t, err, "ssh command failed")
+		})
+
+		// Wait for the prompt or any output really to indicate the command has
+		// started and accepting input on stdin.
+		_ = pty.Peek(ctx, 1)
+
+		// Download the test page
+		pty.WriteLine("ss -xl state listening src /var/roo/daemon.sock | wc -l")
+		pty.ExpectMatch("2")
+
+		// And we're done.
+		pty.WriteLine("exit")
+		<-cmdDone
+	})
+
 	t.Run("FileLogging", func(t *testing.T) {
 		t.Parallel()
 
