diff --git a/studio/backend/auth/storage.py b/studio/backend/auth/storage.py index 9a03f5f542..3233aa05ef 100644 --- a/studio/backend/auth/storage.py +++ b/studio/backend/auth/storage.py @@ -480,6 +480,37 @@ def save_refresh_token( conn.close() +def consume_refresh_token(token: str) -> Optional[Tuple[str, bool]]: + """Atomically validate-and-delete a refresh token for single-use rotation. + + DELETE RETURNING fuses validate and delete into one statement so two + concurrent refresh requests cannot both consume the same token. + """ + token_hash = _hash_token(token) + now = datetime.now(timezone.utc).isoformat() + conn = get_connection() + try: + conn.execute( + "DELETE FROM refresh_tokens WHERE expires_at < ?", + (now,), + ) + cur = conn.execute( + """ + DELETE FROM refresh_tokens + WHERE token_hash = ? AND expires_at >= ? + RETURNING username, is_desktop + """, + (token_hash, now), + ) + row = cur.fetchone() + conn.commit() + if row is None: + return None + return row["username"], bool(row["is_desktop"]) + finally: + conn.close() + + def verify_refresh_token(token: str) -> Optional[Tuple[str, bool]]: """ Verify a refresh token and return the username plus desktop marker. diff --git a/studio/backend/core/inference/tools.py b/studio/backend/core/inference/tools.py index 87cc933d4b..70db5477d4 100644 --- a/studio/backend/core/inference/tools.py +++ b/studio/backend/core/inference/tools.py @@ -10,6 +10,7 @@ Supports web search (DuckDuckGo), Python code execution, and terminal commands. import ast import http.client import os +import signal os.environ["UNSLOTH_IS_PRESENT"] = "1" @@ -58,21 +59,37 @@ _MAX_OUTPUT_CHARS = 8000 # truncate long output _BLOCKED_COMMANDS_COMMON = frozenset( { "rm", - "sudo", - "su", "dd", "chmod", "chown", "mkfs", - "shutdown", - "reboot", - "passwd", "mount", "umount", "fdisk", + "sudo", + "su", + "doas", + "pkexec", + "shutdown", + "reboot", + "halt", + "poweroff", "kill", "killall", "pkill", + "passwd", + "curl", + "wget", + "nc", + "ncat", + "netcat", + "socat", + "ssh", + "scp", + "sftp", + "rsync", + "eval", + "source", } ) _BLOCKED_COMMANDS_WIN = frozenset( @@ -221,35 +238,67 @@ def _build_safe_env(workdir: str) -> dict[str, str]: def _sandbox_preexec(): - """Pre-exec hook: drop privilege escalation ability and set resource limits. + """Best-effort sandbox setup for sandboxed subprocesses. - On Linux, applies PR_SET_NO_NEW_PRIVS so sudo/su/pkexec fail at the - kernel level. On Linux and macOS, sets RLIMIT_FSIZE. - No-op on Windows (use creationflags instead). - - Note: RLIMIT_NPROC is intentionally NOT set because Linux enforces it - per real UID, not per process tree, so it would starve the Studio - server and other sessions sharing the same user account. - - All modules and handles are resolved at import time (module level) so - this function does not trigger Python imports in the forked child, - avoiding potential deadlocks in multi-threaded servers. + Modules are resolved at import time so the forked child runs no imports. """ + try: + os.setsid() + except OSError: + pass + + try: + os.umask(0o077) + except OSError: + pass + if _libc is not None: try: - # PR_SET_NO_NEW_PRIVS = 38, arg2 = 1 (enable) - _libc.prctl(38, 1, 0, 0, 0) + _libc.prctl(38, 1, 0, 0, 0) # PR_SET_NO_NEW_PRIVS except (OSError, AttributeError): - pass # Not available (container, old kernel, etc.) + pass + + try: + _libc.prctl(1, 9, 0, 0, 0) # PR_SET_PDEATHSIG = SIGKILL + except (OSError, AttributeError): + pass + + # CLONE_NEWNET intentionally not applied: where userns is enabled it + # blocks all egress, including allowlisted hosts. Network policy is + # enforced by the AST host check and the bash blocklist. if _resource is not None: + # RLIMIT_NPROC is per-real-UID, so the cap is well above normal usage. + try: + nproc = int(os.environ.get("UNSLOTH_STUDIO_SANDBOX_NPROC", "10000")) + _resource.setrlimit(_resource.RLIMIT_NPROC, (nproc, nproc)) + except (ValueError, OSError, AttributeError): + pass try: - # Limit file size to 100MB (prevents disk filling) _resource.setrlimit( _resource.RLIMIT_FSIZE, (100 * 1024 * 1024, 100 * 1024 * 1024) ) except (ValueError, OSError): pass + try: + as_bytes = ( + int(os.environ.get("UNSLOTH_STUDIO_SANDBOX_AS_GB", "8")) + * 1024 + * 1024 + * 1024 + ) + _resource.setrlimit(_resource.RLIMIT_AS, (as_bytes, as_bytes)) + except (ValueError, OSError, AttributeError): + pass + try: + cpu_s = int(os.environ.get("UNSLOTH_STUDIO_SANDBOX_CPU_S", "600")) + _resource.setrlimit(_resource.RLIMIT_CPU, (cpu_s, cpu_s)) + except (ValueError, OSError, AttributeError): + pass + try: + _resource.setrlimit(_resource.RLIMIT_NOFILE, (1024, 1024)) + except (ValueError, OSError, AttributeError): + pass def _get_shell_cmd(command: str) -> list[str]: @@ -265,25 +314,36 @@ def _get_shell_cmd(command: str) -> list[str]: _workdirs: dict[str, str] = {} +# Non-matching session_ids collapse to ``_invalid`` to block cross-session escapes. +_SESSION_ID_RE = re.compile(r"\A[A-Za-z0-9_\-]{1,64}\Z") + + def _get_workdir(session_id: str | None = None) -> str: - """Return (and lazily create) a persistent working directory for tool execution.""" + """Return a per-session sandbox dir at mode 0o700.""" global _workdirs key = session_id or "_default" if key not in _workdirs or not os.path.isdir(_workdirs[key]): home = os.path.expanduser("~") sandbox_root = os.path.join(home, "studio_sandbox") - if session_id: - # Sanitize: strip path separators and parent-dir references - safe_id = os.path.basename(session_id.replace("..", "")) - if not safe_id: - safe_id = "_invalid" - workdir = os.path.join(sandbox_root, safe_id) - # Verify resolved path stays under sandbox root - if not os.path.realpath(workdir).startswith(os.path.realpath(sandbox_root)): + if session_id and _SESSION_ID_RE.match(session_id): + workdir = os.path.join(sandbox_root, session_id) + if not os.path.realpath(workdir).startswith( + os.path.realpath(sandbox_root) + os.sep + ): workdir = os.path.join(sandbox_root, "_invalid") + elif session_id: + workdir = os.path.join(sandbox_root, "_invalid") else: workdir = os.path.join(sandbox_root, "_default") os.makedirs(workdir, exist_ok = True) + try: + os.chmod(sandbox_root, 0o700) + except OSError: + pass + try: + os.chmod(workdir, 0o700) + except OSError: + pass _workdirs[key] = workdir return _workdirs[key] @@ -932,7 +992,12 @@ def _check_signal_escape_patterns(code: str): isinstance(shell_node, ast.Constant) and shell_node.value is False ) - if shell_func in _STRING_SHELL_FUNCS or not shell_safe: + # Dynamic shell-exec args (chr/format/concat bypasses). + if ( + shell_func in _STRING_SHELL_FUNCS + or shell_func in _SHELL_EXEC_FUNCS + or not shell_safe + ): def _is_safe_literal(n): if _extract_string_from_node(n) is not None: @@ -1006,15 +1071,418 @@ def _check_signal_escape_patterns(code: str): if visitor.imports_signal and not signal_tampering: warnings.append("Code imports 'signal' module - review manually for safety") + # Static host policy: block metadata hosts and any literal host outside + # the trusted allowlist; uploads blocked regardless of host. Dynamic hosts + # are caught by the bash blocklist instead. + network_calls: list[dict] = [] + sensitive_file_reads: list[dict] = [] + _NETWORK_FQ_PREFIXES = ( + "socket.socket", + "socket.create_connection", + "socket.getaddrinfo", + "urllib.request.urlopen", + "urllib.request.urlretrieve", + "urllib3.", + "requests.get", + "requests.post", + "requests.put", + "requests.delete", + "requests.patch", + "requests.head", + "requests.request", + "requests.Session", + "http.client.HTTPConnection", + "http.client.HTTPSConnection", + "httpx.get", + "httpx.post", + "httpx.put", + "httpx.patch", + "httpx.delete", + "httpx.request", + "httpx.Client", + "httpx.AsyncClient", + "aiohttp.ClientSession", + ) + _UPLOAD_HTTP_METHODS = ( + "requests.post", + "requests.put", + "requests.patch", + "requests.delete", + "requests.request", + "httpx.post", + "httpx.put", + "httpx.patch", + "httpx.delete", + "httpx.request", + "urllib.request.urlopen", + "urllib.request.Request", + ) + _UPLOAD_HF_FQ = ( + "huggingface_hub.upload_file", + "huggingface_hub.upload_folder", + "huggingface_hub.upload_large_folder", + "huggingface_hub.create_commit", + ) + _UPLOAD_HF_METHODS = frozenset( + { + "upload_file", + "upload_folder", + "upload_large_folder", + "create_commit", + } + ) + # Cloud-metadata / link-local hosts. + _METADATA_HOST_LITERALS = { + "169.254.169.254", + "fd00:ec2::254", + "metadata.google.internal", + "metadata", + "metadata.tencentyun.com", + "100.100.100.200", + "100.100.100.110", + "169.254.170.2", + "169.254.170.23", + } + _METADATA_HOST_PREFIXES = ( + "169.254.", + "100.64.", + ) + # Allowlist kept explicit so each entry is auditable. + _TRUSTED_PUBLIC_HOST_LITERALS = frozenset( + { + # search + "www.google.com", + "google.com", + "www.bing.com", + "bing.com", + "duckduckgo.com", + "html.duckduckgo.com", + # encyclopedic / reference + "wikipedia.org", + "www.wikipedia.org", + "wikimedia.org", + "www.wikimedia.org", + "wikidata.org", + "www.wikidata.org", + "commons.wikimedia.org", + "www.britannica.com", + "openlibrary.org", + "www.openstreetmap.org", + # ML / dev / data + "huggingface.co", + "hf.co", + "github.com", + "api.github.com", + "raw.githubusercontent.com", + "gist.github.com", + "docs.github.com", + "pypi.org", + "files.pythonhosted.org", + "www.npmjs.com", + "registry.npmjs.org", + "crates.io", + "static.crates.io", + # docs + "docs.python.org", + "python.org", + "www.python.org", + "developer.mozilla.org", + "developer.apple.com", + "learn.microsoft.com", + "docs.docker.com", + "pytorch.org", + "docs.pytorch.org", + "tensorflow.org", + "www.tensorflow.org", + "numpy.org", + "pandas.pydata.org", + "scipy.org", + "scikit-learn.org", + "matplotlib.org", + "fastapi.tiangolo.com", + "starlette.io", + # academic + "arxiv.org", + "export.arxiv.org", + "scholar.google.com", + "openreview.net", + "semanticscholar.org", + "www.semanticscholar.org", + "biorxiv.org", + "www.biorxiv.org", + "medrxiv.org", + "www.medrxiv.org", + "pubmed.ncbi.nlm.nih.gov", + "www.ncbi.nlm.nih.gov", + # Q&A / community + "stackoverflow.com", + "stackexchange.com", + "askubuntu.com", + "superuser.com", + "serverfault.com", + # standards + "www.w3.org", + "tools.ietf.org", + "datatracker.ietf.org", + "www.rfc-editor.org", + # reputable news + "www.bbc.com", + "www.bbc.co.uk", + "www.reuters.com", + "apnews.com", + "www.nature.com", + "www.science.org", + # government / open data + "data.gov", + "catalog.data.gov", + "www.census.gov", + "www.nasa.gov", + "data.nasa.gov", + "www.cdc.gov", + "www.nih.gov", + "www.who.int", + # weather / time + "api.weather.gov", + "worldtimeapi.org", + } + ) + _TRUSTED_PUBLIC_HOST_SUFFIXES = ( + ".wikipedia.org", + ".wikimedia.org", + ".wiktionary.org", + ".wikibooks.org", + ".wikiquote.org", + ".wikisource.org", + ".wikiversity.org", + ".wikivoyage.org", + ".stackexchange.com", + ".hf.co", + ".huggingface.co", + ".githubusercontent.com", + ".github.io", + ".arxiv.org", + ".readthedocs.io", + ".readthedocs.org", + ) + _SENSITIVE_FILE_PREFIXES = ( + "/etc/passwd", + "/etc/shadow", + "/etc/sudoers", + "/etc/ssh/", + ) + _SENSITIVE_FILE_RE = re.compile( + r"^/proc/(?:self|\d+)/(?:environ|cmdline|task/\d+/environ)$" + ) + + def _normalize_host(host: str) -> str: + if not host: + return "" + h = host.strip().lower().rstrip(".") + if "@" in h: + h = h.split("@", 1)[1] + if h.startswith("[") and "]" in h: + h = h[1 : h.index("]")] + elif h.count(":") == 1: + h = h.split(":", 1)[0] + return h + + def _is_metadata_host(host: str) -> bool: + h = _normalize_host(host) + if not h: + return False + if h in _METADATA_HOST_LITERALS: + return True + if any(h.startswith(p) for p in _METADATA_HOST_PREFIXES): + return True + return False + + def _is_trusted_host(host: str) -> bool: + h = _normalize_host(host) + if not h: + return False + if h in _TRUSTED_PUBLIC_HOST_LITERALS: + return True + return any(h.endswith(s) for s in _TRUSTED_PUBLIC_HOST_SUFFIXES) + + def _call_is_upload_shape(node: ast.Call, fq: str) -> bool: + """True for statically obvious upload shapes (files=, data=open(), bytes literal).""" + if fq in _UPLOAD_HF_FQ: + return True + if fq not in _UPLOAD_HTTP_METHODS: + return False + for kw in node.keywords or []: + if kw.arg == "files": + return True + if kw.arg == "data": + v = kw.value + if ( + isinstance(v, ast.Call) + and isinstance(v.func, ast.Name) + and v.func.id == "open" + ): + return True + if isinstance(v, ast.Constant) and isinstance( + v.value, (bytes, bytearray) + ): + return True + return False + + def _method_call_is_hf_upload(node: ast.Call) -> bool: + """True for HfApi upload method names on any receiver.""" + return ( + isinstance(node.func, ast.Attribute) + and node.func.attr in _UPLOAD_HF_METHODS + ) + + class NetworkAndIoVisitor(ast.NodeVisitor): + def visit_Call(self, node): + parts: list[str] = [] + cur = node.func + while isinstance(cur, ast.Attribute): + parts.insert(0, cur.attr) + cur = cur.value + if isinstance(cur, ast.Name): + parts.insert(0, cur.id) + fq = ".".join(parts) if parts else "" + + if _method_call_is_hf_upload(node): + network_calls.append( + { + "type": "upload_blocked", + "line": getattr(node, "lineno", -1), + "description": ("Blocked: file upload disallowed in sandbox"), + } + ) + + # Direct sock.connect((host, port)) bypasses the FQ-prefix branch below. + if ( + isinstance(node.func, ast.Attribute) + and node.func.attr == "connect" + and node.args + ): + a0 = node.args[0] + host_lit = None + if isinstance(a0, ast.Tuple) and a0.elts: + e0 = a0.elts[0] + if isinstance(e0, ast.Constant) and isinstance(e0.value, str): + host_lit = e0.value + elif isinstance(a0, ast.Constant) and isinstance(a0.value, str): + host_lit = a0.value + if host_lit: + if _is_metadata_host(host_lit): + network_calls.append( + { + "type": "metadata_host_blocked", + "line": getattr(node, "lineno", -1), + "description": "Blocked: cloud-metadata host", + } + ) + elif not _is_trusted_host(host_lit): + network_calls.append( + { + "type": "untrusted_host_blocked", + "line": getattr(node, "lineno", -1), + "description": ( + "Blocked: host not in sandbox allowlist; " + "use an allowed informational source" + ), + } + ) + + if fq and any(fq.startswith(p) for p in _NETWORK_FQ_PREFIXES): + # 1) Upload-shape check (host-independent). + if _call_is_upload_shape(node, fq): + network_calls.append( + { + "type": "upload_blocked", + "line": getattr(node, "lineno", -1), + "description": ( + "Blocked: file upload disallowed in sandbox" + ), + } + ) + + # 2) Extract literal host (URL string or (host, port) tuple). + host_arg = None + url_arg = None + if node.args: + a0 = node.args[0] + if isinstance(a0, ast.Constant) and isinstance(a0.value, str): + url_arg = a0.value + elif isinstance(a0, ast.Tuple) and a0.elts: + e0 = a0.elts[0] + if isinstance(e0, ast.Constant) and isinstance(e0.value, str): + host_arg = e0.value + if url_arg and host_arg is None: + m = re.match(r"^\w+://([^/?#]+)", url_arg) + if m: + host_arg = m.group(1) + + if host_arg: + if _is_metadata_host(host_arg): + network_calls.append( + { + "type": "metadata_host_blocked", + "line": getattr(node, "lineno", -1), + "description": "Blocked: cloud-metadata host", + } + ) + elif not _is_trusted_host(host_arg): + network_calls.append( + { + "type": "untrusted_host_blocked", + "line": getattr(node, "lineno", -1), + "description": ( + "Blocked: host not in sandbox allowlist; " + "use an allowed informational source" + ), + } + ) + + is_open_call = ( + (isinstance(node.func, ast.Name) and node.func.id == "open") + or fq in ("io.open", "pathlib.Path.open") + or fq.endswith(".open") + ) + if is_open_call and node.args: + a0 = node.args[0] + path_lit = None + if isinstance(a0, ast.Constant) and isinstance(a0.value, str): + path_lit = a0.value + if path_lit: + flagged = False + if any(path_lit.startswith(p) for p in _SENSITIVE_FILE_PREFIXES): + flagged = True + elif _SENSITIVE_FILE_RE.match(path_lit): + flagged = True + if flagged: + sensitive_file_reads.append( + { + "type": "sensitive_file_read", + "line": getattr(node, "lineno", -1), + "description": ( + f"open({path_lit!r}) targets a host identity / " + "credential file; sandboxed code may not read it" + ), + } + ) + self.generic_visit(node) + + NetworkAndIoVisitor().visit(tree) + is_safe = ( len(signal_tampering) == 0 and len(exception_catching) == 0 and len(shell_escapes) == 0 + and len(network_calls) == 0 + and len(sensitive_file_reads) == 0 ) return is_safe, { "signal_tampering": signal_tampering, "exception_catching": exception_catching, "shell_escapes": shell_escapes, + "network_calls": network_calls, + "sensitive_file_reads": sensitive_file_reads, "warnings": warnings, } @@ -1041,7 +1509,21 @@ def _check_code_safety(code: str) -> str | None: exception_reasons = [ item.get("description", "") for item in info.get("exception_catching", []) ] - all_reasons = [r for r in reasons + shell_reasons + exception_reasons if r] + network_reasons = [ + item.get("description", "") for item in info.get("network_calls", []) + ] + file_reasons = [ + item.get("description", "") for item in info.get("sensitive_file_reads", []) + ] + all_reasons = [ + r + for r in reasons + + shell_reasons + + exception_reasons + + network_reasons + + file_reasons + if r + ] if all_reasons: return ( f"Error: unsafe code detected ({'; '.join(all_reasons)}). " @@ -1051,11 +1533,31 @@ def _check_code_safety(code: str) -> str | None: return None +def _kill_process_tree(proc) -> None: + """SIGKILL the setsid process group; fall back to single-pid kill.""" + if proc.poll() is not None: + return + try: + pgid = os.getpgid(proc.pid) + except (ProcessLookupError, PermissionError): + pgid = None + if pgid is not None: + try: + os.killpg(pgid, signal.SIGKILL) + return + except (ProcessLookupError, PermissionError): + pass + try: + proc.kill() + except (ProcessLookupError, PermissionError): + pass + + def _cancel_watcher(proc, cancel_event, poll_interval = 0.2): """Daemon thread that kills a process when cancel_event is set.""" while proc.poll() is None: if cancel_event is not None and cancel_event.is_set(): - proc.kill() + _kill_process_tree(proc) return cancel_event.wait(poll_interval) if cancel_event else None @@ -1126,8 +1628,11 @@ def _python_exec( try: output, _ = proc.communicate(timeout = timeout) except subprocess.TimeoutExpired: - proc.kill() - proc.communicate() + _kill_process_tree(proc) + try: + proc.communicate(timeout = 5) + except subprocess.TimeoutExpired: + pass return _truncate(f"Execution timed out after {timeout} seconds.") if cancel_event is not None and cancel_event.is_set(): @@ -1211,8 +1716,11 @@ def _bash_exec( try: output, _ = proc.communicate(timeout = timeout) except subprocess.TimeoutExpired: - proc.kill() - proc.communicate() + _kill_process_tree(proc) + try: + proc.communicate(timeout = 5) + except subprocess.TimeoutExpired: + pass return _truncate(f"Execution timed out after {timeout} seconds.") if cancel_event is not None and cancel_event.is_set(): diff --git a/studio/backend/core/training/training.py b/studio/backend/core/training/training.py index 72b13c3225..549d733252 100644 --- a/studio/backend/core/training/training.py +++ b/studio/backend/core/training/training.py @@ -17,7 +17,9 @@ Pattern follows core/data_recipe/jobs/manager.py. import json as _json import math import multiprocessing as mp +import os import queue +import shutil import threading import time import structlog @@ -33,9 +35,54 @@ from utils.native_path_leases import ( native_path_secret_removed_for_child_start, run_without_native_path_secret, ) +from utils.paths import outputs_root logger = get_logger(__name__) + +def _cleanup_cancelled_checkpoints(output_dir: str | os.PathLike) -> None: + """Remove ``checkpoint-`` subdirs after a cancelled run. + Only paths whose realpath is under outputs_root are touched.""" + out = Path(output_dir) + if not out.exists(): + return + try: + out_real = out.resolve() + out_root_real = Path(outputs_root()).resolve() + except OSError: + return + try: + out_real.relative_to(out_root_real) + except ValueError: + # Refuse to delete anything outside the configured outputs root. + logger.warning( + "Skipping checkpoint cleanup - %s is not under outputs_root %s", + out_real, + out_root_real, + ) + return + removed = 0 + for entry in out.iterdir() if out.is_dir() else []: + if not entry.is_dir(): + continue + name = entry.name + if not name.startswith("checkpoint-"): + continue + tail = name[len("checkpoint-") :] + if not tail.isdigit(): + continue + try: + shutil.rmtree(entry, ignore_errors = False) + removed += 1 + except OSError as exc: + logger.warning("Could not remove %s: %s", entry, exc) + logger.info( + "Cancelled-run cleanup removed %d checkpoint dir(s) under %s", + removed, + out, + ) + + _CTX = mp.get_context("spawn") # Plot styling constants @@ -316,6 +363,8 @@ class TrainingBackend: ) self._proc.terminate() proc = self._proc + cancelled = self._cancel_requested + output_dir = self._output_dir if proc is not None: proc.join(timeout = 5.0) @@ -328,6 +377,17 @@ class TrainingBackend: if self._pump_thread is not None and self._pump_thread.is_alive(): self._pump_thread.join(timeout = 8.0) + # Drop checkpoint-* dirs on explicit cancel only; stop-and-save + # keeps its artifacts. + if cancelled and output_dir: + try: + _cleanup_cancelled_checkpoints(output_dir) + except Exception: + logger.exception( + "Failed to clean up cancelled-run checkpoints under %s", + output_dir, + ) + def is_training_active(self) -> bool: """Check if training is currently active.""" with self._lock: diff --git a/studio/backend/main.py b/studio/backend/main.py index 2fba2756e7..81964c98e0 100644 --- a/studio/backend/main.py +++ b/studio/backend/main.py @@ -104,7 +104,7 @@ if os.getenv("ENVIRONMENT_TYPE", "production") == "production": # warnings.filterwarnings("ignore", category=DeprecationWarning) # warnings.filterwarnings("ignore", module="triton.*") -from fastapi import Depends, FastAPI, Request +from fastapi import Depends, FastAPI, HTTPException, Request from fastapi.middleware.cors import CORSMiddleware from fastapi.staticfiles import StaticFiles from fastapi.responses import FileResponse, HTMLResponse, Response @@ -260,6 +260,181 @@ logger = LogConfig.setup_logging( app.add_middleware(LoggingMiddleware) + +# Web-search favicons load from *.gstatic.com; everything else is same-origin. +from starlette.middleware.base import BaseHTTPMiddleware # noqa: E402 +from starlette.requests import Request as _StarletteRequest # noqa: E402 + + +_CSP_SCRIPT_NONCE_HEADER = "x-internal-script-nonce" + + +def _build_csp(script_nonce: "str | None" = None) -> str: + script_src = "script-src 'self'" + if script_nonce: + script_src += f" 'nonce-{script_nonce}'" + return ( + "default-src 'self'; " + "img-src 'self' data: blob: https://t0.gstatic.com " + "https://t1.gstatic.com https://t2.gstatic.com " + "https://t3.gstatic.com; " + "connect-src 'self'; " + "style-src 'self' 'unsafe-inline'; " + f"{script_src}; " + "font-src 'self' data:; " + "frame-ancestors 'none'; " + "form-action 'self'; " + "base-uri 'self'" + ) + + +class SecurityHeadersMiddleware(BaseHTTPMiddleware): + """Set baseline security headers; splice per-response inline-script nonces into CSP.""" + + async def dispatch(self, request: _StarletteRequest, call_next): + response = await call_next(request) + # Strip the internal nonce hand-off header so it never reaches the client. + nonce = response.headers.get(_CSP_SCRIPT_NONCE_HEADER) + if nonce is not None: + del response.headers[_CSP_SCRIPT_NONCE_HEADER] + response.headers.setdefault("Content-Security-Policy", _build_csp(nonce)) + response.headers.setdefault("X-Frame-Options", "DENY") + response.headers.setdefault("X-Content-Type-Options", "nosniff") + response.headers.setdefault("Referrer-Policy", "no-referrer") + response.headers.setdefault( + "Permissions-Policy", + "camera=(), microphone=(), geolocation=(), interest-cohort=()", + ) + response.headers["server"] = "unsloth-studio" + return response + + +app.add_middleware(SecurityHeadersMiddleware) + + +# Cap upload body on protected POSTs; default 500 MB, env-tunable. +import json as _json_for_413 # noqa: E402 + + +_MAX_BODY_BYTES = int(os.environ.get("UNSLOTH_STUDIO_MAX_BODY_MB", "500")) * 1024 * 1024 +_BODY_PROTECTED_PREFIXES = ( + "/v1/chat/completions", + "/v1/completions", + "/api/inference", + "/api/data-recipe", + "/api/datasets", + "/api/train", + "/api/export", +) + + +async def _send_413(send, total_bytes: int) -> None: + payload = _json_for_413.dumps( + { + "detail": ( + f"Request body too large " + f"({total_bytes:,} bytes; max {_MAX_BODY_BYTES:,})." + ) + }, + ).encode("utf-8") + await send( + { + "type": "http.response.start", + "status": 413, + "headers": [ + (b"content-type", b"application/json"), + (b"content-length", str(len(payload)).encode("ascii")), + ], + } + ) + await send({"type": "http.response.body", "body": payload, "more_body": False}) + + +class MaxBodyMiddleware: + """Reject oversized bodies on protected POST/PUT/PATCH; raw ASGI so chunked uploads cannot bypass the cap.""" + + def __init__(self, app, max_bytes: int, protected_prefixes: tuple): + self.app = app + self.max_bytes = max_bytes + self.protected_prefixes = protected_prefixes + + async def __call__(self, scope, receive, send): + if scope["type"] != "http": + await self.app(scope, receive, send) + return + method = scope.get("method", "").upper() + path = scope.get("path", "") + if method not in ("POST", "PUT", "PATCH") or not any( + path.startswith(p) for p in self.protected_prefixes + ): + await self.app(scope, receive, send) + return + + declared = None + for name, value in scope.get("headers", []): + if name == b"content-length": + try: + declared = int(value.decode("latin-1")) + except (ValueError, UnicodeDecodeError): + declared = None + break + if declared is not None and declared > self.max_bytes: + await _send_413(send, declared) + return + + chunks: list = [] + total = 0 + while True: + msg = await receive() + mtype = msg.get("type") + if mtype == "http.disconnect": + return + if mtype != "http.request": + # Mid-stream unexpected frame: forwarding would corrupt downstream. + return + body = msg.get("body", b"") or b"" + if body: + total += len(body) + if total > self.max_bytes: + await _send_413(send, total) + return + chunks.append(body) + if not msg.get("more_body", False): + break + + replayed = {"sent": False} + + async def replay_receive(): + if not replayed["sent"]: + replayed["sent"] = True + return { + "type": "http.request", + "body": b"".join(chunks), + "more_body": False, + } + # After replay, fall through so http.disconnect still propagates. + return await receive() + + await self.app(scope, replay_receive, send) + + +app.add_middleware( + MaxBodyMiddleware, + max_bytes = _MAX_BODY_BYTES, + protected_prefixes = _BODY_PROTECTED_PREFIXES, +) + + +from starlette.responses import RedirectResponse as _RedirectResponse # noqa: E402 + + +@app.get("/recipes", include_in_schema = False) +@app.get("/recipes/{rest:path}", include_in_schema = False) +async def _recipes_redirect(rest: str = ""): + target = "/data-recipes" + (("/" + rest) if rest else "") + return _RedirectResponse(url = target, status_code = 308) + + # CORS middleware _api_only = os.environ.get("UNSLOTH_API_ONLY") == "1" _cors_origins = ["*"] @@ -311,14 +486,35 @@ app.include_router( @app.get("/api/health") -async def health_check(): - """Health check endpoint""" - platform_map = {"darwin": "mac", "win32": "windows", "linux": "linux"} - device_type = platform_map.get(sys.platform, sys.platform) - - return { +async def health_check(request: Request): + """Liveness only; full diagnostic dict gated on a valid bearer.""" + minimal = { "status": "healthy", "timestamp": datetime.now().isoformat(), + } + auth = request.headers.get("authorization", "") + if not auth.lower().startswith("bearer "): + return minimal + try: + from auth.authentication import get_current_subject as _gcs + from fastapi.security import HTTPAuthorizationCredentials + + creds = HTTPAuthorizationCredentials( + scheme = "Bearer", credentials = auth.split(" ", 1)[1] + ) + # Must await: a bare coroutine is truthy and would skip the auth check. + subject = await _gcs(creds) + except HTTPException: + return minimal + except Exception: + return minimal + if not subject: + return minimal + + platform_map = {"darwin": "mac", "win32": "windows", "linux": "linux"} + device_type = platform_map.get(sys.platform, sys.platform) + return { + **minimal, "service": "Unsloth UI Backend", "version": UNSLOTH_VERSION, "studio_version": STUDIO_VERSION, @@ -328,9 +524,7 @@ async def health_check(): "desktop_manageability_version": 1, "supports_desktop_auth": True, "supports_desktop_backend_ownership": True, - # why: launchers compare against an install-time hash so a sibling - # Studio on the same port is rejected; hex digest avoids leaking the - # raw install path on -H 0.0.0.0. + # Hex digest of the install path; launchers reject sibling Studios on the same port. "studio_root_id": _studio_root_id(), "native_path_leases_supported": native_path_leases_supported(), **({"desktop_owner": owner} if (owner := _desktop_owner()) else {}), @@ -463,21 +657,22 @@ def _strip_crossorigin(html_bytes: bytes) -> bytes: return html.encode("utf-8") -def _inject_bootstrap(html_bytes: bytes, app: FastAPI) -> bytes: - """Inject bootstrap credentials into HTML when password change is required. +def _inject_bootstrap(html_bytes: bytes, app: FastAPI): + """Inject bootstrap credentials when password change is pending. - The script tag is only injected while the default admin account still - has ``must_change_password=True``. Once the user changes the password - the HTML is served clean — no credentials leak. + Returns ``(html_bytes, script_nonce_or_None)``. Callers must forward + the nonce via ``_CSP_SCRIPT_NONCE_HEADER`` so the inline script is + not blocked by CSP. """ import json as _json + import secrets as _secrets if not storage.requires_password_change(storage.DEFAULT_ADMIN_USERNAME): - return html_bytes + return html_bytes, None bootstrap_pw = getattr(app.state, "bootstrap_password", None) if not bootstrap_pw: - return html_bytes + return html_bytes, None payload = _json.dumps( { @@ -485,10 +680,11 @@ def _inject_bootstrap(html_bytes: bytes, app: FastAPI) -> bytes: "password": bootstrap_pw, } ) - tag = f"" + 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)"