Re-work JWT Validation and Generation to use public/private key and official claims (apache#46981)

Since requests to Task Execution API can originate  out-of-"cluser" so to
speak, this PR re-works the JWTSigner class so that it is possible to use
public/private keys (vs just a simple pre-shared secret). This is useful in
many cloud environments where many companies have security requirements that
ingress gateways must validate the JWT tokens on the way in, and the only way
of doing this is with public keys

So that we don't have two different ways of generating JWT tokens I have
totally replaced the old "JWTSigner" class (which it turns out didn't have any
unit test of its own, it was only tested indirectly through test_serve_logs
etc).

As part of this change I have also changed the JWT that was generated by the
SimpleAuthManager and the AwsAuthManager (the only two we have that use JWT)
to use the offical `sub` (subject) clain to place the user identifer rather
than a custom claim name.

And although it might seem slightly strange at first, I have made the
JWTValidator an async class internally. (Hence `avalidated_claims` -- the `a`
prefix signifies async, much like `aclose` or `aread` on HTTPX async
responses). This allows us to periodically refresh the JWK document if
configured in the background, and using asgiref's async_to_sync means we only
have one version.

Co-authored-by: Jens Scheffler <jscheffl@apache.org>

Author: ashb

airflow/api_fastapi/auth/managers/base_auth_manager.py

@@ -19,18 +19,24 @@
 
 import logging
 from abc import ABCMeta, abstractmethod
+from functools import cache
 from typing import TYPE_CHECKING, Any, Generic, TypeVar
 
 from jwt import InvalidTokenError
 from sqlalchemy import select
 
 from airflow.api_fastapi.auth.managers.models.base_user import BaseUser
 from airflow.api_fastapi.auth.managers.models.resource_details import BackfillDetails, DagDetails
+from airflow.api_fastapi.auth.tokens import (
+    JWTGenerator,
+    JWTValidator,
+    get_sig_validation_args,
+    get_signing_args,
+)
 from airflow.api_fastapi.common.types import ExtraMenuItem, MenuItem
 from airflow.configuration import conf
 from airflow.models import DagModel
 from airflow.typing_compat import Literal
-from airflow.utils.jwt_signer import JWTSigner, get_signing_key
 from airflow.utils.log.logging_mixin import LoggingMixin
 from airflow.utils.session import NEW_SESSION, provide_session
 
@@ -86,24 +92,24 @@ def deserialize_user(self, token: dict[str, Any]) -> T:
 
     @abstractmethod
     def serialize_user(self, user: T) -> dict[str, Any]:
-        """Create a dict from a user object."""
+        """Create a subject and extra claims dict from a user object."""
 
-    def get_user_from_token(self, token: str) -> BaseUser:
+    async def get_user_from_token(self, token: str) -> BaseUser:
         """Verify the JWT token is valid and create a user object from it if valid."""
         try:
-            payload: dict[str, Any] = self._get_token_signer().verify_token(token)
+            payload: dict[str, Any] = await self._get_token_validator().avalidated_claims(token)
             return self.deserialize_user(payload)
         except InvalidTokenError as e:
-            log.error("JWT token is not valid")
+            log.error("JWT token is not valid: %s", e)
             raise e
 
-    def get_jwt_token(
-        self, user: T, *, expiration_time_in_seconds: int = conf.getint("api", "auth_jwt_expiration_time")
+    def generate_jwt(
+        self, user: T, *, expiration_time_in_seconds: int = conf.getint("api_auth", "jwt_expiration_time")
     ) -> str:
         """Return the JWT token from a user object."""
-        return self._get_token_signer(
-            expiration_time_in_seconds=expiration_time_in_seconds
-        ).generate_signed_token(self.serialize_user(user))
+        return self._get_token_signer(expiration_time_in_seconds=expiration_time_in_seconds).generate(
+            self.serialize_user(user)
+        )
 
     @abstractmethod
     def get_url_login(self, **kwargs) -> str:
@@ -450,19 +456,35 @@ def get_extra_menu_items(self, *, user: T) -> list[ExtraMenuItem]:
         """
         return []
 
-    @staticmethod
+    @classmethod
+    @cache
     def _get_token_signer(
-        expiration_time_in_seconds: int = conf.getint("api", "auth_jwt_expiration_time"),
-    ) -> JWTSigner:
+        cls,
+        expiration_time_in_seconds: int = conf.getint("api_auth", "jwt_expiration_time"),
+    ) -> JWTGenerator:
         """
         Return the signer used to sign JWT token.
 
         :meta private:
 
         :param expiration_time_in_seconds: expiration time in seconds of the token
         """
-        return JWTSigner(
-            secret_key=get_signing_key("api", "auth_jwt_secret"),
-            expiration_time_in_seconds=expiration_time_in_seconds,
-            audience="front-apis",
+        return JWTGenerator(
+            **get_signing_args(),
+            valid_for=expiration_time_in_seconds,
+            audience=conf.get("api", "jwt_audience", fallback="apache-airflow"),
+        )
+
+    @classmethod
+    @cache
+    def _get_token_validator(cls) -> JWTValidator:
+        """
+        Return the signer used to sign JWT token.
+
+        :meta private:
+        """
+        return JWTValidator(
+            **get_sig_validation_args(),
+            leeway=conf.getint("api_auth", "jwt_leeway"),
+            audience=conf.get("api_auth", "jwt_audience", fallback="apache-airflow"),
         )

airflow/api_fastapi/auth/managers/simple/routes/login.py

@@ -62,7 +62,7 @@ def create_token_all_admins() -> RedirectResponse:
         username="Anonymous",
         role="ADMIN",
     )
-    url = urljoin(conf.get("api", "base_url"), f"?token={get_auth_manager().get_jwt_token(user)}")
+    url = urljoin(conf.get("api", "base_url"), f"?token={get_auth_manager().generate_jwt(user)}")
     return RedirectResponse(url=url)
 
 
@@ -76,5 +76,5 @@ def create_token_cli(
 ) -> LoginResponse:
     """Authenticate the user for the CLI."""
     return SimpleAuthManagerLogin.create_token(
-        body=body, expiration_time_in_sec=conf.getint("api", "auth_jwt_cli_expiration_time")
+        body=body, expiration_time_in_sec=conf.getint("api_auth", "jwt_cli_expiration_time")
     )

airflow/api_fastapi/auth/managers/simple/services/login.py

@@ -31,7 +31,7 @@ class SimpleAuthManagerLogin:
 
     @classmethod
     def create_token(
-        cls, body: LoginBody, expiration_time_in_sec: int = conf.getint("api", "auth_jwt_expiration_time")
+        cls, body: LoginBody, expiration_time_in_sec: int = conf.getint("api_auth", "jwt_expiration_time")
     ) -> LoginResponse:
         """
         Authenticate user with given configuration.
@@ -67,7 +67,7 @@ def create_token(
         )
 
         return LoginResponse(
-            jwt_token=get_auth_manager().get_jwt_token(
+            jwt_token=get_auth_manager().generate_jwt(
                 user=user, expiration_time_in_seconds=expiration_time_in_sec
             )
         )

airflow/api_fastapi/auth/managers/simple/simple_auth_manager.py

@@ -139,10 +139,10 @@ def get_url_login(self, **kwargs) -> str:
         return AUTH_MANAGER_FASTAPI_APP_PREFIX + "/login"
 
     def deserialize_user(self, token: dict[str, Any]) -> SimpleAuthManagerUser:
-        return SimpleAuthManagerUser(username=token["username"], role=token["role"])
+        return SimpleAuthManagerUser(username=token["sub"], role=token["role"])
 
     def serialize_user(self, user: SimpleAuthManagerUser) -> dict[str, Any]:
-        return {"username": user.username, "role": user.role}
+        return {"sub": user.username, "role": user.role}
 
     def is_authorized_configuration(
         self,