feat(device-flow): redirect sso-complete failures to device page

All failure paths in sso_complete now return 302 redirects to
/device?sso_error=... instead of raw werkzeug error pages. Adds
_device_error_redirect helper that clamps codes to the two allowed
values, an outer catch-all try/except that logs + redirects on
unhandled exceptions, and relays any inbound sso_error from Go
enterprise straight through. user_code is now extracted before
the nonce is consumed so replay redirects can still carry it.
This commit is contained in:
GareArc
2026-06-29 19:18:15 -07:00
parent cb35c6fa98
commit febace6bcf
3 changed files with 120 additions and 13 deletions

View File

@ -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(

View File

@ -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"]

View File

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