coder · mtojek · Jul 20, 2023 · Jul 14, 2023 · Jul 14, 2023 · Jul 14, 2023
diff --git a/cli/portforward.go b/cli/portforward.go
@@ -32,7 +32,7 @@ func (r *RootCmd) portForward() *clibase.Cmd {
 	client := new(codersdk.Client)
 	cmd := &clibase.Cmd{
 		Use:     "port-forward <workspace>",
-		Short:   "Forward ports from machine to a workspace",
+		Short:   `Forward ports from a workspace to the local machine. Please note that to achieve reverse port forwarding, utilize the "coder ssh -R" command.`,
 		Aliases: []string{"tunnel"},
 		Long: formatExamples(
 			example{

diff --git a/cli/remoteforward.go b/cli/remoteforward.go
@@ -0,0 +1,104 @@
+package cli
+
+import (
+	"context"
+	"fmt"
+	"io"
+	"net"
+	"regexp"
+	"strconv"
+
+	gossh "golang.org/x/crypto/ssh"
+	"golang.org/x/xerrors"
+
+	"github.com/coder/coder/agent/agentssh"
+)
+
+// cookieAddr is a special net.Addr accepted by sshRemoteForward() which includes a
+// cookie which is written to the connection before forwarding.
+type cookieAddr struct {
+	net.Addr
+	cookie []byte
+}
+
+// Format:
+// remote_port:local_address:local_port
+var remoteForwardRegex = regexp.MustCompile(`^(\d+):(.+):(\d+)$`)
+
+func validateRemoteForward(flag string) bool {
+	return remoteForwardRegex.MatchString(flag)
+}
+
+func parseRemoteForward(flag string) (net.Addr, net.Addr, error) {
+	matches := remoteForwardRegex.FindStringSubmatch(flag)
+
+	remotePort, err := strconv.Atoi(matches[1])
+	if err != nil {
+		return nil, nil, xerrors.Errorf("remote port is invalid: %w", err)
+	}
+	localAddress, err := net.ResolveIPAddr("ip", matches[2])
+	if err != nil {
+		return nil, nil, xerrors.Errorf("local address is invalid: %w", err)
+	}
+	localPort, err := strconv.Atoi(matches[3])
+	if err != nil {
+		return nil, nil, xerrors.Errorf("local port is invalid: %w", err)
+	}
+
+	localAddr := &net.TCPAddr{
+		IP:   localAddress.IP,
+		Port: localPort,
+	}
+
+	remoteAddr := &net.TCPAddr{
+		IP:   net.ParseIP("127.0.0.1"),
+		Port: remotePort,
+	}
+	return localAddr, remoteAddr, nil
+}
+
+// sshRemoteForward starts forwarding connections from a remote listener to a
+// local address via SSH in a goroutine.
+//
+// Accepts a `cookieAddr` as the local address.
+func sshRemoteForward(ctx context.Context, stderr io.Writer, sshClient *gossh.Client, localAddr, remoteAddr net.Addr) (io.Closer, error) {
+	listener, err := sshClient.Listen(remoteAddr.Network(), remoteAddr.String())
+	if err != nil {
+		return nil, xerrors.Errorf("listen on remote SSH address %s: %w", remoteAddr.String(), err)
+	}
+
+	go func() {
+		for {
+			remoteConn, err := listener.Accept()
+			if err != nil {
+				if ctx.Err() == nil {
+					_, _ = fmt.Fprintf(stderr, "Accept SSH listener connection: %+v\n", err)
+				}
+				return
+			}
+
+			go func() {
+				defer remoteConn.Close()
+
+				localConn, err := net.Dial(localAddr.Network(), localAddr.String())
+				if err != nil {
+					_, _ = fmt.Fprintf(stderr, "Dial local address %s: %+v\n", localAddr.String(), err)
+					return
+				}
+				defer localConn.Close()
+
+				if c, ok := localAddr.(cookieAddr); ok {
+					_, err = localConn.Write(c.cookie)
+					if err != nil {
+						_, _ = fmt.Fprintf(stderr, "Write cookie to local connection: %+v\n", err)
+						return
+					}
+				}
+
+				agentssh.Bicopy(ctx, localConn, remoteConn)
+			}()
+		}
+	}()
+
+	return listener, nil
+}
diff --git a/cli/ssh.go b/cli/ssh.go
@@ -6,7 +6,6 @@ import (
 	"errors"
 	"fmt"
 	"io"
-	"net"
 	"net/url"
 	"os"
 	"os/exec"
@@ -27,7 +26,6 @@ import (
 	"cdr.dev/slog"
 	"cdr.dev/slog/sloggers/sloghuman"
 
-	"github.com/coder/coder/agent/agentssh"
 	"github.com/coder/coder/cli/clibase"
 	"github.com/coder/coder/cli/cliui"
 	"github.com/coder/coder/coderd/autobuild/notify"
@@ -53,6 +51,7 @@ func (r *RootCmd) ssh() *clibase.Cmd {
 		waitEnum       string
 		noWait         bool
 		logDirPath     string
+		remoteForward  string
 	)
 	client := new(codersdk.Client)
 	cmd := &clibase.Cmd{
@@ -122,6 +121,16 @@ func (r *RootCmd) ssh() *clibase.Cmd {
 				client.SetLogger(logger)
 			}
 
+			if remoteForward != "" {
+				isValid := validateRemoteForward(remoteForward)
+				if !isValid {
+					return xerrors.Errorf(`invalid format of remote-forward, expected: remote_port:local_address:local_port`)
+				}
+				if isValid && stdio {
+					return xerrors.Errorf(`remote-forward can't be enabled in the stdio mode`)
+				}
+			}
+
 			workspace, workspaceAgent, err := getWorkspaceAndAgent(ctx, inv, client, codersdk.Me, inv.Args[0])
 			if err != nil {
 				return err
@@ -198,6 +207,7 @@ func (r *RootCmd) ssh() *clibase.Cmd {
 			}
 			defer conn.Close()
 			conn.AwaitReachable(ctx)
+
 			stopPolling := tryPollWorkspaceAutostop(ctx, client, workspace)
 			defer stopPolling()
 
@@ -300,6 +310,19 @@ func (r *RootCmd) ssh() *clibase.Cmd {
 				defer closer.Close()
 			}
 
+			if remoteForward != "" {
+				localAddr, remoteAddr, err := parseRemoteForward(remoteForward)
+				if err != nil {
+					return err
+				}
+
+				closer, err := sshRemoteForward(ctx, inv.Stderr, sshClient, localAddr, remoteAddr)
+				if err != nil {
+					return xerrors.Errorf("ssh remote forward: %w", err)
+				}
+				defer closer.Close()
+			}
+
 			stdoutFile, validOut := inv.Stdout.(*os.File)
 			stdinFile, validIn := inv.Stdin.(*os.File)
 			if validOut && validIn && isatty.IsTerminal(stdoutFile.Fd()) {
@@ -424,6 +447,13 @@ func (r *RootCmd) ssh() *clibase.Cmd {
 			FlagShorthand: "l",
 			Value:         clibase.StringOf(&logDirPath),
 		},
+		{
+			Flag:          "remote-forward",
+			Description:   "Enable remote port forwarding (remote_port:local_address:local_port).",
+			Env:           "CODER_SSH_REMOTE_FORWARD",
+			FlagShorthand: "R",
+			Value:         clibase.StringOf(&remoteForward),
+		},
 	}
 	return cmd
 }
@@ -568,8 +598,15 @@ func getWorkspaceAndAgent(ctx context.Context, inv *clibase.Invocation, client *
 // of the CLI running simultaneously.
 func tryPollWorkspaceAutostop(ctx context.Context, client *codersdk.Client, workspace codersdk.Workspace) (stop func()) {
 	lock := flock.New(filepath.Join(os.TempDir(), "coder-autostop-notify-"+workspace.ID.String()))
-	condition := notifyCondition(ctx, client, workspace.ID, lock)
-	return notify.Notify(condition, workspacePollInterval, autostopNotifyCountdown...)
+	conditionCtx, cancelCondition := context.WithCancel(ctx)
+	condition := notifyCondition(conditionCtx, client, workspace.ID, lock)
+	stopFunc := notify.Notify(condition, workspacePollInterval, autostopNotifyCountdown...)
+	return func() {
+		// With many "ssh" processes running, `lock.TryLockContext` can be hanging until the context canceled.
+		// Without this cancellation, a CLI process with failed remote-forward could be hanging indefinitely.
+		cancelCondition()
+		stopFunc()
+	}
 }
 
 // Notify the user if the workspace is due to shutdown.
@@ -752,56 +789,3 @@ func remoteGPGAgentSocket(sshClient *gossh.Client) (string, error) {
 
 	return string(bytes.TrimSpace(remoteSocket)), nil
 }
-
-// cookieAddr is a special net.Addr accepted by sshForward() which includes a
-// cookie which is written to the connection before forwarding.
-type cookieAddr struct {
-	net.Addr
-	cookie []byte
-}
-
-// sshForwardRemote starts forwarding connections from a remote listener to a
-// local address via SSH in a goroutine.
-//
-// Accepts a `cookieAddr` as the local address.
-func sshForwardRemote(ctx context.Context, stderr io.Writer, sshClient *gossh.Client, localAddr, remoteAddr net.Addr) (io.Closer, error) {
-	listener, err := sshClient.Listen(remoteAddr.Network(), remoteAddr.String())
-	if err != nil {
-		return nil, xerrors.Errorf("listen on remote SSH address %s: %w", remoteAddr.String(), err)
-	}
-
-	go func() {
-		for {
-			remoteConn, err := listener.Accept()
-			if err != nil {
-				if ctx.Err() == nil {
-					_, _ = fmt.Fprintf(stderr, "Accept SSH listener connection: %+v\n", err)
-				}
-				return
-			}
-
-			go func() {
-				defer remoteConn.Close()
-
-				localConn, err := net.Dial(localAddr.Network(), localAddr.String())
-				if err != nil {
-					_, _ = fmt.Fprintf(stderr, "Dial local address %s: %+v\n", localAddr.String(), err)
-					return
-				}
-				defer localConn.Close()
-
-				if c, ok := localAddr.(cookieAddr); ok {
-					_, err = localConn.Write(c.cookie)
-					if err != nil {
-						_, _ = fmt.Fprintf(stderr, "Write cookie to local connection: %+v\n", err)
-						return
-					}
-				}
-
-				agentssh.Bicopy(ctx, localConn, remoteConn)
-			}()
-		}
-	}()
-
-	return listener, nil
-}
diff --git a/cli/ssh_other.go b/cli/ssh_other.go
@@ -44,5 +44,5 @@ func forwardGPGAgent(ctx context.Context, stderr io.Writer, sshClient *gossh.Cli
 		Net:  "unix",
 	}
 
-	return sshForwardRemote(ctx, stderr, sshClient, localAddr, remoteAddr)
+	return sshRemoteForward(ctx, stderr, sshClient, localAddr, remoteAddr)
 }
diff --git a/cli/ssh_test.go b/cli/ssh_test.go
@@ -10,6 +10,8 @@ import (
 	"fmt"
 	"io"
 	"net"
+	"net/http"
+	"net/http/httptest"
 	"os"
 	"os/exec"
 	"path/filepath"
@@ -408,6 +410,58 @@ func TestSSH(t *testing.T) {
 		<-cmdDone
 	})
 
+	t.Run("RemoteForward", func(t *testing.T) {
+		if runtime.GOOS == "windows" {
+			t.Skip("Test not supported on windows")
+		}
+
+		t.Parallel()
+
+		httpServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			w.Write([]byte("hello world"))
+		}))
+		defer httpServer.Close()
+
+		client, workspace, agentToken := setupWorkspaceForAgent(t, nil)
+
+		agentClient := agentsdk.New(client.URL)
+		agentClient.SetSessionToken(agentToken)
+		agentCloser := agent.New(agent.Options{
+			Client: agentClient,
+			Logger: slogtest.Make(t, nil).Named("agent"),
+		})
+		defer agentCloser.Close()
+
+		ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong)
+		defer cancel()
+
+		inv, root := clitest.New(t,
+			"ssh",
+			workspace.Name,
+			"--remote-forward",
+			"8222:"+httpServer.Listener.Addr().String(),
+		)
+		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("curl localhost:8222")
+		pty.ExpectMatch("hello world")
+
+		// And we're done.
+		pty.WriteLine("exit")
+		<-cmdDone
+	})
+
 	t.Run("FileLogging", func(t *testing.T) {
 		t.Parallel()
 

diff --git a/cli/ssh_windows.go b/cli/ssh_windows.go
@@ -101,5 +101,5 @@ func forwardGPGAgent(ctx context.Context, stderr io.Writer, sshClient *gossh.Cli
 		Net:  "unix",
 	}
 
-	return sshForwardRemote(ctx, stderr, sshClient, localAddr, remoteAddr)
+	return sshRemoteForward(ctx, stderr, sshClient, localAddr, remoteAddr)
 }
diff --git a/cli/testdata/coder_--help.golden b/cli/testdata/coder_--help.golden
@@ -21,7 +21,9 @@ Coder v0.0.0-devel — A tool for provisioning self-hosted development environme
     logout            Unauthenticate your local session
     netcheck          Print network debug information for DERP and STUN
     ping              Ping a workspace
-    port-forward      Forward ports from machine to a workspace
+    port-forward      Forward ports from a workspace to the local machine.
+                      Please note that to achieve reverse port forwarding,
+                      utilize the "coder ssh -R" command.
     publickey         Output your Coder public key used for Git operations
     rename            Rename a workspace
     reset-password    Directly connect to the database to reset a user's

diff --git a/cli/testdata/coder_port-forward_--help.golden b/cli/testdata/coder_port-forward_--help.golden
@@ -1,6 +1,7 @@
 Usage: coder port-forward [flags] <workspace>
 
-Forward ports from machine to a workspace
+Forward ports from a workspace to the local machine. Please note that to achieve
+reverse port forwarding, utilize the "coder ssh -R" command.
 
 Aliases: tunnel
 

diff --git a/cli/testdata/coder_ssh_--help.golden b/cli/testdata/coder_ssh_--help.golden
@@ -27,6 +27,9 @@ Start a shell into a workspace
           behavior as non-blocking.
           DEPRECATED: Use --wait instead.
 
+  -R, --remote-forward string, $CODER_SSH_REMOTE_FORWARD
+          Enable remote port forwarding (remote_port:local_address:local_port).
+
       --stdio bool, $CODER_SSH_STDIO
           Specifies whether to emit SSH output over stdin/stdout.