| 171 | |
| 172 | |
| 173 | class WorkloadIdentityAuth: |
| 174 | def __init__( |
| 175 | self, |
| 176 | *, |
| 177 | workload_identity: WorkloadIdentity, |
| 178 | token_exchange_url: str = DEFAULT_TOKEN_EXCHANGE_URL, |
| 179 | ): |
| 180 | self.workload_identity = workload_identity |
| 181 | self.token_exchange_url = token_exchange_url |
| 182 | |
| 183 | self._cached_token: str | None = None |
| 184 | self._cached_token_expires_at_monotonic: float | None = None |
| 185 | self._cached_token_refresh_at_monotonic: float | None = None |
| 186 | self._refreshing: bool = False |
| 187 | self._lock = threading.Lock() |
| 188 | self._condition = threading.Condition(self._lock) |
| 189 | |
| 190 | def get_token(self) -> str: |
| 191 | with self._lock: |
| 192 | while self._refreshing and self._token_unusable(): |
| 193 | self._condition.wait() |
| 194 | |
| 195 | if not self._token_unusable() and not self._needs_refresh(): |
| 196 | return cast(str, self._cached_token) |
| 197 | |
| 198 | if self._refreshing: |
| 199 | while self._refreshing: |
| 200 | self._condition.wait() |
| 201 | token = self._cached_token # type: ignore[unreachable] |
| 202 | if self._token_unusable(): |
| 203 | raise RuntimeError("Token is unusable after refresh completed") |
| 204 | return cast(str, token) |
| 205 | |
| 206 | self._refreshing = True |
| 207 | |
| 208 | try: |
| 209 | self._perform_refresh() |
| 210 | with self._lock: |
| 211 | if self._token_unusable(): |
| 212 | raise RuntimeError("Token is unusable after refresh completed") |
| 213 | return cast(str, self._cached_token) |
| 214 | finally: |
| 215 | with self._lock: |
| 216 | self._refreshing = False |
| 217 | self._condition.notify_all() |
| 218 | |
| 219 | async def get_token_async(self) -> str: |
| 220 | return await to_thread(self.get_token) |
| 221 | |
| 222 | def invalidate_token(self) -> None: |
| 223 | with self._lock: |
| 224 | self._cached_token = None |
| 225 | self._cached_token_expires_at_monotonic = None |
| 226 | self._cached_token_refresh_at_monotonic = None |
| 227 | |
| 228 | def _perform_refresh(self) -> None: |
| 229 | token_data = self._fetch_token_from_exchange() |
| 230 | now = time.monotonic() |