diff --git a/studio/backend/core/inference/tools.py b/studio/backend/core/inference/tools.py index b23372b766e..6c148e92b00 100644 --- a/studio/backend/core/inference/tools.py +++ b/studio/backend/core/inference/tools.py @@ -14,6 +14,8 @@ os.environ["UNSLOTH_IS_PRESENT"] = "1" import random +import re +import shlex import ssl import subprocess import sys @@ -27,11 +29,235 @@ _EXEC_TIMEOUT = 300 # 5 minutes +# Pre-import modules used in _sandbox_preexec at module level so that +# the preexec_fn closure does not trigger the import machinery in the +# forked child (which can deadlock in multi-threaded servers). +_libc = None +if sys.platform == "linux": + try: + import ctypes + import ctypes.util + + _libc_name = ctypes.util.find_library("c") + if _libc_name: + _libc = ctypes.CDLL(_libc_name, use_errno = True) + except (OSError, AttributeError): + pass + +_resource = None +if sys.platform != "win32": + try: + import resource as _resource + except ImportError: + pass + # Strict raster-image allowlist for sandbox file serving. # No .svg (XSS risk via embedded scripts), no .html, no .pdf. _IMAGE_EXTS = frozenset({".png", ".jpg", ".jpeg", ".gif", ".webp", ".bmp"}) _MAX_OUTPUT_CHARS = 8000 # truncate long output -_BASH_BLOCKED_WORDS = {"rm", "sudo", "dd", "chmod", "mkfs", "shutdown", "reboot"} +_BLOCKED_COMMANDS_COMMON = frozenset( + { + "rm", + "sudo", + "su", + "dd", + "chmod", + "chown", + "mkfs", + "shutdown", + "reboot", + "passwd", + "mount", + "umount", + "fdisk", + "kill", + "killall", + "pkill", + } +) +_BLOCKED_COMMANDS_WIN = frozenset( + { + "rmdir", + "takeown", + "icacls", + "runas", + "powershell", + "pwsh", + } +) +_BLOCKED_COMMANDS = ( + _BLOCKED_COMMANDS_COMMON | _BLOCKED_COMMANDS_WIN + if sys.platform == "win32" + else _BLOCKED_COMMANDS_COMMON +) + + +def _find_blocked_commands(command: str) -> set[str]: + """Detect blocked commands using shlex tokenization and regex scanning. + + Catches: full paths (/usr/bin/sudo), quoted strings ("sudo"), + split-quotes (su""do), backslash escapes (\\rm), and command-position + words after ;, |, &&, $(). + """ + blocked = set() + + # 1. shlex tokenization (handles quotes, escapes, concatenation) + try: + tokens = ( + shlex.split(command) + if sys.platform != "win32" + else shlex.split(command, posix = False) + ) + except ValueError: + tokens = command.split() + + for token in tokens: + base = os.path.basename(token).lower() + # Strip common Windows executable extensions so that + # runas.exe, shutdown.bat, etc. match the blocklist. + stem, ext = os.path.splitext(base) + if ext in {".exe", ".com", ".bat", ".cmd"}: + base = stem + if base in _BLOCKED_COMMANDS: + blocked.add(base) + + # 2. Regex: catch blocked words at shell command boundaries + # (semicolons, pipes, &&, ||, backticks, $(), <(), subshells, newlines) + # Uses a single combined pattern for all blocked words. + # Handles optional Unix path prefix (/usr/bin/) and Windows drive + # letter prefix (C:\Windows\...\). + lowered = command.lower() + if _BLOCKED_COMMANDS: + words_alt = "|".join(re.escape(w) for w in sorted(_BLOCKED_COMMANDS)) + pattern = ( + rf"(?:^|[;&|`\n(]\s*|[$]\(\s*|<\(\s*)" + rf"(?:[\w./\\-]*/|[a-zA-Z]:[/\\][\w./\\-]*)?" + rf"({words_alt})(?:\.(?:exe|com|bat|cmd))?\b" + ) + blocked.update(re.findall(pattern, lowered)) + + # 3. Check for nested shell invocations (bash -c 'sudo whoami', + # bash -lc '...', bash --login -c '...', cmd /c '...'). + # When a -c or /c flag is found, look backwards for a shell name + # (skipping intermediate flags like --login, -l, -x) and recursively + # scan the nested command string. + _SHELLS = {"bash", "sh", "zsh", "dash", "ksh", "csh", "tcsh", "fish"} + _SHELLS_WIN = {"cmd", "cmd.exe"} + for i, token in enumerate(tokens): + tok_lower = token.lower() + # Match -c exactly, or combined flags ending in c (e.g. -lc, -xc) + is_unix_c = tok_lower == "-c" or ( + tok_lower.startswith("-") + and tok_lower.endswith("c") + and not tok_lower.startswith("--") + ) + is_win_c = tok_lower == "/c" + if not (is_unix_c or is_win_c) or i < 1 or i + 1 >= len(tokens): + continue + # Look backwards past any flags to find the shell binary. + # On Unix, flags start with - (skip those). On Windows, flags + # start with / but so do absolute paths, so only skip short + # single-char /X flags (not /bin/bash style paths). + for j in range(i - 1, -1, -1): + prev = tokens[j] + if prev.startswith("-"): + continue # skip Unix flags like --login, -l + if is_win_c and prev.startswith("/") and len(prev) <= 3: + continue # skip Windows flags like /s, /q (not /bin/bash) + prev_base = os.path.basename(prev).lower() + if is_unix_c and prev_base in _SHELLS: + blocked |= _find_blocked_commands(tokens[i + 1]) + elif is_win_c and prev_base in _SHELLS_WIN: + blocked |= _find_blocked_commands(tokens[i + 1]) + break # stop at first non-flag token + + return blocked + + +def _build_safe_env(workdir: str) -> dict[str, str]: + """Build a minimal, credential-free environment for sandboxed subprocesses. + + Strips HF_TOKEN, WANDB_API_KEY, AWS_*, GH_TOKEN, LD_PRELOAD, DYLD_*, etc. + Preserves the active Python interpreter and virtualenv directories in PATH + so that pip, uv, and packages installed in the Studio runtime remain + accessible. + """ + # Start with the directory containing the running Python interpreter + # so that subprocess calls to 'python', 'pip', etc. resolve to the + # same environment the Studio server is running in. + exe_dir = os.path.dirname(sys.executable) + path_entries = [exe_dir] if exe_dir else [] + + # If a virtualenv is active, include its bin/Scripts directory. + venv = os.environ.get("VIRTUAL_ENV") + if venv: + venv_bin = os.path.join(venv, "Scripts" if sys.platform == "win32" else "bin") + if venv_bin not in path_entries: + path_entries.append(venv_bin) + + if sys.platform == "win32": + sysroot = os.environ.get("SystemRoot", r"C:\Windows") + path_entries.extend([os.path.join(sysroot, "System32"), sysroot]) + else: + path_entries.extend(["/usr/local/bin", "/usr/bin", "/bin"]) + + # Deduplicate while preserving order + deduped = list(dict.fromkeys(p for p in path_entries if p)) + + env = { + "PATH": os.pathsep.join(deduped), + "HOME": workdir, + "TMPDIR": workdir, + "LANG": os.environ.get("LANG", "C.UTF-8"), + "TERM": "dumb", + "PYTHONIOENCODING": "utf-8", + } + if venv: + env["VIRTUAL_ENV"] = venv + # Windows needs SystemRoot for Python/subprocess to work + if sys.platform == "win32": + env["SystemRoot"] = os.environ.get("SystemRoot", r"C:\Windows") + return env + + +def _sandbox_preexec(): + """Pre-exec hook: drop privilege escalation ability and set resource limits. + + 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. + """ + if _libc is not None: + try: + # 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.) + + if _resource is not None: + try: + # Limit file size to 100MB (prevents disk filling) + _resource.setrlimit( + _resource.RLIMIT_FSIZE, (100 * 1024 * 1024, 100 * 1024 * 1024) + ) + except (ValueError, OSError): + pass + + +def _get_shell_cmd(command: str) -> list[str]: + """Return the platform-appropriate shell invocation for a command string.""" + if sys.platform == "win32": + return ["cmd", "/c", command] + return ["bash", "-c", command] + # Per-session working directories so each chat thread gets its own sandbox. # Falls back to a shared ~/studio_sandbox/ for API callers without a session_id. @@ -428,6 +654,7 @@ def _check_signal_escape_patterns(code: str): signal_tampering = [] exception_catching = [] + shell_escapes = [] warnings = [] def _ast_name_matches(node, names): @@ -445,10 +672,84 @@ def _ast_name_matches(node, names): return full_name in names return False + # Dangerous os/subprocess functions that can execute shell commands + _SHELL_EXEC_FUNCS = frozenset( + { + "os.system", + "os.popen", + "os.popen2", + "os.popen3", + "os.popen4", + "os.execl", + "os.execle", + "os.execlp", + "os.execlpe", + "os.execv", + "os.execve", + "os.execvp", + "os.execvpe", + "os.spawnl", + "os.spawnle", + "os.spawnlp", + "os.spawnlpe", + "os.spawnv", + "os.spawnve", + "os.spawnvp", + "os.spawnvpe", + "os.posix_spawn", + "os.posix_spawnp", + "subprocess.run", + "subprocess.call", + "subprocess.check_call", + "subprocess.check_output", + "subprocess.Popen", + "subprocess.getoutput", + "subprocess.getstatusoutput", + } + ) + + def _extract_string_from_node(node): + """Extract a plain string value from an AST node, if it is a constant.""" + if isinstance(node, ast.Constant) and isinstance(node.value, str): + return node.value + return None + + def _extract_strings_from_list(node): + """Extract string elements from an AST List or Tuple node.""" + if isinstance(node, (ast.List, ast.Tuple)): + parts = [] + for elt in node.elts: + s = _extract_string_from_node(elt) + if s is not None: + parts.append(s) + return parts + return [] + + # Keyword argument names that carry command content (as opposed to + # control flags like check=True, text=True, capture_output=True). + _CMD_KWARGS = frozenset({"args", "command", "executable", "path", "file"}) + + def _check_args_for_blocked(args_nodes): + """Check if any call arguments contain blocked commands.""" + found = set() + for arg in args_nodes: + s = _extract_string_from_node(arg) + if s is not None: + found |= _find_blocked_commands(s) + strs = _extract_strings_from_list(arg) + for s in strs: + found |= _find_blocked_commands(s) + return found + class SignalEscapeVisitor(ast.NodeVisitor): def __init__(self): self.imports_signal = False self.signal_aliases = {"signal"} + self.os_aliases = {"os"} + self.subprocess_aliases = {"subprocess"} + # Maps bare function names to their fully-qualified form + # for from-import tracking (e.g. "system" -> "os.system") + self.shell_exec_aliases: dict[str, str] = {} self.loop_depth = 0 def visit_Import(self, node): @@ -457,6 +758,10 @@ def visit_Import(self, node): self.imports_signal = True if alias.asname: self.signal_aliases.add(alias.asname) + elif alias.name == "os": + self.os_aliases.add(alias.asname or "os") + elif alias.name == "subprocess": + self.subprocess_aliases.add(alias.asname or "subprocess") self.generic_visit(node) def visit_ImportFrom(self, node): @@ -474,6 +779,16 @@ def visit_ImportFrom(self, node): "alarm", ): self.signal_aliases.add(alias.asname or alias.name) + elif node.module in ("os", "subprocess"): + if node.module == "os": + self.os_aliases.add("os") + else: + self.subprocess_aliases.add("subprocess") + # Track from-imports of dangerous functions + for alias in node.names: + fq = f"{node.module}.{alias.name}" + if fq in _SHELL_EXEC_FUNCS: + self.shell_exec_aliases[alias.asname or alias.name] = fq self.generic_visit(node) def visit_While(self, node): @@ -538,6 +853,111 @@ def visit_Call(self, node): "description": "Modifies signal mask (may block SIGALRM)", } ) + + # --- Shell escape detection --- + # Resolve the fully qualified function name for os.*/subprocess.* + shell_func = None + if isinstance(func, ast.Attribute): + if isinstance(func.value, ast.Name): + if func.value.id in self.os_aliases: + shell_func = f"os.{func.attr}" + elif func.value.id in self.subprocess_aliases: + shell_func = f"subprocess.{func.attr}" + elif isinstance(func, ast.Name): + # Check from-import aliases: from os import system; system(...) + shell_func = self.shell_exec_aliases.get(func.id) + + if shell_func and shell_func in _SHELL_EXEC_FUNCS: + # Expand **kwargs dicts to inspect their keys + expanded_kwargs: dict[str, ast.AST] = {} + has_opaque_kwargs = False + for kw in node.keywords: + if kw.arg is not None: + expanded_kwargs[kw.arg] = kw.value + elif isinstance(kw.value, ast.Dict): + for k, v in zip(kw.value.keys, kw.value.values): + key = _extract_string_from_node(k) if k else None + if key is not None: + expanded_kwargs[key] = v + else: + has_opaque_kwargs = True + + cmd_kw_values = [ + v for k, v in expanded_kwargs.items() if k in _CMD_KWARGS + ] + all_call_args = list(node.args) + cmd_kw_values + blocked_in_args = _check_args_for_blocked(all_call_args) + + if has_opaque_kwargs: + # Can't inspect dynamic **kwargs -- flag as unsafe + shell_escapes.append( + { + "type": "shell_escape_dynamic", + "line": node.lineno, + "description": ( + f"{shell_func}() called with dynamic **kwargs" + ), + } + ) + elif blocked_in_args: + shell_escapes.append( + { + "type": "shell_escape", + "line": node.lineno, + "description": ( + f"{shell_func}() invokes blocked command(s): " + f"{', '.join(sorted(blocked_in_args))}" + ), + } + ) + else: + # Only flag dynamic args for functions that interpret + # strings as shell commands, or when shell= might be + # enabled. Treat any non-literal-False shell= value + # as potentially True (conservative). + _STRING_SHELL_FUNCS = frozenset( + { + "os.system", + "os.popen", + "os.popen2", + "os.popen3", + "os.popen4", + "subprocess.getoutput", + "subprocess.getstatusoutput", + } + ) + shell_node = expanded_kwargs.get("shell") + shell_safe = shell_node is None or ( + isinstance(shell_node, ast.Constant) + and shell_node.value is False + ) + if shell_func in _STRING_SHELL_FUNCS or not shell_safe: + + def _is_safe_literal(n): + if _extract_string_from_node(n) is not None: + return True + if isinstance(n, (ast.List, ast.Tuple)): + return all( + _extract_string_from_node(e) is not None + for e in n.elts + ) + return False + + has_non_literal = any( + not _is_safe_literal(a) for a in all_call_args + ) + if has_non_literal: + shell_escapes.append( + { + "type": "shell_escape_dynamic", + "line": node.lineno, + "description": ( + f"{shell_func}() called with non-literal " + f"shell command (potential shell escape)" + ), + } + ) + self.generic_visit(node) def visit_ExceptHandler(self, node): @@ -553,7 +973,12 @@ def visit_ExceptHandler(self, node): } ) elif isinstance(node.type, ast.Name): - if node.type.id in ("TimeoutError", "BaseException", "Exception"): + # Only flag BaseException and TimeoutError, NOT Exception. + # except Exception does not catch SystemExit or + # KeyboardInterrupt, so it cannot suppress timeout + # enforcement. Flagging Exception causes false positives + # on normal error-handling patterns. + if node.type.id in ("TimeoutError", "BaseException"): exception_catching.append( { "type": f"catches_{node.type.id}_in_loop", @@ -564,7 +989,7 @@ def visit_ExceptHandler(self, node): elif isinstance(node.type, ast.Tuple): for elt in node.type.elts: if isinstance(elt, ast.Name): - if elt.id in ("TimeoutError", "BaseException", "Exception"): + if elt.id in ("TimeoutError", "BaseException"): exception_catching.append( { "type": f"catches_{elt.id}_in_loop", @@ -580,10 +1005,15 @@ def visit_ExceptHandler(self, node): if visitor.imports_signal and not signal_tampering: warnings.append("Code imports 'signal' module - review manually for safety") - is_safe = len(signal_tampering) == 0 and len(exception_catching) == 0 + is_safe = ( + len(signal_tampering) == 0 + and len(exception_catching) == 0 + and len(shell_escapes) == 0 + ) return is_safe, { "signal_tampering": signal_tampering, "exception_catching": exception_catching, + "shell_escapes": shell_escapes, "warnings": warnings, } @@ -604,10 +1034,18 @@ def _check_code_safety(code: str) -> str | None: reasons = [ item.get("description", "") for item in info.get("signal_tampering", []) ] - return ( - f"Error: unsafe code detected ({'; '.join(reasons)}). " - f"Please remove signal manipulation from your code." - ) + shell_reasons = [ + item.get("description", "") for item in info.get("shell_escapes", []) + ] + 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] + if all_reasons: + return ( + f"Error: unsafe code detected ({'; '.join(all_reasons)}). " + f"Please remove unsafe patterns from your code." + ) return None @@ -662,13 +1100,20 @@ def _python_exec( with os.fdopen(fd, "w") as f: f.write(code) - proc = subprocess.Popen( - [sys.executable, tmp_path], + safe_env = _build_safe_env(workdir) + popen_kwargs = dict( stdout = subprocess.PIPE, stderr = subprocess.STDOUT, text = True, cwd = workdir, + env = safe_env, ) + if sys.platform != "win32": + popen_kwargs["preexec_fn"] = _sandbox_preexec + else: + popen_kwargs["creationflags"] = subprocess.CREATE_NO_WINDOW + + proc = subprocess.Popen([sys.executable, tmp_path], **popen_kwargs) # Spawn cancel watcher if we have a cancel event if cancel_event is not None: @@ -734,21 +1179,27 @@ def _bash_exec( if not command or not command.strip(): return "No command provided." - # Block dangerous commands - tokens = set(command.lower().split()) - blocked = tokens & _BASH_BLOCKED_WORDS + # Block dangerous commands (shlex + regex based) + blocked = _find_blocked_commands(command) if blocked: return f"Blocked command(s) for safety: {', '.join(sorted(blocked))}" try: workdir = _get_workdir(session_id) - proc = subprocess.Popen( - ["bash", "-c", command], + safe_env = _build_safe_env(workdir) + popen_kwargs = dict( stdout = subprocess.PIPE, stderr = subprocess.STDOUT, text = True, cwd = workdir, + env = safe_env, ) + if sys.platform != "win32": + popen_kwargs["preexec_fn"] = _sandbox_preexec + else: + popen_kwargs["creationflags"] = subprocess.CREATE_NO_WINDOW + + proc = subprocess.Popen(_get_shell_cmd(command), **popen_kwargs) if cancel_event is not None: watcher = threading.Thread(