mirror of
https://github.com/langgenius/dify.git
synced 2026-07-14 08:57:02 +08:00
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:
@ -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(
|
||||
|
||||
@ -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"]
|
||||
|
||||
@ -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():
|
||||
|
||||
Reference in New Issue
Block a user