fix(BRE2-824): validate SSH port and allow re-entry

Author: patelspratik

pkg/cmd/enablessh/enablessh.go

@@ -115,7 +115,7 @@ func enableSSH(
 		t.Vprint("")
 		port, err = register.PromptSSHPort(t)
 		if err != nil {
-			return fmt.Errorf("SSH port: %w", err)
+			return fmt.Errorf("invalid SSH port: %w", err)
 		}
 
 		if err := register.OpenSSHPort(ctx, t, deps.nodeClients, tokenProvider, reg, port); err != nil {

pkg/cmd/register/register.go

@@ -250,7 +250,7 @@ func runRegister(ctx context.Context, t *terminal.Terminal, s RegisterStore, opt
 			return fmt.Errorf("failed to determine current Linux user: %w", err)
 		}
 		if err := grantSSHAccessWithPort(ctx, t, deps, s, reg, brevUser, osUser, sshPortForGrant, opts.interactive, opts.skipConfirm); err != nil {
-			t.Vprintf("  Warning: SSH access not granted: %v\n", err)
+			t.Vprintf("  %s\n", t.Yellow(fmt.Sprintf("Warning: %v", err)))
 		}
 	}
 
@@ -465,7 +465,7 @@ func grantSSHAccessWithPort(ctx context.Context, t *terminal.Terminal, deps regi
 		t.Vprint("")
 		port, err = PromptSSHPort(t)
 		if err != nil {
-			return fmt.Errorf("SSH port: %w", err)
+			return fmt.Errorf("invalid SSH port: %w", err)
 		}
 	} else {
 		t.Vprintf("  %s %s\n", t.Green(fmt.Sprintf("%-14s", "SSH port:")), t.BoldBlue(fmt.Sprintf("%d", port)))

pkg/cmd/register/register_test.go

@@ -980,7 +980,7 @@ func Test_runRegister_NoNameAlreadyRegistered(t *testing.T) {
 	}
 }
 
-func Test_runRegister_OpenSSHPort(t *testing.T) { // nolint:funlen // test
+func Test_runRegister_OpenSSHPort(t *testing.T) { // nolint:funlen, gocyclo, gocognit // test
 	tests := []struct {
 		name   string
 		port   int32
@@ -1035,6 +1035,23 @@ func Test_runRegister_OpenSSHPort(t *testing.T) { // nolint:funlen // test
 				}
 			},
 		},
+		{
+			name: "InvalidPortNoAPICall",
+			port: 99999,
+			verify: func(t *testing.T, openReq *nodev1.OpenPortRequest, _ *nodev1.GrantNodeSSHAccessRequest, regStore *mockRegistrationStore, err error) {
+				t.Helper()
+				if err != nil {
+					t.Fatalf("registration should succeed even when SSH port is invalid (soft error), got: %v", err)
+				}
+				if openReq != nil {
+					t.Error("expected OpenPort NOT to be called for invalid port")
+				}
+				exists, _ := regStore.Exists()
+				if !exists {
+					t.Error("expected registration to still exist after invalid port")
+				}
+			},
+		},
 		{
 			name: "GrantRequestHasNoPort",
 			port: 22,

pkg/cmd/register/sshkeys.go

@@ -177,6 +177,9 @@ func OpenSSHPort(
 	reg *DeviceRegistration,
 	port int32,
 ) error {
+	if port < 1 || port > 65535 {
+		return fmt.Errorf("invalid SSH port %d: port must be between 1 and 65535", port)
+	}
 	client := nodeClients.NewNodeClient(tokenProvider, config.GlobalConfig.GetBrevPublicAPIURL())
 	_, err := client.OpenPort(ctx, connect.NewRequest(&nodev1.OpenPortRequest{
 		ExternalNodeId: reg.ExternalNodeID,
@@ -275,30 +278,34 @@ func SetTestSSHPort(port int32) { testSSHPort = &port }
 func ClearTestSSHPort() { testSSHPort = nil }
 
 // PromptSSHPort prompts the user for the target SSH port, defaulting to 22 if
-// they press Enter or leave it empty. Returns an error for invalid port numbers.
-// Uses a single print + stdin read to avoid promptui rendering the label twice.
+// they press Enter or leave it empty. Re-prompts on invalid input until a valid
+// port is provided. Only returns an error for unrecoverable I/O failures.
 func PromptSSHPort(t *terminal.Terminal) (int32, error) {
 	if testSSHPort != nil {
 		return *testSSHPort, nil
 	}
-	t.Vprintf("  %s ", t.Green("SSH port (default 22):"))
 	reader := bufio.NewReader(os.Stdin)
-	line, err := reader.ReadString('\n')
-	if err != nil {
-		return 0, fmt.Errorf("reading input: %w", err)
-	}
-	portStr := strings.TrimSpace(line)
-	if portStr == "" {
-		return defaultSSHPort, nil
-	}
-	n, err := strconv.ParseUint(portStr, 10, 16)
-	if err != nil {
-		return 0, fmt.Errorf("invalid port %q: %w", portStr, err)
-	}
-	if n < 1 || n > 65535 {
-		return 0, fmt.Errorf("port must be between 1 and 65535, got %d", n)
+	for {
+		t.Vprintf("  %s ", t.Green("SSH port (default 22):"))
+		line, err := reader.ReadString('\n')
+		if err != nil {
+			return 0, fmt.Errorf("reading input: %w", err)
+		}
+		portStr := strings.TrimSpace(line)
+		if portStr == "" {
+			return defaultSSHPort, nil
+		}
+		n, err := strconv.ParseUint(portStr, 10, 16)
+		if err != nil {
+			t.Vprintf("  %s\n", t.Yellow(fmt.Sprintf("Invalid port %q: port must be a number between 1 and 65535", portStr)))
+			continue
+		}
+		if n < 1 || n > 65535 { // explicit gate even though ParseUint will already fail values out of range
+			t.Vprintf("  %s\n", t.Yellow(fmt.Sprintf("Invalid port %q: port must be between 1 and 65535", portStr)))
+			continue
+		}
+		return int32(n), nil
 	}
-	return int32(n), nil
 }
 
 // InstallAuthorizedKey appends the given public key to the user's

pkg/cmd/register/sshkeys_test.go

@@ -6,6 +6,8 @@ import (
 	"path/filepath"
 	"strings"
 	"testing"
+
+	"github.com/brevdev/brev-cli/pkg/terminal"
 )
 
 func tempUser(t *testing.T) *user.User {
@@ -215,6 +217,61 @@ func TestInstallAuthorizedKey_IncludesUserID(t *testing.T) {
 	}
 }
 
+// --- PromptSSHPort ---
+
+func promptSSHPortWithInput(t *testing.T, input string) (int32, error) {
+	t.Helper()
+
+	r, w, err := os.Pipe()
+	if err != nil {
+		t.Fatalf("creating pipe: %v", err)
+	}
+	defer r.Close()
+
+	// Write input and close writer so ReadString sees EOF after newline.
+	if _, err := w.WriteString(input); err != nil {
+		t.Fatalf("writing to pipe: %v", err)
+	}
+	w.Close()
+
+	origStdin := os.Stdin
+	os.Stdin = r
+	defer func() { os.Stdin = origStdin }()
+
+	ClearTestSSHPort() // ensure we go through the real path
+	term := terminal.New()
+	return PromptSSHPort(term)
+}
+
+func TestPromptSSHPort(t *testing.T) {
+	tests := []struct {
+		name  string
+		input string
+		want  int32
+	}{
+		{"Default", "\n", 22},
+		{"CustomPort", "2222\n", 2222},
+		{"MinPort", "1\n", 1},
+		{"MaxPort", "65535\n", 65535},
+		{"RetryAfterOutOfRange", "99999\n22\n", 22},
+		{"RetryAfterZero", "0\n443\n", 443},
+		{"RetryAfterNonNumeric", "abc\n8080\n", 8080},
+		{"RetryAfterNegative", "-1\n22\n", 22},
+		{"RetryMultipleThenValid", "foo\n99999\n22\n", 22},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			port, err := promptSSHPortWithInput(t, tt.input)
+			if err != nil {
+				t.Fatalf("unexpected error: %v", err)
+			}
+			if port != tt.want {
+				t.Errorf("expected port %d, got %d", tt.want, port)
+			}
+		})
+	}
+}
+
 func TestInstallAuthorizedKey_EmptyUserID_UsesPrefix(t *testing.T) {
 	u := tempUser(t)