diff --git a/studio/backend/core/inference/tools.py b/studio/backend/core/inference/tools.py index 87cc933d4b..eb8a781bc9 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" @@ -57,22 +58,43 @@ _IMAGE_EXTS = frozenset({".png", ".jpg", ".jpeg", ".gif", ".webp", ".bmp"}) _MAX_OUTPUT_CHARS = 8000 # truncate long output _BLOCKED_COMMANDS_COMMON = frozenset( { + # filesystem mutation / device control "rm", - "sudo", - "su", "dd", "chmod", "chown", "mkfs", - "shutdown", - "reboot", - "passwd", "mount", "umount", "fdisk", + # privilege escalation + "sudo", + "su", + "doas", + "pkexec", + # process control / host control + "shutdown", + "reboot", + "halt", + "poweroff", "kill", "killall", "pkill", + "passwd", + # network egress tools (defense in depth even with netns isolation) + "curl", + "wget", + "nc", + "ncat", + "netcat", + "socat", + "ssh", + "scp", + "sftp", + "rsync", + # shell-internal exec primitives + "eval", + "source", } ) _BLOCKED_COMMANDS_WIN = frozenset( @@ -221,35 +243,115 @@ def _build_safe_env(workdir: str) -> dict[str, str]: def _sandbox_preexec(): - """Pre-exec hook: drop privilege escalation ability and set resource limits. + """Pre-exec hook for every sandboxed subprocess (bash + python paths). - 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). + Sequence (each step is best-effort, failures are swallowed so the + sandbox still launches on kernels / configurations that lack a + feature): - 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. + 1. ``os.setsid()`` -> child becomes its own session leader; the + supervisor uses ``os.killpg(getpgid(...))`` so grandchildren die + too (closes the timeout-leak that left bash backgrounded children + alive after ``proc.kill()``). + 2. ``os.umask(0o077)`` -> files the child writes default to 0o600. + 3. ``PR_SET_NO_NEW_PRIVS`` -> sudo/su/pkexec are kernel-blocked. + 4. ``PR_SET_PDEATHSIG = SIGKILL`` -> orphaned children die when + Studio exits. + 5. ``unshare(CLONE_NEWNET)`` -> best-effort private network namespace + so the child has only ``lo``. Blocks AWS IMDS (169.254.169.254), + RFC1918, link-local, and all outbound. Requires unprivileged user + namespaces (kernel.unprivileged_userns_clone=1). Falls through + silently when denied so the sandbox still runs on locked-down + hosts; ``_BLOCKED_COMMANDS`` (curl/wget/nc/ssh/...) and the AST + network-call denylist are the defense-in-depth fallback. + 6. ``setrlimit(RLIMIT_NPROC, (256, 256))`` -> caps the per-real-UID + process count. The existing comment notes this is per-UID rather + than per-process-tree, so a misbehaving session can starve other + sessions of new processes - but that is still strictly better + than the prior fork-bomb-the-host behaviour. Tune via + ``UNSLOTH_STUDIO_SANDBOX_NPROC`` if it pinches. + 7. ``setrlimit(RLIMIT_AS / RLIMIT_CPU / RLIMIT_NOFILE)`` -> guard + against memory bombs and runaway loops. 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. """ + # 1. Become process-group leader so we can killpg the whole tree. + try: + os.setsid() + except OSError: + pass + + # 2. Restrictive umask for files the child writes. + try: + os.umask(0o077) + except OSError: + pass + if _libc is not None: try: - # PR_SET_NO_NEW_PRIVS = 38, arg2 = 1 (enable) + # 3. PR_SET_NO_NEW_PRIVS = 38, arg2 = 1 (enable) _libc.prctl(38, 1, 0, 0, 0) except (OSError, AttributeError): - pass # Not available (container, old kernel, etc.) + pass + + try: + # 4. PR_SET_PDEATHSIG = 1, arg2 = SIGKILL (9) + _libc.prctl(1, 9, 0, 0, 0) + except (OSError, AttributeError): + pass + + try: + # 5. unshare(CLONE_NEWNET=0x40000000) - best-effort private netns. + CLONE_NEWNET = 0x40000000 + rc = _libc.unshare(CLONE_NEWNET) + if rc != 0: + # Permission denied or unsupported - silently continue; + # the blocklist + AST checks remain the fallback. + pass + except (OSError, AttributeError): + pass if _resource is not None: + # 6. Cap per-real-UID process count. + # NOTE: Linux's RLIMIT_NPROC counts LWPs (kernel threads), not + # processes. Studio + llama-server already run several hundred + # LWPs as ``ubuntu``, so a low cap (256) starves legitimate bash + # forks ("bash: fork: retry: Resource temporarily unavailable"). + # The default 10000 is well above the normal LWP count but stops + # true fork bombs (which spawn millions) before they take down + # the host. Pair with the setsid + killpg cleanup at timeout to + # tear down any survivors. Tune via UNSLOTH_STUDIO_SANDBOX_NPROC. + try: + nproc = int(os.environ.get("UNSLOTH_STUDIO_SANDBOX_NPROC", "10000")) + _resource.setrlimit(_resource.RLIMIT_NPROC, (nproc, nproc)) + except (ValueError, OSError, AttributeError): + pass + # File size: 100 MB (preserved from original). try: - # Limit file size to 100MB (prevents disk filling) _resource.setrlimit( _resource.RLIMIT_FSIZE, (100 * 1024 * 1024, 100 * 1024 * 1024) ) except (ValueError, OSError): pass + # 7a. Address space cap (8 GB by default). + 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 + # 7b. CPU time cap (300 s by default - same as the supervisor timeout). + try: + cpu_s = int(os.environ.get("UNSLOTH_STUDIO_SANDBOX_CPU_S", "300")) + _resource.setrlimit(_resource.RLIMIT_CPU, (cpu_s, cpu_s)) + except (ValueError, OSError, AttributeError): + pass + # 7c. Open files cap. + try: + _resource.setrlimit(_resource.RLIMIT_NOFILE, (1024, 1024)) + except (ValueError, OSError, AttributeError): + pass def _get_shell_cmd(command: str) -> list[str]: @@ -265,25 +367,52 @@ def _get_shell_cmd(command: str) -> list[str]: _workdirs: dict[str, str] = {} +# Strict session-id charset. Anything else collapses to ``_invalid`` so +# cross-session escapes ("..", "/", null bytes, control chars, weird +# unicode) cannot synthesise a session dir that overlaps another's. +_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 (and lazily create) a persistent working directory for tool execution. + + Each session directory is created with mode 0o700 so other UIDs + cannot read it. The sandbox root itself is also tightened to 0o700 + on every call (cheap, idempotent, defensive). ``session_id`` is + matched against a strict regex; non-matching values bucket into a + shared ``_invalid`` dir rather than synthesising a partially- + sanitised name that could collide with a real session. + """ 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" + if session_id and _SESSION_ID_RE.match(session_id): + safe_id = session_id 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)): + # Belt + suspenders: realpath must land under realpath(sandbox_root). + if not os.path.realpath(workdir).startswith( + os.path.realpath(sandbox_root) + os.sep + ): workdir = os.path.join(sandbox_root, "_invalid") + elif session_id: + # Non-matching session_id (contains slashes, "..", null, etc.) + workdir = os.path.join(sandbox_root, "_invalid") else: workdir = os.path.join(sandbox_root, "_default") os.makedirs(workdir, exist_ok = True) + # Restrict so other local UIDs cannot read this session's files. + # The sandbox-root parent is also chmodded so previously-loose + # 0o775 directories tighten on next process boot. + 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 +1061,17 @@ 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: + # Always flag dynamic args to ANY shell-exec primitive. + # The previous gate let + # ``subprocess.Popen([chr(115)+...], shell=False)`` slip + # through (shell_safe=True and shell_func is not in + # _STRING_SHELL_FUNCS). Now any non-literal arg to + # a function in _SHELL_EXEC_FUNCS is flagged. + 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 +1145,199 @@ 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") + # ------------------------------------------------------------------ + # Narrow-band IO denylist. The goal is to keep ALL legitimate tool + # use cases working (web_search, public-internet fetches, file + # reads in user-space) while blocking the precise high-impact + # attack surfaces: + # + # 1. Cloud-metadata IPs (AWS / GCP / Alibaba / Azure IMDS). + # 2. Static reads of identity / credential system files + # (/etc/passwd, /etc/shadow, /etc/sudoers, /proc/*/environ). + # + # All other network calls (public DNS / public IPs) and file reads + # (any user dir, model output dir, /tmp, /etc/os-release, etc.) are + # allowed - tool-calling models legitimately read files and fetch + # URLs supplied by the user. + # ------------------------------------------------------------------ + 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.Client", + "httpx.AsyncClient", + "aiohttp.ClientSession", + ) + # Cloud-metadata / link-local / loopback-to-known-metadata. + _METADATA_HOST_LITERALS = { + "169.254.169.254", # AWS / standard + "fd00:ec2::254", # AWS IPv6 IMDS + "metadata.google.internal", + "metadata", + "metadata.tencentyun.com", + "100.100.100.200", # Alibaba ECS + "100.100.100.110", + "169.254.170.2", # ECS task metadata + "169.254.170.23", + } + _METADATA_HOST_PREFIXES = ( + "169.254.", + "100.64.", # CGNAT (often used internally) + ) + _SENSITIVE_FILE_PREFIXES = ( + "/etc/passwd", + "/etc/shadow", + "/etc/sudoers", + "/etc/ssh/", + ) + # /proc//environ and similar are matched via regex below. + _SENSITIVE_FILE_RE = re.compile( + r"^/proc/(?:self|\d+)/(?:environ|cmdline|task/\d+/environ)$" + ) + + def _is_metadata_host(host: str) -> bool: + if not host: + return False + if host in _METADATA_HOST_LITERALS: + return True + if any(host.startswith(p) for p in _METADATA_HOST_PREFIXES): + return True + return False + + class NetworkAndIoVisitor(ast.NodeVisitor): + def visit_Call(self, node): + # Fully-qualified attribute path. + 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 "" + + # Two-step pattern: any ``.connect((host, port))`` call where + # the literal host is a metadata IP. This catches the + # ``s = socket.socket(); s.connect(("169.254.169.254", 80))`` + # shape the FQ-prefix check below misses. + 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 and _is_metadata_host(host_lit): + network_calls.append( + { + "type": "metadata_host_blocked", + "line": getattr(node, "lineno", -1), + "description": ( + f".connect(({host_lit!r}, ...)) targets cloud-metadata; " + "sandboxed code may not call link-local / metadata IPs" + ), + } + ) + + if fq and any(fq.startswith(p) for p in _NETWORK_FQ_PREFIXES): + # Try to read the host literal if present. ONLY block when + # the literal host is a known metadata address. Dynamic / + # public hosts go through (legitimate tool use). + 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 + # Extract host from a URL literal if present. + 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 and _is_metadata_host(host_arg): + network_calls.append( + { + "type": "metadata_host_blocked", + "line": getattr(node, "lineno", -1), + "description": ( + f"{fq}({host_arg!r}) targets cloud-metadata; " + "sandboxed code may not call link-local / metadata IPs" + ), + } + ) + + # Static reads of identity / cred files. + 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 +1364,16 @@ 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 +1383,38 @@ def _check_code_safety(code: str) -> str | None: return None +def _kill_process_tree(proc) -> None: + """SIGKILL the whole process group, then fall back to single-pid kill. + + Pairs with the ``os.setsid()`` in :func:`_sandbox_preexec` so that + bash-backgrounded children, Python ``os.fork()`` survivors, and any + other grandchildren die with the supervised parent. Without this, + ``proc.kill()`` only signals the immediate pid and lets the rest of + the process group continue (the original fork-bomb risk). + """ + 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 +1485,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 +1573,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():