dara-core 1.13.1__py3-none-any.whl → 1.14.0__py3-none-any.whl

dara/core/auth/base.py

@@ -18,7 +18,7 @@ limitations under the License.
 import abc
 from typing import Any, ClassVar, Dict, Union
 
-from fastapi import Response
+from fastapi import HTTPException, Response
 from pydantic import BaseModel
 from typing_extensions import TypedDict
 
@@ -91,6 +91,18 @@ class BaseAuthConfig(BaseModel, abc.ABC):
         :param token: encoded token
         """
 
+    def refresh_token(self, old_token: TokenData, refresh_token: str) -> tuple[str, str]:
+        """
+        Create a new session token and refresh token from a refresh token.
+
+        Note: the new issued session token should include the same session_id as the old token
+
+        :param old_token: old session token data
+        :param refresh_token: encoded refresh token
+        :return: new session token, new refresh token
+        """
+        raise HTTPException(400, f'Auth config {self.__class__.__name__} does not support token refresh')
+
     def revoke_token(self, token: str, response: Response) -> Union[SuccessResponse, RedirectResponse]:
         """
         Revoke a session token.

dara/core/auth/routes.py

@@ -16,9 +16,18 @@ limitations under the License.
 """
 
 from inspect import iscoroutinefunction
+from typing import Union, cast
 
 import jwt
-from fastapi import APIRouter, Depends, HTTPException, Request, Response
+from fastapi import (
+    APIRouter,
+    BackgroundTasks,
+    Cookie,
+    Depends,
+    HTTPException,
+    Request,
+    Response,
+)
 from fastapi.security import HTTPAuthorizationCredentials, HTTPBearer
 
 from dara.core.auth.base import BaseAuthConfig
@@ -31,6 +40,7 @@ from dara.core.auth.definitions import (
     AuthError,
     SessionRequestBody,
 )
+from dara.core.auth.utils import cached_refresh_token, decode_token
 from dara.core.logging import dev_logger
 
 auth_router = APIRouter()
@@ -103,6 +113,71 @@ async def _revoke_session(response: Response, credentials: HTTPAuthorizationCred
     raise HTTPException(status_code=400, detail=BAD_REQUEST_ERROR('No auth credentials passed'))
 
 
+@auth_router.post('/refresh-token')
+async def handle_refresh_token(
+    response: Response,
+    background_tasks: BackgroundTasks,
+    dara_refresh_token: Union[str, None] = Cookie(default=None),
+    credentials: HTTPAuthorizationCredentials = Depends(HTTPBearer()),
+):
+    """
+    Given a refresh token, issues a new session token and refresh token cookie.
+
+    :param response: FastAPI response object
+    :param dara_refresh_token: refresh token cookie
+    :param settings: env settings object
+    """
+    if dara_refresh_token is None:
+        raise HTTPException(status_code=400, detail=BAD_REQUEST_ERROR('No refresh token provided'))
+
+    # Check scheme is correct
+    if credentials.scheme != 'Bearer':
+        raise HTTPException(
+            status_code=400,
+            detail=BAD_REQUEST_ERROR(
+                'Invalid authentication scheme, previous Bearer token must be included in the refresh request'
+            ),
+        )
+
+    from dara.core.internal.registries import auth_registry
+
+    auth_config: BaseAuthConfig = auth_registry.get('auth_config')
+
+    try:
+        # decode the old token ignoring expiry date
+        old_token_data = decode_token(credentials.credentials, options={'verify_exp': False})
+
+        # Refresh logic up to implementation - passing in old token data so session_id can be preserved
+        session_token, refresh_token = await cached_refresh_token(
+            auth_config.refresh_token, old_token_data, dara_refresh_token
+        )
+
+        # Using 'Strict' as it is only used for the refresh-token endpoint so cross-site requests are not expected
+        response.set_cookie(
+            key='dara_refresh_token', value=refresh_token, secure=True, httponly=True, samesite='strict'
+        )
+        return {'token': session_token}
+    except BaseException as e:
+        # Regardless of exception type, clear the refresh token cookie
+        response.delete_cookie('dara_refresh_token')
+        headers = {'set-cookie': response.headers['set-cookie']}
+
+        # If an explicit HTTPException was raised, re-raise it with the cookie header
+        if isinstance(e, HTTPException):
+            dev_logger.error('Auth Error', error=e)
+            e.headers = headers
+            raise e
+
+        # Explicitly handle expired signature error
+        if isinstance(e, jwt.ExpiredSignatureError):
+            dev_logger.error('Expired Token Signature', error=e)
+            raise HTTPException(status_code=401, detail=EXPIRED_TOKEN_ERROR, headers=headers)
+
+        # Otherwise show a generic invalid token error
+        dev_logger.error('Invalid Token', error=cast(Exception, e))
+        raise HTTPException(status_code=401, detail=INVALID_TOKEN_ERROR, headers=headers)
+
+
 # Request to retrieve a session token from the backend. The app does this on startup.
 @auth_router.post('/session')
 async def _get_session(body: SessionRequestBody):

dara/core/auth/utils.py

@@ -15,11 +15,13 @@ See the License for the specific language governing permissions and
 limitations under the License.
 """
 
+import asyncio
 import uuid
 from datetime import datetime, timedelta, timezone
-from typing import List, Optional, Union
+from typing import Any, Callable, Dict, List, Optional, Tuple, Union
 
 import jwt
+from anyio import to_thread
 
 from dara.core.auth.definitions import (
     EXPIRED_TOKEN_ERROR,
@@ -33,12 +35,15 @@ from dara.core.internal.settings import get_settings
 from dara.core.logging import dev_logger
 
 
-def decode_token(token: str) -> TokenData:
+def decode_token(token: str, **kwargs) -> TokenData:
     """
     Decode a JWT token
+
+    :param token: the JWT token to decode
+    :param kwargs: additional arguments to pass to the jwt.decode function
     """
     try:
-        return TokenData.parse_obj(jwt.decode(token, get_settings().jwt_secret, algorithms=[JWT_ALGO]))
+        return TokenData.parse_obj(jwt.decode(token, get_settings().jwt_secret, algorithms=[JWT_ALGO], **kwargs))
     except jwt.ExpiredSignatureError:
         raise AuthError(code=401, detail=EXPIRED_TOKEN_ERROR)
     except jwt.DecodeError:
@@ -52,11 +57,13 @@ def sign_jwt(
     groups: List[str],
     id_token: Optional[str] = None,
     exp: Optional[Union[datetime, int]] = None,
+    session_id: Optional[str] = None,
 ):
     """
     Create a new Dara JWT token
     """
-    session_id = str(uuid.uuid4())
+    if session_id is None:
+        session_id = str(uuid.uuid4())
 
     # Default expiry is 1 day unless specified
     if exp is None:
@@ -95,3 +102,120 @@ def get_user_data():
         )
 
     return user_data
+
+
+class AsyncTokenRefreshCache:
+    """
+    An asynchronous cache for token refresh operations that handles concurrent requests
+    and provides time-based cache invalidation.
+
+    This cache is designed to prevent multiple simultaneous refresh attempts with the
+    same refresh token, while also providing a short-term cache to reduce unnecessary
+    token refreshes from multiple tabs/windows.
+    """
+
+    def __init__(self, ttl_seconds: int = 5):
+        self.cache: Dict[str, Tuple[Any, datetime]] = {}
+        self.locks: Dict[str, asyncio.Lock] = {}
+        self.locks_lock = asyncio.Lock()
+        self.ttl = timedelta(seconds=ttl_seconds)
+
+    async def _get_or_create_lock(self, key: str) -> asyncio.Lock:
+        """
+        Get an existing lock for the given key or create a new one if it doesn't exist.
+
+        This method is thread-safe and ensures that only one lock exists per key.
+
+        :param key: The key to get or create a lock for.
+        """
+
+        async with self.locks_lock:
+            if key not in self.locks:
+                self.locks[key] = asyncio.Lock()
+            return self.locks[key]
+
+    def _cleanup_old_entries(self):
+        """
+        Remove expired entries from both the cache and locks dictionaries.
+
+        This method is called before each cache access to prevent memory leaks
+        from accumulated expired entries.
+        """
+        current_time = datetime.now()
+        expired_keys = [key for key, (_, timestamp) in self.cache.items() if current_time - timestamp > self.ttl]
+        for key in expired_keys:
+            self.cache.pop(key, None)
+            # We can modify self.locks here because we're always under an async lock when calling this
+            self.locks.pop(key, None)
+
+    def get_cached_value(self, key: str) -> Tuple[Any, bool]:
+        """
+        Retrieve a value from the cache if it exists and hasn't expired.
+
+        :param key: The key to retrieve from the cache.
+        :return: A tuple containing the value and a boolean indicating whether the value was found.
+        """
+        self._cleanup_old_entries()
+        if key in self.cache:
+            value, timestamp = self.cache[key]
+            if datetime.now() - timestamp <= self.ttl:
+                return value, True
+        return None, False
+
+    def set_cached_value(self, key: str, value: Any):
+        """
+        Set a value in the cache with the current timestamp.
+
+        :param key: The key to set in the cache.
+        :param value: The value to set in the cache.
+        """
+        self.cache[key] = (value, datetime.now())
+
+    def clear(self):
+        """
+        Clear the cache and locks dictionaries.
+        """
+        self.cache.clear()
+        self.locks.clear()
+
+
+token_refresh_cache = AsyncTokenRefreshCache(ttl_seconds=5)
+"""
+Shared token refresh cache instance
+"""
+
+
+async def cached_refresh_token(
+    do_refresh_token: Callable[[TokenData, str], Tuple[str, str]], old_token_data: TokenData, refresh_token: str
+):
+    """
+    A utility to run a token refresh method with caching to prevent multiple concurrent refreshes
+    and short-term caching to reduce unnecessary refreshes from multiple tabs/windows.
+
+    :param do_refresh_token: The function to perform the token refresh
+    :param old_token_data: The old token data
+    :param refresh_token: The refresh token to use
+    """
+    cache_key = refresh_token
+
+    # check for cache hit
+    cached_result, found = token_refresh_cache.get_cached_value(cache_key)
+    if found:
+        return cached_result
+
+    # cache miss, acquire lock so only one call for given refresh_token is allowed
+    lock = await token_refresh_cache._get_or_create_lock(cache_key)
+
+    async with lock:
+        # check cache again in case another call already refreshed the token while we were waiting
+        cached_result, found = token_refresh_cache.get_cached_value(cache_key)
+        if found:
+            return cached_result
+
+        # Run the refresh function
+        result = await to_thread.run_sync(do_refresh_token, old_token_data, refresh_token)
+
+        # update cache
+        token_refresh_cache.set_cached_value(cache_key, result)
+
+        return result

dara/core/internal/routing.py

@@ -208,9 +208,11 @@ def create_router(config: Configuration):
             'application_name': get_settings().project_name,
         }
 
-    @core_api_router.get('/auth-components')
-    async def get_auth_components():  # pylint: disable=unused-variable
-        return config.auth_config.component_config
+    @core_api_router.get('/auth-config')
+    async def get_auth_config():  # pylint: disable=unused-variable
+        return {
+            'auth_components': config.auth_config.component_config.dict(),
+        }
 
     @core_api_router.get('/components', dependencies=[Depends(verify_session)])
     async def get_components(name: Optional[str] = None):  # pylint: disable=unused-variable

dara/core/internal/websocket.py

@@ -167,9 +167,10 @@ class WebSocketHandler:
 
     def __init__(self, channel_id: str):
         send_stream, receive_stream = create_memory_object_stream[ServerMessage](math.inf)
-        self.channel_id = channel_id
-        self.send_stream = send_stream
         self.receive_stream = receive_stream
+        self.send_stream = send_stream
+
+        self.channel_id = channel_id
         self.pending_responses = {}
 
     async def send_message(self, message: ServerMessage):
@@ -446,17 +447,20 @@ async def ws_handler(websocket: WebSocket, token: Optional[str] = Query(default=
     else:
         sessions_registry.set(user_identifier, {token_content.session_id})
 
-    # Set Auth context vars for the WS connection
-    USER.set(
-        UserData(
-            identity_id=token_content.identity_id,
-            identity_name=token_content.identity_name,
-            identity_email=token_content.identity_email,
-            groups=token_content.groups,
+    def update_context(token_data: TokenData):
+        USER.set(
+            UserData(
+                identity_id=token_data.identity_id,
+                identity_name=token_data.identity_name,
+                identity_email=token_data.identity_email,
+                groups=token_data.groups,
+            )
         )
-    )
-    SESSION_ID.set(token_content.session_id)
-    ID_TOKEN.set(token_content.id_token)
+        SESSION_ID.set(token_data.session_id)
+        ID_TOKEN.set(token_data.id_token)
+
+    # Set initial Auth context vars for the WS connection
+    update_context(token_content)
 
     # Change protocol from http to ws - from this point exceptions can't be raised
     await websocket.accept()
@@ -491,6 +495,12 @@ async def ws_handler(websocket: WebSocket, token: Optional[str] = Query(default=
                         # Heartbeat to keep connection alive
                         if data['type'] == 'ping':
                             await websocket.send_json({'type': 'pong', 'message': None})
+                        elif data['type'] == 'token_update':
+                            try:
+                                # update Auth context vars for the WS connection
+                                update_context(decode_token(data['message']))
+                            except Exception as e:
+                                eng_logger.error('Error updating token data', error=e)
                         else:
                             try:
                                 parsed_data = parse_obj_as(ClientMessage, data)