Split device code login into authorization and token exchange steps

libs/auth/u2m/device_code.go

@@ -32,14 +32,27 @@ 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
+
+	// DeviceCode identifies the login in the token exchange.
+	DeviceCode string
 }
 
 // maxConfirmationAttempts is how many confirmation codes the user can enter
@@ -90,44 +103,117 @@ func (a *PersistentAuth) DeviceCodeChallenge(ctx context.Context, confirm Confir
 
 	da, err := a.requestDeviceAuthorization(ctx, cfg)
 	if err != nil {
-		return nil, fmt.Errorf("device authorization: %w", oauthError(err))
+		return nil, err
 	}
 
-	authorization := DeviceAuthorization{
-		UserCode:        da.UserCode,
-		VerificationURI: da.VerificationURI,
-	}
 	var lastErr error
 	for attempt := range maxConfirmationAttempts {
-		code, err := confirm(ctx, authorization, lastErr, maxConfirmationAttempts-attempt)
+		code, err := confirm(ctx, *da, lastErr, maxConfirmationAttempts-attempt)
 		if err != nil {
 			return nil, err
 		}
 
-		t, err := a.exchangeDeviceCode(ctx, cfg, da.DeviceCode, normalizeCode(code))
+		t, err := a.exchangeConfirmationCode(ctx, cfg, da.DeviceCode, 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 {
+		lastErr = err
+		if !errors.Is(err, ErrInvalidConfirmationCode) {
 			return nil, fmt.Errorf("token exchange: %w", lastErr)
 		}
 	}
 	return nil, fmt.Errorf("%w. Start login from scratch", lastErr)
 }
 
-// requestDeviceAuthorization makes the RFC 8628 device authorization request.
+// RequestDeviceAuthorization starts a device code login by making the RFC
+// 8628 device authorization request. The caller shows the user code and the
+// verification URI to the user and later passes the device code to
+// ExchangeDeviceCode. Only workspace hosts are supported.
+func (a *PersistentAuth) RequestDeviceAuthorization(ctx context.Context) (*DeviceAuthorization, error) {
+	cfg, err := a.deviceCodeConfig()
+	if err != nil {
+		return nil, err
+	}
+	return a.requestDeviceAuthorization(ctx, cfg)
+}
+
+// ExchangeDeviceCode completes a device code login started by
+// RequestDeviceAuthorization with a single token request, without retries.
+// 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. The caller is responsible for storing the
+// returned token.
+func (a *PersistentAuth) ExchangeDeviceCode(ctx context.Context, deviceCode, confirmationCode string) (*oauth2.Token, error) {
+	cfg, err := a.deviceCodeConfig()
+	if err != nil {
+		return nil, err
+	}
+	t, err := a.exchangeConfirmationCode(ctx, cfg, deviceCode, confirmationCode)
+	if err != nil {
+		return nil, fmt.Errorf("token exchange: %w", err)
+	}
+	return t, nil
+}
+
+// requestDeviceAuthorization makes the RFC 8628 device authorization request
+// and returns what the user needs to approve the login.
+func (a *PersistentAuth) requestDeviceAuthorization(ctx context.Context, cfg *oauth2.Config) (*DeviceAuthorization, error) {
+	da, err := a.postDeviceAuthorization(ctx, cfg)
+	if err != nil {
+		return nil, fmt.Errorf("device authorization: %w", oauthError(err))
+	}
+	return &DeviceAuthorization{
+		UserCode:        da.UserCode,
+		VerificationURI: da.VerificationURI,
+		DeviceCode:      da.DeviceCode,
+	}, nil
+}
+
+// exchangeConfirmationCode exchanges the device code and the confirmation
+// code for a token, and marks the errors for a rejected code.
+func (a *PersistentAuth) exchangeConfirmationCode(ctx context.Context, cfg *oauth2.Config, deviceCode, confirmationCode string) (*oauth2.Token, error) {
+	t, err := a.exchangeDeviceCode(ctx, 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}
+	}
+}
+
+// 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}
+}
+
+// postDeviceAuthorization makes the RFC 8628 device authorization request.
 // Unlike oauth2.Config.DeviceAuth, it sends the request with ctx, so that the
 // login timeout and cancellation stop a stalled request.
-func (a *PersistentAuth) requestDeviceAuthorization(ctx context.Context, cfg *oauth2.Config) (*oauth2.DeviceAuthResponse, error) {
+func (a *PersistentAuth) postDeviceAuthorization(ctx context.Context, cfg *oauth2.Config) (*oauth2.DeviceAuthResponse, error) {
 	body, err := a.postForm(ctx, cfg.Endpoint.DeviceAuthURL, url.Values{
 		"client_id": {cfg.ClientID},
 		"scope":     {strings.Join(cfg.Scopes, " ")},
@@ -235,10 +321,10 @@ func (a *PersistentAuth) postForm(ctx context.Context, endpoint string, form url
 	return body, nil
 }
 
-// normalizeCode lets users enter a code the way it is displayed. Databricks
+// NormalizeCode lets users enter a code the way it is displayed. Databricks
 // codes consist of 8 uppercase alphanumeric characters and are compared
 // case-sensitively, but are often shown as two groups of four separated by a
 // dash (e.g. FTFA-QQKY).
-func normalizeCode(code string) string {
+func NormalizeCode(code string) string {
 	return strings.NewReplacer("-", "", " ", "").Replace(strings.ToUpper(strings.TrimSpace(code)))
 }

libs/auth/u2m/device_code_test.go

@@ -133,6 +133,7 @@ func TestDeviceCodeChallenge(t *testing.T) {
 	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)
@@ -401,6 +402,88 @@ func TestDeviceCodeChallenge_AccountHost(t *testing.T) {
 	assert.EqualError(t, err, "device code login is only supported for workspace hosts")
 }
 
+func TestRequestDeviceAuthorization(t *testing.T) {
+	f := &fakeDeviceServer{expiresIn: 600}
+	p := newDeviceCodeAuth(t, f, true)
+
+	da, err := p.RequestDeviceAuthorization(t.Context())
+	require.NoError(t, err)
+	assert.Equal(t, &DeviceAuthorization{
+		UserCode:        testUserCode,
+		VerificationURI: "https://workspace.test/device/verification",
+		DeviceCode:      testDeviceCode,
+	}, da)
+	assert.Equal(t, int32(0), f.tokenRequests.Load())
+}
+
+func TestRequestDeviceAuthorization_Errors(t *testing.T) {
+	_, err := newDeviceCodeAuth(t, &fakeDeviceServer{}, false).RequestDeviceAuthorization(t.Context())
+	assert.ErrorIs(t, err, ErrDeviceCodeNotSupported)
+
+	_, err = newDeviceCodeAuth(t, &fakeDeviceServer{
+		deviceAuthStatus: http.StatusBadRequest,
+		deviceAuthBody:   `{"error": "unauthorized_client", "error_description": "Client is not allowed to use device code"}`,
+	}, true).RequestDeviceAuthorization(t.Context())
+	assert.EqualError(t, err, "device authorization: Client is not allowed to use device code (error code: unauthorized_client)")
+}
+
+func TestExchangeDeviceCode(t *testing.T) {
+	f := &fakeDeviceServer{}
+	p := newDeviceCodeAuth(t, f, true)
+
+	tok, err := p.ExchangeDeviceCode(t.Context(), testDeviceCode, " ftfa-qqky\n")
+	require.NoError(t, err)
+	assert.Equal(t, "access-token", tok.AccessToken)
+	assert.Equal(t, "refresh-token", tok.RefreshToken)
+	assert.Equal(t, int32(1), f.tokenRequests.Load())
+}
+
+func TestExchangeDeviceCode_Errors(t *testing.T) {
+	tests := []struct {
+		name                        string
+		server                      *fakeDeviceServer
+		wantErr                     string
+		wantInvalidConfirmationCode bool
+		wantInvalidDeviceCode       bool
+	}{
+		{
+			name:                        "wrong confirmation code",
+			server:                      &fakeDeviceServer{},
+			wantErr:                     "token exchange: Invalid confirmation code (error code: invalid_request)",
+			wantInvalidConfirmationCode: true,
+		},
+		{
+			name: "unknown device code",
+			server: &fakeDeviceServer{
+				tokenStatus: http.StatusBadRequest,
+				tokenBody:   `{"error": "invalid_request", "error_description": "Invalid device code"}`,
+			},
+			wantErr:               "token exchange: Invalid device code (error code: invalid_request)",
+			wantInvalidDeviceCode: true,
+		},
+		{
+			name: "token request rate limited",
+			server: &fakeDeviceServer{
+				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)",
+		},
+	}
+	for _, tc := range tests {
+		t.Run(tc.name, func(t *testing.T) {
+			p := newDeviceCodeAuth(t, tc.server, true)
+			_, err := p.ExchangeDeviceCode(t.Context(), testDeviceCode, "AAAAAAAA")
+			assert.EqualError(t, err, tc.wantErr)
+			assert.Equal(t, tc.wantInvalidConfirmationCode, errors.Is(err, ErrInvalidConfirmationCode))
+			assert.Equal(t, tc.wantInvalidDeviceCode, errors.Is(err, ErrInvalidDeviceCode))
+			_, ok := errors.AsType[*httpclient.HttpError](err)
+			assert.True(t, ok, "want *httpclient.HttpError in %v", err)
+			assert.Equal(t, int32(1), tc.server.tokenRequests.Load())
+		})
+	}
+}
+
 func TestNormalizeCode(t *testing.T) {
 	tests := []struct {
 		in   string
@@ -412,6 +495,6 @@ func TestNormalizeCode(t *testing.T) {
 		{"  FTFA QQKY \n", "FTFAQQKY"},
 	}
 	for _, tc := range tests {
-		assert.Equal(t, tc.want, normalizeCode(tc.in), tc.in)
+		assert.Equal(t, tc.want, NormalizeCode(tc.in), tc.in)
 	}
 }