Coverage for /home/airflow/.local/lib/python3.12/site-packages/airflow/api_fastapi/auth/middlewares/refresh_token.py: 49%

1# 

2# Licensed to the Apache Software Foundation (ASF) under one 

3# or more contributor license agreements. See the NOTICE file 

4# distributed with this work for additional information 

5# regarding copyright ownership. The ASF licenses this file 

6# to you under the Apache License, Version 2.0 (the 

7# "License"); you may not use this file except in compliance 

8# with the License. You may obtain a copy of the License at 

9# 

10# http://www.apache.org/licenses/LICENSE-2.0 

11# 

12# Unless required by applicable law or agreed to in writing, 

13# software distributed under the License is distributed on an 

14# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY 

15# KIND, either express or implied. See the License for the 

16# specific language governing permissions and limitations 

17# under the License. 

18from __future__ import annotations 

19 

20from fastapi import HTTPException, Request 

21from fastapi.responses import JSONResponse 

22from starlette.middleware.base import BaseHTTPMiddleware 

23 

24from airflow.api_fastapi.app import get_auth_manager, get_cookie_path 

25from airflow.api_fastapi.auth.managers.base_auth_manager import COOKIE_NAME_JWT_TOKEN 

26from airflow.api_fastapi.auth.managers.exceptions import AuthManagerRefreshTokenExpiredException 

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

28from airflow.api_fastapi.core_api.security import ( 

29 USER_INJECTED_BY_TRUSTED_MIDDLEWARE, 

30 resolve_user_from_token, 

31) 

32from airflow.configuration import conf 

33 

34 

35class JWTRefreshMiddleware(BaseHTTPMiddleware): 

36 """ 

37 Middleware to handle JWT token refresh. 

38 

39 This middleware: 

40 1. Extracts JWT token from cookies and build the user from the token 

41 2. Calls ``refresh_user`` method from auth manager with the user 

42 3. If ``refresh_user`` returns a user, generate a JWT token based upon this user and send it in the 

43 response as cookie 

44 """ 

45 

46 async def dispatch(self, request: Request, call_next): 

47 new_token = None 

48 current_token = request.cookies.get(COOKIE_NAME_JWT_TOKEN) 

49 try: 

50 if current_token is not None: 50 ↛ 51line 50 didn't jump to line 51 because the condition on line 50 was never true

51 try: 

52 new_user, current_user = await self._refresh_user(current_token) 

53 if user := (new_user or current_user): 

54 # Stamp the trust sentinel alongside the user so `get_user()` 

55 # can distinguish this trusted assignment from a stray write 

56 # by unrelated middleware. 

57 request.state.user = user 

58 request.state.user_authenticated_via = USER_INJECTED_BY_TRUSTED_MIDDLEWARE 

59 if new_user: 

60 # If we created a new user, serialize it and set it as a cookie 

61 new_token = get_auth_manager().generate_jwt(new_user) 

62 except (HTTPException, AuthManagerRefreshTokenExpiredException): 

63 # Receive a HTTPException when the Airflow token is expired 

64 # Receive a AuthManagerRefreshTokenExpiredException when the potential underlying refresh 

65 # token used by the auth manager is expired 

66 new_token = "" 

67 

68 response = await call_next(request) 

69 

70 if new_token is not None: 70 ↛ 71line 70 didn't jump to line 71 because the condition on line 70 was never true

71 cookie_path = get_cookie_path() 

72 secure = request.base_url.scheme == "https" or bool(conf.get("api", "ssl_cert", fallback="")) 

73 response.set_cookie( 

74 COOKIE_NAME_JWT_TOKEN, 

75 new_token, 

76 path=cookie_path, 

77 httponly=True, 

78 secure=secure, 

79 samesite="lax", 

80 max_age=0 if new_token == "" else None, 

81 ) 

82 # Clear any stale _token cookie at root path "/". 

83 # Older Airflow instances may have set the cookie there; 

84 # without this, the root-path cookie keeps being sent on 

85 # every request, causing an infinite redirect loop. 

86 if cookie_path != "/": 

87 response.delete_cookie( 

88 key=COOKIE_NAME_JWT_TOKEN, 

89 path="/", 

90 httponly=True, 

91 secure=secure, 

92 samesite="lax", 

93 ) 

94 except HTTPException as exc: 

95 # If any HTTPException is raised during user resolution or refresh, return it as response 

96 return JSONResponse(status_code=exc.status_code, content={"detail": exc.detail}) 

97 return response 

98 

99 @staticmethod 

100 async def _refresh_user(current_token: str) -> tuple[BaseUser | None, BaseUser | None]: 

101 user = await resolve_user_from_token(current_token) 

102 return get_auth_manager().refresh_user(user=user), user