feat: Support for comma-separation and ranges in port-forward

Fixes #3766
coder · mafredri · Oct 3, 2022 · Sep 23, 2022 · Sep 23, 2022 · Sep 23, 2022
commit a3587a2453da406aed46e00db121850f3515284f
diff --git a/cli/portforward.go b/cli/portforward.go
@@ -18,6 +18,7 @@ import (
 
 	"cdr.dev/slog"
 	"github.com/coder/coder/agent"
+	"github.com/coder/coder/cli/cliflag"
 	"github.com/coder/coder/cli/cliui"
 	"github.com/coder/coder/codersdk"
 )
@@ -45,6 +46,10 @@ func portForward() *cobra.Command {
 				Description: "Port forward multiple TCP ports and a UDP port",
 				Command:     "coder port-forward <workspace> --tcp 8080:8080 --tcp 9000:3000 --udp 5353:53",
 			},
+			example{
+				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",
+			},
 		),
 		RunE: func(cmd *cobra.Command, args []string) error {
 			ctx, cancel := context.WithCancel(cmd.Context())
@@ -164,8 +169,8 @@ func portForward() *cobra.Command {
 		},
 	}
 
-	cmd.Flags().StringArrayVarP(&tcpForwards, "tcp", "p", []string{}, "Forward a TCP port from the workspace to the local machine")
-	cmd.Flags().StringArrayVar(&udpForwards, "udp", []string{}, "Forward a UDP port from the workspace to the local machine. The UDP connection has TCP-like semantics to support stateful UDP protocols")
+	cliflag.StringArrayVarP(cmd.Flags(), &tcpForwards, "tcp", "p", "CODER_PORT_FORWARD_TCP", nil, "Forward TCP port(s) from the workspace to the local machine")
+	cliflag.StringArrayVarP(cmd.Flags(), &udpForwards, "udp", "", "CODER_PORT_FORWARD_UDP", nil, "Forward UDP port(s) from the workspace to the local machine. The UDP connection has TCP-like semantics to support stateful UDP protocols")
 	return cmd
 }
 
@@ -242,32 +247,40 @@ type portForwardSpec struct {
 func parsePortForwards(tcpSpecs, udpSpecs []string) ([]portForwardSpec, error) {
 	specs := []portForwardSpec{}
 
-	for _, spec := range tcpSpecs {
-		local, remote, err := parsePortPort(spec)
-		if err != nil {
-			return nil, xerrors.Errorf("failed to parse TCP port-forward specification %q: %w", spec, err)
-		}
+	for _, specEntry := range tcpSpecs {
+		for _, spec := range strings.Split(specEntry, ",") {
+			ports, err := parseSrcDestPorts(spec)
+			if err != nil {
+				return nil, xerrors.Errorf("failed to parse TCP port-forward specification %q: %w", spec, err)
+			}
 
-		specs = append(specs, portForwardSpec{
-			listenNetwork: "tcp",
-			listenAddress: fmt.Sprintf("127.0.0.1:%v", local),
-			dialNetwork:   "tcp",
-			dialAddress:   fmt.Sprintf("127.0.0.1:%v", remote),
-		})
+			for _, port := range ports {
+				specs = append(specs, portForwardSpec{
+					listenNetwork: "tcp",
+					listenAddress: fmt.Sprintf("127.0.0.1:%v", port.local),
+					dialNetwork:   "tcp",
+					dialAddress:   fmt.Sprintf("127.0.0.1:%v", port.remote),
+				})
+			}
+		}
 	}
 
-	for _, spec := range udpSpecs {
-		local, remote, err := parsePortPort(spec)
-		if err != nil {
-			return nil, xerrors.Errorf("failed to parse UDP port-forward specification %q: %w", spec, err)
-		}
+	for _, specEntry := range udpSpecs {
+		for _, spec := range strings.Split(specEntry, ",") {
+			ports, err := parseSrcDestPorts(spec)
+			if err != nil {
+				return nil, xerrors.Errorf("failed to parse UDP port-forward specification %q: %w", spec, err)
+			}
 
-		specs = append(specs, portForwardSpec{
-			listenNetwork: "udp",
-			listenAddress: fmt.Sprintf("127.0.0.1:%v", local),
-			dialNetwork:   "udp",
-			dialAddress:   fmt.Sprintf("127.0.0.1:%v", remote),
-		})
+			for _, port := range ports {
+				specs = append(specs, portForwardSpec{
+					listenNetwork: "udp",
+					listenAddress: fmt.Sprintf("127.0.0.1:%v", port.local),
+					dialNetwork:   "udp",
+					dialAddress:   fmt.Sprintf("127.0.0.1:%v", port.remote),
+				})
+			}
+		}
 	}
 
 	// Check for duplicate entries.
@@ -295,24 +308,72 @@ func parsePort(in string) (uint16, error) {
 	return uint16(port), nil
 }
 
-func parsePortPort(in string) (local uint16, remote uint16, err error) {
+type parsedSrcDestPort struct {
+	local, remote uint16
+}
+
+func parseSrcDestPorts(in string) ([]parsedSrcDestPort, error) {
 	parts := strings.Split(in, ":")
 	if len(parts) > 2 {
-		return 0, 0, xerrors.Errorf("invalid port specification %q", in)
+		return nil, xerrors.Errorf("invalid port specification %q", in)
 	}
 	if len(parts) == 1 {
 		// Duplicate the single part
 		parts = append(parts, parts[0])
 	}
+	if !strings.Contains(parts[0], "-") {
+		local, err := parsePort(parts[0])
+		if err != nil {
+			return nil, xerrors.Errorf("parse local port from %q: %w", in, err)
+		}
+		remote, err := parsePort(parts[1])
+		if err != nil {
+			return nil, xerrors.Errorf("parse remote port from %q: %w", in, err)
+		}
 
-	local, err = parsePort(parts[0])
+		return []parsedSrcDestPort{{local: local, remote: remote}}, nil
+	}
+
+	local, err := parsePortRange(parts[0])
 	if err != nil {
-		return 0, 0, xerrors.Errorf("parse local port from %q: %w", in, err)
+		return nil, xerrors.Errorf("parse local port range from %q: %w", in, err)
 	}
-	remote, err = parsePort(parts[1])
+	remote, err := parsePortRange(parts[1])
 	if err != nil {
-		return 0, 0, xerrors.Errorf("parse remote port from %q: %w", in, err)
+		return nil, xerrors.Errorf("parse remote port range from %q: %w", in, err)
+	}
+	if len(local) != len(remote) {
+		return nil, xerrors.Errorf("port ranges must be the same length, got %d ports forwarded to %d ports", len(local), len(remote))
+	}
+	var out []parsedSrcDestPort
+	for i := range local {
+		out = append(out, parsedSrcDestPort{
+			local:  local[i],
+			remote: remote[i],
+		})
 	}
+	return out, nil
+}
 
-	return local, remote, nil
+func parsePortRange(in string) ([]uint16, error) {
+	parts := strings.Split(in, "-")
+	if len(parts) != 2 {
+		return nil, xerrors.Errorf("invalid port range specification %q", in)
+	}
+	start, err := parsePort(parts[0])
+	if err != nil {
+		return nil, xerrors.Errorf("parse range start port from %q: %w", in, err)
+	}
+	end, err := parsePort(parts[1])
+	if err != nil {
+		return nil, xerrors.Errorf("parse range end port from %q: %w", in, err)
+	}
+	if end < start {
+		return nil, xerrors.Errorf("range end port %v is less than start port %v", end, start)
+	}
+	var ports []uint16
+	for i := start; i <= end; i++ {
+		ports = append(ports, i)
+	}
+	return ports, nil
 }
diff --git a/cli/portforward_internal_test.go b/cli/portforward_internal_test.go
@@ -0,0 +1,88 @@
+package cli
+
+import (
+	"fmt"
+	"strings"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+func Test_parsePortForwards(t *testing.T) {
+	portForwardSpecToString := func(v []portForwardSpec) (out []string) {
+		for _, p := range v {
+			require.Equal(t, p.listenNetwork, p.dialNetwork)
+			out = append(out, fmt.Sprintf("%s:%s", strings.Replace(p.listenAddress, "127.0.0.1:", "", 1), strings.Replace(p.dialAddress, "127.0.0.1:", "", 1)))
+		}
+		return out
+	}
+	type args struct {
+		tcpSpecs []string
+		udpSpecs []string
+	}
+	tests := []struct {
+		name    string
+		args    args
+		want    []string
+		wantErr bool
+	}{
+		{
+			name: "TCP mixed ports and ranges",
+			args: args{
+				tcpSpecs: []string{
+					"8000,8080:8081,9000-9002,9003-9004:9005-9006",
+					"10000",
+				},
+			},
+			want: []string{
+				"8000:8000",
+				"8080:8081",
+				"9000:9000",
+				"9001:9001",
+				"9002:9002",
+				"9003:9005",
+				"9004:9006",
+				"10000:10000",
+			},
+		},
+		{
+			name: "UDP with port range",
+			args: args{
+				udpSpecs: []string{"8000,8080-8081"},
+			},
+			want: []string{
+				"8000:8000",
+				"8080:8080",
+				"8081:8081",
+			},
+		},
+		{
+			name: "Bad port range",
+			args: args{
+				tcpSpecs: []string{"8000-7000"},
+			},
+			wantErr: true,
+		},
+		{
+			name: "Bad dest port range",
+			args: args{
+				tcpSpecs: []string{"8080-8081:9080-9082"},
+			},
+			wantErr: true,
+		},
+	}
+	for _, tt := range tests {
+		tt := tt
+		t.Run(tt.name, func(t *testing.T) {
+			t.Parallel()
+
+			got, err := parsePortForwards(tt.args.tcpSpecs, tt.args.udpSpecs)
+			if (err != nil) != tt.wantErr {
+				t.Fatalf("parsePortForwards() error = %v, wantErr %v", err, tt.wantErr)
+				return
+			}
+			gotStrings := portForwardSpecToString(got)
+			require.Equal(t, tt.want, gotStrings)
+		})
+	}
+}