This is an automated email from the ASF dual-hosted git repository.
vincbeck pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new 7cae03be33f Speed up API server requests by keeping JWT checks off the
event loop (#74856)
7cae03be33f is described below
commit 7cae03be33f3d3162926f314db32bfe2c05f9a4b
Author: safaehar <[email protected]>
AuthorDate: Mon Oct 19 15:59:10 2026 +0200
Speed up API server requests by keeping JWT checks off the event loop
(#74856)
Static UI assets, the health check and the version endpoint never read the
user, yet every one of these requests carrying the session cookie decoded
the
JWT and queried the revoked token table. The middleware also added a
BaseHTTPMiddleware task and body stream to every request, including the
large
streamed UI bundle, so a cold UI load serialized on each API server worker.
Auth managers may block in refresh_user (for example when calling an
identity
provider), which stalled every concurrent request on the worker.
---
.../api_fastapi/auth/managers/base_auth_manager.py | 4 +-
.../api_fastapi/auth/middlewares/refresh_token.py | 154 +++++---
.../src/airflow/api_fastapi/core_api/app.py | 6 +-
.../auth/managers/test_base_auth_manager.py | 28 ++
.../auth/middlewares/test_refresh_token.py | 387 +++++++++++----------
airflow-core/tests/unit/api_fastapi/test_app.py | 10 +
6 files changed, 348 insertions(+), 241 deletions(-)
diff --git
a/airflow-core/src/airflow/api_fastapi/auth/managers/base_auth_manager.py
b/airflow-core/src/airflow/api_fastapi/auth/managers/base_auth_manager.py
index 77c0ef9a141..94b96ba5af2 100644
--- a/airflow-core/src/airflow/api_fastapi/auth/managers/base_auth_manager.py
+++ b/airflow-core/src/airflow/api_fastapi/auth/managers/base_auth_manager.py
@@ -29,6 +29,7 @@ from typing import TYPE_CHECKING, Any, Generic, Literal,
TypeVar
from jwt import InvalidTokenError
from sqlalchemy import select
+from starlette.concurrency import run_in_threadpool
from airflow.api_fastapi.auth.managers.models.base_user import BaseUser
from airflow.api_fastapi.auth.managers.models.resource_details import (
@@ -162,7 +163,8 @@ class BaseAuthManager(Generic[T], LoggingMixin,
metaclass=ABCMeta):
log.error("JWT token is not valid: %s", e)
raise e
- if (jti := payload.get("jti")) and RevokedToken.is_revoked(jti):
+ # A synchronous DB lookup, kept off the event loop
+ if (jti := payload.get("jti")) and await
run_in_threadpool(RevokedToken.is_revoked, jti):
raise InvalidTokenError("Token has been revoked")
try:
diff --git
a/airflow-core/src/airflow/api_fastapi/auth/middlewares/refresh_token.py
b/airflow-core/src/airflow/api_fastapi/auth/middlewares/refresh_token.py
index 7c9bfdcb10a..2e97e6c25c1 100644
--- a/airflow-core/src/airflow/api_fastapi/auth/middlewares/refresh_token.py
+++ b/airflow-core/src/airflow/api_fastapi/auth/middlewares/refresh_token.py
@@ -17,9 +17,12 @@
# under the License.
from __future__ import annotations
-from fastapi import HTTPException, Request
+from typing import TYPE_CHECKING
+
+from fastapi import HTTPException, Request, Response
from fastapi.responses import JSONResponse
-from starlette.middleware.base import BaseHTTPMiddleware
+from starlette.concurrency import run_in_threadpool
+from starlette.datastructures import MutableHeaders
from airflow.api_fastapi.app import get_auth_manager, get_cookie_path
from airflow.api_fastapi.auth.managers.base_auth_manager import
COOKIE_NAME_JWT_TOKEN
@@ -31,8 +34,15 @@ from airflow.api_fastapi.core_api.security import (
)
from airflow.configuration import conf
+if TYPE_CHECKING:
+ from starlette.types import ASGIApp, Message, Receive, Scope, Send
+
+# Unauthenticated endpoints that never read the user, so resolving it there
would only cost a
+# JWT decode and a revocation lookup on every health check and UI boot.
+UNAUTHENTICATED_PATHS = frozenset(["/api/v2/monitor/health",
"/api/v2/version"])
-class JWTRefreshMiddleware(BaseHTTPMiddleware):
+
+class JWTRefreshMiddleware:
"""
Middleware to handle JWT token refresh.
@@ -41,62 +51,100 @@ class JWTRefreshMiddleware(BaseHTTPMiddleware):
2. Calls ``refresh_user`` method from auth manager with the user
3. If ``refresh_user`` returns a user, generate a JWT token based upon
this user and send it in the
response as cookie
+
+ Requests to ``UNAUTHENTICATED_PATHS`` or under ``skip_path_prefixes``
(paths relative to the
+ root path) are passed through untouched.
"""
- async def dispatch(self, request: Request, call_next):
- new_token = None
+ def __init__(self, app: ASGIApp, skip_path_prefixes: tuple[str, ...] = ())
-> None:
+ self.app = app
+ self.skip_path_prefixes = skip_path_prefixes
+
+ async def __call__(self, scope: Scope, receive: Receive, send: Send) ->
None:
+ if scope["type"] != "http" or self._is_skipped(scope):
+ await self.app(scope, receive, send)
+ return
+
+ request = Request(scope)
current_token = request.cookies.get(COOKIE_NAME_JWT_TOKEN)
+ if current_token is None:
+ await self.app(scope, receive, send)
+ return
+
+ new_token = None
+ try:
+ new_user, current_user = await self._refresh_user(current_token)
+ if user := (new_user or current_user):
+ # Stamp the trust sentinel alongside the user so `get_user()`
+ # can distinguish this trusted assignment from a stray write
+ # by unrelated middleware.
+ request.state.user = user
+ request.state.user_authenticated_via =
USER_INJECTED_BY_TRUSTED_MIDDLEWARE
+ if new_user:
+ # If we created a new user, serialize it and set it as a cookie
+ new_token = get_auth_manager().generate_jwt(new_user)
+ except (HTTPException, AuthManagerRefreshTokenExpiredException):
+ # Receive a HTTPException when the Airflow token is expired
+ # Receive a AuthManagerRefreshTokenExpiredException when the
potential underlying refresh
+ # token used by the auth manager is expired
+ new_token = ""
+
+ cookie_headers = self._build_cookie_headers(request, new_token) if
new_token is not None else []
+ response_started = False
+
+ async def send_with_cookies(message: Message) -> None:
+ nonlocal response_started
+ if message["type"] == "http.response.start":
+ response_started = True
+ headers = MutableHeaders(scope=message)
+ for value in cookie_headers:
+ headers.append("set-cookie", value)
+ await send(message)
+
try:
- if current_token is not None:
- try:
- new_user, current_user = await
self._refresh_user(current_token)
- if user := (new_user or current_user):
- # Stamp the trust sentinel alongside the user so
`get_user()`
- # can distinguish this trusted assignment from a stray
write
- # by unrelated middleware.
- request.state.user = user
- request.state.user_authenticated_via =
USER_INJECTED_BY_TRUSTED_MIDDLEWARE
- if new_user:
- # If we created a new user, serialize it and set it as
a cookie
- new_token = get_auth_manager().generate_jwt(new_user)
- except (HTTPException,
AuthManagerRefreshTokenExpiredException):
- # Receive a HTTPException when the Airflow token is expired
- # Receive a AuthManagerRefreshTokenExpiredException when
the potential underlying refresh
- # token used by the auth manager is expired
- new_token = ""
-
- response = await call_next(request)
-
- if new_token is not None:
- cookie_path = get_cookie_path()
- secure = request.base_url.scheme == "https" or
bool(conf.get("api", "ssl_cert", fallback=""))
- response.set_cookie(
- COOKIE_NAME_JWT_TOKEN,
- new_token,
- path=cookie_path,
- httponly=True,
- secure=secure,
- samesite="lax",
- max_age=0 if new_token == "" else None,
- )
- # Clear any stale _token cookie at root path "/".
- # Older Airflow instances may have set the cookie there;
- # without this, the root-path cookie keeps being sent on
- # every request, causing an infinite redirect loop.
- if cookie_path != "/":
- response.delete_cookie(
- key=COOKIE_NAME_JWT_TOKEN,
- path="/",
- httponly=True,
- secure=secure,
- samesite="lax",
- )
+ await self.app(scope, receive, send_with_cookies)
except HTTPException as exc:
- # If any HTTPException is raised during user resolution or
refresh, return it as response
- return JSONResponse(status_code=exc.status_code,
content={"detail": exc.detail})
- return response
+ if response_started:
+ raise
+ # An HTTPException escaping the inner middlewares is returned as a
response
+ response = JSONResponse(status_code=exc.status_code,
content={"detail": exc.detail})
+ await response(scope, receive, send)
+
+ def _is_skipped(self, scope: Scope) -> bool:
+ path = scope["path"].removeprefix(scope.get("root_path", ""))
+ return path in UNAUTHENTICATED_PATHS or
path.startswith(self.skip_path_prefixes)
+
+ @staticmethod
+ def _build_cookie_headers(request: Request, new_token: str) -> list[str]:
+ cookie_path = get_cookie_path()
+ secure = request.base_url.scheme == "https" or bool(conf.get("api",
"ssl_cert", fallback=""))
+ # Only used to render the Set-Cookie header values
+ response = Response()
+ response.set_cookie(
+ COOKIE_NAME_JWT_TOKEN,
+ new_token,
+ path=cookie_path,
+ httponly=True,
+ secure=secure,
+ samesite="lax",
+ max_age=0 if new_token == "" else None,
+ )
+ # Clear any stale _token cookie at root path "/".
+ # Older Airflow instances may have set the cookie there;
+ # without this, the root-path cookie keeps being sent on
+ # every request, causing an infinite redirect loop.
+ if cookie_path != "/":
+ response.delete_cookie(
+ key=COOKIE_NAME_JWT_TOKEN,
+ path="/",
+ httponly=True,
+ secure=secure,
+ samesite="lax",
+ )
+ return [value.decode("latin-1") for key, value in response.raw_headers
if key == b"set-cookie"]
@staticmethod
async def _refresh_user(current_token: str) -> tuple[BaseUser | None,
BaseUser | None]:
user = await resolve_user_from_token(current_token)
- return get_auth_manager().refresh_user(user=user), user
+ # Auth managers may do blocking I/O here (e.g. calling an identity
provider)
+ return await run_in_threadpool(get_auth_manager().refresh_user,
user=user), user
diff --git a/airflow-core/src/airflow/api_fastapi/core_api/app.py
b/airflow-core/src/airflow/api_fastapi/core_api/app.py
index d3a66a5b68d..82dfa727ad2 100644
--- a/airflow-core/src/airflow/api_fastapi/core_api/app.py
+++ b/airflow-core/src/airflow/api_fastapi/core_api/app.py
@@ -38,6 +38,8 @@ log = logging.getLogger(__name__)
_AIRFLOW_PATH = Path(__file__).parents[3]
+_STATIC_URL_PATH = "/static"
+
# Language codes ("en", "zh-CN") and namespaces ("common") of the UI
translation files. Anchored
# so a path parameter cannot smuggle in path separators or ``..`` and escape
the locales directory.
_I18N_SEGMENT = re.compile(r"^[A-Za-z0-9_-]+$")
@@ -168,7 +170,7 @@ def init_views(app: FastAPI) -> None:
)
app.mount(
- "/static",
+ _STATIC_URL_PATH,
StaticFiles(
directory=directory,
html=True,
@@ -282,7 +284,7 @@ def init_middlewares(app: FastAPI) -> None:
from airflow.api_fastapi.app import get_auth_manager
from airflow.api_fastapi.auth.middlewares.refresh_token import
JWTRefreshMiddleware
- app.add_middleware(JWTRefreshMiddleware)
+ app.add_middleware(JWTRefreshMiddleware,
skip_path_prefixes=(f"{_STATIC_URL_PATH}/",))
for middleware_cls, middleware_kwargs in
get_auth_manager().get_fastapi_middlewares():
app.add_middleware(middleware_cls, **middleware_kwargs)
diff --git
a/airflow-core/tests/unit/api_fastapi/auth/managers/test_base_auth_manager.py
b/airflow-core/tests/unit/api_fastapi/auth/managers/test_base_auth_manager.py
index 49e3699f203..18a9ceeecd2 100644
---
a/airflow-core/tests/unit/api_fastapi/auth/managers/test_base_auth_manager.py
+++
b/airflow-core/tests/unit/api_fastapi/auth/managers/test_base_auth_manager.py
@@ -16,6 +16,7 @@
# under the License.
from __future__ import annotations
+import threading
import warnings
from typing import TYPE_CHECKING, Any
from unittest.mock import AsyncMock, MagicMock, Mock, create_autospec, patch
@@ -335,6 +336,33 @@ class TestBaseAuthManager:
mock_is_revoked.assert_called_once_with("some-jti")
+ @patch(
+
"airflow.api_fastapi.auth.managers.base_auth_manager.BaseAuthManager._get_token_validator",
+ autospec=True,
+ )
+ @patch.object(EmptyAuthManager, "deserialize_user")
+ @pytest.mark.asyncio
+ async def test_get_user_from_token_checks_revocation_off_the_event_loop(
+ self, mock_deserialize_user, mock__get_token_validator, auth_manager
+ ):
+ """The revocation lookup is a synchronous DB query, so it must not
block the event loop."""
+ signer = AsyncMock(spec=JWTValidator)
+ signer.avalidated_claims.return_value = {"jti": "some-jti"}
+ mock__get_token_validator.return_value = signer
+ lookup_threads = []
+
+ def is_revoked(jti):
+ lookup_threads.append(threading.get_ident())
+ return False
+
+ with patch(
+ "airflow.models.revoked_token.RevokedToken.is_revoked",
autospec=True, side_effect=is_revoked
+ ):
+ await auth_manager.get_user_from_token("token")
+
+ assert lookup_threads
+ assert lookup_threads != [threading.get_ident()]
+
@patch(
"airflow.api_fastapi.auth.managers.base_auth_manager.BaseAuthManager._get_token_validator",
autospec=True,
diff --git
a/airflow-core/tests/unit/api_fastapi/auth/middlewares/test_refresh_token.py
b/airflow-core/tests/unit/api_fastapi/auth/middlewares/test_refresh_token.py
index 9114ce0d938..346c1e8dc61 100644
--- a/airflow-core/tests/unit/api_fastapi/auth/middlewares/test_refresh_token.py
+++ b/airflow-core/tests/unit/api_fastapi/auth/middlewares/test_refresh_token.py
@@ -17,125 +17,199 @@
# under the License.
from __future__ import annotations
-from unittest.mock import AsyncMock, MagicMock, patch
+import threading
+from unittest.mock import MagicMock, patch
import pytest
-from fastapi import HTTPException, Request, Response
+from fastapi import FastAPI, HTTPException, Request
+from fastapi.responses import PlainTextResponse, StreamingResponse
+from fastapi.staticfiles import StaticFiles
+from fastapi.testclient import TestClient
-from airflow.api_fastapi.auth.managers.base_auth_manager import
COOKIE_NAME_JWT_TOKEN
+from airflow.api_fastapi.auth.managers.base_auth_manager import
COOKIE_NAME_JWT_TOKEN, BaseAuthManager
from airflow.api_fastapi.auth.managers.models.base_user import BaseUser
from airflow.api_fastapi.auth.middlewares.refresh_token import
JWTRefreshMiddleware
+from airflow.api_fastapi.core_api.security import
USER_INJECTED_BY_TRUSTED_MIDDLEWARE
+
+MODULE = "airflow.api_fastapi.auth.middlewares.refresh_token"
+
+
+def _build_app(static_dir=None, root_path: str = "") -> tuple[FastAPI, dict]:
+ """Return an app wrapped by the middleware and a dict capturing what the
endpoint saw."""
+ seen: dict = {}
+ app = FastAPI(root_path=root_path)
+
+ @app.get("/api/v2/dags")
+ def dags(request: Request):
+ seen["user"] = getattr(request.state, "user", None)
+ seen["authenticated_via"] = getattr(request.state,
"user_authenticated_via", None)
+ return PlainTextResponse("ok")
+
+ @app.get("/api/v2/monitor/health")
+ def health():
+ return PlainTextResponse("healthy")
+
+ @app.get("/api/v2/version")
+ def version():
+ return PlainTextResponse("3.x")
+
+ @app.get("/api/v2/stream")
+ def stream():
+ return StreamingResponse(iter([b"chunk1-", b"chunk2-", b"chunk3"]),
media_type="text/plain")
+
+ if static_dir is not None:
+ app.mount("/static", StaticFiles(directory=static_dir),
name="webapp_static_folder")
+
+ app.add_middleware(JWTRefreshMiddleware, skip_path_prefixes=("/static/",))
+ return app, seen
+
+
[email protected]
+def mock_auth_manager():
+ with patch(f"{MODULE}.get_auth_manager", autospec=True) as
mock_get_auth_manager:
+ auth_manager = MagicMock(spec=BaseAuthManager)
+ auth_manager.refresh_user.return_value = None
+ mock_get_auth_manager.return_value = auth_manager
+ yield auth_manager
+
+
[email protected]
+def mock_resolve_user():
+ with patch(f"{MODULE}.resolve_user_from_token", autospec=True) as
mock_resolve:
+ mock_resolve.return_value = MagicMock(spec=BaseUser)
+ yield mock_resolve
+
+
[email protected](autouse=True)
+def no_ssl_cert():
+ with patch(f"{MODULE}.conf", autospec=True) as mock_conf:
+ mock_conf.get.return_value = ""
+ yield mock_conf
class TestJWTRefreshMiddleware:
- @pytest.fixture
- def middleware(self):
- return JWTRefreshMiddleware(app=MagicMock())
-
- @pytest.fixture
- def mock_request(self):
- request = MagicMock(spec=Request)
- request.cookies = {}
- request.state = MagicMock()
- return request
-
- @pytest.fixture
- def mock_user(self):
- return MagicMock(spec=BaseUser)
-
- @patch.object(JWTRefreshMiddleware, "_refresh_user")
- @pytest.mark.asyncio
- async def test_dispatch_no_token(self, mock_refresh_user, middleware,
mock_request):
- call_next = AsyncMock(return_value=Response())
-
- await middleware.dispatch(mock_request, call_next)
-
- call_next.assert_called_once_with(mock_request)
- mock_refresh_user.assert_not_called()
-
- @patch.object(
- JWTRefreshMiddleware,
- "_refresh_user",
- side_effect=HTTPException(status_code=403, detail="Invalid JWT token"),
- )
- @pytest.mark.asyncio
- async def test_dispatch_invalid_token(self, mock_refresh_user, middleware,
mock_request):
- mock_request.cookies = {COOKIE_NAME_JWT_TOKEN: "valid_token"}
- call_next = AsyncMock(return_value=Response(status_code=401))
-
- response = await middleware.dispatch(mock_request, call_next)
- assert response.status_code == 401
- assert '_token=""; HttpOnly; Max-Age=0; Path=/; SameSite=lax' in
response.headers.get("set-cookie")
-
-
@patch("airflow.api_fastapi.auth.middlewares.refresh_token.get_auth_manager")
-
@patch("airflow.api_fastapi.auth.middlewares.refresh_token.resolve_user_from_token")
- @pytest.mark.asyncio
- async def test_dispatch_no_refreshed_token(
- self, mock_resolve_user_from_token, mock_get_auth_manager, middleware,
mock_request, mock_user
- ):
- mock_request.cookies = {COOKIE_NAME_JWT_TOKEN: "valid_token"}
- mock_resolve_user_from_token.return_value = mock_user
- mock_auth_manager = MagicMock()
- mock_get_auth_manager.return_value = mock_auth_manager
- mock_auth_manager.refresh_user.return_value = None
+ def test_no_token(self, mock_resolve_user, mock_auth_manager):
+ app, _ = _build_app()
- call_next = AsyncMock(return_value=Response())
- await middleware.dispatch(mock_request, call_next)
+ response = TestClient(app).get("/api/v2/dags")
- call_next.assert_called_once_with(mock_request)
- mock_resolve_user_from_token.assert_called_once_with("valid_token")
+ assert response.status_code == 200
+ assert "set-cookie" not in response.headers
+ mock_resolve_user.assert_not_called()
+
+ def test_invalid_token_clears_cookie(self, mock_resolve_user,
mock_auth_manager):
+ mock_resolve_user.side_effect = HTTPException(status_code=403,
detail="Invalid JWT token")
+ app, seen = _build_app()
+
+ response = TestClient(app, cookies={COOKIE_NAME_JWT_TOKEN:
"invalid_token"}).get("/api/v2/dags")
+
+ assert response.status_code == 200
+ assert seen["user"] is None
+ assert '_token=""; HttpOnly; Max-Age=0; Path=/; SameSite=lax' in
response.headers["set-cookie"]
+ mock_auth_manager.refresh_user.assert_not_called()
+
+ def test_no_refreshed_user(self, mock_resolve_user, mock_auth_manager):
+ app, seen = _build_app()
+
+ response = TestClient(app, cookies={COOKIE_NAME_JWT_TOKEN:
"valid_token"}).get("/api/v2/dags")
+
+ assert "set-cookie" not in response.headers
+ assert seen["user"] is mock_resolve_user.return_value
+ assert seen["authenticated_via"] is USER_INJECTED_BY_TRUSTED_MIDDLEWARE
+ mock_resolve_user.assert_called_once_with("valid_token")
mock_auth_manager.generate_jwt.assert_not_called()
-
@patch("airflow.api_fastapi.auth.middlewares.refresh_token.resolve_user_from_token")
- @pytest.mark.asyncio
- async def test_dispatch_expired_token(self, mock_resolve_user_from_token,
middleware, mock_request):
- mock_request.cookies = {COOKIE_NAME_JWT_TOKEN: "invalid_token"}
- mock_resolve_user_from_token.side_effect =
HTTPException(status_code=403)
-
- call_next = AsyncMock(return_value=Response())
- await middleware.dispatch(mock_request, call_next)
-
- call_next.assert_called_once_with(mock_request)
- mock_resolve_user_from_token.assert_called_once_with("invalid_token")
-
- @pytest.mark.asyncio
-
@patch("airflow.api_fastapi.auth.middlewares.refresh_token.get_auth_manager")
-
@patch("airflow.api_fastapi.auth.middlewares.refresh_token.resolve_user_from_token")
- @patch("airflow.api_fastapi.auth.middlewares.refresh_token.conf")
- async def test_dispatch_with_refreshed_user(
- self,
- mock_conf,
- mock_resolve_user_from_token,
- mock_get_auth_manager,
- middleware,
- mock_request,
- mock_user,
- ):
+ def test_refreshed_user(self, mock_resolve_user, mock_auth_manager):
refreshed_user = MagicMock(spec=BaseUser)
- mock_request.cookies = {COOKIE_NAME_JWT_TOKEN: "valid_token"}
- mock_resolve_user_from_token.return_value = mock_user
- mock_auth_manager = MagicMock()
- mock_get_auth_manager.return_value = mock_auth_manager
mock_auth_manager.refresh_user.return_value = refreshed_user
mock_auth_manager.generate_jwt.return_value = "new_token"
- mock_conf.get.return_value = ""
-
- call_next = AsyncMock(return_value=Response())
- response = await middleware.dispatch(mock_request, call_next)
+ app, seen = _build_app()
- from airflow.api_fastapi.core_api.security import
USER_INJECTED_BY_TRUSTED_MIDDLEWARE
+ response = TestClient(app, cookies={COOKIE_NAME_JWT_TOKEN:
"valid_token"}).get("/api/v2/dags")
- assert mock_request.state.user == refreshed_user
# JWTRefreshMiddleware must stamp the trust sentinel so `get_user()`
accepts
# the cached user. Without the marker, the auth dependency falls
through
# to a fresh JWT validation, defeating the refresh.
- assert mock_request.state.user_authenticated_via is
USER_INJECTED_BY_TRUSTED_MIDDLEWARE
- call_next.assert_called_once_with(mock_request)
- mock_resolve_user_from_token.assert_called_once_with("valid_token")
- mock_auth_manager.refresh_user.assert_called_once_with(user=mock_user)
+ assert seen["user"] is refreshed_user
+ assert seen["authenticated_via"] is USER_INJECTED_BY_TRUSTED_MIDDLEWARE
+
mock_auth_manager.refresh_user.assert_called_once_with(user=mock_resolve_user.return_value)
mock_auth_manager.generate_jwt.assert_called_once_with(refreshed_user)
- set_cookie_headers = response.headers.get("set-cookie", "")
- assert f"{COOKIE_NAME_JWT_TOKEN}=new_token" in set_cookie_headers
+ assert response.headers["set-cookie"] == "_token=new_token; HttpOnly;
Path=/; SameSite=lax"
+
+ def test_refresh_user_runs_off_the_event_loop(self, mock_resolve_user,
mock_auth_manager):
+ """``refresh_user`` is synchronous and may do blocking I/O, so it must
not block the loop."""
+ refresh_threads = []
+ mock_auth_manager.refresh_user.side_effect = lambda user:
refresh_threads.append(
+ threading.get_ident()
+ )
+ loop_threads = []
+ mock_resolve_user.side_effect = lambda token:
loop_threads.append(threading.get_ident())
+ app, _ = _build_app()
+
+ TestClient(app, cookies={COOKIE_NAME_JWT_TOKEN:
"valid_token"}).get("/api/v2/dags")
+
+ assert len(refresh_threads) == 1
+ assert refresh_threads != loop_threads
+
+ def test_streaming_response_passthrough(self, mock_resolve_user,
mock_auth_manager):
+ mock_auth_manager.refresh_user.return_value = MagicMock(spec=BaseUser)
+ mock_auth_manager.generate_jwt.return_value = "new_token"
+ app, _ = _build_app()
+
+ response = TestClient(app, cookies={COOKIE_NAME_JWT_TOKEN:
"valid_token"}).get("/api/v2/stream")
+
+ assert response.content == b"chunk1-chunk2-chunk3"
+ assert response.headers["content-type"].startswith("text/plain")
+ assert "_token=new_token" in response.headers["set-cookie"]
+
+ def test_http_exception_from_inner_app_returned_as_response(self,
mock_resolve_user, mock_auth_manager):
+ async def raising_app(scope, receive, send):
+ raise HTTPException(status_code=418, detail="teapot")
+
+ app = JWTRefreshMiddleware(raising_app)
+
+ response = TestClient(app, cookies={COOKIE_NAME_JWT_TOKEN:
"valid_token"}).get("/api/v2/dags")
+
+ assert response.status_code == 418
+ assert response.json() == {"detail": "teapot"}
+
+ @pytest.mark.parametrize("root_path", ["", "/team-a"])
+ @pytest.mark.parametrize("path", ["/api/v2/monitor/health",
"/api/v2/version", "/static/main.js"])
+ def test_skipped_paths(self, mock_resolve_user, mock_auth_manager,
tmp_path, path, root_path):
+ """Static assets and unauthenticated endpoints never decode the JWT or
hit the DB."""
+ (tmp_path / "main.js").write_text("console.log('hi');")
+ app, _ = _build_app(static_dir=tmp_path, root_path=root_path)
+ client = TestClient(app, root_path=root_path,
cookies={COOKIE_NAME_JWT_TOKEN: "valid_token"})
+
+ response = client.get(f"{root_path}{path}")
+
+ assert response.status_code == 200
+ assert "set-cookie" not in response.headers
+ mock_resolve_user.assert_not_called()
+ mock_auth_manager.refresh_user.assert_not_called()
+
+ @pytest.mark.parametrize(
+ ("headers", "expected_status", "expected_body"),
+ [
+ pytest.param({}, 200, b"0123456789", id="full"),
+ pytest.param({"range": "bytes=2-4"}, 206, b"234", id="range"),
+ ],
+ )
+ def test_static_files_served_unchanged(
+ self, mock_resolve_user, tmp_path, headers, expected_status,
expected_body
+ ):
+ (tmp_path / "main.js").write_bytes(b"0123456789")
+ app, _ = _build_app(static_dir=tmp_path)
+ client = TestClient(app, cookies={COOKIE_NAME_JWT_TOKEN:
"valid_token"})
+
+ response = client.get("/static/main.js", headers=headers)
+
+ assert response.status_code == expected_status
+ assert response.content == expected_body
+ assert "etag" in response.headers
+ assert "last-modified" in response.headers
+ revalidated = client.get("/static/main.js", headers={"if-none-match":
response.headers["etag"]})
+ assert revalidated.status_code == 304
@pytest.mark.parametrize(
("scheme", "ssl_cert", "expected_secure"),
@@ -146,103 +220,46 @@ class TestJWTRefreshMiddleware:
pytest.param("http", "", False, id="http-no-local-ssl-cert"),
],
)
-
@patch("airflow.api_fastapi.auth.middlewares.refresh_token.get_auth_manager")
-
@patch("airflow.api_fastapi.auth.middlewares.refresh_token.resolve_user_from_token")
- @patch("airflow.api_fastapi.auth.middlewares.refresh_token.conf")
- @pytest.mark.asyncio
- async def test_dispatch_cookie_secure_flag(
- self,
- mock_conf,
- mock_resolve_user_from_token,
- mock_get_auth_manager,
- middleware,
- mock_request,
- mock_user,
- scheme,
- ssl_cert,
- expected_secure,
+ def test_cookie_secure_flag(
+ self, mock_resolve_user, mock_auth_manager, no_ssl_cert, scheme,
ssl_cert, expected_secure
):
"""The cookie Secure flag must follow the request scheme as well as
the local ssl_cert."""
- refreshed_user = MagicMock(spec=BaseUser)
- mock_request.cookies = {COOKIE_NAME_JWT_TOKEN: "valid_token"}
- mock_request.base_url = MagicMock(scheme=scheme)
- mock_resolve_user_from_token.return_value = mock_user
- mock_auth_manager = MagicMock()
- mock_get_auth_manager.return_value = mock_auth_manager
- mock_auth_manager.refresh_user.return_value = refreshed_user
+ mock_auth_manager.refresh_user.return_value = MagicMock(spec=BaseUser)
mock_auth_manager.generate_jwt.return_value = "new_token"
- mock_conf.get.return_value = ssl_cert
-
- call_next = AsyncMock(return_value=Response())
- response = await middleware.dispatch(mock_request, call_next)
-
- set_cookie_headers = response.headers.get("set-cookie", "")
- if expected_secure:
- assert "Secure" in set_cookie_headers
- else:
- assert "Secure" not in set_cookie_headers
-
-
@patch("airflow.api_fastapi.auth.middlewares.refresh_token.get_cookie_path",
return_value="/team-a/")
-
@patch("airflow.api_fastapi.auth.middlewares.refresh_token.get_auth_manager")
-
@patch("airflow.api_fastapi.auth.middlewares.refresh_token.resolve_user_from_token")
- @patch("airflow.api_fastapi.auth.middlewares.refresh_token.conf")
- @pytest.mark.asyncio
- async def test_dispatch_cookie_uses_subpath(
- self,
- mock_conf,
- mock_resolve_user_from_token,
- mock_get_auth_manager,
- mock_cookie_path,
- middleware,
- mock_request,
- mock_user,
- ):
- """When a subpath is configured, set_cookie must include it as
path=."""
- refreshed_user = MagicMock(spec=BaseUser)
- mock_request.cookies = {COOKIE_NAME_JWT_TOKEN: "valid_token"}
- mock_resolve_user_from_token.return_value = mock_user
- mock_auth_manager = MagicMock()
- mock_get_auth_manager.return_value = mock_auth_manager
- mock_auth_manager.refresh_user.return_value = refreshed_user
- mock_auth_manager.generate_jwt.return_value = "new_token"
- mock_conf.get.return_value = ""
+ no_ssl_cert.get.return_value = ssl_cert
+ app, _ = _build_app()
+ client = TestClient(
+ app, base_url=f"{scheme}://testserver",
cookies={COOKIE_NAME_JWT_TOKEN: "valid_token"}
+ )
- call_next = AsyncMock(return_value=Response())
- response = await middleware.dispatch(mock_request, call_next)
+ response = client.get("/api/v2/dags")
- set_cookie_headers = response.headers.getlist("set-cookie")
- assert any("Path=/team-a/" in h for h in set_cookie_headers)
- # Stale root-path cookie must also be cleared
- assert any(
- "Path=/" in h and "Path=/team-a/" not in h and "Max-Age=0" in h
for h in set_cookie_headers
- )
+ assert ("Secure" in response.headers["set-cookie"]) is expected_secure
-
@patch("airflow.api_fastapi.auth.middlewares.refresh_token.get_cookie_path",
return_value="/team-a/")
- @patch.object(
- JWTRefreshMiddleware,
- "_refresh_user",
- side_effect=HTTPException(status_code=403, detail="Invalid JWT token"),
+ @pytest.mark.parametrize(
+ "refresh_outcome",
+ [
+ pytest.param("refreshed", id="refreshed-token"),
+ pytest.param("invalid", id="invalid-token"),
+ ],
)
- @patch("airflow.api_fastapi.auth.middlewares.refresh_token.conf")
- @pytest.mark.asyncio
- async def test_dispatch_invalid_token_clears_root_cookie(
- self,
- mock_conf,
- mock_refresh_user,
- mock_cookie_path,
- middleware,
- mock_request,
+ @patch(f"{MODULE}.get_cookie_path", autospec=True, return_value="/team-a/")
+ def test_cookie_uses_subpath_and_clears_root_cookie(
+ self, mock_cookie_path, mock_resolve_user, mock_auth_manager,
refresh_outcome
):
- """When a stale _token exists at root path, clearing must target both
the subpath and root."""
- mock_request.cookies = {COOKIE_NAME_JWT_TOKEN: "stale_root_token"}
- mock_conf.get.return_value = ""
+ """With a subpath, the cookie is scoped to it and a stale root-path
cookie is cleared."""
+ if refresh_outcome == "refreshed":
+ mock_auth_manager.refresh_user.return_value =
MagicMock(spec=BaseUser)
+ mock_auth_manager.generate_jwt.return_value = "new_token"
+ else:
+ mock_resolve_user.side_effect = HTTPException(status_code=403,
detail="Invalid JWT token")
+ app, _ = _build_app()
- call_next = AsyncMock(return_value=Response(status_code=401))
- response = await middleware.dispatch(mock_request, call_next)
+ response = TestClient(app, cookies={COOKIE_NAME_JWT_TOKEN:
"token"}).get("/api/v2/dags")
- set_cookie_headers = response.headers.getlist("set-cookie")
- # Expect two delete cookies: one at the subpath and one at root "/"
- assert any("Path=/team-a/" in h and "Max-Age=0" in h for h in
set_cookie_headers)
+ set_cookie_headers = response.headers.get_list("set-cookie")
+ assert len(set_cookie_headers) == 2
+ assert any("Path=/team-a/" in h for h in set_cookie_headers)
assert any(
"Path=/" in h and "Path=/team-a/" not in h and "Max-Age=0" in h
for h in set_cookie_headers
)
diff --git a/airflow-core/tests/unit/api_fastapi/test_app.py
b/airflow-core/tests/unit/api_fastapi/test_app.py
index c994f90a0d5..c5d5e33222f 100644
--- a/airflow-core/tests/unit/api_fastapi/test_app.py
+++ b/airflow-core/tests/unit/api_fastapi/test_app.py
@@ -31,6 +31,7 @@ from starlette.middleware.gzip import GZipMiddleware
import airflow.api_fastapi.app as app_module
import airflow.plugins_manager as plugins_manager
from airflow import settings
+from airflow.api_fastapi.auth.middlewares.refresh_token import
JWTRefreshMiddleware
from airflow.api_fastapi.common.http_access_log import HttpAccessLogMiddleware
from airflow.api_fastapi.common.http_metrics import HttpMetricsMiddleware
from airflow.utils.session import create_session_async
@@ -212,6 +213,15 @@ def
test_api_metrics_middleware_skipped_without_a_metrics_backend(apps):
assert HttpMetricsMiddleware not in installed
+def test_jwt_refresh_middleware_skips_static_mount():
+ """Static UI assets are served without resolving the user from the JWT
cookie."""
+ [middleware] = [
+ m for m in app_module.create_app(apps="core").user_middleware if m.cls
is JWTRefreshMiddleware
+ ]
+
+ assert middleware.kwargs == {"skip_path_prefixes": ("/static/",)}
+
+
def test_catch_all_route_last():
"""
Ensure the catch all route that returns the initial html is the last route
in the fastapi app.
