Skip to content

Make JWTRefreshMiddleware extensible via BaseAuthManager.get_jwt_refresh_middleware() #70783

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
pierrejeambrun merged 1 commit into from
Oct 22, 2026
Merged

Make JWTRefreshMiddleware extensible via BaseAuthManager.get_jwt_refresh_middleware() #70783

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
18 changes: 18 additions & 0 deletions airflow-core/src/airflow/api_fastapi/app.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -45,6 +45,8 @@
from airflow.utils.providers_configuration_loader import providers_configuration_loaded

if TYPE_CHECKING:
from fastapi import Request

from airflow.api_fastapi.auth.managers.base_auth_manager import BaseAuthManager

API_BASE_URL = conf.get("api", "base_url", fallback="")
Expand All @@ -65,6 +67,22 @@ def get_cookie_path() -> str:
return API_ROOT_PATH or "/"


def request_cookie_is_secure(request: Request) -> bool:
"""
Return whether a FastAPI request cookie should be tagged as HTTP secure.

:param request: FastAPI Request object

usage:
```python
secure = request_cookie_is_secure(request)

response.set_cookie("mycookie", "myvalue", path=get_cookie_path(), secure=secure, httponly=True)
```
"""
return request.base_url.scheme == "https" or bool(conf.get("api", "ssl_cert", fallback=""))


# Fast API apps mounted under these prefixes are not allowed
RESERVED_URL_PREFIXES = ["/api/v2", "/ui", "/execution", "/auth", "/pluginsv2"]

Expand Down
14 changes: 14 additions & 0 deletions airflow-core/src/airflow/api_fastapi/auth/managers/base_auth_manager.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -185,6 +185,20 @@ def get_fastapi_middlewares(self) -> list[tuple[_MiddlewareFactory[Any], dict[st
"""
return []

def get_jwt_refresh_middleware(self) -> list[tuple[_MiddlewareFactory[Any], dict[str, Any]]]:
Comment thread
stephen-bracken marked this conversation as resolved.
Resolved
"""
Return the JWTRefreshMiddleware to refresh the Airflow JWT token.

Middlewares added by this method will be initialised before middleware from
``get_fastapi_middlewares()``. To disable JWT refresh return an empty list.
To add more middleware, simply add more tuples here in the same format as
``get_fastapi_middlewares()``.
"""
from airflow.api_fastapi.auth.middlewares.refresh_token import JWTRefreshMiddleware
Comment thread
stephen-bracken marked this conversation as resolved.
Resolved
from airflow.api_fastapi.core_api.app import _STATIC_URL_PATH

return [(JWTRefreshMiddleware, {"skip_path_prefixes": (f"{_STATIC_URL_PATH}/",)})]

def generate_jwt(
self, user: T, *, expiration_time_in_seconds: int = conf.getint("api_auth", "jwt_expiration_time")
) -> str:
Expand Down
77 changes: 61 additions & 16 deletions airflow-core/src/airflow/api_fastapi/auth/middlewares/refresh_token.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -24,19 +24,19 @@
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.app import get_auth_manager, get_cookie_path, request_cookie_is_secure
from airflow.api_fastapi.auth.managers.base_auth_manager import COOKIE_NAME_JWT_TOKEN
from airflow.api_fastapi.auth.managers.exceptions import AuthManagerRefreshTokenExpiredException
from airflow.api_fastapi.auth.managers.models.base_user import BaseUser
from airflow.api_fastapi.core_api.security import (
USER_INJECTED_BY_TRUSTED_MIDDLEWARE,
resolve_user_from_token,
)
from airflow.configuration import conf

if TYPE_CHECKING:
from starlette.types import ASGIApp, Message, Receive, Scope, Send

from airflow.api_fastapi.auth.managers.models.base_user import BaseUser

# 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"])
Expand Down Expand Up @@ -72,27 +72,29 @@ async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
await self.app(scope, receive, send)
return

new_token = None
new_user = None
session_invalidated = False
try:
new_user, current_user = await self._refresh_user(current_token)
new_user, current_user = await self.refresh_user(request)
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 = ""
session_invalidated = True

token_cookie, cleanup_cookies = (
self._build_cookie_headers(request, new_token) if new_token is not None else (None, [])
)
if session_invalidated or new_user is not None:
extra_cookies = self.get_extra_cookies(request, new_user)
token_cookie, cleanup_cookies = self.build_cookie_headers(
request, await self.generate_new_user_token(request, new_user), extra_cookies
)
else:
token_cookie, cleanup_cookies = (None, [])
response_started = False

async def send_with_cookies(message: Message) -> None:
Expand Down Expand Up @@ -124,10 +126,34 @@ def _is_skipped(self, scope: Scope) -> bool:
return path in UNAUTHENTICATED_PATHS or path.startswith(self.skip_path_prefixes)

@staticmethod
def _build_cookie_headers(request: Request, new_token: str) -> tuple[str, list[str]]:
"""Return the ``Set-Cookie`` value of the token and those clearing stale token cookies."""
def get_extra_cookies(request: Request, new_user: BaseUser | None) -> dict[str, str]:
"""
Get extra cookies from the request context to set in a new token response.

Setting a cookie to "" will cause it to be cleared by build_cookie_headers()
"""
return {}

@staticmethod
async def generate_new_user_token(request: Request, new_user: BaseUser | None) -> str:
"""Generate a new Airflow JWT token from the user model."""
return get_auth_manager().generate_jwt(new_user) if new_user else ""

@staticmethod
def build_cookie_headers(
request: Request, new_token: str, extra_cookies: dict[str, str]
) -> tuple[str, list[str]]:
"""
Return the ``Set-Cookie`` value of the token and those clearing stale token cookies.

Any cookies with a nil value ("") will be cleared with Max-Age=0

:param request: FastAPI Request
:param new_token: New Airflow JWT token
:param extra_cookies: Additional cookie values to set
"""
cookie_path = get_cookie_path()
secure = request.base_url.scheme == "https" or bool(conf.get("api", "ssl_cert", fallback=""))
secure = request_cookie_is_secure(request)
# Only used to render the Set-Cookie header values
response = Response()
response.set_cookie(
Expand All @@ -151,13 +177,32 @@ def _build_cookie_headers(request: Request, new_token: str) -> tuple[str, list[s
secure=secure,
samesite="lax",
)
# Add any additional cookies required
for cookie, value in extra_cookies.items():
response.set_cookie(
cookie,
value,
path=cookie_path,
httponly=True,
secure=secure,
samesite="lax",
max_age=0 if value == "" else None,
)
token_cookie, *cleanup_cookies = (
value.decode("latin-1") for key, value in response.raw_headers if key == b"set-cookie"
)
return token_cookie, cleanup_cookies

@staticmethod
async def _refresh_user(current_token: str) -> tuple[BaseUser | None, BaseUser | None]:
async def refresh_user(request: Request) -> tuple[BaseUser | None, BaseUser | None]:
"""
Refresh the logged in user using the current JWT Token.

If the user is not authenticated, return ``None, None``

:param request: FastAPI Request
"""
current_token = request.cookies.get(COOKIE_NAME_JWT_TOKEN)
user = await resolve_user_from_token(current_token)
# 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
8 changes: 5 additions & 3 deletions airflow-core/src/airflow/api_fastapi/core_api/app.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -304,11 +304,13 @@ def init_config(app: FastAPI) -> None:

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, skip_path_prefixes=(f"{_STATIC_URL_PATH}/",))
auth_manager = get_auth_manager()

for middleware_cls, middleware_kwargs in get_auth_manager().get_fastapi_middlewares():
for middleware_cls, middleware_kwargs in (
*auth_manager.get_jwt_refresh_middleware(),
*auth_manager.get_fastapi_middlewares(),
):
app.add_middleware(middleware_cls, **middleware_kwargs)

# GZipMiddleware must be inside HttpMetricsMiddleware and HttpAccessLogMiddleware so that both
Expand Down
10 changes: 10 additions & 0 deletions airflow-core/tests/unit/api_fastapi/auth/managers/test_base_auth_manager.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -36,8 +36,10 @@
TeamDetails,
VariableDetails,
)
from airflow.api_fastapi.auth.middlewares.refresh_token import JWTRefreshMiddleware
from airflow.api_fastapi.auth.tokens import JWTGenerator, JWTValidator
from airflow.api_fastapi.common.types import MenuItem
from airflow.api_fastapi.core_api.app import _STATIC_URL_PATH
from airflow.exceptions import RemovedInAirflow4Warning
from airflow.models.pool import Pool
from airflow.models.team import Team
Expand Down Expand Up @@ -932,3 +934,11 @@ def test_is_authorized_hitl_task(
user = BaseAuthManagerUserTest(name=user_id)
result = auth_manager.is_authorized_hitl_task(assigned_users=assigned_users, user=user)
assert result == expected

def test_get_jwt_refresh_middleware(self, auth_manager):
assert auth_manager.get_jwt_refresh_middleware() == [
(JWTRefreshMiddleware, {"skip_path_prefixes": (f"{_STATIC_URL_PATH}/",)})
]

def test_get_fastapi_middlewares(self, auth_manager):
Comment thread
stephen-bracken marked this conversation as resolved.
Resolved
assert auth_manager.get_fastapi_middlewares() == []
81 changes: 71 additions & 10 deletions airflow-core/tests/unit/api_fastapi/auth/middlewares/test_refresh_token.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -32,6 +32,8 @@
from airflow.api_fastapi.auth.middlewares.refresh_token import JWTRefreshMiddleware
from airflow.api_fastapi.core_api.security import USER_INJECTED_BY_TRUSTED_MIDDLEWARE

from tests_common.test_utils.config import conf_vars

MODULE = "airflow.api_fastapi.auth.middlewares.refresh_token"


Expand Down Expand Up @@ -97,9 +99,8 @@ def mock_resolve_user():

@pytest.fixture(autouse=True)
def no_ssl_cert():
with patch(f"{MODULE}.conf", autospec=True) as mock_conf:
mock_conf.get.return_value = ""
yield mock_conf
with conf_vars({("api", "ssl_cert"): ""}):
yield


class TestJWTRefreshMiddleware:
Expand Down Expand Up @@ -272,18 +273,18 @@ def test_static_files_served_unchanged(
],
)
def test_cookie_secure_flag(
self, mock_resolve_user, mock_auth_manager, no_ssl_cert, scheme, ssl_cert, expected_secure
self, mock_resolve_user, mock_auth_manager, scheme, ssl_cert, expected_secure
):
"""The cookie Secure flag must follow the request scheme as well as the local ssl_cert."""
mock_auth_manager.refresh_user.return_value = MagicMock(spec=BaseUser)
mock_auth_manager.generate_jwt.return_value = "new_token"
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"}
)
with conf_vars({("api", "ssl_cert"): ssl_cert}):
app, _ = _build_app()
client = TestClient(
app, base_url=f"{scheme}://testserver", cookies={COOKIE_NAME_JWT_TOKEN: "valid_token"}
)

response = client.get("/api/v2/dags")
response = client.get("/api/v2/dags")

assert ("Secure" in response.headers["set-cookie"]) is expected_secure

Expand Down Expand Up @@ -314,3 +315,63 @@ def test_cookie_uses_subpath_and_clears_root_cookie(
assert any(
"Path=/" in h and "Path=/team-a/" not in h and "Max-Age=0" in h for h in set_cookie_headers
)

@pytest.mark.asyncio
@pytest.mark.parametrize("new_user", [(MagicMock(),), (None,)])
async def test_generate_new_user_token(self, mock_auth_manager, new_user):
mock_auth_manager.generate_jwt.return_value = "new_token"
app, seen = _build_app()
middleware = JWTRefreshMiddleware(app)
new_token = await middleware.generate_new_user_token(MagicMock(), new_user)
if new_user is None:
assert new_token == ""
else:
assert new_token == "new_token"

@patch(f"{MODULE}.JWTRefreshMiddleware.get_extra_cookies")
def test_refresh_user_with_extra_cookies(
self, mock_get_extra_cookies: MagicMock, mock_resolve_user, mock_auth_manager
):
mock_get_extra_cookies.return_value = {"foo": "bar"}
refreshed_user = MagicMock(spec=BaseUser)
mock_auth_manager.refresh_user.return_value = refreshed_user
mock_auth_manager.generate_jwt.return_value = "new_token"
app, seen = _build_app()

response = TestClient(app, cookies={COOKIE_NAME_JWT_TOKEN: "valid_token"}).get("/api/v2/dags")

# 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 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)
assert (
response.headers["set-cookie"]
== "_token=new_token; HttpOnly; Path=/; SameSite=lax, foo=bar; HttpOnly; Path=/; SameSite=lax"
)

@patch(f"{MODULE}.JWTRefreshMiddleware.get_extra_cookies")
def test_refresh_user_clear_extra_cookies(
self, mock_get_extra_cookies: MagicMock, mock_resolve_user, mock_auth_manager
):
mock_get_extra_cookies.return_value = {"foo": ""}
refreshed_user = MagicMock(spec=BaseUser)
mock_auth_manager.refresh_user.return_value = refreshed_user
mock_auth_manager.generate_jwt.return_value = "new_token"
app, seen = _build_app()

response = TestClient(app, cookies={COOKIE_NAME_JWT_TOKEN: "valid_token"}).get("/api/v2/dags")

# 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 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)
assert (
response.headers["set-cookie"]
== '_token=new_token; HttpOnly; Path=/; SameSite=lax, foo=""; HttpOnly; Max-Age=0; Path=/; SameSite=lax'
)
4 changes: 2 additions & 2 deletions airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_auth.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -367,7 +367,7 @@ def test_logout_clears_cookie_when_logout_url_redirects(self, logout_client):
# The refresh middleware is stubbed out so that only the logout route can clear the cookie.
with (
patch(
"airflow.api_fastapi.auth.middlewares.refresh_token.JWTRefreshMiddleware._refresh_user",
"airflow.api_fastapi.auth.middlewares.refresh_token.JWTRefreshMiddleware.refresh_user",
new=AsyncMock(return_value=(None, None)),
),
patch.object(auth_manager, "get_url_logout", return_value="http://external/logout"),
Expand All @@ -389,7 +389,7 @@ def test_logout_clears_cookie_when_revocation_cannot_be_recorded(self, logout_cl
logout_client.cookies.set(COOKIE_NAME_JWT_TOKEN, token_str)
with (
patch(
"airflow.api_fastapi.auth.middlewares.refresh_token.JWTRefreshMiddleware._refresh_user",
"airflow.api_fastapi.auth.middlewares.refresh_token.JWTRefreshMiddleware.refresh_user",
new=AsyncMock(return_value=(None, None)),
),
patch.object(auth_manager, "get_url_logout", return_value="http://external/logout"),
Expand Down
Loading