refactor: narrow current user SSH key contract

Author: patelspratik

pkg/cmd/copy/copy.go

@@ -35,7 +35,7 @@ type CopyStore interface {
 	GetWorkspaces(organizationID string, options *store.GetWorkspacesOptions) ([]entity.Workspace, error)
 	StartWorkspace(workspaceID string) (*entity.Workspace, error)
 	GetWorkspace(workspaceID string) (*entity.Workspace, error)
-	GetCurrentUserKeys() (*entity.UserKeys, error)
+	GetCurrentUserSSHPrivateKey() (string, error)
 	GetAccessToken() (string, error)
 }
 

pkg/cmd/proxy/proxy.go

@@ -24,7 +24,7 @@ type ProxyStore interface {
 	huproxyclient.HubProxyStore
 	GetWorkspace(workspaceID string) (*entity.Workspace, error)
 	WritePrivateKey(pem string) error
-	GetCurrentUserKeys() (*entity.UserKeys, error)
+	GetCurrentUserSSHPrivateKey() (string, error)
 }
 
 func NewCmdProxy(t *terminal.Terminal, store ProxyStore) *cobra.Command {
@@ -149,11 +149,11 @@ func checkWorkspaceInfraVersionOrErr(workspace *entity.Workspace) error {
 }
 
 func WriteUserPrivateKey(store ProxyStore) error {
-	keys, err := store.GetCurrentUserKeys()
+	privateKey, err := store.GetCurrentUserSSHPrivateKey()
 	if err != nil {
 		return breverrors.WrapAndTrace(err)
 	}
-	err = store.WritePrivateKey(keys.PrivateKey)
+	err = store.WritePrivateKey(privateKey)
 	if err != nil {
 		return breverrors.WrapAndTrace(err)
 	}

pkg/cmd/proxy/proxy_test.go

@@ -3,11 +3,44 @@ package proxy
 import (
 	"testing"
 
+	"github.com/brevdev/brev-cli/pkg/entity"
 	"github.com/hashicorp/go-version"
 	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
 )
 
+type privateKeyProxyStore struct {
+	privateKey        string
+	writtenPrivateKey string
+}
+
+func (s *privateKeyProxyStore) GetAuthTokens() (*entity.AuthTokens, error) {
+	return nil, nil
+}
+
+func (s *privateKeyProxyStore) GetWorkspace(string) (*entity.Workspace, error) {
+	return nil, nil
+}
+
+func (s *privateKeyProxyStore) GetCurrentUserSSHPrivateKey() (string, error) {
+	return s.privateKey, nil
+}
+
+func (s *privateKeyProxyStore) WritePrivateKey(privateKey string) error {
+	s.writtenPrivateKey = privateKey
+	return nil
+}
+
 func TestVersionParsing(t *testing.T) {
 	_, err := version.NewVersion("abadfjladsf")
 	assert.NotNil(t, err)
 }
+
+func TestWriteUserPrivateKeyWritesCurrentUsersPrivateKey(t *testing.T) {
+	store := &privateKeyProxyStore{privateKey: "private"}
+
+	err := WriteUserPrivateKey(store)
+
+	require.NoError(t, err)
+	require.Equal(t, "private", store.writtenPrivateKey)
+}

pkg/cmd/refresh/refresh.go

@@ -29,7 +29,7 @@ type RefreshStore interface {
 	ssh.ConfigUpdaterStore
 	ssh.SSHConfigurerV2Store
 	GetCurrentUser() (*entity.User, error)
-	GetCurrentUserKeys() (*entity.UserKeys, error)
+	GetCurrentUserSSHPrivateKey() (string, error)
 	GetActiveOrganizationOrDefault() (*entity.Organization, error)
 	GetAccessToken() (string, error)
 	Chmod(string, fs.FileMode) error
@@ -155,12 +155,12 @@ func GetConfigUpdater(store RefreshStore) (*ssh.ConfigUpdater, error) {
 		return nil, breverrors.WrapAndTrace(err)
 	}
 
-	keys, err := store.GetCurrentUserKeys()
+	privateKey, err := store.GetCurrentUserSSHPrivateKey()
 	if err != nil {
 		return nil, breverrors.WrapAndTrace(err)
 	}
 
-	cu := ssh.NewConfigUpdater(store, configs, keys.PrivateKey)
+	cu := ssh.NewConfigUpdater(store, configs, privateKey)
 	cu.ExternalNodes = getExternalNodeSSHEntries(store)
 
 	return cu, nil

pkg/cmd/runtasks/runtasks.go

@@ -48,7 +48,7 @@ type RunTasksStore interface {
 	ssh.SSHConfigurerV2Store
 	tasks.RunTaskAsDaemonStore
 	GetCurrentUser() (*entity.User, error)
-	GetCurrentUserKeys() (*entity.UserKeys, error)
+	GetCurrentUserSSHPrivateKey() (string, error)
 }
 
 func RunTasks(_ *terminal.Terminal, store RunTasksStore, detached bool) error {
@@ -77,12 +77,12 @@ func getDefaultTasks(store RunTasksStore) ([]tasks.Task, error) {
 	}
 
 	// get private key and set here
-	keys, err := store.GetCurrentUserKeys()
+	privateKey, err := store.GetCurrentUserSSHPrivateKey()
 	if err != nil {
 		return nil, breverrors.WrapAndTrace(err)
 	}
 
-	cu := ssh.NewConfigUpdater(store, configs, keys.PrivateKey)
+	cu := ssh.NewConfigUpdater(store, configs, privateKey)
 
 	return []tasks.Task{cu}, nil
 }

pkg/cmd/tasks/tasks.go

@@ -19,7 +19,7 @@ type TaskMap map[string]tasks.Task
 var all bool // used for run command
 
 type TaskStore interface {
-	GetCurrentUserKeys() (*entity.UserKeys, error)
+	GetCurrentUserSSHPrivateKey() (string, error)
 	CopyBin(targetBin string) error
 	WriteString(path, data string) error
 	GetOrCreateFile(path string) (afero.File, error)

pkg/ssh/tasks.go

@@ -2,15 +2,14 @@ package ssh
 
 import (
 	"github.com/brevdev/brev-cli/pkg/autostartconf"
-	"github.com/brevdev/brev-cli/pkg/entity"
 	breverrors "github.com/brevdev/brev-cli/pkg/errors"
 	"github.com/brevdev/brev-cli/pkg/tasks"
 )
 
 type SSHConfigurerTaskStore interface {
 	ConfigUpdaterStore
 	SSHConfigurerV2Store
-	GetCurrentUserKeys() (*entity.UserKeys, error)
+	GetCurrentUserSSHPrivateKey() (string, error)
 }
 
 type SSHConfigurerTask struct {
@@ -29,12 +28,12 @@ func (sct SSHConfigurerTask) Run() error {
 		return breverrors.WrapAndTrace(err)
 	}
 
-	keys, err := sct.Store.GetCurrentUserKeys()
+	privateKey, err := sct.Store.GetCurrentUserSSHPrivateKey()
 	if err != nil {
 		return breverrors.WrapAndTrace(err)
 	}
 
-	cu := NewConfigUpdater(sct.Store, configs, keys.PrivateKey)
+	cu := NewConfigUpdater(sct.Store, configs, privateKey)
 	err = tasks.RunTasks([]tasks.Task{cu})
 	if err != nil {
 		return breverrors.WrapAndTrace(err)

pkg/store/devplane.go

@@ -17,7 +17,7 @@ import (
 
 type devPlaneUserService interface {
 	GetCurrentUser(context.Context) (*nodev1.User, string, error)
-	GetCurrentUserKeys(context.Context) (*entity.UserKeys, error)
+	GetCurrentUserSSHPrivateKey(context.Context) (string, error)
 	GetUser(context.Context, string) (*nodev1.User, error)
 }
 
@@ -88,8 +88,8 @@ func (s *generatedDevPlaneUserService) GetCurrentUser(ctx context.Context) (*nod
 	return response.Msg.GetUser(), "", nil
 }
 
-func (s *generatedDevPlaneUserService) GetCurrentUserKeys(context.Context) (*entity.UserKeys, error) {
-	return nil, fmt.Errorf("GetCurrentUserSSHKeys requires the published Set 1 Buf revision")
+func (s *generatedDevPlaneUserService) GetCurrentUserSSHPrivateKey(context.Context) (string, error) {
+	return "", fmt.Errorf("GetCurrentUserSSHPrivateKey requires the published Set 1 Buf revision")
 }
 
 func (s *generatedDevPlaneUserService) GetUser(ctx context.Context, userID string) (*nodev1.User, error) {

pkg/store/devplane_migration_test.go

@@ -13,7 +13,7 @@ import (
 type mockDevPlaneUserService struct {
 	currentUser  *nodev1.User
 	publicSSHKey string
-	keys         *entity.UserKeys
+	privateKey   string
 	getUser      *nodev1.User
 	gotUserID    string
 }
@@ -22,8 +22,8 @@ func (m *mockDevPlaneUserService) GetCurrentUser(context.Context) (*nodev1.User,
 	return m.currentUser, m.publicSSHKey, nil
 }
 
-func (m *mockDevPlaneUserService) GetCurrentUserKeys(context.Context) (*entity.UserKeys, error) {
-	return m.keys, nil
+func (m *mockDevPlaneUserService) GetCurrentUserSSHPrivateKey(context.Context) (string, error) {
+	return m.privateKey, nil
 }
 
 func (m *mockDevPlaneUserService) GetUser(_ context.Context, userID string) (*nodev1.User, error) {
@@ -73,7 +73,7 @@ func TestIdentityStoresUseDevPlaneServices(t *testing.T) {
 	users := &mockDevPlaneUserService{
 		currentUser:  apiUser,
 		publicSSHKey: "ssh-rsa public",
-		keys:         &entity.UserKeys{PublicKey: "public", PrivateKey: "private"},
+		privateKey:   "private",
 		getUser:      apiUser,
 	}
 	store := MakeMockAuthHTTPStore().withDevPlaneServices(&devPlaneServices{user: users})
@@ -88,16 +88,23 @@ func TestIdentityStoresUseDevPlaneServices(t *testing.T) {
 	require.Equal(t, entity.Admin, current.GlobalUserType)
 	require.Equal(t, true, current.OnboardingData["usedCLI"])
 
-	keys, err := store.GetCurrentUserKeys()
+	privateKey, err := store.GetCurrentUserSSHPrivateKey()
 	require.NoError(t, err)
-	require.Equal(t, "private", keys.PrivateKey)
+	require.Equal(t, "private", privateKey)
 
 	target, err := store.GetUserByID("user-2")
 	require.NoError(t, err)
 	require.Equal(t, "user-1", target.ID)
 	require.Equal(t, "user-2", users.gotUserID)
 }
 
+func TestGeneratedDevPlaneUserServiceGetCurrentUserSSHPrivateKeyFailsClosedUntilBufUpdate(t *testing.T) {
+	privateKey, err := (&generatedDevPlaneUserService{}).GetCurrentUserSSHPrivateKey(context.Background())
+
+	require.Empty(t, privateKey)
+	require.EqualError(t, err, "GetCurrentUserSSHPrivateKey requires the published Set 1 Buf revision")
+}
+
 func TestOrganizationStoresUseCurrentUserAccessesAndDevPlane(t *testing.T) {
 	users := &mockDevPlaneUserService{currentUser: &nodev1.User{
 		UserId: "user-1",

pkg/store/user.go

@@ -59,10 +59,10 @@ func (s AuthHTTPStore) GetCurrentUserID() (string, error) {
 	return user.ID, nil
 }
 
-func (s AuthHTTPStore) GetCurrentUserKeys() (*entity.UserKeys, error) {
-	result, err := s.devPlaneServiceClients().user.GetCurrentUserKeys(context.Background())
+func (s AuthHTTPStore) GetCurrentUserSSHPrivateKey() (string, error) {
+	result, err := s.devPlaneServiceClients().user.GetCurrentUserSSHPrivateKey(context.Background())
 	if err != nil {
-		return nil, breverrors.WrapAndTrace(err)
+		return "", breverrors.WrapAndTrace(err)
 	}
 	return result, nil
 }