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