feat(cli): allow specifying the listen address in coder port-forward (coder#7635)

coadler · web-flow · commit 5eb41e879461 · 2023-05-24T15:38:40.000-04:00
diff --git a/cli/portforward.go b/cli/portforward.go
@@ -4,6 +4,7 @@ import (
 	"context"
 	"fmt"
 	"net"
+	"net/netip"
 	"os"
 	"os/signal"
 	"strconv"
@@ -47,6 +48,10 @@ func (r *RootCmd) portForward() *clibase.Cmd {
 				Description: "Port forward multiple ports (TCP or UDP) in condensed syntax",
 				Command:     "coder port-forward <workspace> --tcp 8080,9000:3000,9090-9092,10000-10002:10010-10012",
 			},
+			example{
+				Description: "Port forward specifying the local address to bind to",
+				Command:     "coder port-forward <workspace> --tcp 1.2.3.4:8080:8080",
+			},
 		),
 		Middleware: clibase.Chain(
 			clibase.RequireNArgs(1),
@@ -255,9 +260,9 @@ func parsePortForwards(tcpSpecs, udpSpecs []string) ([]portForwardSpec, error) {
 			for _, port := range ports {
 				specs = append(specs, portForwardSpec{
 					listenNetwork: "tcp",
-					listenAddress: fmt.Sprintf("127.0.0.1:%v", port.local),
+					listenAddress: port.local.String(),
 					dialNetwork:   "tcp",
-					dialAddress:   fmt.Sprintf("127.0.0.1:%v", port.remote),
+					dialAddress:   port.remote.String(),
 				})
 			}
 		}
@@ -273,9 +278,9 @@ func parsePortForwards(tcpSpecs, udpSpecs []string) ([]portForwardSpec, error) {
 			for _, port := range ports {
 				specs = append(specs, portForwardSpec{
 					listenNetwork: "udp",
-					listenAddress: fmt.Sprintf("127.0.0.1:%v", port.local),
+					listenAddress: port.local.String(),
 					dialNetwork:   "udp",
-					dialAddress:   fmt.Sprintf("127.0.0.1:%v", port.remote),
+					dialAddress:   port.remote.String(),
 				})
 			}
 		}
@@ -307,29 +312,57 @@ func parsePort(in string) (uint16, error) {
 }
 
 type parsedSrcDestPort struct {
-	local, remote uint16
+	local, remote netip.AddrPort
 }
 
 func parseSrcDestPorts(in string) ([]parsedSrcDestPort, error) {
-	parts := strings.Split(in, ":")
-	if len(parts) > 2 {
-		return nil, xerrors.Errorf("invalid port specification %q", in)
-	}
-	if len(parts) == 1 {
+	var (
+		err        error
+		parts      = strings.Split(in, ":")
+		localAddr  = netip.AddrFrom4([4]byte{127, 0, 0, 1})
+		remoteAddr = netip.AddrFrom4([4]byte{127, 0, 0, 1})
+	)
+
+	switch len(parts) {
+	case 1:
 		// Duplicate the single part
 		parts = append(parts, parts[0])
+	case 2:
+		// Check to see if the first part is an IP address.
+		_localAddr, err := netip.ParseAddr(parts[0])
+		if err != nil {
+			break
+		}
+		// The first part is the local address, so duplicate the port.
+		localAddr = _localAddr
+		parts = []string{parts[1], parts[1]}
+
+	case 3:
+		_localAddr, err := netip.ParseAddr(parts[0])
+		if err != nil {
+			return nil, xerrors.Errorf("invalid port specification %q; invalid ip %q: %w", in, parts[0], err)
+		}
+		localAddr = _localAddr
+		parts = parts[1:]
+
+	default:
+		return nil, xerrors.Errorf("invalid port specification %q", in)
 	}
+
 	if !strings.Contains(parts[0], "-") {
-		local, err := parsePort(parts[0])
+		localPort, err := parsePort(parts[0])
 		if err != nil {
 			return nil, xerrors.Errorf("parse local port from %q: %w", in, err)
 		}
-		remote, err := parsePort(parts[1])
+		remotePort, err := parsePort(parts[1])
 		if err != nil {
 			return nil, xerrors.Errorf("parse remote port from %q: %w", in, err)
 		}
 
-		return []parsedSrcDestPort{{local: local, remote: remote}}, nil
+		return []parsedSrcDestPort{{
+			local:  netip.AddrPortFrom(localAddr, localPort),
+			remote: netip.AddrPortFrom(remoteAddr, remotePort),
+		}}, nil
 	}
 
 	local, err := parsePortRange(parts[0])
@@ -346,8 +379,8 @@ func parseSrcDestPorts(in string) ([]parsedSrcDestPort, error) {
 	var out []parsedSrcDestPort
 	for i := range local {
 		out = append(out, parsedSrcDestPort{
-			local:  local[i],
-			remote: remote[i],
+			local:  netip.AddrPortFrom(localAddr, local[i]),
+			remote: netip.AddrPortFrom(remoteAddr, remote[i]),
 		})
 	}
 	return out, nil
diff --git a/cli/portforward_test.go b/cli/portforward_test.go
@@ -104,6 +104,25 @@ func TestPortForward(t *testing.T) {
 				return l.Addr().String(), port
 			},
 		},
+		{
+			name:    "TCPWithAddress",
+			network: "tcp",
+			flag:    "--tcp=%v:%v",
+			setupRemote: func(t *testing.T) net.Listener {
+				l, err := net.Listen("tcp", "127.0.0.1:0")
+				require.NoError(t, err, "create TCP listener")
+				return l
+			},
+			setupLocal: func(t *testing.T) (string, string) {
+				l, err := net.Listen("tcp", "127.0.0.1:0")
+				require.NoError(t, err, "create TCP listener to generate random port")
+				defer l.Close()
+
+				_, port, err := net.SplitHostPort(l.Addr().String())
+				require.NoErrorf(t, err, "split TCP address %q", l.Addr().String())
+				return l.Addr().String(), fmt.Sprint("0.0.0.0:", port)
+			},
+		},
 	}
 
 	// Setup agent once to be shared between test-cases (avoid expensive
diff --git a/cli/testdata/coder_port-forward_--help.golden b/cli/testdata/coder_port-forward_--help.golden
@@ -22,6 +22,10 @@ Aliases: tunnel
 
      [40m [0m[91;40m$ coder port-forward <workspace> --tcp 8080,9000:3000,9090-9092,10000-10002:10010-10012[0m[40m [0m
 
+  - Port forward specifying the local address to bind to:                       
+
+     [40m [0m[91;40m$ coder port-forward <workspace> --tcp 1.2.3.4:8080:8080[0m[40m [0m
+
 [1mOptions[0m
   -p, --tcp string-array, $CODER_PORT_FORWARD_TCP
           Forward TCP port(s) from the workspace to the local machine.
diff --git a/docs/cli/port-forward.md b/docs/cli/port-forward.md
@@ -34,6 +34,10 @@ coder port-forward [flags] <workspace>
   - Port forward multiple ports (TCP or UDP) in condensed syntax:
 
       $ coder port-forward <workspace> --tcp 8080,9000:3000,9090-9092,10000-10002:10010-10012
+
+  - Port forward specifying the local address to bind to:
+
+      $ coder port-forward <workspace> --tcp 1.2.3.4:8080:8080
 ```
 
 ## Options

-Original file line number
+Diff line change
      [40m [0m[91;40m$ coder port-forward <workspace> --tcp 8080,9000:3000,9090-9092,10000-10002:10010-10012[0m[40m [0m
 +  - Port forward specifying the local address to bind to:
++
 +     [40m [0m[91;40m$ coder port-forward <workspace> --tcp 1.2.3.4:8080:8080[0m[40m [0m
++
 [1mOptions[0m
   -p, --tcp string-array, $CODER_PORT_FORWARD_TCP
           Forward TCP port(s) from the workspace to the local machine.