backend: Add refresh token logic to OAuth providers

docs/auth_guide.md

@@ -20,10 +20,21 @@ import secrets
 print(secrets.token_hex(32))
 ```
 
-## Configuring your OAuth app's Redirect URI
+## Configuring your OAuth app
+
+### Redirect URI
 
 When configuring your OAuth apps, make sure to whitelist the Redirect URI to the frontend endpoint, it should look like 
-`<FRONTEND_HOST>/auth/<STRATEGY_NAME>`. For example, your Redirect URI will be `http://localhost:4000/auth/google` if you're running the GoogleOAuth class locally.
+`<FRONTEND_HOST>/auth/<STRATEGY_NAME>`. For example, your Redirect URI will be `http://localhost:4000/auth/google` if you're running the GoogleOAuth class locally. The strategy name is defined in the `NAME` class attribute.
+
+### Enabling Refresh Tokens
+
+To enable refresh tokens, you must implement the `get_refresh_token_params()` method in your auth strategy class. This should return a dictionary containing key-value pairs that contain the query parameters the auth provider needs to return a refresh token. For example, if your auth provider requires a `?scope=offline` query parameter, you should add:
+
+```python
+def get_refresh_token_params(self):
+    return {"scope": "offline"}
+```
 
 ## Enabling Proof of Key Code Exchange (PKCE)
 

src/backend/config/auth.py

@@ -1,18 +1,23 @@
 import os
 import sys
+from typing import Union
 
 from dotenv import load_dotenv
 from fastapi import Depends, HTTPException, status
 from fastapi.security import HTTPAuthorizationCredentials, HTTPBearer
 
 from backend.services.auth import BasicAuthentication, GoogleOAuth, OpenIDConnect
+from backend.services.auth.strategies.base import (
+    BaseAuthenticationStrategy,
+    BaseOAuthStrategy,
+)
 
 load_dotenv()
 
 SKIP_AUTH = os.getenv("SKIP_AUTH", None)
 # Add Auth strategy classes here to enable them
 # Ex: [BasicAuthentication]
-ENABLED_AUTH_STRATEGIES = []
+ENABLED_AUTH_STRATEGIES = [BasicAuthentication, GoogleOAuth]
 if "pytest" in sys.modules or SKIP_AUTH == "true":
     ENABLED_AUTH_STRATEGIES = []
 
@@ -49,6 +54,15 @@ def is_authentication_enabled() -> bool:
     return False
 
 
+def get_auth_strategy(
+    strategy_name: str,
+) -> Union[BaseAuthenticationStrategy, BaseOAuthStrategy]:
+    if strategy_name not in ENABLED_AUTH_STRATEGY_MAPPING.keys():
+        return None
+
+    return ENABLED_AUTH_STRATEGY_MAPPING[strategy_name]
+
+
 async def get_auth_strategy_endpoints() -> None:
     """
     Fetches the endpoints for each enabled strategy.

src/backend/main.py

@@ -24,6 +24,7 @@
 from backend.routers.snapshot import router as snapshot_router
 from backend.routers.tool import router as tool_router
 from backend.routers.user import router as user_router
+from backend.services.auth.request_validators import UPDATE_TOKEN_HEADER
 from backend.services.logger import LoggingMiddleware, get_logger
 from backend.services.metrics import MetricsMiddleware
 
@@ -75,6 +76,7 @@ def create_app():
         allow_credentials=True,
         allow_methods=["*"],
         allow_headers=["*"],
+        expose_headers=[UPDATE_TOKEN_HEADER],
     )
     app.add_middleware(LoggingMiddleware)
     app.add_middleware(MetricsMiddleware)

src/backend/middleware/__init__.py

@@ -0,0 +1,5 @@
+from backend.middleware.logging import LoggingMiddleware
+
+__all__ = [
+    "LoggingMiddleware",
+]

src/backend/middleware/logging.py

@@ -0,0 +1,14 @@
+import logging
+import time
+
+from starlette.middleware.base import BaseHTTPMiddleware
+from starlette.requests import Request
+
+
+class LoggingMiddleware(BaseHTTPMiddleware):
+    async def dispatch(self, request: Request, call_next):
+        start_time = time.time()
+        response = await call_next(request)
+
+        logging.info(f"{request.method} {request.url.path}\n{request.headers}")
+        return response

src/backend/routers/auth.py

@@ -7,7 +7,7 @@
 from fastapi.responses import RedirectResponse
 from starlette.requests import Request
 
-from backend.config.auth import ENABLED_AUTH_STRATEGY_MAPPING
+from backend.config.auth import ENABLED_AUTH_STRATEGY_MAPPING, get_auth_strategy
 from backend.config.routers import RouterName
 from backend.config.tools import AVAILABLE_TOOLS
 from backend.crud import blacklist as blacklist_crud
@@ -49,6 +49,11 @@ def get_strategies() -> list[ListAuthStrategy]:
                     if hasattr(strategy_instance, "get_authorization_endpoint")
                     else None
                 ),
+                "refresh_token_params": (
+                    strategy_instance.get_refresh_token_params()
+                    if hasattr(strategy_instance, "get_refresh_token_params")
+                    else None
+                ),
                 "pkce_enabled": (
                     strategy_instance.get_pkce_enabled()
                     if hasattr(strategy_instance, "get_pkce_enabled")
@@ -80,13 +85,12 @@ async def login(request: Request, login: Login, session: DBSessionDep):
     strategy_name = login.strategy
     payload = login.payload
 
-    if not is_enabled_authentication_strategy(strategy_name):
+    strategy = get_auth_strategy(strategy_name)
+    if not strategy:
         raise HTTPException(
             status_code=422, detail=f"Invalid Authentication strategy: {strategy_name}."
         )
 
-    # Check that the payload required is given
-    strategy = ENABLED_AUTH_STRATEGY_MAPPING[strategy_name]
     strategy_payload = strategy.get_required_payload()
     if not set(strategy_payload).issubset(payload.keys()):
         missing_keys = [key for key in strategy_payload if key not in payload.keys()]
@@ -102,7 +106,7 @@ async def login(request: Request, login: Login, session: DBSessionDep):
             detail=f"Error performing {strategy_name} authentication with payload: {payload}.",
         )
 
-    token = JWTService().create_and_encode_jwt(user)
+    token = JWTService().create_and_encode_jwt(user, strategy_name)
 
     return {"token": token}
 
@@ -165,7 +169,7 @@ async def authorize(
     # Get or create user, then set session user
     user = get_or_create_user(session, userinfo)
 
-    token = JWTService().create_and_encode_jwt(user)
+    token = JWTService().create_and_encode_jwt(user, strategy_name)
 
     return {"token": token}
 

src/backend/services/auth/jwt.py

@@ -4,15 +4,23 @@
 import uuid
 
 import jwt
+from fastapi import Depends
+from sqlalchemy.orm import Session
 
+from backend.database_models import Blacklist, get_session
 from backend.services.logger import get_logger
 
 logger = get_logger()
 
 
 class JWTService:
     ISSUER = "cohere-toolkit"
-    EXPIRY_DAYS = 90
+    # AUTH_EXPIRY_MONTHS = 3
+    # JWT_EXPIRY_HOURS = 72
+    # REFRESH_AVAILABILITY_HOURS = 36
+    AUTH_EXPIRY_SECONDS = 60 * 2
+    JWT_EXPIRY_SECONDS = 30
+    REFRESH_AVAILABILITY_SECONDS = 15
     ALGORITHM = "HS256"
 
     def __init__(self):
@@ -25,29 +33,56 @@ def __init__(self):
 
         self.secret_key = secret_key
 
-    def create_and_encode_jwt(self, user: dict) -> str:
+    def create_and_encode_jwt(self, user: dict, strategy_name: str, **kwargs) -> str:
         """
         Creates a payload based on user info and creates a JWT token.
 
         Args:
             user (dict): User data.
+            strategy_name (str): Name of the authentication strategy.
+            kwargs: Additional payload data. These parameters will override normal payload data.
 
         Returns:
             str: JWT token.
         """
-        now = datetime.datetime.utcnow()
+        now = datetime.datetime.now(datetime.timezone.utc)
         payload = {
             "iss": self.ISSUER,
             "iat": now,
-            "exp": now + datetime.timedelta(days=self.EXPIRY_DAYS),
+            # "exp": now + datetime.timedelta(hours=self.AUTH_EXPIRY_HOURS),
+            "exp": now + datetime.timedelta(seconds=self.JWT_EXPIRY_SECONDS),
             "jti": str(uuid.uuid4()),
+            "strategy": strategy_name,
             "context": user,
+            **kwargs,
         }
 
         token = jwt.encode(payload, self.secret_key, self.ALGORITHM)
 
         return token
 
+    def refresh_jwt(self, token: str) -> str:
+        """
+        Refreshes a given JWT token. Callers should check if the token is expired before calling this method.
+
+        Args:
+            token (str): JWT token.
+
+        Returns:
+            str: New JWT token.
+        """
+        decoded_payload = self.decode_jwt(token)
+
+        if not decoded_payload:
+            return None
+
+        # Create new token with same payload
+        return self.create_and_encode_jwt(
+            decoded_payload["context"],
+            decoded_payload["strategy"],
+            iat=decoded_payload["iat"],
+        )
+
     def decode_jwt(self, token: str) -> dict:
         """
         Decodes a given JWT token.
@@ -64,8 +99,77 @@ def decode_jwt(self, token: str) -> dict:
             )
             return decoded_payload
         except jwt.ExpiredSignatureError:
-            logger.warning("Token has expired.")
-            return None
+            logger.warning("JWT Token has expired.")
+            decoded_payload = jwt.decode(
+                token,
+                self.secret_key,
+                algorithms=[self.ALGORITHM],
+                options={"verify_exp": False},
+            )
+            return decoded_payload
         except jwt.InvalidTokenError:
-            logger.warning("Invalid token.")
+            logger.warning("JWT Token is invalid.")
             return None
+
+    @staticmethod
+    def check_validity(payload: dict, session: Session = Depends(get_session)) -> str:
+        """
+        Checks a given JWT payload for its validity.
+
+        Args:
+            payload (dict): JWT payload.
+
+        Returns:
+            str: One of the following strings depending on the validity:
+                - "valid": Token is valid.
+                - "refreshable": Token is valid and within the refresh availability window.
+                - "expired": Token is expired or blacklisted.
+                - "invalid": Token is invalid.
+        """
+
+        if payload is None or any(
+            [
+                "context" not in payload,
+                "jti" not in payload,
+                "exp" not in payload,
+                "strategy" not in payload,
+                "iat" not in payload,
+            ]
+        ):
+            return "invalid"
+
+        now = datetime.datetime.now(datetime.timezone.utc)
+
+        # Check if token is blacklisted
+        blacklist = (
+            session.query(Blacklist)
+            .filter(Blacklist.token_id == payload["jti"])
+            .first()
+        )
+
+        if blacklist is not None:
+            logger.warning("JWT payload is blacklisted.")
+            return "expired"
+
+        # Check if token is expired; either we're past the expiry time or iat is more than 3 months ago
+        expiry_datetime = datetime.datetime.fromtimestamp(
+            payload["exp"], datetime.timezone.utc
+        )
+        issued_datetime = datetime.datetime.fromtimestamp(
+            payload["iat"], datetime.timezone.utc
+        )
+
+        # if now > expiry_datetime or now > (issued_datetime + datetime.timedelta(months=JWTService.AUTH_EXPIRY_MONTHS)):
+        if now > expiry_datetime or now > (
+            issued_datetime + datetime.timedelta(seconds=JWTService.AUTH_EXPIRY_SECONDS)
+        ):
+            return "expired"
+
+        # if now > (expiry_datetime - datetime.timedelta(hours=JWTService.REFRESH_AVAILABILITY_HOURS)):
+        if now > (
+            expiry_datetime
+            - datetime.timedelta(seconds=JWTService.REFRESH_AVAILABILITY_SECONDS)
+        ):
+            return "refreshable"
+
+        return "valid"