mirror of
https://github.com/langgenius/dify.git
synced 2026-07-15 01:17:04 +08:00
f903b980ea
Validate user_code against a [A-Z0-9-] charset before placing it in the /device redirect, closing the CodeQL url-redirection finding (the value can no longer carry path/scheme separators). Split the compound assert flagged by ruff PT018.
175 lines
5.9 KiB
Python
175 lines
5.9 KiB
Python
"""SSO-branch device-flow endpoints under /openapi/v1/oauth/device/."""
|
|
|
|
import builtins
|
|
from unittest.mock import MagicMock, patch
|
|
|
|
import pytest
|
|
from flask import Flask
|
|
from flask.views import MethodView
|
|
|
|
from controllers.openapi import bp as openapi_bp
|
|
from controllers.openapi.oauth_device_sso import (
|
|
approval_context,
|
|
approve_external,
|
|
sso_complete,
|
|
sso_initiate,
|
|
)
|
|
|
|
if not hasattr(builtins, "MethodView"):
|
|
builtins.MethodView = MethodView # type: ignore[attr-defined]
|
|
|
|
|
|
@pytest.fixture
|
|
def openapi_app() -> Flask:
|
|
app = Flask(__name__)
|
|
app.config["TESTING"] = True
|
|
app.register_blueprint(openapi_bp)
|
|
return app
|
|
|
|
|
|
def _rule(app: Flask, path: str):
|
|
return next(r for r in app.url_map.iter_rules() if r.rule == path)
|
|
|
|
|
|
def test_sso_initiate_registered(openapi_app: Flask):
|
|
rules = {r.rule for r in openapi_app.url_map.iter_rules()}
|
|
assert "/openapi/v1/oauth/device/sso-initiate" in rules
|
|
|
|
|
|
def test_sso_complete_registered(openapi_app: Flask):
|
|
rules = {r.rule for r in openapi_app.url_map.iter_rules()}
|
|
assert "/openapi/v1/oauth/device/sso-complete" in rules
|
|
|
|
|
|
def test_approval_context_registered(openapi_app: Flask):
|
|
rules = {r.rule for r in openapi_app.url_map.iter_rules()}
|
|
assert "/openapi/v1/oauth/device/approval-context" in rules
|
|
|
|
|
|
def test_approve_external_registered(openapi_app: Flask):
|
|
rules = {r.rule for r in openapi_app.url_map.iter_rules()}
|
|
assert "/openapi/v1/oauth/device/approve-external" in rules
|
|
|
|
|
|
def test_sso_initiate_dispatches_to_function(openapi_app: Flask):
|
|
rule = _rule(openapi_app, "/openapi/v1/oauth/device/sso-initiate")
|
|
assert openapi_app.view_functions[rule.endpoint] is sso_initiate
|
|
|
|
|
|
def test_sso_complete_dispatches_to_function(openapi_app: Flask):
|
|
rule = _rule(openapi_app, "/openapi/v1/oauth/device/sso-complete")
|
|
assert openapi_app.view_functions[rule.endpoint] is sso_complete
|
|
|
|
|
|
def test_approval_context_dispatches_to_function(openapi_app: Flask):
|
|
rule = _rule(openapi_app, "/openapi/v1/oauth/device/approval-context")
|
|
assert openapi_app.view_functions[rule.endpoint] is approval_context
|
|
|
|
|
|
def test_approve_external_dispatches_to_function(openapi_app: Flask):
|
|
rule = _rule(openapi_app, "/openapi/v1/oauth/device/approve-external")
|
|
assert openapi_app.view_functions[rule.endpoint] is approve_external
|
|
|
|
|
|
def test_sso_complete_idp_callback_url_uses_canonical_path():
|
|
"""sso_initiate hardcodes the IdP callback URL — must point at the
|
|
canonical /openapi/v1/ path so IdP-side ACS configuration matches.
|
|
"""
|
|
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"]
|