feat(login): Implement secure post-login redirection

Return users to their original requested page after login, with robust same-origin validation to prevent open redirects. Co-Author Agent Hero <agent-hero@neurocis.ai>

neurocis committed Jun 2, 2026 at 21:54 UTC 6609632bf6c94d53768a7d8f82bf9b8a8810c851
5 files changed +109 -10
helpers/api.py
+32 -1
@@ -1,6 +1,7 @@
1 from abc import abstractmethod
2 import json
3 import threading
4 +from urllib.parse import urlsplit
5 from functools import wraps
6 from pathlib import Path
7 from typing import Union, Dict, Any
@@ -102,6 +103,36 @@ class ApiHandler:
103 from helpers.network import is_loopback_address
104
105
106 +def is_safe_next_url(value: str | None) -> bool:
107 + """Return True when value is a safe same-origin redirect target."""
108 + if not value:
109 + return False
110 + if "\r" in value or "\n" in value:
111 + return False
112 +
113 + parsed = urlsplit(value)
114 + if parsed.scheme or parsed.netloc:
115 + return False
116 +
117 + # Require an absolute path within this origin, but reject protocol-relative URLs.
118 + return parsed.path.startswith("/") and not parsed.path.startswith("//")
119 +
120 +
121 +def get_safe_next_url(value: str | None, fallback: str | None = None) -> str | None:
122 + """Return value if it is a safe next URL, otherwise return a safe fallback."""
123 + if is_safe_next_url(value):
124 + return value
125 + if is_safe_next_url(fallback):
126 + return fallback
127 + return None
128 +
129 +
130 +def get_current_request_next_url() -> str:
131 + """Return the current request path/query as a safe relative redirect target."""
132 + next_url = request.full_path if request.query_string else request.path
133 + return get_safe_next_url(next_url, url_for("serve_index")) or url_for("serve_index")
134 +
135 +
136 def requires_api_key(f):
137 @wraps(f)
138 async def decorated(*args, **kwargs):
@@ -142,7 +173,7 @@ def requires_auth(f):
173 if not user_pass_hash:
174 return await f(*args, **kwargs)
175 if session.get("authentication") != user_pass_hash:
145 - return redirect(url_for("login_handler"))
176 + return redirect(url_for("login_handler", next=get_current_request_next_url()))
177 return await f(*args, **kwargs)
178
179 return decorated
helpers/ui_server.py
+9 -3
@@ -26,7 +26,7 @@ from werkzeug.wrappers.request import Request as WerkzeugRequest
26 import socketio # type: ignore[import-untyped]
27
28 from helpers import dotenv, fasta2a_server, files, git, login, mcp_server, runtime
29 -from helpers.api import register_api_route, requires_auth
29 +from helpers.api import get_safe_next_url, register_api_route, requires_auth
30 from helpers.extension import extensible
31 from helpers.files import get_abs_path
32 from helpers.print_style import PrintStyle
@@ -200,19 +200,25 @@ class UiRouteHandlers:
200 @extensible
201 async def login_handler(self):
202 error = None
203 + fallback_url = url_for("serve_index")
204 + next_url = get_safe_next_url(
205 + request.form.get("next") if request.method == "POST" else request.args.get("next"),
206 + fallback_url,
207 + )
208 +
209 if request.method == "POST":
210 user = dotenv.get_dotenv_value("AUTH_LOGIN")
211 password = dotenv.get_dotenv_value("AUTH_PASSWORD")
212
213 if request.form["username"] == user and request.form["password"] == password:
214 session["authentication"] = login.get_credentials_hash()
209 - return redirect(url_for("serve_index"))
215 + return redirect(next_url or fallback_url)
216 else:
217 await asyncio.sleep(1)
218 error = "Invalid Credentials. Please try again."
219
220 login_page_content = files.read_file("webui/login.html")
215 - return render_template_string(login_page_content, error=error)
221 + return render_template_string(login_page_content, error=error, next=next_url)
222
223 @extensible
224 async def logout_handler(self):
tests/test_http_auth_csrf.py
+49
@@ -122,3 +122,52 @@ def test_http_csrf_accepts_valid_cookie(monkeypatch) -> None:
122 _set_csrf_cookie(client, "csrf-4")
123 response = client.get("/secure")
124 assert response.status_code == 200
125 +
126 +
127 +def test_safe_next_url_accepts_plugin_page_path() -> None:
128 + from helpers.api import get_safe_next_url, is_safe_next_url
129 +
130 + target = "/plugins/a0_voqualizer/webui/voqualizer.html"
131 + assert is_safe_next_url(target)
132 + assert get_safe_next_url(target, "/") == target
133 +
134 +
135 +def test_safe_next_url_preserves_query_string() -> None:
136 + from helpers.api import get_safe_next_url
137 +
138 + target = "/plugins/a0_voqualizer/webui/voqualizer.html?context=rlO1iMV7"
139 + assert get_safe_next_url(target, "/") == target
140 +
141 +
142 +def test_safe_next_url_rejects_external_and_protocol_relative_urls() -> None:
143 + from helpers.api import get_safe_next_url, is_safe_next_url
144 +
145 + fallback = "/"
146 + for value in [
147 + "https://evil.example/plugins/a0_voqualizer/webui/voqualizer.html",
148 + "//evil.example/plugins/a0_voqualizer/webui/voqualizer.html",
149 + "javascript:alert(1)",
150 + "/safe\nLocation: https://evil.example",
151 + ]:
152 + assert not is_safe_next_url(value)
153 + assert get_safe_next_url(value, fallback) == fallback
154 +
155 +
156 +def test_auth_redirect_includes_original_path_and_query(monkeypatch) -> None:
157 + from run_ui import requires_auth
158 +
159 + monkeypatch.setattr("helpers.login.get_credentials_hash", lambda: "hash")
160 +
161 + app = _make_app()
162 +
163 + @app.get("/plugins/a0_voqualizer/webui/voqualizer.html")
164 + @requires_auth
165 + async def voqualizer_page():
166 + return Response("ok", status=200)
167 +
168 + client = app.test_client()
169 + response = client.get("/plugins/a0_voqualizer/webui/voqualizer.html?context=rlO1iMV7")
170 + assert response.status_code == 302
171 + location = response.headers["Location"]
172 + assert location.startswith("/login?next=")
173 + assert "%2Fplugins%2Fa0_voqualizer%2Fwebui%2Fvoqualizer.html%3Fcontext%3DrlO1iMV7" in location
webui/js/api.js
+15 -5
@@ -266,10 +266,20 @@ function _normalizeApiUrl(url) {
266 }
267
268 function redirect(response) {
269 - if (!(response.redirected && response.url.endsWith("/login"))) return false;
269 + if (!response.redirected) return false;
270 +
271 const _redirectUrl = new URL(response.url);
271 - if (_redirectUrl.origin === window.location.origin) {
272 - window.location.href = response.url;
272 + if (
273 + _redirectUrl.origin === window.location.origin &&
274 + _redirectUrl.pathname === "/login"
275 + ) {
276 + const currentUrl = `${window.location.pathname}${window.location.search}${window.location.hash}`;
277 + if (currentUrl && currentUrl !== "/login") {
278 + _redirectUrl.searchParams.set("next", currentUrl);
279 + }
280 + window.location.href = _redirectUrl.toString();
281 + return true;
282 }
274 - return true;
275 -}
\ No newline at end of file
283 +
284 + return false;
285 +}
webui/login.html
+4 -1
@@ -9,7 +9,10 @@
9 </head>
10 <body>
11 <div class="login-container">
12 - <form class="login-form" method="POST" action="/login">
12 + <form class="login-form" method="POST" action="/login{% if next %}?next={{ next|urlencode }}{% endif %}">
13 + {% if next %}
14 + <input type="hidden" name="next" value="{{ next }}">
15 + {% endif %}
16 <img src="/public/splash.jpg" alt="Agent Zero Logo" class="logo">
17 <h2>Agent Zero</h2>
18 <div class="input-group">