async def test_event_hubs_send_override_token_refresh_window(
live_eventhub_config):
uri = "sb://{}/{}".format(live_eventhub_config['hostname'],
live_eventhub_config['event_hub'])
target = "amqps://{}/{}/Partitions/0".format(
live_eventhub_config['hostname'], live_eventhub_config['event_hub'])
token = None
async def get_token():
nonlocal token
return _AccessToken(token, expiry)
jwt_auth = authentication.JWTTokenAsync(
uri,
uri,
get_token,
refresh_window=300 # set refresh window to be 5 mins
)
send_client = uamqp.SendClientAsync(target, auth=jwt_auth, debug=False)
# use token of which the valid remaining time < refresh window
expiry = int(time.time()) + (60 * 4 + 30) # 4.5 minutes
token = utils.create_sas_token(live_eventhub_config['key_name'].encode(),
live_eventhub_config['access_key'].encode(),
uri.encode(),
expiry=timedelta(minutes=4, seconds=30))
for _ in range(3):
message = uamqp.message.Message(body='Hello World')
await send_client.send_message_async(message)
auth_status = constants.CBSAuthStatus(jwt_auth._cbs_auth.get_status())
assert auth_status == constants.CBSAuthStatus.RefreshRequired
# update token, the valid remaining time > refresh window
expiry = int(time.time()) + (60 * 5 + 30) # 5.5 minutes
token = utils.create_sas_token(live_eventhub_config['key_name'].encode(),
live_eventhub_config['access_key'].encode(),
uri.encode(),
expiry=timedelta(minutes=5, seconds=30))
for _ in range(3):
message = uamqp.message.Message(body='Hello World')
await send_client.send_message_async(message)
auth_status = constants.CBSAuthStatus(jwt_auth._cbs_auth.get_status())
assert auth_status == constants.CBSAuthStatus.Ok
await send_client.close_async()
def handle_token(self):
"""This function is called periodically to check the status of the current
token if there is one, and request a new one if needed.
If the token request fails, it will be retried according to the retry policy.
A token refresh will be attempted if the token will expire soon.
This function will return a tuple of two booleans. The first represents whether
the token authentication has not completed within it's given timeout window. The
second indicates whether the token negotiation is still in progress.
:raises: ~uamqp.errors.AuthenticationException if the token authentication fails.
:raises: ~uamqp.errors.TokenExpired if the token has expired and cannot be
refreshed.
:returns: tuple[bool, bool]
"""
timeout = False
in_progress = False
self._lock.acquire()
try:
auth_status = self._cbs_auth.get_status()
auth_status = constants.CBSAuthStatus(auth_status)
if auth_status == constants.CBSAuthStatus.Error:
if self.retries >= self._retry_policy.retries: # pylint: disable=no-member
_logger.warning("Authentication Put-Token failed. Retries exhausted.")
raise errors.TokenAuthFailure(*self._cbs_auth.get_failure_info())
else:
_logger.info("Authentication Put-Token failed. Retrying.")
self.retries += 1 # pylint: disable=no-member
time.sleep(self._retry_policy.backoff)
self._cbs_auth.authenticate()
in_progress = True
elif auth_status == constants.CBSAuthStatus.Failure:
errors.AuthenticationException("Failed to open CBS authentication link.")
elif auth_status == constants.CBSAuthStatus.Expired:
raise errors.TokenExpired("CBS Authentication Expired.")
elif auth_status == constants.CBSAuthStatus.Timeout:
timeout = True
elif auth_status == constants.CBSAuthStatus.InProgress:
in_progress = True
elif auth_status == constants.CBSAuthStatus.RefreshRequired:
_logger.info("Token will expire soon - attempting to refresh.")
self.update_token()
self._cbs_auth.refresh(self.token, int(self.expires_at))
elif auth_status == constants.CBSAuthStatus.Idle:
self._cbs_auth.authenticate()
in_progress = True
elif auth_status != constants.CBSAuthStatus.Ok:
raise errors.AuthenticationException("Invalid auth state.")
except ValueError as e:
raise errors.AuthenticationException(
"Token authentication failed: {}".format(e))
except:
raise
finally:
self._lock.release()
return timeout, in_progress
async def handle_token_async(self):
"""This coroutine is called periodically to check the status of the current
token if there is one, and request a new one if needed.
If the token request fails, it will be retried according to the retry policy.
A token refresh will be attempted if the token will expire soon.
This function will return a tuple of two booleans. The first represents whether
the token authentication has not completed within it's given timeout window. The
second indicates whether the token negotiation is still in progress.
:raises: ~uamqp.errors.AuthenticationException if the token authentication fails.
:raises: ~uamqp.errors.TokenExpired if the token has expired and cannot be
refreshed.
:rtype: tuple[bool, bool]
"""
# pylint: disable=protected-access
timeout = False
in_progress = False
try:
await self._connection.lock_async()
if self._connection._closing or self._connection._error:
return timeout, in_progress
auth_status = self._cbs_auth.get_status()
auth_status = constants.CBSAuthStatus(auth_status)
if auth_status == constants.CBSAuthStatus.Error:
if self.retries >= self._retry_policy.retries: # pylint: disable=no-member
_logger.warning(
"Authentication Put-Token failed. Retries exhausted.")
raise errors.TokenAuthFailure(
*self._cbs_auth.get_failure_info())
error_code, error_description = self._cbs_auth.get_failure_info(
)
_logger.info("Authentication status: %r, description: %r",
error_code, error_description)
_logger.info("Authentication Put-Token failed. Retrying.")
self.retries += 1 # pylint: disable=no-member
await asyncio.sleep(self._retry_policy.backoff)
self._cbs_auth.authenticate()
in_progress = True
elif auth_status == constants.CBSAuthStatus.Failure:
raise errors.AuthenticationException(
"Failed to open CBS authentication link.")
elif auth_status == constants.CBSAuthStatus.Expired:
raise errors.TokenExpired("CBS Authentication Expired.")
elif auth_status == constants.CBSAuthStatus.Timeout:
timeout = True
elif auth_status == constants.CBSAuthStatus.InProgress:
in_progress = True
elif auth_status == constants.CBSAuthStatus.RefreshRequired:
_logger.info(
"Token on connection %r will expire soon - attempting to refresh.",
self._connection.container_id)
await self.update_token()
self._cbs_auth.refresh(self.token, int(self.expires_at))
elif auth_status == constants.CBSAuthStatus.Idle:
self._cbs_auth.authenticate()
in_progress = True
elif auth_status != constants.CBSAuthStatus.Ok:
raise ValueError("Invalid auth state.")
except asyncio.TimeoutError:
_logger.debug(
"CBS auth timed out while waiting for lock acquisition.")
return None, None
except ValueError as e:
raise errors.AuthenticationException(
"Token authentication failed: {}".format(e))
finally:
self._connection.release_async()
return timeout, in_progress