[stephen-bracken]

closes: #70720

Make JWTRefreshMiddleware extensible via BaseAuthManager.get_jwt_refresh_middleware()

Refactor the initialisation flow for the JWTRefreshMiddleware to give greater control to auth managers over how this middleware gets initialised.

By moving the JWTRefreshMiddleware initialisation to BaseAuthManager.get_jwt_refresh_middleware(), this allows auth managers to override the JWTRefreshMiddleware behaviour by inheriting from it and using the inherited class in get_fastapi_middlewares instead.

• Changed the _refresh_user() and _build_cookie_headers() methods to public refresh_user() and build_cookie_headers() methods to allow accessing any cookies or state information from the request and setting cookies in the response.
• get_extra_cookies() allows subclasses to add / clear additional cookies in the response when a new user token is issued.
• generate_new_user_token() allows subclasses to control how JWTs are issued for new users.

Also adds the airflow.api_fastapi.app.request_cookie_is_secure() helper to standardise setting HTTP secure cookies

Was generative AI tooling used to co-author this PR?

• Yes (please specify the tool below)
• No

---

[vincbeck]

CI is failing

[stephen-bracken]

@pierrejeambrun

I believe this is breaking the because core no longer hardcode JWTRefreshMiddleware, but old provider doesn't call super().get_fastapi_middlewares().

I have added BaseAuthManager.use_jwt_middleware (default True) that will trigger the old refresh behaviour where necessary. This allows auth managers to control the order in which the middleware from get_jwt_refresh_middleware() gets initialised without breaking old auth flows.

[vincbeck]

I am trying to understand, why it is breaking? Do we really need BaseAuthManager.use_jwt_middleware?

[stephen-bracken]

@vincbeck The original implementation of this middleware made an implicit assumption that the Airflow JWT can fully reconstruct the user model for the request. The JWTRefreshMiddleware gets initialised as the first middleware in the FastAPI app and this behaviour is not modifiable by the auth manager. See initialisation here:

airflow/airflow-core/src/airflow/api_fastapi/core_api/app.py

Line 309 in 4fb12b3

app.add_middleware(JWTRefreshMiddleware, skip_path_prefixes=(f"{_STATIC_URL_PATH}/",))

That assumption is broken for auth managers that use additional cookies, e.g. in #70800 I added KeycloakJWTMiddleware to the KeycloakAuthManager that manages a separate JWT for Keycloak in a separate cookie. There are good reasons to do this like reducing the size of the Airflow JWT claims, but the repercussion is that the Airflow JWT no longer contains a full user model, as the Keycloak JWTs have to be injected from separate cookies.
There may be other scenarios where auth managers want to control how and where the JWT gets refreshed, so adding a use_jwt_middleware property to the BaseAuthManager allows other auth managers to disable the default behaviour which adds the JWTRefreshMiddleware to the app as the first middleware, and take full control of the user request validation process. The intent of this PR is to allow the KeycloakAuthManager to override the base JWTRefreshMiddleware with the KeycloakJWTMiddleware or similar so that we don't have to break the auth manager interface as badly in KeycloakAuthManager.

Adding the use_jwt_middleware property to the BaseAuthManager with default value True preserves the existing JWT behaviour for any auth managers that do not need this override.

[stephen-bracken]

The getattr guard is what preserves the behaviour for backwards compatibility:

https://github.com/apache/airflow/pull/70783/changes#diff-da26ba2c49d5bc5f117895c11ad7fea6e24615eccd3d27c3300a5026195b5cd8R310

[vincbeck]

But do we really need this use_jwt_middleware? If an auth manager do not want to use JWTRefreshMiddleware, it can just override get_jwt_refresh_middleware and return an empty list?

[stephen-bracken]

I see your point - I can modify the init procedure to prepend the jwt middleware if present.

[vincbeck]

LGTM but I would like to have @pierrejeambrun's review

[pierrejeambrun]

Thanks for iterating on this, I think the current shape is much better: keeping the JWT middleware in its own overridable method keeps old providers working on the new core, and preserves the middleware order. A few things to fix before it's ready, the first one is breaking the API server (and the standalone CI job). Could you also update the PR description? It still mentions use_jwt_middleware and _set_new_token.

[stephen-bracken]

I've fixed the CI issues and added a get_extra_cookies() method that auth manager middlewares can subclass to set / clear extra cookies in new token responses (e.g. Keycloak tokens)