Skip to content

Commit 6e4c50b

Browse files
patelspratikpatelspratik
committed
fixing org set
1 parent c097c8d commit 6e4c50b

10 files changed

Lines changed: 192 additions & 14 deletions

File tree

‎pkg/auth/auth.go‎

Lines changed: 25 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -107,6 +107,10 @@ type TokenProvider interface {
107107
GetAccessToken() (string, error)
108108
}
109109

110+
type APIKeyOrgProvider interface {
111+
GetAuthTokens() (*entity.AuthTokens, error)
112+
}
113+
110114
func IsBrevAPIKey(token string) bool {
111115
return strings.HasPrefix(strings.TrimSpace(token), BrevAPIKeyPrefix)
112116
}
@@ -119,6 +123,21 @@ func IsAPIKeyAuthStore(tokenProvider TokenProvider) bool {
119123
return IsBrevAPIKey(token)
120124
}
121125

126+
func GetAPIKeyOrgID(authTokensProvider APIKeyOrgProvider) (string, error) {
127+
tokens, err := authTokensProvider.GetAuthTokens()
128+
if err != nil {
129+
return "", breverrors.WrapAndTrace(err)
130+
}
131+
if tokens == nil {
132+
return "", breverrors.NewValidationError("api key auth requires an org id; run brev login --api-key <api-key> --org-id <org-id>")
133+
}
134+
orgID := strings.TrimSpace(tokens.APIKeyOrgID)
135+
if orgID == "" {
136+
return "", breverrors.NewValidationError("api key auth requires an org id; run brev login --api-key <api-key> --org-id <org-id>")
137+
}
138+
return orgID, nil
139+
}
140+
122141
func NewAuth(authStore AuthStore, oauth OAuth) *Auth {
123142
return &Auth{
124143
authStore: authStore,
@@ -248,14 +267,18 @@ func (t Auth) LoginWithToken(token string) error {
248267
return nil
249268
}
250269

251-
func (t Auth) LoginWithAPIKey(apiKey string) error {
270+
func (t Auth) LoginWithAPIKey(apiKey string, orgID string) error {
252271
apiKey = strings.TrimSpace(apiKey)
253272
if apiKey == "" {
254273
return breverrors.NewValidationError("api key is empty")
255274
}
256275
if !strings.HasPrefix(apiKey, BrevAPIKeyPrefix) {
257276
return breverrors.NewValidationError(fmt.Sprintf("api key must start with %s", BrevAPIKeyPrefix))
258277
}
278+
orgID = strings.TrimSpace(orgID)
279+
if orgID == "" {
280+
return breverrors.NewValidationError("org-id is required with api-key")
281+
}
259282

260283
tokens, err := t.getSavedTokensOrNil()
261284
if err != nil {
@@ -265,6 +288,7 @@ func (t Auth) LoginWithAPIKey(apiKey string) error {
265288
tokens = &entity.AuthTokens{}
266289
}
267290
tokens.APIKey = apiKey
291+
tokens.APIKeyOrgID = orgID
268292

269293
err = t.authStore.SaveAuthTokens(*tokens)
270294
if err != nil {

‎pkg/auth/auth_test.go‎

Lines changed: 18 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -136,11 +136,12 @@ func TestLoginWithAPIKey_SavesTypedCredential(t *testing.T) {
136136
oauth: &MockOauth{},
137137
}
138138

139-
err := a.LoginWithAPIKey("brev_api_test-key")
139+
err := a.LoginWithAPIKey("brev_api_test-key", "org-test")
140140
assert.NoError(t, err)
141141
assert.True(t, s.didSave)
142142
assert.Equal(t, entity.AuthTokens{
143-
APIKey: "brev_api_test-key",
143+
APIKey: "brev_api_test-key",
144+
APIKeyOrgID: "org-test",
144145
}, s.saved)
145146
}
146147

@@ -154,12 +155,13 @@ func TestLoginWithAPIKey_PreservesExistingJWT(t *testing.T) {
154155
oauth: &MockOauth{},
155156
}
156157

157-
err := a.LoginWithAPIKey("brev_api_test-key")
158+
err := a.LoginWithAPIKey("brev_api_test-key", "org-test")
158159
assert.NoError(t, err)
159160
assert.Equal(t, entity.AuthTokens{
160161
AccessToken: "existing-jwt",
161162
RefreshToken: "existing-refresh",
162163
APIKey: "brev_api_test-key",
164+
APIKeyOrgID: "org-test",
163165
}, s.saved)
164166
}
165167

@@ -170,7 +172,19 @@ func TestLoginWithAPIKey_EmptyKeyReturnsError(t *testing.T) {
170172
oauth: &MockOauth{},
171173
}
172174

173-
err := a.LoginWithAPIKey("")
175+
err := a.LoginWithAPIKey("", "org-test")
176+
assert.Error(t, err)
177+
assert.False(t, s.didSave)
178+
}
179+
180+
func TestLoginWithAPIKey_EmptyOrgIDReturnsError(t *testing.T) {
181+
s := MockAuthStore{}
182+
a := Auth{
183+
authStore: &s,
184+
oauth: &MockOauth{},
185+
}
186+
187+
err := a.LoginWithAPIKey("brev_api_test-key", "")
174188
assert.Error(t, err)
175189
assert.False(t, s.didSave)
176190
}

‎pkg/cmd/login/login.go‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -48,7 +48,7 @@ type LoginStore interface {
4848
type Auth interface {
4949
Login(skipBrowser bool) (*auth.LoginTokens, error)
5050
LoginWithToken(token string) error
51-
LoginWithAPIKey(apiKey string) error
51+
LoginWithAPIKey(apiKey string, orgID string) error
5252
}
5353

5454
// loginStore must be a no prompt store
@@ -218,7 +218,7 @@ func (o LoginOptions) doApiKeyLogin(t *terminal.Terminal, loginToken string, api
218218
if orgID == "" {
219219
return breverrors.NewValidationError("org-id is required with api-key")
220220
}
221-
if err := o.Auth.LoginWithAPIKey(apiKey); err != nil {
221+
if err := o.Auth.LoginWithAPIKey(apiKey, orgID); err != nil {
222222
return breverrors.WrapAndTrace(err)
223223
}
224224
if err := o.LoginStore.SetDefaultOrganization(&entity.Organization{

‎pkg/cmd/login/login_test.go‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@ import (
1515
type mockLoginAuth struct {
1616
apiKeyCalls int
1717
apiKey string
18+
apiKeyOrgID string
1819
tokenCalls int
1920
loginCalls int
2021
}
@@ -29,9 +30,10 @@ func (m *mockLoginAuth) LoginWithToken(_ string) error {
2930
return nil
3031
}
3132

32-
func (m *mockLoginAuth) LoginWithAPIKey(apiKey string) error {
33+
func (m *mockLoginAuth) LoginWithAPIKey(apiKey string, orgID string) error {
3334
m.apiKeyCalls++
3435
m.apiKey = apiKey
36+
m.apiKeyOrgID = orgID
3537
return nil
3638
}
3739

@@ -118,6 +120,7 @@ func TestRunLoginWithAPIKey_SavesKeyAndOrgWithoutUserOrBackendOrgCalls(t *testin
118120
require.NoError(t, err)
119121
assert.Equal(t, 1, auth.apiKeyCalls)
120122
assert.Equal(t, "brev_api_test-key", auth.apiKey)
123+
assert.Equal(t, "org-test", auth.apiKeyOrgID)
121124
assert.Equal(t, 1, loginStore.setDefaultOrgCalls)
122125
require.NotNil(t, loginStore.defaultOrg)
123126
assert.Equal(t, "org-test", loginStore.defaultOrg.ID)

‎pkg/cmd/ls/ls.go‎

Lines changed: 3 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -44,6 +44,7 @@ type LsStore interface {
4444
GetWorkspace(workspaceID string) (*entity.Workspace, error)
4545
GetOrganizations(options *store.GetOrganizationsOptions) ([]entity.Organization, error)
4646
GetAccessToken() (string, error)
47+
GetAuthTokens() (*entity.AuthTokens, error)
4748
GetInstanceTypes(includeCPU bool) (*gpusearch.InstanceTypesResponse, error)
4849
hello.HelloStore
4950
}
@@ -138,14 +139,11 @@ func getOrgForRunLs(lsStore LsStore, orgflag string, apiKeyAuth bool) (*entity.O
138139
if orgflag != "" {
139140
return nil, breverrors.NewValidationError("api key auth is scoped to the org saved during login; --org is not supported")
140141
}
141-
cachedOrg, err := lsStore.GetCachedActiveOrganizationOrNil()
142+
orgID, err := auth.GetAPIKeyOrgID(lsStore)
142143
if err != nil {
143144
return nil, breverrors.WrapAndTrace(err)
144145
}
145-
if cachedOrg == nil || cachedOrg.ID == "" {
146-
return nil, breverrors.NewValidationError("api key auth requires an active org; run brev login --api-key <api-key> --org-id <org-id>")
147-
}
148-
return cachedOrg, nil
146+
return &entity.Organization{ID: orgID, Name: orgID}, nil
149147
}
150148

151149
if orgflag != "" {

‎pkg/cmd/ls/ls_test.go‎

Lines changed: 58 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,8 @@ type mockLsStore struct {
2222
orgs []entity.Organization
2323
workspaces []entity.Workspace
2424
accessToken string
25+
authTokens *entity.AuthTokens
26+
workspaceOrgID string
2527
currentUserCalls int
2628
getOrganizationsCall int
2729
}
@@ -38,6 +40,10 @@ func (m *mockLsStore) GetAccessToken() (string, error) {
3840
return "tok", nil
3941
}
4042

43+
func (m *mockLsStore) GetAuthTokens() (*entity.AuthTokens, error) {
44+
return m.authTokens, nil
45+
}
46+
4147
func (m *mockLsStore) GetWorkspace(_ string) (*entity.Workspace, error) {
4248
return nil, nil
4349
}
@@ -50,7 +56,8 @@ func (m *mockLsStore) GetCachedActiveOrganizationOrNil() (*entity.Organization,
5056
return m.org, nil
5157
}
5258

53-
func (m *mockLsStore) GetWorkspaces(_ string, _ *store.GetWorkspacesOptions) ([]entity.Workspace, error) {
59+
func (m *mockLsStore) GetWorkspaces(orgID string, _ *store.GetWorkspacesOptions) ([]entity.Workspace, error) {
60+
m.workspaceOrgID = orgID
5461
return m.workspaces, nil
5562
}
5663

@@ -90,6 +97,7 @@ func newTestStore() *mockLsStore {
9097
func TestRunLs_APIKeyJSONSkipsUserAndOrgList(t *testing.T) {
9198
s := newTestStore()
9299
s.accessToken = "brev_api_test-key"
100+
s.authTokens = &entity.AuthTokens{APIKey: "brev_api_test-key", APIKeyOrgID: "org1"}
93101
s.workspaces = []entity.Workspace{
94102
{
95103
ID: "ws1",
@@ -130,6 +138,55 @@ func TestRunLs_APIKeyJSONSkipsUserAndOrgList(t *testing.T) {
130138
}
131139
}
132140

141+
func TestRunLs_APIKeyUsesCredentialOrgNotCachedActiveOrg(t *testing.T) {
142+
s := newTestStore()
143+
s.accessToken = "brev_api_test-key"
144+
s.authTokens = &entity.AuthTokens{APIKey: "brev_api_test-key", APIKeyOrgID: "org-login"}
145+
s.org = &entity.Organization{ID: "org-set", Name: "set-org"}
146+
s.workspaces = []entity.Workspace{
147+
{
148+
ID: "ws1",
149+
Name: "api-key-instance",
150+
Status: entity.Running,
151+
CreatedByUserID: "other-user",
152+
},
153+
}
154+
term := terminal.New()
155+
156+
out := captureStdout(t, func() {
157+
err := RunLs(term, s, nil, "", false, true)
158+
if err != nil {
159+
t.Fatalf("RunLs returned error: %v", err)
160+
}
161+
})
162+
163+
if !strings.Contains(out, "api-key-instance") {
164+
t.Fatalf("expected workspace output, got %s", out)
165+
}
166+
if s.workspaceOrgID != "org-login" {
167+
t.Fatalf("expected API key credential org org-login, got %s", s.workspaceOrgID)
168+
}
169+
}
170+
171+
func TestRunLs_APIKeyRequiresCredentialOrg(t *testing.T) {
172+
s := newTestStore()
173+
s.accessToken = "brev_api_test-key"
174+
s.authTokens = &entity.AuthTokens{APIKey: "brev_api_test-key"}
175+
term := terminal.New()
176+
177+
err := RunLs(term, s, nil, "", false, true)
178+
179+
if err == nil {
180+
t.Fatal("expected missing API key org error, got nil")
181+
}
182+
if !strings.Contains(err.Error(), "api key auth requires an org id") {
183+
t.Fatalf("expected API key org validation error, got %v", err)
184+
}
185+
if s.workspaceOrgID != "" {
186+
t.Fatalf("expected no workspace call, got org %s", s.workspaceOrgID)
187+
}
188+
}
189+
133190
// captureStdout runs fn while capturing stdout and returns the output.
134191
func captureStdout(t *testing.T, fn func()) string {
135192
t.Helper()

‎pkg/cmd/set/set.go‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@ package set
44
import (
55
"fmt"
66

7+
"github.com/brevdev/brev-cli/pkg/auth"
78
"github.com/brevdev/brev-cli/pkg/cmd/cmderrors"
89
"github.com/brevdev/brev-cli/pkg/cmd/completions"
910
"github.com/brevdev/brev-cli/pkg/cmdcontext"
@@ -21,6 +22,7 @@ type SetStore interface {
2122
GetOrganizations(options *store.GetOrganizationsOptions) ([]entity.Organization, error)
2223
GetServerSockFile() string
2324
GetCurrentWorkspaceID() (string, error)
25+
GetAccessToken() (string, error)
2426
}
2527

2628
func NewCmdSet(t *terminal.Terminal, loginSetStore SetStore, noLoginSetStore SetStore) *cobra.Command {
@@ -59,6 +61,9 @@ func set(orgName string, setStore SetStore) error {
5961
if workspaceID != "" {
6062
return fmt.Errorf("can not set orgs in a workspace")
6163
}
64+
if auth.IsAPIKeyAuthStore(setStore) {
65+
return breverrors.NewValidationError("api key auth is scoped to the org saved during login; run brev login --api-key <api-key> --org-id <org-id> to change it")
66+
}
6267
orgs, err := setStore.GetOrganizations(&store.GetOrganizationsOptions{Name: orgName})
6368
if err != nil {
6469
return breverrors.WrapAndTrace(err)

‎pkg/cmd/set/set_test.go‎

Lines changed: 74 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1 +1,75 @@
11
package set
2+
3+
import (
4+
"strings"
5+
"testing"
6+
7+
"github.com/brevdev/brev-cli/pkg/entity"
8+
"github.com/brevdev/brev-cli/pkg/store"
9+
)
10+
11+
type mockSetStore struct {
12+
accessToken string
13+
workspaceID string
14+
orgs []entity.Organization
15+
getOrganizations int
16+
setDefaultOrgCalls int
17+
defaultOrganization *entity.Organization
18+
}
19+
20+
func (m *mockSetStore) GetWorkspaces(_ string, _ *store.GetWorkspacesOptions) ([]entity.Workspace, error) {
21+
return nil, nil
22+
}
23+
24+
func (m *mockSetStore) GetActiveOrganizationOrDefault() (*entity.Organization, error) {
25+
return nil, nil
26+
}
27+
28+
func (m *mockSetStore) GetCurrentUser() (*entity.User, error) {
29+
return nil, nil
30+
}
31+
32+
func (m *mockSetStore) SetDefaultOrganization(org *entity.Organization) error {
33+
m.setDefaultOrgCalls++
34+
m.defaultOrganization = org
35+
return nil
36+
}
37+
38+
func (m *mockSetStore) GetOrganizations(_ *store.GetOrganizationsOptions) ([]entity.Organization, error) {
39+
m.getOrganizations++
40+
return m.orgs, nil
41+
}
42+
43+
func (m *mockSetStore) GetServerSockFile() string {
44+
return ""
45+
}
46+
47+
func (m *mockSetStore) GetCurrentWorkspaceID() (string, error) {
48+
return m.workspaceID, nil
49+
}
50+
51+
func (m *mockSetStore) GetAccessToken() (string, error) {
52+
return m.accessToken, nil
53+
}
54+
55+
func TestSetRejectsAPIKeyAuth(t *testing.T) {
56+
s := &mockSetStore{
57+
accessToken: "brev_api_test-key",
58+
orgs: []entity.Organization{{ID: "org-other", Name: "other-org"}},
59+
}
60+
61+
err := set("other-org", s)
62+
63+
if err == nil {
64+
t.Fatal("expected API key auth set error, got nil")
65+
}
66+
if !strings.Contains(err.Error(), "api key auth is scoped") {
67+
t.Fatalf("expected API key auth validation error, got %v", err)
68+
}
69+
if s.getOrganizations != 0 {
70+
t.Fatalf("expected set to skip org lookup, got %d calls", s.getOrganizations)
71+
}
72+
if s.setDefaultOrgCalls != 0 {
73+
t.Fatalf("expected set to skip default org write, got %d calls", s.setDefaultOrgCalls)
74+
}
75+
}

‎pkg/entity/entity.go‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,7 @@ type AuthTokens struct {
2828
AccessToken string `json:"access_token"`
2929
RefreshToken string `json:"refresh_token"`
3030
APIKey string `json:"api_key,omitempty"`
31+
APIKeyOrgID string `json:"api_key_org_id,omitempty"`
3132
}
3233

3334
type IDEConfig struct {

0 commit comments

Comments
 (0)