Coverage for /home/airflow/.local/lib/python3.12/site-packages/airflow/api_fastapi/execution_api/security.py: 68%

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

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

3# distributed with this work for additional information 

4# regarding copyright ownership. The ASF licenses this file 

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

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

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

8# 

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

10# 

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

12# software distributed under the License is distributed on an 

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

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

15# specific language governing permissions and limitations 

16# under the License. 

17 

18""" 

19Execution API security: JWT validation, token scopes, and route-level access control. 

20 

21Token types (``TokenType``): 

22 

23``"execution"`` 

24 Default scope, accepted by all endpoints. Short-lived, automatically 

25 refreshed by ``JWTReissueMiddleware``. 

26 

27``"workload"`` 

28 Restricted scope, only accepted on routes that opt in via 

29 ``Security(require_auth, scopes=["token:workload"])``. 

30 

31Tokens without a ``scope`` claim default to ``"execution"`` for backwards 

32compatibility (``claims.setdefault("scope", "execution")``). 

33 

34Enforcement flow: 

35 1. ``JWTBearer.__call__`` validates the JWT once per request (crypto + 

36 signature verification), caching the result on the ASGI request scope. 

37 Subsequent FastAPI dependency resolutions and Cadwyn replays return 

38 the cache. 

39 2. ``require_auth`` is the Security dependency on routers. It receives 

40 the token from ``JWTBearer`` and enforces: 

41 - Token type against the route's ``allowed_token_types`` (precomputed 

42 by ``ExecutionAPIRoute`` from ``token:*`` Security scopes). 

43 - ``ti:self`` scope — checks that the JWT ``sub`` matches the 

44 ``{task_instance_id}`` path parameter. 

45 3. ``ExecutionAPIRoute`` precomputes ``allowed_token_types`` from 

46 ``token:*`` Security scopes at route registration time. Routes 

47 without explicit ``token:*`` scopes default to execution-only. 

48 

49Why ``ExecutionAPIRoute`` is needed: 

50 FastAPI resolves router-level ``Security()`` dependencies from outermost 

51 to innermost. A ``token:workload`` scope on an inner endpoint would need 

52 to *relax* the outer router's default execution-only restriction, but 

53 ``SecurityScopes`` only accumulate additively — an outer dependency 

54 cannot see scopes declared by inner ones. ``ExecutionAPIRoute`` solves 

55 this by inspecting the **merged** dependency list at route registration 

56 time (after ``include_router`` has combined all parent and child 

57 dependencies) and precomputing the full ``allowed_token_types`` set. 

58 ``require_auth`` then reads this precomputed set from the matched route 

59 at request time, avoiding the ordering problem entirely. 

60 

61 Any router whose routes need non-default token type policies must use 

62 ``route_class=ExecutionAPIRoute``. Routers that only need the default 

63 (execution-only) can use the standard route class — ``require_auth`` 

64 falls back to ``{"execution"}`` when the attribute is absent. 

65""" 

66 

67# Disable future annotations in this file to work around https://github.com/fastapi/fastapi/issues/13056 

68# ruff: noqa: I002 

69 

70from typing import Any, get_args 

71 

72import structlog 

73from fastapi import Depends, HTTPException, Request, status 

74from fastapi.params import Security as SecurityParam 

75from fastapi.routing import APIRoute 

76from fastapi.security import HTTPBearer, SecurityScopes 

77from pydantic import ValidationError 

78from sqlalchemy import select 

79 

80from airflow.api_fastapi.auth.tokens import JWTValidator 

81from airflow.api_fastapi.execution_api.datamodels.token import TIClaims, TIToken, TokenScope 

82from airflow.api_fastapi.execution_api.deps import DepContainer 

83 

84log = structlog.get_logger(logger_name=__name__) 

85 

86VALID_TOKEN_TYPES: frozenset[str] = frozenset(get_args(TokenScope)) 

87 

88_REQUEST_SCOPE_TOKEN_KEY = "ti_token" 

89 

90 

91class JWTBearer(HTTPBearer): 

92 """ 

93 Validates JWT tokens for the Execution API. 

94 

95 Performs cryptographic validation once per request and caches the result 

96 on the ASGI request scope. Subsequent resolutions (FastAPI dependency 

97 dedup or Cadwyn replays) return the cached token. 

98 

99 This dependency handles ONLY crypto validation and token construction. 

100 All route-specific authorization (token type, ti:self) is handled by 

101 ``require_auth``. 

102 """ 

103 

104 def __init__(self, required_claims: dict[str, Any] | None = None): 

105 super().__init__(auto_error=False) 

106 self.required_claims = required_claims or {} 

107 

108 async def __call__( # type: ignore[override] 

109 self, 

110 request: Request, 

111 services=DepContainer, 

112 ) -> TIToken | None: 

113 # Return cached token (handles both FastAPI dependency dedup and Cadwyn replays). 

114 if cached := request.scope.get(_REQUEST_SCOPE_TOKEN_KEY): 

115 return cached 

116 

117 # First resolution — full cryptographic validation. 

118 creds = await super().__call__(request) 

119 if not creds: 119 ↛ 120line 119 didn't jump to line 120 because the condition on line 119 was never true

120 raise HTTPException(status_code=status.HTTP_401_UNAUTHORIZED, detail="Missing auth token") 

121 

122 validator: JWTValidator = await services.aget(JWTValidator) 

123 

124 try: 

125 claims = await validator.avalidated_claims(creds.credentials, dict(self.required_claims)) 

126 except Exception: 

127 log.warning("Failed to validate JWT", exc_info=True) 

128 raise HTTPException(status_code=status.HTTP_403_FORBIDDEN, detail="Invalid auth token") 

129 

130 claims.setdefault("scope", "execution") 

131 

132 try: 

133 claim_model = TIClaims(**claims) 

134 except ValidationError as err: 

135 log.warning("JWT claims did not match task identity token schema", exc_info=True) 

136 raise HTTPException(status_code=status.HTTP_403_FORBIDDEN, detail=f"Invalid auth token: {err}") 

137 

138 token = TIToken(id=claims["sub"], claims=claim_model) 

139 request.scope[_REQUEST_SCOPE_TOKEN_KEY] = token 

140 return token 

141 

142 

143_jwt_bearer = JWTBearer() 

144 

145 

146async def require_auth( 

147 security_scopes: SecurityScopes, 

148 request: Request, 

149 token: TIToken = Depends(_jwt_bearer), 

150) -> TIToken: 

151 """ 

152 Security dependency that enforces token type and ``ti:self`` scope. 

153 

154 Used via ``Security(require_auth)`` on routers. ``SecurityScopes`` are 

155 accumulated by FastAPI from all parent ``Security()`` declarations. 

156 

157 Token type enforcement reads ``route.allowed_token_types`` (precomputed 

158 by ``ExecutionAPIRoute``) or defaults to ``{"execution"}``. 

159 """ 

160 token_scope = token.claims.scope 

161 

162 if token_scope not in VALID_TOKEN_TYPES: 162 ↛ 163line 162 didn't jump to line 163 because the condition on line 162 was never true

163 log.warning("Invalid token scope in claims", token_scope=token_scope, path=request.url.path) 

164 raise HTTPException( 

165 status_code=status.HTTP_403_FORBIDDEN, 

166 detail=f"Invalid token scope: {token_scope}", 

167 ) 

168 

169 route = request.scope.get("route") 

170 allowed_token_types = getattr(route, "allowed_token_types", frozenset({"execution"})) 

171 

172 if token_scope not in allowed_token_types: 172 ↛ 173line 172 didn't jump to line 173 because the condition on line 172 was never true

173 log.warning( 

174 "Token type not allowed for endpoint", 

175 token_scope=token_scope, 

176 allowed_types=sorted(allowed_token_types), 

177 path=request.url.path, 

178 ) 

179 raise HTTPException( 

180 status_code=status.HTTP_403_FORBIDDEN, 

181 detail=f"Token type '{token_scope}' not allowed for this endpoint. " 

182 f"Allowed types: {', '.join(sorted(allowed_token_types))}", 

183 ) 

184 

185 if "ti:self" in security_scopes.scopes: 

186 ti_self_id = str(request.path_params["task_instance_id"]) 

187 if str(token.id) != ti_self_id: 187 ↛ 188line 187 didn't jump to line 188 because the condition on line 187 was never true

188 raise HTTPException( 

189 status_code=status.HTTP_403_FORBIDDEN, 

190 detail="Token subject does not match task instance ID", 

191 ) 

192 elif "ct:self" in security_scopes.scopes: 192 ↛ 193line 192 didn't jump to line 193 because the condition on line 192 was never true

193 ct_self_id = str(request.path_params["connection_test_id"]) 

194 if str(token.id) != ct_self_id: 

195 raise HTTPException( 

196 status_code=status.HTTP_403_FORBIDDEN, 

197 detail="Token subject does not match connection test ID", 

198 ) 

199 

200 return token 

201 

202 

203CurrentTIToken: TIToken = Depends(require_auth) 

204 

205 

206class ExecutionAPIRoute(APIRoute): 

207 """ 

208 Custom route class that precomputes allowed token types from Security scopes. 

209 

210 Scopes prefixed with ``token:`` (e.g., ``token:execution``, ``token:workload``) 

211 are extracted at route registration time and stored as ``allowed_token_types``. 

212 If no ``token:*`` scopes are declared, defaults to ``{"execution"}``. 

213 

214 ``require_auth`` reads ``route.allowed_token_types`` at request time. 

215 """ 

216 

217 allowed_token_types: frozenset[str] 

218 

219 def __init__(self, *args: Any, **kwargs: Any) -> None: 

220 super().__init__(*args, **kwargs) 

221 

222 all_scopes: set[str] = set() 

223 for dep in self.dependencies: 

224 if isinstance(dep, SecurityParam): 

225 all_scopes.update(dep.scopes or []) 

226 

227 token_scopes = {s.removeprefix("token:") for s in all_scopes if s.startswith("token:")} 

228 

229 if token_scopes and not token_scopes <= VALID_TOKEN_TYPES: 229 ↛ 230line 229 didn't jump to line 230 because the condition on line 229 was never true

230 invalid = token_scopes - VALID_TOKEN_TYPES 

231 raise ValueError(f"Invalid token types in Security scopes: {invalid}") 

232 

233 self.allowed_token_types = frozenset(token_scopes) if token_scopes else frozenset({"execution"}) 

234 

235 

236async def get_team_name_dep(token=CurrentTIToken) -> str | None: 

237 """Return the team name associated to the task (if any).""" 

238 from airflow.configuration import conf 

239 

240 if not conf.getboolean("core", "multi_team"): 240 ↛ 243line 240 didn't jump to line 243 because the condition on line 240 was always true

241 return None 

242 

243 from airflow.utils.session import create_session_async 

244 

245 async with create_session_async() as session: 

246 return await session.scalar(_team_name_for_ti_stmt(token.id)) 

247 

248 

249def get_team_name_for_ti(ti_id, session) -> str | None: 

250 """ 

251 Return the team name associated to the task (if any), using a sync session. 

252 

253 Sync counterpart to :func:`get_team_name_dep` for callers that already hold a 

254 SQLAlchemy session (e.g., the ``ti_run`` endpoint). No-op when multi-team is disabled. 

255 """ 

256 from airflow.configuration import conf 

257 

258 if not conf.getboolean("core", "multi_team"): 258 ↛ 260line 258 didn't jump to line 260 because the condition on line 258 was always true

259 return None 

260 return session.scalar(_team_name_for_ti_stmt(ti_id)) 

261 

262 

263def _team_name_for_ti_stmt(ti_id): 

264 """Build the select statement resolving ``TaskInstance.id -> Team.name``.""" 

265 from airflow.models import DagModel, TaskInstance 

266 from airflow.models.dagbundle import DagBundleModel 

267 from airflow.models.team import Team 

268 

269 return ( 

270 select(Team.name) 

271 .select_from(TaskInstance) 

272 .join(DagModel, DagModel.dag_id == TaskInstance.dag_id) 

273 .join(DagBundleModel, DagBundleModel.name == DagModel.bundle_name) 

274 .join(DagBundleModel.teams) 

275 .where(TaskInstance.id == ti_id) 

276 ) 

277 

278 

279def _team_name_for_dag_stmt(dag_id): 

280 """Build the select statement resolving ``DagModel.dag_id -> Team.name``.""" 

281 from airflow.models import DagModel 

282 from airflow.models.dagbundle import DagBundleModel 

283 from airflow.models.team import Team 

284 

285 return ( 

286 select(Team.name) 

287 .select_from(DagModel) 

288 .join(DagBundleModel, DagBundleModel.name == DagModel.bundle_name) 

289 .join(DagBundleModel.teams) 

290 .where(DagModel.dag_id == dag_id) 

291 )