Refac/improve api auth

pyproject.toml

@@ -13,6 +13,7 @@ dependencies = [
   "loguru>=0.7.3",
   "pydantic>=2.11.7",
   "pydantic-settings>=2.10.1",
+  "pyjwt>=2.10.1",
   "pytz>=2025.2",
   "ray[serve]==2.53.0",
   "sentry-sdk[fastapi]>=2.33.0",

pytest.ini

@@ -29,6 +29,6 @@ env =
     ; plugins__proxy__force_stream_apis=[]
     sentry__enable=false
     test__silent=true
-    auth__rules={{"/api/v1/echo":["i_am_general_auth_keys"],"/api/v1/proxy/remote":["i_am_local_proxy_auth_keys"]}}
+    auth__rules={{"/api/v1/openapi.json":["i_am_docs_key"],"/api/v1/echo":["i_am_general_auth_keys"],"/api/v1/proxy/remote":["i_am_local_proxy_auth_keys"]}}
 
 asyncio_mode = auto

ruff.toml

@@ -79,7 +79,7 @@ ignore = [
 ]
 
 [lint.per-file-ignores]
-"tests/**/*.py" = ["S101", "ANN201", "ANN001", "ANN002"]
+"tests/**/*.py" = ["S101", "ANN201", "ANN001", "ANN002", "ANN003"]
 
 [lint.flake8-builtins]
 builtins-ignorelist = ["input", "id", "bytes", "type"]

src/framex/config.py

@@ -1,3 +1,4 @@
+import secrets
 from typing import Any, Literal
 
 from pydantic import BaseModel, Field
@@ -45,24 +46,44 @@ class ServerConfig(BaseModel):
     excluded_log_paths: list[str] = Field(default_factory=list)
     ingress_config: dict[str, Any] = Field(default_factory=lambda: {"max_ongoing_requests": 60})
 
-    # docs config
-    docs_user: str = "admin"
-    docs_password: str = ""
-
-    def model_post_init(self, __context: Any) -> None:  # pragma: no cover
-        if self.docs_password == "":
-            self.docs_password = "admin"  # noqa: S105
-            from framex.log import logger
-
-            logger.warning("No docs_password set, fallback to default password: admin")
-
 
 class TestConfig(BaseModel):
     disable_record_request: bool = False
     silent: bool = False
 
 
+class OauthConfig(BaseModel):
+    client_id: str = ""
+    client_secret: str = ""
+    authorization_url: str = ""
+    redirect_uri: str = ""
+    base_url: str = ""
+
+    token_url: str = ""
+    user_info_url: str = ""
+    app_url: str = ""
+
+    jwt_secret: str = ""
+    jwt_algorithm: str = "HS256"
+
+    @property
+    def call_back_url(self) -> str:
+        return f"{self.app_url}{self.redirect_uri}"
+
+    def model_post_init(self, context: Any) -> None:
+        super().model_post_init(context)
+        if not self.authorization_url:
+            self.authorization_url = f"{self.base_url}/oauth/authorize"
+        if not self.token_url:
+            self.token_url = f"{self.base_url}/oauth/token"
+        if not self.user_info_url:
+            self.user_info_url = f"{self.base_url}/api/v4/user"
+        if not self.jwt_secret:
+            self.jwt_secret = secrets.token_urlsafe(32)
+
+
 class AuthConfig(BaseModel):
+    oauth: OauthConfig | None = Field(default=None)
     rules: dict[str, list[str]] = Field(default_factory=dict)
 
     def _is_url_protected(self, url: str) -> bool:

src/framex/driver/application.py

@@ -1,7 +1,6 @@
 """Module containing FastAPI instance related functions and classes."""
 
 import json
-import secrets
 from collections.abc import Callable
 from contextlib import asynccontextmanager
 from datetime import UTC, datetime
@@ -13,7 +12,6 @@
 from fastapi.openapi.docs import get_redoc_html, get_swagger_ui_html
 from fastapi.openapi.utils import get_openapi
 from fastapi.responses import HTMLResponse
-from fastapi.security import HTTPBasic, HTTPBasicCredentials
 from starlette import status
 from starlette.concurrency import iterate_in_threadpool
 from starlette.exceptions import HTTPException
@@ -24,6 +22,7 @@
 
 from framex.config import settings
 from framex.consts import API_STR, DOCS_URL, OPENAPI_URL, PROJECT_NAME, REDOC_URL, VERSION
+from framex.driver.auth import authenticate, oauth_callback
 from framex.utils import format_uptime
 
 FRAME_START_TIME = datetime.now(tz=UTC)
@@ -95,22 +94,12 @@ async def _on_start(deployment: DeploymentHandle) -> None:
         redirect_slashes=False,
     )
 
-    security = HTTPBasic(realm="Swagger Docs")
-
-    def authenticate(credentials: HTTPBasicCredentials = Depends(security)) -> str:
-        correct_username = secrets.compare_digest(credentials.username, settings.server.docs_user)
-        correct_password = secrets.compare_digest(credentials.password, settings.server.docs_password)
-
-        if not (correct_username and correct_password):
-            from framex.log import logger
-
-            logger.warning(f"Failed authentication attempt for user: {credentials.username}")
-            raise HTTPException(
-                status_code=status.HTTP_401_UNAUTHORIZED,
-                detail="Incorrect user or password",
-                headers={"WWW-Authenticate": "Basic"},
-            )
-        return credentials.username
+    if settings.auth.oauth:
+        application.add_api_route(
+            settings.auth.oauth.redirect_uri,
+            oauth_callback,
+            methods=["GET"],
+        )
 
     @application.get(DOCS_URL, include_in_schema=False)
     async def get_documentation(_: Annotated[str, Depends(authenticate)]) -> HTMLResponse:

src/framex/driver/auth.py

@@ -0,0 +1,128 @@
+from datetime import UTC, datetime, timedelta
+
+import httpx
+import jwt
+from fastapi import Depends, HTTPException, Response, status
+from fastapi.responses import RedirectResponse
+from fastapi.security import APIKeyHeader
+from starlette.requests import Request
+
+from framex.config import settings
+from framex.consts import DOCS_URL
+
+api_key_header = APIKeyHeader(name="Authorization", auto_error=False)
+
+
+def create_jwt(payload: dict) -> str:
+    if not settings.auth.oauth:
+        raise RuntimeError("OAuth not configured")
+
+    now_utc = datetime.now(UTC)
+
+    payload.update(
+        {
+            "iat": int(now_utc.timestamp()),
+            "exp": int((now_utc + timedelta(hours=24)).timestamp()),
+        }
+    )
+    return jwt.encode(payload, settings.auth.oauth.jwt_secret, algorithm=settings.auth.oauth.jwt_algorithm)
+
+
+def auth_jwt(request: Request) -> bool:
+    if not settings.auth.oauth:
+        return False
+
+    token = request.cookies.get("token")
+    if not token:
+        return False
+
+    try:
+        jwt.decode(
+            token,
+            settings.auth.oauth.jwt_secret,
+            algorithms=[settings.auth.oauth.jwt_algorithm],
+        )
+        return True
+    except (jwt.InvalidTokenError, jwt.ExpiredSignatureError):
+        return False
+
+
+def authenticate(request: Request, api_key: str | None = Depends(api_key_header)) -> None:
+    if settings.auth.oauth:
+        if token := request.cookies.get("token"):
+            try:
+                jwt.decode(
+                    token,
+                    settings.auth.oauth.jwt_secret,
+                    algorithms=[settings.auth.oauth.jwt_algorithm],
+                )
+                return
+
+            except Exception as e:
+                from framex.log import logger
+
+                logger.warning(f"JWT decode failed: {e}")
+
+        if api_key and api_key in (settings.auth.get_auth_keys(request.url.path) or []):
+            return
+
+        raise HTTPException(
+            status_code=status.HTTP_302_FOUND,
+            headers={
+                "Location": (
+                    f"{settings.auth.oauth.authorization_url}"
+                    f"?client_id={settings.auth.oauth.client_id}"
+                    "&response_type=code"
+                    f"&redirect_uri={settings.auth.oauth.call_back_url}"
+                    "&scope=read_user"
+                )
+            },
+        )
+
+
+async def oauth_callback(code: str) -> Response:
+    if not settings.auth.oauth:  # pragma: no cover
+        raise RuntimeError("OAuth not configured")
+
+    async with httpx.AsyncClient() as client:
+        resp = await client.post(
+            settings.auth.oauth.token_url,
+            data={
+                "client_id": settings.auth.oauth.client_id,
+                "client_secret": settings.auth.oauth.client_secret,
+                "grant_type": "authorization_code",
+                "code": code,
+                "redirect_uri": settings.auth.oauth.call_back_url,
+            },
+        )
+        resp.raise_for_status()
+        token_resp = resp.json()
+
+        if not (auth_token := token_resp.get("access_token")):  # pragma: no cover
+            raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail="GitLab token exchange failed")
+
+        resp = await client.get(
+            settings.auth.oauth.user_info_url,
+            headers={"Authorization": f"Bearer {auth_token}"},
+            timeout=10,
+        )
+        resp.raise_for_status()
+        user_resp = resp.json()
+
+    if not (username := user_resp.get("username")):  # pragma: no cover
+        raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail="Failed to get GitLab user")
+
+    user_info = {
+        "message": f"Welcome {username}",
+        "username": username,
+        "email": user_resp.get("email"),
+    }
+
+    res = RedirectResponse(url=DOCS_URL, status_code=status.HTTP_302_FOUND)
+    res.set_cookie(
+        "token",
+        create_jwt(user_info),
+        httponly=True,
+        samesite="lax",
+    )
+    return res

src/framex/driver/ingress.py

@@ -2,24 +2,23 @@
 from enum import Enum
 from typing import Any
 
-from fastapi import Depends, HTTPException, Response, status
+from fastapi import Depends, HTTPException, Request, Response, status
 from fastapi.responses import JSONResponse, StreamingResponse
 from fastapi.routing import APIRoute
-from fastapi.security import APIKeyHeader
 from pydantic import create_model
 from ray.serve.handle import DeploymentHandle
 from starlette.routing import Route
 
 from framex.adapter import get_adapter
 from framex.consts import BACKEND_NAME
 from framex.driver.application import create_fastapi_application
+from framex.driver.auth import api_key_header, auth_jwt
 from framex.driver.decorator import api_ingress
 from framex.log import setup_logger
 from framex.plugin.model import ApiType, PluginApi
 from framex.utils import escape_tag, safe_error_message
 
 app = create_fastapi_application()
-api_key_header = APIKeyHeader(name="Authorization", auto_error=True)
 
 
 @app.get("/health")
@@ -114,8 +113,8 @@ async def route_handler(response: Response, model: Model = Depends()) -> Any:  #
             if auth_keys is not None:
                 logger.debug(f"API({path}) with tags {tags} requires auth.")
 
-                def _verify_api_key(api_key: str = Depends(api_key_header)) -> None:
-                    if api_key not in auth_keys:
+                def _verify_api_key(request: Request, api_key: str | None = Depends(api_key_header)) -> None:
+                    if (api_key is None or api_key not in auth_keys) and (not auth_jwt(request)):
                         logger.error(f"Unauthorized access attempt with API Key({api_key}) for API({path})")
                         raise HTTPException(
                             status_code=status.HTTP_401_UNAUTHORIZED,

src/framex/plugins/proxy/__init__.py

@@ -295,7 +295,7 @@ async def call_proxy_function(self, func_name: str, data: str) -> str:
             else:
                 kwargs = decode_kwargs
             tag = "remote" if proxy_func.is_remote else "local"
-            logger.info(f"Calling proxy function[{tag}]: {decode_func_name},  kwargs: {decode_kwargs}")
+            logger.info(f"Calling proxy function[{tag}]: {decode_func_name}")
             res = await proxy_func.func(**kwargs)
             return res if proxy_func.is_remote else cache_encode(res)
         raise RuntimeError(f"Proxy function({decode_func_name}) not registered")

tests/api/test_echo.py

@@ -16,8 +16,8 @@ def test_echo(client: TestClient):
 def test_echo_with_no_api_key(client: TestClient):
     params = {"message": "hello world"}
     res = client.get(f"{API_STR}/echo", params=params).json()
-    assert res["status"] == 403
-    assert res["message"] == "Not authenticated"
+    assert res["status"] == 401
+    assert res["message"] == "Invalid API Key(None) for API(/api/v1/echo)"
 
 
 def test_echo_with_error_api_key(client: TestClient):

tests/driver/test_application.py

@@ -8,7 +8,7 @@
 from starlette import status
 from starlette.exceptions import HTTPException
 
-from framex.consts import API_STR, DOCS_URL, OPENAPI_URL
+from framex.consts import API_STR
 from framex.driver.application import create_fastapi_application
 
 
@@ -33,45 +33,6 @@ def test_cors_middleware_configured(self):
         assert "CORSMiddleware" in middleware_classes
 
 
-class TestAuthenticationEndpoints:
-    @pytest.fixture
-    def app(self):
-        return create_fastapi_application()
-
-    @pytest.fixture
-    def client(self, app):
-        return TestClient(app)
-
-    def test_docs_requires_auth(self, client):
-        response = client.get(DOCS_URL)
-        assert response.status_code == status.HTTP_401_UNAUTHORIZED
-
-    def test_docs_wrong_credentials(self, client):
-        response = client.get(DOCS_URL, auth=("bad", "bad"))
-        assert response.status_code == status.HTTP_401_UNAUTHORIZED
-
-    def test_docs_correct_credentials(self, client):
-        from framex.config import settings
-
-        response = client.get(
-            DOCS_URL,
-            auth=(settings.server.docs_user, settings.server.docs_password),
-        )
-        assert response.status_code == status.HTTP_200_OK
-        assert "text/html" in response.headers["content-type"]
-
-    def test_openapi_correct_credentials(self, client):
-        from framex.config import settings
-
-        response = client.get(
-            OPENAPI_URL,
-            auth=(settings.server.docs_user, settings.server.docs_password),
-        )
-        assert response.status_code == status.HTTP_200_OK
-        data = response.json()
-        assert data["info"]["title"] == "FrameX API"
-
-
 class TestExceptionHandlers:
     @pytest.fixture
     def app(self):
@@ -146,15 +107,6 @@ async def endpoint() -> Any:
         assert data["data"] == {"result": "ok"}
         assert "timestamp" in data
 
-    def test_docs_not_wrapped(self, client):
-        from framex.config import settings
-
-        response = client.get(
-            DOCS_URL,
-            auth=(settings.server.docs_user, settings.server.docs_password),
-        )
-        assert "text/html" in response.headers["content-type"]
-
 
 class TestLifespanBehavior:
     @patch("framex.config.settings")