BRE2-987: Use SSH Access instead of Public DNS for 'brev shell'

pkg/cmd/refresh/refresh.go

@@ -6,7 +6,6 @@ import (
 	"fmt"
 	"io"
 	"io/fs"
-	"log"
 	"sync"
 
 	nodev1 "buf.build/gen/go/brevdev/devplane/protocolbuffers/go/devplaneapi/v1"
@@ -160,7 +159,7 @@ func GetConfigUpdater(store RefreshStore) (*ssh.ConfigUpdater, error) {
 		return nil, breverrors.WrapAndTrace(err)
 	}
 
-	cu := ssh.NewConfigUpdater(store, configs, keys.PrivateKey)
+	cu := ssh.NewConfigUpdater(workspaceSSHStore{RefreshStore: store}, configs, keys.PrivateKey)
 	cu.ExternalNodes = getExternalNodeSSHEntries(store)
 
 	return cu, nil
@@ -171,13 +170,11 @@ func GetConfigUpdater(store RefreshStore) (*ssh.ConfigUpdater, error) {
 func getExternalNodeSSHEntries(store RefreshStore) []ssh.ExternalNodeSSHEntry {
 	org, err := store.GetActiveOrganizationOrDefault()
 	if err != nil {
-		log.Printf("external nodes: skipping (no org): %v", err)
 		return nil
 	}
 
 	user, err := store.GetCurrentUser()
 	if err != nil {
-		log.Printf("external nodes: skipping (no user): %v", err)
 		return nil
 	}
 
@@ -186,7 +183,6 @@ func getExternalNodeSSHEntries(store RefreshStore) []ssh.ExternalNodeSSHEntry {
 		OrganizationId: org.ID,
 	}))
 	if err != nil {
-		log.Printf("external nodes: skipping (list failed): %v", err)
 		return nil
 	}
 

pkg/cmd/refresh/sshaccess.go

@@ -0,0 +1,159 @@
+package refresh
+
+import (
+	"context"
+	"fmt"
+	"log"
+	"time"
+
+	devplanev1 "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"
+)
+
+const sshAccessLookupTimeout = 10 * time.Second
+
+type environmentSSHClient interface {
+	GetEnvironment(context.Context, *connect.Request[devplanev1.GetEnvironmentRequest]) (*connect.Response[devplanev1.GetEnvironmentResponse], error)
+	GetNetworkInfo(context.Context, *connect.Request[devplanev1.EnvironmentServiceGetNetworkInfoRequest]) (*connect.Response[devplanev1.EnvironmentServiceGetNetworkInfoResponse], error)
+}
+
+type workspaceSSHStore struct {
+	RefreshStore
+}
+
+func (s workspaceSSHStore) GetContextWorkspaces() ([]entity.Workspace, error) {
+	workspaces, err := s.RefreshStore.GetContextWorkspaces()
+	if err != nil {
+		return nil, breverrors.WrapAndTrace(err)
+	}
+
+	user, err := s.GetCurrentUser()
+	if err != nil {
+		log.Printf("workspace SSH access: using legacy configuration (current user lookup failed): %v", err)
+		return workspaces, nil
+	}
+
+	client := register.NewEnvironmentServiceClient(s, config.GlobalConfig.GetBrevPublicAPIURL())
+	ctx, cancel := context.WithTimeout(context.Background(), sshAccessLookupTimeout)
+	defer cancel()
+
+	return enrichWorkspacesWithSSHAccess(ctx, client, user.ID, workspaces), nil
+}
+
+func enrichWorkspacesWithSSHAccess(ctx context.Context, client environmentSSHClient, userID string, workspaces []entity.Workspace) []entity.Workspace {
+	for i := range workspaces {
+		if workspaces[i].Status != entity.Running {
+			continue
+		}
+
+		workspace, err := resolveWorkspaceSSH(ctx, client, userID, workspaces[i])
+		if err != nil {
+			log.Printf("workspace SSH access: using legacy configuration for %s: %v", workspaces[i].ID, err)
+			continue
+		}
+		workspaces[i] = workspace
+	}
+	return workspaces
+}
+
+func resolveWorkspaceSSH(
+	ctx context.Context,
+	client environmentSSHClient,
+	userID string,
+	workspace entity.Workspace,
+) (entity.Workspace, error) {
+	environmentRes, err := client.GetEnvironment(ctx, connect.NewRequest(&devplanev1.GetEnvironmentRequest{
+		EnvironmentId: workspace.ID,
+		AttachedDataOptions: &devplanev1.GetEnvironmentAttachedDataOptions{
+			Instance:  true,
+			SshAccess: true,
+		},
+	}))
+	if err != nil {
+		return workspace, fmt.Errorf("get environment: %w", err)
+	}
+	if environmentRes == nil || environmentRes.Msg == nil || environmentRes.Msg.GetEnvironment() == nil {
+		return workspace, fmt.Errorf("get environment: empty response")
+	}
+
+	// Use the ssh access information to determine the SSH target, rather than the public IP, DNS, or other information
+	// returned by the initial workspace query.
+	environment := environmentRes.Msg.GetEnvironment()
+	access := findUserSSHAccess(environment.GetSshAccess(), userID)
+	if access == nil {
+		return workspace, nil
+	}
+	if access.GetLinuxUser() == "" {
+		return workspace, fmt.Errorf("SSH access has no Linux user")
+	}
+
+	// Using the port ID, we can lookup the network information to get the hostname and port number of the SSH target.
+	networkRes, err := client.GetNetworkInfo(ctx, connect.NewRequest(&devplanev1.EnvironmentServiceGetNetworkInfoRequest{
+		EnvironmentId: workspace.ID,
+	}))
+	if err != nil {
+		return workspace, fmt.Errorf("get network info: %w", err)
+	}
+	if networkRes == nil || networkRes.Msg == nil {
+		return workspace, fmt.Errorf("get network info: empty response")
+	}
+
+	port := findNetworkPort(networkRes.Msg.GetNetworkInfo(), access.GetPortId())
+	if port == nil {
+		return workspace, fmt.Errorf("SSH access port %q not found", access.GetPortId())
+	}
+	if port.GetHostname() == "" || port.GetPortNumber() == 0 {
+		return workspace, fmt.Errorf("SSH access port %q has no endpoint", access.GetPortId())
+	}
+
+	// Honor the SSH access information as the source of truth for the SSH target.
+	workspace.SSHHostname = port.GetHostname()
+	workspace.SSHPort = int(port.GetPortNumber())
+	workspace.SSHUser = access.GetLinuxUser()
+	workspace.SSHProxyHostname = ""
+
+	// To support the "--host" fallback, preserve the legacy hostname information returned by the initial workspace query.
+	if providerHostname := providerSSHHostname(environment.GetInstance(), port.GetHostname()); providerHostname != "" {
+		workspace.HostSSHHostname = providerHostname
+		workspace.HostSSHProxyHostname = ""
+	}
+	return workspace, nil
+}
+
+func findUserSSHAccess(accesses []*devplanev1.SSHAccess, userID string) *devplanev1.SSHAccess {
+	// Technically it is possible that multiple SSHAccess entries exist for a single user+environment. This is typically only
+	// for external nodes, which allow for multiple sshd processes to be targeted. For the normal environment flow, the below
+	// is a best-effort approach to find the *first* applicable entry.
+	for _, access := range accesses {
+		if access.GetUserId() == userID && access.GetPortId() != "" {
+			return access
+		}
+	}
+	return nil
+}
+
+func findNetworkPort(networkInfo *devplanev1.EnvironmentNetworkInfo, portID string) *devplanev1.Port {
+	for _, port := range networkInfo.GetPorts() {
+		if port.GetPortId() == portID {
+			return port
+		}
+	}
+	return nil
+}
+
+func providerSSHHostname(instance *devplanev1.Instance, workloadHostname string) string {
+	if instance == nil {
+		return ""
+	}
+	for _, hostname := range []string{instance.GetPublicIp(), instance.GetPublicDns(), instance.GetHostname()} {
+		if hostname != "" && hostname != workloadHostname {
+			return hostname
+		}
+	}
+	return ""
+}

pkg/cmd/refresh/sshaccess_test.go

@@ -0,0 +1,140 @@
+package refresh
+
+import (
+	"context"
+	"errors"
+	"testing"
+
+	devplanev1 "buf.build/gen/go/brevdev/devplane/protocolbuffers/go/devplaneapi/v1"
+	"connectrpc.com/connect"
+	"github.com/google/go-cmp/cmp"
+
+	"github.com/brevdev/brev-cli/pkg/entity"
+)
+
+type stubEnvironmentSSHClient struct {
+	environment    *devplanev1.Environment
+	networkInfo    *devplanev1.EnvironmentNetworkInfo
+	err            error
+	request        *devplanev1.GetEnvironmentRequest
+	networkRequest *devplanev1.EnvironmentServiceGetNetworkInfoRequest
+}
+
+func (s *stubEnvironmentSSHClient) GetEnvironment(
+	_ context.Context,
+	req *connect.Request[devplanev1.GetEnvironmentRequest],
+) (*connect.Response[devplanev1.GetEnvironmentResponse], error) {
+	s.request = req.Msg
+	if s.err != nil {
+		return nil, s.err
+	}
+	return connect.NewResponse(&devplanev1.GetEnvironmentResponse{Environment: s.environment}), nil
+}
+
+func (s *stubEnvironmentSSHClient) GetNetworkInfo(
+	_ context.Context,
+	req *connect.Request[devplanev1.EnvironmentServiceGetNetworkInfoRequest],
+) (*connect.Response[devplanev1.EnvironmentServiceGetNetworkInfoResponse], error) {
+	s.networkRequest = req.Msg
+	return connect.NewResponse(&devplanev1.EnvironmentServiceGetNetworkInfoResponse{
+		NetworkInfo: s.networkInfo,
+	}), nil
+}
+
+func TestEnrichWorkspacesWithSSHAccess_UsesCurrentUsersPort(t *testing.T) {
+	workspace := entity.Workspace{
+		ID:                   "env-1",
+		Name:                 "container-env",
+		DNS:                  "legacy.example.com",
+		Status:               entity.Running,
+		SSHUser:              "ubuntu",
+		SSHPort:              22,
+		HostSSHUser:          "ubuntu",
+		HostSSHPort:          22,
+		SSHProxyHostname:     "legacy-proxy.example.com",
+		HostSSHProxyHostname: "legacy-host-proxy.example.com",
+	}
+	client := &stubEnvironmentSSHClient{
+		environment: &devplanev1.Environment{
+			Instance: &devplanev1.Instance{
+				SshHostname: "203.0.113.10",
+				SshPort:     22,
+				PublicIp:    "203.0.113.10",
+			},
+			SshAccess: []*devplanev1.SSHAccess{
+				{UserId: "other-user", LinuxUser: "wrong-user", PortId: "other-port"},
+				{UserId: "user-1", LinuxUser: "root", PortId: "ssh-port"},
+			},
+		},
+		networkInfo: &devplanev1.EnvironmentNetworkInfo{
+			Ports: []*devplanev1.Port{
+				{PortId: "other-port", Hostname: strPtr("wrong.example.com"), PortNumber: 49999},
+				{PortId: "ssh-port", Hostname: strPtr("skybridge.example.com"), PortNumber: 41234, ServerPort: 22},
+			},
+		},
+	}
+
+	got := enrichWorkspacesWithSSHAccess(context.Background(), client, "user-1", []entity.Workspace{workspace})
+	want := workspace
+	want.SSHHostname = "skybridge.example.com"
+	want.SSHPort = 41234
+	want.SSHUser = "root"
+	want.SSHProxyHostname = ""
+	want.HostSSHHostname = "203.0.113.10"
+	want.HostSSHProxyHostname = ""
+
+	if diff := cmp.Diff([]entity.Workspace{want}, got); diff != "" {
+		t.Fatalf("unexpected workspace (-want +got): %s", diff)
+	}
+	if client.request.GetEnvironmentId() != workspace.ID {
+		t.Fatalf("requested environment %q, want %q", client.request.GetEnvironmentId(), workspace.ID)
+	}
+	options := client.request.GetAttachedDataOptions()
+	if !options.GetInstance() || !options.GetSshAccess() || options.GetSysUsers() {
+		t.Fatalf("missing SSH attachment options: %+v", options)
+	}
+	if client.networkRequest.GetEnvironmentId() != workspace.ID {
+		t.Fatalf("requested network environment %q, want %q", client.networkRequest.GetEnvironmentId(), workspace.ID)
+	}
+}
+
+func TestEnrichWorkspacesWithSSHAccess_FallsBackOnError(t *testing.T) {
+	workspace := entity.Workspace{
+		ID:      "env-1",
+		Name:    "legacy-env",
+		DNS:     "legacy.example.com",
+		Status:  entity.Running,
+		SSHUser: "ubuntu",
+		SSHPort: 22,
+	}
+	client := &stubEnvironmentSSHClient{err: errors.New("dev-plane unavailable")}
+
+	got := enrichWorkspacesWithSSHAccess(context.Background(), client, "user-1", []entity.Workspace{workspace})
+	if diff := cmp.Diff([]entity.Workspace{workspace}, got); diff != "" {
+		t.Fatalf("legacy workspace changed (-want +got): %s", diff)
+	}
+}
+
+func TestEnrichWorkspacesWithSSHAccess_FallsBackWithoutPortBackedAccess(t *testing.T) {
+	workspace := entity.Workspace{
+		ID:      "env-1",
+		DNS:     "legacy.example.com",
+		Status:  entity.Running,
+		SSHUser: "ubuntu",
+		SSHPort: 22,
+	}
+	client := &stubEnvironmentSSHClient{environment: &devplanev1.Environment{
+		Instance: &devplanev1.Instance{},
+		SshAccess: []*devplanev1.SSHAccess{
+			{UserId: "user-1", LinuxUser: "root"},
+		},
+	}}
+
+	got := enrichWorkspacesWithSSHAccess(context.Background(), client, "user-1", []entity.Workspace{workspace})
+	if diff := cmp.Diff([]entity.Workspace{workspace}, got); diff != "" {
+		t.Fatalf("legacy workspace changed (-want +got): %s", diff)
+	}
+	if client.networkRequest != nil {
+		t.Fatal("network info should not be fetched without port-backed access")
+	}
+}

pkg/cmd/register/rpcclient.go

@@ -53,6 +53,14 @@ func NewNodeServiceClient(provider externalnode.TokenProvider, baseURL string) n
 	)
 }
 
+// NewEnvironmentServiceClient creates an authenticated ConnectRPC EnvironmentServiceClient.
+func NewEnvironmentServiceClient(provider externalnode.TokenProvider, baseURL string) nodev1connect.EnvironmentServiceClient {
+	return nodev1connect.NewEnvironmentServiceClient(
+		newAuthenticatedHTTPClient(provider),
+		baseURL,
+	)
+}
+
 // toProtoNodeSpec converts the local HardwareProfile (used for collection, display,
 // persistence) to the generated proto NodeSpec for RPC calls.
 func toProtoNodeSpec(hw *HardwareProfile) *nodev1.NodeSpec {