80d3434d61* Studio: require signed capability tokens for /p preview links The public /p preview routes added in #6486 run model load and chat generation as the admin user with no authentication. The only gate is the preview ref, a deterministic outputs-root path (run or run/checkpoint) that is guessable rather than secret. On a network-reachable Studio (--secure tunnel or -H 0.0.0.0), an unauthenticated caller who guesses a ref can consume GPU and probe a private fine-tuned checkpoint. Make the share link an unguessable, revocable capability: - Sign the canonical ref with a dedicated server-side secret (HMAC-SHA256, stored in app_secrets, independent of the JWT/login secret). - Require a valid token on every /p chat, models, and page request before resolving a checkpoint or loading a model; missing or invalid tokens get a generic 404 so the surface never confirms a ref exists. - Accept the token via ?k= (browser link and preview page) or Authorization: Bearer (OpenAI-compatible clients). - Rotate the secret to revoke every outstanding link (POST /api/settings/preview-links/rotate). - Clamp preview generation (max_tokens/max_completion_tokens <= 1024, n = 1) and set Referrer-Policy: no-referrer on the page so the token is not leaked via Referer. Training history hands the authenticated owner the signed token, and the copy-link button builds /p/{ref}?k={sig}. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: honor a lower caller token limit in the preview clamp Codex review: when only the legacy max_tokens was sent, the clamp left max_completion_tokens at the 1024 default, and _effective_max_tokens prefers max_completion_tokens, so a request like max_tokens=16 could still generate up to 1024 tokens. Derive one effective limit (max_completion_tokens wins, else the legacy max_tokens) and pin both fields to it so a caller's lower limit is kept. * Studio: add preview kill switch, rate limit, and revoke-links UI Follow-ups to the /p preview capability work: - Public-sharing kill switch: a persisted setting (default on) gates the public /p surface. When off, every preview request 404s even with a valid token, and the owner UI stops offering share links. GET/PUT /api/settings/preview-sharing; enforced in _verify_or_404. - Per-IP rate limit on the preview chat route: a coarse in-process sliding-window limiter (20 req/min/IP) returns 429 + Retry-After before the GPU lock is taken. Client IP honors X-Forwarded-For only when UNSLOTH_STUDIO_TRUST_FORWARDED is set, matching the login limiter's trust model. - Settings UI: a "Preview sharing" section with the public-sharing toggle and a "Revoke all preview links" button (confirm dialog) that rotates the secret. Tests cover the kill switch (404 when off), the 429 path, the sliding window, client-IP trust behavior, and the setting default. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: fix preview-fields sharing arg and refresh sigs after revoke Codex review: - P1: get_training_run_detail and update_training_run called _preview_fields with only output_dir after it gained a required sharing_on parameter, raising a 500 TypeError once get_run succeeded. Pass get_preview_sharing_enabled() at both sites; add a detail-endpoint regression test. - P2: after rotating the preview secret from settings, the history grid still held stale preview_sig values, so a freshly copied link would 404. Emit emitTrainingRunsChanged() after a successful revoke so the grid refetches freshly signed refs. * Studio: harden preview sharing controls (Codex review) - Fail closed: a read failure on the preview-sharing kill switch now returns False instead of defaulting to enabled, so an unavailable settings DB can't reopen the public surface. A missing key still defaults to enabled. - Per-IP rate limit behind the managed Cloudflare tunnel: client_ip now honors CF-Connecting-IP when the socket peer is loopback, so tunneled visitors are keyed by their real IP instead of collapsing onto the local cloudflared peer. - GET /p no longer mints key/share_url when sharing is disabled; it returns sharing_enabled=false so clients don't distribute links that 404. - Settings UI: toggling public sharing emits the training-runs-changed event so the history grid shows/hides Copy preview link without a manual refresh. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: harden preview rate limiter and IP keying (Opus review) From a two-agent review of the PR: - Rate limiter no longer evicts an active bucket when the table is full: a flood of distinct keys could otherwise cycle out a throttled bucket and reset its counter. Evict only aged-out buckets; if the table is full of live clients, fail closed (deny the new key) instead. - client_ip keys on the rightmost (proxy-appended) X-Forwarded-For hop when the trust env is set; the leftmost is client-spoofable. Documented the append/overwrite-proxy assumption. - _verify_or_404 checks the capability token before the kill-switch DB read, so unauthenticated /p spam can't be used as an unbounded settings-DB sink and the response is identical regardless of the sharing on/off state. Tests: nested run/checkpoint happy path + wrong-ref rejection, the eviction fail-closed behavior, and route-level coverage for the rotate / preview-sharing settings endpoints. --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
75 lines
3.3 KiB
Python
75 lines
3.3 KiB
Python
# SPDX-License-Identifier: AGPL-3.0-only
|
|
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
|
|
|
"""Unit + rotation coverage for `/p` preview capability tokens.
|
|
|
|
The token turns a guessable preview ref into an unguessable bearer capability:
|
|
it must round-trip for the ref it was signed for, reject tampering / wrong refs,
|
|
and stop verifying once the signing secret is rotated (link revocation).
|
|
"""
|
|
|
|
from pathlib import Path
|
|
import sys
|
|
import types as _types
|
|
|
|
|
|
_BACKEND_DIR = str(Path(__file__).resolve().parent.parent)
|
|
if _BACKEND_DIR not in sys.path:
|
|
sys.path.insert(0, _BACKEND_DIR)
|
|
|
|
# Mirror the other preview tests: avoid the heavy real `loggers` handlers.
|
|
_loggers_stub = _types.ModuleType("loggers")
|
|
_loggers_stub.get_logger = lambda name: __import__("logging").getLogger(name)
|
|
sys.modules.setdefault("loggers", _loggers_stub)
|
|
|
|
import auth.storage as storage
|
|
import utils.preview_token as preview_token
|
|
|
|
|
|
_S1 = b"secret-one-aaaaaaaaaaaaaaaaaaaaaaaa"
|
|
_S2 = b"secret-two-bbbbbbbbbbbbbbbbbbbbbbbb"
|
|
|
|
|
|
def test_sign_verify_roundtrip(monkeypatch):
|
|
monkeypatch.setattr(preview_token, "get_or_create_preview_link_secret", lambda: _S1)
|
|
token = preview_token.sign_preview_ref("run/checkpoint-1")
|
|
assert preview_token.verify_preview_ref("run/checkpoint-1", token)
|
|
# URL-safe, unpadded, and high-entropy (SHA-256 -> 43 base64url chars).
|
|
assert "=" not in token and "/" not in token and "+" not in token
|
|
assert len(token) >= 40
|
|
|
|
|
|
def test_missing_or_tampered_token_rejected(monkeypatch):
|
|
monkeypatch.setattr(preview_token, "get_or_create_preview_link_secret", lambda: _S1)
|
|
token = preview_token.sign_preview_ref("demorun")
|
|
assert not preview_token.verify_preview_ref("demorun", None)
|
|
assert not preview_token.verify_preview_ref("demorun", "")
|
|
flipped = token[:-1] + ("A" if token[-1] != "A" else "B")
|
|
assert not preview_token.verify_preview_ref("demorun", flipped)
|
|
# A token minted for one ref does not unlock another.
|
|
assert not preview_token.verify_preview_ref("otherrun", token)
|
|
# A non-ASCII token is invalid, not a crash (the route would 500 otherwise).
|
|
assert not preview_token.verify_preview_ref("demorun", "tøken-é")
|
|
|
|
|
|
def test_secret_change_invalidates_token(monkeypatch):
|
|
monkeypatch.setattr(preview_token, "get_or_create_preview_link_secret", lambda: _S1)
|
|
token = preview_token.sign_preview_ref("demorun")
|
|
monkeypatch.setattr(preview_token, "get_or_create_preview_link_secret", lambda: _S2)
|
|
assert not preview_token.verify_preview_ref("demorun", token)
|
|
|
|
|
|
def test_rotation_revokes_links(tmp_path, monkeypatch):
|
|
# Exercise the real storage helpers against a throwaway auth.db.
|
|
monkeypatch.setattr(storage, "DB_PATH", tmp_path / "auth.db")
|
|
monkeypatch.setattr(storage, "_preview_link_secret_cache", None)
|
|
|
|
token = preview_token.sign_preview_ref("demorun")
|
|
assert preview_token.verify_preview_ref("demorun", token)
|
|
# Secret persists across calls (the link keeps working until rotated).
|
|
assert preview_token.verify_preview_ref("demorun", token)
|
|
|
|
storage.rotate_preview_link_secret()
|
|
# Old shared link is revoked; a freshly minted one works.
|
|
assert not preview_token.verify_preview_ref("demorun", token)
|
|
assert preview_token.verify_preview_ref("demorun", preview_token.sign_preview_ref("demorun"))
|