feat(device-flow): redirect SSO-complete failures to a dedicated device error view (#38185)

api/controllers/openapi/oauth_device_sso.py

@@ -13,8 +13,10 @@ handlers do redirects + cookie kwargs that don't fit the Resource shape.
 from __future__ import annotations
 
 import logging
+import re
 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 +76,21 @@ 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"}
+
+# user_code only ever reaches the redirect as a urlencoded query value; the
+# charset bound additionally forbids the path/scheme separators a redirection
+# attack would need, so an untrusted value cannot escape the fixed /device path.
+_USER_CODE_RE = re.compile(r"\A[A-Z0-9-]{1,16}\Z")
+
+
+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 and _USER_CODE_RE.match(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 +151,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 +173,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 +200,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(

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,97 @@ 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"]
+
+
+def test_device_error_redirect_drops_malformed_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", "https://evil.example/")
+    loc = resp.headers["Location"]
+    assert loc.startswith("/device?")
+    assert "user_code=" not in loc
+    assert "evil" not in loc
+
+
+# ---------------------------------------------------------------------------
+# 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
+    assert "sso_error=sso_failed" in loc
+    assert "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"]

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