Require authentication for RL control-plane routes

lightllm/server/api_cli.py

@@ -923,7 +923,9 @@ def add_cli_args(parser: argparse.ArgumentParser) -> argparse.ArgumentParser:
         action="store_true",
         default=False,
         help="""enable RL control plane (HTTP APIs, router rl_rpyc, model RlBackendOps).
-        When disabled (default), RL routes/services are not started.""",
+        When disabled (default), RL routes/services are not started.
+        When enabled, set LIGHTLLM_RL_API_KEY and send it as an Authorization: Bearer token
+        to the RL control-plane routes.""",
     )
     parser.add_argument(
         "--enable_torch_memory_saver",

lightllm/server/api_http.py

@@ -547,8 +547,15 @@ async def metrics() -> Response:
 
 # RL 控制面接口（abort / pause / flush / memory / weight update），见 api_http_rl.py
 from .api_http_rl import router as rl_router
+from .rl_auth import RLStartArgs, mount_rl_router, validate_rl_auth
+
+_rl_router_mounted = False
+
+
+def _mount_rl_router(args: RLStartArgs) -> None:
+    global _rl_router_mounted
+    _rl_router_mounted = mount_rl_router(app, rl_router, args, _rl_router_mounted)
 
-app.include_router(rl_router)
 
 # PD 分离控制面接口（worker 发现、P/D 注册与 KV 状态上报），见 api_http_pd.py
 from .api_http_pd import router as pd_router
@@ -595,7 +602,10 @@ async def shutdown():
 async def startup_event():
     logger.info("server start up")
     loop = asyncio.get_event_loop()
-    g_objs.set_args(get_env_start_args())
+    args = get_env_start_args()
+    validate_rl_auth(args)
+    g_objs.set_args(args)
+    _mount_rl_router(args)
     loop.create_task(g_objs.httpserver_manager.handle_loop())
     logger.info(f"server start up ok, loop use is {asyncio.get_event_loop()}")
     return

lightllm/server/api_http_rl.py

@@ -12,7 +12,7 @@
 
 from http import HTTPStatus
 
-from fastapi import APIRouter, Request
+from fastapi import APIRouter, Depends, Request
 from fastapi.responses import JSONResponse, Response
 
 from lightllm.server.io_struct import (
@@ -30,10 +30,11 @@
 from lightllm.utils.log_utils import init_logger
 
 from .api_errors import create_error_response
+from .rl_auth import require_rl_auth
 
 logger = init_logger(__name__)
 
-router = APIRouter()
+router = APIRouter(dependencies=[Depends(require_rl_auth)])
 
 
 async def handle_request_common(request_obj, handler):

lightllm/server/api_start.py

@@ -33,6 +33,7 @@
     get_default_disable_chunked_mm_input,
 )
 from lightllm.utils.dist_check_utils import auto_configure_allreduce_flags_from_args
+from .rl_auth import validate_rl_auth
 
 logger = init_logger(__name__)
 
@@ -43,6 +44,7 @@ def _set_envs_and_config(args: StartArgs):
 
 def _launch_subprocesses(args: StartArgs):
     _validate_pd_features(args)
+    validate_rl_auth(args)
     _set_envs_and_config(args)
 
     if args.target_vocab_topk_sampling is not None:

lightllm/server/rl_auth.py

@@ -0,0 +1,52 @@
+"""Authentication helpers for the RL control-plane HTTP routes."""
+
+import os
+import secrets
+from typing import Protocol
+
+from fastapi import APIRouter, FastAPI, HTTPException, Request
+
+RL_API_KEY_ENV = "LIGHTLLM_RL_API_KEY"
+
+
+class RLStartArgs(Protocol):
+    enable_rl: bool
+
+
+def get_rl_api_key() -> str | None:
+    """Return the configured RL control-plane key without logging it."""
+    key = os.getenv(RL_API_KEY_ENV)
+    return key if key else None
+
+
+def validate_rl_auth(args: RLStartArgs) -> None:
+    """Reject an RL-enabled server that has no control-plane credential."""
+    if args.enable_rl and get_rl_api_key() is None:
+        raise ValueError(f"{RL_API_KEY_ENV} must be set when --enable_rl is enabled")
+
+
+def mount_rl_router(app: FastAPI, router: APIRouter, args: RLStartArgs, mounted: bool) -> bool:
+    """Mount the RL router once and only for an RL-enabled server."""
+    if args.enable_rl and not mounted:
+        app.include_router(router)
+        return True
+    return mounted
+
+
+def require_rl_auth(request: Request) -> None:
+    """Require a valid bearer token for an RL control-plane request."""
+    configured_key = get_rl_api_key()
+    if configured_key is None:
+        raise HTTPException(status_code=503, detail="RL control-plane authentication is not configured")
+
+    authorization_parts = request.headers.get("authorization", "").split()
+    if (
+        len(authorization_parts) != 2
+        or authorization_parts[0].lower() != "bearer"
+        or not secrets.compare_digest(authorization_parts[1].encode("utf-8"), configured_key.encode("utf-8"))
+    ):
+        raise HTTPException(
+            status_code=401,
+            detail="Invalid RL control-plane credentials",
+            headers={"WWW-Authenticate": "Bearer"},
+        )

unit_tests/server/test_rl_api_auth.py

@@ -0,0 +1,117 @@
+import sys
+from types import ModuleType
+from types import SimpleNamespace
+
+import pytest
+from fastapi import FastAPI, HTTPException
+from fastapi.testclient import TestClient
+from starlette.requests import Request
+
+from lightllm.server.api_http_rl import router as rl_router
+from lightllm.server.rl_auth import (
+    RL_API_KEY_ENV,
+    mount_rl_router,
+    require_rl_auth,
+    validate_rl_auth,
+)
+
+
+def _request(headers=None):
+    raw_headers = [(key.lower().encode(), value.encode()) for key, value in (headers or {}).items()]
+    return Request({"type": "http", "headers": raw_headers})
+
+
+def _route_paths(routes):
+    for route in routes:
+        path = getattr(route, "path", None)
+        if path is not None:
+            yield path
+        yield from _route_paths(getattr(route, "routes", []))
+        original_router = getattr(route, "original_router", None)
+        if original_router is not None:
+            yield from _route_paths(original_router.routes)
+
+
+def test_rl_routes_are_mounted_only_when_enabled(monkeypatch):
+    app = FastAPI()
+
+    mounted = mount_rl_router(app, rl_router, SimpleNamespace(enable_rl=False), False)
+    assert "/pause_generation" not in set(_route_paths(app.routes))
+
+    monkeypatch.setenv(RL_API_KEY_ENV, "control-secret")
+    mounted = mount_rl_router(app, rl_router, SimpleNamespace(enable_rl=True), mounted)
+    assert "/pause_generation" in set(_route_paths(app.routes))
+    assert mount_rl_router(app, rl_router, SimpleNamespace(enable_rl=True), mounted) is True
+
+
+def test_mounted_rl_routes_reject_unauthenticated_requests(monkeypatch):
+    monkeypatch.setenv(RL_API_KEY_ENV, "control-secret")
+    app = FastAPI()
+    app.include_router(rl_router)
+
+    response = TestClient(app).post("/pause_generation")
+
+    assert response.status_code == 401
+    assert response.headers["www-authenticate"] == "Bearer"
+
+
+def test_every_rl_route_requires_auth():
+    assert rl_router.routes
+    assert all(
+        any(dependency.call is require_rl_auth for dependency in route.dependant.dependencies)
+        for route in rl_router.routes
+    )
+
+
+def test_mounted_rl_route_dispatches_with_valid_bearer_token(monkeypatch):
+    monkeypatch.setenv(RL_API_KEY_ENV, "control-secret")
+    calls = []
+
+    async def pause_generation():
+        calls.append("pause")
+
+    api_http_module = ModuleType("lightllm.server.api_http")
+    api_http_module.g_objs = SimpleNamespace(httpserver_manager=SimpleNamespace(pause_generation=pause_generation))
+    monkeypatch.setitem(sys.modules, "lightllm.server.api_http", api_http_module)
+
+    app = FastAPI()
+    app.include_router(rl_router)
+    response = TestClient(app).post(
+        "/pause_generation",
+        headers={"Authorization": "Bearer control-secret"},
+    )
+
+    assert response.status_code == 200
+    assert calls == ["pause"]
+
+
+def test_rl_auth_accepts_only_the_configured_bearer_token(monkeypatch):
+    monkeypatch.setenv(RL_API_KEY_ENV, "control-secret")
+
+    assert require_rl_auth(_request({"authorization": "Bearer control-secret"})) is None
+    assert require_rl_auth(_request({"authorization": "bEaReR control-secret"})) is None
+
+    with pytest.raises(HTTPException) as error:
+        require_rl_auth(_request({"authorization": "Bearer wrong-secret"}))
+    assert error.value.status_code == 401
+
+    with pytest.raises(HTTPException) as error:
+        require_rl_auth(_request({"authorization": "Bearer incorrect-" + chr(233)}))
+    assert error.value.status_code == 401
+
+
+def test_rl_auth_rejects_missing_configuration(monkeypatch):
+    monkeypatch.delenv(RL_API_KEY_ENV, raising=False)
+
+    with pytest.raises(HTTPException) as error:
+        require_rl_auth(_request({"authorization": "Bearer control-secret"}))
+    assert error.value.status_code == 503
+
+
+def test_rl_startup_requires_a_control_plane_key(monkeypatch):
+    monkeypatch.delenv(RL_API_KEY_ENV, raising=False)
+    with pytest.raises(ValueError, match=RL_API_KEY_ENV):
+        validate_rl_auth(SimpleNamespace(enable_rl=True))
+
+    monkeypatch.setenv(RL_API_KEY_ENV, "control-secret")
+    validate_rl_auth(SimpleNamespace(enable_rl=True))