jason810496 commented on code in PR #61339:
URL: https://github.com/apache/airflow/pull/61339#discussion_r2753115355
##########
airflow-core/src/airflow/api_fastapi/core_api/routes/public/auth.py:
##########
@@ -16,15 +16,22 @@
# under the License.
from __future__ import annotations
+import logging
Review Comment:
```suggestion
import structlog
```
##########
airflow-core/src/airflow/api_fastapi/core_api/routes/public/auth.py:
##########
@@ -55,6 +62,25 @@ def logout(request: Request, auth_manager: AuthManagerDep)
-> RedirectResponse:
if logout_url:
return RedirectResponse(logout_url)
+ # Revoke the current token before deleting the cookie
+ token_str = request.cookies.get(COOKIE_NAME_JWT_TOKEN)
+ if token_str:
+ try:
+ header = jwt.get_unverified_header(token_str)
+ claims = jwt.decode(
+ token_str,
+ options={"verify_exp": False, "verify_signature": False},
+ algorithms=[header.get("alg", "HS256")],
+ )
+ jti = claims.get("jti")
+ exp = claims.get("exp")
+ if jti and exp:
+ from airflow.models.revoked_token import RevokedToken
+
+ RevokedToken.revoke(jti, datetime.fromtimestamp(exp,
tz=timezone.utc))
+ except (jwt.InvalidTokenError, SQLAlchemyError):
+ log.warning("Failed to revoke token during logout", exc_info=True)
Review Comment:
So maybe we could move this part as a new method of `JWTValidator`.
##########
airflow-core/src/airflow/api_fastapi/core_api/routes/public/auth.py:
##########
@@ -55,6 +62,25 @@ def logout(request: Request, auth_manager: AuthManagerDep)
-> RedirectResponse:
if logout_url:
return RedirectResponse(logout_url)
+ # Revoke the current token before deleting the cookie
+ token_str = request.cookies.get(COOKIE_NAME_JWT_TOKEN)
+ if token_str:
Review Comment:
```suggestion
if(token_str := request.cookies.get(COOKIE_NAME_JWT_TOKEN)):
```
##########
airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_auth.py:
##########
@@ -100,3 +101,39 @@ def test_should_respond_307(
if delete_cookies:
cookies = response.headers.get_list("set-cookie")
assert any(f"{COOKIE_NAME_JWT_TOKEN}=" in c for c in cookies)
+
Review Comment:
We could separate the new tests to `TestLogout` class.
```suggestion
class TestLogout(TestAuthEndpoint):
```
##########
airflow-core/src/airflow/api_fastapi/auth/managers/base_auth_manager.py:
##########
@@ -139,6 +139,12 @@ async def get_user_from_token(self, token: str) ->
BaseUser:
log.error("JWT token is not valid: %s", e)
raise e
+ from airflow.models.revoked_token import RevokedToken
Review Comment:
It's fine to move to top level as `RevokedToken` will be used for each
request, so we don't need `RevokedToken` at all.
##########
airflow-core/src/airflow/api_fastapi/auth/managers/base_auth_manager.py:
##########
@@ -139,6 +139,12 @@ async def get_user_from_token(self, token: str) ->
BaseUser:
log.error("JWT token is not valid: %s", e)
raise e
+ from airflow.models.revoked_token import RevokedToken
+
+ jti = payload.get("jti")
+ if jti and RevokedToken.is_revoked(jti):
+ raise InvalidTokenError("Token has been revoked")
Review Comment:
```suggestion
if (jti := payload.get("jti")) and RevokedToken.is_revoked(jti):
raise InvalidTokenError("Token has been revoked")
```
##########
airflow-core/src/airflow/api_fastapi/core_api/routes/public/auth.py:
##########
@@ -16,15 +16,22 @@
# under the License.
from __future__ import annotations
+import logging
+from datetime import datetime, timezone
+
+import jwt
from fastapi import HTTPException, Request, status
from fastapi.responses import RedirectResponse
+from sqlalchemy.exc import SQLAlchemyError
from airflow.api_fastapi.auth.managers.base_auth_manager import
COOKIE_NAME_JWT_TOKEN
from airflow.api_fastapi.common.router import AirflowRouter
from airflow.api_fastapi.core_api.openapi.exceptions import
create_openapi_http_exception_doc
from airflow.api_fastapi.core_api.security import AuthManagerDep, is_safe_url
from airflow.configuration import conf
+log = logging.getLogger(__name__)
Review Comment:
```suggestion
log = structlog.getLogger(__name__)
```
##########
airflow-core/src/airflow/api_fastapi/core_api/routes/public/auth.py:
##########
@@ -55,6 +62,25 @@ def logout(request: Request, auth_manager: AuthManagerDep)
-> RedirectResponse:
if logout_url:
return RedirectResponse(logout_url)
+ # Revoke the current token before deleting the cookie
+ token_str = request.cookies.get(COOKIE_NAME_JWT_TOKEN)
+ if token_str:
+ try:
+ header = jwt.get_unverified_header(token_str)
+ claims = jwt.decode(
+ token_str,
+ options={"verify_exp": False, "verify_signature": False},
+ algorithms=[header.get("alg", "HS256")],
+ )
+ jti = claims.get("jti")
+ exp = claims.get("exp")
Review Comment:
It seems to be duplicated of what `JWTValidator.validated_claims` has
implemented.
##########
airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_auth.py:
##########
@@ -100,3 +101,39 @@ def test_should_respond_307(
if delete_cookies:
cookies = response.headers.get_list("set-cookie")
assert any(f"{COOKIE_NAME_JWT_TOKEN}=" in c for c in cookies)
+
+ @patch("airflow.models.revoked_token.RevokedToken")
+ def test_logout_blacklists_token(self, mock_revoked_token, test_client):
+ """Test that logout blacklists the JWT token's jti."""
+ test_client.app.state.auth_manager.get_url_logout.return_value = None
+ token_payload = {"sub": "admin", "jti": "test-jti-123", "exp":
9999999999}
+ token_str = jwt.encode(token_payload, "secret", algorithm="HS256")
+
+ test_client.cookies.set(COOKIE_NAME_JWT_TOKEN, token_str)
+ response = test_client.get("/auth/logout", follow_redirects=False)
+
+ assert response.status_code == 307
+ mock_revoked_token.revoke.assert_called_once()
+ call_args = mock_revoked_token.revoke.call_args
+ assert call_args[0][0] == "test-jti-123"
+
+ @patch("airflow.models.revoked_token.RevokedToken")
+ def test_logout_without_cookie_does_not_blacklist(self,
mock_revoked_token, test_client):
+ """Test that logout without a cookie does not attempt to blacklist."""
+ test_client.app.state.auth_manager.get_url_logout.return_value = None
+
+ response = test_client.get("/auth/logout", follow_redirects=False)
+
+ assert response.status_code == 307
+ mock_revoked_token.revoke.assert_not_called()
+
+ @patch("airflow.models.revoked_token.RevokedToken")
+ def test_logout_with_malformed_cookie_does_not_blacklist(self,
mock_revoked_token, test_client):
+ """Test that logout with a malformed cookie does not raise and does
not blacklist."""
+ test_client.app.state.auth_manager.get_url_logout.return_value = None
+ test_client.cookies.set(COOKIE_NAME_JWT_TOKEN, "not-a-valid-jwt")
+
+ response = test_client.get("/auth/logout", follow_redirects=False)
+
+ assert response.status_code == 307
+ mock_revoked_token.revoke.assert_not_called()
Review Comment:
Additionally, it would be nice to add test cases that don't mock the
`airflow.models.revoked_token.RevokedToken` and check the real behavior of with
DB, so we need to add `clear_db_revoke_token` in
`devel-common/src/tests_common/test_utils/db.py`.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
