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..e960c6dd2b9 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.NewDeviceCodeLogin(ctx) + 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..f8c11ff9c0c 100644 --- a/cmd/auth/login_device_code_test.go +++ b/cmd/auth/login_device_code_test.go @@ -5,8 +5,10 @@ import ( "context" "errors" "io" + "net/http" "runtime" "strings" + "sync/atomic" "testing" "time" @@ -17,6 +19,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 +307,119 @@ 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) { + // 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()) + }) + } +} + +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, + }, + { + name: "no code", + closeStdin: true, + wantErr: "reading confirmation code: EOF", + wantTokenRequests: 0, + }, + } + 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 51f77f8e189..5feb16f0839 100644 --- a/libs/auth/u2m/device_code.go +++ b/libs/auth/u2m/device_code.go @@ -32,31 +32,32 @@ 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. -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") } @@ -65,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) } @@ -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) { - cfg, err := a.deviceCodeConfig() +// DeviceCodeLogin makes the requests of a login using the Databricks variant +// of the RFC 8628 device authorization grant. Get one with +// PersistentAuth.NewDeviceCodeLogin. +type DeviceCodeLogin struct { + auth *PersistentAuth + cfg *oauth2.Config +} + +// 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(ctx context.Context) (*DeviceCodeLogin, error) { + cfg, err := a.deviceCodeConfig(ctx) 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..27968fd7d34 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...).NewDeviceCodeLogin(t.Context()) + 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,28 @@ 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 + wantTokenRequests int32 + wantInvalidDeviceCode bool }{ { name: "device authorization rejected", @@ -175,7 +191,8 @@ 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)", + wantTokenRequests: 1, }, { name: "token request rate limited", @@ -183,7 +200,8 @@ 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)", + wantTokenRequests: 1, }, { name: "unknown device code", @@ -191,24 +209,26 @@ 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)", + wantTokenRequests: 1, + 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)) + 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) + 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 +241,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_WrongConfirmationCode(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 +293,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 TestNewDeviceCodeLogin_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.NewDeviceCodeLogin(t.Context()) assert.ErrorIs(t, err, ErrDeviceCodeNotSupported) } -func TestDeviceCodeChallenge_Resources(t *testing.T) { +func TestNewDeviceCodeLogin_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.NewDeviceCodeLogin(t.Context()) assert.EqualError(t, err, "device code login does not support resource indicators") } -func TestDeviceCodeChallenge_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(), @@ -394,10 +349,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.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) }