fix: persist refreshable OAuth sessions

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

Author: lox

cmd/auth/login.go

@@ -73,6 +73,35 @@ func LoginWithToken(f *factory.Factory, org, token string) error {
 	return nil
 }
 
+// LoginWithSession stores an OAuth session for an organization in the system keychain.
+func LoginWithSession(f *factory.Factory, org string, session *oauth.Session) error {
+	if org == "" {
+		return errors.New("organization cannot be empty")
+	}
+	if session == nil || session.AccessToken == "" {
+		return errors.New("oauth session must include an access token")
+	}
+
+	kr := keyring.New()
+	if !kr.IsAvailable() {
+		return errors.New("system keychain is not available; cannot store token")
+	}
+	if err := kr.SetSession(org, session); err != nil {
+		return fmt.Errorf("failed to store token in keychain: %w", err)
+	}
+	fmt.Println("Token stored securely in system keychain.")
+
+	if err := f.Config.EnsureOrganization(org); err != nil {
+		return fmt.Errorf("failed to register organization in config: %w", err)
+	}
+
+	if err := f.Config.SelectOrganization(org, f.GitRepository != nil); err != nil {
+		return fmt.Errorf("failed to select organization: %w", err)
+	}
+
+	return nil
+}
+
 func (c *LoginCmd) Run(kongCtx *kong.Context, globals cli.GlobalFlags) error {
 	f, err := factory.New(factory.WithDebug(globals.EnableDebug()))
 	if err != nil {
@@ -95,8 +124,7 @@ func (c *LoginCmd) Run(kongCtx *kong.Context, globals cli.GlobalFlags) error {
 	// Create OAuth flow
 	cfg := &oauth.Config{
 		// Host default handled via NewFlow, omitted to allow usage of BUILDKITE_HOST
-		ClientID: oauth.DefaultClientID,
-		Scopes:   c.Scopes,
+		Scopes: c.Scopes,
 	}
 
 	flow, err := oauth.NewFlow(cfg)
@@ -135,23 +163,13 @@ func (c *LoginCmd) Run(kongCtx *kong.Context, globals cli.GlobalFlags) error {
 		return fmt.Errorf("token exchange failed: %w", err)
 	}
 
-	// Resolve org from the API using the new token
-	client, err := buildkite.NewOpts(buildkite.WithTokenAuth(tokenResp.AccessToken))
-	if err != nil {
-		return fmt.Errorf("failed to create API client: %w", err)
-	}
-
-	orgs, _, err := client.Organizations.List(ctx, nil)
+	org, err := resolveOrganizationFromToken(ctx, f.Config.RESTAPIEndpoint(), tokenResp.AccessToken)
 	if err != nil {
-		return fmt.Errorf("failed to list organizations: %w", err)
-	}
-	if len(orgs) == 0 {
-		return fmt.Errorf("no organizations found for this token")
+		return err
 	}
 
-	org := orgs[0]
-
-	if err := LoginWithToken(f, org.Slug, tokenResp.AccessToken); err != nil {
+	session := tokenResp.Session(cfg.Host, time.Now())
+	if err := LoginWithSession(f, org.Slug, session); err != nil {
 		return err
 	}
 
@@ -160,3 +178,23 @@ func (c *LoginCmd) Run(kongCtx *kong.Context, globals cli.GlobalFlags) error {
 
 	return nil
 }
+
+func resolveOrganizationFromToken(ctx context.Context, baseURL, token string) (*buildkite.Organization, error) {
+	client, err := buildkite.NewOpts(
+		buildkite.WithBaseURL(baseURL),
+		buildkite.WithTokenAuth(token),
+	)
+	if err != nil {
+		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)
+	}
+	if len(orgs) == 0 {
+		return nil, fmt.Errorf("no organizations found for this token")
+	}
+
+	return &orgs[0], nil
+}

cmd/auth/login_test.go

@@ -0,0 +1,34 @@
+package auth
+
+import (
+	"context"
+	"encoding/json"
+	"net/http"
+	"net/http/httptest"
+	"testing"
+)
+
+func TestResolveOrganizationFromTokenUsesConfiguredBaseURL(t *testing.T) {
+	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		if got := r.Header.Get("Authorization"); got != "Bearer bkua_test_token" {
+			t.Fatalf("Authorization = %q, want Bearer bkua_test_token", got)
+		}
+
+		w.Header().Set("Content-Type", "application/json")
+		if err := json.NewEncoder(w).Encode([]map[string]any{{"slug": "test-org"}}); err != nil {
+			t.Fatalf("Encode returned error: %v", err)
+		}
+	}))
+	defer server.Close()
+
+	org, err := resolveOrganizationFromToken(context.Background(), server.URL, "bkua_test_token")
+	if err != nil {
+		t.Fatalf("resolveOrganizationFromToken returned error: %v", err)
+	}
+	if org == nil {
+		t.Fatal("resolveOrganizationFromToken returned nil organization")
+	}
+	if org.Slug != "test-org" {
+		t.Fatalf("Slug = %q, want test-org", org.Slug)
+	}
+}

internal/config/config.go

@@ -6,19 +6,24 @@
 package config
 
 import (
+	"context"
 	"errors"
 	"fmt"
 	"io"
 	"maps"
+	"net/url"
 	"os"
 	"path/filepath"
 	"runtime"
 	"slices"
 	"strconv"
+	"strings"
 	"sync"
+	"time"
 
 	"github.com/buildkite/cli/v3/internal/pipeline"
 	"github.com/buildkite/cli/v3/pkg/keyring"
+	"github.com/buildkite/cli/v3/pkg/oauth"
 	buildkite "github.com/buildkite/go-buildkite/v4"
 	git "github.com/go-git/go-git/v5"
 	"github.com/goccy/go-yaml"
@@ -137,8 +142,14 @@ func (conf *Config) APITokenForOrg(org string) string {
 
 	kr := keyring.New()
 	if kr.IsAvailable() {
-		if token, err := kr.Get(org); err == nil && token != "" {
-			return token
+		if session, err := kr.GetSession(org); err == nil && session != nil && session.AccessToken != "" {
+			refreshedSession, refreshErr := conf.refreshOAuthSession(org, kr, session)
+			if refreshedSession != nil && refreshedSession.AccessToken != "" {
+				if refreshErr != nil {
+					fmt.Fprintf(os.Stderr, "Warning: failed to refresh OAuth token for %q: %v\n", org, refreshErr)
+				}
+				return refreshedSession.AccessToken
+			}
 		}
 	}
 
@@ -213,12 +224,33 @@ func (conf *Config) GetGraphQLEndpoint() string {
 func (conf *Config) RESTAPIEndpoint() string {
 	value := os.Getenv("BUILDKITE_REST_API_ENDPOINT")
 	if value != "" {
-		return value
+		return normaliseRESTAPIEndpoint(value)
 	}
 
 	return buildkite.DefaultBaseURL
 }
 
+func normaliseRESTAPIEndpoint(value string) string {
+	parsed, err := url.Parse(value)
+	if err != nil {
+		return value
+	}
+
+	trimmedPath := strings.TrimRight(parsed.Path, "/")
+	if !strings.HasSuffix(trimmedPath, "/v2") {
+		return value
+	}
+
+	parsed.Path = strings.TrimSuffix(trimmedPath, "/v2")
+	if parsed.Path == "" {
+		parsed.Path = "/"
+	} else {
+		parsed.Path += "/"
+	}
+
+	return parsed.String()
+}
+
 func (conf *Config) PagerDisabled() bool {
 	if v, ok := lookupBoolEnv("BUILDKITE_NO_PAGER"); ok {
 		return v
@@ -517,3 +549,24 @@ func (conf *Config) writeUser() error {
 func (conf *Config) writeLocal() error {
 	return writeFileConfig(conf.fs, conf.localPath, conf.local)
 }
+
+func (conf *Config) refreshOAuthSession(org string, kr *keyring.Keyring, session *oauth.Session) (*oauth.Session, error) {
+	if session == nil || session.AccessToken == "" {
+		return nil, nil
+	}
+	if !session.NeedsRefresh(time.Now()) {
+		return session, nil
+	}
+
+	refreshedToken, err := oauth.RefreshAccessToken(context.Background(), &oauth.Config{Host: session.Host}, session.RefreshToken, session.Scope)
+	if err != nil {
+		return session, err
+	}
+
+	refreshedSession := session.Update(refreshedToken, time.Now())
+	if err := kr.SetSession(org, refreshedSession); err != nil {
+		return refreshedSession, err
+	}
+
+	return refreshedSession, nil
+}

internal/config/config_test.go

@@ -1,10 +1,14 @@
 package config
 
 import (
+	"context"
+	"net/http"
+	"net/http/httptest"
 	"os"
 	"path/filepath"
 	"testing"
 
+	buildkite "github.com/buildkite/go-buildkite/v4"
 	"github.com/spf13/afero"
 )
 
@@ -100,6 +104,33 @@ func TestConfig(t *testing.T) {
 		}
 	})
 
+	t.Run("RESTAPIEndpoint strips version suffix from env override", func(t *testing.T) {
+		server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			if r.URL.Path != "/v2/organizations" {
+				t.Fatalf("request path = %q, want /v2/organizations", r.URL.Path)
+			}
+
+			w.Header().Set("Content-Type", "application/json")
+			_, _ = w.Write([]byte(`[]`))
+		}))
+		defer server.Close()
+
+		setEnv(t, "BUILDKITE_REST_API_ENDPOINT", server.URL+"/v2")
+
+		conf := New(afero.NewMemMapFs(), nil)
+		client, err := buildkite.NewOpts(
+			buildkite.WithBaseURL(conf.RESTAPIEndpoint()),
+			buildkite.WithTokenAuth("bkua_test_token"),
+		)
+		if err != nil {
+			t.Fatalf("NewOpts returned error: %v", err)
+		}
+
+		if _, _, err := client.Organizations.List(context.Background(), nil); err != nil {
+			t.Fatalf("Organizations.List returned error: %v", err)
+		}
+	})
+
 	t.Run("loadFileConfig returns error on invalid yaml", func(t *testing.T) {
 		fs := afero.NewMemMapFs()
 		path := filepath.Join(t.TempDir(), "bk.yaml")

internal/config/oauth_refresh_test.go

@@ -0,0 +1,85 @@
+package config
+
+import (
+	"encoding/json"
+	"net/http"
+	"net/http/httptest"
+	"testing"
+	"time"
+
+	"github.com/buildkite/cli/v3/pkg/keyring"
+	"github.com/buildkite/cli/v3/pkg/oauth"
+	"github.com/spf13/afero"
+)
+
+func TestAPITokenForOrgRefreshesStoredOAuthSession(t *testing.T) {
+	keyring.MockForTesting()
+	t.Setenv(oauth.EnvClientID, "env-client")
+	t.Setenv(oauth.LegacyEnvClientID, "")
+
+	var requests int
+	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		requests++
+
+		if r.URL.Path != "/oauth/token" {
+			t.Fatalf("unexpected token path %q", r.URL.Path)
+		}
+		if err := r.ParseForm(); err != nil {
+			t.Fatalf("ParseForm returned error: %v", err)
+		}
+		if got := r.Form.Get("grant_type"); got != "refresh_token" {
+			t.Fatalf("grant_type = %q, want refresh_token", got)
+		}
+		if got := r.Form.Get("refresh_token"); got != "bkrt_old_refresh" {
+			t.Fatalf("refresh_token = %q, want bkrt_old_refresh", got)
+		}
+
+		w.Header().Set("Content-Type", "application/json")
+		if err := json.NewEncoder(w).Encode(oauth.TokenResponse{
+			AccessToken:  "bkua_refreshed_access",
+			RefreshToken: "bkrt_rotated_refresh",
+			TokenType:    "Bearer",
+			Scope:        "read_user read_organizations",
+			ExpiresIn:    3600,
+		}); err != nil {
+			t.Fatalf("Encode returned error: %v", err)
+		}
+	}))
+	defer server.Close()
+
+	conf := New(afero.NewMemMapFs(), nil)
+	kr := keyring.New()
+	if err := kr.SetSession("test-org", &oauth.Session{
+		Version:      oauth.SessionVersion,
+		Host:         server.URL,
+		AccessToken:  "bkua_expired_access",
+		RefreshToken: "bkrt_old_refresh",
+		TokenType:    "Bearer",
+		Scope:        "read_user read_organizations",
+		ExpiresAt:    time.Now().Add(-time.Minute),
+	}); err != nil {
+		t.Fatalf("SetSession returned error: %v", err)
+	}
+
+	token := conf.APITokenForOrg("test-org")
+	if token != "bkua_refreshed_access" {
+		t.Fatalf("APITokenForOrg() = %q, want bkua_refreshed_access", token)
+	}
+	if requests != 1 {
+		t.Fatalf("refresh requests = %d, want 1", requests)
+	}
+
+	session, err := kr.GetSession("test-org")
+	if err != nil {
+		t.Fatalf("GetSession returned error: %v", err)
+	}
+	if session.AccessToken != "bkua_refreshed_access" {
+		t.Fatalf("stored AccessToken = %q, want bkua_refreshed_access", session.AccessToken)
+	}
+	if session.RefreshToken != "bkrt_rotated_refresh" {
+		t.Fatalf("stored RefreshToken = %q, want bkrt_rotated_refresh", session.RefreshToken)
+	}
+	if !session.ExpiresAt.After(time.Now()) {
+		t.Fatalf("stored ExpiresAt = %s, want a future timestamp", session.ExpiresAt)
+	}
+}