diff --git a/api/controllers/openapi/oauth_device_sso.py b/api/controllers/openapi/oauth_device_sso.py index 0218d143309..f7f6feaa7a9 100644 --- a/api/controllers/openapi/oauth_device_sso.py +++ b/api/controllers/openapi/oauth_device_sso.py @@ -15,6 +15,7 @@ from __future__ import annotations import logging import secrets from dataclasses import dataclass +from urllib.parse import urlencode from flask import jsonify, make_response, redirect, request from pydantic import ValidationError @@ -74,6 +75,16 @@ STATE_ENVELOPE_TTL_SECONDS = 15 * 60 # Canonical sso-complete path. IdP-side ACS callback URL must point here. _SSO_COMPLETE_PATH = "/openapi/v1/oauth/device/sso-complete" +_ALLOWED_SSO_ERRORS = {"sso_failed", "email_belongs_to_dify_account"} + + +def _device_error_redirect(code: str, user_code: str | None = None): + safe_code = code if code in _ALLOWED_SSO_ERRORS else "sso_failed" + params: dict[str, str] = {"sso_error": safe_code} + if user_code: + params["user_code"] = user_code + return redirect(f"/device?{urlencode(params)}", code=302) + def _trusted_origin() -> str: base = (dify_config.CONSOLE_API_URL or "").rstrip("/") @@ -134,9 +145,21 @@ def sso_initiate(): @bp.route("/oauth/device/sso-complete", methods=["GET"]) @enterprise_only def sso_complete(): + try: + return _sso_complete_impl() + except Exception: + logger.exception("sso-complete: unhandled") + return _device_error_redirect("sso_failed") + + +def _sso_complete_impl(): + inbound_error = request.args.get("sso_error") + if inbound_error: + return _device_error_redirect(inbound_error, request.args.get("user_code")) + blob = request.args.get("sso_assertion") if not blob: - raise BadRequest("sso_assertion required") + return _device_error_redirect("sso_failed") keyset = jws.KeySet.from_shared_secret() @@ -144,25 +167,26 @@ def sso_complete(): raw_claims = jws.verify(keyset, blob, expected_aud=jws.AUD_EXT_SUBJECT_ASSERTION) except jws.VerifyError as e: logger.warning("sso-complete: rejected assertion: %s", e) - raise BadRequest("invalid_sso_assertion") from e + return _device_error_redirect("sso_failed") try: claims = ExtSubjectAssertionClaims.model_validate(raw_claims) except ValidationError as e: logger.warning("sso-complete: claim shape invalid: %s", e) - raise BadRequest("invalid_sso_assertion") from e - - if not consume_sso_assertion_nonce(redis_client, claims.nonce): - raise BadRequest("invalid_sso_assertion") + return _device_error_redirect("sso_failed") user_code = claims.user_code.strip().upper() + + if not consume_sso_assertion_nonce(redis_client, claims.nonce): + return _device_error_redirect("sso_failed", user_code) + store = DeviceFlowRedis(redis_client) found = store.load_by_user_code(user_code) if found is None: - raise Conflict("user_code_not_pending") + return _device_error_redirect("sso_failed", user_code) _, state = found if state.status is not DeviceFlowStatus.PENDING: - raise Conflict("user_code_not_pending") + return _device_error_redirect("sso_failed", user_code) if AccountService.has_active_account_with_email(db.session, claims.email): _emit_external_rejection_audit( @@ -170,7 +194,7 @@ def sso_complete(): _RejectedClaims(subject_email=claims.email, subject_issuer=claims.issuer), reason="email_belongs_to_dify_account", ) - return redirect("/device?sso_error=email_belongs_to_dify_account", code=302) + return _device_error_redirect("email_belongs_to_dify_account", user_code) iss = _trusted_origin() cookie_value, _ = mint_approval_grant( diff --git a/api/tests/unit_tests/controllers/openapi/test_device_sso.py b/api/tests/unit_tests/controllers/openapi/test_device_sso.py index 0125c583f04..d7e92b07a79 100644 --- a/api/tests/unit_tests/controllers/openapi/test_device_sso.py +++ b/api/tests/unit_tests/controllers/openapi/test_device_sso.py @@ -1,6 +1,7 @@ """SSO-branch device-flow endpoints under /openapi/v1/oauth/device/.""" import builtins +from unittest.mock import MagicMock, patch import pytest from flask import Flask @@ -77,3 +78,83 @@ def test_sso_complete_idp_callback_url_uses_canonical_path(): from controllers.openapi import oauth_device_sso assert oauth_device_sso._SSO_COMPLETE_PATH == "/openapi/v1/oauth/device/sso-complete" + + +# --------------------------------------------------------------------------- +# _device_error_redirect helper +# --------------------------------------------------------------------------- + + +def test_device_error_redirect_builds_relative_location(): + from controllers.openapi import oauth_device_sso + + app = Flask(__name__) + with app.test_request_context(): + resp = oauth_device_sso._device_error_redirect("sso_failed", "ABCD-1234") + assert resp.status_code == 302 + loc = resp.headers["Location"] + assert loc.startswith("/device?") + assert "sso_error=sso_failed" in loc + assert "user_code=ABCD-1234" in loc + + +def test_device_error_redirect_clamps_unknown_code(): + from controllers.openapi import oauth_device_sso + + app = Flask(__name__) + with app.test_request_context(): + resp = oauth_device_sso._device_error_redirect("totally-bogus") + assert "sso_error=sso_failed" in resp.headers["Location"] + + +def test_device_error_redirect_keeps_email_special_case(): + from controllers.openapi import oauth_device_sso + + app = Flask(__name__) + with app.test_request_context(): + resp = oauth_device_sso._device_error_redirect("email_belongs_to_dify_account", "ABCD-1234") + assert "sso_error=email_belongs_to_dify_account" in resp.headers["Location"] + + +def test_device_error_redirect_omits_empty_user_code(): + from controllers.openapi import oauth_device_sso + + app = Flask(__name__) + with app.test_request_context(): + resp = oauth_device_sso._device_error_redirect("sso_failed") + assert "user_code=" not in resp.headers["Location"] + + +# --------------------------------------------------------------------------- +# sso_complete redirect behaviour +# --------------------------------------------------------------------------- + + +def _ee_features(): + from services.feature_service import LicenseStatus + + m = MagicMock() + m.license.status = LicenseStatus.ACTIVE + return m + + +@patch("libs.device_flow_security.FeatureService.get_system_features") +def test_sso_complete_relays_inbound_sso_error(ee_feat, openapi_app): + ee_feat.return_value = _ee_features() + client = openapi_app.test_client() + resp = client.get( + "/openapi/v1/oauth/device/sso-complete?sso_error=sso_failed&user_code=ABCD-1234", + follow_redirects=False, + ) + assert resp.status_code == 302 + loc = resp.headers["Location"] + assert "/device?" in loc and "sso_error=sso_failed" in loc and "user_code=ABCD-1234" in loc + + +@patch("libs.device_flow_security.FeatureService.get_system_features") +def test_sso_complete_missing_assertion_redirects_generic(ee_feat, openapi_app): + ee_feat.return_value = _ee_features() + client = openapi_app.test_client() + resp = client.get("/openapi/v1/oauth/device/sso-complete", follow_redirects=False) + assert resp.status_code == 302 + assert "sso_error=sso_failed" in resp.headers["Location"] diff --git a/api/tests/unit_tests/controllers/openapi/test_oauth_sso_claims.py b/api/tests/unit_tests/controllers/openapi/test_oauth_sso_claims.py index 9eb4cf01f41..9551a37142f 100644 --- a/api/tests/unit_tests/controllers/openapi/test_oauth_sso_claims.py +++ b/api/tests/unit_tests/controllers/openapi/test_oauth_sso_claims.py @@ -34,8 +34,9 @@ def test_sso_complete_rejects_assertion_missing_email(ee_feat, jws_mod, app: Fla jws_mod.VerifyError = Exception client = app.test_client() - resp = client.get("/openapi/v1/oauth/device/sso-complete?sso_assertion=blob") - assert resp.status_code == 400, resp.data + resp = client.get("/openapi/v1/oauth/device/sso-complete?sso_assertion=blob", follow_redirects=False) + assert resp.status_code == 302, resp.data + assert "sso_error=sso_failed" in resp.headers["Location"] @patch("controllers.openapi.oauth_device_sso.jws") @@ -48,8 +49,9 @@ def test_sso_complete_rejects_assertion_empty_issuer(ee_feat, jws_mod, app: Flas jws_mod.VerifyError = Exception client = app.test_client() - resp = client.get("/openapi/v1/oauth/device/sso-complete?sso_assertion=blob") - assert resp.status_code == 400 + resp = client.get("/openapi/v1/oauth/device/sso-complete?sso_assertion=blob", follow_redirects=False) + assert resp.status_code == 302 + assert "sso_error=sso_failed" in resp.headers["Location"] def test_verify_approval_grant_raises_on_missing_field():