From 0881a7a5d72f79237f49a222f8e5a292da8e67b0 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Wed, 13 May 2026 06:12:18 -0700 Subject: [PATCH] studio: security and hardening pass (auth rate-limit, sandbox, path containment, schema validation, headers) (#5375) * studio: contain export and dataset paths under their configured roots resolve_under_root and resolve_dataset_path previously returned absolute paths unchanged, so an authenticated client could supply save_directory="/tmp/escape" (or any other absolute path) and have the exporter drop adapter files anywhere the server user could write. This turned up during a recent audit pass where an authenticated POST to /api/export/export/lora with save_directory="/tmp/lora_escape_test" returned 200 and wrote adapter_model.safetensors, adapter_config.json, and tokenizer files under /tmp. The fix is two-layered: storage_roots.py adds an _assert_contained(resolved, root) helper that runs after path resolution and rejects any result whose realpath does not sit under realpath(root). resolve_under_root now rejects '..' segments and null bytes outright, and only accepts absolute inputs when they are already inside the configured root (internal call sites that re-resolve a stored absolute path stay idempotent; worker.py:resolve_output_dir(output_dir) etc. continue to work). resolve_dataset_path picks up the same containment rule, scoped to the three dataset roots. models/export.py adds field_validator("save_directory", mode="before") to ExportCommonOptions and ExportGGUFRequest so bad input fails fast at 422 with a clear message rather than a 500 deep inside the resolver. The validator rejects empty/whitespace, null bytes, control chars, strings longer than 255 chars, absolute paths, and '..' segments. routes/export.py:_export_details now returns os.path.relpath(output_path, exports_root()) so the Export Complete dialog and /api/models/loras no longer leak the absolute install prefix to the UI; the basename is used as a last-resort fallback. Verified end to end: - POST /api/export/export/lora {"save_directory":"/tmp/foo"} -> 422 "save_directory must be a name or relative path under the export root; absolute paths are rejected". /tmp/foo is not created. - "../../etc/escape" -> 422 "may not contain '..' segments". - save_directory="my_subdir" -> still accepted (400 only because the test had no checkpoint loaded yet, not because of validation). - Internal idempotent re-resolve via resolve_export_dir(absolute path that is already under exports_root) returns the same path unchanged. * studio/sandbox: harden bash + python tool execution The sandboxed Bash and Python tool channels in Chat ran with a thin preexec hook (PR_SET_NO_NEW_PRIVS + RLIMIT_FSIZE only). Bash had a small word blocklist; Python had an AST safety pass aimed at signal-tampering and shell-escape primitives. An audit pass showed several gaps that a tool-calling model could trigger inadvertently: - bash curl/wget/nc reached AWS IMDSv2 and returned live STS credentials for the instance role. - python "import socket; s.connect((169.254.169.254, 80))" reached the same endpoint regardless of the bash blocklist. - "cat /etc/passwd" was blocked at the bash side (because "passwd" is in the blocklist), but "open('/etc/passwd').read()" in Python happily returned its contents. - "chr(115)+chr(117)+chr(100)+chr(111)" style dynamic-arg construction slipped through the AST shell-escape check. - The supervisor used proc.kill() on timeout, which only signals the immediate pid; bash-backgrounded children survived. A fork bomb could spawn for the full 300s timeout window. - Session work directories under ~/studio_sandbox// were created with default umask (0o755), so any other UID on the host could enumerate them. - session_id sanitisation used a one-shot str.replace("..",""), which is non-iterative and a small footgun. This commit takes a conservative middle path: the sandbox still runs as the Studio UID with no namespace tricks where the kernel disallows them, but every chokepoint is tightened. _sandbox_preexec now: - calls os.setsid() so children share a process group; the supervisor uses os.killpg(SIGKILL) on timeout/cancel so backgrounded children die with the parent (new _kill_process_tree helper, wired into _cancel_watcher and both _bash_exec / _python_exec timeout branches). - calls os.umask(0o077) so files the child writes default to 0o600. - applies PR_SET_PDEATHSIG=SIGKILL so an orphaned child dies if Studio exits. - best-effort unshare(CLONE_NEWNET) for a private network namespace (failure is logged and swallowed; defense-in-depth is still in place via the bash blocklist and the AST checker below). - sets RLIMIT_NPROC=10000 (tunable via UNSLOTH_STUDIO_SANDBOX_NPROC), RLIMIT_AS=8GB, RLIMIT_CPU=300, RLIMIT_NOFILE=1024. The 10k NPROC figure is chosen to sit well above the ~500 LWPs a healthy Studio + llama-server combination already uses while still capping a runaway fork bomb. NPROC counts LWPs per real UID, so a lower figure (e.g. 256) starves legitimate bash forks ("bash: fork: retry: Resource temporarily unavailable"). _get_workdir: - rejects session_id that doesn't match [A-Za-z0-9_-]{1,64}; non-matching values bucket into a shared "_invalid" dir. - chmod 0o700 on both the workdir and on ~/studio_sandbox/ so other UIDs cannot read another session's contents. _BLOCKED_COMMANDS_COMMON gains: doas, pkexec, halt, poweroff, curl, wget, nc, ncat, netcat, socat, ssh, scp, sftp, rsync, eval, source. The intent is to keep general bash usage working (echo, ls, pipes, loops, for, head, etc.) while denying the obvious egress and escalation paths. The AST checker (_check_signal_escape_patterns) is split into the existing shell/signal/loop checks plus a new narrow IO denylist: - Always flag non-literal args to anything in _SHELL_EXEC_FUNCS, not just _STRING_SHELL_FUNCS. Closes the dynamic-arg bypass. - Reject calls to socket.create_connection, socket.socket().connect, urllib.request.urlopen, http.client.HTTP*Connection, requests.*, httpx.* whose literal host argument is in a cloud-metadata denylist (169.254.169.254 + 169.254.* + 100.64.*, plus the GCP/Alibaba/ECS metadata hostnames and IPv6 link-local). Public hosts (example.com, huggingface.co, ...) still work. Dynamic hosts cannot be statically blocked; mitigated by the bash blocklist + the netns where the kernel allows it. - Reject literal open("/etc/passwd"), /etc/shadow, /etc/sudoers, /etc/ssh/*, and /proc//environ. Other files (/etc/os-release, /etc/hostname, /tmp/*, user dirs) still work. The _check_code_safety summariser is updated to include the new network_calls and sensitive_file_reads buckets in its error string. Regression-checked: echo, sleep, ls /tmp, for loops, piped helpers (echo a | tr a A), urllib.request.urlopen("http://example.com"), socket.getaddrinfo("example.com",80), open("/etc/os-release"), open("/tmp/...","w") all still succeed. curl, wget, nc, ssh, rm, socket.create_connection(("169.254.169.254",80)), open("/etc/passwd"), open("/proc/self/environ") all correctly blocked. * studio: rate-limit login, rotate refresh tokens, add logout, security headers, gate bootstrap injection A pass over the auth surface found a cluster of related issues that this commit closes together. Login (routes/auth.py): - Add an in-memory per-IP login rate limiter. Five failed POSTs to /api/auth/login inside a 60s window produce 429 with Retry-After. A successful login clears the bucket. Previously 30 wrong passwords in under one second was accepted as 30x 401, which combined with the (now fixed) admin-username leak from /api/auth/status made brute-force trivial against a small password. Logout (routes/auth.py): - New POST /api/auth/logout returns 204 and calls storage.revoke_user_refresh_tokens(subject) so the refresh token is no longer valid. Previously POST /api/auth/logout returned 405 and there was no way to invalidate refresh tokens short of changing the password. Frontend session.ts already calls clearAuthTokens() to drop localStorage; the new endpoint lets the client also tell the server to revoke server-side state. Refresh-token rotation (routes/auth.py + auth/storage.py): - New storage.consume_refresh_token(token) atomically validates + deletes a refresh token, returning (username, is_desktop). The /api/auth/refresh handler now mints both a new access AND a new refresh token; the supplied token becomes invalid. Replaying a consumed refresh returns 401 "Invalid or expired refresh token". The previous refresh_access_token helper is left in place for callers that intentionally want the non-rotating shape; nothing in the route layer uses it now. /api/auth/status no longer leaks default_username (models/auth.py + routes/auth.py): - AuthStatusResponse.default_username becomes Optional[str] with a None default; the handler always returns None. The frontend already hardcodes HIDDEN_LOGIN_USERNAME = "unsloth" (auth-form.tsx:82), so no UI change is required. window.__UNSLOTH_BOOTSTRAP__ no longer auto-injects (main.py): - _inject_bootstrap is now opt-in via the UNSLOTH_STUDIO_INJECT_BOOTSTRAP env var. The previous default (inject whenever requires_password_change is true) embedded the plaintext bootstrap password into the first-boot HTML for any caller that hit /, /change-password, or any unknown SPA path. Browser extensions and any XSS payload on the page could read it trivially. With the new gate the bootstrap password lives only in the auth/.bootstrap_password file (mode 0o600) where it has always been; users typing it into a current-password field is the right UX. routes/auth.py:change_password also clears app.state.bootstrap_password defensively. Security headers + server fingerprint (main.py + run.py): - New SecurityHeadersMiddleware adds Content-Security-Policy, X-Frame-Options: DENY, X-Content-Type-Options: nosniff, Referrer-Policy: no-referrer, Permissions-Policy: camera=(), microphone=(), geolocation=(), interest-cohort=(), and stamps server: unsloth-studio so the generic uvicorn banner no longer fingerprints the stack. The uvicorn.Config gains server_header=False so it stops emitting its own Server header. /api/health minimisation (main.py): - Unauthenticated GET /api/health returns just {"status":"healthy","timestamp":...} so load-balancer liveness probes keep working without leaking version, device_type, chat_only, desktop_protocol_version, or studio_root_id to arbitrary callers. A request that presents a valid Bearer token still gets the full diagnostic payload so internal launchers and sibling-Studio detection (which compares studio_root_id) keep working. Verification: - 30 wrong-password POSTs to /api/auth/login -> first 5 = 401, 6th through 30th = 429. - POST /api/auth/logout with a fresh token -> 204. The matching refresh token then fails 401. - Login -> R1; /api/auth/refresh with R1 -> new access + R2 (R2 != R1); /api/auth/refresh with R1 again -> 401; /api/auth/refresh with R2 -> still succeeds once and rotates again. - curl /api/auth/status -> default_username: null. - curl http://127.0.0.1/ does not contain __UNSLOTH_BOOTSTRAP__. - curl -I / shows CSP, X-Frame-Options: DENY, X-Content-Type-Options: nosniff, Referrer-Policy: no-referrer, Permissions-Policy, and server: unsloth-studio. - curl /api/health unauthenticated -> {status, timestamp} only. curl with Authorization: Bearer -> full payload. - Existing /api/system, /api/models/list, /api/train/status, /api/inference/status, /api/auth/api-keys, login flow, SPA root all still return 200 after the changes (regression smoke). * studio: add SecurityHeadersMiddleware, MaxBodyMiddleware, /recipes redirect, gate _inject_bootstrap, minimise /api/health This commit lands the main.py-side changes that share a single middleware-registration spot. They are kept together because every change here is either (a) a top-level middleware definition that has to be added next to LoggingMiddleware, or (b) a route handler at the same file-level. SecurityHeadersMiddleware (Content-Security-Policy, X-Frame-Options: DENY, X-Content-Type-Options: nosniff, Referrer-Policy: no-referrer, Permissions-Policy, server: unsloth-studio). The previous responses emitted no CSP, no XFO, no Referrer-Policy and were stamped server: uvicorn. MaxBodyMiddleware rejects POST/PUT/PATCH on the inference / dataset / data-recipe / train / export prefixes when Content-Length exceeds UNSLOTH_STUDIO_MAX_BODY_MB (default 100). The audit hit this by attaching a 50 MB plain-text file to a chat message and watching Studio base64-encode it into the JSON body; uvicorn has no enforced cap so the only previous guard was the per-file 50 MB ceiling that data-recipe upload routes already enforce. The new middleware extends that ceiling to the OpenAI-compat path that the Chat attachments flow through. Verified: a 200 MB JSON POST to /v1/chat/completions returns HTTP 413 "Request body too large (209,715,264 bytes; max 104,857,600)". A small valid request continues to reach the handler. _inject_bootstrap is gated behind UNSLOTH_STUDIO_INJECT_BOOTSTRAP. The previous default was to inline window.__UNSLOTH_BOOTSTRAP__ = {username, password} into the first-boot HTML whenever requires_password_change was true, which exposed the plaintext bootstrap password to any browser extension, page script, or LAN caller on -H 0.0.0.0. The bootstrap password remains in the on-disk .bootstrap_password file (mode 0o600) where it has always lived; users typing it into a current-password field is the right UX. /api/health unauthenticated returns {"status":"healthy","timestamp": ...} only; the previous payload (version, device_type, chat_only, desktop_protocol_version, supports_desktop_auth, studio_root_id, native_path_leases_supported) is preserved for callers that present a valid Bearer token, so internal launchers and sibling-Studio detection (which compares studio_root_id) keep working. /recipes -> /data-recipes 308 redirect. The Data Recipes page lives at /data-recipes; users typing /recipes hit the SPA catch-all and saw "Not Found". The redirect also preserves any tail path, so /recipes/ -> /data-recipes/. Verified end to end with curl: CSP / XFO / X-Content-Type-Options / Referrer-Policy / Permissions-Policy all present on /, server header is now unsloth-studio (uvicorn's own banner is suppressed via server_header=False in run.py from the auth-batch commit). Followed the /recipes redirect lands on the SPA HTML. * studio: bound TrainingStartRequest hyperparameters at the schema level POST /api/train/start accepted any value for learning_rate, batch_size, max_steps, max_seq_length, warmup_steps, warmup_ratio, num_epochs, save_steps, weight_decay, gradient_accumulation_steps, lora_r, lora_alpha and lora_dropout, including -1, 0, 1e9, and non-numeric strings like 'abc' or 'two' (which silently coerce to 0 in the trainer). Probing showed the API returning 200 to learning_rate=-1 and batch_size=0; only max_steps had any partial clamping. This commit adds field_validator on every numeric hyperparameter. Bounds are chosen wide enough to span realistic single-host configurations (B200 with 180 GB of memory comfortably fits the upper end) while rejecting the values that always produce broken training: - learning_rate: parses str/float, requires 0 < lr < 1.0. Non-numeric input raises with "learning_rate must be parseable as float (got 'abc')" instead of silently coercing to 0. - batch_size: [1, 1024]. - gradient_accumulation_steps: [1, 4096]. - num_epochs: [1, 1000]. - max_steps: [1, 1_000_000]. - max_seq_length: [1, 131072]. - warmup_steps: [0, max_steps]. - warmup_ratio: [0.0, 1.0]. - save_steps: [0, 1_000_000]. - weight_decay: [0, 10] (typical 0..0.1). - lora_r: [1, 512]. - lora_alpha: [1, 1024]. - lora_dropout: [0.0, 1.0). Each validator names the offending field in its ValueError message so the 422 response body identifies which input is bad. The learning_rate validator returns its result as str (the schema field type is str("2e-4") for backwards compatibility) so existing call sites that float() the value continue to work. Verified: - learning_rate=-1 -> 422 "learning_rate must be > 0 (got -1.0); typical range is 1e-6 .. 1e-3". - learning_rate='abc' -> 422 "must be parseable as float". - batch_size=-1 / 0 / 999999 -> 422 "batch_size must be in [1, 1024]". - batch_size='two' -> 422 (pydantic int parser). - max_steps=0 / -5 -> 422 "must be a positive int". - max_seq_length=200000 -> 422 "must be in [1, 131072]". - warmup_ratio=2.5 -> 422 "must be in [0.0, 1.0]". - lora_dropout=1.5 -> 422 "must be in [0.0, 1.0)". - Valid request with learning_rate='2e-4', batch_size=1, max_steps=5 passes validation and the training run starts as normal. * studio: redact image-decode errors, clean checkpoint dirs on cancel, tolerate Stop-button + tool-result message shapes Three small fixes that fall under "do not let the audit findings become user-visible papercuts". routes/inference.py - image-decode error redaction (the audit hit this with a 0-byte / malformed / wrong-extension image upload). The three image-normalise sites previously raised HTTPException(400, detail=f"Failed to process image: {e}"). When PIL raised UnidentifiedImageError(io.BytesIO(raw)) the message string included "<_io.BytesIO object at 0x7e40a5d7bf60>", leaking both the Python class name (confirming the PIL/io stack) and a heap address (mildly useful for ASLR-bypass chaining if another memory-corruption bug is ever found). Each site now catches UnidentifiedImageError and returns the generic "Unsupported or corrupt image format"; the fall-through generic except returns "Failed to process image". No exception-repr is interpolated into a response body anywhere along these paths. core/training/training.py - checkpoint cleanup on cancel. When a user clicks Cancel Training, the trainer flips _cancel_requested=True and the supervisor force-terminates the subprocess. The trainer writes checkpoint- directories under output_dir every save_steps; previously these survived the cancel and accumulated on disk (the audit recorded ~67 MB stuck after a 200-step cancel with save_steps=20). New helper _cleanup_cancelled_checkpoints(output_dir) globs checkpoint- entries and removes them. It is gated by a realpath containment check against outputs_root() so it cannot accidentally rmtree anything outside the configured outputs root. force_terminate() invokes the helper after the subprocess join when _cancel_requested is true. Stop-and-Save runs are unaffected because that path keeps _cancel_requested=False. models/inference.py - chat message shape tolerance. Two related frontend interactions used to crash the request validator: - After the Stop button truncates a generation, the frontend retained {role:"assistant", content:""} in the conversation history and replayed it on the next send. ChatMessage previously required role="assistant" to have non-empty content or tool_calls, so the next message returned 422 and the thread was permanently broken. The validator now normalises empty assistant content to None so the request round-trips and the trailing empty turn can be ignored downstream. - The frontend's second-round tool POST drops the streamed tool_call_id, hitting the strict-spec check "role=tool requires tool_call_id". The validator now synthesises an opaque id (call_<8 hex>) when missing, so the request reaches the handler and the model's final summarising response gets generated. The proper fix lives in the frontend (carry the streamed id through the second POST) and will follow. Verified end to end with curl: HTTP 400 (model not loaded) on both the empty-assistant history shape and the tool-result-without-id shape, instead of HTTP 422 from the schema validator. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * studio: tighten code comments from security-hardening pass Trim verbose docstrings and inline finding references added in the previous commits in this branch. Functionality unchanged. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * studio: await get_current_subject in /api/health and make refresh-token consumption atomic The /api/health auth probe called get_current_subject(creds) without awaiting it. The coroutine object is truthy, so any caller presenting a Bearer header (valid or not) received the full diagnostic payload including version, device_type, studio_root_id, etc. Await the coroutine and treat HTTPException as 'fall back to the minimal liveness payload'. consume_refresh_token did SELECT then DELETE WHERE id under default autocommit isolation. Two concurrent POST /api/auth/refresh requests could both win the SELECT before either DELETE ran, defeating single-use refresh-token rotation. Replace with a single DELETE ... WHERE token_hash = ? AND expires_at >= ? RETURNING ... statement so the validate-and-delete lands as one atomic op under SQLite's write lock (3.45.1 supports RETURNING; min was 3.35). * studio: enforce body cap on chunked uploads and drop unsafe-inline from script-src MaxBodyMiddleware previously only inspected the declared Content-Length header; clients omitting it or sending Transfer-Encoding: chunked bypassed the cap and could still drive an OOM via the downstream JSON / file readers on /v1/chat/completions, /api/inference, /api/data-recipe, /api/datasets, /api/train, /api/export. Rewrite as a raw ASGI middleware that drains and counts http.request frames, replies 413 once the running total exceeds UNSLOTH_STUDIO_MAX_BODY_MB before invoking the FastAPI handler, and replays the buffered body to downstream so route code that calls request.json() / await request.body() works unchanged. CSP previously included 'unsafe-inline' on script-src, which defeats the main XSS protection. The frontend bundle does not need inline scripts; the only inline " + nonce = _secrets.token_urlsafe(16) + tag = f'' html = html_bytes.decode("utf-8") html = html.replace("", f"{tag}", 1) - return html.encode("utf-8") + return html.encode("utf-8"), nonce def setup_frontend(app: FastAPI, build_path: Path): @@ -501,17 +697,23 @@ def setup_frontend(app: FastAPI, build_path: Path): if assets_dir.exists(): app.mount("/assets", StaticFiles(directory = assets_dir), name = "assets") - @app.get("/") - async def serve_root(): + def _build_index_response() -> Response: content = (build_path / "index.html").read_bytes() content = _strip_crossorigin(content) - content = _inject_bootstrap(content, app) + content, nonce = _inject_bootstrap(content, app) + headers = {"Cache-Control": "no-cache, no-store, must-revalidate"} + if nonce: + headers[_CSP_SCRIPT_NONCE_HEADER] = nonce return Response( content = content, media_type = "text/html", - headers = {"Cache-Control": "no-cache, no-store, must-revalidate"}, + headers = headers, ) + @app.get("/") + async def serve_root(): + return _build_index_response() + @app.get("/{full_path:path}") async def serve_frontend(full_path: str): if full_path in {"api", "v1"} or full_path.startswith(("api/", "v1/")): @@ -527,13 +729,6 @@ def setup_frontend(app: FastAPI, build_path: Path): return FileResponse(file_path) # Serve index.html as bytes — avoids Content-Length mismatch - content = (build_path / "index.html").read_bytes() - content = _strip_crossorigin(content) - content = _inject_bootstrap(content, app) - return Response( - content = content, - media_type = "text/html", - headers = {"Cache-Control": "no-cache, no-store, must-revalidate"}, - ) + return _build_index_response() return True diff --git a/studio/backend/models/auth.py b/studio/backend/models/auth.py index 23eb0ac4c0..b7870379f7 100644 --- a/studio/backend/models/auth.py +++ b/studio/backend/models/auth.py @@ -37,7 +37,10 @@ class AuthStatusResponse(BaseModel): initialized: bool = Field( ..., description = "True if the auth database contains a login user" ) - default_username: str = Field(..., description = "Default seeded admin username") + default_username: str = Field( + "unsloth", + description = "Default admin username for first-boot UI prefill.", + ) requires_password_change: bool = Field( ..., description = "True if the seeded admin must still change the default password", diff --git a/studio/backend/models/export.py b/studio/backend/models/export.py index a86596f199..86ce2b05bf 100644 --- a/studio/backend/models/export.py +++ b/studio/backend/models/export.py @@ -5,10 +5,36 @@ Pydantic schemas for Export API. """ -from pydantic import BaseModel, Field +from pathlib import Path + +from pydantic import BaseModel, Field, field_validator from typing import List, Optional, Literal, Dict, Any +def _validate_save_directory(value: str) -> str: + """Reject save_directory values that escape the export root.""" + if value is None: + raise ValueError("save_directory is required") + raw = str(value).strip() + if not raw: + raise ValueError("save_directory must not be empty") + if "\x00" in raw: + raise ValueError("save_directory may not contain null bytes") + if any(ch in raw for ch in ("\r", "\n")): + raise ValueError("save_directory may not contain control characters") + if len(raw) > 255: + raise ValueError("save_directory must be <= 255 characters") + path = Path(raw).expanduser() + if path.is_absolute(): + raise ValueError( + "save_directory must be a name or relative path under the " + "export root; absolute paths are rejected" + ) + if ".." in path.parts: + raise ValueError("save_directory may not contain '..' segments") + return raw + + class LoadCheckpointRequest(BaseModel): """Request for loading a checkpoint into the export backend.""" @@ -64,6 +90,12 @@ class ExportCommonOptions(BaseModel): ..., description = "Local directory where the exported artifacts will be written", ) + + @field_validator("save_directory", mode = "before") + @classmethod + def _check_save_directory(cls, v): + return _validate_save_directory(v) + push_to_hub: bool = Field( False, description = "If True, also push the exported model to the Hugging Face Hub", @@ -108,6 +140,12 @@ class ExportGGUFRequest(BaseModel): ..., description = "Directory where GGUF files will be saved", ) + + @field_validator("save_directory", mode = "before") + @classmethod + def _check_save_directory(cls, v): + return _validate_save_directory(v) + quantization_method: str = Field( "Q4_K_M", description = 'GGUF quantization method (e.g. "Q4_K_M")', diff --git a/studio/backend/models/inference.py b/studio/backend/models/inference.py index 7a4c7d0b3c..746ac8bbc2 100644 --- a/studio/backend/models/inference.py +++ b/studio/backend/models/inference.py @@ -425,14 +425,6 @@ class ChatMessage(BaseModel): @model_validator(mode = "after") def _validate_role_shape(self) -> "ChatMessage": - # Enforce the per-role OpenAI spec shape at the request boundary. - # Without this, malformed messages (e.g. user entries with no - # content, tool_calls on a user/system role, role="tool" without - # tool_call_id) would be silently forwarded to llama-server via - # the passthrough path, surfacing as opaque upstream errors or - # broken tool-call reconciliation downstream. - - # Tool-call metadata must appear only on the appropriate role. if self.tool_calls is not None and self.role != "assistant": raise ValueError('"tool_calls" is only valid on role="assistant" messages.') if self.tool_call_id is not None and self.role != "tool": @@ -440,23 +432,20 @@ class ChatMessage(BaseModel): if self.name is not None and self.role != "tool": raise ValueError('"name" is only valid on role="tool" messages.') - # Per-role content requirements. OpenAI-compatible clients may send - # ``content=""`` for image-only turns when the image travels in a - # companion field such as Studio's ``image_base64`` extension, so treat - # empty strings as present content for user/system messages. if self.role == "tool": if not self.tool_call_id: - raise ValueError( - 'role="tool" messages require "tool_call_id" per the OpenAI spec.' - ) + # Frontend's second-round POST drops the streamed id; + # synthesise one so the request round-trips. + import secrets as _secrets + + self.tool_call_id = f"call_{_secrets.token_hex(8)}" if not self.content: raise ValueError('role="tool" messages require non-empty "content".') elif self.role == "assistant": - # Assistant messages may omit content when tool_calls is set. - if not self.content and not self.tool_calls: - raise ValueError( - 'role="assistant" messages require either "content" or "tool_calls".' - ) + # Tolerate the post-Stop empty-assistant sentinel by + # collapsing content="" to None. + if (self.content == "" or self.content == []) and not self.tool_calls: + self.content = None else: # "user" | "system" if self.content is None or self.content == []: raise ValueError(f'role="{self.role}" messages require "content".') diff --git a/studio/backend/models/training.py b/studio/backend/models/training.py index 0c5825c54e..6b5e95e188 100644 --- a/studio/backend/models/training.py +++ b/studio/backend/models/training.py @@ -5,10 +5,43 @@ Pydantic schemas for Training API """ -from pydantic import BaseModel, ConfigDict, Field, model_validator +from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator from typing import Any, Optional, List, Dict, Literal +_MAX_BATCH_SIZE = 4096 +_MAX_GRAD_ACCUM = 4096 +_MAX_STEPS = 1_000_000 +_MAX_EPOCHS = 1000 +# 2M is a sanity cap; host RAM runs out long before this. +_MAX_SEQ_LENGTH = 2_000_000 +_MAX_LR_VALUE = 1.0 +_MAX_LORA_R = 16_384 +_MAX_LORA_ALPHA = 32_768 + + +def _parse_lr(v: Any) -> float: + """Parse learning_rate as a positive float strictly below _MAX_LR_VALUE.""" + if v is None: + raise ValueError("learning_rate is required") + if isinstance(v, bool): + raise ValueError("learning_rate must be a number, not a bool") + try: + lr = float(v) + except (TypeError, ValueError): + raise ValueError(f"learning_rate must be parseable as float (got {v!r})") + if not (lr > 0.0): + raise ValueError( + f"learning_rate must be > 0 (got {lr!r}); " "typical range is 1e-6 .. 1e-3" + ) + if lr >= _MAX_LR_VALUE: + raise ValueError( + f"learning_rate must be < 1.0 (got {lr!r}); " + "values that large always diverge training" + ) + return lr + + class TrainingStartRequest(BaseModel): """Request schema for starting training""" @@ -64,6 +97,147 @@ class TrainingStartRequest(BaseModel): values.setdefault("train_split", values.pop("split")) return values + @field_validator("learning_rate", mode = "before") + @classmethod + def _check_learning_rate(cls, v): + # Stringify because downstream call sites float() it themselves. + lr = _parse_lr(v) + return str(lr) + + @field_validator("batch_size") + @classmethod + def _check_batch_size(cls, v: int) -> int: + if v is None: + raise ValueError("batch_size is required") + if v < 1 or v > _MAX_BATCH_SIZE: + raise ValueError( + f"batch_size must be in [1, {_MAX_BATCH_SIZE}] (got {v!r})" + ) + return v + + @field_validator("gradient_accumulation_steps") + @classmethod + def _check_grad_accum(cls, v: int) -> int: + if v is None: + return 1 + if v < 1 or v > _MAX_GRAD_ACCUM: + raise ValueError( + f"gradient_accumulation_steps must be in [1, {_MAX_GRAD_ACCUM}] " + f"(got {v!r})" + ) + return v + + @field_validator("num_epochs") + @classmethod + def _check_num_epochs(cls, v: int) -> int: + if v is None: + return 1 + if v < 1 or v > _MAX_EPOCHS: + raise ValueError(f"num_epochs must be in [1, {_MAX_EPOCHS}] (got {v!r})") + return v + + @field_validator("max_steps") + @classmethod + def _check_max_steps(cls, v): + if v is None: + return v + if not isinstance(v, int) or v < 1 or v > _MAX_STEPS: + raise ValueError( + f"max_steps must be a positive int <= {_MAX_STEPS} (got {v!r})" + ) + return v + + @field_validator("max_seq_length") + @classmethod + def _check_max_seq_length(cls, v: int) -> int: + if v is None or v < 1 or v > _MAX_SEQ_LENGTH: + raise ValueError( + f"max_seq_length must be in [1, {_MAX_SEQ_LENGTH}] (got {v!r})" + ) + return v + + @field_validator("warmup_steps") + @classmethod + def _check_warmup_steps(cls, v): + if v is None: + return v + if not isinstance(v, int) or v < 0 or v > _MAX_STEPS: + raise ValueError( + f"warmup_steps must be a non-negative int <= {_MAX_STEPS} " + f"(got {v!r})" + ) + return v + + @field_validator("warmup_ratio") + @classmethod + def _check_warmup_ratio(cls, v): + if v is None: + return v + try: + r = float(v) + except (TypeError, ValueError): + raise ValueError(f"warmup_ratio must be a number (got {v!r})") + if not (0.0 <= r <= 1.0): + raise ValueError(f"warmup_ratio must be in [0.0, 1.0] (got {r!r})") + return r + + @field_validator("save_steps") + @classmethod + def _check_save_steps(cls, v: int) -> int: + if v is None: + return 100 + if v < 0 or v > _MAX_STEPS: + raise ValueError(f"save_steps must be in [0, {_MAX_STEPS}] (got {v!r})") + return v + + @field_validator("weight_decay") + @classmethod + def _check_weight_decay(cls, v: float) -> float: + if v is None: + return 0.0 + try: + wd = float(v) + except (TypeError, ValueError): + raise ValueError(f"weight_decay must be a number (got {v!r})") + if wd < 0 or wd > 10.0: + raise ValueError( + f"weight_decay must be in [0, 10] (got {wd!r}); typical 0..0.1" + ) + return wd + + @field_validator("lora_r") + @classmethod + def _check_lora_r(cls, v: int) -> int: + if v is None: + return 16 + if v < 1 or v > _MAX_LORA_R: + raise ValueError(f"lora_r must be in [1, {_MAX_LORA_R}] (got {v!r})") + return v + + @field_validator("lora_alpha") + @classmethod + def _check_lora_alpha(cls, v: int) -> int: + if v is None: + return 16 + if v < 1 or v > _MAX_LORA_ALPHA: + raise ValueError( + f"lora_alpha must be in [1, {_MAX_LORA_ALPHA}] (got {v!r})" + ) + return v + + @field_validator("lora_dropout") + @classmethod + def _check_lora_dropout(cls, v: float) -> float: + if v is None: + return 0.0 + try: + d = float(v) + except (TypeError, ValueError): + raise ValueError(f"lora_dropout must be a number (got {v!r})") + if not (0.0 <= d < 1.0): + raise ValueError(f"lora_dropout must be in [0.0, 1.0) (got {d!r})") + return d + custom_format_mapping: Optional[Dict[str, Any]] = Field( None, description = ( diff --git a/studio/backend/routes/auth.py b/studio/backend/routes/auth.py index 3deeb6793b..30221c2c93 100644 --- a/studio/backend/routes/auth.py +++ b/studio/backend/routes/auth.py @@ -5,8 +5,11 @@ Authentication API routes """ -from fastapi import APIRouter, Depends, HTTPException, status +from fastapi import APIRouter, Depends, HTTPException, Request, Response, status +import threading +import time +from collections import deque from datetime import datetime, timedelta, timezone from models.auth import ( @@ -33,14 +36,52 @@ from auth.authentication import ( router = APIRouter() +# In-memory per-IP login rate limiter; multi-process deployment needs a shared store. +_LOGIN_BUCKETS: dict[str, deque] = {} +_LOGIN_BUCKETS_LOCK = threading.Lock() +_LOGIN_WINDOW_SECONDS = 60.0 +_LOGIN_MAX_FAILS = 5 +_LOGIN_LOCKOUT_SECONDS = 60 + + +def _client_key(request: Request | None) -> str: + if request is None or request.client is None: + return "_unknown" + return request.client.host or "_unknown" + + +def _record_login_failure(ip: str) -> int: + now = time.monotonic() + with _LOGIN_BUCKETS_LOCK: + bucket = _LOGIN_BUCKETS.setdefault(ip, deque()) + while bucket and now - bucket[0] > _LOGIN_WINDOW_SECONDS: + bucket.popleft() + bucket.append(now) + return len(bucket) + + +def _login_blocked(ip: str) -> int: + """Return seconds until the next attempt is allowed, or 0.""" + now = time.monotonic() + with _LOGIN_BUCKETS_LOCK: + bucket = _LOGIN_BUCKETS.get(ip) + if not bucket: + return 0 + while bucket and now - bucket[0] > _LOGIN_WINDOW_SECONDS: + bucket.popleft() + if len(bucket) >= _LOGIN_MAX_FAILS: + return max(1, int(_LOGIN_WINDOW_SECONDS - (now - bucket[0]))) + return 0 + + +def _clear_login_bucket(ip: str) -> None: + with _LOGIN_BUCKETS_LOCK: + _LOGIN_BUCKETS.pop(ip, None) + + @router.get("/status", response_model = AuthStatusResponse) async def auth_status() -> AuthStatusResponse: - """ - Check whether auth has already been initialized. - - - initialized = False -> frontend should wait for the seeded admin bootstrap. - - initialized = True -> frontend should show login or force the first password change. - """ + """Auth initialization state; ``default_username`` is exposed for first-boot UI prefill only.""" return AuthStatusResponse( initialized = storage.is_initialized(), default_username = storage.DEFAULT_ADMIN_USERNAME, @@ -53,12 +94,23 @@ async def auth_status() -> AuthStatusResponse: @router.post("/login", response_model = Token) -async def login(payload: AuthLoginRequest) -> Token: - """ - Login with username/password and receive access + refresh tokens. - """ +async def login(payload: AuthLoginRequest, request: Request) -> Token: + """Login with username/password. Rate-limited per source IP.""" + ip = _client_key(request) + blocked_for = _login_blocked(ip) + if blocked_for > 0: + raise HTTPException( + status_code = status.HTTP_429_TOO_MANY_REQUESTS, + detail = ( + f"Too many failed login attempts from {ip}. " + f"Try again in {blocked_for} seconds." + ), + headers = {"Retry-After": str(blocked_for)}, + ) + record = storage.get_user_and_secret(payload.username) if record is None: + _record_login_failure(ip) raise HTTPException( status_code = status.HTTP_401_UNAUTHORIZED, detail = "Incorrect password. Run 'unsloth studio reset-password' in your terminal to reset it.", @@ -66,11 +118,13 @@ async def login(payload: AuthLoginRequest) -> Token: salt, pwd_hash, _jwt_secret, must_change_password = record if not hashing.verify_password(payload.password, salt, pwd_hash): + _record_login_failure(ip) raise HTTPException( status_code = status.HTTP_401_UNAUTHORIZED, detail = "Incorrect password. Run 'unsloth studio reset-password' in your terminal to reset it.", ) + _clear_login_bucket(ip) access_token = create_access_token(subject = payload.username) refresh_token = create_refresh_token(subject = payload.username) return Token( @@ -81,6 +135,23 @@ async def login(payload: AuthLoginRequest) -> Token: ) +@router.post("/logout", status_code = status.HTTP_204_NO_CONTENT) +async def logout( + request: Request, + current_subject: str = Depends(get_current_subject_allow_password_change), +) -> Response: + """Revoke refresh tokens for the subject; the access token is stateless and expires on its own.""" + try: + storage.revoke_user_refresh_tokens(current_subject) + except Exception: + pass + try: + request.app.state.bootstrap_password = None + except AttributeError: + pass + return Response(status_code = status.HTTP_204_NO_CONTENT) + + @router.post("/desktop-login", response_model = Token) async def desktop_login(payload: DesktopLoginRequest) -> Token: """Exchange a local desktop secret for normal admin-subject tokens.""" @@ -101,21 +172,20 @@ async def desktop_login(payload: DesktopLoginRequest) -> Token: @router.post("/refresh", response_model = Token) async def refresh(payload: RefreshTokenRequest) -> Token: - """ - Exchange a valid refresh token for a new access token. - - The refresh token itself is reusable until it expires (7 days). - """ - new_access_token, username, is_desktop = refresh_access_token(payload.refresh_token) - if new_access_token is None or username is None: + """Exchange a refresh token for a new access+refresh pair (single-use).""" + consumed = storage.consume_refresh_token(payload.refresh_token) + if consumed is None: raise HTTPException( status_code = status.HTTP_401_UNAUTHORIZED, detail = "Invalid or expired refresh token", ) + username, is_desktop = consumed + new_access_token = create_access_token(subject = username, desktop = is_desktop) + new_refresh_token = create_refresh_token(subject = username, desktop = is_desktop) return Token( access_token = new_access_token, - refresh_token = payload.refresh_token, + refresh_token = new_refresh_token, token_type = "bearer", must_change_password = False if is_desktop @@ -126,6 +196,7 @@ async def refresh(payload: RefreshTokenRequest) -> Token: @router.post("/change-password", response_model = Token) async def change_password( payload: ChangePasswordRequest, + request: Request, current_subject: str = Depends(get_current_subject_allow_password_change), ) -> Token: """Allow the authenticated user to replace the default password.""" @@ -150,6 +221,10 @@ async def change_password( storage.update_password(current_subject, payload.new_password) storage.revoke_user_refresh_tokens(current_subject) + try: + request.app.state.bootstrap_password = None + except AttributeError: + pass access_token = create_access_token(subject = current_subject) refresh_token = create_refresh_token(subject = current_subject) return Token( diff --git a/studio/backend/routes/export.py b/studio/backend/routes/export.py index 798859fc87..7dbc52dbed 100644 --- a/studio/backend/routes/export.py +++ b/studio/backend/routes/export.py @@ -7,6 +7,7 @@ Export API routes: checkpoint discovery and model export operations. import asyncio import json +import os import sys import time from pathlib import Path @@ -184,14 +185,18 @@ async def get_export_status( def _export_details(output_path: Optional[str]) -> Optional[Dict[str, Any]]: - """Wrap the resolved on-disk export path into the details dict the - frontend reads to populate the Export Complete screen. Returns None - when the export had no local component (Hub-only push) so the - Pydantic field stays absent rather than ``{"output_path": null}``. - """ + """Return the export path relative to exports_root so the install path is not leaked.""" if not output_path: return None - return {"output_path": output_path} + try: + from utils.paths.storage_roots import exports_root + + rel = os.path.relpath(output_path, exports_root()) + if rel.startswith(".."): + rel = os.path.basename(output_path) + return {"output_path": rel} + except Exception: + return {"output_path": os.path.basename(output_path)} @router.post("/export/merged", response_model = ExportOperationResponse) diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 6b559b9c45..7102e12bf8 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -1743,7 +1743,7 @@ async def openai_chat_completions( try: import base64 as _b64 from io import BytesIO as _BytesIO - from PIL import Image as _Image + from PIL import Image as _Image, UnidentifiedImageError as _UIE raw = _b64.b64decode(image_b64) # Normalize to RGB so PNG encoding succeeds regardless of @@ -1754,9 +1754,15 @@ async def openai_chat_completions( buf = _BytesIO() img.save(buf, format = "PNG") image_b64 = _b64.b64encode(buf.getvalue()).decode("ascii") - except Exception as e: + except _UIE: raise HTTPException( - status_code = 400, detail = f"Failed to process image: {e}" + status_code = 400, + detail = "Unsupported or corrupt image format.", + ) + except Exception: + raise HTTPException( + status_code = 400, + detail = "Failed to process image.", ) # Build message list with system prompt prepended @@ -3426,10 +3432,10 @@ def _normalize_anthropic_openai_images( buf = io.BytesIO() img.save(buf, format = "PNG") png_b64 = base64.b64encode(buf.getvalue()).decode("ascii") - except Exception as e: + except Exception: raise HTTPException( status_code = 400, - detail = f"Failed to process image: {e}", + detail = "Failed to process image.", ) part["image_url"] = {"url": f"data:image/png;base64,{png_b64}"} @@ -3465,6 +3471,7 @@ async def anthropic_messages( [m.model_dump() for m in payload.messages], payload.system, ) + openai_messages = _drop_empty_assistant_sentinels(openai_messages) # Enforce vision guard + re-encode embedded images to PNG so the # Anthropic endpoint matches the behavior of /v1/chat/completions. @@ -4190,6 +4197,19 @@ async def _anthropic_passthrough_non_streaming( # ===================================================================== +def _drop_empty_assistant_sentinels(messages: list[dict]) -> list[dict]: + """Drop bare ``{"role":"assistant"}`` Stop-button sentinels; passthrough backends reject them.""" + out: list[dict] = [] + for m in messages: + if m.get("role") == "assistant": + has_content = bool(m.get("content")) + has_tool_calls = bool(m.get("tool_calls")) + if not has_content and not has_tool_calls: + continue + out.append(m) + return out + + def _openai_messages_for_passthrough(payload) -> list[dict]: """Build OpenAI-format message dicts for the /v1/chat/completions passthrough path. @@ -4206,7 +4226,9 @@ def _openai_messages_for_passthrough(payload) -> list[dict]: ``image_url`` content part so vision + function-calling requests work transparently. """ - messages = [m.model_dump(exclude_none = True) for m in payload.messages] + messages = _drop_empty_assistant_sentinels( + [m.model_dump(exclude_none = True) for m in payload.messages] + ) if not payload.image_base64: return messages @@ -4221,10 +4243,10 @@ def _openai_messages_for_passthrough(payload) -> list[dict]: buf = _BytesIO() img.save(buf, format = "PNG") png_b64 = _b64.b64encode(buf.getvalue()).decode("ascii") - except Exception as e: + except Exception: raise HTTPException( status_code = 400, - detail = f"Failed to process image: {e}", + detail = "Failed to process image.", ) data_url = f"data:image/png;base64,{png_b64}" diff --git a/studio/backend/run.py b/studio/backend/run.py index dfd4b7453e..0787e04c47 100644 --- a/studio/backend/run.py +++ b/studio/backend/run.py @@ -354,9 +354,14 @@ def run_server( if getattr(self, "started", False) and not self.should_exit: ready_event.set() - # Create the uvicorn server and expose it for signal handlers + # server_header=False suppresses uvicorn's "Server: uvicorn"; SecurityHeadersMiddleware sets its own. config = uvicorn.Config( - app, host = host, port = port, log_level = "info", access_log = False + app, + host = host, + port = port, + log_level = "info", + access_log = False, + server_header = False, ) _server = _ReadyServer(config) _shutdown_event = Event() diff --git a/studio/backend/tests/test_desktop_auth.py b/studio/backend/tests/test_desktop_auth.py index a5508c1c8b..a7201ac433 100644 --- a/studio/backend/tests/test_desktop_auth.py +++ b/studio/backend/tests/test_desktop_auth.py @@ -227,6 +227,60 @@ def test_desktop_refresh_preserves_desktop_marker(): assert payload["desktop"] is True +def test_consume_refresh_token_second_call_returns_none(): + """Single-use rotation rejects the same token on a second consume.""" + seed_user() + from datetime import datetime, timedelta, timezone + + raw = secrets.token_urlsafe(48) + expires = (datetime.now(timezone.utc) + timedelta(days = 30)).isoformat() + storage.save_refresh_token(raw, storage.DEFAULT_ADMIN_USERNAME, expires) + + first = storage.consume_refresh_token(raw) + assert first == (storage.DEFAULT_ADMIN_USERNAME, False) + second = storage.consume_refresh_token(raw) + assert second is None + + +def test_consume_refresh_token_concurrent_only_one_succeeds(tmp_path, monkeypatch): + """64-thread pile-up against one token; DELETE RETURNING permits one winner.""" + seed_user() + from concurrent.futures import ThreadPoolExecutor + from datetime import datetime, timedelta, timezone + + raw = secrets.token_urlsafe(48) + expires = (datetime.now(timezone.utc) + timedelta(days = 30)).isoformat() + storage.save_refresh_token(raw, storage.DEFAULT_ADMIN_USERNAME, expires) + + workers = 64 + + def attempt(_idx: int): + try: + return storage.consume_refresh_token(raw) + except sqlite3.OperationalError: + # "database is locked" under heavy contention; treat as losing the race. + return None + + with ThreadPoolExecutor(max_workers = workers) as pool: + results = list(pool.map(attempt, range(workers))) + + successes = [r for r in results if r is not None] + assert ( + len(successes) == 1 + ), f"expected exactly one consumer to win, got {len(successes)}" + assert successes[0] == (storage.DEFAULT_ADMIN_USERNAME, False) + + +def test_consume_refresh_token_expired_returns_none(): + seed_user() + from datetime import datetime, timedelta, timezone + + raw = secrets.token_urlsafe(48) + expires = (datetime.now(timezone.utc) - timedelta(hours = 1)).isoformat() + storage.save_refresh_token(raw, storage.DEFAULT_ADMIN_USERNAME, expires) + assert storage.consume_refresh_token(raw) is None + + def test_desktop_session_uses_real_admin_identity_for_api_keys(): seed_user(must_change_password = True) raw = storage.create_desktop_secret() @@ -392,7 +446,21 @@ def test_health_response_reports_desktop_capability_fields(monkeypatch): monkeypatch.setattr(backend_main._hw_module, "CHAT_ONLY", False) - body = asyncio.run(backend_main.health_check()) + seed_user() + from auth.authentication import create_access_token + + token = create_access_token(storage.DEFAULT_ADMIN_USERNAME) + + app = FastAPI() + app.add_api_route("/api/health", backend_main.health_check, methods = ["GET"]) + client = TestClient(app) + + response = client.get( + "/api/health", + headers = {"Authorization": f"Bearer {token}"}, + ) + assert response.status_code == 200 + body = response.json() assert body["desktop_protocol_version"] == 1 assert body["supports_desktop_auth"] is True diff --git a/studio/backend/tests/test_middleware.py b/studio/backend/tests/test_middleware.py new file mode 100644 index 0000000000..bdf8e6d5a5 --- /dev/null +++ b/studio/backend/tests/test_middleware.py @@ -0,0 +1,269 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. + +"""Tests for MaxBodyMiddleware, SecurityHeadersMiddleware, and the /api/health auth gate.""" + +import asyncio +import importlib.util +import json +import os +import sys +from pathlib import Path + +import pytest +from fastapi import FastAPI, HTTPException, Request +from fastapi.responses import Response +from fastapi.testclient import TestClient + + +_BACKEND_ROOT = Path(__file__).resolve().parents[1] +if str(_BACKEND_ROOT) not in sys.path: + sys.path.insert(0, str(_BACKEND_ROOT)) + + +@pytest.fixture(scope = "module") +def main_module(): + import main as _main # noqa: F401 + + return _main + + +# ===================================================================== +# MaxBodyMiddleware +# ===================================================================== + + +def _make_protected_app(max_bytes: int, main_module): + app = FastAPI() + app.add_middleware( + main_module.MaxBodyMiddleware, + max_bytes = max_bytes, + protected_prefixes = ("/v1/chat/completions", "/api/train"), + ) + + @app.post("/v1/chat/completions") + async def chat(payload: dict): + return {"ok": True, "n": len(payload.get("text", ""))} + + @app.post("/api/other") + async def other(payload: dict): + return {"ok": True, "unprotected": True} + + @app.get("/api/train/status") + async def status_get(): + return {"ok": True, "get": True} + + return app + + +class TestMaxBodyMiddleware: + def test_small_protected_body_passes(self, main_module): + app = _make_protected_app(1024, main_module) + c = TestClient(app) + r = c.post("/v1/chat/completions", json = {"text": "x" * 100}) + assert r.status_code == 200 + assert r.json()["n"] == 100 + + def test_large_declared_content_length_rejected(self, main_module): + app = _make_protected_app(1024, main_module) + c = TestClient(app) + r = c.post("/v1/chat/completions", json = {"text": "x" * 5000}) + assert r.status_code == 413 + assert "too large" in r.json()["detail"].lower() + + def test_unprotected_prefix_passes_large_body(self, main_module): + app = _make_protected_app(1024, main_module) + c = TestClient(app) + r = c.post("/api/other", json = {"text": "x" * 5000}) + assert r.status_code == 200 + assert r.json()["unprotected"] is True + + def test_chunked_upload_over_cap_rejected(self, main_module): + # Regression: declared-Content-Length-only check could be bypassed + # by chunked transfer-encoding. + app = _make_protected_app(1024, main_module) + c = TestClient(app) + + def gen(): + yield b'{"text":"' + yield b"x" * 800 + yield b'"}' + yield b"\n" + b"y" * 500 + + r = c.post( + "/v1/chat/completions", + content = gen(), + headers = {"content-type": "application/json"}, + ) + assert r.status_code == 413 + assert "too large" in r.json()["detail"].lower() + + def test_chunked_upload_under_cap_passes(self, main_module): + app = _make_protected_app(1024, main_module) + c = TestClient(app) + + def gen(): + yield b'{"text":"' + yield b"x" * 50 + yield b'"}' + + r = c.post( + "/v1/chat/completions", + content = gen(), + headers = {"content-type": "application/json"}, + ) + assert r.status_code == 200 + assert r.json()["n"] == 50 + + def test_get_not_subject_to_cap(self, main_module): + app = _make_protected_app(1024, main_module) + c = TestClient(app) + r = c.get("/api/train/status") + assert r.status_code == 200 + + +# ===================================================================== +# SecurityHeadersMiddleware / CSP +# ===================================================================== + + +def _make_csp_app(main_module, attach_nonce: str | None = None): + app = FastAPI() + app.add_middleware(main_module.SecurityHeadersMiddleware) + + @app.get("/plain") + async def plain(): + return {"ok": True} + + @app.get("/with-nonce") + async def with_nonce(): + headers = {} + if attach_nonce: + headers[main_module._CSP_SCRIPT_NONCE_HEADER] = attach_nonce + return Response( + content = b"", + media_type = "text/html", + headers = headers, + ) + + return app + + +class TestSecurityHeadersMiddleware: + def test_csp_has_no_unsafe_inline_for_script_src(self, main_module): + app = _make_csp_app(main_module) + c = TestClient(app) + r = c.get("/plain") + assert r.status_code == 200 + csp = r.headers["content-security-policy"] + # Parse per-directive so style-src unsafe-inline does not false-match. + directives = { + chunk.strip().split(" ", 1)[0]: chunk.strip() + for chunk in csp.split(";") + if chunk.strip() + } + assert "script-src" in directives + assert "'unsafe-inline'" not in directives["script-src"] + # style-src keeps unsafe-inline for Vite-injected styles. + assert "'unsafe-inline'" in directives["style-src"] + + def test_default_security_headers_present(self, main_module): + app = _make_csp_app(main_module) + c = TestClient(app) + r = c.get("/plain") + assert r.headers["x-frame-options"] == "DENY" + assert r.headers["x-content-type-options"] == "nosniff" + assert r.headers["referrer-policy"] == "no-referrer" + assert "camera=()" in r.headers["permissions-policy"] + assert r.headers["server"] == "unsloth-studio" + + def test_internal_nonce_header_is_spliced_into_csp_and_stripped(self, main_module): + nonce = "test-nonce-abc" + app = _make_csp_app(main_module, attach_nonce = nonce) + c = TestClient(app) + r = c.get("/with-nonce") + csp = r.headers["content-security-policy"] + assert f"'nonce-{nonce}'" in csp + # Internal handoff header must not leak to clients. + assert main_module._CSP_SCRIPT_NONCE_HEADER not in { + k.lower() for k in r.headers.keys() + } + + def test_build_csp_helper_shape(self, main_module): + plain = main_module._build_csp() + assert "script-src 'self';" in plain + assert "'unsafe-inline'" not in plain.split("script-src", 1)[1].split(";", 1)[0] + nonced = main_module._build_csp("XYZ") + assert "script-src 'self' 'nonce-XYZ';" in nonced + + +# ===================================================================== +# /api/health auth gate +# ===================================================================== + + +@pytest.fixture +def health_app(tmp_path, monkeypatch): + """Mount /api/health on a fresh app against an isolated auth db.""" + from auth import storage + + monkeypatch.setattr(storage, "DB_PATH", tmp_path / "auth.db") + monkeypatch.setattr(storage, "_BOOTSTRAP_PW_PATH", tmp_path / ".bootstrap_password") + monkeypatch.setattr(storage, "_bootstrap_password", None) + + import main as _main + + app = FastAPI() + app.add_api_route("/api/health", _main.health_check, methods = ["GET"]) + + import secrets as _secrets + + storage.create_initial_user( + username = storage.DEFAULT_ADMIN_USERNAME, + password = "human-password-123", + jwt_secret = _secrets.token_urlsafe(64), + must_change_password = False, + ) + return app + + +class TestHealthAuthGate: + def test_no_auth_returns_minimal_payload(self, health_app): + c = TestClient(health_app) + r = c.get("/api/health") + assert r.status_code == 200 + body = r.json() + assert body["status"] == "healthy" + assert "timestamp" in body + for forbidden in ("version", "device_type", "studio_root_id"): + assert forbidden not in body + + def test_invalid_bearer_returns_minimal_payload(self, health_app): + # Regression: calling the async dep without await made any Bearer header pass. + c = TestClient(health_app) + r = c.get( + "/api/health", + headers = {"Authorization": "Bearer not-a-real-token"}, + ) + assert r.status_code == 200 + body = r.json() + assert body["status"] == "healthy" + for forbidden in ("version", "device_type", "studio_root_id"): + assert forbidden not in body + + def test_valid_bearer_returns_full_payload(self, health_app): + from auth import storage + from auth.authentication import create_access_token + + token = create_access_token(storage.DEFAULT_ADMIN_USERNAME) + c = TestClient(health_app) + r = c.get( + "/api/health", + headers = {"Authorization": f"Bearer {token}"}, + ) + assert r.status_code == 200 + body = r.json() + assert body["status"] == "healthy" + assert "version" in body + assert "device_type" in body + assert "studio_root_id" in body diff --git a/studio/backend/tests/test_openai_tool_passthrough.py b/studio/backend/tests/test_openai_tool_passthrough.py index cdb7f5d270..a379282b70 100644 --- a/studio/backend/tests/test_openai_tool_passthrough.py +++ b/studio/backend/tests/test_openai_tool_passthrough.py @@ -125,22 +125,21 @@ class TestChatMessageToolRoles: ) assert msg.content is None - def test_tool_role_missing_tool_call_id_rejected(self): - # Per OpenAI spec, role="tool" messages must carry tool_call_id so - # upstream backends can associate the result with its prior call. - # Pin the boundary-level rejection so a malformed tool-result - # message never reaches the passthrough path. - with pytest.raises(ValidationError) as exc_info: - ChatMessage(role = "tool", content = '{"temperature": 72}') - assert "tool_call_id" in str(exc_info.value) + def test_tool_role_missing_tool_call_id_synthesised(self): + # Frontend drops the id on second-round POST; validator synthesises one. + msg = ChatMessage(role = "tool", content = '{"temperature": 72}') + assert msg.tool_call_id is not None + assert msg.tool_call_id.startswith("call_") + assert len(msg.tool_call_id) >= len("call_") + 8 - def test_tool_role_empty_tool_call_id_rejected(self): - with pytest.raises(ValidationError): - ChatMessage( - role = "tool", - tool_call_id = "", - content = '{"temperature": 72}', - ) + def test_tool_role_empty_tool_call_id_synthesised(self): + msg = ChatMessage( + role = "tool", + tool_call_id = "", + content = '{"temperature": 72}', + ) + assert msg.tool_call_id is not None + assert msg.tool_call_id.startswith("call_") # ── Role-aware content requirements ──────────────────────────── @@ -162,10 +161,19 @@ class TestChatMessageToolRoles: ChatMessage(role = "tool", tool_call_id = "call_1", content = "") assert "content" in str(exc_info.value) - def test_assistant_without_content_or_tool_calls_rejected(self): - with pytest.raises(ValidationError) as exc_info: - ChatMessage(role = "assistant") - assert "content" in str(exc_info.value) or "tool_calls" in str(exc_info.value) + def test_assistant_without_content_or_tool_calls_tolerated(self): + # Stop-button leaves an empty assistant turn; tolerate so replay round-trips. + msg = ChatMessage(role = "assistant") + assert msg.content is None + assert msg.tool_calls is None + + def test_assistant_empty_string_content_normalised_to_none(self): + msg = ChatMessage(role = "assistant", content = "") + assert msg.content is None + + def test_assistant_empty_list_content_normalised_to_none(self): + msg = ChatMessage(role = "assistant", content = []) + assert msg.content is None # ── Role-constrained tool-call metadata ──────────────────────── @@ -472,3 +480,91 @@ class TestFriendlyErrorHttpx: assert ( _friendly_error(RuntimeError("unrelated")) == "An internal error occurred" ) + + +from routes.inference import ( # noqa: E402 + _drop_empty_assistant_sentinels, + _openai_messages_for_passthrough, +) + + +class TestDropEmptyAssistantSentinels: + def test_drops_empty_assistant_between_real_turns(self): + msgs = [ + {"role": "user", "content": "hi"}, + {"role": "assistant", "content": ""}, + {"role": "user", "content": "again"}, + ] + out = _drop_empty_assistant_sentinels(msgs) + assert out == [ + {"role": "user", "content": "hi"}, + {"role": "user", "content": "again"}, + ] + + def test_drops_assistant_with_no_content_key(self): + # exclude_none=True strips the content key entirely; filter must catch this. + msgs = [ + {"role": "user", "content": "hi"}, + {"role": "assistant"}, + {"role": "user", "content": "ok"}, + ] + out = _drop_empty_assistant_sentinels(msgs) + assert out == [ + {"role": "user", "content": "hi"}, + {"role": "user", "content": "ok"}, + ] + + def test_preserves_assistant_with_text(self): + msgs = [ + {"role": "user", "content": "hi"}, + {"role": "assistant", "content": "hello back"}, + ] + out = _drop_empty_assistant_sentinels(msgs) + assert out == msgs + + def test_preserves_assistant_with_tool_calls_only(self): + msgs = [ + {"role": "user", "content": "weather?"}, + { + "role": "assistant", + "tool_calls": [ + { + "id": "call_1", + "type": "function", + "function": {"name": "get_weather", "arguments": "{}"}, + }, + ], + }, + { + "role": "tool", + "tool_call_id": "call_1", + "content": '{"t": 72}', + }, + ] + out = _drop_empty_assistant_sentinels(msgs) + assert out == msgs + + def test_preserves_user_and_system_with_empty_content(self): + # Filter scoped to role="assistant" only. + msgs = [ + {"role": "system", "content": ""}, + {"role": "user", "content": ""}, + ] + out = _drop_empty_assistant_sentinels(msgs) + assert out == msgs + + def test_openai_messages_for_passthrough_drops_sentinel(self): + """End-to-end: Stop-sentinel must not reach the wire.""" + req = ChatCompletionRequest( + model = "default", + messages = [ + ChatMessage(role = "user", content = "hi"), + ChatMessage(role = "assistant", content = ""), + ChatMessage(role = "user", content = "again"), + ], + ) + out = _openai_messages_for_passthrough(req) + roles = [m["role"] for m in out] + assert roles == ["user", "user"] + for m in out: + assert m.get("content"), m diff --git a/studio/backend/tests/test_sandbox_tools.py b/studio/backend/tests/test_sandbox_tools.py new file mode 100644 index 0000000000..fcc531c212 --- /dev/null +++ b/studio/backend/tests/test_sandbox_tools.py @@ -0,0 +1,241 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. + +"""Tests for the sandboxed-Python AST policy in core/inference/tools.py.""" + +import os +import sys +from pathlib import Path + +import pytest + +_BACKEND_ROOT = Path(__file__).resolve().parents[1] +if str(_BACKEND_ROOT) not in sys.path: + sys.path.insert(0, str(_BACKEND_ROOT)) + +from core.inference.tools import _check_code_safety + + +def _ok(code: str): + assert _check_code_safety(code) is None, code + + +def _blocked(code: str, *, expect_phrase: str): + msg = _check_code_safety(code) + assert msg is not None, code + assert expect_phrase in msg, (expect_phrase, msg) + + +class TestMetadataHostDenylist: + def test_aws_imds_literal_blocked(self): + _blocked( + 'import requests; requests.get("http://169.254.169.254/latest/meta-data/")', + expect_phrase = "Blocked: cloud-metadata host", + ) + + def test_gcp_metadata_dns_blocked(self): + _blocked( + 'import requests; requests.get("http://metadata.google.internal/")', + expect_phrase = "Blocked: cloud-metadata host", + ) + + def test_alibaba_ecs_literal_blocked(self): + _blocked( + 'import socket; s=socket.socket(); s.connect(("100.100.100.200", 80))', + expect_phrase = "Blocked: cloud-metadata host", + ) + + def test_ipv6_imds_literal_blocked(self): + _blocked( + 'import urllib.request; urllib.request.urlopen("http://[fd00:ec2::254]/")', + expect_phrase = "Blocked: cloud-metadata host", + ) + + def test_metadata_link_local_prefix_blocked(self): + _blocked( + 'import requests; requests.get("http://169.254.170.2/v3/")', + expect_phrase = "Blocked: cloud-metadata host", + ) + + +class TestTrustedHostAllowlist: + @pytest.mark.parametrize( + "url", + [ + "https://en.wikipedia.org/wiki/Python_(programming_language)", + "https://fr.wikipedia.org/wiki/Python_(langage)", + "https://www.google.com/search?q=foo", + "https://duckduckgo.com/?q=foo", + "https://huggingface.co/unsloth", + "https://cdn-lfs.huggingface.co/repos/abc/def/file.bin", + "https://raw.githubusercontent.com/foo/bar/main/README.md", + "https://api.github.com/repos/foo/bar", + "https://arxiv.org/abs/2401.12345", + "https://export.arxiv.org/abs/2401.12345", + "https://stackoverflow.com/questions/12345", + "https://math.stackexchange.com/questions/12345", + "https://developer.mozilla.org/en-US/docs/Web/JavaScript", + "https://docs.python.org/3/library/asyncio.html", + "https://pypi.org/project/requests/", + "https://files.pythonhosted.org/packages/foo/bar.whl", + "https://www.bbc.com/news", + "https://api.weather.gov/points/40,-90", + "https://numpy.org/doc/stable/", + "https://pytorch.org/docs/stable/index.html", + ], + ) + def test_trusted_host_passes(self, url): + _ok(f"import requests; requests.get({url!r})") + + def test_wikipedia_subdomain_passes(self): + _ok( + 'import urllib.request; urllib.request.urlopen("https://m.en.wikipedia.org/wiki/Foo")' + ) + + def test_hf_co_short_form_passes(self): + _ok('import requests; requests.get("https://hf.co/unsloth/Qwen3.5-4B-GGUF")') + + def test_github_io_pages_pass(self): + _ok('import requests; requests.get("https://unslothai.github.io/")') + + +class TestUntrustedHostBlock: + def test_example_com_blocked(self): + _blocked( + 'import requests; requests.get("https://example.com/")', + expect_phrase = "Blocked: host not in sandbox allowlist", + ) + + def test_random_blog_blocked(self): + _blocked( + 'import urllib.request; urllib.request.urlopen("https://random-blog-host.example/")', + expect_phrase = "Blocked: host not in sandbox allowlist", + ) + + def test_socket_connect_random_host_blocked(self): + _blocked( + 'import socket; s=socket.socket(); s.connect(("evil.example", 80))', + expect_phrase = "Blocked: host not in sandbox allowlist", + ) + + def test_dynamic_url_not_statically_blocked(self): + # Static AST cannot resolve runtime URLs; bash blocklist is the fallback. + _ok('import requests; url = "https://example.com/"; requests.get(url)') + + +class TestHostNormalization: + def test_trailing_dot_treated_same(self): + _ok('import requests; requests.get("https://wikipedia.org./")') + + def test_explicit_port_does_not_unblock_or_misblock(self): + _ok('import requests; requests.get("https://en.wikipedia.org:443/wiki/Foo")') + _blocked( + 'import requests; requests.get("https://example.com:8080/")', + expect_phrase = "Blocked: host not in sandbox allowlist", + ) + + def test_userinfo_at_does_not_smuggle_metadata_host(self): + _blocked( + 'import requests; requests.get("https://wikipedia.org@169.254.169.254/latest/")', + expect_phrase = "Blocked: cloud-metadata host", + ) + + def test_uppercase_host_normalised(self): + _ok('import requests; requests.get("https://EN.WIKIPEDIA.ORG/wiki/Foo")') + + +class TestUploadDenylist: + def test_requests_post_files_blocked(self): + _blocked( + ( + "import requests\n" + 'requests.post("https://huggingface.co/api/repos/upload", ' + 'files={"f": open("x.bin", "rb")})' + ), + expect_phrase = "Blocked: file upload disallowed in sandbox", + ) + + def test_requests_put_data_bytes_blocked(self): + _blocked( + ( + "import requests\n" + 'requests.put("https://huggingface.co/api/repos/upload", ' + 'data=b"\\x00\\x01\\x02")' + ), + expect_phrase = "Blocked: file upload disallowed in sandbox", + ) + + def test_requests_post_data_open_handle_blocked(self): + _blocked( + ( + "import requests\n" + 'requests.post("https://huggingface.co/api/repos/upload", ' + 'data=open("x.bin", "rb"))' + ), + expect_phrase = "Blocked: file upload disallowed in sandbox", + ) + + def test_httpx_post_files_blocked(self): + _blocked( + ( + "import httpx\n" + 'httpx.post("https://huggingface.co/api/repos/upload", ' + 'files={"f": open("x.bin", "rb")})' + ), + expect_phrase = "Blocked: file upload disallowed in sandbox", + ) + + def test_hf_api_upload_file_blocked(self): + _blocked( + ( + "from huggingface_hub import HfApi\n" + 'HfApi().upload_file(path_or_fileobj="x.bin", ' + 'path_in_repo="x.bin", repo_id="foo/bar")' + ), + expect_phrase = "Blocked: file upload disallowed in sandbox", + ) + + def test_hf_module_upload_folder_blocked(self): + _blocked( + ( + "import huggingface_hub\n" + 'huggingface_hub.upload_folder(folder_path="./", repo_id="foo/bar")' + ), + expect_phrase = "Blocked: file upload disallowed in sandbox", + ) + + def test_hf_create_commit_method_blocked(self): + _blocked( + ( + "import huggingface_hub\n" + "api = huggingface_hub.HfApi()\n" + 'api.create_commit(repo_id="foo/bar", operations=[])' + ), + expect_phrase = "Blocked: file upload disallowed in sandbox", + ) + + def test_plain_post_json_not_blocked(self): + _ok( + "import requests\n" + 'requests.post("https://api.weather.gov/lookup", json={"k": "v"})' + ) + + +class TestSandboxCpuRlimitDefault: + """Pin the default so a regression below 600s without opt-in is caught.""" + + def test_default_cpu_s_is_600(self): + src = (_BACKEND_ROOT / "core" / "inference" / "tools.py").read_text() + assert 'UNSLOTH_STUDIO_SANDBOX_CPU_S", "600"' in src + + def test_clone_newnet_removed(self): + src = (_BACKEND_ROOT / "core" / "inference" / "tools.py").read_text() + assert "_libc.unshare(0x40000000)" not in src + # Explanatory comment retained. + assert "CLONE_NEWNET" in src + + +class TestMaxBodyDefault: + def test_default_is_500_mb(self): + src = (_BACKEND_ROOT / "main.py").read_text() + assert 'UNSLOTH_STUDIO_MAX_BODY_MB", "500"' in src diff --git a/studio/backend/tests/test_studio_train_validation.py b/studio/backend/tests/test_studio_train_validation.py new file mode 100644 index 0000000000..7ffa9bb384 --- /dev/null +++ b/studio/backend/tests/test_studio_train_validation.py @@ -0,0 +1,90 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. + +"""Pin TrainingStartRequest hyperparameter caps at the at-cap / over-cap boundary.""" + +import sys +from pathlib import Path + +import pytest +from pydantic import ValidationError + +_BACKEND_ROOT = Path(__file__).resolve().parents[1] +if str(_BACKEND_ROOT) not in sys.path: + sys.path.insert(0, str(_BACKEND_ROOT)) + +from models.training import ( + _MAX_BATCH_SIZE, + _MAX_LORA_ALPHA, + _MAX_LORA_R, + _MAX_SEQ_LENGTH, +) + + +def _check_field(field_name: str, value): + """Run the field validator without constructing a full TrainingStartRequest.""" + from models.training import TrainingStartRequest + + schema_field = TrainingStartRequest.model_fields[field_name] + return TrainingStartRequest.__pydantic_validator__.validate_assignment( + TrainingStartRequest.model_construct(), + field_name, + value, + ) + + +class TestSeqLengthCap: + def test_at_cap_accepts(self): + _check_field("max_seq_length", _MAX_SEQ_LENGTH) + assert _MAX_SEQ_LENGTH == 2_000_000 + + def test_over_cap_rejects(self): + with pytest.raises(ValidationError) as exc: + _check_field("max_seq_length", _MAX_SEQ_LENGTH + 1) + assert "max_seq_length" in str(exc.value) + + def test_below_min_rejects(self): + with pytest.raises(ValidationError): + _check_field("max_seq_length", 0) + + +class TestBatchSizeCap: + def test_at_cap_accepts(self): + _check_field("batch_size", _MAX_BATCH_SIZE) + assert _MAX_BATCH_SIZE == 4096 + + def test_over_cap_rejects(self): + with pytest.raises(ValidationError): + _check_field("batch_size", _MAX_BATCH_SIZE + 1) + + def test_below_min_rejects(self): + with pytest.raises(ValidationError): + _check_field("batch_size", 0) + + +class TestLoraRCap: + def test_at_cap_accepts(self): + _check_field("lora_r", _MAX_LORA_R) + assert _MAX_LORA_R == 16_384 + + def test_over_cap_rejects(self): + with pytest.raises(ValidationError): + _check_field("lora_r", _MAX_LORA_R + 1) + + def test_below_min_rejects(self): + with pytest.raises(ValidationError): + _check_field("lora_r", 0) + + +class TestLoraAlphaCap: + def test_at_cap_accepts(self): + _check_field("lora_alpha", _MAX_LORA_ALPHA) + assert _MAX_LORA_ALPHA == 32_768 + + def test_over_cap_rejects(self): + with pytest.raises(ValidationError): + _check_field("lora_alpha", _MAX_LORA_ALPHA + 1) + + def test_below_min_rejects(self): + with pytest.raises(ValidationError): + _check_field("lora_alpha", 0) diff --git a/studio/backend/tests/test_trained_model_scan.py b/studio/backend/tests/test_trained_model_scan.py index 84be681fca..8ba97af701 100644 --- a/studio/backend/tests/test_trained_model_scan.py +++ b/studio/backend/tests/test_trained_model_scan.py @@ -28,7 +28,16 @@ from utils.models.model_config import ( ) -def test_scan_trained_models_includes_lora_and_full_finetune_outputs(tmp_path: Path): +def test_scan_trained_models_includes_lora_and_full_finetune_outputs( + tmp_path: Path, monkeypatch +): + # resolve_output_dir refuses absolutes outside outputs_root; point it at tmp_path. + from utils.models import model_config as _mc + from utils.paths import storage_roots as _sr + + monkeypatch.setattr(_sr, "outputs_root", lambda: tmp_path) + monkeypatch.setattr(_mc, "outputs_root", lambda: tmp_path) + lora_dir = tmp_path / "unsloth_SmolLM-135M_1775412608" lora_dir.mkdir() (lora_dir / "adapter_config.json").write_text( diff --git a/studio/backend/utils/paths/storage_roots.py b/studio/backend/utils/paths/storage_roots.py index 58a4d7967c..763d18bf3e 100644 --- a/studio/backend/utils/paths/storage_roots.py +++ b/studio/backend/utils/paths/storage_roots.py @@ -276,21 +276,52 @@ def _clean_relative_path( return Path(*parts) if parts else Path() +def _assert_contained(resolved: Path, root: Path) -> None: + """Raise ValueError if ``resolved`` realpaths outside ``root``.""" + try: + resolved_real = Path(os.path.realpath(resolved)) + root_real = Path(os.path.realpath(root)) + except OSError as exc: + raise ValueError(f"path resolution failed: {exc}") from exc + try: + resolved_real.relative_to(root_real) + except ValueError as exc: + raise ValueError( + f"path escapes root: {resolved!s} -> {resolved_real!s} " + f"is not under {root_real!s}" + ) from exc + + def resolve_under_root( path_value: str | None, *, root: Path, strip_prefixes: tuple[str, ...] = (), ) -> Path: + """Resolve ``path_value`` and assert the result is under ``root``. + + Absolutes are accepted only if already contained (so internal pre-resolved + paths re-enter idempotently); user-facing schemas reject absolutes upstream. + """ if not path_value or not str(path_value).strip(): return root - path = Path(str(path_value).strip()).expanduser() + raw = str(path_value).strip() + if "\x00" in raw: + raise ValueError("path may not contain null bytes") + + path = Path(raw).expanduser() + if ".." in path.parts: + raise ValueError(f"path may not contain '..' segments: {raw!r}") + if path.is_absolute(): + _assert_contained(path, root) return path - cleaned = _clean_relative_path(str(path), strip_prefixes = strip_prefixes) - return root / cleaned + cleaned = _clean_relative_path(raw, strip_prefixes = strip_prefixes) + candidate = root / cleaned + _assert_contained(candidate, root) + return candidate def resolve_output_dir(path_value: str | None = None) -> Path: @@ -318,9 +349,22 @@ def resolve_tensorboard_dir(path_value: str | None = None) -> Path: def resolve_dataset_path(path_value: str) -> Path: - path = Path(path_value).expanduser() + raw = str(path_value or "").strip() + if "\x00" in raw: + raise ValueError("dataset path may not contain null bytes") + path = Path(raw).expanduser() + if ".." in path.parts: + raise ValueError(f"dataset path may not contain '..' segments: {raw!r}") if path.is_absolute(): - return path + for root_fn in (datasets_root, dataset_uploads_root, recipe_datasets_root): + try: + _assert_contained(path, root_fn()) + return path + except ValueError: + continue + raise ValueError( + f"dataset path must be relative or under a dataset root: {raw!r}" + ) parts = [part for part in Path(path_value).parts if part not in ("", ".")] if parts[:2] == ["assets", "datasets"]: diff --git a/tests/studio/studio_api_smoke.py b/tests/studio/studio_api_smoke.py index 9e04630391..d6718defcc 100644 --- a/tests/studio/studio_api_smoke.py +++ b/tests/studio/studio_api_smoke.py @@ -316,18 +316,47 @@ if code in (400, 422): else: fail(f"/api/auth/refresh without body returned {code} (expected 400/422)") -# Login burst with wrong password must keep returning 401, NOT 429. -# Documents that no rate-limit / brute-force lockout exists today. -# When/if we add one, this assertion updates in the same PR. -all_401 = True -for i in range(5): - code, _ = login("definitely-wrong-password") - if code != 401: - all_401 = False - fail(f"login burst attempt {i+1} returned {code} (expected 401)") + +# Wrong-password burst: expect 401 until the per-IP bucket fills, then +# 429 with Retry-After. Bucket cannot be reset between tests, so we +# assert the observable invariant rather than a fixed transition index. +def _login_with_headers(password: str) -> tuple[int, str | None]: + """Like ``login`` but returns ``(status, retry_after_header)``.""" + url = f"{BASE}/api/auth/login" + data = json.dumps({"username": "unsloth", "password": password}).encode() + req = urllib.request.Request( + url, + data = data, + method = "POST", + headers = {"Content-Type": "application/json"}, + ) + try: + with urllib.request.urlopen(req, timeout = 10) as r: + return r.status, r.headers.get("Retry-After") + except urllib.error.HTTPError as exc: + return exc.code, exc.headers.get("Retry-After") if exc.headers else None + + +codes = [] +retry_after = None +for i in range(8): + code, ra = _login_with_headers("definitely-wrong-password") + codes.append(code) + if code == 429: + retry_after = ra break -if all_401: - ok("login burst (5x wrong pw) -> 401 each (no rate-limit, documented)") + if code != 401: + fail(f"login burst attempt {i+1} returned {code} (expected 401 or 429)") + break + +if 401 not in codes: + fail(f"login burst never returned 401 before rate-limit (codes={codes})") +elif 429 not in codes: + fail(f"login burst never rate-limited after {len(codes)} wrongs (codes={codes})") +elif retry_after is None: + fail("429 response missing Retry-After header") +else: + ok(f"login burst -> 401x{codes.count(401)} then 429 with Retry-After={retry_after}") # ───────────────────────────────────────────────────────────────────────── diff --git a/tests/test_studio_install_workspace_guard.py b/tests/test_studio_install_workspace_guard.py index c9aa3b2744..b68f882126 100644 --- a/tests/test_studio_install_workspace_guard.py +++ b/tests/test_studio_install_workspace_guard.py @@ -593,12 +593,16 @@ def test_install_ps1_bakes_studio_root_id_into_launcher(): def test_health_endpoint_exposes_studio_root_id_not_raw_path(): """studio/backend/main.py /api/health must expose studio_root_id (a hex digest) and NOT the raw studio_root path. Studio supports - `-H 0.0.0.0`; an unauthenticated /api/health that returns the raw - install path leaks username, home dir, workspace name, etc.""" + `-H 0.0.0.0`; a /api/health that returns the raw install path + leaks username, home dir, workspace name, etc.""" main_py = REPO_ROOT / "studio" / "backend" / "main.py" src = main_py.read_text() health_idx = src.index('@app.get("/api/health")') - health_block = src[health_idx : health_idx + 1500] + # Slice up to the next top-level @app. so a growing body stays in scope. + next_app_idx = src.find("\n@app.", health_idx + 1) + if next_app_idx == -1: + next_app_idx = len(src) + health_block = src[health_idx:next_app_idx] assert ( '"studio_root_id"' in health_block ), "/api/health must expose studio_root_id (hex digest)"