diff --git a/pyproject.toml b/pyproject.toml index 3ac25cb266..278247ae3e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -128,6 +128,12 @@ otel = [ "opentelemetry-instrumentation-fastapi>=0.41b0", "opentelemetry-instrumentation-httpx>=0.41b0", ] +guard = [ + # Optional fastapi-guard security middleware + # (api/guard_integration.py; default-off) + "fastapi-guard>=8.1.0,<9.0.0", +] + router = [ "tokenizers>=0.21.0,<0.23.0", ] diff --git a/xinference/api/guard_integration.py b/xinference/api/guard_integration.py new file mode 100644 index 0000000000..ef02352329 --- /dev/null +++ b/xinference/api/guard_integration.py @@ -0,0 +1,200 @@ +# Copyright 2022-2026 Xinference Holdings Pte. Ltd +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Optional fastapi-guard security middleware wiring for the RESTful API. + +fastapi-guard (https://github.com/Guard-Core/fastapi-guard) provides IP +block/allow lists, rate limiting with auto-ban, user-agent blocking, +penetration-attempt detection, passive mode, optional Redis-backed shared +state, and optional IPInfo geo/cloud-provider lookups. + +This complements the existing protections without replacing them: the +``XINFERENCE_ALLOWED_IPS`` allowlist stays first, and the auth rate limiter +(bans on failed logins) remains auth-scoped. The guard adds general +per-request rate limiting, blocklists, user-agent blocking, and URL/query/ +header screening. + +Client identity is zero-trust: ``X-Forwarded-For`` is honored only from +proxies explicitly listed in ``XINFERENCE_GUARD_TRUSTED_PROXIES``, by both +Uvicorn (which would otherwise rewrite the client address for any peer +before middleware runs) and guard. Prompt bodies are not signature-scanned +by default (model input is data, not commands this API executes); opt into +that with ``XINFERENCE_GUARD_SCAN_BODY``. The default security headers keep +same-origin microphone/camera capture working for the Web UI. + +Everything here is opt-in: with ``XINFERENCE_GUARD_ENABLED`` unset (the +default), ``attach_guard`` is a no-op and the server behaves exactly as +before. Install the extra with ``pip install "xinference[guard]"``. +""" + +from __future__ import annotations + +import os +from typing import TYPE_CHECKING, Any + +if TYPE_CHECKING: + from fastapi import FastAPI + from guard import SecurityConfig + +DEFAULT_EXCLUDED_PATHS = "/health,/docs,/redoc,/openapi.json" + + +def _csv(raw: str | None) -> tuple[str, ...]: + if not raw: + return () + return tuple(item.strip() for item in raw.split(",") if item.strip()) + + +def _env_int(name: str, default: int) -> int: + raw = os.environ.get(name) + if raw is None or raw == "": + return default + return int(raw) + + +def _env_bool(name: str) -> bool: + return os.environ.get(name, "").strip().lower() in ("1", "true", "yes") + + +def guard_enabled() -> bool: + """Whether XINFERENCE_GUARD_ENABLED is set (read at call time so forked + child processes and late env wiring are seen).""" + return _env_bool("XINFERENCE_GUARD_ENABLED") + + +def uvicorn_forwarded_allow_ips() -> str: + """The Uvicorn ``forwarded_allow_ips`` value matching guard's proxy trust. + + ``RESTfulAPI.serve()`` launches Uvicorn with proxy headers enabled. With + the previous default ``forwarded_allow_ips="*"``, Uvicorn rewrote + ``scope["client"]`` from ANY peer's ``X-Forwarded-For`` before any + middleware ran, so a direct client could forge its address and defeat + guard's blocklists and per-IP limiter (and the allowlist/audit layers + reading the same field). When the guard is enabled, only proxies listed + in ``XINFERENCE_GUARD_TRUSTED_PROXIES`` may rewrite the client address; + with none listed, Uvicorn trusts loopback only. Guard's own + ``trusted_proxies`` reads the same variable, so Uvicorn, guard, and the + existing allowlist/audit consumers resolve one client identity. + """ + if not guard_enabled(): + return "*" + return ( + ",".join(_csv(os.environ.get("XINFERENCE_GUARD_TRUSTED_PROXIES"))) + or "127.0.0.1" + ) + + +def _build_guard_config() -> SecurityConfig: + """Build the SecurityConfig from the XINFERENCE_GUARD_* environment variables. + + Environment variables are read at call time (the same reason + ``is_auth_advanced`` reads at call time: forked child processes). + """ + from guard import SecurityConfig + + kwargs: dict[str, Any] = { + "enable_rate_limiting": True, + "rate_limit": _env_int("XINFERENCE_GUARD_RATE_LIMIT", 100), + "rate_limit_window": _env_int("XINFERENCE_GUARD_RATE_LIMIT_WINDOW", 60), + "enable_ip_banning": True, + "auto_ban_threshold": _env_int("XINFERENCE_GUARD_AUTO_BAN_THRESHOLD", 10), + "auto_ban_duration": _env_int("XINFERENCE_GUARD_AUTO_BAN_DURATION", 300), + "enable_penetration_detection": True, + # Model input is data, not commands this API executes: signature- + # scanning prompt bodies false-positives on ordinary code-generation + # and security-analysis conversations (SQL/HTML/shell text inside + # chat messages). Detection keeps screening the URL path, query + # params, and headers; reading and scanning bodies is opt-in via + # XINFERENCE_GUARD_SCAN_BODY. + "detection_scan_body": _env_bool("XINFERENCE_GUARD_SCAN_BODY"), + # The dependency's default Permissions-Policy locks the microphone + # to none, which silently breaks the Web UI's in-page audio recorder + # (frontend audio-recorder-upload.tsx calls getUserMedia({audio: + # true}) on the same origin). Keep the hardening but let same-origin + # media capture work; operators can override the whole policy. + "security_headers": { + "enabled": True, + "hsts": {"max_age": 31536000, "include_subdomains": True, "preload": False}, + "csp": None, + "frame_options": "SAMEORIGIN", + "content_type_options": "nosniff", + "xss_protection": "1; mode=block", + "referrer_policy": "strict-origin-when-cross-origin", + "permissions_policy": os.environ.get("XINFERENCE_GUARD_PERMISSIONS_POLICY") + or "geolocation=(), microphone=(self), camera=(self)", + "custom": None, + }, + # In-memory state unless a Redis URL is configured: never implicitly + # depend on a Redis server being reachable. The RESTful API runs as a + # single process per supervisor, matching the existing auth rate + # limiter's per-process design; set XINFERENCE_GUARD_REDIS_URL for + # multi-instance deployments behind xinference-router. + "enable_redis": False, + "passive_mode": _env_bool("XINFERENCE_GUARD_PASSIVE_MODE"), + "blacklist": _csv(os.environ.get("XINFERENCE_GUARD_BLOCKED_IPS")), + "blocked_user_agents": list( + _csv(os.environ.get("XINFERENCE_GUARD_BLOCKED_USER_AGENTS")) + ), + # Zero-trust default: no proxies are trusted, so forwarded headers + # are ignored and the connecting peer is the client (matching the + # Uvicorn trust configured by uvicorn_forwarded_allow_ips). List + # real ingress addresses to enable X-Forwarded-For resolution. + "trusted_proxies": _csv(os.environ.get("XINFERENCE_GUARD_TRUSTED_PROXIES")), + "trusted_proxy_depth": _env_int("XINFERENCE_GUARD_TRUSTED_PROXY_DEPTH", 1), + "exclude_paths": list( + _csv(os.environ.get("XINFERENCE_GUARD_EXCLUDED_PATHS")) + or DEFAULT_EXCLUDED_PATHS.split(",") + ), + } + + if allowed_ips := _csv(os.environ.get("XINFERENCE_GUARD_ALLOWED_IPS")): + kwargs["whitelist"] = allowed_ips + if blocked_countries := _csv(os.environ.get("XINFERENCE_GUARD_BLOCKED_COUNTRIES")): + kwargs["blocked_countries"] = frozenset(blocked_countries) + if allowed_countries := _csv(os.environ.get("XINFERENCE_GUARD_ALLOWED_COUNTRIES")): + kwargs["whitelist_countries"] = frozenset(allowed_countries) + if cloud_providers := _csv( + os.environ.get("XINFERENCE_GUARD_BLOCK_CLOUD_PROVIDERS") + ): + kwargs["block_cloud_providers"] = frozenset(cloud_providers) + + if redis_url := os.environ.get("XINFERENCE_GUARD_REDIS_URL"): + kwargs["enable_redis"] = True + kwargs["redis_url"] = redis_url + kwargs["redis_prefix"] = "xinference_guard:" + + if ipinfo_token := os.environ.get("XINFERENCE_GUARD_IPINFO_TOKEN"): + kwargs["ipinfo_token"] = ipinfo_token + + return SecurityConfig(**kwargs) + + +def attach_guard(app: FastAPI) -> None: + """Attach the fastapi-guard middleware when XINFERENCE_GUARD_ENABLED is set. + + No-op unless the env flag is on. Fails loudly when the flag is on but + the package is missing: a misconfiguration must never silently disable + security. + """ + if not guard_enabled(): + return + try: + from guard import SecurityMiddleware + except ImportError as exc: + raise ImportError( + "XINFERENCE_GUARD_ENABLED requires fastapi-guard. " + 'Install it with: pip install "xinference[guard]"' + ) from exc + + app.add_middleware(SecurityMiddleware, config=_build_guard_config()) diff --git a/xinference/api/restful_api.py b/xinference/api/restful_api.py index 31c68cbc63..cd5b4fe72d 100644 --- a/xinference/api/restful_api.py +++ b/xinference/api/restful_api.py @@ -1161,6 +1161,18 @@ async def body_stream() -> AsyncIterator[bytes]: ) def serve(self, logging_conf: Optional[dict] = None): + # Optional fastapi-guard security middleware (no-op when disabled). + # Attached FIRST so it nests INSIDE both CORS and the audited IP + # middleware (Starlette runs the last-added middleware outermost): + # the effective chain is audit/IP -> CORS -> guard -> app, so guard's + # own rejections (403/429) leave with CORS headers (cross-origin + # browser clients can read them) and still pass through the audit + # wrapper keeping X-Request-ID, while the XINFERENCE_ALLOWED_IPS + # allowlist stays first in line on the way in. + from .guard_integration import attach_guard, uvicorn_forwarded_allow_ips + + attach_guard(self._app) + self._app.add_middleware( CORSMiddleware, allow_origins=["*"], @@ -1331,7 +1343,12 @@ async def internal_exception_handler(request: Request, exc: Exception): port=self._port, log_config=logging_conf, proxy_headers=True, - forwarded_allow_ips="*", + # With the guard enabled, only explicitly trusted proxies may + # rewrite the client address from X-Forwarded-For; trusting any + # peer here would let a direct client forge its identity before + # guard (or the allowlist above) ever sees the request. + # No-op ("*") when the guard is disabled. + forwarded_allow_ips=uvicorn_forwarded_allow_ips(), # Slow HTTP DoS (Slowloris) protection, see xinference/core/http_protocol.py http=create_hardened_http_protocol(XINFERENCE_HTTP_REQUEST_TIMEOUT), timeout_keep_alive=XINFERENCE_HTTP_TIMEOUT_KEEP_ALIVE, diff --git a/xinference/api/tests/test_guard_integration.py b/xinference/api/tests/test_guard_integration.py new file mode 100644 index 0000000000..49c4787774 --- /dev/null +++ b/xinference/api/tests/test_guard_integration.py @@ -0,0 +1,497 @@ +# Copyright 2022-2026 Xinference Holdings Pte. Ltd +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Unit tests for the optional fastapi-guard middleware wiring.""" + +import asyncio +import itertools +import os +import sys +import threading +import time +import unittest + +_IP = itertools.count(1) + +try: + import guard # noqa: F401 + + _GUARD_AVAILABLE = True +except ImportError: + _GUARD_AVAILABLE = False + + +def _unique_ip() -> str: + n = next(_IP) + return f"198.51.{n // 250}.{(n % 250) + 1}" + + +def _restore_env(saved): + for name in list(os.environ): + if name.startswith("XINFERENCE_GUARD_"): + os.environ.pop(name, None) + os.environ.update(saved) + + +class TestGuardIntegration(unittest.TestCase): + def setUp(self): + self._saved = { + k: os.environ.pop(k) + for k in list(os.environ) + if k.startswith("XINFERENCE_GUARD_") + } + + def tearDown(self): + _restore_env(self._saved) + + def _build_app(self): + from fastapi import FastAPI + + from xinference.api.guard_integration import attach_guard + + app = FastAPI() + + @app.get("/ping") + async def ping(): + return {"ok": True} + + attach_guard(app) + return app + + async def _get(self, client_ip): + from httpx import ASGITransport, AsyncClient + + transport = ASGITransport(app=self._build_app(), client=(client_ip, 50000)) + async with AsyncClient( + transport=transport, base_url="http://testserver" + ) as client: + return (await client.get("/ping")).status_code + + def test_disabled_by_default(self): + from fastapi import FastAPI + + from xinference.api.guard_integration import attach_guard + + app = FastAPI() + attach_guard(app) + self.assertEqual(len(app.user_middleware), 0) + + def test_fail_loud_without_package(self): + from fastapi import FastAPI + + from xinference.api.guard_integration import attach_guard + + if _GUARD_AVAILABLE: + self.skipTest("fastapi-guard is installed") + os.environ["XINFERENCE_GUARD_ENABLED"] = "true" + with self.assertRaisesRegex(ImportError, r"xinference\[guard\]"): + attach_guard(FastAPI()) + + @unittest.skipUnless(_GUARD_AVAILABLE, "fastapi-guard not installed") + def test_blocked_ip_is_rejected(self): + blocked = _unique_ip() + os.environ["XINFERENCE_GUARD_ENABLED"] = "true" + os.environ["XINFERENCE_GUARD_BLOCKED_IPS"] = blocked + self.assertEqual(asyncio.run(self._get(blocked)), 403) + + @unittest.skipUnless(_GUARD_AVAILABLE, "fastapi-guard not installed") + def test_rate_limit_returns_429(self): + client_ip = _unique_ip() + os.environ["XINFERENCE_GUARD_ENABLED"] = "true" + os.environ["XINFERENCE_GUARD_RATE_LIMIT"] = "2" + os.environ["XINFERENCE_GUARD_RATE_LIMIT_WINDOW"] = "60" + codes = [asyncio.run(self._get(client_ip)) for _ in range(3)] + self.assertEqual(codes, [200, 200, 429]) + + @unittest.skipUnless(_GUARD_AVAILABLE, "fastapi-guard not installed") + def test_passive_mode_never_blocks(self): + blocked = _unique_ip() + os.environ["XINFERENCE_GUARD_ENABLED"] = "true" + os.environ["XINFERENCE_GUARD_PASSIVE_MODE"] = "true" + os.environ["XINFERENCE_GUARD_BLOCKED_IPS"] = blocked + self.assertEqual(asyncio.run(self._get(blocked)), 200) + + @unittest.skipUnless(_GUARD_AVAILABLE, "fastapi-guard not installed") + def test_redis_stays_off_without_config(self): + from xinference.api.guard_integration import _build_guard_config + + os.environ["XINFERENCE_GUARD_ENABLED"] = "true" + config = _build_guard_config() + self.assertFalse(config.enable_redis) + + def test_uvicorn_forwarded_allow_ips_follows_guard_trust(self): + from xinference.api.guard_integration import uvicorn_forwarded_allow_ips + + # Guard off: serve() keeps the historical behavior untouched. + self.assertEqual(uvicorn_forwarded_allow_ips(), "*") + os.environ["XINFERENCE_GUARD_ENABLED"] = "true" + try: + # Guard on, no proxies listed: loopback only, so a remote peer + # cannot rewrite its own address. + self.assertEqual(uvicorn_forwarded_allow_ips(), "127.0.0.1") + os.environ["XINFERENCE_GUARD_TRUSTED_PROXIES"] = "10.0.0.1, 10.0.0.2" + self.assertEqual(uvicorn_forwarded_allow_ips(), "10.0.0.1,10.0.0.2") + finally: + os.environ.pop("XINFERENCE_GUARD_ENABLED", None) + os.environ.pop("XINFERENCE_GUARD_TRUSTED_PROXIES", None) + + def test_default_config_disables_body_scan_and_proxy_trust(self): + import types + from unittest import mock + + from xinference.api import guard_integration + + captured = {} + + class _StubConfig: + def __init__(self, **kwargs): + captured.update(kwargs) + + stub = types.ModuleType("guard") + stub.SecurityConfig = _StubConfig + os.environ["XINFERENCE_GUARD_ENABLED"] = "true" + with mock.patch.dict(sys.modules, {"guard": stub}): + guard_integration._build_guard_config() + # Model input is data, not payloads: body signature-scanning must be + # opt-in, and no proxy is trusted unless explicitly listed. + self.assertIs(captured["detection_scan_body"], False) + self.assertEqual(captured["trusted_proxies"], ()) + self.assertIn( + "microphone=(self)", + captured["security_headers"]["permissions_policy"], + ) + self.assertIs(captured["enable_redis"], False) + + @staticmethod + def _serve_on_ephemeral_port(app, forwarded_allow_ips): + """Run the app on a real Uvicorn server thread with serve()'s + proxy-header configuration, and return (server, base_url).""" + from uvicorn import Config, Server + + config = Config( + app=app, + host="127.0.0.1", + port=0, + log_level="warning", + proxy_headers=True, + forwarded_allow_ips=forwarded_allow_ips, + lifespan="off", + ) + server = Server(config) + threading.Thread(target=server.run, daemon=True).start() + deadline = time.time() + 10 + while not server.started and time.time() < deadline: + time.sleep(0.02) + host, port = server.servers[0].sockets[0].getsockname()[:2] + return server, f"http://{host}:{int(port)}" + + @staticmethod + def _http_get(url, headers=None): + from httpx import Client + + with Client() as client: + return client.get(url, headers=headers or {}) + + @unittest.skipUnless(_GUARD_AVAILABLE, "fastapi-guard not installed") + def test_untrusted_forwarded_header_cannot_spoof(self): + """Through the real Uvicorn wrapper: a peer on the blocklist stays + blocked even when it forges X-Forwarded-For (the connecting peer is + not a trusted proxy here, so Uvicorn must not rewrite its + address).""" + os.environ["XINFERENCE_GUARD_ENABLED"] = "true" + # The loopback test client must NOT be a trusted proxy for this + # probe: point the trust list elsewhere. + os.environ["XINFERENCE_GUARD_TRUSTED_PROXIES"] = "10.99.99.99" + os.environ["XINFERENCE_GUARD_BLOCKED_IPS"] = "127.0.0.1" + + from xinference.api.guard_integration import uvicorn_forwarded_allow_ips + + server, base_url = self._serve_on_ephemeral_port( + self._build_app(), uvicorn_forwarded_allow_ips() + ) + try: + direct = self._http_get(f"{base_url}/ping") + forged = self._http_get( + f"{base_url}/ping", headers={"X-Forwarded-For": "198.51.100.9"} + ) + finally: + server.should_exit = True + self.assertEqual(direct.status_code, 403) + self.assertEqual(forged.status_code, 403) + + @unittest.skipUnless(_GUARD_AVAILABLE, "fastapi-guard not installed") + def test_configured_proxy_forwarded_header_is_honored(self): + """Loopback listed as a trusted proxy: Uvicorn resolves the client + from X-Forwarded-For, so the guard sees the forwarded address (a + blocked forwarded address is rejected; an unlisted one passes).""" + os.environ["XINFERENCE_GUARD_ENABLED"] = "true" + os.environ["XINFERENCE_GUARD_TRUSTED_PROXIES"] = "127.0.0.1" + blocked = _unique_ip() + os.environ["XINFERENCE_GUARD_BLOCKED_IPS"] = blocked + + from xinference.api.guard_integration import uvicorn_forwarded_allow_ips + + server, base_url = self._serve_on_ephemeral_port( + self._build_app(), uvicorn_forwarded_allow_ips() + ) + try: + via_blocked = self._http_get( + f"{base_url}/ping", headers={"X-Forwarded-For": blocked} + ) + via_allowed = self._http_get( + f"{base_url}/ping", headers={"X-Forwarded-For": _unique_ip()} + ) + finally: + server.should_exit = True + # If the header were ignored, the client would be loopback and both + # would return 200: the 403 proves the forwarded address was used. + self.assertEqual(via_blocked.status_code, 403) + self.assertEqual(via_allowed.status_code, 200) + + @unittest.skipUnless(_GUARD_AVAILABLE, "fastapi-guard not installed") + def test_rotating_untrusted_headers_cannot_evade_rate_limit(self): + """One peer rotating forged forwarded addresses is still one client + to the per-IP limiter (the loopback peer is not a trusted proxy + here, so every forged header is ignored).""" + os.environ["XINFERENCE_GUARD_ENABLED"] = "true" + os.environ["XINFERENCE_GUARD_TRUSTED_PROXIES"] = "10.99.99.99" + os.environ["XINFERENCE_GUARD_RATE_LIMIT"] = "2" + os.environ["XINFERENCE_GUARD_RATE_LIMIT_WINDOW"] = "60" + + from xinference.api.guard_integration import uvicorn_forwarded_allow_ips + + server, base_url = self._serve_on_ephemeral_port( + self._build_app(), uvicorn_forwarded_allow_ips() + ) + try: + codes = [ + self._http_get( + f"{base_url}/ping", + headers={"X-Forwarded-For": f"198.51.100.{n}"}, + ).status_code + for n in (10, 11, 12, 13) + ] + finally: + server.should_exit = True + self.assertEqual(codes, [200, 200, 429, 429]) + + @unittest.skipUnless(_GUARD_AVAILABLE, "fastapi-guard not installed") + def test_guard_rejection_stays_inside_audit_wrapping(self): + """serve() attaches guard before CORS and the audited IP middleware, + so the effective chain is audit/IP -> CORS -> guard -> app: guard + rejections flow back out through CORS (readable cross-origin errors) + and the audit layer (audited, X-Request-ID set), while the existing + IP allowlist still rejects before the guard ever sees the request.""" + os.environ["XINFERENCE_GUARD_ENABLED"] = "true" + blocked = _unique_ip() + os.environ["XINFERENCE_GUARD_BLOCKED_IPS"] = blocked + + from fastapi import FastAPI + from fastapi.responses import PlainTextResponse + from httpx import ASGITransport, AsyncClient + from starlette.middleware.base import BaseHTTPMiddleware + from starlette.middleware.cors import CORSMiddleware + + from xinference.api.guard_integration import attach_guard + + app = FastAPI() + + @app.get("/ping") + async def ping(): + return {"ok": True} + + attach_guard(app) # added first -> innermost + + app.add_middleware( + CORSMiddleware, + allow_origins=["*"], + allow_credentials=True, + allow_methods=["*"], + allow_headers=["*"], + ) + + audited = [] + + class _AuditedIP(BaseHTTPMiddleware): + async def dispatch(self, request, call_next): + audited.append(request.url.path) + + async def check_ip(inner_request): + # stand-in for ip_restriction_middleware: the + # XINFERENCE_ALLOWED_IPS allowlist, ahead of the guard + if inner_request.client.host == "203.0.113.9": + return PlainTextResponse(status_code=403, content="denied") + return await call_next(inner_request) + + # serve()'s audit wrapper calls check_ip as its call_next + response = await check_ip(request) + response.headers["X-Request-ID"] = "req-123" + return response + + app.add_middleware(_AuditedIP) # added last -> outermost, like serve() + + async def _get(ip): + transport = ASGITransport(app=app, client=(ip, 50000)) + async with AsyncClient( + transport=transport, base_url="http://testserver" + ) as client: + return await client.get( + "/ping", headers={"Origin": "https://client.example"} + ) + + rejected = asyncio.run(_get(blocked)) + self.assertEqual(rejected.status_code, 403) + self.assertIn( + rejected.headers.get("access-control-allow-origin"), + ("*", "https://client.example"), + ) + self.assertEqual(rejected.headers["x-request-id"], "req-123") + self.assertEqual(audited, ["/ping"]) + + # allowlist-first: the existing IP check denies before the guard + allowlisted_out = asyncio.run(_get("203.0.113.9")) + self.assertEqual(allowlisted_out.status_code, 403) + self.assertEqual(allowlisted_out.headers["x-request-id"], "req-123") + self.assertEqual(audited, ["/ping", "/ping"]) + + @unittest.skipUnless(_GUARD_AVAILABLE, "fastapi-guard not installed") + def test_rate_limited_response_carries_cors_headers(self): + """A guard-generated 429 passes back out through CORS, so + cross-origin clients read the rate-limit response instead of a + network failure.""" + os.environ["XINFERENCE_GUARD_ENABLED"] = "true" + client_ip = _unique_ip() + os.environ["XINFERENCE_GUARD_RATE_LIMIT"] = "1" + os.environ["XINFERENCE_GUARD_RATE_LIMIT_WINDOW"] = "60" + + from fastapi import FastAPI + from httpx import ASGITransport, AsyncClient + from starlette.middleware.base import BaseHTTPMiddleware + from starlette.middleware.cors import CORSMiddleware + + from xinference.api.guard_integration import attach_guard + + app = FastAPI() + + @app.get("/ping") + async def ping(): + return {"ok": True} + + attach_guard(app) + app.add_middleware( + CORSMiddleware, + allow_origins=["*"], + allow_credentials=True, + allow_methods=["*"], + allow_headers=["*"], + ) + + class _Audit(BaseHTTPMiddleware): + async def dispatch(self, request, call_next): + response = await call_next(request) + response.headers["X-Request-ID"] = "req-429" + return response + + app.add_middleware(_Audit) + + async def _get(): + transport = ASGITransport(app=app, client=(client_ip, 50000)) + async with AsyncClient( + transport=transport, base_url="http://testserver" + ) as client: + return await client.get( + "/ping", headers={"Origin": "https://client.example"} + ) + + asyncio.run(_get()) + limited = asyncio.run(_get()) + self.assertEqual(limited.status_code, 429) + self.assertIn( + limited.headers.get("access-control-allow-origin"), + ("*", "https://client.example"), + ) + self.assertEqual(limited.headers["x-request-id"], "req-429") + + @unittest.skipUnless(_GUARD_AVAILABLE, "fastapi-guard not installed") + def test_inference_bodies_not_scanned_by_default(self): + """Ordinary code/security conversation input is model data, not + payloads executed by this API: it passes with the guard on. Body + signature-scanning is opt-in via XINFERENCE_GUARD_SCAN_BODY.""" + from fastapi import FastAPI + from fastapi.responses import JSONResponse + from httpx import ASGITransport, AsyncClient + + from xinference.api.guard_integration import attach_guard + + os.environ["XINFERENCE_GUARD_ENABLED"] = "true" + prompts = [ + "Explain this SQL query: SELECT * FROM users WHERE name = 'Alice';", + "Show a simple HTML page with .", + "Review this shell command: cat /etc/passwd", + ] + + def _make_app(): + app = FastAPI() + + @app.post("/v1/chat/completions") + async def chat(): + return JSONResponse({"ok": True}) + + attach_guard(app) + return app + + async def _post(content): + transport = ASGITransport(app=_make_app(), client=(_unique_ip(), 50000)) + async with AsyncClient( + transport=transport, base_url="http://testserver" + ) as client: + return ( + await client.post( + "/v1/chat/completions", + json={"messages": [{"role": "user", "content": content}]}, + ) + ).status_code + + for prompt in prompts: + self.assertEqual(asyncio.run(_post(prompt)), 200) + + os.environ["XINFERENCE_GUARD_SCAN_BODY"] = "1" + self.assertEqual(asyncio.run(_post(prompts[0])), 400) + + @unittest.skipUnless(_GUARD_AVAILABLE, "fastapi-guard not installed") + def test_permissions_policy_keeps_webui_microphone(self): + """The default Permissions-Policy must not break the Web UI's + in-page audio recorder (same-origin getUserMedia audio).""" + os.environ["XINFERENCE_GUARD_ENABLED"] = "true" + + async def _headers(): + from httpx import ASGITransport, AsyncClient + + transport = ASGITransport( + app=self._build_app(), client=(_unique_ip(), 50000) + ) + async with AsyncClient( + transport=transport, base_url="http://testserver" + ) as client: + return (await client.get("/ping")).headers + + headers = asyncio.run(_headers()) + self.assertEqual( + headers["permissions-policy"], + "geolocation=(), microphone=(self), camera=(self)", + ) + + +if __name__ == "__main__": + unittest.main()