fix(auth): refresh token before expiration, use locks while refreshing (#9045)

STACKITSDK-412

Author: cgoetz-inovex

core/src/stackit/core/auth_methods/key_auth.py

@@ -79,6 +79,16 @@ def __call__(self, r: Request) -> Request:
             if self.__is_token_expired(self.access_token):
                 if self.refresh_future is None or self.refresh_future.done():
                     self.refresh_future = self.executor.submit(self.__refresh_token)
+                refresh_future = self.refresh_future
+            else:
+                refresh_future = None
+
+        # Do not hold the lock while waiting. The refresh worker acquires it when
+        # updating the access token.
+        if refresh_future is not None:
+            refresh_future.result()
+
+        with self.lock:
             r.headers["Authorization"] = f"Bearer {self.access_token}"
         return r
 
@@ -108,8 +118,9 @@ def __fetch_token_from_endpoint(self) -> None:
             response = requests.post(self.token_endpoint, data=body, timeout=self.timeout)
             response.raise_for_status()
             response_json = response.json()
-            self.access_token = response_json["access_token"]
-            self.refresh_token = response_json["refresh_token"]
+            with self.lock:
+                self.access_token = response_json["access_token"]
+                self.refresh_token = response_json["refresh_token"]
         except requests.RequestException as e:
             raise requests.RequestException("Initial token fetch failed") from e
 
@@ -130,14 +141,17 @@ def token_refresh_task():
         thread.start()
 
     def __refresh_token(self):
-        if self.__is_token_expired(self.refresh_token):
+        with self.lock:
+            refresh_token = self.refresh_token
+
+        if self.__is_token_expired(refresh_token):
             self.__create_initial_token()
             self.__fetch_token_from_endpoint()
             return
 
         body = {
             "grant_type": "refresh_token",
-            "refresh_token": self.refresh_token,
+            "refresh_token": refresh_token,
         }
 
         last_exception = None
@@ -147,24 +161,23 @@ def __refresh_token(self):
                 response.raise_for_status()
                 response_data = response.json()
                 new_token = response_data.get("access_token")
-                self.access_token = new_token
+                with self.lock:
+                    self.access_token = new_token
                 return
             except requests.RequestException as e:
                 last_exception = e
 
         raise requests.RequestException("Token refresh failed after retries") from last_exception
 
-    def __is_token_expired(self, token: str) -> bool:
+    def __is_token_expired(self, token: Optional[str]) -> bool:
         try:
             decoded_token = jwt.decode(token, options={"verify_signature": False})
             exp = decoded_token.get("exp")
-            if exp:
-                return time.time() > (exp + self.EXPIRATION_LEEWAY.total_seconds())
-        except jwt.ExpiredSignatureError:
-            return True
-        except jwt.DecodeError:
+            if exp is None:
+                return True
+            return time.time() > (float(exp) - self.EXPIRATION_LEEWAY.total_seconds())
+        except (jwt.InvalidTokenError, TypeError, ValueError):
             return True
-        return False
 
     def __shutdown(self):
         self.executor.shutdown(wait=False)

core/tests/core/test_auth.py

@@ -1,11 +1,14 @@
+from concurrent.futures import ThreadPoolExecutor
 from pathlib import Path, PurePath
+from threading import Event, Lock, Thread
 
 import pytest
 import json
 import jwt
 import requests
 from unittest.mock import patch, mock_open, Mock
 
+from requests import Request
 from requests.auth import HTTPBasicAuth
 
 from stackit.core.auth_methods.key_auth import KeyAuth, ServiceAccountKey
@@ -294,3 +297,88 @@ def set_initial_token(auth):
                 auth._KeyAuth__refresh_token()
 
             assert mock_post.call_count == KeyAuth.MAX_REFRESH_RETRIES
+
+
+class TestKeyAuth:
+    @pytest.mark.parametrize("refresh_already_failed", [False, True])
+    def test_call_propagates_token_refresh_failure(self, refresh_already_failed):
+        auth = object.__new__(KeyAuth)
+        auth.lock = Lock()
+        auth.access_token = jwt.encode({"exp": 0}, "x" * 32, algorithm="HS256")
+        auth.refresh_token = jwt.encode({"exp": 4_000_000_000}, "x" * 32, algorithm="HS256")
+        auth.token_endpoint = KeyAuth.DEFAULT_TOKEN_ENDPOINT
+        auth.refresh_future = None
+        original_token = auth.access_token
+        request = Request("GET", "https://example.com")
+        refresh_error = requests.RequestException("refresh failed")
+
+        with (
+            ThreadPoolExecutor(max_workers=1) as executor,
+            patch("stackit.core.auth_methods.key_auth.requests.post", side_effect=refresh_error) as mock_post,
+        ):
+            auth.executor = executor
+            if refresh_already_failed:
+                auth.refresh_future = executor.submit(auth._KeyAuth__refresh_token)
+                with pytest.raises(requests.RequestException, match="Token refresh failed after retries"):
+                    auth.refresh_future.result(timeout=1)
+
+            with pytest.raises(requests.RequestException, match="Token refresh failed after retries") as exc_info:
+                auth(request)
+
+            assert exc_info.value.__cause__ is refresh_error
+            assert auth.refresh_future.done()
+            assert mock_post.call_count == KeyAuth.MAX_REFRESH_RETRIES
+            mock_post.assert_called_with(
+                auth.token_endpoint,
+                data={"grant_type": "refresh_token", "refresh_token": auth.refresh_token},
+                timeout=auth.timeout,
+            )
+            assert auth.access_token == original_token
+            assert "Authorization" not in request.headers
+
+    def test_token_is_expired_before_expiration_with_leeway(self):
+        auth = object.__new__(KeyAuth)
+        secret = "x" * 32
+
+        with patch("stackit.core.auth_methods.key_auth.time.time", return_value=1_000):
+            token_inside_leeway = jwt.encode({"exp": 1_240}, secret, algorithm="HS256")
+            token_outside_leeway = jwt.encode({"exp": 1_360}, secret, algorithm="HS256")
+            token_without_expiration = jwt.encode({}, secret, algorithm="HS256")
+            token_with_zero_expiration = jwt.encode({"exp": 0}, secret, algorithm="HS256")
+
+            assert auth._KeyAuth__is_token_expired(token_inside_leeway)
+            assert not auth._KeyAuth__is_token_expired(token_outside_leeway)
+            assert auth._KeyAuth__is_token_expired(token_without_expiration)
+            assert auth._KeyAuth__is_token_expired(token_with_zero_expiration)
+
+    def test_call_waits_for_an_in_progress_refresh_before_setting_header(self):
+        auth = object.__new__(KeyAuth)
+        auth.lock = Lock()
+        auth.access_token = jwt.encode({"exp": 0}, "x" * 32, algorithm="HS256")
+        fresh_token = jwt.encode({"exp": 4_000_000_000}, "x" * 32, algorithm="HS256")
+        refresh_started = Event()
+        release_refresh = Event()
+
+        class BlockingFuture:
+            def done(self):
+                return False
+
+            def result(self):
+                refresh_started.set()
+                release_refresh.wait(timeout=1)
+                with auth.lock:
+                    auth.access_token = fresh_token
+
+        auth.refresh_future = BlockingFuture()
+        request = Request("GET", "https://example.com")
+        call_thread = Thread(target=auth, args=(request,))
+        call_thread.start()
+
+        assert refresh_started.wait(timeout=1)
+        assert call_thread.is_alive()
+
+        release_refresh.set()
+        call_thread.join(timeout=1)
+
+        assert not call_thread.is_alive()
+        assert request.headers["Authorization"] == f"Bearer {fresh_token}"