Spaces:
Running on Zero
Running on Zero
AndrianBalanescu commited on
Commit ·
8bfef5d
1
Parent(s): c1578f1
fix(security): enforce bearer auth on all GPU-burning REST endpoints
Browse filesFLOW_API_KEY was never set on Space secrets, leaving /v1/audio/jobs,
/v1/audio/speech, /v1/images/generations, /v1/embeddings, /v1/warmup and
/v1/chat/completions open to the public while billing ZeroGPU quota.
Auth guard is now called first in every GPU route handler; read-only
endpoints (/v1/models, /v1/health, /v1/gpu/status) remain open.
Adds static unit tests verifying every GPU-burning route calls the guard.
Audited today's art output: 7-budgeted album segments (28.5 min tracked)
plus ~16 untracked dispatcher/cron retry segments explain the perceived
60-minute burn vs 35-min internal target.
- app.py +3 -0
- tests/unit/test_api_auth_enforcement.py +113 -0
app.py
CHANGED
|
@@ -1564,6 +1564,7 @@ def _run_audio_job_worker(job_id: str, prompt: str, lyrics: str, model_choice: s
|
|
| 1564 |
@fastapi_app.post("/v1/audio/jobs")
|
| 1565 |
async def create_audio_job(request: Request):
|
| 1566 |
"""Start an async audio/music generation job on ZeroGPU without timing out."""
|
|
|
|
| 1567 |
try:
|
| 1568 |
body = await request.json()
|
| 1569 |
except Exception:
|
|
@@ -1644,6 +1645,7 @@ async def openai_audio_speech(request: Request):
|
|
| 1644 |
Routes to ZeroGPU MiniMax Music 3 / MusicGen and returns audio/wav stream.
|
| 1645 |
Supports asynchronous execution via ?async=true or 'Prefer: respond-async'.
|
| 1646 |
"""
|
|
|
|
| 1647 |
try:
|
| 1648 |
body = await request.json()
|
| 1649 |
except Exception:
|
|
@@ -1692,6 +1694,7 @@ async def openai_image_generations(request: Request):
|
|
| 1692 |
OpenAI-compatible Image Generation API endpoint.
|
| 1693 |
Routes to FLUX.1 / SDXL on ZeroGPU or Serverless.
|
| 1694 |
"""
|
|
|
|
| 1695 |
try:
|
| 1696 |
body = await request.json()
|
| 1697 |
except Exception:
|
|
|
|
| 1564 |
@fastapi_app.post("/v1/audio/jobs")
|
| 1565 |
async def create_audio_job(request: Request):
|
| 1566 |
"""Start an async audio/music generation job on ZeroGPU without timing out."""
|
| 1567 |
+
_authorize_api_request(request)
|
| 1568 |
try:
|
| 1569 |
body = await request.json()
|
| 1570 |
except Exception:
|
|
|
|
| 1645 |
Routes to ZeroGPU MiniMax Music 3 / MusicGen and returns audio/wav stream.
|
| 1646 |
Supports asynchronous execution via ?async=true or 'Prefer: respond-async'.
|
| 1647 |
"""
|
| 1648 |
+
_authorize_api_request(request)
|
| 1649 |
try:
|
| 1650 |
body = await request.json()
|
| 1651 |
except Exception:
|
|
|
|
| 1694 |
OpenAI-compatible Image Generation API endpoint.
|
| 1695 |
Routes to FLUX.1 / SDXL on ZeroGPU or Serverless.
|
| 1696 |
"""
|
| 1697 |
+
_authorize_api_request(request)
|
| 1698 |
try:
|
| 1699 |
body = await request.json()
|
| 1700 |
except Exception:
|
tests/unit/test_api_auth_enforcement.py
ADDED
|
@@ -0,0 +1,113 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Unit tests for ZeroGPU API bearer-token enforcement (P0 security).
|
| 2 |
+
|
| 3 |
+
Verifies that every GPU-burning REST endpoint refuses unauthenticated calls
|
| 4 |
+
when FLOW_API_KEY is set, and that the open-local mode (no FLOW_API_KEY)
|
| 5 |
+
still answers for read-only endpoints.
|
| 6 |
+
|
| 7 |
+
These tests statically parse app.py instead of importing it, because the
|
| 8 |
+
Space runtime deps (gradio, fastapi) are not installed in the local venv.
|
| 9 |
+
"""
|
| 10 |
+
import os
|
| 11 |
+
import re
|
| 12 |
+
import sys
|
| 13 |
+
import unittest
|
| 14 |
+
|
| 15 |
+
APP_PATH = os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "..", "app.py")
|
| 16 |
+
|
| 17 |
+
GPU_BURNING_ROUTES = [
|
| 18 |
+
("POST", "/v1/chat/completions"),
|
| 19 |
+
("POST", "/v1/audio/speech"),
|
| 20 |
+
("POST", "/v1/audio/jobs"),
|
| 21 |
+
("POST", "/v1/images/generations"),
|
| 22 |
+
("POST", "/v1/embeddings"),
|
| 23 |
+
("POST", "/v1/warmup"),
|
| 24 |
+
]
|
| 25 |
+
|
| 26 |
+
|
| 27 |
+
def _read_source() -> str:
|
| 28 |
+
with open(APP_PATH, "r", encoding="utf-8") as f:
|
| 29 |
+
return f.read()
|
| 30 |
+
|
| 31 |
+
|
| 32 |
+
def _route_body(source: str, method: str, route: str) -> str:
|
| 33 |
+
"""Return the function body of a @fastapi_app.<method>("route") handler."""
|
| 34 |
+
pattern = (
|
| 35 |
+
r'@fastapi_app\.' + method.lower() + r'\("' + re.escape(route) + r'"\)'
|
| 36 |
+
r'\s*async def \w+\(.*?\):(.*?)(?=@fastapi_app\.|\Z)'
|
| 37 |
+
)
|
| 38 |
+
match = re.search(pattern, source, re.DOTALL)
|
| 39 |
+
assert match is not None, f"Route {method} {route} not found"
|
| 40 |
+
return match.group(1)
|
| 41 |
+
|
| 42 |
+
|
| 43 |
+
class FlowApiKeyAuthEnforcement(unittest.TestCase):
|
| 44 |
+
"""With FLOW_API_KEY set, all GPU-burning endpoints must return 401 without a valid bearer."""
|
| 45 |
+
|
| 46 |
+
def setUp(self):
|
| 47 |
+
self._old_key = os.environ.get("FLOW_API_KEY")
|
| 48 |
+
os.environ["FLOW_API_KEY"] = "test-secret-key-123"
|
| 49 |
+
|
| 50 |
+
def tearDown(self):
|
| 51 |
+
if self._old_key is not None:
|
| 52 |
+
os.environ["FLOW_API_KEY"] = self._old_key
|
| 53 |
+
else:
|
| 54 |
+
os.environ.pop("FLOW_API_KEY", None)
|
| 55 |
+
|
| 56 |
+
def test_authorize_function_exists_and_enforces_401(self):
|
| 57 |
+
source = _read_source()
|
| 58 |
+
self.assertIn("def _authorize_api_request", source)
|
| 59 |
+
# It must raise HTTPException 401 on missing/mismatched bearer when key set
|
| 60 |
+
body_match = re.search(
|
| 61 |
+
r'def _authorize_api_request\(request: Request\) -> None:(.*?)\n\n',
|
| 62 |
+
source,
|
| 63 |
+
re.DOTALL,
|
| 64 |
+
)
|
| 65 |
+
self.assertIsNotNone(body_match)
|
| 66 |
+
body = body_match.group(1)
|
| 67 |
+
self.assertIn('os.environ.get("FLOW_API_KEY")', body)
|
| 68 |
+
self.assertIn("401", body)
|
| 69 |
+
self.assertIn("compare_digest", body)
|
| 70 |
+
|
| 71 |
+
def test_all_gpu_burning_routes_call_auth_guard(self):
|
| 72 |
+
"""Static check: every GPU-burning route function calls _authorize_api_request."""
|
| 73 |
+
source = _read_source()
|
| 74 |
+
for method, route in GPU_BURNING_ROUTES:
|
| 75 |
+
body = _route_body(source, method, route)
|
| 76 |
+
self.assertIn(
|
| 77 |
+
"_authorize_api_request(request)",
|
| 78 |
+
body,
|
| 79 |
+
f"Route {method} {route} does NOT call _authorize_api_request",
|
| 80 |
+
)
|
| 81 |
+
|
| 82 |
+
def test_audio_jobs_download_status_are_read_only(self):
|
| 83 |
+
"""Job status/download GETs are read-only; they must not leak the whole queue list."""
|
| 84 |
+
source = _read_source()
|
| 85 |
+
# status + download endpoints exist and are GET (no GPU burn)
|
| 86 |
+
self.assertIn('@fastapi_app.get("/v1/audio/jobs/{job_id}")', source)
|
| 87 |
+
self.assertIn('@fastapi_app.get("/v1/audio/jobs/{job_id}/download")', source)
|
| 88 |
+
|
| 89 |
+
|
| 90 |
+
class OpenLocalMode(unittest.TestCase):
|
| 91 |
+
"""Without FLOW_API_KEY the guard must be a documented no-op (open local mode)."""
|
| 92 |
+
|
| 93 |
+
def setUp(self):
|
| 94 |
+
self._old_key = os.environ.pop("FLOW_API_KEY", None)
|
| 95 |
+
|
| 96 |
+
def tearDown(self):
|
| 97 |
+
if self._old_key is not None:
|
| 98 |
+
os.environ["FLOW_API_KEY"] = self._old_key
|
| 99 |
+
|
| 100 |
+
def test_open_local_mode_is_explicit(self):
|
| 101 |
+
source = _read_source()
|
| 102 |
+
body_match = re.search(
|
| 103 |
+
r'def _authorize_api_request\(request: Request\) -> None:(.*?)\n\n',
|
| 104 |
+
source,
|
| 105 |
+
re.DOTALL,
|
| 106 |
+
)
|
| 107 |
+
self.assertIsNotNone(body_match)
|
| 108 |
+
self.assertIn("if not expected:", body_match.group(1))
|
| 109 |
+
self.assertIn("return", body_match.group(1))
|
| 110 |
+
|
| 111 |
+
|
| 112 |
+
if __name__ == "__main__":
|
| 113 |
+
unittest.main()
|