You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@airflow.apache.org by as...@apache.org on 2025/03/13 11:01:39 UTC
(airflow) 01/01: fixup! Add JWT validation and generation machinery for the Task Execution API to use
This is an automated email from the ASF dual-hosted git repository.
ash pushed a commit to branch generate-task-jwt-tokens
in repository https://gitbox.apache.org/repos/asf/airflow.git
commit 6c77bb2c8bb32d2b0c50544560d8a7ae994df9f3
Author: Ash Berlin-Taylor <as...@apache.org>
AuthorDate: Wed Mar 12 21:18:41 2025 +0000
fixup! Add JWT validation and generation machinery for the Task Execution API to use
---
.../api_fastapi/auth/managers/base_auth_manager.py | 6 ++--
airflow/api_fastapi/auth/tokens.py | 5 ++++
airflow/api_fastapi/core_api/security.py | 6 ++--
airflow/assets/manager.py | 2 +-
airflow/configuration.py | 2 +-
airflow/utils/log/file_task_handler.py | 2 +-
airflow/utils/serve_logs.py | 2 ++
hatch_build.py | 2 +-
.../providers/fab/auth_manager/fab_auth_manager.py | 4 +--
.../unit/fab/auth_manager/test_fab_auth_manager.py | 4 +--
.../auth/managers/test_base_auth_manager.py | 13 +++++----
tests/api_fastapi/auth/test_tokens.py | 2 +-
tests/api_fastapi/conftest.py | 2 +-
tests/api_fastapi/core_api/test_security.py | 25 +++++++++--------
tests/utils/test_serve_logs.py | 32 +++++++++++++---------
15 files changed, 62 insertions(+), 47 deletions(-)
diff --git a/airflow/api_fastapi/auth/managers/base_auth_manager.py b/airflow/api_fastapi/auth/managers/base_auth_manager.py
index 110906268cb..5fe24f85624 100644
--- a/airflow/api_fastapi/auth/managers/base_auth_manager.py
+++ b/airflow/api_fastapi/auth/managers/base_auth_manager.py
@@ -95,13 +95,13 @@ class BaseAuthManager(Generic[T], LoggingMixin, metaclass=ABCMeta):
def serialize_user(self, user: T) -> dict[str, Any]:
"""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_validator().validated_claims(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 generate_jwt(
diff --git a/airflow/api_fastapi/auth/tokens.py b/airflow/api_fastapi/auth/tokens.py
index d76037af77a..a85c9ff8f27 100644
--- a/airflow/api_fastapi/auth/tokens.py
+++ b/airflow/api_fastapi/auth/tokens.py
@@ -405,6 +405,11 @@ class JWTGenerator:
"exp": int(now + self.valid_for),
"iat": now,
}
+ if claims["iss"] is None:
+ del claims["iss"]
+ if claims["aud"] is None:
+ del claims["aud"]
+
if extras is not None:
claims.update(extras)
headers = {"alg": self.algorithm, **(headers or {})}
diff --git a/airflow/api_fastapi/core_api/security.py b/airflow/api_fastapi/core_api/security.py
index 2ae47ce188c..5256d62de98 100644
--- a/airflow/api_fastapi/core_api/security.py
+++ b/airflow/api_fastapi/core_api/security.py
@@ -53,9 +53,9 @@ if TYPE_CHECKING:
oauth2_scheme = OAuth2PasswordBearer(tokenUrl="token")
-def get_user(token_str: Annotated[str, Depends(oauth2_scheme)]) -> BaseUser:
+async def get_user(token_str: Annotated[str, Depends(oauth2_scheme)]) -> BaseUser:
try:
- return get_auth_manager().get_user_from_token(token_str)
+ return await get_auth_manager().get_user_from_token(token_str)
except ExpiredSignatureError:
raise HTTPException(status.HTTP_401_UNAUTHORIZED, "Token Expired")
except InvalidTokenError:
@@ -79,7 +79,7 @@ async def get_user_with_exception_handling(request: Request) -> BaseUser | None:
if not token_str: # Handle None or empty token
return None
- return get_user(token_str)
+ return await get_user(token_str)
def requires_access_dag(
diff --git a/airflow/assets/manager.py b/airflow/assets/manager.py
index ca6cb345dd4..4e5dd37a66b 100644
--- a/airflow/assets/manager.py
+++ b/airflow/assets/manager.py
@@ -17,7 +17,7 @@
# under the License.
from __future__ import annotations
-from collections.abc import Collection, Iterable
+from collections.abc import Collection
from typing import TYPE_CHECKING
from sqlalchemy import exc, or_, select
diff --git a/airflow/configuration.py b/airflow/configuration.py
index cc3517b5790..a5a478bdcb9 100644
--- a/airflow/configuration.py
+++ b/airflow/configuration.py
@@ -95,7 +95,7 @@ def expand_env_var(env_var: str | None) -> str | None:
Repeat and apply `expandvars` and `expanduser` until
interpolation stops having any effect.
"""
- if not env_var:
+ if not env_var or not isinstance(env_var, str):
return env_var
while True:
interpolated = os.path.expanduser(os.path.expandvars(str(env_var)))
diff --git a/airflow/utils/log/file_task_handler.py b/airflow/utils/log/file_task_handler.py
index 16396f17118..4d45216bfa3 100644
--- a/airflow/utils/log/file_task_handler.py
+++ b/airflow/utils/log/file_task_handler.py
@@ -107,7 +107,7 @@ def _fetch_logs_from_service(url, log_relative_path):
response = requests.get(
url,
timeout=timeout,
- headers={"Authorization": signer.generate({"sub": "task-logs", "filename": log_relative_path})},
+ headers={"Authorization": signer.generate({"filename": log_relative_path})},
)
response.encoding = "utf-8"
return response
diff --git a/airflow/utils/serve_logs.py b/airflow/utils/serve_logs.py
index ab364aa977b..a665f1807d4 100644
--- a/airflow/utils/serve_logs.py
+++ b/airflow/utils/serve_logs.py
@@ -45,6 +45,7 @@ logger = logging.getLogger(__name__)
def create_app():
flask_app = Flask(__name__, static_folder=None)
+ leeway = conf.getint("webserver", "log_request_clock_grace", fallback=30)
log_directory = os.path.expanduser(conf.get("logging", "BASE_LOG_FOLDER"))
log_config_class = conf.get("logging", "logging_config_class")
if log_config_class:
@@ -72,6 +73,7 @@ def create_app():
signer = JWTValidator(
issuer=None,
secret_key=conf.get("webserver", "secret_key"),
+ leeway=leeway,
audience="task-instance-logs",
)
diff --git a/hatch_build.py b/hatch_build.py
index 50a5c3bde63..8edbbe6d0ee 100644
--- a/hatch_build.py
+++ b/hatch_build.py
@@ -248,7 +248,7 @@ DEPENDENCIES = [
# Pygments 2.19.0 improperly renders .ini files with dictionaries as values
# See https://github.com/pygments/pygments/issues/2834
"pygments>=2.0.1,!=2.19.0",
- "pyjwt>=2.0.0",
+ "pyjwt>=2.7.0",
"python-daemon>=3.0.0",
"python-dateutil>=2.7.0",
"python-nvd3>=0.15.0",
diff --git a/providers/fab/src/airflow/providers/fab/auth_manager/fab_auth_manager.py b/providers/fab/src/airflow/providers/fab/auth_manager/fab_auth_manager.py
index 7c8672be8c1..9c7b84e758d 100644
--- a/providers/fab/src/airflow/providers/fab/auth_manager/fab_auth_manager.py
+++ b/providers/fab/src/airflow/providers/fab/auth_manager/fab_auth_manager.py
@@ -242,10 +242,10 @@ class FabAuthManager(BaseAuthManager[User]):
def deserialize_user(self, token: dict[str, Any]) -> User:
with create_session() as session:
- return session.get(User, token["sub"])
+ return session.get(User, int(token["sub"]))
def serialize_user(self, user: User) -> dict[str, Any]:
- return {"sub": user.id}
+ return {"sub": str(user.id)}
def is_logged_in(self) -> bool:
"""Return whether the user is logged in."""
diff --git a/providers/fab/tests/unit/fab/auth_manager/test_fab_auth_manager.py b/providers/fab/tests/unit/fab/auth_manager/test_fab_auth_manager.py
index ce7d58e20ec..99b53a1437f 100644
--- a/providers/fab/tests/unit/fab/auth_manager/test_fab_auth_manager.py
+++ b/providers/fab/tests/unit/fab/auth_manager/test_fab_auth_manager.py
@@ -136,13 +136,13 @@ class TestFabAuthManager:
def test_deserialize_user(self, flask_app, auth_manager_with_appbuilder):
user = create_user(flask_app, "test")
- result = auth_manager_with_appbuilder.deserialize_user({"id": user.id})
+ result = auth_manager_with_appbuilder.deserialize_user({"sub": str(user.id)})
assert user == result
def test_serialize_user(self, flask_app, auth_manager_with_appbuilder):
user = create_user(flask_app, "test")
result = auth_manager_with_appbuilder.serialize_user(user)
- assert result == {"id": user.id}
+ assert result == {"sub": str(user.id)}
@pytest.mark.db_test
@mock.patch.object(FabAuthManager, "get_user")
diff --git a/tests/api_fastapi/auth/managers/test_base_auth_manager.py b/tests/api_fastapi/auth/managers/test_base_auth_manager.py
index 557350d78b7..033bc222ad4 100644
--- a/tests/api_fastapi/auth/managers/test_base_auth_manager.py
+++ b/tests/api_fastapi/auth/managers/test_base_auth_manager.py
@@ -17,7 +17,7 @@
from __future__ import annotations
from typing import TYPE_CHECKING, Any
-from unittest.mock import MagicMock, Mock, patch
+from unittest.mock import AsyncMock, MagicMock, Mock, patch
import pytest
@@ -168,19 +168,20 @@ class TestBaseAuthManager:
@patch("airflow.api_fastapi.auth.managers.base_auth_manager.JWTValidator", autospec=True)
@patch.object(EmptyAuthManager, "deserialize_user")
- def test_get_user_from_token(self, mock_deserialize_user, mock_jwt_validator, auth_manager):
+ @pytest.mark.asyncio
+ async def test_get_user_from_token(self, mock_deserialize_user, mock_jwt_validator, auth_manager):
token = "token"
payload = {}
user = BaseAuthManagerUserTest(name="test")
- signer = Mock(spec=JWTValidator)
- signer.validated_claims.return_value = payload
+ signer = AsyncMock(spec=JWTValidator)
+ signer.avalidated_claims.return_value = payload
mock_jwt_validator.return_value = signer
mock_deserialize_user.return_value = user
- result = auth_manager.get_user_from_token(token)
+ result = await auth_manager.get_user_from_token(token)
mock_deserialize_user.assert_called_once_with(payload)
- signer.validated_claims.assert_called_once_with(token)
+ signer.avalidated_claims.assert_called_once_with(token)
assert result == user
@patch("airflow.api_fastapi.auth.managers.base_auth_manager.JWTGenerator", autospec=True)
diff --git a/tests/api_fastapi/auth/test_tokens.py b/tests/api_fastapi/auth/test_tokens.py
index b30e01fb3f8..02d12996d14 100644
--- a/tests/api_fastapi/auth/test_tokens.py
+++ b/tests/api_fastapi/auth/test_tokens.py
@@ -196,7 +196,7 @@ async def test_jwt_wrong_subject(jwt_generator, jwt_validator):
wrong_subject = jwt_generator.generate({"sub": "wrong_subject"})
with pytest.raises(InvalidClaimError, match="Invalid claim: sub"):
await jwt_validator.avalidated_claims(
- wrong_subject, extra_claims={"sub": {"essential": True, "value": "test_subject"}}
+ wrong_subject, required_claims={"sub": {"essential": True, "value": "test_subject"}}
)
diff --git a/tests/api_fastapi/conftest.py b/tests/api_fastapi/conftest.py
index 9d3773018c7..0b26b11ed10 100644
--- a/tests/api_fastapi/conftest.py
+++ b/tests/api_fastapi/conftest.py
@@ -53,7 +53,7 @@ def test_client():
# to make the JWT token always valid for all test cases with time_machine
time_very_before = datetime.datetime(2014, 1, 1, 0, 0, 0)
time_very_after = datetime.datetime.now() + datetime.timedelta(days=1)
- token = auth_manager._get_token_signer().generate_signed_token(
+ token = auth_manager._get_token_signer().generate(
{
"iat": time_very_before,
"nbf": time_very_before,
diff --git a/tests/api_fastapi/core_api/test_security.py b/tests/api_fastapi/core_api/test_security.py
index 84328e650f7..736afa6050a 100644
--- a/tests/api_fastapi/core_api/test_security.py
+++ b/tests/api_fastapi/core_api/test_security.py
@@ -16,7 +16,7 @@
# under the License.
from __future__ import annotations
-from unittest.mock import Mock, patch
+from unittest.mock import AsyncMock, Mock, patch
import pytest
from fastapi import HTTPException
@@ -30,6 +30,7 @@ from airflow.api_fastapi.core_api.security import get_user, is_safe_url, require
from tests_common.test_utils.config import conf_vars
+@pytest.mark.asyncio
class TestFastApiSecurity:
@classmethod
def setup_class(cls):
@@ -44,47 +45,47 @@ class TestFastApiSecurity:
create_app()
@patch("airflow.api_fastapi.core_api.security.get_auth_manager")
- def test_get_user(self, mock_get_auth_manager):
+ async def test_get_user(self, mock_get_auth_manager):
token_str = "test-token"
user = SimpleAuthManagerUser(username="username", role="admin")
- auth_manager = Mock()
+ auth_manager = AsyncMock()
auth_manager.get_user_from_token.return_value = user
mock_get_auth_manager.return_value = auth_manager
- result = get_user(token_str)
+ result = await get_user(token_str)
auth_manager.get_user_from_token.assert_called_once_with(token_str)
assert result == user
@patch("airflow.api_fastapi.core_api.security.get_auth_manager")
- def test_get_user_wrong_token(self, mock_get_auth_manager):
+ async def test_get_user_wrong_token(self, mock_get_auth_manager):
token_str = "test-token"
- auth_manager = Mock()
+ auth_manager = AsyncMock()
auth_manager.get_user_from_token.side_effect = InvalidTokenError()
mock_get_auth_manager.return_value = auth_manager
with pytest.raises(HTTPException, match="Forbidden"):
- get_user(token_str)
+ await get_user(token_str)
auth_manager.get_user_from_token.assert_called_once_with(token_str)
@patch("airflow.api_fastapi.core_api.security.get_auth_manager")
- def test_get_user_expired_token(self, mock_get_auth_manager):
+ async def test_get_user_expired_token(self, mock_get_auth_manager):
token_str = "test-token"
- auth_manager = Mock()
+ auth_manager = AsyncMock()
auth_manager.get_user_from_token.side_effect = ExpiredSignatureError()
mock_get_auth_manager.return_value = auth_manager
with pytest.raises(HTTPException, match="Token Expired"):
- get_user(token_str)
+ await get_user(token_str)
auth_manager.get_user_from_token.assert_called_once_with(token_str)
@patch("airflow.api_fastapi.core_api.security.get_auth_manager")
- def test_requires_access_dag_authorized(self, mock_get_auth_manager):
+ async def test_requires_access_dag_authorized(self, mock_get_auth_manager):
auth_manager = Mock()
auth_manager.is_authorized_dag.return_value = True
mock_get_auth_manager.return_value = auth_manager
@@ -96,7 +97,7 @@ class TestFastApiSecurity:
auth_manager.is_authorized_dag.assert_called_once()
@patch("airflow.api_fastapi.core_api.security.get_auth_manager")
- def test_requires_access_dag_unauthorized(self, mock_get_auth_manager):
+ async def test_requires_access_dag_unauthorized(self, mock_get_auth_manager):
auth_manager = Mock()
auth_manager.is_authorized_dag.return_value = False
mock_get_auth_manager.return_value = auth_manager
diff --git a/tests/utils/test_serve_logs.py b/tests/utils/test_serve_logs.py
index a4a5b3b0ec5..ad2d6ebf7c3 100644
--- a/tests/utils/test_serve_logs.py
+++ b/tests/utils/test_serve_logs.py
@@ -39,7 +39,12 @@ LOG_DATA = "Airflow log data" * 20
@pytest.fixture
def client_without_config(tmp_path):
- with conf_vars({("logging", "base_log_folder"): tmp_path.as_posix()}):
+ with conf_vars(
+ {
+ ("logging", "base_log_folder"): tmp_path.as_posix(),
+ ("api_auth", "jwt_leeyway"): "0",
+ }
+ ):
app = create_app()
yield app.test_client()
@@ -52,7 +57,8 @@ def client_with_config():
(
"logging",
"logging_config_class",
- ): "airflow.config_templates.airflow_local_settings.DEFAULT_LOGGING_CONFIG"
+ ): "airflow.config_templates.airflow_local_settings.DEFAULT_LOGGING_CONFIG",
+ ("api_auth", "jwt_leeyway"): "0",
}
):
app = create_app()
@@ -85,7 +91,7 @@ def sample_log(request, tmp_path):
def jwt_generator(secret_key):
return JWTGenerator(
secret_key=secret_key,
- valid_for=30,
+ valid_for=5,
audience="task-instance-logs",
)
@@ -108,7 +114,7 @@ class TestServeLogs:
response = client.get(
"/log/sample.log",
headers={
- "Authorization": jwt_generator.generate("logs", {"filename": "sample.log"}),
+ "Authorization": jwt_generator.generate({"filename": "sample.log"}),
},
)
assert response.data.decode() == LOG_DATA
@@ -118,14 +124,14 @@ class TestServeLogs:
response = client.get(
"/log/sample.log",
headers={
- "Authorization": jwt_generator.generate("logs", {"filename": "different.log"}),
+ "Authorization": jwt_generator.generate({"filename": "different.log"}),
},
)
assert response.status_code == 403
def test_forbidden_expired(self, client: FlaskClient, jwt_generator):
with time_machine.travel("2010-01-14"):
- token = jwt_generator.generate("logs", {"filename": "sample.log"})
+ token = jwt_generator.generate({"filename": "sample.log"})
assert (
client.get(
"/log/sample.log",
@@ -138,7 +144,7 @@ class TestServeLogs:
def test_forbidden_future(self, client: FlaskClient, jwt_generator):
with time_machine.travel(timezone.utcnow() + timedelta(seconds=3600)):
- token = jwt_generator.generate("logs", {"filename": "sample.log"})
+ token = jwt_generator.generate({"filename": "sample.log"})
assert (
client.get(
"/log/sample.log",
@@ -153,7 +159,7 @@ class TestServeLogs:
print(f"Ts= {timezone.utcnow().timestamp()}")
with time_machine.travel(timezone.utcnow() + timedelta(seconds=1)):
print(f"Ts with travvel = {timezone.utcnow().timestamp()}")
- token = jwt_generator.generate("logs", {"filename": "sample.log"})
+ token = jwt_generator.generate({"filename": "sample.log"})
assert (
client.get(
"/log/sample.log",
@@ -166,7 +172,7 @@ class TestServeLogs:
def test_ok_with_short_past_skew(self, client: FlaskClient, jwt_generator):
with time_machine.travel(timezone.utcnow() - timedelta(seconds=31)):
- token = jwt_generator.generate("logs", {"filename": "sample.log"})
+ token = jwt_generator.generate({"filename": "sample.log"})
assert (
client.get(
"/log/sample.log",
@@ -178,8 +184,8 @@ class TestServeLogs:
)
def test_forbidden_with_long_future_skew(self, client: FlaskClient, jwt_generator):
- with time_machine.travel(timezone.utcnow() + timedelta(seconds=10)):
- token = jwt_generator.generate("logs", {"filename": "sample.log"})
+ with time_machine.travel(timezone.utcnow() + timedelta(seconds=40)):
+ token = jwt_generator.generate({"filename": "sample.log"})
assert (
client.get(
"/log/sample.log",
@@ -192,7 +198,7 @@ class TestServeLogs:
def test_forbidden_with_long_past_skew(self, client: FlaskClient, jwt_generator):
with time_machine.travel(timezone.utcnow() - timedelta(seconds=40)):
- token = jwt_generator.generate("logs", {"filename": "sample.log"})
+ token = jwt_generator.generate({"filename": "sample.log"})
assert (
client.get(
"/log/sample.log",
@@ -208,7 +214,7 @@ class TestServeLogs:
client.get(
"/log/sample.log",
headers={
- "Authorization": different_audience.generate("logs", {"filename": "sample.log"}),
+ "Authorization": different_audience.generate({"filename": "sample.log"}),
},
).status_code
== 403