fix enable-ssh port collision

Author: patelspratik

pkg/cmd/enablessh/enablessh.go

@@ -8,7 +8,11 @@ import (
 	"os/exec"
 	"os/user"
 
+	nodev1 "buf.build/gen/go/brevdev/devplane/protocolbuffers/go/devplaneapi/v1"
+	"connectrpc.com/connect"
+
 	"github.com/brevdev/brev-cli/pkg/cmd/register"
+	"github.com/brevdev/brev-cli/pkg/config"
 	"github.com/brevdev/brev-cli/pkg/entity"
 	breverrors "github.com/brevdev/brev-cli/pkg/errors"
 	"github.com/brevdev/brev-cli/pkg/externalnode"
@@ -99,14 +103,24 @@ func enableSSH(
 	t.Vprintf("  Linux user: %s\n", linuxUsername)
 	t.Vprint("")
 
-	t.Vprint("")
-	port, err := register.PromptSSHPort(t)
+	// Check if the node already has an SSH port allocated (e.g. for another linux user)
+	port, err := existingSSHPort(ctx, deps, tokenProvider, reg)
 	if err != nil {
-		return fmt.Errorf("SSH port: %w", err)
+		t.Vprintf("  %s\n", t.Yellow(fmt.Sprintf("Warning: could not check for existing ports: %v", err)))
 	}
 
-	if err := register.OpenSSHPort(ctx, t, deps.nodeClients, tokenProvider, reg, port); err != nil {
-		return fmt.Errorf("enable SSH failed: %w", err)
+	if port != 0 {
+		t.Vprintf("  Using existing SSH port %d.\n", port)
+	} else {
+		t.Vprint("")
+		port, err = register.PromptSSHPort(t)
+		if err != nil {
+			return fmt.Errorf("SSH port: %w", err)
+		}
+
+		if err := register.OpenSSHPort(ctx, t, deps.nodeClients, tokenProvider, reg, port); err != nil {
+			return fmt.Errorf("enable SSH failed: %w", err)
+		}
 	}
 
 	if err := register.SetupAndRegisterNodeSSHAccess(ctx, t, deps.nodeClients, tokenProvider, reg, brevUser, linuxUsername); err != nil {
@@ -117,6 +131,27 @@ func enableSSH(
 	return nil
 }
 
+// existingSSHPort calls GetNode and returns the PortNumber of an already-allocated
+// SSH port, or 0 if none exists
+func existingSSHPort(ctx context.Context, deps enableSSHDeps, tokenProvider externalnode.TokenProvider, reg *register.DeviceRegistration) (int32, error) {
+	client := deps.nodeClients.NewNodeClient(tokenProvider, config.GlobalConfig.GetBrevPublicAPIURL())
+	resp, err := client.GetNode(ctx, connect.NewRequest(&nodev1.GetNodeRequest{
+		ExternalNodeId: reg.ExternalNodeID,
+		OrganizationId: reg.OrgID,
+	}))
+	if err != nil {
+		return 0, fmt.Errorf("error retrieving node: %w", err)
+	}
+
+	for _, p := range resp.Msg.GetExternalNode().GetPorts() {
+		// TODO if we ever allow more than one SSH port, this should be modified
+		if p.GetProtocol() == nodev1.PortProtocol_PORT_PROTOCOL_SSH {
+			return p.GetPortNumber(), nil
+		}
+	}
+	return 0, nil
+}
+
 // checkSSHDaemon prints a warning if neither "ssh" nor "sshd" systemd services
 // appear to be active. It never returns an error — it is best-effort.
 func checkSSHDaemon(t *terminal.Terminal) {

pkg/cmd/enablessh/enablessh_test.go

@@ -1,13 +1,21 @@
 package enablessh
 
 import (
+	"context"
+	"fmt"
+	"net/http/httptest"
 	"os"
 	"os/user"
 	"path/filepath"
 	"strings"
 	"testing"
 
+	nodev1connect "buf.build/gen/go/brevdev/devplane/connectrpc/go/devplaneapi/v1/devplaneapiv1connect"
+	nodev1 "buf.build/gen/go/brevdev/devplane/protocolbuffers/go/devplaneapi/v1"
+	"connectrpc.com/connect"
+
 	"github.com/brevdev/brev-cli/pkg/cmd/register"
+	"github.com/brevdev/brev-cli/pkg/externalnode"
 )
 
 // tempUser returns a *user.User whose HomeDir points to a temporary directory.
@@ -413,3 +421,104 @@ func Test_InstallThenRemoveSpecificKey_RollbackScenario(t *testing.T) {
 		t.Errorf("Alice's key was removed during Bob's rollback:\n%s", result)
 	}
 }
+
+// --- existingSSHPort ---
+
+type mockNodeClientFactory struct{ serverURL string }
+
+func (m mockNodeClientFactory) NewNodeClient(provider externalnode.TokenProvider, _ string) nodev1connect.ExternalNodeServiceClient {
+	return register.NewNodeServiceClient(provider, m.serverURL)
+}
+
+type mockEnableSSHStore struct {
+	token string
+}
+
+func (m *mockEnableSSHStore) GetCurrentUser() (interface{}, error) { return nil, nil }
+func (m *mockEnableSSHStore) GetAccessToken() (string, error)      { return m.token, nil }
+
+// fakeNodeService implements the server side of ExternalNodeService for testing.
+type fakeNodeService struct {
+	nodev1connect.UnimplementedExternalNodeServiceHandler
+	getNodeFn func(*nodev1.GetNodeRequest) (*nodev1.GetNodeResponse, error)
+}
+
+func (f *fakeNodeService) GetNode(_ context.Context, req *connect.Request[nodev1.GetNodeRequest]) (*connect.Response[nodev1.GetNodeResponse], error) {
+	resp, err := f.getNodeFn(req.Msg)
+	if err != nil {
+		return nil, err
+	}
+	return connect.NewResponse(resp), nil
+}
+
+func startFakeServer(t *testing.T, svc *fakeNodeService) (enableSSHDeps, *httptest.Server) {
+	t.Helper()
+	_, handler := nodev1connect.NewExternalNodeServiceHandler(svc)
+	server := httptest.NewServer(handler)
+	t.Cleanup(server.Close)
+	return enableSSHDeps{
+		nodeClients: mockNodeClientFactory{serverURL: server.URL},
+	}, server
+}
+
+func Test_existingSSHPort(t *testing.T) {
+	ssh := nodev1.PortProtocol_PORT_PROTOCOL_SSH
+	tcp := nodev1.PortProtocol_PORT_PROTOCOL_TCP
+
+	tests := []struct {
+		name     string
+		resp     *nodev1.GetNodeResponse
+		rpcErr   error
+		wantPort int32
+		wantErr  bool
+	}{
+		{
+			name: "ReturnsExistingPort",
+			resp: &nodev1.GetNodeResponse{ExternalNode: &nodev1.ExternalNode{
+				Ports: []*nodev1.Port{{Protocol: ssh, PortNumber: 2222}},
+			}},
+			wantPort: 2222,
+		},
+		{
+			name:     "ReturnsZeroWhenNoPorts",
+			resp:     &nodev1.GetNodeResponse{ExternalNode: &nodev1.ExternalNode{}},
+			wantPort: 0,
+		},
+		{
+			name:    "ReturnsErrorOnRPCFailure",
+			rpcErr:  connect.NewError(connect.CodeInternal, fmt.Errorf("server error")),
+			wantErr: true,
+		},
+		{
+			name: "IgnoresNonSSHPorts",
+			resp: &nodev1.GetNodeResponse{ExternalNode: &nodev1.ExternalNode{
+				Ports: []*nodev1.Port{
+					{Protocol: tcp, PortNumber: 8080},
+					{Protocol: ssh, PortNumber: 3333},
+				},
+			}},
+			wantPort: 3333,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			svc := &fakeNodeService{
+				getNodeFn: func(_ *nodev1.GetNodeRequest) (*nodev1.GetNodeResponse, error) {
+					return tt.resp, tt.rpcErr
+				},
+			}
+			deps, _ := startFakeServer(t, svc)
+			store := &mockEnableSSHStore{token: "tok"}
+			reg := &register.DeviceRegistration{ExternalNodeID: "unode_abc"}
+
+			port, err := existingSSHPort(context.Background(), deps, store, reg)
+			if (err != nil) != tt.wantErr {
+				t.Fatalf("err = %v, wantErr = %v", err, tt.wantErr)
+			}
+			if port != tt.wantPort {
+				t.Errorf("port = %d, want %d", port, tt.wantPort)
+			}
+		})
+	}
+}