cache existing authentication tokens when chaning username+password

when changing `auth.user.username` and `auth.user.password` during runtime, any existing authorization tokens should be cached, so they can be re-used if changed back to the original user.

this is to avoid having to authenticate everytime a user is switched.

tokens are only refresh when the original `refresh_time` has expired.

grizzly/auth/__init__.py

@@ -8,7 +8,7 @@
 from functools import wraps
 from importlib import import_module
 from time import time
-from typing import TYPE_CHECKING, Any, Callable, ClassVar, Dict, Generic, Literal, Optional, Tuple, Type, TypeVar, Union, cast
+from typing import TYPE_CHECKING, Any, Callable, ClassVar, Dict, Generic, Literal, Optional, Set, Tuple, Type, TypeVar, Union, cast
 from urllib.parse import urlparse
 
 from grizzly.tasks import RequestTask
@@ -64,6 +64,8 @@ class GrizzlyHttpAuthClient(Generic[P], metaclass=ABCMeta):
             },
         },
         'metadata': None,
+        '__cached_auth__': {},
+        '__context_change_history__': set(),
     }
     session_started: Optional[float]
     grizzly: GrizzlyContext
@@ -76,6 +78,14 @@ def add_metadata(self, key: str, value: str) -> None:
 
         cast(dict, self._context['metadata']).update({key: value})
 
+    @property
+    def __context_change_history__(self) -> Set[str]:
+        return cast(Set[str], self._context['__context_change_history__'])
+
+    @property
+    def __cached_auth__(self) -> Dict[str, str]:
+        return cast(Dict[str, str], self._context['__cached_auth__'])
+
 
 AuthenticatableFunc = TypeVar('AuthenticatableFunc', bound=Callable[..., GrizzlyResponse])
 

grizzly/auth/aad.py

@@ -752,7 +752,7 @@ def generate_pkcs() -> Tuple[str, str]:
             request_meta = {
                 'request_type': 'AUTH',
                 'response_time': int((time_perf_counter() - start_time) * 1000),
-                'name': f'{scenario_index} {cls.__name__} OAuth2 user token {version}',
+                'name': f'{scenario_index} {cls.__name__} OAuth2 user token {version}: {username_lowercase}',
                 'context': client._context,
                 'response': None,
                 'exception': exception,
@@ -872,7 +872,7 @@ def get_oauth_token(cls, client: GrizzlyHttpAuthClient, pkcs: Optional[Tuple[str
             request_meta = {
                 'request_type': 'AUTH',
                 'response_time': int((time_perf_counter() - start_time) * 1000),
-                'name': f'{scenario_index} {cls.__name__} OAuth2 user token {version}',
+                'name': f'{scenario_index} {cls.__name__} OAuth2 client token {version}: {auth_client_context["id"]}',
                 'context': client._context,
                 'response': None,
                 'exception': exception,

grizzly/users/restapi.py

@@ -36,6 +36,32 @@
 
 See {@pylink grizzly.auth.aad}.
 
+It is possible to change authenticated user during runtime by using {@pylink grizzly.steps.scenario.setup.step_setup_set_context_variable} step expressions inbetween other tasks.
+To change user both `auth.user.username` and `auth.user.password` has to be changed (even though maybe only one of them changes value).
+
+This will then cache the `Authorization` token for the current user, and if changed back to that user there is no need to re-authenticate again, unless `refresh_time` for the first login
+expires.
+
+```gherkin
+Given a user of type "RestApi" load testing "https://api.example.com"
+And repeat for "2" iterations
+And set context variable "user.auth.username" to "bob"
+And set context variable "user.auth.password" to "foobar"
+
+Then get request from endpoint "/api/test"
+
+Given set context variable "user.auth.username" to "alice"
+And set context variable "user.auth.password" to "hello world"
+
+Then get request from endpoint "/api/test"
+
+Given set context variable "user.auth.username" to "bob"
+And set context variable "user.auth.password" to "foobar"
+```
+
+In the above, hypotetical scenario, there will 2 "AAD OAuth2 user token" requests, once for user "bob" and one for user "alice", both done in the first iteration. The second iteration
+the cached authentication tokens will be re-used.
+
 ### Multipart/form-data
 
 RestApi supports posting of multipart/form-data content-type, and in that case additional arguments needs to be passed with the request:
@@ -55,6 +81,7 @@
 import json
 from abc import ABCMeta
 from copy import copy
+from hashlib import sha256
 from html.parser import HTMLParser
 from http.cookiejar import Cookie
 from time import time
@@ -125,6 +152,8 @@ def handle_endtag(self, tag: str) -> None:
             'initialize_uri': None,
         },
     },
+    '__cached_auth__': {},
+    '__context_change_history__': set(),
 })
 class RestApiUser(GrizzlyUser, AsyncRequests, GrizzlyHttpAuthClient, metaclass=RestApiUserMeta):  # type: ignore[misc]
     session_started: Optional[float]
@@ -298,9 +327,69 @@ def _request(self, request: RequestTask, client: FastHttpSession) -> GrizzlyResp
         return (headers, payload)
 
     def add_context(self, context: Dict[str, Any]) -> None:
-        """If context change contains a username we should re-authenticate. This is forced by removing the Authorization header."""
-        # something change in auth context, we need to re-authenticate
-        if context.get('auth', {}).get('user', {}).get('username', None) is not None:
-            safe_del(self.metadata, 'Authorization')
+        """If added context contains changes in `auth`, we should cache current `Authorization` token and force re-auth for a new, if the auth
+        doesn't exist in the cache.
+
+        To force a re-authentication, both auth.user.username and auth.user.password needs be set, even though the actual value is only changed
+        for one of them.
+        """
+        current_username = self._context.get('auth', {}).get('user', {}).get('username', None)
+        current_password = self._context.get('auth', {}).get('user', {}).get('password', None)
+        current_authorization = self.metadata.get('Authorization', None)
+
+        changed_username = context.get('auth', {}).get('user', {}).get('username', None)
+        changed_password = context.get('auth', {}).get('user', {}).get('password', None)
+
+        # check if we're currently have Authorization header connected to a username and password,
+        # and if we're changing either username or password.
+        # if so, we need to cache current username+password token
+        if (
+            current_username is not None
+            and current_password is not None
+            and self.__context_change_history__ == set()
+            and current_authorization is not None
+            and (changed_username is not None or changed_password is not None)
+        ):
+            cache_key_plain = f'{current_username}:{current_password}'
+            cache_key = sha256(cache_key_plain.encode()).hexdigest()
+
+            if cache_key not in self.__cached_auth__:
+                self.__cached_auth__.update({cache_key: current_authorization})
 
         super().add_context(context)
+
+        changed_username_path = 'auth.user.username'
+        changed_password_path = 'auth.user.password'  # noqa: S105
+
+        # update change history if needed
+        context_change_history = self.__context_change_history__
+        if context_change_history != {changed_username_path, changed_password_path}:
+            if changed_username is not None and changed_username_path not in context_change_history:
+                self.__context_change_history__.add(changed_username_path)
+
+            if changed_password is not None and changed_password_path not in context_change_history:
+                self.__context_change_history__.add(changed_password_path)
+
+        # everything needed to force re-auth is not in place
+        if self.__context_change_history__ != {changed_username_path, changed_password_path}:
+            return
+
+        # every context change needed to force re-auth is in place, clear change history
+        self.__context_change_history__.clear()
+
+        username = self._context.get('auth', {}).get('user', {}).get('username', None)
+        password = self._context.get('auth', {}).get('user', {}).get('password', None)
+
+        if username is None or password is None:
+            return
+
+        # check if current username+password has a cached Authorization token
+        cache_key_plain = f'{username}:{password}'
+        cache_key = sha256(cache_key_plain.encode()).hexdigest()
+
+        cached_authorization = self.__cached_auth__.get(cache_key, None)
+
+        if cached_authorization is not None:   # restore from cache
+            self.metadata.update({'Authorization': cached_authorization})
+        else:  # remove from metadata, to force a re-authentication
+            safe_del(self.metadata, 'Authorization')

tests/e2e/test_auth.py

@@ -13,7 +13,7 @@
     from tests.fixtures import End2EndFixture
 
 
-def test_e2e_auth(e2e_fixture: End2EndFixture) -> None:
+def test_e2e_auth_user_token(e2e_fixture: End2EndFixture) -> None:
     if e2e_fixture._distributed:
         pytest.skip('telling the webserver what to expected auth-wise is not as simple when running dist, compare to running local')
 
@@ -27,14 +27,14 @@ def after_feature(context: Context, *_args: Any, **_kwargs: Any) -> None:
         stats = grizzly.state.locust.environment.stats
 
         expectations = [
-            ('001 AAD OAuth2 user token v1.0', 'AUTH', 0, 1),
-            ('001 RestApi auth', 'SCEN', 0, 1),
-            ('001 RestApi auth', 'TSTD', 0, 1),
-            ('001 restapi-echo', 'GET', 0, 1),
-            ('002 AAD OAuth2 user token v1.0', 'AUTH', 0, 1),
-            ('002 HttpClientTask auth', 'SCEN', 0, 1),
-            ('002 HttpClientTask auth', 'TSTD', 0, 1),
-            ('002 httpclient-echo', 'CLTSK', 0, 1),
+            ('001 AAD OAuth2 client token v1.0: dummy-client-id', 'AUTH', 0, 1),
+            ('001 RestApi auth', 'SCEN', 0, 2),
+            ('001 RestApi auth', 'TSTD', 0, 2),
+            ('001 restapi-echo', 'GET', 0, 2),
+            ('002 AAD OAuth2 client token v1.0: dummy-client-id', 'AUTH', 0, 1),
+            ('002 HttpClientTask auth', 'SCEN', 0, 2),
+            ('002 HttpClientTask auth', 'TSTD', 0, 2),
+            ('002 httpclient-echo', 'CLTSK', 0, 2),
         ]
 
         for name, method, expected_num_failures, expected_num_requests in expectations:
@@ -77,7 +77,7 @@ def add_metadata() -> str:
             And set context variable "auth.client.id" to "{e2e_fixture.webserver.auth['client']['id']}"
             And set context variable "auth.client.secret" to "{e2e_fixture.webserver.auth['client']['secret']}"
             {add_metadata()}
-            And repeat for "1" iterations
+            And repeat for "2" iterations
             Then get request with name "restapi-echo" from endpoint "/api/echo"
 
         Scenario: HttpClientTask auth
@@ -86,7 +86,7 @@ def add_metadata() -> str:
             And set context variable "{e2e_fixture.host}/auth.provider" to "http://{e2e_fixture.host}{e2e_fixture.webserver.auth_provider_uri}"
             And set context variable "{e2e_fixture.host}/auth.client.id" to "{e2e_fixture.webserver.auth['client']['id']}"
             And set context variable "{e2e_fixture.host}/auth.client.secret" to "{e2e_fixture.webserver.auth['client']['secret']}"
-            And repeat for "1" iterations
+            And repeat for "2" iterations
             Then get "http://{e2e_fixture.host}/api/echo" with name "httpclient-echo" and save response payload in "foobar"
             {add_metadata()}
         """))

tests/unit/test_grizzly/auth/test___init__.py

@@ -209,13 +209,13 @@ def request(_: RestApiUser, *_args: Any, **_kwargs: Any) -> GrizzlyResponse:
     get_token_mock.assert_called_once_with(parent.user, AuthMethod.USER)
     get_token_mock.reset_mock()
 
-    parent.user.add_context({'auth': {'user': {'username': '[email protected]'}}})
+    parent.user.add_context({'auth': {'user': {'username': '[email protected]', 'password': 'foobar'}}})
     assert 'Authorization' not in parent.user.metadata
     auth_context = parent.user._context.get('auth', None)
     assert auth_context is not None
     assert auth_context.get('user', None) == {
         'username': '[email protected]',
-        'password': 'HemligaArne',
+        'password': 'foobar',
         'otp_secret': None,
         'redirect_uri': '/authenticated',
         'initialize_uri': None,