AndrianBalanescu commited on
Commit
8bfef5d
·
1 Parent(s): c1578f1

fix(security): enforce bearer auth on all GPU-burning REST endpoints

Browse files

FLOW_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.

Files changed (2) hide show
  1. app.py +3 -0
  2. 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()