From 0ab827cf82f2e644245236f375b0212c7c583143 Mon Sep 17 00:00:00 2001 From: Maksim Zuev Date: Thu, 29 Oct 2026 12:01:00 +0000 Subject: [PATCH 1/6] Split device code login into authorization and token exchange steps libs/auth/u2m now provides the two steps of device code login instead of the whole login. PersistentAuth.DeviceCodeLogin checks that the host supports it and fetches the OAuth endpoints, and its RequestDeviceAuthorization and RequestDeviceAccessToken methods make the two requests. A rejected code matches ErrInvalidConfirmationCode or ErrInvalidDeviceCode. The confirmation code prompt and its retries move from DeviceCodeChallenge to auth login. The login's output does not change. Co-authored-by: Isaac --- cmd/auth/login.go | 4 +- cmd/auth/login_device_code.go | 40 ++++++ cmd/auth/login_device_code_test.go | 22 ++++ libs/auth/u2m/device_code.go | 127 +++++++++++------- libs/auth/u2m/device_code_test.go | 201 +++++++++++------------------ 5 files changed, 217 insertions(+), 177 deletions(-) diff --git a/cmd/auth/login.go b/cmd/auth/login.go index 5c9f322d53d..b2844451e87 100644 --- a/cmd/auth/login.go +++ b/cmd/auth/login.go @@ -343,7 +343,7 @@ a new profile is created. } // Check the host once it is resolved, before setAccountID prompts for // an account ID for account and unified hosts. A workspace without - // device code login enabled is rejected later by DeviceCodeChallenge. + // device code login enabled is rejected later by deviceCodeLogin. if useDeviceCode { // The host prompt accepts an empty answer. Without this check, // login would resolve the host from DATABRICKS_HOST instead. @@ -401,7 +401,7 @@ a new profile is created. var token *oauth2.Token if useDeviceCode { - token, err = persistentAuth.DeviceCodeChallenge(ctx, promptForConfirmationCode) + token, err = deviceCodeLogin(ctx, persistentAuth) } else { token, err = persistentAuth.Challenge() } diff --git a/cmd/auth/login_device_code.go b/cmd/auth/login_device_code.go index 84ca1b03762..0c339eff30e 100644 --- a/cmd/auth/login_device_code.go +++ b/cmd/auth/login_device_code.go @@ -11,8 +11,13 @@ import ( "github.com/databricks/cli/libs/databrickscfg/profile" "github.com/databricks/cli/libs/env" "github.com/spf13/cobra" + "golang.org/x/oauth2" ) +// maxConfirmationAttempts is how many confirmation codes the user can enter +// for a single device code before the login fails. +const maxConfirmationAttempts = 3 + // deviceCodeIncompatibleFlags lists flags that device code login cannot honor. var deviceCodeIncompatibleFlags = []string{ // Resource indicators are sent on the browser authorization request, @@ -57,6 +62,41 @@ func validateDeviceCodeLogin(ctx context.Context, cmd *cobra.Command) error { return nil } +// deviceCodeLogin logs in using the Databricks variant of the RFC 8628 +// device authorization grant. It shows the user where to approve the login +// and exchanges the device code together with the confirmation code that the +// user enters for a token. A rejected confirmation code leaves the device +// code valid, so the user can re-enter it up to maxConfirmationAttempts +// times. +func deviceCodeLogin(ctx context.Context, persistentAuth *u2m.PersistentAuth) (*oauth2.Token, error) { + login, err := persistentAuth.DeviceCodeLogin() + if err != nil { + return nil, err + } + da, err := login.RequestDeviceAuthorization(ctx) + if err != nil { + return nil, err + } + + var lastErr error + for attempt := range maxConfirmationAttempts { + code, err := promptForConfirmationCode(ctx, *da, lastErr, maxConfirmationAttempts-attempt) + if err != nil { + return nil, err + } + + token, err := login.RequestDeviceAccessToken(ctx, da.DeviceCode, code) + if err == nil { + return token, nil + } + lastErr = err + if !errors.Is(err, u2m.ErrInvalidConfirmationCode) { + return nil, fmt.Errorf("token exchange: %w", err) + } + } + return nil, fmt.Errorf("%w. Start login from scratch", lastErr) +} + // promptForConfirmationCode shows the user where to approve the login and // reads the confirmation code from stdin. If the server rejected the previous // code, it shows the server's error and the number of attempts left instead diff --git a/cmd/auth/login_device_code_test.go b/cmd/auth/login_device_code_test.go index d38d2543b7c..2f24f11b9fc 100644 --- a/cmd/auth/login_device_code_test.go +++ b/cmd/auth/login_device_code_test.go @@ -5,6 +5,7 @@ import ( "context" "errors" "io" + "net/http" "runtime" "strings" "testing" @@ -17,6 +18,7 @@ import ( "github.com/databricks/cli/libs/databrickscfg/profile" "github.com/databricks/cli/libs/env" "github.com/databricks/cli/libs/flags" + "github.com/databricks/cli/libs/testserver" "github.com/spf13/cobra" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" @@ -304,3 +306,23 @@ resources = https://workspace.test/ai-gateway/mcp-services/system.ai.github _, err := tm.Result() assert.EqualError(t, err, "^C") } + +// Only a wrong confirmation code can be fixed by entering the code again, so +// any other token exchange error ends the login. +func TestLoginDeviceCodeFailsOnTokenExchangeError(t *testing.T) { + server := testserver.New(t) + // The router keeps the first handler registered for a route, so this one + // replaces the default token endpoint. + server.Handle("POST", "/oidc/v1/token", func(_ testserver.Request) any { + return testserver.Response{ + StatusCode: http.StatusTooManyRequests, + Body: map[string]string{"error": "invalid_request", "error_description": "Device grant rate limit exceeded"}, + } + }) + testserver.AddDefaultHandlers(server) + tm := startDeviceCodeLogin(t, "[test]\nhost = "+server.URL+"\n", "--profile", "test") + tm.WaitFor("Confirmation code") + tm.Type("FTFA-QQKY\n") + _, err := tm.Result() + assert.EqualError(t, err, "token exchange: Device grant rate limit exceeded (error code: invalid_request)") +} diff --git a/libs/auth/u2m/device_code.go b/libs/auth/u2m/device_code.go index 51f77f8e189..a0ca3886791 100644 --- a/libs/auth/u2m/device_code.go +++ b/libs/auth/u2m/device_code.go @@ -32,27 +32,28 @@ const invalidDeviceCodeDescription = "Invalid device code" // a device authorization endpoint in its OIDC discovery document. var ErrDeviceCodeNotSupported = errors.New("device code login is not enabled for this workspace") +// ErrInvalidConfirmationCode matches the token exchange error for a +// confirmation code the server rejects. The device code stays valid, so the +// exchange can be retried with the right code. +var ErrInvalidConfirmationCode = errors.New("invalid confirmation code") + +// ErrInvalidDeviceCode matches the token exchange error for a device code the +// server does not know, e.g. one it discarded after the login expired. The +// login has to be started again. +var ErrInvalidDeviceCode = errors.New("invalid device code") + // DeviceAuthorization holds what the user needs to approve a device code -// login from another device. +// login from another device, and the device code to exchange afterwards. type DeviceAuthorization struct { // UserCode is the code the user enters at VerificationURI. UserCode string // VerificationURI is the page where the user approves the login. VerificationURI string -} -// maxConfirmationAttempts is how many confirmation codes the user can enter -// for a single device code before the login fails. -const maxConfirmationAttempts = 3 - -// ConfirmFunc shows the device authorization to the user and returns the -// confirmation code that the user received after approving the login. When -// the server rejects a confirmation code, it is called again with the -// server's error in lastErr so that the user can re-enter the code. lastErr -// is nil on the first call. attemptsLeft is the number of confirmation codes -// the user can still enter, including this one. -type ConfirmFunc func(ctx context.Context, da DeviceAuthorization, lastErr error, attemptsLeft int) (string, error) + // DeviceCode identifies the login in the token exchange. + DeviceCode string +} // deviceCodeConfig returns the OAuth config for device code login, or an // error if the host does not support it. @@ -75,53 +76,83 @@ func (a *PersistentAuth) deviceCodeConfig() (*oauth2.Config, error) { return cfg, nil } -// DeviceCodeChallenge logs in using the Databricks variant of the RFC 8628 -// device authorization grant. It requests a device authorization, calls -// confirm with the user code and verification URI, and exchanges the device -// code together with the returned confirmation code for a token. A rejected -// confirmation code leaves the device code valid, so the user can re-enter -// it up to maxConfirmationAttempts times. Only workspace hosts are supported. -// The caller is responsible for storing the returned token. -func (a *PersistentAuth) DeviceCodeChallenge(ctx context.Context, confirm ConfirmFunc) (*oauth2.Token, error) { +// DeviceCodeLogin makes the requests of a login using the Databricks variant +// of the RFC 8628 device authorization grant. Get one with +// PersistentAuth.DeviceCodeLogin. +type DeviceCodeLogin struct { + auth *PersistentAuth + cfg *oauth2.Config +} + +// DeviceCodeLogin fetches the OAuth endpoints for a device code login, or +// returns an error if the host does not support it. Only workspace hosts are +// supported. +func (a *PersistentAuth) DeviceCodeLogin() (*DeviceCodeLogin, error) { cfg, err := a.deviceCodeConfig() if err != nil { return nil, err } + return &DeviceCodeLogin{auth: a, cfg: cfg}, nil +} - da, err := a.requestDeviceAuthorization(ctx, cfg) +// RequestDeviceAuthorization starts the login. The caller shows the user code +// and the verification URI to the user and later passes the device code to +// RequestDeviceAccessToken. +func (l *DeviceCodeLogin) RequestDeviceAuthorization(ctx context.Context) (*DeviceAuthorization, error) { + da, err := l.auth.requestDeviceAuthorization(ctx, l.cfg) if err != nil { return nil, fmt.Errorf("device authorization: %w", oauthError(err)) } - - authorization := DeviceAuthorization{ + return &DeviceAuthorization{ UserCode: da.UserCode, VerificationURI: da.VerificationURI, - } - var lastErr error - for attempt := range maxConfirmationAttempts { - code, err := confirm(ctx, authorization, lastErr, maxConfirmationAttempts-attempt) - if err != nil { - return nil, err - } + DeviceCode: da.DeviceCode, + }, nil +} - t, err := a.exchangeDeviceCode(ctx, cfg, da.DeviceCode, normalizeCode(code)) - if err == nil { - return t, nil - } - lastErr = oauthError(err) - // The server answers a wrong confirmation code with HTTP 400 and - // invalid_request. Any other error cannot be fixed by entering the - // code again, including invalid_request with HTTP 429 when the server - // rate limits the token requests for this device code. An unknown - // device code also gets HTTP 400 and invalid_request, and only the - // description tells it apart. If the description changes, the user is - // asked for the code again until the attempts run out. - status, errCode, description := oauthErrorResponse(err) - if status != http.StatusBadRequest || errCode != "invalid_request" || description == invalidDeviceCodeDescription { - return nil, fmt.Errorf("token exchange: %w", lastErr) - } +// RequestDeviceAccessToken completes the login with a single token request. +// The confirmation code is accepted the way the browser displays it. An error +// matches ErrInvalidConfirmationCode or ErrInvalidDeviceCode when the server +// rejects one of the codes. A rejected confirmation code leaves the device +// code valid, so the exchange can be retried with the right code. The caller +// is responsible for storing the returned token. +func (l *DeviceCodeLogin) RequestDeviceAccessToken(ctx context.Context, deviceCode, confirmationCode string) (*oauth2.Token, error) { + t, err := l.auth.exchangeDeviceCode(ctx, l.cfg, deviceCode, normalizeCode(confirmationCode)) + if err == nil { + return t, nil + } + // The server answers a wrong confirmation code with HTTP 400 and + // invalid_request. Any other error cannot be fixed by entering the + // code again, including invalid_request with HTTP 429 when the server + // rate limits the token requests for this device code. An unknown + // device code also gets HTTP 400 and invalid_request, and only the + // description tells it apart. If the description changes, the error is + // reported as a wrong confirmation code. + status, errCode, description := oauthErrorResponse(err) + err = oauthError(err) + switch { + case status != http.StatusBadRequest || errCode != "invalid_request": + return nil, err + case description == invalidDeviceCodeDescription: + return nil, &codeRejectedError{err: err, sentinel: ErrInvalidDeviceCode} + default: + return nil, &codeRejectedError{err: err, sentinel: ErrInvalidConfirmationCode} } - return nil, fmt.Errorf("%w. Start login from scratch", lastErr) +} + +// codeRejectedError shows the server's error for a rejected code and matches +// sentinel with errors.Is. +type codeRejectedError struct { + err error + sentinel error +} + +func (e *codeRejectedError) Error() string { + return e.err.Error() +} + +func (e *codeRejectedError) Unwrap() []error { + return []error{e.err, e.sentinel} } // requestDeviceAuthorization makes the RFC 8628 device authorization request. diff --git a/libs/auth/u2m/device_code_test.go b/libs/auth/u2m/device_code_test.go index f976ef5555d..03cffa0145b 100644 --- a/libs/auth/u2m/device_code_test.go +++ b/libs/auth/u2m/device_code_test.go @@ -119,20 +119,34 @@ func newDeviceCodeAuth(t *testing.T, f *fakeDeviceServer, advertise bool, opts . return p } -func TestDeviceCodeChallenge(t *testing.T) { +func newDeviceCodeLogin(t *testing.T, f *fakeDeviceServer, opts ...PersistentAuthOption) *DeviceCodeLogin { + l, err := newDeviceCodeAuth(t, f, true, opts...).DeviceCodeLogin() + require.NoError(t, err) + return l +} + +// runDeviceCodeLogin runs both steps of a device code login. +func runDeviceCodeLogin(ctx context.Context, l *DeviceCodeLogin, confirmationCode string) (*oauth2.Token, error) { + da, err := l.RequestDeviceAuthorization(ctx) + if err != nil { + return nil, err + } + return l.RequestDeviceAccessToken(ctx, da.DeviceCode, confirmationCode) +} + +func TestDeviceCodeLogin(t *testing.T) { f := &fakeDeviceServer{expiresIn: 600} - p := newDeviceCodeAuth(t, f, true) + l := newDeviceCodeLogin(t, f) - var got DeviceAuthorization - tok, err := p.DeviceCodeChallenge(t.Context(), func(_ context.Context, da DeviceAuthorization, _ error, _ int) (string, error) { - got = da - return " ftfa-qqky\t", nil - }) + got, err := l.RequestDeviceAuthorization(t.Context()) + require.NoError(t, err) + tok, err := l.RequestDeviceAccessToken(t.Context(), got.DeviceCode, " ftfa-qqky\t") require.NoError(t, err) - assert.Equal(t, DeviceAuthorization{ + assert.Equal(t, &DeviceAuthorization{ UserCode: testUserCode, VerificationURI: "https://workspace.test/device/verification", + DeviceCode: testDeviceCode, }, got) assert.Equal(t, "access-token", tok.AccessToken) assert.Equal(t, "refresh-token", tok.RefreshToken) @@ -140,26 +154,27 @@ func TestDeviceCodeChallenge(t *testing.T) { assert.Equal(t, int32(1), f.tokenRequests.Load()) } -func TestDeviceCodeChallenge_ConfirmationAfterInitialExpiry(t *testing.T) { +func TestDeviceCodeLogin_ConfirmationAfterInitialExpiry(t *testing.T) { // The server extends the device code's lifetime once the user approves // the login, so a confirmation code entered after the initial expires_in // must still be exchanged. f := &fakeDeviceServer{expiresIn: 1} - p := newDeviceCodeAuth(t, f, true) + l := newDeviceCodeLogin(t, f) - tok, err := p.DeviceCodeChallenge(t.Context(), func(_ context.Context, _ DeviceAuthorization, _ error, _ int) (string, error) { - time.Sleep(1100 * time.Millisecond) - return testConfirmationCode, nil - }) + da, err := l.RequestDeviceAuthorization(t.Context()) + require.NoError(t, err) + time.Sleep(1100 * time.Millisecond) + tok, err := l.RequestDeviceAccessToken(t.Context(), da.DeviceCode, testConfirmationCode) require.NoError(t, err) assert.Equal(t, "access-token", tok.AccessToken) } -func TestDeviceCodeChallenge_Errors(t *testing.T) { +func TestDeviceCodeLogin_Errors(t *testing.T) { tests := []struct { - name string - server *fakeDeviceServer - wantErr string + name string + server *fakeDeviceServer + wantErr string + wantInvalidDeviceCode bool }{ { name: "device authorization rejected", @@ -175,7 +190,7 @@ func TestDeviceCodeChallenge_Errors(t *testing.T) { tokenStatus: http.StatusBadRequest, tokenBody: `{"error": "invalid_client", "error_description": "Client is not allowed to use device code"}`, }, - wantErr: "token exchange: Client is not allowed to use device code (error code: invalid_client)", + wantErr: "Client is not allowed to use device code (error code: invalid_client)", }, { name: "token request rate limited", @@ -183,7 +198,7 @@ func TestDeviceCodeChallenge_Errors(t *testing.T) { tokenStatus: http.StatusTooManyRequests, tokenBody: `{"error": "invalid_request", "error_description": "Device grant rate limit exceeded"}`, }, - wantErr: "token exchange: Device grant rate limit exceeded (error code: invalid_request)", + wantErr: "Device grant rate limit exceeded (error code: invalid_request)", }, { name: "unknown device code", @@ -191,24 +206,24 @@ func TestDeviceCodeChallenge_Errors(t *testing.T) { tokenStatus: http.StatusBadRequest, tokenBody: `{"error": "invalid_request", "error_description": "Invalid device code"}`, }, - wantErr: "token exchange: Invalid device code (error code: invalid_request)", + wantErr: "Invalid device code (error code: invalid_request)", + wantInvalidDeviceCode: true, }, } for _, tc := range tests { t.Run(tc.name, func(t *testing.T) { tc.server.expiresIn = 600 - p := newDeviceCodeAuth(t, tc.server, true) - _, err := p.DeviceCodeChallenge(t.Context(), func(_ context.Context, _ DeviceAuthorization, _ error, _ int) (string, error) { - return testConfirmationCode, nil - }) + _, err := runDeviceCodeLogin(t.Context(), newDeviceCodeLogin(t, tc.server), testConfirmationCode) assert.EqualError(t, err, tc.wantErr) - // None of these errors asks for the confirmation code again. - assert.LessOrEqual(t, tc.server.tokenRequests.Load(), int32(1)) + // None of these errors can be fixed by entering the confirmation + // code again. + assert.NotErrorIs(t, err, ErrInvalidConfirmationCode) + assert.Equal(t, tc.wantInvalidDeviceCode, errors.Is(err, ErrInvalidDeviceCode)) }) } } -func TestDeviceCodeChallenge_IncompleteDeviceAuthorization(t *testing.T) { +func TestRequestDeviceAuthorization_Incomplete(t *testing.T) { for _, missing := range []string{"device_code", "user_code", "verification_uri"} { t.Run(missing, func(t *testing.T) { resp := map[string]any{ @@ -221,90 +236,50 @@ func TestDeviceCodeChallenge_IncompleteDeviceAuthorization(t *testing.T) { body, err := json.Marshal(resp) require.NoError(t, err) f := &fakeDeviceServer{deviceAuthStatus: http.StatusOK, deviceAuthBody: string(body)} - p := newDeviceCodeAuth(t, f, true) - _, err = p.DeviceCodeChallenge(t.Context(), func(_ context.Context, _ DeviceAuthorization, _ error, _ int) (string, error) { - t.Fatal("confirm must not be called") - return "", nil - }) + _, err = newDeviceCodeLogin(t, f).RequestDeviceAuthorization(t.Context()) assert.EqualError(t, err, "device authorization: server response missing "+missing) assert.Equal(t, int32(0), f.tokenRequests.Load()) }) } } -func TestDeviceCodeChallenge_ErrorKeepsHTTPError(t *testing.T) { +func TestRequestDeviceAccessToken_ErrorKeepsHTTPError(t *testing.T) { f := &fakeDeviceServer{ expiresIn: 600, tokenStatus: http.StatusBadRequest, tokenBody: `{"error": "invalid_client", "error_description": "Client is not allowed to use device code"}`, } - confirm := func(_ context.Context, _ DeviceAuthorization, _ error, _ int) (string, error) { - return testConfirmationCode, nil - } - _, err := newDeviceCodeAuth(t, f, true).DeviceCodeChallenge(t.Context(), confirm) + _, err := newDeviceCodeLogin(t, f).RequestDeviceAccessToken(t.Context(), testDeviceCode, testConfirmationCode) httpErr, ok := errors.AsType[*httpclient.HttpError](err) require.True(t, ok, "want *httpclient.HttpError in %v", err) assert.Equal(t, http.StatusBadRequest, httpErr.StatusCode) // A plain HTTP client hands the error response to the CLI as is. - _, err = newDeviceCodeAuth(t, f, true, WithHttpClient(&http.Client{})).DeviceCodeChallenge(t.Context(), confirm) + _, err = newDeviceCodeLogin(t, f, WithHttpClient(&http.Client{})).RequestDeviceAccessToken(t.Context(), testDeviceCode, testConfirmationCode) retrieveErr, ok := errors.AsType[*oauth2.RetrieveError](err) require.True(t, ok, "want *oauth2.RetrieveError in %v", err) assert.Equal(t, http.StatusBadRequest, retrieveErr.Response.StatusCode) - assert.EqualError(t, err, "token exchange: Client is not allowed to use device code (error code: invalid_client)") + assert.EqualError(t, err, "Client is not allowed to use device code (error code: invalid_client)") } -func TestDeviceCodeChallenge_RetriesWrongConfirmationCode(t *testing.T) { +func TestRequestDeviceAccessToken_RetriesWrongConfirmationCode(t *testing.T) { f := &fakeDeviceServer{expiresIn: 600} - p := newDeviceCodeAuth(t, f, true) - - var lastErrs []string - tok, err := p.DeviceCodeChallenge(t.Context(), func(_ context.Context, _ DeviceAuthorization, lastErr error, _ int) (string, error) { - if lastErr == nil { - lastErrs = append(lastErrs, "") - return "AAAAAAAA", nil - } - lastErrs = append(lastErrs, lastErr.Error()) - return testConfirmationCode, nil - }) + l := newDeviceCodeLogin(t, f) + + _, err := l.RequestDeviceAccessToken(t.Context(), testDeviceCode, "AAAAAAAA") + assert.ErrorIs(t, err, ErrInvalidConfirmationCode) + assert.EqualError(t, err, "Invalid confirmation code (error code: invalid_request)") + _, ok := errors.AsType[*httpclient.HttpError](err) + assert.True(t, ok, "want *httpclient.HttpError in %v", err) + tok, err := l.RequestDeviceAccessToken(t.Context(), testDeviceCode, testConfirmationCode) require.NoError(t, err) assert.Equal(t, "access-token", tok.AccessToken) - assert.Equal(t, []string{"", "Invalid confirmation code (error code: invalid_request)"}, lastErrs) assert.Equal(t, int32(2), f.tokenRequests.Load()) } -func TestDeviceCodeChallenge_TooManyWrongConfirmationCodes(t *testing.T) { - f := &fakeDeviceServer{expiresIn: 600} - p := newDeviceCodeAuth(t, f, true) - - var attemptsLeft []int - _, err := p.DeviceCodeChallenge(t.Context(), func(_ context.Context, _ DeviceAuthorization, _ error, left int) (string, error) { - attemptsLeft = append(attemptsLeft, left) - return "AAAAAAAA", nil - }) - assert.EqualError(t, err, "Invalid confirmation code (error code: invalid_request). Start login from scratch") - assert.Equal(t, []int{3, 2, 1}, attemptsLeft) - assert.Equal(t, int32(maxConfirmationAttempts), f.tokenRequests.Load()) -} - -func TestDeviceCodeChallenge_ConfirmErrorOnRetry(t *testing.T) { - f := &fakeDeviceServer{expiresIn: 600} - p := newDeviceCodeAuth(t, f, true) - confirmErr := errors.New("no input") - - _, err := p.DeviceCodeChallenge(t.Context(), func(_ context.Context, _ DeviceAuthorization, lastErr error, _ int) (string, error) { - if lastErr != nil { - return "", confirmErr - } - return "AAAAAAAA", nil - }) - assert.ErrorIs(t, err, confirmErr) - assert.Equal(t, int32(1), f.tokenRequests.Load()) -} - -func TestDeviceCodeChallenge_DoesNotPoll(t *testing.T) { +func TestRequestDeviceAccessToken_DoesNotPoll(t *testing.T) { // A plain HTTP client hands the error response to the CLI as is, rather // than turning it into an HttpError. A pending login must fail right away // instead of being polled until it is approved. @@ -313,78 +288,53 @@ func TestDeviceCodeChallenge_DoesNotPoll(t *testing.T) { tokenStatus: http.StatusBadRequest, tokenBody: `{"error": "authorization_pending", "error_description": "Login is not approved yet"}`, } - p := newDeviceCodeAuth(t, f, true, WithHttpClient(&http.Client{})) + l := newDeviceCodeLogin(t, f, WithHttpClient(&http.Client{})) - _, err := p.DeviceCodeChallenge(t.Context(), func(_ context.Context, _ DeviceAuthorization, _ error, _ int) (string, error) { - return testConfirmationCode, nil - }) - assert.EqualError(t, err, "token exchange: Login is not approved yet (error code: authorization_pending)") + _, err := l.RequestDeviceAccessToken(t.Context(), testDeviceCode, testConfirmationCode) + assert.EqualError(t, err, "Login is not approved yet (error code: authorization_pending)") assert.Equal(t, int32(1), f.tokenRequests.Load()) } -func TestDeviceCodeChallenge_MissingAccessToken(t *testing.T) { +func TestRequestDeviceAccessToken_MissingAccessToken(t *testing.T) { f := &fakeDeviceServer{ expiresIn: 600, tokenStatus: http.StatusOK, tokenBody: `{"token_type": "Bearer"}`, } - p := newDeviceCodeAuth(t, f, true) - - _, err := p.DeviceCodeChallenge(t.Context(), func(_ context.Context, _ DeviceAuthorization, _ error, _ int) (string, error) { - return testConfirmationCode, nil - }) - assert.EqualError(t, err, "token exchange: server response missing access_token") -} + l := newDeviceCodeLogin(t, f) -func TestDeviceCodeChallenge_ConfirmError(t *testing.T) { - f := &fakeDeviceServer{expiresIn: 600} - p := newDeviceCodeAuth(t, f, true) - confirmErr := errors.New("no input") - - _, err := p.DeviceCodeChallenge(t.Context(), func(_ context.Context, _ DeviceAuthorization, _ error, _ int) (string, error) { - return "", confirmErr - }) - assert.ErrorIs(t, err, confirmErr) - assert.Equal(t, int32(0), f.tokenRequests.Load()) + _, err := l.RequestDeviceAccessToken(t.Context(), testDeviceCode, testConfirmationCode) + assert.EqualError(t, err, "server response missing access_token") } -func TestDeviceCodeChallenge_DeviceAuthorizationCancelled(t *testing.T) { - p := newDeviceCodeAuth(t, &fakeDeviceServer{deviceAuthStalls: true}, true) +func TestRequestDeviceAuthorization_Cancelled(t *testing.T) { + l := newDeviceCodeLogin(t, &fakeDeviceServer{deviceAuthStalls: true}) ctx, cancel := context.WithTimeout(t.Context(), 100*time.Millisecond) defer cancel() start := time.Now() - _, err := p.DeviceCodeChallenge(ctx, func(_ context.Context, _ DeviceAuthorization, _ error, _ int) (string, error) { - t.Fatal("confirm must not be called") - return "", nil - }) + _, err := l.RequestDeviceAuthorization(ctx) assert.ErrorIs(t, err, context.DeadlineExceeded) // The HTTP client's own 30 second timeout also ends the request with // DeadlineExceeded, so check that the context ended it. assert.Less(t, time.Since(start), 10*time.Second) } -func TestDeviceCodeChallenge_NotAdvertised(t *testing.T) { +func TestDeviceCodeLogin_NotAdvertised(t *testing.T) { p := newDeviceCodeAuth(t, &fakeDeviceServer{}, false) - _, err := p.DeviceCodeChallenge(t.Context(), func(_ context.Context, _ DeviceAuthorization, _ error, _ int) (string, error) { - t.Fatal("confirm must not be called") - return "", nil - }) + _, err := p.DeviceCodeLogin() assert.ErrorIs(t, err, ErrDeviceCodeNotSupported) } -func TestDeviceCodeChallenge_Resources(t *testing.T) { +func TestDeviceCodeLogin_Resources(t *testing.T) { p := newDeviceCodeAuth(t, &fakeDeviceServer{}, true, WithResources([]string{"https://resource.test"})) - _, err := p.DeviceCodeChallenge(t.Context(), func(_ context.Context, _ DeviceAuthorization, _ error, _ int) (string, error) { - t.Fatal("confirm must not be called") - return "", nil - }) + _, err := p.DeviceCodeLogin() assert.EqualError(t, err, "device code login does not support resource indicators") } -func TestDeviceCodeChallenge_AccountHost(t *testing.T) { +func TestDeviceCodeLogin_AccountHost(t *testing.T) { arg, err := NewBasicAccountOAuthArgument("https://accounts.cloud.databricks.test", "xyz") require.NoError(t, err) p, err := NewPersistentAuth(t.Context(), @@ -394,10 +344,7 @@ func TestDeviceCodeChallenge_AccountHost(t *testing.T) { require.NoError(t, err) defer p.Close() - _, err = p.DeviceCodeChallenge(t.Context(), func(_ context.Context, _ DeviceAuthorization, _ error, _ int) (string, error) { - t.Fatal("confirm must not be called") - return "", nil - }) + _, err = p.DeviceCodeLogin() assert.EqualError(t, err, "device code login is only supported for workspace hosts") } From 33b2751817185b3dc53e8917ea4f8b122d01e64b Mon Sep 17 00:00:00 2001 From: Maksim Zuev Date: Thu, 29 Oct 2026 12:51:07 +0000 Subject: [PATCH 2/6] Test confirmation code retries in auth login The retries moved from u2m to auth login without unit tests. The new test enters a wrong code followed by the right one, three wrong codes, and a wrong code followed by no input, and checks the number of token requests for each. Also rename PersistentAuth.DeviceCodeLogin to NewDeviceCodeLogin, which had the same name as the DeviceCodeLogin type it returns, rename TestRequestDeviceAccessToken_RetriesWrongConfirmationCode to TestRequestDeviceAccessToken_WrongConfirmationCode because it no longer retries, and check the token requests in TestDeviceCodeLogin_Errors again. Co-authored-by: Isaac --- cmd/auth/login_device_code.go | 2 +- cmd/auth/login_device_code_test.go | 65 ++++++++++++++++++++++++++++++ libs/auth/u2m/device_code.go | 6 +-- libs/auth/u2m/device_code_test.go | 17 ++++---- 4 files changed, 78 insertions(+), 12 deletions(-) diff --git a/cmd/auth/login_device_code.go b/cmd/auth/login_device_code.go index 0c339eff30e..40bd55347db 100644 --- a/cmd/auth/login_device_code.go +++ b/cmd/auth/login_device_code.go @@ -69,7 +69,7 @@ func validateDeviceCodeLogin(ctx context.Context, cmd *cobra.Command) error { // code valid, so the user can re-enter it up to maxConfirmationAttempts // times. func deviceCodeLogin(ctx context.Context, persistentAuth *u2m.PersistentAuth) (*oauth2.Token, error) { - login, err := persistentAuth.DeviceCodeLogin() + login, err := persistentAuth.NewDeviceCodeLogin() if err != nil { return nil, err } diff --git a/cmd/auth/login_device_code_test.go b/cmd/auth/login_device_code_test.go index 2f24f11b9fc..d19591de382 100644 --- a/cmd/auth/login_device_code_test.go +++ b/cmd/auth/login_device_code_test.go @@ -8,6 +8,7 @@ import ( "net/http" "runtime" "strings" + "sync/atomic" "testing" "time" @@ -326,3 +327,67 @@ func TestLoginDeviceCodeFailsOnTokenExchangeError(t *testing.T) { _, err := tm.Result() assert.EqualError(t, err, "token exchange: Device grant rate limit exceeded (error code: invalid_request)") } + +// A wrong confirmation code leaves the device code valid, so the user can +// enter the code again until the attempts run out. +func TestLoginDeviceCodeWrongConfirmationCode(t *testing.T) { + // The prompt for each attempt, which the test waits for before typing + // the next code. + prompts := []string{"Confirmation code", "(2 attempts left)", "(1 attempt left)"} + tests := []struct { + name string + codes []string + // closeStdin ends the input after the codes, so that the next + // prompt fails. + closeStdin bool + wantErr string + wantTokenRequests int32 + }{ + { + name: "right code after a wrong one", + codes: []string{"AAAA-AAAA", "FTFA-QQKY"}, + wantTokenRequests: 2, + }, + { + name: "too many wrong codes", + codes: []string{"AAAA-AAAA", "BBBB-BBBB", "CCCC-CCCC"}, + wantErr: "Invalid confirmation code (error code: invalid_request). Start login from scratch", + wantTokenRequests: 3, + }, + { + name: "no code after a wrong one", + codes: []string{"AAAA-AAAA"}, + closeStdin: true, + wantErr: "reading confirmation code: EOF", + wantTokenRequests: 1, + }, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + server := testserver.New(t) + testserver.AddDefaultHandlers(server) + var tokenRequests atomic.Int32 + server.RequestCallback = func(req *testserver.Request) { + if req.URL.Path == "/oidc/v1/token" { + tokenRequests.Add(1) + } + } + tm := startDeviceCodeLogin(t, "[test]\nhost = "+server.URL+"\n", "--profile", "test") + for i, code := range tt.codes { + tm.WaitFor(prompts[i]) + tm.Type(code + "\n") + } + if tt.closeStdin { + tm.WaitFor(prompts[len(tt.codes)]) + tm.Close() + } + _, err := tm.Result() + if tt.wantErr != "" { + assert.EqualError(t, err, tt.wantErr) + } else { + assert.NoError(t, err) + } + assert.Equal(t, tt.wantTokenRequests, tokenRequests.Load()) + }) + } +} diff --git a/libs/auth/u2m/device_code.go b/libs/auth/u2m/device_code.go index a0ca3886791..46f3a8c762e 100644 --- a/libs/auth/u2m/device_code.go +++ b/libs/auth/u2m/device_code.go @@ -78,16 +78,16 @@ func (a *PersistentAuth) deviceCodeConfig() (*oauth2.Config, error) { // DeviceCodeLogin makes the requests of a login using the Databricks variant // of the RFC 8628 device authorization grant. Get one with -// PersistentAuth.DeviceCodeLogin. +// PersistentAuth.NewDeviceCodeLogin. type DeviceCodeLogin struct { auth *PersistentAuth cfg *oauth2.Config } -// DeviceCodeLogin fetches the OAuth endpoints for a device code login, or +// NewDeviceCodeLogin fetches the OAuth endpoints for a device code login, or // returns an error if the host does not support it. Only workspace hosts are // supported. -func (a *PersistentAuth) DeviceCodeLogin() (*DeviceCodeLogin, error) { +func (a *PersistentAuth) NewDeviceCodeLogin() (*DeviceCodeLogin, error) { cfg, err := a.deviceCodeConfig() if err != nil { return nil, err diff --git a/libs/auth/u2m/device_code_test.go b/libs/auth/u2m/device_code_test.go index 03cffa0145b..b7a6fec3f4e 100644 --- a/libs/auth/u2m/device_code_test.go +++ b/libs/auth/u2m/device_code_test.go @@ -120,7 +120,7 @@ func newDeviceCodeAuth(t *testing.T, f *fakeDeviceServer, advertise bool, opts . } func newDeviceCodeLogin(t *testing.T, f *fakeDeviceServer, opts ...PersistentAuthOption) *DeviceCodeLogin { - l, err := newDeviceCodeAuth(t, f, true, opts...).DeviceCodeLogin() + l, err := newDeviceCodeAuth(t, f, true, opts...).NewDeviceCodeLogin() require.NoError(t, err) return l } @@ -215,6 +215,7 @@ func TestDeviceCodeLogin_Errors(t *testing.T) { tc.server.expiresIn = 600 _, err := runDeviceCodeLogin(t.Context(), newDeviceCodeLogin(t, tc.server), testConfirmationCode) assert.EqualError(t, err, tc.wantErr) + assert.LessOrEqual(t, tc.server.tokenRequests.Load(), int32(1)) // None of these errors can be fixed by entering the confirmation // code again. assert.NotErrorIs(t, err, ErrInvalidConfirmationCode) @@ -264,7 +265,7 @@ func TestRequestDeviceAccessToken_ErrorKeepsHTTPError(t *testing.T) { assert.EqualError(t, err, "Client is not allowed to use device code (error code: invalid_client)") } -func TestRequestDeviceAccessToken_RetriesWrongConfirmationCode(t *testing.T) { +func TestRequestDeviceAccessToken_WrongConfirmationCode(t *testing.T) { f := &fakeDeviceServer{expiresIn: 600} l := newDeviceCodeLogin(t, f) @@ -320,21 +321,21 @@ func TestRequestDeviceAuthorization_Cancelled(t *testing.T) { assert.Less(t, time.Since(start), 10*time.Second) } -func TestDeviceCodeLogin_NotAdvertised(t *testing.T) { +func TestNewDeviceCodeLogin_NotAdvertised(t *testing.T) { p := newDeviceCodeAuth(t, &fakeDeviceServer{}, false) - _, err := p.DeviceCodeLogin() + _, err := p.NewDeviceCodeLogin() assert.ErrorIs(t, err, ErrDeviceCodeNotSupported) } -func TestDeviceCodeLogin_Resources(t *testing.T) { +func TestNewDeviceCodeLogin_Resources(t *testing.T) { p := newDeviceCodeAuth(t, &fakeDeviceServer{}, true, WithResources([]string{"https://resource.test"})) - _, err := p.DeviceCodeLogin() + _, err := p.NewDeviceCodeLogin() assert.EqualError(t, err, "device code login does not support resource indicators") } -func TestDeviceCodeLogin_AccountHost(t *testing.T) { +func TestNewDeviceCodeLogin_AccountHost(t *testing.T) { arg, err := NewBasicAccountOAuthArgument("https://accounts.cloud.databricks.test", "xyz") require.NoError(t, err) p, err := NewPersistentAuth(t.Context(), @@ -344,7 +345,7 @@ func TestDeviceCodeLogin_AccountHost(t *testing.T) { require.NoError(t, err) defer p.Close() - _, err = p.DeviceCodeLogin() + _, err = p.NewDeviceCodeLogin() assert.EqualError(t, err, "device code login is only supported for workspace hosts") } From b19d8b7f78bee0137b6e405db5f3d501c93f998c Mon Sep 17 00:00:00 2001 From: Maksim Zuev Date: Thu, 29 Oct 2026 13:34:02 +0000 Subject: [PATCH 3/6] Test a device code login that gets no confirmation code The retry test covered missing input only after a wrong code. The new case closes stdin at the first prompt and checks that no token request is sent. Co-authored-by: Isaac --- cmd/auth/login_device_code_test.go | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/cmd/auth/login_device_code_test.go b/cmd/auth/login_device_code_test.go index d19591de382..871f4a3053b 100644 --- a/cmd/auth/login_device_code_test.go +++ b/cmd/auth/login_device_code_test.go @@ -361,6 +361,12 @@ func TestLoginDeviceCodeWrongConfirmationCode(t *testing.T) { wantErr: "reading confirmation code: EOF", wantTokenRequests: 1, }, + { + name: "no code", + closeStdin: true, + wantErr: "reading confirmation code: EOF", + wantTokenRequests: 0, + }, } for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { From 4ea745d6e3e312d26a561f4a255b4692838ea7a0 Mon Sep 17 00:00:00 2001 From: Maksim Zuev Date: Thu, 29 Oct 2026 15:18:16 +0000 Subject: [PATCH 4/6] Pass a context to NewDeviceCodeLogin NewDeviceCodeLogin fetches the OAuth endpoints, and used the context that PersistentAuth was created with. It now takes the context of the call, like the two steps of the login, and passes it to oauth2Config. Co-authored-by: Isaac --- cmd/auth/login_device_code.go | 2 +- libs/auth/u2m/device_code.go | 8 ++++---- libs/auth/u2m/device_code_test.go | 8 ++++---- libs/auth/u2m/persistent_auth.go | 12 ++++++------ libs/auth/u2m/persistent_auth_test.go | 4 ++-- 5 files changed, 17 insertions(+), 17 deletions(-) diff --git a/cmd/auth/login_device_code.go b/cmd/auth/login_device_code.go index 40bd55347db..e960c6dd2b9 100644 --- a/cmd/auth/login_device_code.go +++ b/cmd/auth/login_device_code.go @@ -69,7 +69,7 @@ func validateDeviceCodeLogin(ctx context.Context, cmd *cobra.Command) error { // code valid, so the user can re-enter it up to maxConfirmationAttempts // times. func deviceCodeLogin(ctx context.Context, persistentAuth *u2m.PersistentAuth) (*oauth2.Token, error) { - login, err := persistentAuth.NewDeviceCodeLogin() + login, err := persistentAuth.NewDeviceCodeLogin(ctx) if err != nil { return nil, err } diff --git a/libs/auth/u2m/device_code.go b/libs/auth/u2m/device_code.go index 46f3a8c762e..5feb16f0839 100644 --- a/libs/auth/u2m/device_code.go +++ b/libs/auth/u2m/device_code.go @@ -57,7 +57,7 @@ type DeviceAuthorization struct { // deviceCodeConfig returns the OAuth config for device code login, or an // error if the host does not support it. -func (a *PersistentAuth) deviceCodeConfig() (*oauth2.Config, error) { +func (a *PersistentAuth) deviceCodeConfig(ctx context.Context) (*oauth2.Config, error) { if _, ok := a.oAuthArgument.(WorkspaceOAuthArgument); !ok { return nil, errors.New("device code login is only supported for workspace hosts") } @@ -66,7 +66,7 @@ func (a *PersistentAuth) deviceCodeConfig() (*oauth2.Config, error) { if len(a.resources) > 0 { return nil, errors.New("device code login does not support resource indicators") } - cfg, err := a.oauth2Config() + cfg, err := a.oauth2Config(ctx) if err != nil { return nil, fmt.Errorf("fetching oauth config: %w", err) } @@ -87,8 +87,8 @@ type DeviceCodeLogin struct { // NewDeviceCodeLogin fetches the OAuth endpoints for a device code login, or // returns an error if the host does not support it. Only workspace hosts are // supported. -func (a *PersistentAuth) NewDeviceCodeLogin() (*DeviceCodeLogin, error) { - cfg, err := a.deviceCodeConfig() +func (a *PersistentAuth) NewDeviceCodeLogin(ctx context.Context) (*DeviceCodeLogin, error) { + cfg, err := a.deviceCodeConfig(ctx) if err != nil { return nil, err } diff --git a/libs/auth/u2m/device_code_test.go b/libs/auth/u2m/device_code_test.go index b7a6fec3f4e..1e30a94df48 100644 --- a/libs/auth/u2m/device_code_test.go +++ b/libs/auth/u2m/device_code_test.go @@ -120,7 +120,7 @@ func newDeviceCodeAuth(t *testing.T, f *fakeDeviceServer, advertise bool, opts . } func newDeviceCodeLogin(t *testing.T, f *fakeDeviceServer, opts ...PersistentAuthOption) *DeviceCodeLogin { - l, err := newDeviceCodeAuth(t, f, true, opts...).NewDeviceCodeLogin() + l, err := newDeviceCodeAuth(t, f, true, opts...).NewDeviceCodeLogin(t.Context()) require.NoError(t, err) return l } @@ -324,14 +324,14 @@ func TestRequestDeviceAuthorization_Cancelled(t *testing.T) { func TestNewDeviceCodeLogin_NotAdvertised(t *testing.T) { p := newDeviceCodeAuth(t, &fakeDeviceServer{}, false) - _, err := p.NewDeviceCodeLogin() + _, err := p.NewDeviceCodeLogin(t.Context()) assert.ErrorIs(t, err, ErrDeviceCodeNotSupported) } func TestNewDeviceCodeLogin_Resources(t *testing.T) { p := newDeviceCodeAuth(t, &fakeDeviceServer{}, true, WithResources([]string{"https://resource.test"})) - _, err := p.NewDeviceCodeLogin() + _, err := p.NewDeviceCodeLogin(t.Context()) assert.EqualError(t, err, "device code login does not support resource indicators") } @@ -345,7 +345,7 @@ func TestNewDeviceCodeLogin_AccountHost(t *testing.T) { require.NoError(t, err) defer p.Close() - _, err = p.NewDeviceCodeLogin() + _, err = p.NewDeviceCodeLogin(t.Context()) assert.EqualError(t, err, "device code login is only supported for workspace hosts") } diff --git a/libs/auth/u2m/persistent_auth.go b/libs/auth/u2m/persistent_auth.go index de0df97d47e..1e3709ef04d 100644 --- a/libs/auth/u2m/persistent_auth.go +++ b/libs/auth/u2m/persistent_auth.go @@ -412,7 +412,7 @@ func (a *PersistentAuth) refresh(oldToken *oauth2.Token) (*oauth2.Token, error) if oldToken.RefreshToken == "" { return nil, ErrMissingRefreshToken } - cfg, err := a.oauth2Config() + cfg, err := a.oauth2Config(a.ctx) if err != nil { return nil, err } @@ -460,7 +460,7 @@ func (a *PersistentAuth) Challenge() (*oauth2.Token, error) { // the callback server is not created, we need to close the listener manually. defer a.Close() - cfg, err := a.oauth2Config() + cfg, err := a.oauth2Config(a.ctx) if err != nil { return nil, fmt.Errorf("fetching oauth config: %w", err) } @@ -591,20 +591,20 @@ func (a *PersistentAuth) resolveScopes() []string { } // oauth2Config returns the OAuth2 configuration for the given OAuthArgument. -func (a *PersistentAuth) oauth2Config() (*oauth2.Config, error) { +func (a *PersistentAuth) oauth2Config(ctx context.Context) (*oauth2.Config, error) { scopes := a.resolveScopes() var endpoints *OAuthAuthorizationServer var err error switch argg := a.oAuthArgument.(type) { case WorkspaceOAuthArgument: - endpoints, err = a.endpointSupplier.GetWorkspaceOAuthEndpoints(a.ctx, argg.GetWorkspaceHost()) + endpoints, err = a.endpointSupplier.GetWorkspaceOAuthEndpoints(ctx, argg.GetWorkspaceHost()) case AccountOAuthArgument: endpoints, err = a.endpointSupplier.GetAccountOAuthEndpoints( - a.ctx, argg.GetAccountHost(), argg.GetAccountId(), + ctx, argg.GetAccountHost(), argg.GetAccountId(), ) case UnifiedOAuthArgument: - endpoints, err = a.endpointSupplier.GetUnifiedOAuthEndpoints(a.ctx, argg.GetHost(), argg.GetAccountId()) + endpoints, err = a.endpointSupplier.GetUnifiedOAuthEndpoints(ctx, argg.GetHost(), argg.GetAccountId()) case DiscoveryOAuthArgument: return nil, fmt.Errorf("discovery OAuthArgument %T is only supported with WithDiscoveryLogin during Challenge; after discovery, construct a WorkspaceOAuthArgument with the discovered host", a.oAuthArgument) default: diff --git a/libs/auth/u2m/persistent_auth_test.go b/libs/auth/u2m/persistent_auth_test.go index 1817ff778ae..74eefc3fad0 100644 --- a/libs/auth/u2m/persistent_auth_test.go +++ b/libs/auth/u2m/persistent_auth_test.go @@ -159,7 +159,7 @@ func TestPersistentAuthClientID(t *testing.T) { if err != nil { t.Fatalf("NewPersistentAuth(): %v", err) } - cfg, err := p.oauth2Config() + cfg, err := p.oauth2Config(t.Context()) if err != nil { t.Fatalf("oauth2Config(): %v", err) } @@ -206,7 +206,7 @@ func TestPersistentAuthResources(t *testing.T) { if err != nil { t.Fatalf("NewPersistentAuth(): %v", err) } - cfg, err := p.oauth2Config() + cfg, err := p.oauth2Config(t.Context()) if err != nil { t.Fatalf("oauth2Config(): %v", err) } From 78267844d5e1a8c7e7db86bb7365420a41904d5e Mon Sep 17 00:00:00 2001 From: Maksim Zuev Date: Thu, 29 Oct 2026 15:18:41 +0000 Subject: [PATCH 5/6] Check the exact number of token requests in TestDeviceCodeLogin_Errors Each case now states how many token requests it expects: none when the device authorization is rejected, and one otherwise. The previous upper bound of one also passed if a case that should reach the token endpoint never did. Co-authored-by: Isaac --- libs/auth/u2m/device_code_test.go | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/libs/auth/u2m/device_code_test.go b/libs/auth/u2m/device_code_test.go index 1e30a94df48..27968fd7d34 100644 --- a/libs/auth/u2m/device_code_test.go +++ b/libs/auth/u2m/device_code_test.go @@ -174,6 +174,7 @@ func TestDeviceCodeLogin_Errors(t *testing.T) { name string server *fakeDeviceServer wantErr string + wantTokenRequests int32 wantInvalidDeviceCode bool }{ { @@ -190,7 +191,8 @@ func TestDeviceCodeLogin_Errors(t *testing.T) { tokenStatus: http.StatusBadRequest, tokenBody: `{"error": "invalid_client", "error_description": "Client is not allowed to use device code"}`, }, - wantErr: "Client is not allowed to use device code (error code: invalid_client)", + wantErr: "Client is not allowed to use device code (error code: invalid_client)", + wantTokenRequests: 1, }, { name: "token request rate limited", @@ -198,7 +200,8 @@ func TestDeviceCodeLogin_Errors(t *testing.T) { tokenStatus: http.StatusTooManyRequests, tokenBody: `{"error": "invalid_request", "error_description": "Device grant rate limit exceeded"}`, }, - wantErr: "Device grant rate limit exceeded (error code: invalid_request)", + wantErr: "Device grant rate limit exceeded (error code: invalid_request)", + wantTokenRequests: 1, }, { name: "unknown device code", @@ -207,6 +210,7 @@ func TestDeviceCodeLogin_Errors(t *testing.T) { tokenBody: `{"error": "invalid_request", "error_description": "Invalid device code"}`, }, wantErr: "Invalid device code (error code: invalid_request)", + wantTokenRequests: 1, wantInvalidDeviceCode: true, }, } @@ -215,7 +219,7 @@ func TestDeviceCodeLogin_Errors(t *testing.T) { tc.server.expiresIn = 600 _, err := runDeviceCodeLogin(t.Context(), newDeviceCodeLogin(t, tc.server), testConfirmationCode) assert.EqualError(t, err, tc.wantErr) - assert.LessOrEqual(t, tc.server.tokenRequests.Load(), int32(1)) + assert.Equal(t, tc.wantTokenRequests, tc.server.tokenRequests.Load()) // None of these errors can be fixed by entering the confirmation // code again. assert.NotErrorIs(t, err, ErrInvalidConfirmationCode) From 84cdd4ec64ab8dbe9d670e562c0d1b00566b5669 Mon Sep 17 00:00:00 2001 From: Maksim Zuev Date: Thu, 29 Oct 2026 15:19:47 +0000 Subject: [PATCH 6/6] Test that auth login stops at an unknown device code The server answers both a wrong confirmation code and an unknown device code with HTTP 400 and invalid_request, and only a wrong confirmation code lets the user enter the code again. Only the u2m tests checked that an unknown device code is told apart, so auth login could ask for the code again with no test failing. TestLoginDeviceCodeFailsOnTokenExchangeError now covers an unknown device code next to the rate limit, and checks that each sends a single token request. Co-authored-by: Isaac --- cmd/auth/login_device_code_test.go | 60 +++++++++++++++++++++--------- 1 file changed, 43 insertions(+), 17 deletions(-) diff --git a/cmd/auth/login_device_code_test.go b/cmd/auth/login_device_code_test.go index 871f4a3053b..f8c11ff9c0c 100644 --- a/cmd/auth/login_device_code_test.go +++ b/cmd/auth/login_device_code_test.go @@ -311,25 +311,51 @@ resources = https://workspace.test/ai-gateway/mcp-services/system.ai.github // Only a wrong confirmation code can be fixed by entering the code again, so // any other token exchange error ends the login. func TestLoginDeviceCodeFailsOnTokenExchangeError(t *testing.T) { - server := testserver.New(t) - // The router keeps the first handler registered for a route, so this one - // replaces the default token endpoint. - server.Handle("POST", "/oidc/v1/token", func(_ testserver.Request) any { - return testserver.Response{ - StatusCode: http.StatusTooManyRequests, - Body: map[string]string{"error": "invalid_request", "error_description": "Device grant rate limit exceeded"}, - } - }) - testserver.AddDefaultHandlers(server) - tm := startDeviceCodeLogin(t, "[test]\nhost = "+server.URL+"\n", "--profile", "test") - tm.WaitFor("Confirmation code") - tm.Type("FTFA-QQKY\n") - _, err := tm.Result() - assert.EqualError(t, err, "token exchange: Device grant rate limit exceeded (error code: invalid_request)") + // Neither error can be fixed by entering the confirmation code again, so + // the login fails after one token request instead of asking again. + tests := []struct { + name string + status int + description string + wantErr string + }{ + { + name: "rate limited", + status: http.StatusTooManyRequests, + description: "Device grant rate limit exceeded", + wantErr: "token exchange: Device grant rate limit exceeded (error code: invalid_request)", + }, + { + name: "unknown device code", + status: http.StatusBadRequest, + description: "Invalid device code", + wantErr: "token exchange: Invalid device code (error code: invalid_request)", + }, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + server := testserver.New(t) + var tokenRequests atomic.Int32 + // The router keeps the first handler registered for a route, so + // this one replaces the default token endpoint. + server.Handle("POST", "/oidc/v1/token", func(_ testserver.Request) any { + tokenRequests.Add(1) + return testserver.Response{ + StatusCode: tt.status, + Body: map[string]string{"error": "invalid_request", "error_description": tt.description}, + } + }) + testserver.AddDefaultHandlers(server) + tm := startDeviceCodeLogin(t, "[test]\nhost = "+server.URL+"\n", "--profile", "test") + tm.WaitFor("Confirmation code") + tm.Type("FTFA-QQKY\n") + _, err := tm.Result() + assert.EqualError(t, err, tt.wantErr) + assert.Equal(t, int32(1), tokenRequests.Load()) + }) + } } -// A wrong confirmation code leaves the device code valid, so the user can -// enter the code again until the attempts run out. func TestLoginDeviceCodeWrongConfirmationCode(t *testing.T) { // The prompt for each attempt, which the test waits for before typing // the next code.