From 660a0acf7c0028e16072f396787e32971e14faac Mon Sep 17 00:00:00 2001 From: David Perl Date: Thu, 30 Jul 2026 15:05:55 +0200 Subject: [PATCH] feat: get auth from dispatch --- .../default/beamline_dispatch.py | 20 +++++++++++++-- src/aare/beamline_dispatch/protocols.py | 7 +++++- .../simulated/beamline_dispatch.py | 12 +++++++-- .../x06da/beamline_dispatch.py | 4 +-- .../x06sa/beamline_dispatch.py | 4 +-- .../x10sa/beamline_dispatch.py | 4 +-- src/aare/daq/auth.py | 25 ++++++------------- src/aare/daq/server.py | 21 +++++++--------- 8 files changed, 56 insertions(+), 41 deletions(-) diff --git a/src/aare/beamline_dispatch/default/beamline_dispatch.py b/src/aare/beamline_dispatch/default/beamline_dispatch.py index 84c91139..24aac23f 100644 --- a/src/aare/beamline_dispatch/default/beamline_dispatch.py +++ b/src/aare/beamline_dispatch/default/beamline_dispatch.py @@ -1,4 +1,20 @@ -from aare.beamline_dispatch.protocols import BeamlineDispatch +import os + +from aare.beamline_dispatch.protocols import AuthDispatch, BeamlineDispatch -class DefaultDispatch(BeamlineDispatch): ... +class DefaultAuthDispatch(AuthDispatch): + def get_jwt_key(self) -> str: + if (key := os.environ.get("JWT_AAREDAQ_KEY")) is None: + raise Exception( + "JWT_AAREDAQ_KEY environment variable not set, cannot guarantee safe authentication." + ) + return key + + +class DefaultDispatch(BeamlineDispatch): + """Default implementation for anything which can vary between beamlines and/or simulation. + Should be safe and fail rather than assuming anything.""" + + def __init__(self) -> None: + self.auth = DefaultAuthDispatch() diff --git a/src/aare/beamline_dispatch/protocols.py b/src/aare/beamline_dispatch/protocols.py index bcce4970..4aaddb41 100644 --- a/src/aare/beamline_dispatch/protocols.py +++ b/src/aare/beamline_dispatch/protocols.py @@ -1,4 +1,9 @@ from typing import Protocol -class BeamlineDispatch(Protocol): ... +class AuthDispatch(Protocol): + def get_jwt_key(self) -> str: ... + + +class BeamlineDispatch(Protocol): + auth: AuthDispatch diff --git a/src/aare/beamline_dispatch/simulated/beamline_dispatch.py b/src/aare/beamline_dispatch/simulated/beamline_dispatch.py index 47bb9989..0d285dfb 100644 --- a/src/aare/beamline_dispatch/simulated/beamline_dispatch.py +++ b/src/aare/beamline_dispatch/simulated/beamline_dispatch.py @@ -1,4 +1,12 @@ -from aare.beamline_dispatch.protocols import BeamlineDispatch +from aare.beamline_dispatch.default.beamline_dispatch import DefaultDispatch +from aare.beamline_dispatch.protocols import AuthDispatch -class SimulatedDispatch(BeamlineDispatch): ... +class SimulatedAuthDispatch(AuthDispatch): + def get_jwt_key(self) -> str: + return "Ns93ijN8VHv4ybvXaGNEDKUb3Sif4m4MYpfcEBIcs1h" + + +class SimulatedDispatch(DefaultDispatch): + def __init__(self) -> None: + self.auth = SimulatedAuthDispatch() diff --git a/src/aare/beamline_dispatch/x06da/beamline_dispatch.py b/src/aare/beamline_dispatch/x06da/beamline_dispatch.py index 89d52a33..7aab4be3 100644 --- a/src/aare/beamline_dispatch/x06da/beamline_dispatch.py +++ b/src/aare/beamline_dispatch/x06da/beamline_dispatch.py @@ -1,4 +1,4 @@ -from aare.beamline_dispatch.protocols import BeamlineDispatch +from aare.beamline_dispatch.default.beamline_dispatch import DefaultDispatch -class X06daDispatch(BeamlineDispatch): ... +class X06daDispatch(DefaultDispatch): ... diff --git a/src/aare/beamline_dispatch/x06sa/beamline_dispatch.py b/src/aare/beamline_dispatch/x06sa/beamline_dispatch.py index 67007ae9..fcae2cec 100644 --- a/src/aare/beamline_dispatch/x06sa/beamline_dispatch.py +++ b/src/aare/beamline_dispatch/x06sa/beamline_dispatch.py @@ -1,4 +1,4 @@ -from aare.beamline_dispatch.protocols import BeamlineDispatch +from aare.beamline_dispatch.default.beamline_dispatch import DefaultDispatch -class X06saDispatch(BeamlineDispatch): ... +class X06saDispatch(DefaultDispatch): ... diff --git a/src/aare/beamline_dispatch/x10sa/beamline_dispatch.py b/src/aare/beamline_dispatch/x10sa/beamline_dispatch.py index f8391e4e..421dbbf3 100644 --- a/src/aare/beamline_dispatch/x10sa/beamline_dispatch.py +++ b/src/aare/beamline_dispatch/x10sa/beamline_dispatch.py @@ -1,4 +1,4 @@ -from aare.beamline_dispatch.protocols import BeamlineDispatch +from aare.beamline_dispatch.default.beamline_dispatch import DefaultDispatch -class X10saDispatch(BeamlineDispatch): ... +class X10saDispatch(DefaultDispatch): ... diff --git a/src/aare/daq/auth.py b/src/aare/daq/auth.py index f7fbf429..a6bf8099 100644 --- a/src/aare/daq/auth.py +++ b/src/aare/daq/auth.py @@ -9,6 +9,7 @@ from datetime import UTC, datetime, timedelta from typing import List import jwt +from aarecommon.config.beamline import MXBeamline, mx_beamline from aarecommon.errors.exception_handler import ( AuthenticationException, AuthErrorCode, @@ -20,15 +21,11 @@ from fastapi import Depends, Request from fastapi.security import OAuth2PasswordBearer from pydantic import BaseModel +from aare.beamline_dispatch.protocols import AuthDispatch from aare.daq.config import BeamlineConfig logger = logging.getLogger("aareDAQ") -if os.environ.get("JWT_AAREDAQ_KEY") is None: - raise Exception( - "JWT_AAREDAQ_KEY environment variable not set, cannot guarantee safe authentication." - ) -SECRET_KEY = os.environ.get("JWT_AAREDAQ_KEY") ALGORITHM = "HS256" ACCESS_TOKEN_EXPIRE_MINUTES = 24 * 60 * 7 # 1 week @@ -48,6 +45,11 @@ class TokenData(BaseModel): staff: bool = False +def init_jwt_key(dispatch: AuthDispatch): + global SECRET_KEY + SECRET_KEY = dispatch.get_jwt_key() + + def create_access_token(token: TokenData): to_encode = token.model_dump() to_encode.update({"exp": datetime.now(UTC) + timedelta(minutes=ACCESS_TOKEN_EXPIRE_MINUTES)}) @@ -55,19 +57,6 @@ def create_access_token(token: TokenData): return encoded_jwt -def _is_loopback(host: str | None) -> bool: - if host is None: - return False - try: - addr = ipaddress.ip_address(host) - # Unwrap IPv4-mapped IPv6 (e.g. ::ffff:127.0.0.1) before checking - if isinstance(addr, ipaddress.IPv6Address) and addr.ipv4_mapped is not None: - return addr.ipv4_mapped.is_loopback - return addr.is_loopback - except ValueError: - return False - - def authenticate_from_proxy_header(request: Request) -> str: """ Extract the pre-authenticated username from the trusted Apache proxy header. diff --git a/src/aare/daq/server.py b/src/aare/daq/server.py index 068304a6..5cfc0710 100644 --- a/src/aare/daq/server.py +++ b/src/aare/daq/server.py @@ -3,7 +3,6 @@ import hmac import importlib import json import os -import random import time from contextlib import asynccontextmanager from typing import AsyncGenerator, Optional @@ -48,6 +47,8 @@ from fastapi.security import OAuth2PasswordBearer from starlette.responses import StreamingResponse from uvicorn.workers import UvicornWorker # deprecated shim, present in pinned 0.34.2 +from aare.beamline_dispatch.beamline_dispatch import get_beamline_dispatch +from aare.beamline_dispatch.protocols import BeamlineDispatch from aare.daq import auth from aare.daq.config import BeamlineConfig from aare.daq.config_model import LocalContactConfigModel @@ -60,9 +61,9 @@ logger = setup_logger("aareDAQ") oauth2_scheme = OAuth2PasswordBearer(tokenUrl="token") # ── Per-worker state: populated inside the lifespan, after fork ── -bl = None -cfg = None -daq = None +cfg: BeamlineConfig +daq: AareDAQ +bl_dispatch: BeamlineDispatch _all_pgroups_cache: dict[str, tuple[list[str], float]] = {} _ALL_PGROUPS_TTL_S = 60.0 # adjust TTL as needed @@ -93,9 +94,11 @@ async def lifespan(application: FastAPI): All stateful / connection-opening initialisation belongs here so that each worker gets its own fresh Redis, BEC, EPICS, and TELL connections. """ - global bl, cfg, daq - await asyncio.sleep(random.uniform(0.5, 3.0)) + global cfg, daq, bl_dispatch + logger.info(f"Worker {os.getpid()} setting up JWT authentication...") + bl_dispatch = get_beamline_dispatch() + auth.init_jwt_key(bl_dispatch.auth) logger.info(f"Worker {os.getpid()} starting initialisation...") # ── Core objects (Redis, EPICS PVs, BEC, TELL, JFJoch, etc.) ── @@ -203,9 +206,6 @@ def _get_automation_progress_state() -> dict: """ Read automation progress state from shared config/Redis storage. """ - if cfg is None: - return {"seq": 0, "progress": None} - try: return cfg.get_automation_progress_state() except Exception as e: @@ -220,9 +220,6 @@ def _push_automation_progress(progress: AutomationProgress) -> None: Args: progress: Current automation progress model. """ - if cfg is None: - return - try: state = cfg.set_automation_progress_state(progress) logger.info(