From fabc0de700e4fbb8d052cd154eedfcbefa9ed0a4 Mon Sep 17 00:00:00 2001 From: Carlo Goetz Date: Mon, 31 Aug 2026 14:39:38 +0200 Subject: [PATCH] fix(auth): refresh token before expiration, use locks while refreshing --- .../src/stackit/core/auth_methods/key_auth.py | 37 +++++++++----- core/tests/core/test_auth.py | 51 +++++++++++++++++++ 2 files changed, 76 insertions(+), 12 deletions(-) diff --git a/core/src/stackit/core/auth_methods/key_auth.py b/core/src/stackit/core/auth_methods/key_auth.py index 8cd135470..6f24eaff2 100644 --- a/core/src/stackit/core/auth_methods/key_auth.py +++ b/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) diff --git a/core/tests/core/test_auth.py b/core/tests/core/test_auth.py index cf679b7d9..ec5ce96f7 100644 --- a/core/tests/core/test_auth.py +++ b/core/tests/core/test_auth.py @@ -1,4 +1,5 @@ from pathlib import Path, PurePath +from threading import Event, Lock, Thread import pytest import json @@ -6,6 +7,7 @@ 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 +296,52 @@ def set_initial_token(auth): auth._KeyAuth__refresh_token() assert mock_post.call_count == KeyAuth.MAX_REFRESH_RETRIES + + +class TestKeyAuth: + 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}"