feat: integration of Keycloak

Author: JanssenBrm

app/auth.py

@@ -1,17 +1,69 @@
-from fastapi_keycloak import FastAPIKeycloak
+import jwt
+from fastapi import Depends, HTTPException, WebSocket, status
+from fastapi.security import OAuth2AuthorizationCodeBearer
+from fastapi_keycloak import OIDCUser
+from jwt import PyJWKClient
+from loguru import logger
+
 from .config.settings import settings
 
-# create the FastAPIKeycloak instance — used to protect routes
-# The server_url must include trailing slash for library
-keycloak = FastAPIKeycloak(
-    server_url=str(settings.keycloak_server_url),
-    client_id=settings.keycloak_client_id,
-    client_secret=settings.keycloak_client_secret,
-    admin_client_secret=settings.keycloak_client_secret,  # optional for admin operations
-    realm=settings.keycloak_realm,
-    callback_uri="http://localhost:8000/callback",  # for auth code flow if needed
+# Keycloak OIDC info
+KEYCLOAK_BASE_URL = f"https://{settings.keycloak_host}/realms/{settings.keycloak_realm}"
+JWKS_URL = f"{KEYCLOAK_BASE_URL}/protocol/openid-connect/certs"
+ALGORITHM = "RS256"
+
+
+# Keycloak OIDC endpoints
+oauth2_scheme = OAuth2AuthorizationCodeBearer(
+    authorizationUrl=f"https://{settings.keycloak_host}/realms/{settings.keycloak_realm}/"
+    "protocol/openid-connect/auth",
+    tokenUrl=f"https://{settings.keycloak_host}/realms/{settings.keycloak_realm}/"
+    "protocol/openid-connect/token",
 )
 
-# expose a helper dependency for current user
-get_current_user = keycloak.get_current_user
-get_current_active_user = keycloak.get_current_user
+# PyJWT helper to fetch and cache keys
+jwks_client = PyJWKClient(JWKS_URL, cache_keys=True)
+
+
+def _decode_token(token: str):
+    try:
+        signing_key = jwks_client.get_signing_key_from_jwt(token).key
+        payload = jwt.decode(
+            token,
+            signing_key,
+            algorithms=[ALGORITHM],
+            issuer=KEYCLOAK_BASE_URL,
+        )
+        return payload
+    except Exception:
+        raise HTTPException(
+            status_code=status.HTTP_401_UNAUTHORIZED,
+            detail="Could not validate credentials",
+        )
+
+
+def get_current_user_id(token: str = Depends(oauth2_scheme)):
+    user: OIDCUser = _decode_token(token)
+    return user["sub"]
+
+
+async def websocket_authenticate(websocket: WebSocket) -> str | None:
+    """
+    Authenticate a WebSocket connection using a JWT token from query params.
+    Returns the ID of the authenticated user payload if valid, otherwise closes the connection.
+    """
+    logger.debug("Authenticating websocket")
+    token = websocket.query_params.get("token")
+    if not token:
+        logger.error("Token is missing from websocket authentication")
+        await websocket.close(code=1008, reason="Missing token")
+        return None
+
+    try:
+        user_id = get_current_user_id(token)
+        await websocket.accept()
+        return user_id
+    except Exception as e:
+        logger.error(f"Invalid token in websocket authentication: {e}")
+        await websocket.close(code=1008, reason="Invalid token")
+        return None

app/config/settings.py

@@ -1,4 +1,4 @@
-from pydantic import AnyHttpUrl, Field
+from pydantic import Field
 from pydantic_settings import BaseSettings, SettingsConfigDict
 
 
@@ -13,9 +13,9 @@ class Settings(BaseSettings):
     env: str = Field(default="development", json_schema_extra={"env": "APP_ENV"})
 
     # Keycloak / OIDC
-    keycloak_server_url: AnyHttpUrl = Field(
-        default=AnyHttpUrl("https://localhost"),
-        json_schema_extra={"env": "KEYCLOAK_SERVER_URL"},
+    keycloak_host: str = Field(
+        default=str("localhost"),
+        json_schema_extra={"env": "KEYCLOAK_HOST"},
     )
     keycloak_realm: str = Field(default="", json_schema_extra={"env": "KEYCLOAK_REALM"})
     keycloak_client_id: str = Field(

app/main.py

@@ -3,9 +3,9 @@
 from app.middleware.correlation_id import add_correlation_id
 from app.platforms.dispatcher import load_processing_platforms
 from app.services.tiles.base import load_grids
-from .config.logger import setup_logging
-from .config.settings import settings
-from .routers import jobs_status, unit_jobs, health, tiles, upscale_tasks
+from app.config.logger import setup_logging
+from app.config.settings import settings
+from app.routers import jobs_status, unit_jobs, health, tiles, upscale_tasks
 
 setup_logging()
 
@@ -20,9 +20,6 @@
 
 app.middleware("http")(add_correlation_id)
 
-# Register Keycloak - must be done after FastAPI app creation
-# keycloak.register(app, prefix="/auth")  # mounts OIDC endpoints for login if needed
-
 # include routers
 app.include_router(tiles.router)
 app.include_router(jobs_status.router)

app/routers/jobs_status.py

@@ -10,6 +10,7 @@
 from app.schemas.websockets import WSStatusMessage
 from app.services.processing import get_processing_jobs_by_user_id
 from app.services.upscaling import get_upscaling_tasks_by_user_id
+from app.auth import get_current_user_id, websocket_authenticate
 
 router = APIRouter()
 
@@ -21,7 +22,7 @@
 )
 async def get_jobs_status(
     db: Session = Depends(get_db),
-    user: str = "foobar",
+    user: str = Depends(get_current_user_id),
 ) -> JobsStatusResponse:
     """
     Return combined list of upscaling tasks and processing jobs for the authenticated user.
@@ -37,14 +38,16 @@ async def get_jobs_status(
     "/ws/jobs_status",
 )
 async def ws_jobs_status(
-    websocket: WebSocket, user: str = "foobar", interval: int = 10
+    websocket: WebSocket,
+    interval: int = 10,
 ):
     """
     Return combined list of upscaling tasks and processing jobs for the authenticated user.
     """
 
-    await websocket.accept()
-    logger.debug(f"WebSocket connected for user {user}")
+    user = await websocket_authenticate(websocket)
+    if not user:
+        return
 
     await websocket.send_json(
         WSStatusMessage(type="init", message="Starting status stream").model_dump()

app/routers/unit_jobs.py

@@ -3,6 +3,7 @@
 from loguru import logger
 from sqlalchemy.orm import Session
 
+from app.auth import get_current_user_id
 from app.database.db import get_db
 from app.schemas.enum import OutputFormatEnum, ProcessTypeEnum
 from app.schemas.unit_job import (
@@ -99,7 +100,7 @@ async def create_unit_job(
         ),
     ],
     db: Session = Depends(get_db),
-    user: str = "foobar",
+    user: str = Depends(get_current_user_id),
 ) -> ProcessingJobSummary:
     """Create a new processing job with the provided data."""
     try:
@@ -118,7 +119,9 @@ async def create_unit_job(
     responses={404: {"description": "Processing job not found"}},
 )
 async def get_job(
-    job_id: int, db: Session = Depends(get_db), user: str = "foobar"
+    job_id: int,
+    db: Session = Depends(get_db),
+    user: str = Depends(get_current_user_id),
 ) -> ProcessingJob:
     job = get_processing_job_by_user_id(db, job_id, user)
     if not job:
@@ -136,7 +139,9 @@ async def get_job(
     responses={404: {"description": "Processing job not found"}},
 )
 async def get_job_results(
-    job_id: int, db: Session = Depends(get_db), user: str = "foobar"
+    job_id: int,
+    db: Session = Depends(get_db),
+    user: str = Depends(get_current_user_id),
 ) -> Collection | None:
     try:
         result = get_processing_job_results(db, job_id, user)

app/routers/upscale_tasks.py

@@ -14,6 +14,7 @@
 from loguru import logger
 from sqlalchemy.orm import Session
 
+from app.auth import get_current_user_id, websocket_authenticate
 from app.database.db import SessionLocal, get_db
 from app.schemas.enum import OutputFormatEnum, ProcessTypeEnum
 from app.schemas.unit_job import (
@@ -106,7 +107,7 @@ async def create_upscale_task(
     ],
     background_tasks: BackgroundTasks,
     db: Session = Depends(get_db),
-    user: str = "foobar",
+    user: str = Depends(get_current_user_id),
 ) -> UpscalingTaskSummary:
     """Create a new upscaling job with the provided data."""
     try:
@@ -133,7 +134,9 @@ async def create_upscale_task(
     responses={404: {"description": "Upscale task not found"}},
 )
 async def get_upscale_task(
-    task_id: int, db: Session = Depends(get_db), user: str = "foobar"
+    task_id: int,
+    db: Session = Depends(get_db),
+    user: str = Depends(get_current_user_id),
 ) -> UpscalingTask:
     job = get_upscaling_task_by_user_id(db, task_id, user)
     if not job:
@@ -151,10 +154,12 @@ async def get_upscale_task(
 async def ws_task_status(
     websocket: WebSocket,
     task_id: int,
-    user: str = "foobar",
     interval: int = 10,
 ):
-    await websocket.accept()
+    user = await websocket_authenticate(websocket)
+    if not user:
+        return
+
     logger.info("WebSocket connected", extra={"user": user, "task_id": task_id})
 
     try:

env.example

@@ -1,8 +1,6 @@
 # Keycloak
-KEYCLOAK_SERVER_URL=
+KEYCLOAK_HOST=
 KEYCLOAK_REALM=
-KEYCLOAK_CLIENT_ID=
-KEYCLOAK_CLIENT_SECRET=
 
 # App
 APP_NAME="APEx Dispatch API"