diff --git a/README.md b/README.md index ef7c4a8..657d14a 100644 --- a/README.md +++ b/README.md @@ -169,6 +169,7 @@ Client.build_context( regex_max_time_limit=100, remote=RemoteOptions( cert_path='./certs/ca.pem', + auto_renew_token=True, connect_timeout=0.3, read_timeout=5.0, write_timeout=5.0, @@ -201,6 +202,7 @@ switcher = Client.get_switcher() | Option | Type | Description | Default | |--------|------|-------------|---------| | `cert_path` | `str` | Path to custom certificate for secure API connections | `None` | +| `auto_renew_token` | `bool` | Automatically renew the auth token in the background before it expires | `False` | | `connect_timeout` | `float` | Max seconds to establish a remote connection before failing fast | `0.3` | | `read_timeout` | `float` | Max seconds to wait for remote response data | `5.0` | | `write_timeout` | `float` | Max seconds to send remote request data | `5.0` | diff --git a/switcher_client/lib/globals/global_context.py b/switcher_client/lib/globals/global_context.py index 5194c9b..88231a9 100644 --- a/switcher_client/lib/globals/global_context.py +++ b/switcher_client/lib/globals/global_context.py @@ -12,6 +12,8 @@ DEFAULT_REMOTE_READ_TIMEOUT = 5.0 DEFAULT_REMOTE_WRITE_TIMEOUT = 5.0 DEFAULT_REMOTE_POOL_TIMEOUT = 5.0 +DEFAULT_REMOTE_CERT_PATH = None +DEFAULT_REMOTE_AUTO_RENEW_TOKEN = False DEFAULT_TEST_MODE = False @dataclass @@ -19,13 +21,15 @@ class RemoteOptions: """ :param cert_path: The path to the SSL certificate file for secure connections. If not set, it will use the default system certificates + :param auto_renew_token: When enabled, it will automatically renew the token before it expires. :param connect_timeout: Max time in seconds to establish a remote connection. Lower values help silent mode trip faster when the upstream is unavailable :param read_timeout: Max time in seconds to wait for a remote response body :param write_timeout: Max time in seconds to send a remote request body :param pool_timeout: Max time in seconds to wait for a pooled connection """ - cert_path: Optional[str] = None + cert_path: Optional[str] = DEFAULT_REMOTE_CERT_PATH + auto_renew_token: bool = DEFAULT_REMOTE_AUTO_RENEW_TOKEN connect_timeout: float = DEFAULT_REMOTE_CONNECT_TIMEOUT read_timeout: float = DEFAULT_REMOTE_READ_TIMEOUT write_timeout: float = DEFAULT_REMOTE_WRITE_TIMEOUT diff --git a/switcher_client/lib/remote_auth.py b/switcher_client/lib/remote_auth.py index 7410d29..93b85dc 100644 --- a/switcher_client/lib/remote_auth.py +++ b/switcher_client/lib/remote_auth.py @@ -1,5 +1,9 @@ +import threading + from time import time +from functools import partial from datetime import datetime +from typing import Optional from switcher_client.lib.remote import Remote from switcher_client.lib.globals.global_context import Context @@ -14,9 +18,15 @@ class RemoteAuth: __context: Context = Context.empty() __retry_options: RetryOptions + __auto_renew_timer: Optional[threading.Timer] = None + __auto_renew_generation: int = 0 + __auto_renew_lock = threading.Lock() + __auto_renew_buffer_seconds = 5.0 + __auto_renew_min_delay_seconds = 1.0 @staticmethod def init(context: Context): + RemoteAuth._stop_auto_renew() RemoteAuth.__context = context GlobalAuth.init() @@ -29,9 +39,7 @@ def set_retry_options(silent_mode: str): @staticmethod def auth(): - token, exp = Remote.auth(RemoteAuth.__context) - GlobalAuth.set_token(token) - GlobalAuth.set_exp(exp) + RemoteAuth._refresh_token() @staticmethod def check_health(): @@ -72,3 +80,78 @@ def is_valid(): errors = [name for name, value in required_fields if not value] if errors: raise ValueError(f"Something went wrong: Missing or empty required fields ({', '.join(errors)})") + + @staticmethod + def _refresh_token(schedule_next: bool = True, generation: Optional[int] = None): + # Prevent token refresh if the generation is stale + if not RemoteAuth._is_current_generation(generation): + return + + token, exp = Remote.auth(RemoteAuth.__context) + + # Timer is valid, another thread may have already updated the token, so we should not overwrite it + if not RemoteAuth._is_current_generation(generation): + return + + GlobalAuth.set_token(token) + GlobalAuth.set_exp(exp) + + if schedule_next: + RemoteAuth._schedule_auto_renew(exp) + + @staticmethod + def _schedule_auto_renew(exp: str): + if not RemoteAuth.__context.options.remote.auto_renew_token: + return + + delay = RemoteAuth._get_auto_renew_delay(exp) + with RemoteAuth.__auto_renew_lock: + RemoteAuth.__auto_renew_generation += 1 + current_generation = RemoteAuth.__auto_renew_generation + previous_timer = RemoteAuth.__auto_renew_timer + + timer = threading.Timer(delay, partial(RemoteAuth._auto_renew, current_generation)) + timer.daemon = True + + with RemoteAuth.__auto_renew_lock: + RemoteAuth.__auto_renew_timer = timer + + if previous_timer is not None: + previous_timer.cancel() + + timer.start() + + @staticmethod + def _auto_renew(generation: int): + try: + RemoteAuth._refresh_token(generation=generation) + except Exception: # pylint: disable=broad-exception-caught + RemoteAuth._stop_auto_renew(generation) + + @staticmethod + def _stop_auto_renew(generation: Optional[int] = None): + with RemoteAuth.__auto_renew_lock: + if generation is not None and generation != RemoteAuth.__auto_renew_generation: + return + + RemoteAuth.__auto_renew_generation += 1 + timer = RemoteAuth.__auto_renew_timer + RemoteAuth.__auto_renew_timer = None + + if timer is not None: + timer.cancel() + + @staticmethod + def _is_current_generation(generation: Optional[int]) -> bool: + if generation is None: + return True + + with RemoteAuth.__auto_renew_lock: + return generation == RemoteAuth.__auto_renew_generation + + @staticmethod + def _get_auto_renew_delay(exp: str) -> float: + expiration = float(exp) + remaining = expiration - time() + delay = remaining - RemoteAuth.__auto_renew_buffer_seconds + return max(delay, RemoteAuth.__auto_renew_min_delay_seconds) diff --git a/tests/remote/__init__.py b/tests/remote/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/remote/test_remote_auth.py b/tests/remote/test_remote_auth.py new file mode 100644 index 0000000..8fb08f6 --- /dev/null +++ b/tests/remote/test_remote_auth.py @@ -0,0 +1,90 @@ +import threading + +from unittest.mock import Mock, patch + +from tests.helpers import given_context + +from switcher_client.lib.globals.global_auth import GlobalAuth +from switcher_client.lib.globals.global_context import ContextOptions, RemoteOptions +from switcher_client.lib.remote_auth import RemoteAuth + +def test_refresh_token_skips_auth_for_stale_generation(): + """ Should ignore stale background refresh attempts before calling remote auth """ + + # given + given_context(options=ContextOptions( + remote=RemoteOptions(auto_renew_token=True) + )) + GlobalAuth.set_token('[current_token]') + GlobalAuth.set_exp('9999999999') + + # test + with patch('switcher_client.lib.remote_auth.Remote.auth') as mock_auth: + RemoteAuth._refresh_token(generation=_stale_generation()) + + assert GlobalAuth.get_token() == '[current_token]' + mock_auth.assert_not_called() + +def test_refresh_token_discards_stale_result_after_auth_returns(): + """ Should discard auth results when the renewer becomes stale during the request """ + + # given + given_context(options=ContextOptions( + remote=RemoteOptions(auto_renew_token=True) + )) + GlobalAuth.set_token('[current_token]') + GlobalAuth.set_exp('9999999999') + + auth_started = threading.Event() + release_auth = threading.Event() + current_generation = _current_generation() + + def fake_auth(_): + auth_started.set() + release_auth.wait(timeout=2) + return '[new_token]', '9999999999' + + with patch('switcher_client.lib.remote_auth.Remote.auth', side_effect=fake_auth) as mock_auth, \ + patch.object(RemoteAuth, '_schedule_auto_renew') as mock_schedule: + refresh_thread = threading.Thread( + target=RemoteAuth._refresh_token, + kwargs={'generation': current_generation} + ) + refresh_thread.start() + + assert auth_started.wait(timeout=1) + RemoteAuth._stop_auto_renew() + release_auth.set() + refresh_thread.join(timeout=2) + + assert not refresh_thread.is_alive() + assert GlobalAuth.get_token() == '[current_token]' + assert mock_auth.call_count == 1 + mock_schedule.assert_not_called() + +def test_stop_auto_renew_ignores_stale_generation(): + """ Should not cancel the active renewer when stop is requested with a stale generation """ + + # given + given_context(options=ContextOptions( + remote=RemoteOptions(auto_renew_token=True) + )) + active_timer = Mock() + current_generation = _current_generation() + RemoteAuth._RemoteAuth__auto_renew_timer = active_timer # type: ignore + + # test + RemoteAuth._stop_auto_renew(generation=_stale_generation()) + + assert _current_generation() == current_generation + assert RemoteAuth._RemoteAuth__auto_renew_timer is active_timer # type: ignore + active_timer.cancel.assert_not_called() + +# Helpers + +def _current_generation() -> int: + return RemoteAuth._RemoteAuth__auto_renew_generation # type: ignore + +def _stale_generation() -> int: + current_generation = _current_generation() + return current_generation - 1 if current_generation > 0 else current_generation + 1 diff --git a/tests/test_switcher_remote.py b/tests/test_switcher_remote.py index 6485747..dc643c5 100644 --- a/tests/test_switcher_remote.py +++ b/tests/test_switcher_remote.py @@ -1,5 +1,6 @@ import pytest import time +import threading import httpx from unittest.mock import Mock, patch @@ -114,6 +115,82 @@ def test_remote_renew_token(httpx_mock): switcher.is_on('MY_SWITCHER') assert GlobalAuth.get_token() == '[new_token]' +def test_remote_autorenew_token(httpx_mock): + """ Should refresh the token before it expires in the background """ + + # given + given_auth(httpx_mock, status=200, token='[token_1]', exp=int(round(time.time())) + 2) + given_auth(httpx_mock, status=200, token='[token_2]', exp=int(round(time.time())) + 3600) + given_check_criteria(httpx_mock, response={'result': True}) + given_context(options=ContextOptions( + remote=RemoteOptions(auto_renew_token=True) + )) + + switcher = Client.get_switcher() + + # test + switcher.is_on('MY_SWITCHER') + assert GlobalAuth.get_token() == '[token_1]' + time.sleep(3) # Wait for the token to expire and auto-renew + assert GlobalAuth.get_token() == '[token_2]' + +def test_remote_autorenew_token_disabled(httpx_mock): + """ Should not create a background renewer when auto renew is disabled """ + + # given + given_auth(httpx_mock, status=200, token='[token_1]', exp=int(round(time.time())) + 2) + given_check_criteria(httpx_mock, response={'result': True}) + given_context(options=ContextOptions( + remote=RemoteOptions(auto_renew_token=False) + )) + + with patch('switcher_client.lib.remote_auth.threading.Timer') as mock_timer: + switcher = Client.get_switcher() + + # test + switcher.is_on('MY_SWITCHER') + assert GlobalAuth.get_token() == '[token_1]' + time.sleep(3) + assert GlobalAuth.get_token() == '[token_1]' + mock_timer.assert_not_called() + +def test_remote_autorenew_token_failure(httpx_mock): + """ Should stop auto renew after background failure and restart after foreground auth """ + + # given + real_timer = threading.Timer + given_auth(httpx_mock, status=200, token='[token_1]', exp=int(round(time.time())) + 6) + given_auth(httpx_mock, status=500, token=None, exp=int(round(time.time())) + 3600) + given_auth(httpx_mock, status=200, token='[token_2]', exp=int(round(time.time())) + 13) + given_auth(httpx_mock, status=200, token='[token_3]', exp=int(round(time.time())) + 3600) + given_check_criteria(httpx_mock, response={'result': True}) + given_check_criteria(httpx_mock, response={'result': True}) + given_context(options=ContextOptions( + remote=RemoteOptions(auto_renew_token=True) + )) + + with patch('switcher_client.lib.remote_auth.threading.Timer', + side_effect=lambda *args, **kwargs: real_timer(*args, **kwargs)) as mock_timer: + switcher = Client.get_switcher() + + # test + assert switcher.is_on('MY_SWITCHER') + assert GlobalAuth.get_token() == '[token_1]' + assert mock_timer.call_count == 1 + + time.sleep(2) + assert GlobalAuth.get_token() == '[token_1]' + assert mock_timer.call_count == 1 + + time.sleep(5) + assert switcher.is_on('MY_SWITCHER') + assert GlobalAuth.get_token() == '[token_2]' + assert mock_timer.call_count == 2 + + time.sleep(2) + assert GlobalAuth.get_token() == '[token_3]' + assert mock_timer.call_count == 3 + def test_remote_with_remote_required_request(httpx_mock): """ Should call the remote API with success for required remote request"""