371 lines
12 KiB
Python
371 lines
12 KiB
Python
import asyncio
|
|
import logging
|
|
import threading
|
|
from datetime import datetime, timezone
|
|
from time import sleep
|
|
from typing import Any, Awaitable, Callable, Union
|
|
|
|
from redis.auth.err import RequestTokenErr, TokenRenewalErr
|
|
from redis.auth.idp import IdentityProviderInterface
|
|
from redis.auth.token import TokenResponse
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class CredentialsListener:
|
|
"""
|
|
Listeners that will be notified on events related to credentials.
|
|
Accepts callbacks and awaitable callbacks.
|
|
"""
|
|
|
|
def __init__(self):
|
|
self._on_next = None
|
|
self._on_error = None
|
|
|
|
@property
|
|
def on_next(self) -> Union[Callable[[Any], None], Awaitable]:
|
|
return self._on_next
|
|
|
|
@on_next.setter
|
|
def on_next(self, callback: Union[Callable[[Any], None], Awaitable]) -> None:
|
|
self._on_next = callback
|
|
|
|
@property
|
|
def on_error(self) -> Union[Callable[[Exception], None], Awaitable]:
|
|
return self._on_error
|
|
|
|
@on_error.setter
|
|
def on_error(self, callback: Union[Callable[[Exception], None], Awaitable]) -> None:
|
|
self._on_error = callback
|
|
|
|
|
|
class RetryPolicy:
|
|
def __init__(self, max_attempts: int, delay_in_ms: float):
|
|
self.max_attempts = max_attempts
|
|
self.delay_in_ms = delay_in_ms
|
|
|
|
def get_max_attempts(self) -> int:
|
|
"""
|
|
Retry attempts before exception will be thrown.
|
|
|
|
:return: int
|
|
"""
|
|
return self.max_attempts
|
|
|
|
def get_delay_in_ms(self) -> float:
|
|
"""
|
|
Delay between retries in seconds.
|
|
|
|
:return: int
|
|
"""
|
|
return self.delay_in_ms
|
|
|
|
|
|
class TokenManagerConfig:
|
|
def __init__(
|
|
self,
|
|
expiration_refresh_ratio: float,
|
|
lower_refresh_bound_millis: int,
|
|
token_request_execution_timeout_in_ms: int,
|
|
retry_policy: RetryPolicy,
|
|
):
|
|
self._expiration_refresh_ratio = expiration_refresh_ratio
|
|
self._lower_refresh_bound_millis = lower_refresh_bound_millis
|
|
self._token_request_execution_timeout_in_ms = (
|
|
token_request_execution_timeout_in_ms
|
|
)
|
|
self._retry_policy = retry_policy
|
|
|
|
def get_expiration_refresh_ratio(self) -> float:
|
|
"""
|
|
Represents the ratio of a token's lifetime at which a refresh should be triggered. # noqa: E501
|
|
For example, a value of 0.75 means the token should be refreshed
|
|
when 75% of its lifetime has elapsed (or when 25% of its lifetime remains).
|
|
|
|
:return: float
|
|
"""
|
|
|
|
return self._expiration_refresh_ratio
|
|
|
|
def get_lower_refresh_bound_millis(self) -> int:
|
|
"""
|
|
Represents the minimum time in milliseconds before token expiration
|
|
to trigger a refresh, in milliseconds.
|
|
This value sets a fixed lower bound for when a token refresh should occur,
|
|
regardless of the token's total lifetime.
|
|
If set to 0 there will be no lower bound and the refresh will be triggered
|
|
based on the expirationRefreshRatio only.
|
|
|
|
:return: int
|
|
"""
|
|
return self._lower_refresh_bound_millis
|
|
|
|
def get_token_request_execution_timeout_in_ms(self) -> int:
|
|
"""
|
|
Represents the maximum time in milliseconds to wait
|
|
for a token request to complete.
|
|
|
|
:return: int
|
|
"""
|
|
return self._token_request_execution_timeout_in_ms
|
|
|
|
def get_retry_policy(self) -> RetryPolicy:
|
|
"""
|
|
Represents the retry policy for token requests.
|
|
|
|
:return: RetryPolicy
|
|
"""
|
|
return self._retry_policy
|
|
|
|
|
|
class TokenManager:
|
|
def __init__(
|
|
self, identity_provider: IdentityProviderInterface, config: TokenManagerConfig
|
|
):
|
|
self._idp = identity_provider
|
|
self._config = config
|
|
self._next_timer = None
|
|
self._listener = None
|
|
self._init_timer = None
|
|
self._retries = 0
|
|
|
|
def __del__(self):
|
|
logger.info("Token manager are disposed")
|
|
self.stop()
|
|
|
|
def start(
|
|
self,
|
|
listener: CredentialsListener,
|
|
skip_initial: bool = False,
|
|
) -> Callable[[], None]:
|
|
self._listener = listener
|
|
|
|
try:
|
|
loop = asyncio.get_running_loop()
|
|
except RuntimeError:
|
|
# Run loop in a separate thread to unblock main thread.
|
|
loop = asyncio.new_event_loop()
|
|
thread = threading.Thread(
|
|
target=_start_event_loop_in_thread, args=(loop,), daemon=True
|
|
)
|
|
thread.start()
|
|
|
|
# Event to block for initial execution.
|
|
init_event = asyncio.Event()
|
|
self._init_timer = loop.call_later(
|
|
0, self._renew_token, skip_initial, init_event
|
|
)
|
|
logger.info("Token manager started")
|
|
|
|
# Blocks in thread-safe manner.
|
|
asyncio.run_coroutine_threadsafe(init_event.wait(), loop).result()
|
|
return self.stop
|
|
|
|
async def start_async(
|
|
self,
|
|
listener: CredentialsListener,
|
|
block_for_initial: bool = False,
|
|
initial_delay_in_ms: float = 0,
|
|
skip_initial: bool = False,
|
|
) -> Callable[[], None]:
|
|
self._listener = listener
|
|
|
|
loop = asyncio.get_running_loop()
|
|
init_event = asyncio.Event()
|
|
|
|
# Wraps the async callback with async wrapper to schedule with loop.call_later()
|
|
wrapped = _async_to_sync_wrapper(
|
|
loop, self._renew_token_async, skip_initial, init_event
|
|
)
|
|
self._init_timer = loop.call_later(initial_delay_in_ms / 1000, wrapped)
|
|
logger.info("Token manager started")
|
|
|
|
if block_for_initial:
|
|
await init_event.wait()
|
|
|
|
return self.stop
|
|
|
|
def stop(self):
|
|
if self._init_timer is not None:
|
|
self._init_timer.cancel()
|
|
if self._next_timer is not None:
|
|
self._next_timer.cancel()
|
|
|
|
def acquire_token(self, force_refresh=False) -> TokenResponse:
|
|
try:
|
|
token = self._idp.request_token(force_refresh)
|
|
except RequestTokenErr as e:
|
|
if self._retries < self._config.get_retry_policy().get_max_attempts():
|
|
self._retries += 1
|
|
sleep(self._config.get_retry_policy().get_delay_in_ms() / 1000)
|
|
return self.acquire_token(force_refresh)
|
|
else:
|
|
raise e
|
|
|
|
self._retries = 0
|
|
return TokenResponse(token)
|
|
|
|
async def acquire_token_async(self, force_refresh=False) -> TokenResponse:
|
|
try:
|
|
token = self._idp.request_token(force_refresh)
|
|
except RequestTokenErr as e:
|
|
if self._retries < self._config.get_retry_policy().get_max_attempts():
|
|
self._retries += 1
|
|
await asyncio.sleep(
|
|
self._config.get_retry_policy().get_delay_in_ms() / 1000
|
|
)
|
|
return await self.acquire_token_async(force_refresh)
|
|
else:
|
|
raise e
|
|
|
|
self._retries = 0
|
|
return TokenResponse(token)
|
|
|
|
def _calculate_renewal_delay(self, expire_date: float, issue_date: float) -> float:
|
|
delay_for_lower_refresh = self._delay_for_lower_refresh(expire_date)
|
|
delay_for_ratio_refresh = self._delay_for_ratio_refresh(expire_date, issue_date)
|
|
delay = min(delay_for_ratio_refresh, delay_for_lower_refresh)
|
|
|
|
return 0 if delay < 0 else delay / 1000
|
|
|
|
def _delay_for_lower_refresh(self, expire_date: float):
|
|
return (
|
|
expire_date
|
|
- self._config.get_lower_refresh_bound_millis()
|
|
- (datetime.now(timezone.utc).timestamp() * 1000)
|
|
)
|
|
|
|
def _delay_for_ratio_refresh(self, expire_date: float, issue_date: float):
|
|
token_ttl = expire_date - issue_date
|
|
refresh_before = token_ttl - (
|
|
token_ttl * self._config.get_expiration_refresh_ratio()
|
|
)
|
|
|
|
return (
|
|
expire_date
|
|
- refresh_before
|
|
- (datetime.now(timezone.utc).timestamp() * 1000)
|
|
)
|
|
|
|
def _renew_token(
|
|
self, skip_initial: bool = False, init_event: asyncio.Event = None
|
|
):
|
|
"""
|
|
Task to renew token from identity provider.
|
|
Schedules renewal tasks based on token TTL.
|
|
"""
|
|
|
|
try:
|
|
token_res = self.acquire_token(force_refresh=True)
|
|
delay = self._calculate_renewal_delay(
|
|
token_res.get_token().get_expires_at_ms(),
|
|
token_res.get_token().get_received_at_ms(),
|
|
)
|
|
|
|
if token_res.get_token().is_expired():
|
|
raise TokenRenewalErr("Requested token is expired")
|
|
|
|
if self._listener.on_next is None:
|
|
logger.warning(
|
|
"No registered callback for token renewal task. Renewal cancelled"
|
|
)
|
|
return
|
|
|
|
if not skip_initial:
|
|
try:
|
|
self._listener.on_next(token_res.get_token())
|
|
except Exception as e:
|
|
raise TokenRenewalErr(e)
|
|
|
|
if delay <= 0:
|
|
return
|
|
|
|
loop = asyncio.get_running_loop()
|
|
self._next_timer = loop.call_later(delay, self._renew_token)
|
|
logger.info(f"Next token renewal scheduled in {delay} seconds")
|
|
return token_res
|
|
except Exception as e:
|
|
if self._listener.on_error is None:
|
|
raise e
|
|
|
|
self._listener.on_error(e)
|
|
finally:
|
|
if init_event:
|
|
init_event.set()
|
|
|
|
async def _renew_token_async(
|
|
self, skip_initial: bool = False, init_event: asyncio.Event = None
|
|
):
|
|
"""
|
|
Async task to renew tokens from identity provider.
|
|
Schedules renewal tasks based on token TTL.
|
|
"""
|
|
|
|
try:
|
|
token_res = await self.acquire_token_async(force_refresh=True)
|
|
delay = self._calculate_renewal_delay(
|
|
token_res.get_token().get_expires_at_ms(),
|
|
token_res.get_token().get_received_at_ms(),
|
|
)
|
|
|
|
if token_res.get_token().is_expired():
|
|
raise TokenRenewalErr("Requested token is expired")
|
|
|
|
if self._listener.on_next is None:
|
|
logger.warning(
|
|
"No registered callback for token renewal task. Renewal cancelled"
|
|
)
|
|
return
|
|
|
|
if not skip_initial:
|
|
try:
|
|
await self._listener.on_next(token_res.get_token())
|
|
except Exception as e:
|
|
raise TokenRenewalErr(e)
|
|
|
|
if delay <= 0:
|
|
return
|
|
|
|
loop = asyncio.get_running_loop()
|
|
wrapped = _async_to_sync_wrapper(loop, self._renew_token_async)
|
|
logger.info(f"Next token renewal scheduled in {delay} seconds")
|
|
loop.call_later(delay, wrapped)
|
|
except Exception as e:
|
|
if self._listener.on_error is None:
|
|
raise e
|
|
|
|
await self._listener.on_error(e)
|
|
finally:
|
|
if init_event:
|
|
init_event.set()
|
|
|
|
|
|
def _async_to_sync_wrapper(loop, coro_func, *args, **kwargs):
|
|
"""
|
|
Wraps an asynchronous function so it can be used with loop.call_later.
|
|
|
|
:param loop: The event loop in which the coroutine will be executed.
|
|
:param coro_func: The coroutine function to wrap.
|
|
:param args: Positional arguments to pass to the coroutine function.
|
|
:param kwargs: Keyword arguments to pass to the coroutine function.
|
|
:return: A regular function suitable for loop.call_later.
|
|
"""
|
|
|
|
def wrapped():
|
|
# Schedule the coroutine in the event loop
|
|
asyncio.ensure_future(coro_func(*args, **kwargs), loop=loop)
|
|
|
|
return wrapped
|
|
|
|
|
|
def _start_event_loop_in_thread(event_loop: asyncio.AbstractEventLoop):
|
|
"""
|
|
Starts event loop in a thread.
|
|
Used to be able to schedule tasks using loop.call_later.
|
|
|
|
:param event_loop:
|
|
:return:
|
|
"""
|
|
asyncio.set_event_loop(event_loop)
|
|
event_loop.run_forever()
|