Split device code login into authorization and token exchange steps

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()
 		}

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

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())
+		})
+	}
+}

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.