diff --git a/cycode/cyclient/base_token_auth_client.py b/cycode/cyclient/base_token_auth_client.py index ec315e7d..cdadf98d 100644 --- a/cycode/cyclient/base_token_auth_client.py +++ b/cycode/cyclient/base_token_auth_client.py @@ -1,3 +1,5 @@ +import secrets +import time from abc import ABC, abstractmethod from threading import Lock from typing import Any, Optional @@ -5,6 +7,7 @@ import arrow from requests import Response +from cycode.cli.exceptions.custom_exceptions import HttpUnauthorizedError from cycode.cli.user_settings.credentials_manager import CredentialsManager from cycode.cli.user_settings.jwt_creator import JwtCreator from cycode.cyclient.cycode_client import CycodeClient @@ -15,6 +18,11 @@ b'JWT Token validation failed', ] +# Identity provider brute-force protection rejects logins for the same user that land within +# milliseconds of each other, so when several processes mint at once all but one are refused. +_MINT_CONFLICT_RETRY_MIN_MS = 50 +_MINT_CONFLICT_RETRY_SPREAD_MS = 100 + class BaseTokenAuthClient(CycodeClient, ABC): """Base client for token-based authentication flows with cached JWTs.""" @@ -49,7 +57,20 @@ def refresh_access_token_if_needed(self) -> None: self._load_token_from_disk() if self._has_valid_token(): return - self.refresh_access_token() + + try: + self.refresh_access_token() + except HttpUnauthorizedError: + # Processes sharing one cached token all expire at the same instant, so a burst of + # them mints together and the identity provider refuses all but the first as a + # too-fast login. The winner persists a usable token, so prefer re-reading it over + # minting again. The wait is randomized to keep the losers from colliding a second + # time. A genuinely invalid token still raises on the retry. + time.sleep((_MINT_CONFLICT_RETRY_MIN_MS + secrets.randbelow(_MINT_CONFLICT_RETRY_SPREAD_MS)) / 1000) + self._load_token_from_disk() + if self._has_valid_token(): + return + self.refresh_access_token() def _has_valid_token(self) -> bool: return self._access_token is not None and self._expires_in is not None and arrow.utcnow() < self._expires_in diff --git a/tests/cyclient/test_token_based_client.py b/tests/cyclient/test_token_based_client.py index 4c3dd4c5..ffb1f0c1 100644 --- a/tests/cyclient/test_token_based_client.py +++ b/tests/cyclient/test_token_based_client.py @@ -1,6 +1,9 @@ import arrow +import pytest import responses +from pyfakefs.fake_filesystem import FakeFilesystem +from cycode.cli.exceptions.custom_exceptions import HttpUnauthorizedError from cycode.cyclient.cycode_token_based_client import CycodeTokenBasedClient from tests.conftest import _EXPECTED_API_TOKEN, create_token_based_client @@ -66,6 +69,51 @@ def test_access_token_cached_creator_changed( assert client2._expires_in is None +@responses.activate +def test_access_token_mint_conflict_prefers_token_persisted_by_another_process( + api_token_url: str, fs: FakeFilesystem +) -> None: + client = create_token_based_client() + + def _refuse_while_another_process_wins(_request: object) -> tuple: + # the process that won the race persists its token while this one is being refused + client._credentials_manager.update_access_token( + _EXPECTED_API_TOKEN, arrow.utcnow().shift(hours=1).timestamp(), client._create_jwt_creator() + ) + return 401, {}, '' + + responses.add_callback(responses.POST, api_token_url, callback=_refuse_while_another_process_wins) + + assert client.get_access_token() == _EXPECTED_API_TOKEN + assert len(responses.calls) == 1 + + +@responses.activate +def test_access_token_mint_conflict_retries_when_no_other_process_won( + api_token_url: str, api_token_response: responses.Response, fs: FakeFilesystem +) -> None: + client = create_token_based_client() + + responses.add(responses.Response(method=responses.POST, url=api_token_url, status=401)) + responses.add(api_token_response) + + assert client.get_access_token() == _EXPECTED_API_TOKEN + assert len(responses.calls) == 2 + + +@responses.activate +def test_access_token_mint_conflict_raises_when_retry_is_refused_too(api_token_url: str, fs: FakeFilesystem) -> None: + client = create_token_based_client() + + responses.add(responses.Response(method=responses.POST, url=api_token_url, status=401)) + responses.add(responses.Response(method=responses.POST, url=api_token_url, status=401)) + + with pytest.raises(HttpUnauthorizedError): + client.get_access_token() + + assert len(responses.calls) == 2 + + @responses.activate def test_access_token_invalidation( token_based_client: CycodeTokenBasedClient, api_token_response: responses.Response