address greptile review feedback (greploop iteration 1)

- Fix XSS: add _html_safe_json() to escape <, >, &, ' for safe script embedding
- Remove duplicate /ui/mcp/oauth/callback route from discoverable_endpoints
- Add test_mcp_oauth_callback_html_xss_safe; update callback tests for _build_mcp_oauth_callback_html

Made-with: Cursor
This commit is contained in:
Sameer Kankute 2026-03-09 16:05:31 +05:30
parent 6f49645eb5
commit 9568d5de4f
2 changed files with 47 additions and 33 deletions

View file

@ -420,6 +420,20 @@ async def callback(code: str, state: str):
)
def _html_safe_json(value: object) -> str:
"""
JSON-serialize a value for safe embedding inside HTML <script> tags.
Escapes characters that could break out of the script block (XSS).
"""
s = json.dumps(value)
return (
s.replace("<", "\\u003c")
.replace(">", "\\u003e")
.replace("&", "\\u0026")
.replace("'", "\\u0027")
)
def _build_mcp_oauth_callback_html(code: Optional[str], state: Optional[str]) -> str:
"""
Build the inline HTML page for the MCP OAuth callback.
@ -432,8 +446,8 @@ def _build_mcp_oauth_callback_html(code: Optional[str], state: Optional[str]) ->
Served as a backend route so it works even when the static-files
mount for /ui is unavailable (e.g. read-only container filesystems).
"""
code_json = json.dumps(code)
state_json = json.dumps(state)
code_json = _html_safe_json(code)
state_json = _html_safe_json(state)
return f"""<!DOCTYPE html>
<html lang="en">
<head>
@ -506,25 +520,8 @@ def _build_mcp_oauth_callback_html(code: Optional[str], state: Optional[str]) ->
</html>"""
@router.get("/ui/mcp/oauth/callback", include_in_schema=False)
async def mcp_oauth_ui_callback(
code: Optional[str] = None,
state: Optional[str] = None,
) -> HTMLResponse:
"""
OAuth callback landing page for the MCP UI flow.
After the external OAuth provider redirects the browser to /callback,
LiteLLM decodes the state and performs a second redirect to this page
(the original redirect_uri supplied by the UI).
This backend route serves an inline HTML page that stores the OAuth
result in browser storage and redirects the user back to the LiteLLM
dashboard. It acts as a reliable fallback for environments where the
static-file mount for /ui is unavailable (read-only container
filesystems, Kubernetes deployments, etc.).
"""
return HTMLResponse(content=_build_mcp_oauth_callback_html(code, state))
# The /ui/mcp/oauth/callback route is registered in proxy_server.py (before
# the StaticFiles mount) to ensure it takes precedence over static files.
# ------------------------------

View file

@ -1131,8 +1131,7 @@ def test_get_request_base_url_comprehensive(
# ---------------------------------------------------------------------------
@pytest.mark.asyncio
async def test_mcp_oauth_ui_callback_returns_html():
def test_mcp_oauth_ui_callback_returns_html():
"""
The /ui/mcp/oauth/callback endpoint must return an HTML page regardless
of whether the static files mount is available. This is the fallback
@ -1140,18 +1139,16 @@ async def test_mcp_oauth_ui_callback_returns_html():
"""
try:
from litellm.proxy._experimental.mcp_server.discoverable_endpoints import (
mcp_oauth_ui_callback,
_build_mcp_oauth_callback_html,
)
except ImportError:
pytest.skip("MCP discoverable endpoints not available")
response = await mcp_oauth_ui_callback(
body = _build_mcp_oauth_callback_html(
code="test_auth_code_123",
state="test_state_456",
)
assert response.status_code == 200
body = response.body.decode()
assert "<!DOCTYPE html>" in body
# Payload values are embedded as JSON literals in the JS
assert "test_auth_code_123" in body
@ -1161,24 +1158,44 @@ async def test_mcp_oauth_ui_callback_returns_html():
assert "litellm-mcp-oauth-return-url" in body
@pytest.mark.asyncio
async def test_mcp_oauth_ui_callback_handles_missing_params():
def test_mcp_oauth_ui_callback_handles_missing_params():
"""
The /ui/mcp/oauth/callback endpoint must handle missing code/state
gracefully (e.g. direct navigation or error redirect from provider).
"""
try:
from litellm.proxy._experimental.mcp_server.discoverable_endpoints import (
mcp_oauth_ui_callback,
_build_mcp_oauth_callback_html,
)
except ImportError:
pytest.skip("MCP discoverable endpoints not available")
response = await mcp_oauth_ui_callback(code=None, state=None)
body = _build_mcp_oauth_callback_html(code=None, state=None)
assert response.status_code == 200
body = response.body.decode()
assert "<!DOCTYPE html>" in body
# null values should be safely embedded as JSON null
assert "null" in body
def test_mcp_oauth_callback_html_xss_safe():
"""
OAuth code/state must be escaped for safe embedding in <script> tags.
"""
try:
from litellm.proxy._experimental.mcp_server.discoverable_endpoints import (
_build_mcp_oauth_callback_html,
)
except ImportError:
pytest.skip("MCP discoverable endpoints not available")
# Malicious payload that could break out of script tag
body = _build_mcp_oauth_callback_html(
code='</script><img src=x onerror=alert(1)>',
state="normal_state",
)
# Payload must be escaped: < -> \u003c, > -> \u003e (safe for <script> embedding)
# The code value in the JSON should show \\u003c/script\\u003e, not raw </script>
assert '\\u003c/script\\u003e' in body
assert '\\u003cimg' in body