@wendellresearch-ai/benchwarmers 0.1.1 → 0.2.1

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (96) hide show
  1. package/AGENTS.md +147 -147
  2. package/BENCHMARKS.md +262 -262
  3. package/GATEWAY.md +272 -271
  4. package/NOTICE +28 -28
  5. package/PRIVACY.md +46 -46
  6. package/README.md +368 -275
  7. package/SECURITY.md +83 -83
  8. package/TERMS.md +22 -22
  9. package/USAGE.md +170 -155
  10. package/VERSION +1 -1
  11. package/bin/benchwarmers.js +16 -72
  12. package/cli.py +666 -0
  13. package/examples/config/gateway.json +5 -5
  14. package/examples/config/model-roles.json +11 -11
  15. package/gateway.py +755 -712
  16. package/launcher.py +120 -0
  17. package/mcp_server.py +266 -0
  18. package/metrics/alerts.py +116 -116
  19. package/metrics/aliases.py +41 -41
  20. package/metrics/bench_collections.py +141 -141
  21. package/metrics/catalog.py +295 -295
  22. package/metrics/client.py +394 -373
  23. package/metrics/cliproxy.py +191 -191
  24. package/metrics/connect.py +691 -691
  25. package/metrics/cost.py +152 -152
  26. package/metrics/decide.py +227 -227
  27. package/metrics/discover.py +318 -318
  28. package/metrics/endpoints.py +436 -436
  29. package/metrics/guard.py +617 -592
  30. package/metrics/learn.py +221 -221
  31. package/metrics/litellm_hub.py +206 -206
  32. package/metrics/model_roles.py +84 -84
  33. package/metrics/offload.py +166 -166
  34. package/metrics/owui.py +205 -205
  35. package/metrics/paths.py +60 -60
  36. package/metrics/route.py +579 -574
  37. package/metrics/scheduler.py +275 -275
  38. package/metrics/sources.py +167 -167
  39. package/metrics/wrong.py +144 -144
  40. package/package.json +68 -64
  41. package/perf/README.md +89 -89
  42. package/perf/run.py +261 -261
  43. package/scripts/auto_retire.py +367 -367
  44. package/scripts/check.py +108 -104
  45. package/scripts/check_deprecations.py +378 -378
  46. package/scripts/check_themes.py +103 -103
  47. package/scripts/init.py +124 -124
  48. package/scripts/list_served.py +78 -78
  49. package/scripts/notice.py +78 -78
  50. package/scripts/onboard.py +401 -331
  51. package/scripts/provider_tags.py +129 -129
  52. package/scripts/reprefix.py +136 -136
  53. package/scripts/sync_config.py +105 -105
  54. package/scripts/test_auth.py +255 -255
  55. package/scripts/test_cli.py +189 -0
  56. package/scripts/test_client.py +233 -233
  57. package/scripts/test_connect.py +222 -222
  58. package/scripts/test_first_run.py +180 -180
  59. package/scripts/test_hubs.py +328 -328
  60. package/scripts/test_launcher.py +67 -0
  61. package/scripts/test_llm_log.py +85 -0
  62. package/scripts/test_mcp.py +178 -0
  63. package/scripts/test_merged.py +214 -193
  64. package/scripts/test_no_owui.py +291 -291
  65. package/scripts/test_owui.py +90 -90
  66. package/scripts/test_serve_edit.py +448 -386
  67. package/scripts/test_verdict.py +98 -98
  68. package/scripts/test_vision.py +228 -228
  69. package/serve.py +3904 -3662
  70. package/suite/mine.py +154 -154
  71. package/suite/pace.py +145 -145
  72. package/suite/run.py +1046 -1040
  73. package/suite/seed.py +421 -421
  74. package/suite/tasks/claim-alias-unbound/fixtures/onboard.py +286 -286
  75. package/suite/tasks/claim-error-text-string/fixtures/decide.py +187 -187
  76. package/suite/tasks/claim-list-body/fixtures/gateway.py +235 -235
  77. package/suite/tasks/claim-member-nominated-twice/fixtures/gateway.py +284 -284
  78. package/suite/tasks/claim-nominee-uses-slot/fixtures/gateway.py +284 -284
  79. package/suite/tasks/claim-silent-exit/fixtures/gateway.py +235 -235
  80. package/suite/tasks/claim-sse-hang/fixtures/gateway.py +235 -235
  81. package/suite/tasks/claim-unknown-nominee-crash/fixtures/gateway.py +284 -284
  82. package/suite/tasks/rebench-matchms-809-bug/task.json +23 -23
  83. package/suite/tasks/rebench-matchms-809-control/task.json +27 -27
  84. package/suite/tasks/rebench-more-itertools-1028-bug/task.json +23 -23
  85. package/suite/tasks/rebench-more-itertools-1028-control/task.json +27 -27
  86. package/suite/tasks/rebench-narwhals-2864-bug/task.json +23 -23
  87. package/suite/tasks/rebench-narwhals-2864-control/task.json +27 -27
  88. package/suite/tasks/rebench-quantecon-py-769-bug/task.json +23 -23
  89. package/suite/tasks/rebench-quantecon-py-769-control/task.json +27 -27
  90. package/suite/tasks/rebench-scrapy-6867-bug/task.json +23 -23
  91. package/suite/tasks/rebench-scrapy-6867-control/task.json +27 -27
  92. package/ui/bench.css +430 -424
  93. package/ui/bench.js +1756 -1555
  94. package/ui/themes/dark.css +35 -35
  95. package/ui/themes/hardwood.css +52 -52
  96. package/ui/themes/rainforest.css +32 -32
package/metrics/guard.py CHANGED
@@ -1,592 +1,617 @@
1
- """One request check for both servers, the dashboard (serve.py) and the gateway.
2
-
3
- python metrics/guard.py # self-check
4
-
5
- Both bind to 127.0.0.1 by default, which keeps other machines out but not
6
- other web pages: any site open in the browser can send requests to localhost.
7
-
8
- - **Host allow-list, on every request.** DNS rebinding points a site's own
9
- name at 127.0.0.1, so the browser treats our replies as that site's and
10
- lets its script read them (prompts and answers, endpoints, keys' names).
11
- The Host header still carries the site's name, so only our own names, port
12
- included, are answered. More (a Docker service name, a proxy) come from
13
- BENCH_ALLOWED_HOSTS, comma-separated host:port, and BENCHWARMERS_PUBLIC_URL.
14
- - **On POST, JSON only and never cross-site.** A page can send a plain form
15
- or text/plain to us without asking, but not application/json: that needs
16
- a CORS preflight. A browser also stamps Sec-Fetch-Site and Origin, which a
17
- page cannot remove; cross-site or a foreign Origin is refused. API clients
18
- (SDKs, curl, agents, Open WebUI's backend) send neither and are
19
- unaffected. BENCH_ALLOWED_ORIGINS adds a browser UI on another port.
20
- - **CORS on the gateway's routes only**, for a browser app calling the
21
- gateway: a preflight is answered, and Access-Control-Allow-Origin echoed,
22
- only for an allowed origin (BENCH_ALLOWED_ORIGINS, BENCHWARMERS_PUBLIC_URL),
23
- which may then be cross-site. Never `*`, and never Allow-Credentials: a
24
- gateway key travels in a header the page sets, not in a cookie. The
25
- dashboard answers no preflight, so it stays same-origin.
26
- - **A body-size cap**, so one request cannot make a server read gigabytes.
27
-
28
- WHO MAY ASK (RELEASE-PLAN Phase 5, "Token login before Docker")
29
- ---------------------------------------------------------------
30
- Decided by ROUTE, never by what a request happens to carry, so an API call
31
- with a bearer never meets the login cookie and a browser's cookie never opens
32
- the gateway:
33
-
34
- - the gateway's routes (`/v1/...`, `/bench/...`, on either port) take a
35
- GATEWAY KEY as `Authorization: Bearer <key>`: named, revocable, kept only
36
- as a SHA-256 in `config/gateway-keys.json`, shown once when made. A key
37
- never opens the dashboard, so a chat UI or agent given one cannot change
38
- settings, read keys or spend through it.
39
- - every other route is the dashboard's and takes the LOGIN COOKIE, set by
40
- posting BENCHWARMERS_ADMIN_TOKEN to /login: `<expiry>.<HMAC of it keyed by
41
- the token>`, HttpOnly, SameSite=Strict, Secure when the public URL is https.
42
- Changing the token signs every browser out. Open without either: /health,
43
- the login page, and Connect's callback (the provider's redirect back is
44
- cross-site, so a Strict cookie never rides on it; its single-use,
45
- unguessable state is what protects it).
46
- - NO ADMIN TOKEN, NO AUTH: on 127.0.0.1 without BENCHWARMERS_ADMIN_TOKEN both
47
- servers answer as they always did (no login, no key), and the checks above
48
- are what keep web pages out. Binding anywhere else refuses to start without
49
- a token (`startup`), so a server on a network always asks.
50
- BENCHWARMERS_REQUIRE_KEYS asks for gateway keys on this machine too (the
51
- dashboard stays open), so each local app has its own name, scope and cap.
52
-
53
- A key may be limited to some collections (`collections`, none = all: 403
54
- otherwise, and /v1/models lists only those) and to a number of chat
55
- completions a minute (`rpm`: 429 with Retry-After). Each records when it was
56
- last used, and the gateway writes its name on every call it makes.
57
- """
58
-
59
- from __future__ import annotations
60
-
61
- import datetime as dt
62
- import hashlib
63
- import hmac
64
- import http.cookies
65
- import ipaddress
66
- import os
67
- import re
68
- import secrets
69
- import sys
70
- import threading
71
- import time
72
- import urllib.parse
73
- from collections import deque
74
-
75
- from paths import DATA # noqa: E402 -- the install's data folder (metrics/paths.py)
76
- import locked # noqa: E402
77
-
78
- KEYS = DATA / "config" / "gateway-keys.json"
79
- KEY_NAME = re.compile(r"^[a-z0-9][a-z0-9-]{1,31}$")
80
- PREFIX = "bwk_" # every gateway key starts so: a secret scanner, or a person, can tell it
81
- COOKIE = "bw_session"
82
- SESSION_S = 7 * 24 * 3600
83
- MIN_TOKEN = 16
84
- GATEWAY_ROUTES = ("/v1/", "/bench/")
85
-
86
-
87
- class Refused(Exception):
88
- """The request is not served. `status` is the HTTP code to answer with,
89
- `headers` any the answer must carry (WWW-Authenticate, Retry-After)."""
90
-
91
- def __init__(self, status: int, message: str, headers: dict | None = None):
92
- super().__init__(message)
93
- self.status = status
94
- self.headers = headers or {}
95
-
96
-
97
- def _env(name: str) -> set[str]:
98
- return {v.strip().lower() for v in os.environ.get(name, "").split(",") if v.strip()}
99
-
100
-
101
- def bind() -> str:
102
- """The address both servers listen on: BENCH_BIND, else this machine only."""
103
- return os.environ.get("BENCH_BIND", "").strip() or "127.0.0.1"
104
-
105
-
106
- def public_url() -> str:
107
- return os.environ.get("BENCHWARMERS_PUBLIC_URL", "").strip().rstrip("/")
108
-
109
-
110
- def _public() -> tuple[str, str]:
111
- """(host, origin) of BENCHWARMERS_PUBLIC_URL, or ("", "") when unset or not
112
- an http(s) URL. ponytail: the host only; serving under a path prefix behind
113
- a proxy is not supported."""
114
- u = urllib.parse.urlsplit(public_url())
115
- if u.scheme not in ("http", "https") or not u.netloc or "@" in u.netloc:
116
- return "", ""
117
- return u.netloc.lower(), f"{u.scheme}://{u.netloc}".lower()
118
-
119
-
120
- def hosts(port: int) -> set[str]:
121
- public = _public()[0]
122
- return ({f"127.0.0.1:{port}", f"localhost:{port}", f"[::1]:{port}"} | _env("BENCH_ALLOWED_HOSTS")
123
- | ({public} if public else set()))
124
-
125
-
126
- def origins(port: int) -> set[str]:
127
- """Browser origins allowed to call: this server's own, BENCH_ALLOWED_ORIGINS
128
- and the public URL's."""
129
- return {f"http://{h}" for h in hosts(port)} | _env("BENCH_ALLOWED_ORIGINS") | ({_public()[1]} - {""})
130
-
131
-
132
- def cors_origin(headers, port: int) -> str:
133
- """The request's Origin, as sent, when it is allowed; else "". What a
134
- gateway reply echoes in Access-Control-Allow-Origin."""
135
- origin = (headers.get("Origin") or "").strip() if headers else ""
136
- return origin if origin.lower() in origins(port) else ""
137
-
138
-
139
- def check(headers, port: int, post: bool, max_body: int, cors: bool = False) -> int:
140
- """Raise Refused, or return how many body bytes to read (0 for a GET).
141
- `cors`: a gateway route, where an allowed origin may be cross-site."""
142
- host = (headers.get("Host") or "").strip().lower()
143
- if host not in hosts(port):
144
- raise Refused(403, f"unknown Host {host or '(none)'!r}; set BENCHWARMERS_PUBLIC_URL"
145
- " or BENCH_ALLOWED_HOSTS to add one")
146
- if not post:
147
- return 0
148
- origin = (headers.get("Origin") or "").strip().lower()
149
- allowed = origins(port)
150
- if (headers.get("Sec-Fetch-Site") or "").lower() == "cross-site" and not (cors and origin in allowed):
151
- raise Refused(403, "cross-site requests are refused")
152
- if origin and origin not in allowed:
153
- raise Refused(403, f"Origin {origin!r} is not allowed; set BENCH_ALLOWED_ORIGINS to add one")
154
- if not (headers.get("Content-Type") or "").lower().startswith("application/json"):
155
- raise Refused(415, "send Content-Type: application/json")
156
- try:
157
- length = int(headers.get("Content-Length") or 0)
158
- except ValueError:
159
- raise Refused(400, "Content-Length is not a number") from None
160
- if length < 0:
161
- raise Refused(400, "Content-Length is negative")
162
- if length > max_body:
163
- raise Refused(413, f"body is {length:,} bytes; the limit is {max_body:,}")
164
- return length
165
-
166
-
167
- # ---- auth: decided by route ------------------------------------------------
168
-
169
- def admin_token() -> str:
170
- return os.environ.get("BENCHWARMERS_ADMIN_TOKEN", "")
171
-
172
-
173
- def auth_on() -> bool:
174
- return bool(admin_token())
175
-
176
-
177
- def keys_required() -> bool:
178
- """Gateway keys are asked for: always with an admin token, and without one
179
- when BENCHWARMERS_REQUIRE_KEYS is set. Any value but 0/false/no/off counts,
180
- so a misspelt "required" fails closed rather than silently open."""
181
- return auth_on() or os.environ.get("BENCHWARMERS_REQUIRE_KEYS", "").strip().lower() \
182
- not in ("", "0", "false", "no", "off")
183
-
184
-
185
- def _loopback(host: str) -> bool:
186
- if host.lower() == "localhost":
187
- return True
188
- try:
189
- return ipaddress.ip_address(host).is_loopback
190
- except ValueError:
191
- return False
192
-
193
-
194
- def startup(address: str) -> None:
195
- """Refuse (SystemExit) a configuration that must not serve: a network
196
- address without an admin token, a guessable token, a public URL that is
197
- not one."""
198
- token = admin_token()
199
- if token and len(token) < MIN_TOKEN:
200
- # ponytail: no login rate limit; this floor is what makes guessing
201
- # hopeless. Add a limit if short tokens ever have to be allowed.
202
- raise SystemExit(f"BENCHWARMERS_ADMIN_TOKEN is {len(token)} characters; use {MIN_TOKEN} or more,"
203
- ' e.g. python -c "import secrets; print(secrets.token_urlsafe(24))"')
204
- if public_url() and not _public()[0]:
205
- raise SystemExit(f"BENCHWARMERS_PUBLIC_URL must be http(s)://host[:port], not {public_url()!r}")
206
- if not token and not _loopback(address):
207
- raise SystemExit(f"refusing to listen on {address}: beyond this machine every request needs"
208
- " a login, so set BENCHWARMERS_ADMIN_TOKEN first (or BENCH_BIND=127.0.0.1)")
209
-
210
-
211
- def gateway_route(path: str) -> bool:
212
- """A route the gateway answers, on its own port or the dashboard's."""
213
- return path.startswith(GATEWAY_ROUTES)
214
-
215
-
216
- def _sign(expires: int) -> bytes:
217
- return hmac.new(admin_token().encode(), f"session|{expires}".encode(), hashlib.sha256).hexdigest().encode()
218
-
219
-
220
- def login(given: str) -> bool:
221
- """True when `given` is the admin token. Digests compared, so neither the
222
- token's length nor its content leaks through timing."""
223
- token = admin_token()
224
- return bool(token) and hmac.compare_digest(hashlib.sha256(given.encode()).digest(),
225
- hashlib.sha256(token.encode()).digest())
226
-
227
-
228
- def session_cookie() -> str:
229
- """The Set-Cookie value for a good login."""
230
- expires = int(time.time()) + SESSION_S
231
- secure = "; Secure" if _public()[1].startswith("https://") else ""
232
- return (f"{COOKIE}={expires}.{_sign(expires).decode()}; Path=/; Max-Age={SESSION_S};"
233
- f" HttpOnly; SameSite=Strict{secure}")
234
-
235
-
236
- def _session_ok(headers) -> bool:
237
- jar = http.cookies.SimpleCookie()
238
- try:
239
- jar.load(headers.get("Cookie") or "")
240
- except http.cookies.CookieError:
241
- return False
242
- morsel = jar.get(COOKIE)
243
- expires, _, sig = (morsel.value if morsel else "").partition(".")
244
- if not (expires.isascii() and expires.isdigit()) or int(expires) < time.time():
245
- return False
246
- return hmac.compare_digest(sig.encode(), _sign(int(expires)))
247
-
248
-
249
- def admin_caller(headers) -> str:
250
- """Who is asking for a dashboard route: "admin" with a good login cookie,
251
- "local" when no admin token is set. Raises Refused(401)."""
252
- if not auth_on():
253
- return "local"
254
- if not _session_ok(headers):
255
- raise Refused(401, "sign in first: /login")
256
- return "admin"
257
-
258
-
259
- def gateway_caller(headers) -> str:
260
- """Who is asking for a gateway route: the gateway key's name, or "local"
261
- when keys are not required (any bearer is ignored then: SDKs send a dummy
262
- one). Raises Refused(401), saying which of missing, malformed or unknown
263
- it was. The cookie is never looked at."""
264
- if not keys_required():
265
- return "local"
266
- challenge = {"WWW-Authenticate": 'Bearer realm="benchwarmers"'} # RFC 6750: every 401 says how
267
- raw = headers.get("Authorization")
268
- if raw is None:
269
- raise Refused(401, "this route needs a gateway key (Authorization: Bearer <key>);"
270
- " make one on the dashboard's Settings page", challenge)
271
- scheme, _, value = raw.strip().partition(" ")
272
- value = value.strip()
273
- if scheme.lower() != "bearer" or not value:
274
- raise Refused(401, "send the gateway key as 'Authorization: Bearer <key>'", challenge)
275
- name = key_name(value)
276
- if not name:
277
- raise Refused(401, "unknown or revoked gateway key" + ("" if value.startswith(PREFIX) else
278
- f" (gateway keys start {PREFIX}; this is some other key)"), challenge)
279
- _touch(name)
280
- return name
281
-
282
-
283
- def scope(name: str) -> list[str] | None:
284
- """The collections key `name` may call: [] for all (and for "local", no
285
- key), or None for none, when its entry is not a list (a hand edit gone
286
- wrong fails closed)."""
287
- entry = gateway_keys().get(name)
288
- allowed = entry.get("collections") if isinstance(entry, dict) else None
289
- if allowed is None or allowed == []: # not `not allowed`: "" or 0 must fail closed
290
- return []
291
- return allowed if isinstance(allowed, list) and all(isinstance(c, str) for c in allowed) else None
292
-
293
-
294
- def may_call(name: str, collection: str) -> bool:
295
- allowed = scope(name)
296
- return allowed == [] or (allowed is not None and collection in allowed)
297
-
298
-
299
- _SPENT: dict[str, deque] = {}
300
- _SPENT_LOCK = threading.Lock()
301
-
302
-
303
- def spend(name: str) -> None:
304
- """Count one chat completion against key `name`'s per-minute cap, if it has
305
- one; Refused(429) with Retry-After over it. Each call fans out to every
306
- seat, so a client stuck in a loop would otherwise empty free tiers in
307
- seconds. ponytail: counted in this process, so the gateway's port and the
308
- dashboard's count apart and a restart forgets; a shared count needs a
309
- store both read."""
310
- entry = gateway_keys().get(name)
311
- cap = entry.get("rpm") if isinstance(entry, dict) else None
312
- if not isinstance(cap, int) or isinstance(cap, bool) or cap < 1:
313
- return
314
- now = time.monotonic()
315
- with _SPENT_LOCK:
316
- recent = _SPENT.setdefault(name, deque())
317
- while recent and recent[0] <= now - 60:
318
- recent.popleft()
319
- if len(recent) >= cap:
320
- wait = int(recent[0] + 60 - now) + 1
321
- raise Refused(429, f"the gateway key {name} may make {cap} requests a minute;"
322
- f" try again in {wait}s", {"Retry-After": str(wait)})
323
- recent.append(now)
324
-
325
-
326
- # ---- gateway keys ------------------------------------------------------------
327
-
328
- def _digest(value: str) -> bytes:
329
- return hashlib.sha256(value.encode()).hexdigest().encode()
330
-
331
-
332
- def gateway_keys() -> dict[str, dict]:
333
- """name -> {"sha256", "created", "hint", "collections", "rpm", "last_used"}.
334
- Never a key itself."""
335
- return locked.read_json(KEYS)
336
-
337
-
338
- def make_key(name: str, collections: list[str] | None = None, rpm: int | None = None) -> str:
339
- """A new key, returned this once: only its SHA-256 is kept, and its last
340
- four characters, so a key found in a log or a repository can be told
341
- apart (the bwk_ prefix says whose it is). A slow hash would add nothing,
342
- the key being 256 random bits rather than a password. `collections`: the
343
- ones it may call (none = all); `rpm`: chat completions a minute (none = no cap)."""
344
- if not KEY_NAME.match(name):
345
- raise ValueError("name: 2-32 lower-case letters, digits or '-', e.g. open-webui")
346
- if name == "local":
347
- raise ValueError("local is what a call without a key is recorded as; pick another name")
348
- collections = sorted(set(collections or []))
349
- if not all(isinstance(c, str) and KEY_NAME.match(c) for c in collections):
350
- raise ValueError("collections: names such as core, as they follow bench/")
351
- if rpm is not None and (not isinstance(rpm, int) or isinstance(rpm, bool) or not 1 <= rpm <= 10_000):
352
- raise ValueError("requests a minute: a whole number from 1 to 10000, or empty for no cap")
353
- value = PREFIX + secrets.token_urlsafe(32)
354
-
355
- def add(data: dict) -> None:
356
- if name in data:
357
- raise ValueError(f"there is already a key named {name}: revoke it first, or pick another name")
358
- data[name] = {"sha256": _digest(value).decode(), "hint": value[-4:],
359
- "created": dt.datetime.now().astimezone().isoformat(timespec="seconds"),
360
- **({"collections": collections} if collections else {}),
361
- **({"rpm": rpm} if rpm else {})}
362
- locked.update(KEYS, add)
363
- return value
364
-
365
-
366
- def _touch(name: str) -> None:
367
- """Note when key `name` was last used, at most once a minute: chat UIs poll
368
- /v1/models, and a write per poll would be a write per page load."""
369
- now = dt.datetime.now().astimezone()
370
-
371
- def stamp(data: dict) -> None:
372
- entry = data.get(name)
373
- if not isinstance(entry, dict):
374
- return # revoked meanwhile
375
- try:
376
- if (now - dt.datetime.fromisoformat(entry.get("last_used", ""))).total_seconds() < 60:
377
- return
378
- except (TypeError, ValueError):
379
- pass
380
- entry["last_used"] = now.isoformat(timespec="seconds")
381
- try:
382
- locked.update(KEYS, stamp)
383
- except (OSError, TimeoutError): # a busy file must not refuse the call
384
- pass
385
-
386
-
387
- def revoke_key(name: str) -> bool:
388
- return locked.update(KEYS, lambda data: data.pop(name, None)) is not None
389
-
390
-
391
- def key_name(value: str) -> str | None:
392
- """The name of the gateway key `value` is, or None. Every stored digest is
393
- compared in constant time, and all of them, whatever matched."""
394
- if not value:
395
- return None
396
- given, found = _digest(value), None
397
- for name, entry in gateway_keys().items():
398
- stored = str(entry.get("sha256", "")).encode() if isinstance(entry, dict) else b""
399
- if hmac.compare_digest(given, stored):
400
- found = name
401
- return found
402
-
403
-
404
- def access(handler, code) -> None:
405
- """One access-log line on stderr: method, route, status, who. Never the
406
- query (a sign-in's code travels in one) or a Connect callback's id. A
407
- request line too broken to parse has no path or command yet."""
408
- route = re.sub(r"^(/connect/callback)/.*", r"\1/...", (getattr(handler, "path", "") or "").split("?", 1)[0])
409
- sys.stderr.write(f"{dt.datetime.now():%Y-%m-%d %H:%M:%S} {getattr(handler, 'command', '') or '-'}"
410
- f" {ascii(route[:200])[1:-1]}"
411
- f" {getattr(code, 'value', code)} {getattr(handler, 'who', '-')}\n")
412
-
413
-
414
- def _selftest() -> None:
415
- import pathlib
416
- import tempfile
417
- global KEYS
418
- ok = {"Host": "127.0.0.1:8123", "Content-Type": "application/json", "Content-Length": "12"}
419
-
420
- def status(port=8123, post=True, cap=100, cors=False, **h):
421
- try:
422
- check({**ok, **h}, port, post, cap, cors)
423
- return 0
424
- except Refused as r:
425
- return r.status
426
- assert status() == 0
427
- assert status(Host="localhost:8123") == 0 and status(Host="LOCALHOST:8123") == 0
428
- assert status(Host="evil.example:8123") == 403 # DNS rebinding
429
- assert status(Host="127.0.0.1") == 403 # port is part of the name
430
- assert status(port=8122) == 403 # the other server's port
431
- assert status(post=False, Host="evil.example") == 403 # GETs too
432
- assert status(post=False, **{"Content-Type": ""}) == 0 # a GET needs no JSON type
433
- assert status(**{"Content-Type": "text/plain"}) == 415 # the form a page can send freely
434
- assert status(Origin="http://evil.example") == 403
435
- assert status(Origin="http://127.0.0.1:8123") == 0
436
- assert status(**{"Sec-Fetch-Site": "cross-site"}) == 403
437
- assert status(**{"Sec-Fetch-Site": "same-origin"}) == 0
438
- assert status(**{"Content-Length": "101"}) == 413
439
- assert status(**{"Content-Length": "x"}) == 400
440
- assert status(**{"Content-Length": "-1"}) == 400
441
- saved = {k: os.environ.pop(k, None) for k in ("BENCH_ALLOWED_HOSTS", "BENCH_ALLOWED_ORIGINS",
442
- "BENCHWARMERS_PUBLIC_URL", "BENCHWARMERS_ADMIN_TOKEN",
443
- "BENCHWARMERS_REQUIRE_KEYS")}
444
- try:
445
- os.environ["BENCH_ALLOWED_HOSTS"] = "bench:8123"
446
- os.environ["BENCH_ALLOWED_ORIGINS"] = "http://localhost:3000"
447
- assert status(Host="bench:8123") == 0
448
- assert status(Origin="http://localhost:3000") == 0
449
- # The public URL is a Host and an Origin, https included.
450
- os.environ["BENCHWARMERS_PUBLIC_URL"] = "https://Bench.example.com/"
451
- assert status(Host="bench.example.com", Origin="https://bench.example.com") == 0
452
- assert status(Host="bench.example.com", Origin="https://evil.example") == 403
453
- # CORS: on a gateway route an allowed origin may be cross-site; the dashboard's rule stands.
454
- site = {"Sec-Fetch-Site": "cross-site"}
455
- assert status(cors=True, Origin="http://localhost:3000", **site) == 0
456
- assert status(Origin="http://localhost:3000", **site) == 403
457
- assert status(cors=True, Origin="http://evil.example", **site) == 403
458
- assert status(cors=True, **site) == 403 # no Origin: not a CORS call
459
- assert cors_origin({"Origin": "http://localhost:3000"}, 8123) == "http://localhost:3000"
460
- assert cors_origin({"Origin": "http://evil.example"}, 8123) == "" and cors_origin({}, 8123) == ""
461
-
462
- # Binding: this machine needs nothing; anywhere else needs a long enough token.
463
- def refused(address: str) -> bool:
464
- try:
465
- startup(address)
466
- return False
467
- except SystemExit:
468
- return True
469
- assert not refused("127.0.0.1") and not refused("localhost") and not refused("::1")
470
- assert refused("0.0.0.0") and refused("") and refused("192.0.2.7")
471
- os.environ["BENCHWARMERS_ADMIN_TOKEN"] = "short"
472
- assert refused("127.0.0.1")
473
- token = os.environ["BENCHWARMERS_ADMIN_TOKEN"] = "t" * 24
474
- assert not refused("0.0.0.0")
475
- os.environ["BENCHWARMERS_PUBLIC_URL"] = "bench.example.com" # no scheme
476
- assert refused("0.0.0.0")
477
- os.environ["BENCHWARMERS_PUBLIC_URL"] = "https://bench.example.com"
478
-
479
- # The login cookie: only the token makes one; tampered, expired or foreign ones fail.
480
- assert login(token) and not login("t" * 23) and not login("")
481
- cookie = session_cookie()
482
- assert "HttpOnly" in cookie and "SameSite=Strict" in cookie and "Secure" in cookie and token not in cookie
483
- value = cookie.split(";", 1)[0]
484
- assert admin_caller({"Cookie": value}) == "admin"
485
- assert admin_caller({"Cookie": f"theme=x; {value}"}) == "admin"
486
- for bad in ("", value + "0", value.replace(f"{COOKIE}=", f"{COOKIE}=1"), f"{COOKIE}=1.abc",
487
- f"{COOKIE}=²." + "0" * 64, f"{COOKIE}=0.{_sign(0).decode()}", 'x="'):
488
- try:
489
- admin_caller({"Cookie": bad})
490
- raise AssertionError(f"accepted cookie {bad!r}")
491
- except Refused as r:
492
- assert r.status == 401
493
- os.environ["BENCHWARMERS_ADMIN_TOKEN"] = "u" * 24 # a new token signs everyone out
494
- try:
495
- admin_caller({"Cookie": value})
496
- raise AssertionError("an old token's cookie still works")
497
- except Refused:
498
- pass
499
-
500
- # Gateway keys: shown once, kept hashed, constant-time, revocable; never a cookie's job.
501
- with tempfile.TemporaryDirectory() as tmp:
502
- real, KEYS = KEYS, pathlib.Path(tmp) / "gateway-keys.json"
503
- try:
504
- key = make_key("chat-ui")
505
- assert key.startswith("bwk_") and key not in KEYS.read_text(encoding="utf-8")
506
- assert gateway_keys()["chat-ui"]["sha256"] == hashlib.sha256(key.encode()).hexdigest()
507
- for bad in ("Bad Name", "x", "chat-ui"): # form, length, taken
508
- try:
509
- make_key(bad)
510
- raise AssertionError(f"made a key named {bad!r}")
511
- except ValueError:
512
- pass
513
- other = make_key("agent")
514
- assert key_name(key) == "chat-ui" and key_name(other) == "agent"
515
- assert key_name("") is None and key_name(key + "x") is None and key_name("bwk_é") is None
516
- assert gateway_caller({"Authorization": f"Bearer {key}"}) == "chat-ui"
517
- assert gateway_caller({"Authorization": f"bearer {key} "}) == "chat-ui"
518
- entry = gateway_keys()["chat-ui"]
519
- assert entry["hint"] == key[-4:] and entry["last_used"] and "rpm" not in entry, entry
520
- # Each refusal says which it was, and how to authenticate (RFC 6750).
521
- for headers, said in (({}, "needs a gateway key"), ({"Cookie": value}, "needs a gateway key"),
522
- ({"Authorization": key}, "Bearer <key>"), ({"Authorization": "Bearer "}, "Bearer <key>"),
523
- ({"Authorization": f"Basic {key}"}, "Bearer <key>"),
524
- ({"Authorization": "Bearer sk-dummy"}, "some other key"),
525
- ({"Authorization": f"Bearer {key}x"}, "unknown or revoked")):
526
- try:
527
- gateway_caller(headers)
528
- raise AssertionError(f"gateway let in {headers}")
529
- except Refused as r:
530
- assert r.status == 401 and said in str(r), (headers, str(r))
531
- assert r.headers["WWW-Authenticate"].startswith("Bearer"), r.headers
532
- # Scope and cap, set when made: checked per collection, and per minute.
533
- for name, kw in (("scoped", {"collections": ["Core"]}), ("scoped", {"rpm": 0}),
534
- ("scoped", {"rpm": True}), ("scoped", {"rpm": "5"}), ("local", {})):
535
- try:
536
- make_key(name, **kw)
537
- raise AssertionError(f"made a key with {kw}")
538
- except ValueError:
539
- pass
540
- make_key("scoped", ["vision", "core", "core"], 2)
541
- assert gateway_keys()["scoped"]["collections"] == ["core", "vision"]
542
- assert may_call("scoped", "vision") and not may_call("scoped", "reviewers")
543
- assert may_call("chat-ui", "reviewers") and may_call("local", "reviewers")
544
- spend("scoped")
545
- spend("scoped")
546
- try:
547
- spend("scoped")
548
- raise AssertionError("a third request in a minute on a cap of 2")
549
- except Refused as r:
550
- assert r.status == 429 and 1 <= int(r.headers["Retry-After"]) <= 61, r.headers
551
- for _ in range(5):
552
- spend("chat-ui") # no cap
553
- for bad in ("core", "", 0, False, [1]): # a bad edit fails closed
554
- locked.update(KEYS, lambda d: d["scoped"].__setitem__("collections", bad))
555
- assert scope("scoped") is None and not may_call("scoped", "core"), bad
556
- try: # a key never opens the dashboard
557
- admin_caller({"Authorization": f"Bearer {key}"})
558
- raise AssertionError("a gateway key opened the dashboard")
559
- except Refused:
560
- pass
561
- assert revoke_key("chat-ui") and not revoke_key("chat-ui") and key_name(key) is None
562
- assert key_name(other) == "agent"
563
- # No admin token: as before keys existed, whatever is sent.
564
- del os.environ["BENCHWARMERS_ADMIN_TOKEN"]
565
- assert gateway_caller({}) == "local" and gateway_caller({"Authorization": "Bearer x"}) == "local"
566
- assert admin_caller({}) == "local"
567
- # ...unless keys are required on this machine too; a misspelt value still asks.
568
- for v in ("1", "required", "TRUE"):
569
- os.environ["BENCHWARMERS_REQUIRE_KEYS"] = v
570
- assert gateway_caller({"Authorization": f"Bearer {other}"}) == "agent"
571
- try:
572
- gateway_caller({"Authorization": "Bearer x"})
573
- raise AssertionError(f"BENCHWARMERS_REQUIRE_KEYS={v} let a dummy bearer in")
574
- except Refused as r:
575
- assert r.status == 401
576
- assert admin_caller({}) == "local" # the dashboard stays open
577
- os.environ["BENCHWARMERS_REQUIRE_KEYS"] = "off"
578
- assert not keys_required() and gateway_caller({}) == "local"
579
- finally:
580
- KEYS = real
581
- assert gateway_route("/v1/models") and gateway_route("/bench/verdict")
582
- assert not gateway_route("/api/endpoints") and not gateway_route("/v1") and not gateway_route("//v1/x")
583
- finally:
584
- for k, v in saved.items():
585
- os.environ.pop(k, None)
586
- if v is not None:
587
- os.environ[k] = v
588
- print("ok")
589
-
590
-
591
- if __name__ == "__main__":
592
- _selftest()
1
+ """One request check for both servers, the dashboard (serve.py) and the gateway.
2
+
3
+ python metrics/guard.py # self-check
4
+
5
+ Both bind to 127.0.0.1 by default, which keeps other machines out but not
6
+ other web pages: any site open in the browser can send requests to localhost.
7
+
8
+ - **Host allow-list, on every request.** DNS rebinding points a site's own
9
+ name at 127.0.0.1, so the browser treats our replies as that site's and
10
+ lets its script read them (prompts and answers, endpoints, keys' names).
11
+ The Host header still carries the site's name, so only our own names, port
12
+ included, are answered. More (a Docker service name, a proxy) come from
13
+ BENCH_ALLOWED_HOSTS, comma-separated host:port, and BENCHWARMERS_PUBLIC_URL.
14
+ - **On POST, JSON only and never cross-site.** A page can send a plain form
15
+ or text/plain to us without asking, but not application/json: that needs
16
+ a CORS preflight. A browser also stamps Sec-Fetch-Site and Origin, which a
17
+ page cannot remove; cross-site or a foreign Origin is refused. API clients
18
+ (SDKs, curl, agents, Open WebUI's backend) send neither and are
19
+ unaffected. BENCH_ALLOWED_ORIGINS adds a browser UI on another port.
20
+ - **CORS on the gateway's routes only**, for a browser app calling the
21
+ gateway: a preflight is answered, and Access-Control-Allow-Origin echoed,
22
+ only for an allowed origin (BENCH_ALLOWED_ORIGINS, BENCHWARMERS_PUBLIC_URL),
23
+ which may then be cross-site. Never `*`, and never Allow-Credentials: a
24
+ gateway key travels in a header the page sets, not in a cookie. The
25
+ dashboard answers no preflight, so it stays same-origin.
26
+ - **A body-size cap**, so one request cannot make a server read gigabytes.
27
+
28
+ WHO MAY ASK (RELEASE-PLAN Phase 5, "Token login before Docker")
29
+ ---------------------------------------------------------------
30
+ Decided by ROUTE, never by what a request happens to carry, so an API call
31
+ with a bearer never meets the login cookie and a browser's cookie never opens
32
+ the gateway:
33
+
34
+ - the gateway's routes (`/v1/...`, `/bench/...`, on either port) take a
35
+ GATEWAY KEY as `Authorization: Bearer <key>`: named, revocable, kept only
36
+ as a SHA-256 in `config/gateway-keys.json`, shown once when made. A key
37
+ never opens the dashboard, so a chat UI or agent given one cannot change
38
+ settings, read keys or spend through it.
39
+ - every other route is the dashboard's and takes the LOGIN COOKIE, set by
40
+ posting BENCHWARMERS_ADMIN_TOKEN to /login: `<expiry>.<HMAC of it keyed by
41
+ the token>`, HttpOnly, SameSite=Strict, Secure when the public URL is https.
42
+ Changing the token signs every browser out. Open without either: /health,
43
+ the login page, and Connect's callback (the provider's redirect back is
44
+ cross-site, so a Strict cookie never rides on it; its single-use,
45
+ unguessable state is what protects it).
46
+ - NO ADMIN TOKEN, NO AUTH: on 127.0.0.1 without BENCHWARMERS_ADMIN_TOKEN both
47
+ servers answer as they always did (no login, no key), and the checks above
48
+ are what keep web pages out. Binding anywhere else refuses to start without
49
+ a token (`startup`), so a server on a network always asks.
50
+ BENCHWARMERS_REQUIRE_KEYS asks for gateway keys on this machine too (the
51
+ dashboard stays open), so each local app has its own name, scope and cap.
52
+
53
+ A key may be limited to some collections (`collections`, none = all: 403
54
+ otherwise, and /v1/models lists only those) and to a number of chat
55
+ completions a minute (`rpm`: 429 with Retry-After). Each records when it was
56
+ last used, and the gateway writes its name on every call it makes.
57
+ """
58
+
59
+ from __future__ import annotations
60
+
61
+ import datetime as dt
62
+ import hashlib
63
+ import hmac
64
+ import http.cookies
65
+ import ipaddress
66
+ import os
67
+ import re
68
+ import secrets
69
+ import sys
70
+ import threading
71
+ import time
72
+ import urllib.parse
73
+ from collections import deque
74
+
75
+ from paths import DATA # noqa: E402 -- the install's data folder (metrics/paths.py)
76
+ import locked # noqa: E402
77
+
78
+ KEYS = DATA / "config" / "gateway-keys.json"
79
+ KEY_NAME = re.compile(r"^[a-z0-9][a-z0-9-]{1,31}$")
80
+ PREFIX = "bwk_" # every gateway key starts so: a secret scanner, or a person, can tell it
81
+ COOKIE = "bw_session"
82
+ SESSION_S = 7 * 24 * 3600
83
+ MIN_TOKEN = 16
84
+ GATEWAY_ROUTES = ("/v1/", "/bench/")
85
+
86
+
87
+ class Refused(Exception):
88
+ """The request is not served. `status` is the HTTP code to answer with,
89
+ `headers` any the answer must carry (WWW-Authenticate, Retry-After)."""
90
+
91
+ def __init__(self, status: int, message: str, headers: dict | None = None):
92
+ super().__init__(message)
93
+ self.status = status
94
+ self.headers = headers or {}
95
+
96
+
97
+ def _env(name: str) -> set[str]:
98
+ return {v.strip().lower() for v in os.environ.get(name, "").split(",") if v.strip()}
99
+
100
+
101
+ def bind() -> str:
102
+ """The address both servers listen on: BENCH_BIND, else this machine only."""
103
+ return os.environ.get("BENCH_BIND", "").strip() or "127.0.0.1"
104
+
105
+
106
+ def public_url() -> str:
107
+ return os.environ.get("BENCHWARMERS_PUBLIC_URL", "").strip().rstrip("/")
108
+
109
+
110
+ def _public() -> tuple[str, str]:
111
+ """(host, origin) of BENCHWARMERS_PUBLIC_URL, or ("", "") when unset or not
112
+ an http(s) URL. ponytail: the host only; serving under a path prefix behind
113
+ a proxy is not supported."""
114
+ u = urllib.parse.urlsplit(public_url())
115
+ if u.scheme not in ("http", "https") or not u.netloc or "@" in u.netloc:
116
+ return "", ""
117
+ return u.netloc.lower(), f"{u.scheme}://{u.netloc}".lower()
118
+
119
+
120
+ def hosts(port: int) -> set[str]:
121
+ public = _public()[0]
122
+ return ({f"127.0.0.1:{port}", f"localhost:{port}", f"[::1]:{port}"} | _env("BENCH_ALLOWED_HOSTS")
123
+ | ({public} if public else set()))
124
+
125
+
126
+ def origins(port: int) -> set[str]:
127
+ """Browser origins allowed to call: this server's own, BENCH_ALLOWED_ORIGINS
128
+ and the public URL's."""
129
+ return {f"http://{h}" for h in hosts(port)} | _env("BENCH_ALLOWED_ORIGINS") | ({_public()[1]} - {""})
130
+
131
+
132
+ def cors_origin(headers, port: int) -> str:
133
+ """The request's Origin, as sent, when it is allowed; else "". What a
134
+ gateway reply echoes in Access-Control-Allow-Origin."""
135
+ origin = (headers.get("Origin") or "").strip() if headers else ""
136
+ return origin if origin.lower() in origins(port) else ""
137
+
138
+
139
+ def check(headers, port: int, post: bool, max_body: int, cors: bool = False) -> int:
140
+ """Raise Refused, or return how many body bytes to read (0 for a GET).
141
+ `cors`: a gateway route, where an allowed origin may be cross-site."""
142
+ host = (headers.get("Host") or "").strip().lower()
143
+ if host not in hosts(port):
144
+ raise Refused(403, f"unknown Host {host or '(none)'!r}; set BENCHWARMERS_PUBLIC_URL"
145
+ " or BENCH_ALLOWED_HOSTS to add one")
146
+ if not post:
147
+ return 0
148
+ origin = (headers.get("Origin") or "").strip().lower()
149
+ allowed = origins(port)
150
+ if (headers.get("Sec-Fetch-Site") or "").lower() == "cross-site" and not (cors and origin in allowed):
151
+ raise Refused(403, "cross-site requests are refused")
152
+ if origin and origin not in allowed:
153
+ raise Refused(403, f"Origin {origin!r} is not allowed; set BENCH_ALLOWED_ORIGINS to add one")
154
+ if not (headers.get("Content-Type") or "").lower().startswith("application/json"):
155
+ raise Refused(415, "send Content-Type: application/json")
156
+ try:
157
+ length = int(headers.get("Content-Length") or 0)
158
+ except ValueError:
159
+ raise Refused(400, "Content-Length is not a number") from None
160
+ if length < 0:
161
+ raise Refused(400, "Content-Length is negative")
162
+ if length > max_body:
163
+ raise Refused(413, f"body is {length:,} bytes; the limit is {max_body:,}")
164
+ return length
165
+
166
+
167
+ # ---- auth: decided by route ------------------------------------------------
168
+
169
+ def admin_token() -> str:
170
+ return os.environ.get("BENCHWARMERS_ADMIN_TOKEN", "")
171
+
172
+
173
+ def auth_on() -> bool:
174
+ return bool(admin_token())
175
+
176
+
177
+ def keys_required() -> bool:
178
+ """Gateway keys are asked for: always with an admin token, and without one
179
+ when BENCHWARMERS_REQUIRE_KEYS is set. Any value but 0/false/no/off counts,
180
+ so a misspelt "required" fails closed rather than silently open."""
181
+ return auth_on() or os.environ.get("BENCHWARMERS_REQUIRE_KEYS", "").strip().lower() \
182
+ not in ("", "0", "false", "no", "off")
183
+
184
+
185
+ def _loopback(host: str) -> bool:
186
+ if host.lower() == "localhost":
187
+ return True
188
+ try:
189
+ return ipaddress.ip_address(host).is_loopback
190
+ except ValueError:
191
+ return False
192
+
193
+
194
+ def startup(address: str) -> None:
195
+ """Refuse (SystemExit) a configuration that must not serve: a network
196
+ address without an admin token, a guessable token, a public URL that is
197
+ not one."""
198
+ token = admin_token()
199
+ if token and len(token) < MIN_TOKEN:
200
+ # ponytail: no login rate limit; this floor is what makes guessing
201
+ # hopeless. Add a limit if short tokens ever have to be allowed.
202
+ raise SystemExit(f"BENCHWARMERS_ADMIN_TOKEN is {len(token)} characters; use {MIN_TOKEN} or more,"
203
+ ' e.g. python -c "import secrets; print(secrets.token_urlsafe(24))"')
204
+ if public_url() and not _public()[0]:
205
+ raise SystemExit(f"BENCHWARMERS_PUBLIC_URL must be http(s)://host[:port], not {public_url()!r}")
206
+ if not token and not _loopback(address):
207
+ raise SystemExit(f"refusing to listen on {address}: beyond this machine every request needs"
208
+ " a login, so set BENCHWARMERS_ADMIN_TOKEN first (or BENCH_BIND=127.0.0.1)")
209
+
210
+
211
+ def gateway_route(path: str) -> bool:
212
+ """A route the gateway answers, on its own port or the dashboard's."""
213
+ return path.startswith(GATEWAY_ROUTES)
214
+
215
+
216
+ def _sign(expires: int) -> bytes:
217
+ return hmac.new(admin_token().encode(), f"session|{expires}".encode(), hashlib.sha256).hexdigest().encode()
218
+
219
+
220
+ def login(given: str) -> bool:
221
+ """True when `given` is the admin token. Digests compared, so neither the
222
+ token's length nor its content leaks through timing."""
223
+ token = admin_token()
224
+ return bool(token) and hmac.compare_digest(hashlib.sha256(given.encode()).digest(),
225
+ hashlib.sha256(token.encode()).digest())
226
+
227
+
228
+ def session_cookie() -> str:
229
+ """The Set-Cookie value for a good login."""
230
+ expires = int(time.time()) + SESSION_S
231
+ secure = "; Secure" if _public()[1].startswith("https://") else ""
232
+ return (f"{COOKIE}={expires}.{_sign(expires).decode()}; Path=/; Max-Age={SESSION_S};"
233
+ f" HttpOnly; SameSite=Strict{secure}")
234
+
235
+
236
+ def _session_ok(headers) -> bool:
237
+ jar = http.cookies.SimpleCookie()
238
+ try:
239
+ jar.load(headers.get("Cookie") or "")
240
+ except http.cookies.CookieError:
241
+ return False
242
+ morsel = jar.get(COOKIE)
243
+ expires, _, sig = (morsel.value if morsel else "").partition(".")
244
+ if not (expires.isascii() and expires.isdigit()) or int(expires) < time.time():
245
+ return False
246
+ return hmac.compare_digest(sig.encode(), _sign(int(expires)))
247
+
248
+
249
+ def admin_caller(headers) -> str:
250
+ """Who is asking for a dashboard route: "admin" with a good login cookie
251
+ or the admin token itself as `Authorization: Bearer` (the `bench` CLI),
252
+ "local" when no admin token is set. Raises Refused(401).
253
+
254
+ The bearer is safe beside the cookie: a page cannot make a browser send an
255
+ Authorization header cross-site without a preflight, which the dashboard
256
+ never answers. A gateway key is not the admin token, so it still never
257
+ opens a dashboard route."""
258
+ if not auth_on():
259
+ return "local"
260
+ scheme, _, given = (headers.get("Authorization") or "").strip().partition(" ")
261
+ if scheme.lower() == "bearer" and given.strip():
262
+ if login(given.strip()):
263
+ return "admin"
264
+ raise Refused(401, "that bearer is not the admin token (a gateway key opens /v1 only)")
265
+ if not _session_ok(headers):
266
+ raise Refused(401, "sign in first: /login, or send the admin token as Authorization: Bearer")
267
+ return "admin"
268
+
269
+
270
+ def gateway_caller(headers) -> str:
271
+ """Who is asking for a gateway route: the gateway key's name, or "local"
272
+ when keys are not required (any bearer is ignored then: SDKs send a dummy
273
+ one). Raises Refused(401), saying which of missing, malformed or unknown
274
+ it was. The cookie is never looked at."""
275
+ if not keys_required():
276
+ return "local"
277
+ challenge = {"WWW-Authenticate": 'Bearer realm="benchwarmers"'} # RFC 6750: every 401 says how
278
+ raw = headers.get("Authorization")
279
+ if raw is None:
280
+ raise Refused(401, "this route needs a gateway key (Authorization: Bearer <key>);"
281
+ " make one on the dashboard's Settings page", challenge)
282
+ scheme, _, value = raw.strip().partition(" ")
283
+ value = value.strip()
284
+ if scheme.lower() != "bearer" or not value:
285
+ raise Refused(401, "send the gateway key as 'Authorization: Bearer <key>'", challenge)
286
+ name = key_name(value)
287
+ if not name:
288
+ raise Refused(401, "unknown or revoked gateway key" + ("" if value.startswith(PREFIX) else
289
+ f" (gateway keys start {PREFIX}; this is some other key)"), challenge)
290
+ _touch(name)
291
+ return name
292
+
293
+
294
+ def scope(name: str) -> list[str] | None:
295
+ """The collections key `name` may call: [] for all (and for "local", no
296
+ key), or None for none, when its entry is not a list (a hand edit gone
297
+ wrong fails closed)."""
298
+ entry = gateway_keys().get(name)
299
+ allowed = entry.get("collections") if isinstance(entry, dict) else None
300
+ if allowed is None or allowed == []: # not `not allowed`: "" or 0 must fail closed
301
+ return []
302
+ return allowed if isinstance(allowed, list) and all(isinstance(c, str) for c in allowed) else None
303
+
304
+
305
+ def may_call(name: str, collection: str) -> bool:
306
+ allowed = scope(name)
307
+ return allowed == [] or (allowed is not None and collection in allowed)
308
+
309
+
310
+ _SPENT: dict[str, deque] = {}
311
+ _SPENT_LOCK = threading.Lock()
312
+
313
+
314
+ def spend(name: str) -> None:
315
+ """Count one chat completion against key `name`'s per-minute cap, if it has
316
+ one; Refused(429) with Retry-After over it. Each call fans out to every
317
+ seat, so a client stuck in a loop would otherwise empty free tiers in
318
+ seconds. ponytail: counted in this process, so the gateway's port and the
319
+ dashboard's count apart and a restart forgets; a shared count needs a
320
+ store both read."""
321
+ entry = gateway_keys().get(name)
322
+ cap = entry.get("rpm") if isinstance(entry, dict) else None
323
+ if not isinstance(cap, int) or isinstance(cap, bool) or cap < 1:
324
+ return
325
+ now = time.monotonic()
326
+ with _SPENT_LOCK:
327
+ recent = _SPENT.setdefault(name, deque())
328
+ while recent and recent[0] <= now - 60:
329
+ recent.popleft()
330
+ if len(recent) >= cap:
331
+ wait = int(recent[0] + 60 - now) + 1
332
+ raise Refused(429, f"the gateway key {name} may make {cap} requests a minute;"
333
+ f" try again in {wait}s", {"Retry-After": str(wait)})
334
+ recent.append(now)
335
+
336
+
337
+ # ---- gateway keys ------------------------------------------------------------
338
+
339
+ def _digest(value: str) -> bytes:
340
+ return hashlib.sha256(value.encode()).hexdigest().encode()
341
+
342
+
343
+ def gateway_keys() -> dict[str, dict]:
344
+ """name -> {"sha256", "created", "hint", "collections", "rpm", "last_used"}.
345
+ Never a key itself."""
346
+ return locked.read_json(KEYS)
347
+
348
+
349
+ def make_key(name: str, collections: list[str] | None = None, rpm: int | None = None) -> str:
350
+ """A new key, returned this once: only its SHA-256 is kept, and its last
351
+ four characters, so a key found in a log or a repository can be told
352
+ apart (the bwk_ prefix says whose it is). A slow hash would add nothing,
353
+ the key being 256 random bits rather than a password. `collections`: the
354
+ ones it may call (none = all); `rpm`: chat completions a minute (none = no cap)."""
355
+ if not KEY_NAME.match(name):
356
+ raise ValueError("name: 2-32 lower-case letters, digits or '-', e.g. open-webui")
357
+ if name == "local":
358
+ raise ValueError("local is what a call without a key is recorded as; pick another name")
359
+ collections = sorted(set(collections or []))
360
+ if not all(isinstance(c, str) and KEY_NAME.match(c) for c in collections):
361
+ raise ValueError("collections: names such as core, as they follow bench/")
362
+ if rpm is not None and (not isinstance(rpm, int) or isinstance(rpm, bool) or not 1 <= rpm <= 10_000):
363
+ raise ValueError("requests a minute: a whole number from 1 to 10000, or empty for no cap")
364
+ value = PREFIX + secrets.token_urlsafe(32)
365
+
366
+ def add(data: dict) -> None:
367
+ if name in data:
368
+ raise ValueError(f"there is already a key named {name}: revoke it first, or pick another name")
369
+ data[name] = {"sha256": _digest(value).decode(), "hint": value[-4:],
370
+ "created": dt.datetime.now().astimezone().isoformat(timespec="seconds"),
371
+ **({"collections": collections} if collections else {}),
372
+ **({"rpm": rpm} if rpm else {})}
373
+ locked.update(KEYS, add)
374
+ return value
375
+
376
+
377
+ def _touch(name: str) -> None:
378
+ """Note when key `name` was last used, at most once a minute: chat UIs poll
379
+ /v1/models, and a write per poll would be a write per page load."""
380
+ now = dt.datetime.now().astimezone()
381
+
382
+ def stamp(data: dict) -> None:
383
+ entry = data.get(name)
384
+ if not isinstance(entry, dict):
385
+ return # revoked meanwhile
386
+ try:
387
+ if (now - dt.datetime.fromisoformat(entry.get("last_used", ""))).total_seconds() < 60:
388
+ return
389
+ except (TypeError, ValueError):
390
+ pass
391
+ entry["last_used"] = now.isoformat(timespec="seconds")
392
+ try:
393
+ locked.update(KEYS, stamp)
394
+ except (OSError, TimeoutError): # a busy file must not refuse the call
395
+ pass
396
+
397
+
398
+ def revoke_key(name: str) -> bool:
399
+ return locked.update(KEYS, lambda data: data.pop(name, None)) is not None
400
+
401
+
402
+ def key_name(value: str) -> str | None:
403
+ """The name of the gateway key `value` is, or None. Every stored digest is
404
+ compared in constant time, and all of them, whatever matched."""
405
+ if not value:
406
+ return None
407
+ given, found = _digest(value), None
408
+ for name, entry in gateway_keys().items():
409
+ stored = str(entry.get("sha256", "")).encode() if isinstance(entry, dict) else b""
410
+ if hmac.compare_digest(given, stored):
411
+ found = name
412
+ return found
413
+
414
+
415
+ def access(handler, code) -> None:
416
+ """One access-log line on stderr: method, route, status, who. Never the
417
+ query (a sign-in's code travels in one) or a Connect callback's id. A
418
+ request line too broken to parse has no path or command yet."""
419
+ route = re.sub(r"^(/connect/callback)/.*", r"\1/...", (getattr(handler, "path", "") or "").split("?", 1)[0])
420
+ sys.stderr.write(f"{dt.datetime.now():%Y-%m-%d %H:%M:%S} {getattr(handler, 'command', '') or '-'}"
421
+ f" {ascii(route[:200])[1:-1]}"
422
+ f" {getattr(code, 'value', code)} {getattr(handler, 'who', '-')}\n")
423
+
424
+
425
+ def _selftest() -> None:
426
+ import pathlib
427
+ import tempfile
428
+ global KEYS
429
+ ok = {"Host": "127.0.0.1:8123", "Content-Type": "application/json", "Content-Length": "12"}
430
+
431
+ def status(port=8123, post=True, cap=100, cors=False, **h):
432
+ try:
433
+ check({**ok, **h}, port, post, cap, cors)
434
+ return 0
435
+ except Refused as r:
436
+ return r.status
437
+ assert status() == 0
438
+ assert status(Host="localhost:8123") == 0 and status(Host="LOCALHOST:8123") == 0
439
+ assert status(Host="evil.example:8123") == 403 # DNS rebinding
440
+ assert status(Host="127.0.0.1") == 403 # port is part of the name
441
+ assert status(port=8122) == 403 # the other server's port
442
+ assert status(post=False, Host="evil.example") == 403 # GETs too
443
+ assert status(post=False, **{"Content-Type": ""}) == 0 # a GET needs no JSON type
444
+ assert status(**{"Content-Type": "text/plain"}) == 415 # the form a page can send freely
445
+ assert status(Origin="http://evil.example") == 403
446
+ assert status(Origin="http://127.0.0.1:8123") == 0
447
+ assert status(**{"Sec-Fetch-Site": "cross-site"}) == 403
448
+ assert status(**{"Sec-Fetch-Site": "same-origin"}) == 0
449
+ assert status(**{"Content-Length": "101"}) == 413
450
+ assert status(**{"Content-Length": "x"}) == 400
451
+ assert status(**{"Content-Length": "-1"}) == 400
452
+ saved = {k: os.environ.pop(k, None) for k in ("BENCH_ALLOWED_HOSTS", "BENCH_ALLOWED_ORIGINS",
453
+ "BENCHWARMERS_PUBLIC_URL", "BENCHWARMERS_ADMIN_TOKEN",
454
+ "BENCHWARMERS_REQUIRE_KEYS")}
455
+ try:
456
+ os.environ["BENCH_ALLOWED_HOSTS"] = "bench:8123"
457
+ os.environ["BENCH_ALLOWED_ORIGINS"] = "http://localhost:3000"
458
+ assert status(Host="bench:8123") == 0
459
+ assert status(Origin="http://localhost:3000") == 0
460
+ # The public URL is a Host and an Origin, https included.
461
+ os.environ["BENCHWARMERS_PUBLIC_URL"] = "https://Bench.example.com/"
462
+ assert status(Host="bench.example.com", Origin="https://bench.example.com") == 0
463
+ assert status(Host="bench.example.com", Origin="https://evil.example") == 403
464
+ # CORS: on a gateway route an allowed origin may be cross-site; the dashboard's rule stands.
465
+ site = {"Sec-Fetch-Site": "cross-site"}
466
+ assert status(cors=True, Origin="http://localhost:3000", **site) == 0
467
+ assert status(Origin="http://localhost:3000", **site) == 403
468
+ assert status(cors=True, Origin="http://evil.example", **site) == 403
469
+ assert status(cors=True, **site) == 403 # no Origin: not a CORS call
470
+ assert cors_origin({"Origin": "http://localhost:3000"}, 8123) == "http://localhost:3000"
471
+ assert cors_origin({"Origin": "http://evil.example"}, 8123) == "" and cors_origin({}, 8123) == ""
472
+
473
+ # Binding: this machine needs nothing; anywhere else needs a long enough token.
474
+ def refused(address: str) -> bool:
475
+ try:
476
+ startup(address)
477
+ return False
478
+ except SystemExit:
479
+ return True
480
+ assert not refused("127.0.0.1") and not refused("localhost") and not refused("::1")
481
+ assert refused("0.0.0.0") and refused("") and refused("192.0.2.7")
482
+ os.environ["BENCHWARMERS_ADMIN_TOKEN"] = "short"
483
+ assert refused("127.0.0.1")
484
+ token = os.environ["BENCHWARMERS_ADMIN_TOKEN"] = "t" * 24
485
+ assert not refused("0.0.0.0")
486
+ os.environ["BENCHWARMERS_PUBLIC_URL"] = "bench.example.com" # no scheme
487
+ assert refused("0.0.0.0")
488
+ os.environ["BENCHWARMERS_PUBLIC_URL"] = "https://bench.example.com"
489
+
490
+ # The login cookie: only the token makes one; tampered, expired or foreign ones fail.
491
+ assert login(token) and not login("t" * 23) and not login("")
492
+ cookie = session_cookie()
493
+ assert "HttpOnly" in cookie and "SameSite=Strict" in cookie and "Secure" in cookie and token not in cookie
494
+ value = cookie.split(";", 1)[0]
495
+ assert admin_caller({"Cookie": value}) == "admin"
496
+ assert admin_caller({"Cookie": f"theme=x; {value}"}) == "admin"
497
+ for bad in ("", value + "0", value.replace(f"{COOKIE}=", f"{COOKIE}=1"), f"{COOKIE}=1.abc",
498
+ f"{COOKIE}=²." + "0" * 64, f"{COOKIE}=0.{_sign(0).decode()}", 'x="'):
499
+ try:
500
+ admin_caller({"Cookie": bad})
501
+ raise AssertionError(f"accepted cookie {bad!r}")
502
+ except Refused as r:
503
+ assert r.status == 401
504
+ # The admin token as a bearer (the bench CLI): it, and only it.
505
+ assert admin_caller({"Authorization": f"Bearer {token}"}) == "admin"
506
+ assert admin_caller({"Authorization": f"bearer {token} "}) == "admin"
507
+ for bad in (f"Bearer {token}x", "Bearer " + "t" * 23, f"Basic {token}"):
508
+ try:
509
+ admin_caller({"Authorization": bad})
510
+ raise AssertionError(f"accepted {bad!r}")
511
+ except Refused as r:
512
+ assert r.status == 401
513
+ try: # a wrong bearer is refused, not saved by a cookie
514
+ admin_caller({"Authorization": f"Bearer {token}x", "Cookie": value})
515
+ raise AssertionError("a wrong bearer fell back to the cookie")
516
+ except Refused:
517
+ pass
518
+ os.environ["BENCHWARMERS_ADMIN_TOKEN"] = "u" * 24 # a new token signs everyone out
519
+ try:
520
+ admin_caller({"Cookie": value})
521
+ raise AssertionError("an old token's cookie still works")
522
+ except Refused:
523
+ pass
524
+
525
+ # Gateway keys: shown once, kept hashed, constant-time, revocable; never a cookie's job.
526
+ with tempfile.TemporaryDirectory() as tmp:
527
+ real, KEYS = KEYS, pathlib.Path(tmp) / "gateway-keys.json"
528
+ try:
529
+ key = make_key("chat-ui")
530
+ assert key.startswith("bwk_") and key not in KEYS.read_text(encoding="utf-8")
531
+ assert gateway_keys()["chat-ui"]["sha256"] == hashlib.sha256(key.encode()).hexdigest()
532
+ for bad in ("Bad Name", "x", "chat-ui"): # form, length, taken
533
+ try:
534
+ make_key(bad)
535
+ raise AssertionError(f"made a key named {bad!r}")
536
+ except ValueError:
537
+ pass
538
+ other = make_key("agent")
539
+ assert key_name(key) == "chat-ui" and key_name(other) == "agent"
540
+ assert key_name("") is None and key_name(key + "x") is None and key_name("bwk_é") is None
541
+ assert gateway_caller({"Authorization": f"Bearer {key}"}) == "chat-ui"
542
+ assert gateway_caller({"Authorization": f"bearer {key} "}) == "chat-ui"
543
+ entry = gateway_keys()["chat-ui"]
544
+ assert entry["hint"] == key[-4:] and entry["last_used"] and "rpm" not in entry, entry
545
+ # Each refusal says which it was, and how to authenticate (RFC 6750).
546
+ for headers, said in (({}, "needs a gateway key"), ({"Cookie": value}, "needs a gateway key"),
547
+ ({"Authorization": key}, "Bearer <key>"), ({"Authorization": "Bearer "}, "Bearer <key>"),
548
+ ({"Authorization": f"Basic {key}"}, "Bearer <key>"),
549
+ ({"Authorization": "Bearer sk-dummy"}, "some other key"),
550
+ ({"Authorization": f"Bearer {key}x"}, "unknown or revoked")):
551
+ try:
552
+ gateway_caller(headers)
553
+ raise AssertionError(f"gateway let in {headers}")
554
+ except Refused as r:
555
+ assert r.status == 401 and said in str(r), (headers, str(r))
556
+ assert r.headers["WWW-Authenticate"].startswith("Bearer"), r.headers
557
+ # Scope and cap, set when made: checked per collection, and per minute.
558
+ for name, kw in (("scoped", {"collections": ["Core"]}), ("scoped", {"rpm": 0}),
559
+ ("scoped", {"rpm": True}), ("scoped", {"rpm": "5"}), ("local", {})):
560
+ try:
561
+ make_key(name, **kw)
562
+ raise AssertionError(f"made a key with {kw}")
563
+ except ValueError:
564
+ pass
565
+ make_key("scoped", ["vision", "core", "core"], 2)
566
+ assert gateway_keys()["scoped"]["collections"] == ["core", "vision"]
567
+ assert may_call("scoped", "vision") and not may_call("scoped", "reviewers")
568
+ assert may_call("chat-ui", "reviewers") and may_call("local", "reviewers")
569
+ spend("scoped")
570
+ spend("scoped")
571
+ try:
572
+ spend("scoped")
573
+ raise AssertionError("a third request in a minute on a cap of 2")
574
+ except Refused as r:
575
+ assert r.status == 429 and 1 <= int(r.headers["Retry-After"]) <= 61, r.headers
576
+ for _ in range(5):
577
+ spend("chat-ui") # no cap
578
+ for bad in ("core", "", 0, False, [1]): # a bad edit fails closed
579
+ locked.update(KEYS, lambda d: d["scoped"].__setitem__("collections", bad))
580
+ assert scope("scoped") is None and not may_call("scoped", "core"), bad
581
+ try: # a key never opens the dashboard
582
+ admin_caller({"Authorization": f"Bearer {key}"})
583
+ raise AssertionError("a gateway key opened the dashboard")
584
+ except Refused:
585
+ pass
586
+ assert revoke_key("chat-ui") and not revoke_key("chat-ui") and key_name(key) is None
587
+ assert key_name(other) == "agent"
588
+ # No admin token: as before keys existed, whatever is sent.
589
+ del os.environ["BENCHWARMERS_ADMIN_TOKEN"]
590
+ assert gateway_caller({}) == "local" and gateway_caller({"Authorization": "Bearer x"}) == "local"
591
+ assert admin_caller({}) == "local"
592
+ # ...unless keys are required on this machine too; a misspelt value still asks.
593
+ for v in ("1", "required", "TRUE"):
594
+ os.environ["BENCHWARMERS_REQUIRE_KEYS"] = v
595
+ assert gateway_caller({"Authorization": f"Bearer {other}"}) == "agent"
596
+ try:
597
+ gateway_caller({"Authorization": "Bearer x"})
598
+ raise AssertionError(f"BENCHWARMERS_REQUIRE_KEYS={v} let a dummy bearer in")
599
+ except Refused as r:
600
+ assert r.status == 401
601
+ assert admin_caller({}) == "local" # the dashboard stays open
602
+ os.environ["BENCHWARMERS_REQUIRE_KEYS"] = "off"
603
+ assert not keys_required() and gateway_caller({}) == "local"
604
+ finally:
605
+ KEYS = real
606
+ assert gateway_route("/v1/models") and gateway_route("/bench/verdict")
607
+ assert not gateway_route("/api/endpoints") and not gateway_route("/v1") and not gateway_route("//v1/x")
608
+ finally:
609
+ for k, v in saved.items():
610
+ os.environ.pop(k, None)
611
+ if v is not None:
612
+ os.environ[k] = v
613
+ print("ok")
614
+
615
+
616
+ if __name__ == "__main__":
617
+ _selftest()