fix: address oauth review findings

Amp-Thread-ID: https://ampcode.com/threads/T-019d3412-2251-7197-b667-d42d29ab0ded

Author: lox

cmd/api/api.go

@@ -116,7 +116,7 @@ func (c *ApiCmd) Run(kongCtx *kong.Context, globals cli.GlobalFlags) error {
 
 	// Create an HTTP client with appropriate configuration
 	client := httpClient.NewClient(
-		f.Config.APIToken(),
+		f.Token,
 		httpClient.WithBaseURL(f.RestAPIClient.BaseURL.String()),
 		httpClient.WithMaxRetries(3),
 		httpClient.WithMaxRetryDelay(60*time.Second),

cmd/auth/login.go

@@ -113,7 +113,7 @@ func LoginWithSession(f *factory.Factory, org string, session *oauth.Session) er
 }
 
 func (c *LoginCmd) Run(kongCtx *kong.Context, globals cli.GlobalFlags) error {
-	f, err := factory.New(factory.WithDebug(globals.EnableDebug()))
+	f, err := factory.New(factory.WithDebug(globals.EnableDebug()), factory.WithoutAPIClients())
 	if err != nil {
 		return err
 	}
@@ -178,23 +178,23 @@ func (c *LoginCmd) Run(kongCtx *kong.Context, globals cli.GlobalFlags) error {
 		return fmt.Errorf("token exchange failed: %w", err)
 	}
 
-	org, err := resolveOrganizationFromToken(ctx, f.Config.RESTAPIEndpoint(), tokenResp.AccessToken)
+	orgs, err := resolveOrganizationsFromToken(ctx, f.Config.RESTAPIEndpoint(), tokenResp.AccessToken)
 	if err != nil {
 		return err
 	}
 
-	session := tokenResp.Session(cfg.Host, time.Now())
-	if err := LoginWithSession(f, org.Slug, session); err != nil {
+	session := tokenResp.Session(cfg.Host, cfg.ClientID, time.Now())
+	if err := storeSessionForOrganizations(f, orgs, session); err != nil {
 		return err
 	}
 
-	fmt.Printf("\n✅ Successfully authenticated with organization %q\n", org.Slug)
+	fmt.Printf("\n✅ Successfully authenticated with organization %q\n", orgs[0].Slug)
 	fmt.Printf("  Scopes: %s\n", tokenResp.Scope)
 
 	return nil
 }
 
-func resolveOrganizationFromToken(ctx context.Context, baseURL, token string) (*buildkite.Organization, error) {
+func resolveOrganizationsFromToken(ctx context.Context, baseURL, token string) ([]buildkite.Organization, error) {
 	client, err := buildkite.NewOpts(
 		buildkite.WithBaseURL(baseURL),
 		buildkite.WithTokenAuth(token),
@@ -203,13 +203,64 @@ func resolveOrganizationFromToken(ctx context.Context, baseURL, token string) (*
 		return nil, fmt.Errorf("failed to create API client: %w", err)
 	}
 
-	orgs, _, err := client.Organizations.List(ctx, nil)
-	if err != nil {
-		return nil, fmt.Errorf("failed to list organizations: %w", err)
+	var allOrgs []buildkite.Organization
+	page := 1
+	for {
+		orgs, resp, err := client.Organizations.List(ctx, &buildkite.OrganizationListOptions{
+			ListOptions: buildkite.ListOptions{Page: page},
+		})
+		if err != nil {
+			return nil, fmt.Errorf("failed to list organizations: %w", err)
+		}
+		allOrgs = append(allOrgs, orgs...)
+		if resp == nil || resp.NextPage == 0 {
+			break
+		}
+		page = resp.NextPage
 	}
-	if len(orgs) == 0 {
+
+	if len(allOrgs) == 0 {
 		return nil, fmt.Errorf("no organizations found for this token")
 	}
 
+	return allOrgs, nil
+}
+
+func resolveOrganizationFromToken(ctx context.Context, baseURL, token string) (*buildkite.Organization, error) {
+	orgs, err := resolveOrganizationsFromToken(ctx, baseURL, token)
+	if err != nil {
+		return nil, err
+	}
+
 	return &orgs[0], nil
 }
+
+func storeSessionForOrganizations(f *factory.Factory, orgs []buildkite.Organization, session *oauth.Session) error {
+	if len(orgs) == 0 {
+		return errors.New("no organizations found for this token")
+	}
+	if err := LoginWithSession(f, orgs[0].Slug, session); err != nil {
+		return err
+	}
+
+	kr := keyring.New()
+	seen := map[string]struct{}{orgs[0].Slug: {}}
+	for _, org := range orgs[1:] {
+		if org.Slug == "" {
+			continue
+		}
+		if _, exists := seen[org.Slug]; exists {
+			continue
+		}
+		seen[org.Slug] = struct{}{}
+
+		if err := kr.SetSession(org.Slug, session); err != nil {
+			return fmt.Errorf("failed to store token in keychain: %w", err)
+		}
+		if err := f.Config.EnsureOrganization(org.Slug); err != nil {
+			return fmt.Errorf("failed to register organization in config: %w", err)
+		}
+	}
+
+	return nil
+}

cmd/auth/login_test.go

@@ -6,6 +6,13 @@ import (
 	"net/http"
 	"net/http/httptest"
 	"testing"
+
+	"github.com/buildkite/cli/v3/internal/config"
+	"github.com/buildkite/cli/v3/pkg/cmd/factory"
+	"github.com/buildkite/cli/v3/pkg/keyring"
+	"github.com/buildkite/cli/v3/pkg/oauth"
+	buildkite "github.com/buildkite/go-buildkite/v4"
+	"github.com/spf13/afero"
 )
 
 func TestResolveOrganizationFromTokenUsesConfiguredBaseURL(t *testing.T) {
@@ -32,3 +39,74 @@ func TestResolveOrganizationFromTokenUsesConfiguredBaseURL(t *testing.T) {
 		t.Fatalf("Slug = %q, want test-org", org.Slug)
 	}
 }
+
+func TestStoreSessionForOrganizationsStoresAllAccessibleOrgs(t *testing.T) {
+	keyring.MockForTesting()
+
+	f := &factory.Factory{
+		Config: config.New(afero.NewMemMapFs(), nil),
+	}
+	session := &oauth.Session{
+		Version:     oauth.SessionVersion,
+		AccessToken: "bkua_access",
+		TokenType:   "Bearer",
+	}
+
+	orgs := []buildkite.Organization{
+		{Slug: "test-org"},
+		{Slug: "other-org"},
+		{Slug: "other-org"},
+	}
+
+	if err := storeSessionForOrganizations(f, orgs, session); err != nil {
+		t.Fatalf("storeSessionForOrganizations returned error: %v", err)
+	}
+
+	kr := keyring.New()
+	for _, slug := range []string{"test-org", "other-org"} {
+		storedSession, err := kr.GetSession(slug)
+		if err != nil {
+			t.Fatalf("GetSession(%q) returned error: %v", slug, err)
+		}
+		if storedSession.AccessToken != "bkua_access" {
+			t.Fatalf("stored access token for %q = %q, want bkua_access", slug, storedSession.AccessToken)
+		}
+	}
+
+	if got := f.Config.OrganizationSlug(); got != "test-org" {
+		t.Fatalf("OrganizationSlug() = %q, want test-org", got)
+	}
+}
+
+func TestResolveOrganizationsFromTokenPaginates(t *testing.T) {
+	var server *httptest.Server
+	server = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		page := r.URL.Query().Get("page")
+		w.Header().Set("Content-Type", "application/json")
+		switch page {
+		case "", "1":
+			w.Header().Set("Link", `<`+server.URL+`/v2/organizations?page=2>; rel="next"`)
+			if err := json.NewEncoder(w).Encode([]map[string]any{{"slug": "org-one"}}); err != nil {
+				t.Fatalf("Encode page 1 returned error: %v", err)
+			}
+		case "2":
+			if err := json.NewEncoder(w).Encode([]map[string]any{{"slug": "org-two"}}); err != nil {
+				t.Fatalf("Encode page 2 returned error: %v", err)
+			}
+		default:
+			t.Fatalf("unexpected page query %q", page)
+		}
+	}))
+	defer server.Close()
+
+	orgs, err := resolveOrganizationsFromToken(context.Background(), server.URL, "bkua_test_token")
+	if err != nil {
+		t.Fatalf("resolveOrganizationsFromToken returned error: %v", err)
+	}
+	if len(orgs) != 2 {
+		t.Fatalf("len(orgs) = %d, want 2", len(orgs))
+	}
+	if orgs[0].Slug != "org-one" || orgs[1].Slug != "org-two" {
+		t.Fatalf("org slugs = [%q %q], want [org-one org-two]", orgs[0].Slug, orgs[1].Slug)
+	}
+}

cmd/auth/logout.go

@@ -7,6 +7,7 @@ import (
 	"github.com/buildkite/cli/v3/internal/cli"
 	"github.com/buildkite/cli/v3/pkg/cmd/factory"
 	"github.com/buildkite/cli/v3/pkg/keyring"
+	"github.com/buildkite/cli/v3/pkg/oauth"
 )
 
 type LogoutCmd struct {
@@ -15,7 +16,7 @@ type LogoutCmd struct {
 }
 
 func (c *LogoutCmd) Run(kongCtx *kong.Context, globals cli.GlobalFlags) error {
-	f, err := factory.New(factory.WithDebug(globals.EnableDebug()))
+	f, err := factory.New(factory.WithDebug(globals.EnableDebug()), factory.WithoutAPIClients())
 	if err != nil {
 		return err
 	}
@@ -59,13 +60,41 @@ func (c *LogoutCmd) logoutOrg(f *factory.Factory) error {
 
 	kr := keyring.New()
 	if kr.IsAvailable() {
+		var currentSession *oauth.Session
+		currentSession, _ = kr.GetSession(org)
 		if err := kr.Delete(org); err != nil {
 			fmt.Printf("Warning: could not remove token from keychain: %v\n", err)
 		} else {
+			c.deleteSiblingOAuthSessions(f, kr, org, currentSession)
 			fmt.Println("Token removed from system keychain.")
 		}
 	}
 
 	fmt.Printf("Logged out of organization %q\n", org)
 	return nil
 }
+
+func (c *LogoutCmd) deleteSiblingOAuthSessions(f *factory.Factory, kr *keyring.Keyring, org string, session *oauth.Session) {
+	if session == nil || session.RefreshToken == "" {
+		return
+	}
+
+	for _, sibling := range f.Config.ConfiguredOrganizations() {
+		if sibling == "" || sibling == org {
+			continue
+		}
+
+		siblingSession, err := kr.GetSession(sibling)
+		if err != nil || siblingSession == nil {
+			continue
+		}
+		if siblingSession.Host != session.Host || siblingSession.ClientID != session.ClientID {
+			continue
+		}
+		if siblingSession.RefreshToken != session.RefreshToken || siblingSession.AccessToken != session.AccessToken {
+			continue
+		}
+
+		_ = kr.Delete(sibling)
+	}
+}

cmd/auth/logout_test.go

@@ -0,0 +1,56 @@
+package auth
+
+import (
+	"testing"
+
+	"github.com/buildkite/cli/v3/internal/config"
+	"github.com/buildkite/cli/v3/pkg/cmd/factory"
+	"github.com/buildkite/cli/v3/pkg/keyring"
+	"github.com/buildkite/cli/v3/pkg/oauth"
+	"github.com/spf13/afero"
+)
+
+func TestLogoutOrgDeletesSiblingOAuthAliases(t *testing.T) {
+	keyring.MockForTesting()
+
+	conf := config.New(afero.NewMemMapFs(), nil)
+	if err := conf.EnsureOrganization("org-a"); err != nil {
+		t.Fatalf("EnsureOrganization org-a returned error: %v", err)
+	}
+	if err := conf.EnsureOrganization("org-b"); err != nil {
+		t.Fatalf("EnsureOrganization org-b returned error: %v", err)
+	}
+	if err := conf.SelectOrganization("org-a", false); err != nil {
+		t.Fatalf("SelectOrganization returned error: %v", err)
+	}
+
+	session := &oauth.Session{
+		Version:      oauth.SessionVersion,
+		Host:         "buildkite.localhost",
+		ClientID:     "buildkite-cli",
+		AccessToken:  "bkua_access",
+		RefreshToken: "bkrt_refresh",
+		TokenType:    "Bearer",
+	}
+
+	kr := keyring.New()
+	if err := kr.SetSession("org-a", session); err != nil {
+		t.Fatalf("SetSession org-a returned error: %v", err)
+	}
+	if err := kr.SetSession("org-b", session); err != nil {
+		t.Fatalf("SetSession org-b returned error: %v", err)
+	}
+
+	cmd := &LogoutCmd{Org: "org-a"}
+	f := &factory.Factory{Config: conf}
+	if err := cmd.logoutOrg(f); err != nil {
+		t.Fatalf("logoutOrg returned error: %v", err)
+	}
+
+	if _, err := kr.GetSession("org-a"); err == nil {
+		t.Fatal("expected org-a session to be deleted")
+	}
+	if _, err := kr.GetSession("org-b"); err == nil {
+		t.Fatal("expected org-b sibling session to be deleted")
+	}
+}

cmd/auth/switch.go

@@ -26,7 +26,7 @@ Examples:
 }
 
 func (c *SwitchCmd) Run(globals cli.GlobalFlags) error {
-	f, err := factory.New(factory.WithDebug(globals.EnableDebug()))
+	f, err := factory.New(factory.WithDebug(globals.EnableDebug()), factory.WithoutAPIClients())
 	if err != nil {
 		return err
 	}

cmd/config/get.go

@@ -36,7 +36,7 @@ func (c *GetCmd) Run() error {
 		return err
 	}
 
-	f, err := factory.New()
+	f, err := factory.New(factory.WithoutAPIClients())
 	if err != nil {
 		return err
 	}

cmd/config/list.go

@@ -24,7 +24,7 @@ Examples:
 }
 
 func (c *ListCmd) Run() error {
-	f, err := factory.New()
+	f, err := factory.New(factory.WithoutAPIClients())
 	if err != nil {
 		return err
 	}

cmd/config/set.go

@@ -57,7 +57,7 @@ func (c *SetCmd) Run() error {
 		return fmt.Errorf("%s can only be set in user config (not --local)", key)
 	}
 
-	f, err := factory.New()
+	f, err := factory.New(factory.WithoutAPIClients())
 	if err != nil {
 		return err
 	}

cmd/config/unset.go

@@ -36,7 +36,7 @@ func (c *UnsetCmd) Run() error {
 		return fmt.Errorf("%s can only be unset from user config (not --local)", key)
 	}
 
-	f, err := factory.New()
+	f, err := factory.New(factory.WithoutAPIClients())
 	if err != nil {
 		return err
 	}