""" id: run_code title: Run code description: Run arbitrary Python or Bash code safely in a gVisor sandbox. author: EtiennePerot author_url: https://github.com/EtiennePerot/safe-code-execution funding_url: https://github.com/EtiennePerot/safe-code-execution version: 0.8.0 license: Apache-2.0 """ # NOTE: If running Open WebUI in a container, you *need* to set up this container to allow sandboxed code execution. # Please read the docs here: # # https://github.com/EtiennePerot/safe-code-execution/blob/master/README.md # # This is an OpenWebUI *function*. It can run code within LLM-generated code blocks. # If you are looking for an OpenWebUI *tool* to allow the LLM to run its own code, # see here instead: https://openwebui.com/t/etienneperot/run_code/ # # See https://github.com/EtiennePerot/safe-code-execution for more info. # # Protip: You can test this manually by running it as a Python script, like so: # (Run this inside the Open WebUI container) # # python3 run_code.py --self_test # # This will simulate that OpenWebUI would do if it asked this tool to evaluate the Python code `print("Hello world!")`. # This can be useful when setting up this tool to verify that it works in your environment. # You can also use it for one-off code execution like this: # # echo 'print("Hello world!")' | python3 run_code.py # import asyncio import argparse import base64 import fcntl import json import hashlib import mimetypes import os import os.path import pydantic import shutil import stat import subprocess import sys import tempfile import time import typing import urllib.parse import urllib.request import uuid import inspect import ctypes import ctypes.util import copy import platform import re import signal import threading import datetime import urllib.error class _Action: class Valves(pydantic.BaseModel): _VALVE_OVERRIDE_ENVIRONMENT_VARIABLE_NAME_PREFIX = "CODE_EVAL_VALVE_OVERRIDE_" NETWORKING_ALLOWED: bool = pydantic.Field( default=True, description=f"Whether to allow network access during code execution; may be overridden by environment variable {_VALVE_OVERRIDE_ENVIRONMENT_VARIABLE_NAME_PREFIX}NETWORKING_ALLOWED.", ) MAX_RUNTIME_SECONDS: int = pydantic.Field( ge=1, default=30, description=f"Maximum number of seconds code is given to run; may be overridden by environment variable {_VALVE_OVERRIDE_ENVIRONMENT_VARIABLE_NAME_PREFIX}MAX_RUNTIME_SECONDS.", ) MAX_RAM_MEGABYTES: int = pydantic.Field( ge=0, default=128, description=f"Maximum number of megabytes that the interpreter has when running. Must run as root with host cgroups writable (`--mount=type=bind,source=/sys/fs/cgroup,target=/sys/fs/cgroup,readonly=false`) for this to work. Set to 0 to disable memory limits. May be overridden by environment variable {_VALVE_OVERRIDE_ENVIRONMENT_VARIABLE_NAME_PREFIX}MAX_RAM_MEGABYTES", ) AUTO_INSTALL: bool = pydantic.Field( default=True, description=f"Whether to automatically install gVisor if not installed on the system; may be overridden by environment variable {_VALVE_OVERRIDE_ENVIRONMENT_VARIABLE_NAME_PREFIX}AUTO_INSTALL. Use the 'HTTPS_PROXY' environment variable to control the proxy used for download.", ) CHECK_FOR_UPDATES: bool = pydantic.Field( default=True, description=f"Whether to automatically check for updates; may be overridden by environment variable {_VALVE_OVERRIDE_ENVIRONMENT_VARIABLE_NAME_PREFIX}CHECK_FOR_UPDATES. Use the 'HTTPS_PROXY' environment variable to control the proxy used for update checks.", ) DEBUG: bool = pydantic.Field( default=False, description=f"Whether to produce debug logs during execution; may be overridden by environment variable {_VALVE_OVERRIDE_ENVIRONMENT_VARIABLE_NAME_PREFIX}DEBUG.", ) MAX_FILES_PER_EXECUTION: int = pydantic.Field( ge=1, default=32, description=f"Maximum number of generated files to allow per code execution; may be overridden by environment variable {_VALVE_OVERRIDE_ENVIRONMENT_VARIABLE_NAME_PREFIX}MAX_FILES_PER_EXECUTION.", ) MAX_FILES_PER_USER: int = pydantic.Field( ge=1, default=4096, description=f"Maximum number of files to keep around for a given user; may be overridden by environment variable {_VALVE_OVERRIDE_ENVIRONMENT_VARIABLE_NAME_PREFIX}MAX_FILES_PER_USER.", ) MAX_MEGABYTES_PER_USER: int = pydantic.Field( ge=1, default=256, description=f"Maximum total size of files to keep around for a given user; may be overridden by environment variable {_VALVE_OVERRIDE_ENVIRONMENT_VARIABLE_NAME_PREFIX}MAX_MEGABYTES_PER_USER.", ) REQUIRE_RESOURCE_LIMITING: bool = pydantic.Field( default=True, description=f"Whether to enforce resource limiting, which requires cgroups v2 to be available; may be overridden by environment variable {_VALVE_OVERRIDE_ENVIRONMENT_VARIABLE_NAME_PREFIX}REQUIRE_RESOURCE_LIMITING.", ) WEB_ACCESSIBLE_DIRECTORY_PATH: str = pydantic.Field( default="$DATA_DIR/cache/functions/run_code", description=f"Path of the directory to write files that should be accessible for user download in. If it begins by '$DATA_DIR', this will be replaced with the DATA_DIR environment variable. The whole field may be overridden by environment variable {_VALVE_OVERRIDE_ENVIRONMENT_VARIABLE_NAME_PREFIX}WEB_ACCESSIBLE_DIRECTORY_PATH.", ) WEB_ACCESSIBLE_DIRECTORY_URL: str = pydantic.Field( default="/cache/functions/run_code", description=f"URL corresponding to WEB_ACCESSIBLE_DIRECTORY_PATH. May start with '/' to make it relative to the Open WebUI serving domain. May be overridden by environment variable {_VALVE_OVERRIDE_ENVIRONMENT_VARIABLE_NAME_PREFIX}WEB_ACCESSIBLE_DIRECTORY_URL.", ) def __init__(self, valves): self.valves = valves for valve_name, valve_value in valves.dict().items(): override = os.getenv( self.valves._VALVE_OVERRIDE_ENVIRONMENT_VARIABLE_NAME_PREFIX + valve_name ) if override is None: continue try: if type(valve_value) is type(True): assert override.lower() in ( "true", "false", ), 'Value must be "true" or "false"' override = override.lower() == "true" elif type(valve_value) is type(42): override = int(override) elif type(valve_value) is type(""): pass else: valve_value_type = type(valve_value) raise ValueError(f"Unknown valve type: {valve_value_type}") except Exception as e: raise ValueError( f"Valve override {self.valves._VALVE_OVERRIDE_ENVIRONMENT_VARIABLE_NAME_PREFIX}{valve_name}={valve_value}: bad value: {e}" ) else: setattr(self.valves, valve_name, override) async def action( self, body: dict, __event_emitter__: typing.Callable[[dict], typing.Any] = None, __id__: typing.Optional[str] = None, __user__: typing.Optional[dict] = None, ) -> typing.Optional[dict]: valves = self.valves debug = valves.DEBUG emitter = EventEmitter(__event_emitter__, debug=debug) update_check_error = None update_check_notice = "" if valves.CHECK_FOR_UPDATES: if UpdateCheck.need_check(): await emitter.status("Checking for updates...") try: newer_version = UpdateCheck.get_newer_version() except UpdateCheck.VersionCheckError as e: update_check_error = e update_check_notice = ( f"\n\n(Failed to check for update to code execution function: {e})" ) else: if newer_version is not None: await emitter.status(f"New version found: {newer_version}") emitter.set_status_prefix(f"[Update available: {newer_version}] ") update_check_notice = f"\n\n(Code execution function update available: [{newer_version}]({UpdateCheck.USER_URL}))" storage = self._UserStorage( storage_root_path=os.path.join( valves.WEB_ACCESSIBLE_DIRECTORY_PATH, "user_files" ), storage_root_url=valves.WEB_ACCESSIBLE_DIRECTORY_URL.rstrip("/") + "/user_files", __user__=__user__, max_files_per_user=valves.MAX_FILES_PER_USER, max_bytes_per_user=valves.MAX_MEGABYTES_PER_USER * 1024 * 1024, ) async def _fail(error_message, status="SANDBOX_ERROR"): if debug: await emitter.fail( f"[DEBUG MODE] {error_message}; body={body}; valves=[{valves}]" ) elif update_check_error is not None: await emitter.fail(f"[{update_check_error}] {error_message}") else: await emitter.fail(error_message) return json.dumps({"status": status, "output": error_message}) await emitter.status("Checking messages for code blocks...") if len(body.get("messages", ())) == 0: return await _fail("No messages in conversation.", status="INVALID_INPUT") last_message = body["messages"][-1] if last_message["role"] != "assistant": return await _fail( "Last message was not from the AI model.", status="INVALID_INPUT" ) split_three_backticks = last_message["content"].split("```") if len(split_three_backticks) < 3: return await _fail( "Last message did not contain code blocks.", status="INVALID_INPUT" ) if len(split_three_backticks) % 2 != 1: return await _fail( "Last message did not contain well-formed code blocks.", status="INVALID_INPUT", ) chosen_code_block = None language = None for code_block in split_three_backticks[-2:0:-2]: if code_block.startswith("python\n") or code_block.startswith("python3\n"): chosen_code_block = code_block language = Sandbox.LANGUAGE_PYTHON if ( code_block.startswith("bash\n") or code_block.startswith("sh\n") or code_block.startswith("shell\n") ): chosen_code_block = code_block language = Sandbox.LANGUAGE_BASH break if chosen_code_block is None: # Try to see if the last code block looks like Python or bash. last_code_block = split_three_backticks[-2] # Look for an interpreter line. first_line = last_code_block.strip().split("\n")[0] if first_line.startswith("#!") and ( first_line.endswith("python") or first_line.endswith("python3") ): chosen_code_block = code_block language = Sandbox.LANGUAGE_PYTHON elif first_line.startswith("#!") and first_line.endswith("sh"): chosen_code_block = code_block language = Sandbox.LANGUAGE_BASH elif any( python_like in last_code_block for python_like in ("import ", "print(", "print ") ): chosen_code_block = code_block language = Sandbox.LANGUAGE_PYTHON elif any( bash_like in last_code_block for bash_like in ("echo ", "if [", "; do", "esac\n") ): chosen_code_block = code_block language = Sandbox.LANGUAGE_BASH if chosen_code_block is None: return await _fail( "Message does not contain code blocks detected as Python or Bash." ) try: max_ram_bytes = None if self.valves.MAX_RAM_MEGABYTES != 0: max_ram_bytes = self.valves.MAX_RAM_MEGABYTES * 1024 * 1024 await emitter.status("Checking if environment supports sandboxing...") Sandbox.check_setup( language=language, auto_install_allowed=self.valves.AUTO_INSTALL, require_resource_limiting=self.valves.REQUIRE_RESOURCE_LIMITING, ) if self.valves.AUTO_INSTALL and Sandbox.runsc_needs_installation(): await emitter.status("Auto-installing gVisor...") Sandbox.install_runsc() await emitter.status("Initializing sandbox configuration...") status = "UNKNOWN" output = None generated_files = [] language_title = language.title() # If the provided code starts/ends with "```SOME_LANGUAGE", remove that. code = chosen_code_block if language == Sandbox.LANGUAGE_PYTHON: code = code.removeprefix("python3") code = code.removeprefix("python") elif language == Sandbox.LANGUAGE_BASH: code = code.removeprefix("shell") code = code.removeprefix("bash") code = code.removeprefix("sh") code = code.strip() with tempfile.TemporaryDirectory(prefix="sandbox_") as tmp_dir: sandbox_storage_path = os.path.join(tmp_dir, "storage") os.makedirs(sandbox_storage_path, mode=0o777) sandbox = Sandbox( tmp_dir=tmp_dir, language=language, code=code, debug=debug, networking_allowed=valves.NETWORKING_ALLOWED, max_runtime_seconds=valves.MAX_RUNTIME_SECONDS, max_ram_bytes=max_ram_bytes, require_resource_limiting=valves.REQUIRE_RESOURCE_LIMITING, persistent_home_dir=sandbox_storage_path, ) await emitter.status( f"Running {language_title} code in gVisor sandbox..." ) try: result = sandbox.run() except Sandbox.ExecutionTimeoutError as e: await emitter.fail( f"Code timed out after {valves.MAX_RUNTIME_SECONDS} seconds" ) status = "TIMEOUT" output = e.stderr except Sandbox.InterruptedExecutionError as e: await emitter.fail("Code used too many resources") status = "INTERRUPTED" output = e.stderr except Sandbox.CodeExecutionError as e: await emitter.fail(f"{language_title}: {e}") status = "ERROR" output = e.stderr else: status = "OK" output = result.stdout or result.stderr output_num_files, _ = storage.measure_directory( sandbox_storage_path ) if output_num_files > valves.MAX_FILES_PER_EXECUTION: output = f"Code produced {output_num_files} files, exceeding per-execution quota of {valves.MAX_FILES_PER_EXECUTION}" await emitter.fail(output) status = "STORAGE_ERROR" elif output_num_files > 0: try: with storage: generated_files = storage.copy( __id__=__id__, intake_path=sandbox_storage_path, ) except self._UserStorage.OutOfStorageException as e: status = "STORAGE_ERROR" output = f"Storage quota exceeded: {e}" await emitter.fail(output) if status == "OK": await emitter.status( status="complete", done=True, description=f"{language_title} code executed successfully.", ) if output: output = output.strip() if debug: def _log(filename: str, log_line: str): print(f"[{filename}] {log_line}", file=sys.stderr) sandbox.debug_logs(_log) if status == "OK": generated_files_output = "" if len(generated_files) > 0: generated_files_output = "* " + "\n* ".join( f.markdown() for f in sorted(generated_files, key=lambda f: f.name) ) if output and len(generated_files) > 0: await emitter.message( f"\n\n---\nI executed this {language_title} code and got:\n```Output\n{output}\n```\n**Files**:\n{generated_files_output}{update_check_notice}" ) elif output and len(generated_files) == 0: await emitter.message( f"\n\n---\nI executed this {language_title} code and got:\n```Output\n{output}\n```{update_check_notice}" ) elif len(generated_files) > 0: await emitter.message( f"\n\n---\nI executed this {language_title} code and it generated these files:\n{generated_files_output}{update_check_notice}" ) else: await emitter.message( f"\n\n---\nI executed this {language_title} code and it ran successfully with no output.{update_check_notice}" ) return json.dumps( { "status": status, "output": output, "generated_files": { f.name: f.markdown() for f in generated_files }, } ) if status == "TIMEOUT": if output: await emitter.message( f"\n\n---\nI executed this {language_title} code and it timed out after {self.valves.MAX_RUNTIME_SECONDS} seconds:\n```Error\n{output}\n```\n{update_check_notice}" ) else: await emitter.message( f"\n\n---\nI executed this {language_title} code and it timed out after {self.valves.MAX_RUNTIME_SECONDS} seconds.\n{update_check_notice}" ) if status == "ERROR" and output: await emitter.message( f"\n\n---\nI executed this {language_title} code and got the following error:\n```Error\n{output}\n```\n{update_check_notice}" ) else: await emitter.message( f"\n\n---\nI executed this {language_title} code but got an unexplained error.\n{update_check_notice}" ) return json.dumps({"status": status, "output": output}) except Sandbox.PlatformNotSupportedException as e: return await _fail(f"Sandbox cannot run on this machine: {e}") except Sandbox.SandboxRuntimeException as e: return await _fail(f"Sandbox runtime failed: {e}") except Sandbox.FixableException as e: return await _fail(f"Environment needs setup work: {e}") except Sandbox.SandboxException as e: return await _fail(f"Sandbox exception: {e}") except Exception as e: return await _fail(f"Unhandled exception: {e}") class _UserFile: MAX_INLINE_URL_SIZE = 0 MAX_INLINE_TEXT_LINES = 128 MAX_INLINE_TEXT_BYTES = 65535 def __init__(self, file_path, file_relative_path, file_url, file_size): self._file_path = file_path self._file_relative_path = file_relative_path self._file_url = file_url self._size_bytes = file_size mimetypes.init() mime_type, _ = mimetypes.guess_type(file_path) if mime_type is None: mime_type, _ = mimetypes.guess_type(file_url) if mime_type is None: # Check if the file is valid UTF-8 text. is_utf8 = True is_empty = True with open(self._file_path, "rb") as f: for line in f: is_empty = is_empty or len(line) > 0 try: line.decode("utf-8") except UnicodeDecodeError: is_utf8 = False mime_type = ( "text/plain" if (is_utf8 and not is_empty) else "application/octet-stream" ) self._mime_type = mime_type self._cached_markdown = None @property def name(self): return self._file_relative_path @property def url(self): if self._size_bytes > 0 and self._size_bytes <= self.MAX_INLINE_URL_SIZE: # Try to use an inline URL. with open(self._file_path, "rb") as f: contents = f.read() inline_url = f"data:{self._mime_type}," + urllib.parse.quote_from_bytes( contents ) inline_base64 = ( f"data:{self._mime_type};base64," + base64.standard_b64encode(contents).decode("ascii") ) shortest_url = ( inline_url if len(inline_url) < len(inline_base64) else inline_base64 ) if len(shortest_url) <= self.MAX_INLINE_URL_SIZE: return shortest_url return self._file_url def _inline_markdown(self): """Render the file as inline text or markdown if small enough; otherwise return None.""" if self._size_bytes > self.MAX_INLINE_TEXT_BYTES: return None if not self._mime_type.startswith("text/"): return None with open(self._file_path, "rb") as f: try: contents = f.read().decode("utf-8") except UnicodeDecodeError: return None lines = contents.split("\n") if len(lines) > self.MAX_INLINE_TEXT_LINES: return None if self._mime_type != "text/markdown": if "```" in contents: return None if contents and contents[-1] == "\n": contents = contents[:-1] return f"\U0001f4c4 [{self.name}]({self.url}):\n```\n{contents}\n```" components = [f"\U0001f4c3 [{self.name}]({self.url}):"] for line in lines: components.append(f"> {line}") if components[-1] == "> ": components = components[:-1] return "\n".join(components) def _markdown(self): if self._size_bytes == 0: return f"\u2049 `{self.name}` (empty)" inline_markdown = self._inline_markdown() if inline_markdown is not None: return inline_markdown icon = "\U0001f4be" if self._mime_type.startswith("text/"): icon = "\U0001f4c4" elif self._mime_type.startswith("image/"): icon = "\U0001f5bc" elif self._mime_type.startswith("audio/"): icon = "\U0001f3b5" elif self._mime_type.startswith("video/"): icon = "\U0001f3ac" size = f"{self._size_bytes} bytes" if self._size_bytes > 1024 * 1024 * 1024: size = f"{self._size_bytes // 1024 // 1024 // 1024} GiB" elif self._size_bytes > 1024 * 1024: size = f"{self._size_bytes // 1024 // 1024} MiB" elif self._size_bytes > 1024: size = f"{self._size_bytes // 1024} KiB" return f"{icon} [{self.name}]({self.url}) ({size})" def markdown(self): if self._cached_markdown is None: self._cached_markdown = self._markdown() return self._cached_markdown class _UserStorage: class StorageException(Exception): """Base class for storage-related exceptions.""" class OutOfStorageException(Exception): """Not enough files or bytes quota.""" class EnvironmentNeedsSetupException(Exception): """Storage is badly configured.""" # Number of zeroes to use in nonces (in per-day storage directories). # This effectively limits the number of nonces usable in a given day. NUM_NONCE_ZEROS = 4 # Free space buffer to keep free on the underlying storage device, # rather than allowing user storage to fill it to the brim. MUST_KEEP_FREE_MARGIN_MEGABYTES = 512 @classmethod def measure_directory(cls, path, predicate=None): """ Measure storage cost of a directory. :param path: Path to the directory to measure. :param predicate: Optional predicate to filter files and directories, called with absolute paths. :return: 2-tuple `(total_files, total_bytes)`. Note that `total_files` counts the number of non-root directories as well, and `total_bytes` also includes storage necessary to store filenames and directory names. """ path = os.path.normpath(os.path.abspath(path)) total_files = 0 total_bytes = 0 try: for dirpath, subdirs, subfiles in os.walk( path, onerror=None, followlinks=False ): dirpath = os.path.normpath(os.path.abspath(dirpath)) for subdir in subdirs: if predicate is None or predicate( os.path.join(dirpath, subdir) ): total_files += 1 total_bytes += len(subdir) for subfile in subfiles: subfile_path = os.path.join(dirpath, subfile) if predicate is not None and not predicate(subfile_path): continue try: subfile_stat = os.stat(subfile_path, follow_symlinks=False) except FileNotFoundError: continue # Likely raced with another execution. if not stat.S_ISREG(subfile_stat.st_mode): continue # Ignore non-regular files. total_files += 1 total_bytes += len(subfile) total_bytes += subfile_stat.st_size except OSError as e: raise cls.EnvironmentNeedsSetupException( f"Failed to explore directory {path} (please adjust permissions): {e}" ) return total_files, total_bytes def __init__( self, storage_root_path, storage_root_url, __user__: typing.Optional[dict] = None, max_files_per_user=None, max_bytes_per_user=None, ): if storage_root_path.startswith("$DATA_DIR" + os.sep): if "DATA_DIR" not in os.environ: data_dir = "/app/backend/data" if not os.path.isdir(data_dir): if os.path.isdir("/app/backend"): os.makedirs(data_dir, mode=0o755) else: raise self.EnvironmentNeedsSetupException( f"DATA_DIR specified in user storage configuration ({storage_root_path}), but not specified in environment, and default path '/app/backend/data' does not exist; please create it or configure user storage directory." ) else: data_dir = os.environ["DATA_DIR"] storage_root_path = os.path.join( data_dir, storage_root_path[len("$DATA_DIR" + os.sep) :].lstrip(os.sep), ) self._storage_root_path = os.path.normpath( os.path.abspath(storage_root_path) ) try: os.makedirs(self._storage_root_path, mode=0o755, exist_ok=True) except OSError as e: raise self.EnvironmentNeedsSetupException( f"User storage directory ({self._storage_root_path}) does not exist and cannot automatically create it ({e}); please create it or reconfigure it." ) self._storage_root_url = storage_root_url.rstrip("/") self._date = time.strftime("%Y/%m/%d") self._max_files_per_user = max_files_per_user self._max_bytes_per_user = max_bytes_per_user self._user = f"anon_{self._date}" if __user__ is not None: if type(__user__) is type({}): self._user = str( "|".join( f"{k}={v}" for k, v in sorted(__user__.items(), key=lambda x: x[0]) ) ) else: self._user = str(__user__) user_hash = hashlib.sha512() user_hash.update(self._user.encode("utf-8")) self._user_hash = ( base64.b32encode(user_hash.digest()).decode("ascii").lower()[:12] ) self._user_path = os.path.join(storage_root_path, self._user_hash) self._lock_fd = None def __enter__(self): assert self._lock_fd is None os.makedirs(self._user_path, mode=0o755, exist_ok=True) lock_fd = os.open( os.path.join(self._user_path, ".lock"), os.O_RDWR | os.O_CREAT | os.O_TRUNC, ) deadline = time.time() + 10 last_exception = None while time.time() < deadline: try: fcntl.flock(lock_fd, fcntl.LOCK_EX | fcntl.LOCK_NB) except (IOError, OSError) as e: last_exception = e time.sleep(0.01) else: self._lock_fd = lock_fd break if self._lock_fd is None: os.close(lock_fd) raise self.StorageException( f"Cannot lock storage directory (too many concurrent code executions?) {last_exception}" ) def __exit__(self, *args, **kwargs): assert self._lock_fd is not None fcntl.flock(self._lock_fd, fcntl.LOCK_UN) os.close(self._lock_fd) self._lock_fd = None def _is_user_file(self, path): """Used as predicate when measuring user storage directories.""" assert path.startswith(self._user_path + os.sep) path = path[len(self._user_path) + len(os.sep) :] # User files are under: # YYYY/MM/DD/NONCE/HASH (5 levels of nesting). # So a real user file has at least 6 components, which means it # must have at least 5 slashes. return path.count(os.sep) >= 5 def copy(self, __id__, intake_path): """ Copy a directory to user storage. Ensure that the user storage has room for the given number of files totaling the given number of bytes. If this is feasible by deleting previous user files, it will do so. Must be done while holding the lock. :param __id__: Chat or action ID. :param intake_path: Path to a directory that should be copied to the user storage. :raises OutOfStorageException: If there is not enough available file or bytes quota. :return: A list of `_UserFile`s from each file copied from `intake_path`. """ assert ( self._lock_fd is not None ), "Cannot perform this operation without holding the lock" want_num_files, want_num_bytes = self.measure_directory(intake_path) if want_num_files == 0: return () # Nothing to copy. if self._max_files_per_user <= want_num_files: raise self.OutOfStorageException( f"Cannot allocate storage for {want_num_files} files; maximum is {self._max_files_per_user} files per user" ) if self._max_bytes_per_user <= want_num_bytes: raise self.OutOfStorageException( f"Cannot allocate storage for {want_num_bytes} bytes; maximum is {self._max_bytes_per_user} bytes per user" ) disk_usage_free = shutil.disk_usage(self._user_path).free if ( disk_usage_free <= want_num_bytes + self.MUST_KEEP_FREE_MARGIN_MEGABYTES * 1024 * 1024 ): raise self.OutOfStorageException( f"Not enough free disk space for {want_num_bytes} bytes; current free space is {disk_usage_free} bytes and must keep at least {self.MUST_KEEP_FREE_MARGIN_MEGABYTES} megabytes free" ) user_root_num_files, user_root_num_bytes = self.measure_directory( self._user_path, predicate=self._is_user_file, ) user_root_remaining_files = self._max_files_per_user - user_root_num_files user_root_remaining_bytes = self._max_bytes_per_user - user_root_num_bytes while ( user_root_remaining_files < want_num_files or user_root_remaining_bytes < want_num_bytes ): oldest_directory = None try: oldest_yyyy = next( iter( sorted( f for f in os.listdir(self._user_path) if len(f) >= 4 and f.isdigit() ) ) ) oldest_mm = next( iter( sorted( f for f in os.listdir( os.path.join(self._user_path, oldest_yyyy) ) if len(f) == 2 and f.isdigit() ) ) ) oldest_dd = next( iter( sorted( f for f in os.listdir( os.path.join( self._user_path, oldest_yyyy, oldest_mm ) ) if len(f) == 2 and f.isdigit() ) ) ) oldest_nonce = next( iter( sorted( f for f in os.listdir( os.path.join( self._user_path, oldest_yyyy, oldest_mm, oldest_dd, ) ) if len(f) == self.NUM_NONCE_ZEROS and f.isdigit() ) ) ) oldest_directory = os.path.join( self._user_path, oldest_yyyy, oldest_mm, oldest_dd, oldest_nonce ) except StopIteration: raise self.OutOfStorageException( f"Cannot find directory to clear in order to make enough room for new user storage ({want_num_files} files, {want_num_bytes} bytes)" ) assert oldest_directory is not None if not shutil.rmtree.avoids_symlink_attacks: raise self.EnvironmentNeedsSetupException( "Only supported on platforms with symlink-attack-resistant rmtree implementations" ) shutil.rmtree(oldest_directory) for parent_directory in ( os.path.join(self._user_path, oldest_yyyy, oldest_mm, oldest_dd), os.path.join(self._user_path, oldest_yyyy, oldest_mm), os.path.join(self._user_path, oldest_yyyy), ): if len(os.listdir(parent_directory)) == 0: os.rmdir(parent_directory) user_root_num_files, user_root_num_bytes = self.measure_directory( self._user_path, predicate=self._is_user_file, ) user_root_remaining_files = ( self._max_files_per_user - user_root_num_files ) user_root_remaining_bytes = ( self._max_bytes_per_user - user_root_num_bytes ) # We now have enough. Find new directory name. path_with_counter = None max_nonce = 10**self.NUM_NONCE_ZEROS - 1 for nonce in range( 1, min(self._max_files_per_user or max_nonce, max_nonce) ): path_with_counter = os.path.join( self._user_path, self._date, str(nonce).zfill(self.NUM_NONCE_ZEROS) ) try: os.makedirs(path_with_counter, mode=0o755, exist_ok=False) except FileExistsError: pass else: break if path_with_counter is None: raise self.OutOfStorageException("No free storage directory available!") id_str = str(__id__) if __id__ is not None else self._date id_hash = hashlib.sha512() id_hash.update(self._user.encode("utf-8")) id_hash.update(b"||||") id_hash.update(self._date.encode("utf-8")) id_hash.update(b"||||") id_hash.update(path_with_counter.encode("utf-8")) id_hash.update(b"||||") id_hash.update(str(uuid.uuid4()).encode("utf-8")) id_hash.update(b"||||") id_hash.update(id_str.encode("utf-8")) id_hash_component = ( base64.b32encode(id_hash.digest()).decode("ascii").lower()[:12] ) final_path = os.path.normpath( os.path.abspath(os.path.join(path_with_counter, id_hash_component)) ) # Now do the copy. # This doesn't use `shutil.copytree` because we explicitly avoid copying anything but regular files. user_files = [] for dirpath, subdirs, subfiles in os.walk( intake_path, onerror=None, followlinks=False ): dirpath = os.path.normpath(os.path.abspath(dirpath)) relative_dirpath = None if dirpath == intake_path: relative_dirpath = "." elif dirpath.startswith(intake_path + os.sep): relative_dirpath = dirpath[len(intake_path) + len(os.sep) :] else: assert False, f"Bad traversal: expected all paths to starts with {intake_path} but got path that does not: {dirpath}" assert relative_dirpath is not None assert not os.path.isabs(relative_dirpath) copy_dirpath = os.path.join(final_path, relative_dirpath) os.makedirs(copy_dirpath, mode=0o755, exist_ok=True) for subfile in subfiles: subfile_path = os.path.join(dirpath, subfile) subfile_relative_path = os.path.normpath( os.path.join(relative_dirpath, subfile) ) assert not os.path.isabs(subfile_relative_path) subfile_stat = os.stat(subfile_path, follow_symlinks=False) if not stat.S_ISREG(subfile_stat.st_mode): continue # Ignore non-regular files. subfile_copy = os.path.join(copy_dirpath, subfile) assert subfile_copy.startswith(self._storage_root_path + os.sep) subfile_url = f"{self._storage_root_url}/" + urllib.parse.quote( subfile_copy[len(self._storage_root_path) + len(os.sep) :], safe=os.sep, ) shutil.move(subfile_path, subfile_copy, copy_function=shutil.copy) user_files.append( _Action._UserFile( file_path=subfile_copy, file_relative_path=subfile_relative_path, file_url=subfile_url, file_size=subfile_stat.st_size, ) ) # We are done. return user_files class Action: Valves = _Action.Valves def __init__(self): self.valves = self.Valves() async def action( self, body: dict, __event_emitter__: typing.Callable[[dict], typing.Any] = None, ) -> typing.Optional[dict]: return await _Action(self.valves).action( body=body, __event_emitter__=__event_emitter__ ) class EventEmitter: """ Helper wrapper for OpenWebUI event emissions. """ def __init__( self, event_emitter: typing.Callable[[dict], typing.Any] = None, debug: bool = False, ): self.event_emitter = event_emitter self._debug = debug self._status_prefix = None def set_status_prefix(self, status_prefix): self._status_prefix = status_prefix async def _emit(self, typ, data): if self._debug: print(f"Emitting {typ} event: {data}", file=sys.stderr) if not self.event_emitter: return None maybe_future = self.event_emitter( { "type": typ, "data": data, } ) if asyncio.isfuture(maybe_future) or inspect.isawaitable(maybe_future): return await maybe_future async def status( self, description="Unknown state", status="in_progress", done=False ): if self._status_prefix is not None: description = f"{self._status_prefix}{description}" await self._emit( "status", { "status": status, "description": description, "done": done, }, ) if not done and len(description) <= 1024: # Emit it again; Open WebUI does not seem to flush this reliably. # Only do it for relatively small statuses; when debug mode is enabled, # this can take up a lot of space. await self._emit( "status", { "status": status, "description": description, "done": done, }, ) async def fail(self, description="Unknown error"): await self.status(description=description, status="error", done=True) async def message(self, content): await self._emit( "message", { "content": content, }, ) async def citation(self, document, metadata, source): await self._emit( "citation", { "document": document, "metadata": metadata, "source": source, }, ) async def code_execution_result(self, output): await self._emit( "code_execution_result", { "output": output, }, ) class Sandbox: """ Sandbox manages a gVisor sandbox's lifecycle. """ # Set of supported programming languages. LANGUAGE_PYTHON = "python" LANGUAGE_BASH = "bash" SUPPORTED_LANGUAGES = [LANGUAGE_PYTHON, LANGUAGE_BASH] # The following directories will be exposed as read-only to the # sandboxed environment. This must contain at least the necessary # files and libraries necessary to run the code interpreter. # Subdirectories of these directories may be hidden by adding them # to the `EMPTY_READ_ONLY_DIRECTORIES` or `EMPTY_WRITABLE_DIRECTORIES` # lists below. EXPOSED_SYSTEM_DIRECTORIES = [ "/bin", "/etc/alternatives", "/etc/ssl/certs", "/lib", "/lib32", "/lib64", "/opt", "/sbin", "/usr", "/var/lib", ] # The following files will be exposed as read-only to the sandboxed # environment. This should contain the set of files necessary by the # code interpreter to function correctly, e.g. `/etc/resolv.conf` # is necessary to properly resolve hosts through DNS. EXPOSED_SYSTEM_FILES = [ "/etc/hosts", "/etc/localtime", "/etc/mime.types", "/etc/nsswitch.conf", "/etc/os-release", "/etc/resolv.conf", "/etc/shells", ] # The following directories will exist in the sandbox environment but # will appear as empty and read-only. # This is useful to have a filesystem that feels like a normal Linux # environment without actually revealing these directories to the # sandbox. EMPTY_READ_ONLY_DIRECTORIES = [ "/etc", "/home", "/lost+found", "/root", "/run", "/run/user", "/sys", "/var", ] # The following directories will exist in the sandbox environment but # will appear as empty and writable. # This is useful to have a filesystem that feels like a normal Linux # environment without actually revealing these directories to the # sandbox. EMPTY_WRITABLE_DIRECTORIES = [ "/dev/shm", "/home/user", "/run/user/1000", "/var/run", "/var/tmp", "/tmp", ] # Static parts of the OCI configuration. OCI_CONFIG_SKELETON = { "ociVersion": "1.0.0", "process": { "user": {"uid": 1000, "gid": 1000}, "args": ["/bin/INVALID"], # Will be filled in. "env": [ # Basic environment variables. "EDITOR=cat", "LANG=C.UTF-8", "LC_ALL=C.UTF-8", "LC_CTYPE=C.UTF-8", "HOME=/home/user", "HOSTNAME=sandbox", "PAGER=cat", "PATH=/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin", "PWD=/home/user", "SHLVL=1", "TERM=xterm", "USER=user", "TZ=" + time.tzname[0], ], "cwd": "/home/user", "capabilities": { # No capabilities whatsoever. "bounding": [], "effective": [], "inheritable": [], "permitted": [], }, "rlimits": [ {"type": "RLIMIT_NOFILE", "hard": 1048576, "soft": 1048576}, ], "noNewPrivileges": True, }, "root": { "path": "/invalid", # Will be filled in. "readonly": True, }, "hostname": "sandbox", "mounts": [ {"destination": "/dev", "type": "dev"}, {"destination": "/proc", "type": "proc"}, ], "linux": { "namespaces": [ {"type": "pid"}, {"type": "ipc"}, {"type": "uts"}, {"type": "mount"}, ], "resources": { "memory": { # `limit` may be be filled in here depending on user configuration. "disableOOMKiller": False, }, }, }, } # The path where the `runsc` binary will be downloaded and installed if # requested. AUTO_INSTALLATION_PATH = "/tmp/gvisor/runsc" # Regular expression for log filename prefixes generated by `runsc`. _LOG_FILENAME_TRUNCATE_RE = re.compile(r"^runsc\.log\.\d{8}-\d{6}(?:\.\d+)?\.") # Other files worth logging when dumping debug logs. _EXTRA_DEBUG_LOG_PATHS = ( "/etc/os-release", "/proc/self/cgroup", "/proc/self/personality", "/proc/self/mountinfo", "/proc/self/setgroups", "/proc/self/status", "/proc/self/uid_map", "/proc/self/gid_map", "/proc/cmdline", "/proc/cpuinfo", "/proc/cgroups", "/proc/mounts", "/proc/version", ) # Other commands worth running when dumping debug logs. _EXTRA_DEBUG_LOG_COMMANDS = ( ("pwd",), ("id",), ("uname", "-a"), ("ls", "-l", "/proc/self/ns"), ("findmnt",), (sys.executable, "--version"), ) # Environment variable used to detect interpreter re-execution. _MARKER_ENVIRONMENT_VARIABLE = "__CODE_EXECUTION_STAGE" # libc bindings. # Populated using `_libc`. _LIBC = None class _Libc: """ Wrapper over libc functions. """ def __init__(self): libc = ctypes.CDLL(ctypes.util.find_library("c"), use_errno=True) libc.mount.argtypes = (ctypes.c_char_p,) self._libc = libc def mount(self, source, target, fs, options): if ( self._libc.mount( source.encode("ascii"), target.encode("ascii"), fs.encode("ascii"), 0, options.encode("ascii"), ) < 0 ): errno = ctypes.get_errno() raise OSError( errno, f"mount({source}, {target}, {fs}, {options}): {os.strerror(errno)}", ) def umount(self, path): if self._libc.umount(path.encode("ascii")) < 0: errno = ctypes.get_errno() raise OSError(errno, f"umount({path}): {os.strerror(errno)}") def unshare(self, flags): if self._libc.unshare(flags) < 0: raise OSError(f"unshare({flags}) failed") class _SelfFile: """ Manages a copy of this file's own contents. """ _CONTENTS = None @classmethod def init(cls): """ Read `__file__` into `cls._CONTENTS`. Must be called during init. """ if cls._CONTENTS is None: with open(__file__, "rb") as self_f: cls._CONTENTS = self_f.read().decode("ascii") @classmethod def contents(cls) -> str: """ Return this file's own contents. """ assert cls._CONTENTS is not None, f"{cls.__name__}.init not called" return cls._CONTENTS class _Switcheroo: """ Management of the switcheroo procedure for running in a usable cgroup namespace and node. """ _CGROUP_ROOT = "/sys/fs/cgroup" _CGROUP_NAME_PREFIX = "codeeval_" _CGROUP_MAX_COUNT = 4096 _CGROUP_SANDBOX_NAME = "sandbox" _CGROUP_SUPERVISOR_NAME = "supervisor" _CGROUP_LEAF = "leaf" def __init__(self, libc, log_path, max_sandbox_ram_bytes, do_resource_limiting): self._libc = libc self._log_path = log_path self._max_sandbox_ram_bytes = max_sandbox_ram_bytes self._do_resource_limiting = do_resource_limiting self._my_euid = None self._my_egid = None self._checkpoint = None self._cgroup_controllers = None self._needed_controllers = set() if max_sandbox_ram_bytes is not None: self._needed_controllers.add("memory") self._initial_cgroup_name = None self._codeeval_cgroup_name = None self._moved = False self._operations = [ # Save EUID and EGID before we move to a new user namespace. ("save_euid", self._save_euid), ("save_egid", self._save_egid), ("unshare_user", self._unshare_user), # Map our current user as being root in the new user namespace. ("write_uid_map", self._write_uid_map), ("write_setgroups", self._write_setgroups), ("write_gid_map", self._write_gid_map), ] if do_resource_limiting: self._operations.extend( ( # cgroupfs's view does not take into account cgroup namespaces. # Weird, right? # This means `/proc/PID/cgroup` will show the namespaced view of # the cgroup that the PID is in, but `/sys/fs/cgroup` will still # contain the whole system cgroup hierarchy regardless of namespace. # Instead, namespaces act as "boundary box" around process movement # requests when writing to cgroup.procs or creating new cgroups. # So our first order of business here is to find out which cgroup we # are running in. We do this by scanning the whole cgroupfs hierarchy # and looking for our PID. This will populate # `self._initial_cgroup_name`. ( "find_self_in_cgroup_hierarchy", self._find_self_in_cgroup_hierarchy, ), # The cgroup nesting rules are complicated, but the short of it is: # A cgroup can either **contain processes** OR **have limits**. # Also, cgroups that contain processes must be leaf nodes. # Also, cgroups that enforce limits must have their parent cgroup # also have the same limit "controller" be active. # So we will have two types of cgroups: # - Leaf nodes with no controllers # - Non-leaf nodes with controllers # So initially, all the processes in the container's initial # namespace need to be moved out to a new leaf node, # otherwise we cannot turn on controllers on the initial # cgroup. # So we will set up the following hierarchy: # /sys/fs/cgroup/$INITIAL: # The cgroup where the container's processes were running # the first time we run any Sandbox in the container. # It may initially have no controllers enabled, but we will # turn them on later. # /sys/fs/cgroup/$INITIAL/leaf: # The cgroup where the container's processes are moved to # from the $INITIAL cgroup upon first run of any Sandbox in # this container. When this code runs again, processes that # are already in `$INITIAL/leaf` are not moved. # /sys/fs/cgroup/$INITIAL/codeeval_$NUM: # A per-Sandbox cgroup that never contains any processes. # It will have controllers enabled on it but will never have # specific limits enforced. # /sys/fs/cgroup/$INITIAL/codeeval_$NUM/sandbox: # A per-Sandbox cgroup that never contains any processes. # It will have controllers enabled on it and will enforce # resource limits for the processes running in its /leaf. # /sys/fs/cgroup/$INITIAL/codeeval_$NUM/sandbox/leaf: # A per-Sandbox cgroup that is running `runsc` (gVisor). # It has no controllers enabled on it, but resources are # being enforced by virtue of being a child of # `$INITIAL/codeeval_$NUM/sandbox` which does enforce limits. # /sys/fs/cgroup/$INITIAL/codeeval_$NUM/supervisor: # A per-Sandbox cgroup that never contains any processes. # It will have controllers enabled on it and will enforce # resource limits for the processes running in its /leaf. # /sys/fs/cgroup/$INITIAL/codeeval_$NUM/supervisor/leaf: # A per-Sandbox cgroup that is running a Python interpreter # that manages the lifetime of the `runsc` process. # It will run `Sandbox.maybe_main`. # It has no controllers enabled on it, but resources are # being enforced by virtue of being a child of # `$INITIAL/codeeval_$NUM/sandbox` which does enforce limits. # # This particular step creates the `$INITIAL/leaf` cgroup. # If already created, it does nothing. ( "create_initial_leaf_cgroup", self._create_initial_leaf_cgroup, ), # Move all processes in `$INITIAL` to `$INITIAL/leaf`. ( "move_initial_cgroup_processes_to_initial_leaf_cgroup", self._move_initial_cgroup_processes_to_initial_leaf_cgroup, ), # Read the cgroup controllers enabled in `$INITIAL`. This acts # as a bounding set on the ones we can enable in any child of it. ("read_cgroup_controllers", self._read_cgroup_controllers), # Cleanup old `$INITIAL/codeeval_*` cgroups that may be lying # around from past runs. ("cleanup_old_cgroups", self._cleanup_old_cgroups), # Create a new `$INITIAL/codeeval_$NUM` cgroup. ("create_codeeval_cgroup", self._create_codeeval_cgroup), # Create a new `$INITIAL/codeeval_$NUM/sandbox` cgroup. ("create_sandbox_cgroup", self._create_sandbox_cgroup), # Create a new `$INITIAL/codeeval_$NUM/sandbox/leaf` cgroup. ( "create_sandbox_leaf_cgroup", self._create_sandbox_leaf_cgroup, ), # Create a new `$INITIAL/codeeval_$NUM/supervisor` cgroup. ("create_supervisor_cgroup", self._create_supervisor_cgroup), # Create a new `$INITIAL/codeeval_$NUM/supervisor/leaf` cgroup. ( "create_supervisor_leaf_cgroup", self._create_supervisor_leaf_cgroup, ), # Add controllers to `$INITIAL`. ( "add_cgroup_controllers_to_root", self._add_cgroup_controllers_to_root, ), # Add controllers to `$INITIAL/codeeval_$NUM`. ( "add_cgroup_controllers_to_codeeval", self._add_cgroup_controllers_to_codeeval, ), # Add controllers to `$INITIAL/codeeval_$NUM/sandbox`. ( "add_cgroup_controllers_to_sandbox", self._add_cgroup_controllers_to_sandbox, ), # Set resource limits on `$INITIAL/codeeval_$NUM`. ("set_sandbox_cgroup_limits", self._set_sandbox_cgroup_limits), # Add controllers to `$INITIAL/codeeval_$NUM/supervisor`. ( "add_cgroup_controllers_to_supervisor", self._add_cgroup_controllers_to_supervisor, ), # Set resource limits on `$INITIAL/codeeval_$NUM/supervisor`. ( "set_supervisor_cgroup_limits", self._set_supervisor_cgroup_limits, ), # Move current process to # `$INITIAL/codeeval_$NUM/supervisor/leaf`. ( "move_process_to_supervisor_leaf", self._move_process_to_supervisor_leaf, ), # Double-check that we have moved to # `$INITIAL/codeeval_$NUM/supervisor/leaf`. ("sanity_check_own_cgroup", self._sanity_check_own_cgroup), ) ) def _status(self): """ Return the current switcheroo status. :return: The last successful operation name, "UNSTARTED" if unstarted, or "OK" if all done, and some information. """ main_status = self._checkpoint if self._checkpoint is None: main_status = "UNSTARTED" if self._checkpoint == self._operations[-1][0]: main_status = "OK" my_pid = os.getpid() status_line = f"{main_status} (euid={self._my_euid} egid={self._my_egid} pid={my_pid} do_resource_limiting={self._do_resource_limiting} initial_cgroup_name={self._initial_cgroup_name} codeeval_cgroup_name={self._codeeval_cgroup_name} controllers={self._cgroup_controllers})" if self._do_resource_limiting: cgroupfs_data = [] for cgroup_components in ( (self._initial_cgroup_name,), (self._initial_cgroup_name, self._CGROUP_LEAF), (self._initial_cgroup_name, self._codeeval_cgroup_name), ( self._initial_cgroup_name, self._codeeval_cgroup_name, self._CGROUP_LEAF, ), ( self._initial_cgroup_name, self._codeeval_cgroup_name, self._CGROUP_SUPERVISOR_NAME, ), ( self._initial_cgroup_name, self._codeeval_cgroup_name, self._CGROUP_SUPERVISOR_NAME, self._CGROUP_LEAF, ), ( self._initial_cgroup_name, self._codeeval_cgroup_name, self._CGROUP_SANDBOX_NAME, ), ( self._initial_cgroup_name, self._codeeval_cgroup_name, self._CGROUP_SANDBOX_NAME, self._CGROUP_LEAF, ), ): if any(c is None for c in cgroup_components): continue file_data = [] for filename in ("procs", "controllers", "subtree_control"): data = None try: with self._open( self._cgroup_path( *(cgroup_components + (f"cgroup.{filename}",)) ), "rb", ) as f: data = f.read().decode("ascii").replace("\n", " ") except Exception as e: data = f"[fail: {e}]" file_data.append(f"{filename}: {data}") cgroup_components_joined = os.sep.join(cgroup_components) file_data_joined = ", ".join(file_data) cgroupfs_data.append( f"{cgroup_components_joined}=[{file_data_joined}]" ) if len(cgroupfs_data) > 0: cgroupfs_data_joined = " ".join(cgroupfs_data) status_line += f" {cgroupfs_data_joined}" return status_line def _cgroup_path(self, *components): assert all( c is not None for c in components ), f"Tried to build cgroup path with not-yet-determined component: {components}" return os.path.join(self._CGROUP_ROOT, *(c for c in components if c)) def _log(self, log_f, message): """ Log a message to `log_f`. :param log_f: Log file object. """ timestamp = time.strftime("%H:%M:%S") status = self._status() log_f.write(f"[{timestamp}] {message} [{status}]\n".encode("utf-8")) def do(self): """ Do the switcheroo. :raises OSError: If anything goes wrong. Progress is saved. """ op_index = -1 for i, (op, _) in enumerate(self._operations): if self._checkpoint == op: op_index = i break with self._open(self._log_path, "ab") as log_f: def do_log(s): return self._log(log_f, s) for op, fn in self._operations[op_index + 1 :]: do_log(f"Starting operation: {op}") errors = [] success = False for attempt in range(1, 4): try: fn() except OSError as e: do_log(f"OSError #{attempt}: {op}: {e}") errors.append(OSError(f"OSError in {op} (#{attempt}): {e}")) except Exception as e: do_log(f"Exception #{attempt}: {op}: {e}") errors.append(OSError(f"{op} failed (#{attempt}): {e}")) else: success = True break time.sleep(0.1) if success: self._checkpoint = op do_log(f"Success: {op}") continue assert len(errors) > 0, "Logic error" first_exception = errors[0] if len(errors) == 1: raise first_exception other_exceptions = "; ".join(str(e) for e in errors[1:]) raise errors[0].__class__( f"{first_exception} (other attempts: {other_exceptions})" ) def _best_effort_remove_cgroup_subtree(self, codeeval_name): for cgroup_components in ( ( self._initial_cgroup_name, codeeval_name, self._CGROUP_SANDBOX_NAME, self._CGROUP_LEAF, ), ( self._initial_cgroup_name, codeeval_name, self._CGROUP_SUPERVISOR_NAME, self._CGROUP_LEAF, ), (self._initial_cgroup_name, codeeval_name, self._CGROUP_SANDBOX_NAME), ( self._initial_cgroup_name, codeeval_name, self._CGROUP_SUPERVISOR_NAME, ), (self._initial_cgroup_name, codeeval_name), ): try: os.rmdir(self._cgroup_path(*cgroup_components)) except OSError: pass def cleanup(self): if self._moved: self._move_process_back() if self._codeeval_cgroup_name is not None: self._best_effort_remove_cgroup_subtree(self._codeeval_cgroup_name) def _open(self, path, mode): try: return open(path, mode) except OSError as e: raise OSError(f"opening {path} mode={mode}: {e}") def _save_euid(self): self._my_euid = os.geteuid() def _save_egid(self): self._my_egid = os.getegid() def _unshare_user(self): Sandbox.unshare( os.CLONE_NEWUSER if "CLONE_NEWUSER" in os.__dict__ else 0x10000000 ) def _write_uid_map(self): with self._open("/proc/self/uid_map", "wb") as uid_map_f: uid_map_f.write(f"0 {self._my_euid} 1\n".encode("ascii")) def _write_setgroups(self): with self._open("/proc/self/setgroups", "wb") as setgroups_f: setgroups_f.write(b"deny") def _write_gid_map(self): with self._open("/proc/self/gid_map", "wb") as gid_map_f: gid_map_f.write(f"0 {self._my_egid} 1\n".encode("ascii")) def _find_self_in_cgroup_hierarchy(self): my_pid = os.getpid() cgroup_root_slash = self._CGROUP_ROOT + os.sep found_cgroup = None num_checked = 0 num_except = 0 sample_exception = None for dirpath, _, subfiles in os.walk( self._CGROUP_ROOT, onerror=None, followlinks=False ): if dirpath != self._CGROUP_ROOT and not dirpath.startswith( cgroup_root_slash ): continue if "cgroup.procs" not in subfiles: continue num_checked += 1 found_pid = False try: with self._open( os.path.join(dirpath, "cgroup.procs"), "rb" ) as cgroup_procs_f: for line in cgroup_procs_f: for pid_str in line.strip().split(b" "): if not pid_str: continue try: pid = int(pid_str) except ValueError: continue if pid == my_pid: found_pid = True break except Exception as e: num_except += 1 if sample_exception is None: sample_exception = e.__class__(f"{dirpath}: {e}") continue if not found_pid: continue current_cgroup = dirpath[len(cgroup_root_slash) :] if found_cgroup is not None: raise OSError( f"Found PID {my_pid} in two separate cgroups: {found_cgroup} and {current_cgroup}; racing with another process?" ) found_cgroup = current_cgroup if found_cgroup is None: raise OSError( f"PID {my_pid} could not be found in any cgroup (checked {num_checked} cgroups, {num_except} exceptions; sample: {sample_exception})" ) if found_cgroup.endswith(os.sep + self._CGROUP_LEAF): found_cgroup = found_cgroup[: -len(os.sep + self._CGROUP_LEAF)] self._initial_cgroup_name = found_cgroup def _read_cgroup_controllers(self): cgroup_controllers = [] with self._open( self._cgroup_path(self._initial_cgroup_name, "cgroup.controllers"), "rb" ) as cgroup_controllers_f: for line in cgroup_controllers_f: for controller in line.strip().split(b" "): if controller and controller not in cgroup_controllers: cgroup_controllers.append(controller.decode("ascii")) self._cgroup_controllers = cgroup_controllers def _cleanup_old_cgroups(self): initial_cgroup_path = self._cgroup_path(self._initial_cgroup_name) for filename in os.listdir(initial_cgroup_path): if not filename.startswith(self._CGROUP_NAME_PREFIX): continue cgroup_path = os.path.join(initial_cgroup_path, filename) if not os.path.isdir(cgroup_path): continue self._best_effort_remove_cgroup_subtree(filename) def _create_initial_leaf_cgroup(self): try: os.mkdir( self._cgroup_path(self._initial_cgroup_name, self._CGROUP_LEAF), mode=0o755, ) except FileExistsError: pass def _create_codeeval_cgroup(self): for counter in range(0, self._CGROUP_MAX_COUNT): codeeval_cgroup_name_candidate = f"{self._CGROUP_NAME_PREFIX}{counter}" cgroup_path = self._cgroup_path( self._initial_cgroup_name, codeeval_cgroup_name_candidate ) try: os.mkdir(cgroup_path, mode=0o755) except FileExistsError: pass else: self._codeeval_cgroup_name = codeeval_cgroup_name_candidate return initial_cgroup_path_prefix = self._cgroup_path( self._initial_cgroup_name, self._CGROUP_NAME_PREFIX ) raise OSError( f"Out of cgroups (tried {initial_cgroup_path_prefix}NUM with NUM from 0 to {self._CGROUP_MAX_COUNT-1} and they all already exist)" ) def _create_sandbox_cgroup(self): os.mkdir( self._cgroup_path( self._initial_cgroup_name, self._codeeval_cgroup_name, self._CGROUP_SANDBOX_NAME, ), mode=0o755, ) def _create_sandbox_leaf_cgroup(self): os.mkdir( self._cgroup_path( self._initial_cgroup_name, self._codeeval_cgroup_name, self._CGROUP_SANDBOX_NAME, self._CGROUP_LEAF, ), mode=0o755, ) def _create_supervisor_cgroup(self): os.mkdir( self._cgroup_path( self._initial_cgroup_name, self._codeeval_cgroup_name, self._CGROUP_SUPERVISOR_NAME, ), mode=0o755, ) def _create_supervisor_leaf_cgroup(self): os.mkdir( self._cgroup_path( self._initial_cgroup_name, self._codeeval_cgroup_name, self._CGROUP_SUPERVISOR_NAME, self._CGROUP_LEAF, ), mode=0o755, ) def _add_cgroup_controllers_to(self, *cgroup_components): add_controllers = tuple( controller for controller in self._cgroup_controllers if controller in self._needed_controllers ) cgroup_components = tuple(cgroup_components) + ("cgroup.subtree_control",) cgroup_subtree_control_path = self._cgroup_path(*cgroup_components) try: with self._open( cgroup_subtree_control_path, "wb" ) as cgroup_subtree_control_f: controllers_data = ( " ".join(f"+{controller}" for controller in add_controllers) + "\n" ) cgroup_subtree_control_f.write(controllers_data.encode("ascii")) except OSError as e: raise OSError( f"Adding controllers {add_controllers} on {cgroup_subtree_control_path}: {e}" ) got_controllers = set() try: with self._open( cgroup_subtree_control_path, "rb" ) as cgroup_subtree_control_f: for line in cgroup_subtree_control_f: for controller in line.strip().split(b" "): if not controller: continue got_controllers.add(controller.decode("ascii")) except OSError as e: raise OSError( f"Reading controllers from {cgroup_subtree_control_path}: {e}" ) assert all( controller in got_controllers for controller in add_controllers ), f"Missing controllers in {cgroup_subtree_control_path}: got {got_controllers} expected {add_controllers}" def _add_cgroup_controllers_to_root(self): return self._add_cgroup_controllers_to(self._initial_cgroup_name) def _add_cgroup_controllers_to_codeeval(self): return self._add_cgroup_controllers_to( self._initial_cgroup_name, self._codeeval_cgroup_name ) def _add_cgroup_controllers_to_sandbox(self): return self._add_cgroup_controllers_to( self._initial_cgroup_name, self._codeeval_cgroup_name, self._CGROUP_SANDBOX_NAME, ) def _add_cgroup_controllers_to_supervisor(self): return self._add_cgroup_controllers_to( self._initial_cgroup_name, self._codeeval_cgroup_name, self._CGROUP_SUPERVISOR_NAME, ) def _move_initial_cgroup_processes_to_initial_leaf_cgroup(self): initial_cgroup_procs_path = self._cgroup_path( self._initial_cgroup_name, "cgroup.procs" ) initial_leaf_cgroup_procs_path = self._cgroup_path( self._initial_cgroup_name, self._CGROUP_LEAF, "cgroup.procs" ) done_zero_pid = False while True: moving_process_pid = None with self._open( initial_cgroup_procs_path, "rb" ) as initial_cgroup_procs_f: for line in initial_cgroup_procs_f: if moving_process_pid is not None: continue for pid_str in line.strip().split(b" "): if not pid_str: continue try: pid = int(pid_str) except ValueError: continue if pid == 0: if done_zero_pid: continue done_zero_pid = True moving_process_pid = pid break if moving_process_pid is None: break with self._open( initial_leaf_cgroup_procs_path, "wb" ) as initial_leaf_cgroup_procs_f: initial_leaf_cgroup_procs_f.write( f"{moving_process_pid}\n".encode("ascii") ) def _move_process_to_supervisor_leaf(self): supervisor_leaf_cgroup_procs_path = self._cgroup_path( self._initial_cgroup_name, self._codeeval_cgroup_name, self._CGROUP_SUPERVISOR_NAME, self._CGROUP_LEAF, "cgroup.procs", ) f = self._open(supervisor_leaf_cgroup_procs_path, "wb") try: f.write(b"0\n") self._moved = True finally: try: f.close() except OSError: pass def _move_process_back(self): initial_leaf_cgroup_procs_path = self._cgroup_path( self._initial_cgroup_name, self._CGROUP_LEAF, "cgroup.procs" ) f = self._open(initial_leaf_cgroup_procs_path, "wb") try: f.write(b"0\n") self._moved = False finally: try: f.close() except OSError: pass def _set_cgroup_limits(self, *cgroup_components): cgroup_components = tuple(cgroup_components) + ("memory.max",) cgroup_memory_max_path = self._cgroup_path(*cgroup_components) if self._max_sandbox_ram_bytes is not None: try: with self._open(cgroup_memory_max_path, "wb") as memory_max_f: memory_max_f.write( f"{self._max_sandbox_ram_bytes}\n".encode("ascii") ) except OSError as e: raise OSError( f"Trying to set max RAM limit to {self._max_sandbox_ram_bytes} bytes: {e}" ) for swap_type in ("swap", "zswap"): cgroup_swap_components = tuple(cgroup_components) + ( f"memory.{swap_type}.max", ) cgroup_swap_path = self._cgroup_path(*cgroup_swap_components) if not os.path.exists(cgroup_swap_path): continue try: with self._open(cgroup_swap_path, "wb") as swap_max_f: swap_max_f.write("0\n".encode("ascii")) except OSError as e: raise OSError( f"Trying to set max {swap_type} limit to 0 bytes: {e}" ) def _set_supervisor_cgroup_limits(self): return self._set_cgroup_limits( self._initial_cgroup_name, self._codeeval_cgroup_name, self._CGROUP_SUPERVISOR_NAME, ) def _set_sandbox_cgroup_limits(self): return self._set_cgroup_limits( self._initial_cgroup_name, self._codeeval_cgroup_name, self._CGROUP_SANDBOX_NAME, ) def _sanity_check_own_cgroup(self): supervisor_cgroup_path = self._cgroup_path( self._initial_cgroup_name, self._codeeval_cgroup_name, self._CGROUP_SUPERVISOR_NAME, ) with self._open("/proc/self/cgroup", "rb") as cgroup_f: cgroup_data = cgroup_f.read().decode("ascii").strip() assert cgroup_data.endswith( os.sep + os.path.join(self._CGROUP_SUPERVISOR_NAME, self._CGROUP_LEAF) ), f"Unexpected self cgroup after moving to {supervisor_cgroup_path}: {cgroup_data}" def move_process_to_sandbox_leaf_cgroup_lambda(self): """ Returns a function that can move the current process to the sandbox leaf cgroup. :return: A function to move the current process to the sandbox cgroup. :raises SandboxException: If not queried after we have already chosen a new cgroup name. """ if not self._do_resource_limiting: return lambda: None if self._codeeval_cgroup_name is None: raise Sandbox.SandboxException( "Tried to move process to sandbox leaf cgroup before we know it" ) def _move(cgroup_path): """Dependency-free preexec_fn-compatible function to move to the given cgroup.procs.""" try: f = open(cgroup_path, "wb") except OSError as e: raise OSError(f"Cannot open cgroup path {cgroup_path}: {e}") try: f.write(b"0\n") except OSError as e: raise OSError(f"Cannot move process to {cgroup_path}: {e}") finally: try: f.close() except OSError: pass clone_newcgroup = ( os.CLONE_NEWCGROUP if "CLONE_NEWCGROUP" in os.__dict__ else 0x2000000 ) if "unshare" in os.__dict__: # Python >= 3.12. try: os.unshare(clone_newcgroup) except OSError as e: raise OSError(f"unshare({clone_newcgroup}) failed: {e}") else: import ctypes libc = ctypes.CDLL(None) libc.unshare.argtypes = [ctypes.c_int] rc = libc.unshare(clone_newcgroup) if rc == -1: raise OSError(f"unshare({clone_newcgroup}) failed") sandbox_cgroup_procs_path = self._cgroup_path( self._initial_cgroup_name, self._codeeval_cgroup_name, self._CGROUP_SANDBOX_NAME, self._CGROUP_LEAF, "cgroup.procs", )[:] return lambda: _move(sandbox_cgroup_procs_path) def monitor_cgroup_resources(self): """ Spawns a background thread that monitors resources, if limiting is enabled. cgroups should be taking care of this, but some systems do not enforce this. So this does the same in userspace. Better than nothing. :return: A function to cancel the monitor thread, if resource limiting is enabled. """ if not self._do_resource_limiting: return lambda: None self_memory_path = self._cgroup_path( self._initial_cgroup_name, self._codeeval_cgroup_name, "memory.peak", ) sandbox_procs_path = self._cgroup_path( self._initial_cgroup_name, self._codeeval_cgroup_name, self._CGROUP_SANDBOX_NAME, self._CGROUP_LEAF, "cgroup.procs", ) def _kill(): new_pids_to_kill = True pids_to_kill = set() while new_pids_to_kill: prev_pids_to_kill_len = len(pids_to_kill) with self._open(sandbox_procs_path, "rb") as cgroup_procs_f: for line in cgroup_procs_f: for pid_str in line.strip().split(b" "): if not pid_str: continue try: pid = int(pid_str) except ValueError: continue if pid != 0: pids_to_kill.add(pid) for pid_to_kill in pids_to_kill: try: os.kill(pid_to_kill, signal.SIGKILL) except Exception: pass new_pids_to_kill = prev_pids_to_kill_len < len(pids_to_kill) def _monitor(): if self._max_sandbox_ram_bytes is not None: try: with self._open(self_memory_path, "rb") as memory_peak_f: memory_peak_bytes = int( memory_peak_f.read().decode("ascii").strip() ) if memory_peak_bytes > self._max_sandbox_ram_bytes: _kill() except Exception as e: print( f"Warning: Failed to enforce code execution RAM: {e}", file=sys.stderr, ) lock = threading.Lock() enabled = [True] def _loop(): while True: time.sleep(0.1) with lock: if not enabled[0]: break _monitor() monitor_thread = threading.Thread( target=_loop, name="Monitor thread for code execution", daemon=True ) monitor_thread.start() def _cancel(): with lock: enabled[0] = False monitor_thread.join() return _cancel class SandboxException(Exception): """ Base class for all exceptions generated by `Sandbox`. """ def __init__(self, *args, **kwargs): self._sandbox_exception_args = tuple(args) self._sandbox_exception_kwargs = dict(kwargs) super().__init__(*args, **kwargs) class PlatformNotSupportedException(SandboxException): """ Raised when the sandbox cannot run on the current platform. The only way to fix this is to run on a different platform. """ class SandboxRuntimeException(SandboxException): """ Raised when the sandbox fails to run properly. This means gVisor itself is failing, not the code in the sandbox. """ class ExecutionError(subprocess.CalledProcessError): """ Raised when the sandboxed code fails to run. This means the sandbox worked, but the code that ran within failed. """ def __init__(self, code, **kwargs): super().__init__(**kwargs) self._code = code self._sandbox_exception_args = () self._sandbox_exception_kwargs = kwargs.copy() self._sandbox_exception_kwargs["code"] = code def __str__(self): super_str = super().__str__() full_code = self._code short_code = full_code.replace("\n", ";") if len(short_code) >= 128: short_code = short_code[:60] + "\u2026" + short_code[-60:] if self.stderr: lines = [ line.strip() for line in self.stderr.split("\n") if line.strip() ] if len(lines) >= 2: first_line, last_line = lines[0], lines[-1] return f"{first_line} [\u2026] {last_line} (`{short_code}`)\n{super_str}\n```\n{full_code}\n```" if len(lines) == 1: first_line = lines[0] return f"{first_line} (`{short_code}`)\n{super_str}\n```\n{full_code}\n```" return f"`{short_code}` failed\n{super_str}\n```\n{full_code}\n```" class CodeExecutionError(ExecutionError): """ Raised when the sandboxed code returns with a non-zero exit code. """ def __str__(self): super_str = super().__str__() return f"Code error: {super_str}" class ExecutionTimeoutError(ExecutionError): """ Raised when the code runs for too long relative to the timeout. """ def __str__(self): super_str = super().__str__() return f"Timeout: {super_str}" class InterruptedExecutionError(ExecutionError): """ Raised when the code runs but is interrupted before it finishes. This could happen from running out of resources. """ def __str__(self): super_str = super().__str__() return f"Interrupted: {super_str}" class FixableException(SandboxException): """ Base class for exceptions which can be addressed by the user. """ class GVisorNotInstalledException(FixableException): """ Raised when gVisor is not installed (`runsc` not found in $PATH). """ class CorruptDownloadException(FixableException): """ Raised when auto-downloading gVisor resulted in a hash mismatch. """ class EnvironmentNeedsSetupException(FixableException): """ Raised when the environment does not give adequate control over the system to run gVisor properly. """ @classmethod def check_platform(cls): """ Verifies that this tool is running on a supported platform. :return: Nothing. :raises PlatformNotSupportedException: If running on an unsupported platform. """ uname = platform.uname() if uname.system != "Linux": raise cls.PlatformNotSupportedException(f"{uname.system} is not supported") if uname.machine not in ("x86_64", "aarch64"): raise cls.PlatformNotSupportedException(f"{uname.machine} is not supported") @classmethod def check_cgroups(cls): """ Verifies that cgroupfs is mounted and usable for resource limiting. :return: Nothing. :raises EnvironmentNeedsSetupException: If cgroupfs is not mounted or unusable for resource limiting. """ if not os.path.exists("/sys/fs/cgroup"): raise cls.EnvironmentNeedsSetupException( "cgroupfs not mounted as /sys/fs/cgroup but necessary for the sandbox to enforce memory limits; please mount it (`--mount=type=bind,source=/sys/fs/cgroup,target=/sys/fs/cgroup,readonly=false`), or disable resource limiting if appropriate" ) if not os.path.exists("/sys/fs/cgroup/cgroup.subtree_control"): raise cls.EnvironmentNeedsSetupException( "/sys/fs/cgroup/cgroup.subtree_control not found; make sure you are using cgroups v2, or disable resource limiting if appropriate" ) # Try to open the file for writing to see if we can actually control cgroups. # They may be mounted read-only, as is default with Docker. try: with open( "/sys/fs/cgroup/cgroup.subtree_control", "wb" ) as subtree_control_f: pass except OSError: if os.geteuid() != 0: raise cls.EnvironmentNeedsSetupException( "This script is not running as root, but it needs to do so in order to enforce resource limits; please run as root, or disable resource limiting if appropriate" ) raise cls.EnvironmentNeedsSetupException( "cgroupfs is not mounted writable but necessary for the sandbox to enforce memory limits; please remount it as writable (`--mount=type=bind,source=/sys/fs/cgroup,target=/sys/fs/cgroup,readonly=false`), or disable resource limiting if appropriate" ) with open("/sys/fs/cgroup/cgroup.controllers", "rb") as subtree_control_f: controllers = subtree_control_f.read().decode("ascii").split(" ") if "memory" not in controllers: raise cls.EnvironmentNeedsSetupException( "cgroupfs does not have the 'memory' controller enabled, necessary to enforce memory limits; please enable it, or disable resource limiting if appropriate" ) @classmethod def check_procfs(cls): """ Verifies that we have an unobstructed view of procfs. :return: Nothing. :raises EnvironmentNeedsSetupException: If procfs is obstructed. """ mount_infos = [] with open("/proc/self/mountinfo", "rb") as mountinfo_f: for line in mountinfo_f: line = line.decode("utf-8").strip() if not line: continue mount_components = line.split(" ") if len(mount_components) < 10: continue hyphen_index = mount_components.index("-") if hyphen_index < 6: continue mount_info = { "mount_path": mount_components[4], "path_within_mount": mount_components[3], "fs_type": mount_components[hyphen_index + 1], } mount_infos.append(mount_info) procfs_mounts = frozenset( m["mount_path"] for m in mount_infos if m["fs_type"] == "proc" and m["path_within_mount"] == "/" ) if len(procfs_mounts) == 0: raise cls.EnvironmentNeedsSetupException( "procfs is not mounted; please mount it" ) obstructed_procfs_mounts = set() for mount_info in mount_infos: for procfs_mount in procfs_mounts: if mount_info["mount_path"].startswith(procfs_mount + os.sep): obstructed_procfs_mounts.add(procfs_mount) for procfs_mount in procfs_mounts: if procfs_mount not in obstructed_procfs_mounts: return # We have at least one unobstructed procfs view. assert len(obstructed_procfs_mounts) > 0, "Logic error" raise cls.EnvironmentNeedsSetupException( "procfs is obstructed; please mount a new procfs mount somewhere in the container, e.g. /proc2 (`--mount=type=bind,source=/proc,target=/proc2,readonly=false,bind-recursive=disabled`)" ) @classmethod def unshare(cls, flags): """ Implementation of `os.unshare` that works on Python < 3.12. :param flags: Flags to pass to the `unshare(2)` system call. :raises OSError: If something goes wrong. """ if "unshare" in os.__dict__: # Python >= 3.12. return os.unshare(flags) # Python <= 3.11: return cls._libc().unshare(flags) @classmethod def check_unshare(cls): """ Verifies that the `unshare(2)` system call is available. :return: Nothing. :raises EnvironmentNeedsSetupException: If `unshare(2)` is not available. """ try: cls.unshare(0) except OSError: raise cls.EnvironmentNeedsSetupException( "`unshare(2)` syscall is unavailable but necessary for the sandbox to isolate itself; please remove the seccomp filter (`--security-opt=seccomp=unconfined`)" ) @classmethod def get_runsc_path(cls): """ Returns the absolute path where the `runsc` binary is installed. :return: Absolute path to `runsc` binary, or `None` if not installed. """ runsc_path = shutil.which("runsc") if runsc_path: return runsc_path if os.path.exists(cls.AUTO_INSTALLATION_PATH): return cls.AUTO_INSTALLATION_PATH return None @classmethod def runsc_needs_installation(cls): """ Checks whether the `runsc` binary is installed. :return: Whether the `runsc` binary is installed. """ return cls.get_runsc_path() is None @classmethod def install_runsc(cls): """ Download and install the `runsc` binary if not already present. :return: Nothing. :raises CorruptDownloadException: If the download resulted in a hash mismatch. """ if not cls.runsc_needs_installation(): return uname = platform.uname() release_url_dir = f"https://storage.googleapis.com/gvisor/releases/release/latest/{uname.machine}" os.makedirs( os.path.dirname(cls.AUTO_INSTALLATION_PATH), mode=0o755, exist_ok=True ) with tempfile.TemporaryDirectory( prefix="sandbox_download_" ) as download_tmp_dir: download_path = os.path.join(download_tmp_dir, "runsc") urllib.request.urlretrieve( url=f"{release_url_dir}/runsc", filename=download_path, ) sha512_raw = urllib.request.urlopen( f"{release_url_dir}/runsc.sha512" ).read() want_sha512 = sha512_raw.decode("ascii").split(" ")[0] runsc_hash = hashlib.sha512() with open(download_path, "rb") as runsc_f: while True: chunk = runsc_f.read(65536) if not chunk: break runsc_hash.update(chunk) runsc_sha512 = runsc_hash.hexdigest() if runsc_sha512 != want_sha512: raise cls.CorruptDownloadException( "gVisor hash mismatch when auto-installing; please install gVisor manually" ) os.chmod(download_path, mode=0o755) os.rename(download_path, cls.AUTO_INSTALLATION_PATH) @classmethod def check_setup( cls, language: str, auto_install_allowed: bool, require_resource_limiting: bool, ): """ Verifies that the environment is compatible with running sandboxes. :param language: The programming language to run. :param auto_install_allowed: Whether auto-installation of `runsc` is allowed. :param require_resource_limiting: Check that the host supports resource limiting via cgroups. :return: Nothing. :raises ValueError: If provided an invalid language name. :raises PlatformNotSupportedException: If running on an unsupported platform. :raises FixableException: If another issue occurs but that can be fixed by the user. """ if language not in cls.SUPPORTED_LANGUAGES: raise ValueError(f"Unsupported language: {language}") if shutil.which("bash") is None: raise cls.EnvironmentNeedsSetupException( "bash is not installed (`bash` binary not found in $PATH); please install it" ) if shutil.which("unshare") is None: raise cls.EnvironmentNeedsSetupException( "unshare is not installed (`unshare` binary not found in $PATH); please install it" ) cls.check_platform() cls.check_unshare() if require_resource_limiting: cls.check_cgroups() cls.check_procfs() if not auto_install_allowed and cls.get_runsc_path() is None: raise cls.GVisorNotInstalledException( "gVisor is not installed (runsc binary not found in $PATH); please install it or enable AUTO_INSTALL valve for auto installation" ) @classmethod def _libc(cls): if cls._LIBC is None: cls._LIBC = cls._Libc() return cls._LIBC @classmethod def main(cls): """ Entry-point for (re-)execution. Must be called during import. May call `sys.exit` if this is intended to be a code evaluation re-execution. """ cls._SelfFile.init() if cls._MARKER_ENVIRONMENT_VARIABLE not in os.environ: return directives = json.load(sys.stdin) try: result = cls(**directives["settings"])._run() except Exception as e: exception_info = { "name": e.__class__.__name__, "str": str(e), } if isinstance(e, cls.SandboxException) or isinstance(e, cls.ExecutionError): exception_info["args"] = e._sandbox_exception_args exception_info["kwargs"] = e._sandbox_exception_kwargs json.dump( { "exception": exception_info, }, sys.stdout, ) else: stdout = result.stdout if type(stdout) is not type(b""): stdout = stdout.encode("utf-8", errors="replace") stderr = result.stderr if type(stderr) is not type(b""): stderr = stderr.encode("utf-8", errors="replace") json.dump( { "result": { "args": result.args, "returncode": result.returncode, "stdout": base64.b64encode(stdout).decode("utf-8"), "stderr": base64.b64encode(stderr).decode("utf-8"), }, }, sys.stdout, ) finally: sys.stdout.flush() sys.exit(0) def __init__( self, tmp_dir: str, language: str, code: str, debug: bool, networking_allowed: bool, max_runtime_seconds: int, max_ram_bytes: typing.Optional[int] = None, require_resource_limiting: bool = False, persistent_home_dir: typing.Optional[str] = None, ): """ Constructor. :param tmp_dir: Temporary directory exclusive to this sandbox. Must outlive the Sandbox object. :param language: The language of the code; must be one of `SUPPORTED_LANGUAGES`. :param code: Arbitrary code that needs to run in the sandbox. :param debug: Whether or not to enable debug-level logging for the sandbox. :param networking_allowed: Whether the code should be given access to the network. :param max_runtime_seconds: How long the code should be allowed to run, in seconds. :param max_ram_bytes: How many bytes of RAM the interpreter should be allowed to use, or `None` for no limit. :param require_resource_limiting: If true, refuse to launch a sandbox if the host doesn't support resource limiting via cgroups. :param persistent_home_dir: Optional directory which will be mapped read-write to this real host directory. """ self._init( { "tmp_dir": tmp_dir, "language": language, "code": code, "debug": debug, "networking_allowed": networking_allowed, "max_runtime_seconds": max_runtime_seconds, "max_ram_bytes": max_ram_bytes, "require_resource_limiting": require_resource_limiting, "persistent_home_dir": persistent_home_dir, } ) def _init(self, settings): self._settings = settings self._tmp_dir = self._settings["tmp_dir"] self._bundle_path = os.path.join(self._tmp_dir, "bundle") self._runtime_root_path = os.path.join(self._tmp_dir, "runtime") self._logs_path = os.path.join(self._tmp_dir, "logs") self._gotmp_dir = os.path.join(self._tmp_dir, "gotmp") self._sandbox_shared_path = os.path.join(self._tmp_dir, "sandbox") self._language = self._settings["language"] self._code = self._settings["code"] self._debug = self._settings["debug"] self._networking_allowed = self._settings["networking_allowed"] self._max_runtime_seconds = self._settings["max_runtime_seconds"] self._max_ram_bytes = self._settings["max_ram_bytes"] self._require_resource_limiting = self._settings[ "require_resource_limiting" ] or all((self._max_ram_bytes is None,)) self._persistent_home_dir = self._settings["persistent_home_dir"] self._sandboxed_command = None self._switcheroo = None def _setup_sandbox(self): """ Set up the sandbox's root filesystem and OCI config prior to execution. Runs in separate forked process. Performs the switcheroo. :raises FixableException: If an issue occurs but that can be fixed by the user. """ # Set up basic configuration options. oci_config = copy.deepcopy(self.OCI_CONFIG_SKELETON) if self._max_ram_bytes: oci_config["linux"]["resources"]["memory"]["limit"] = self._max_ram_bytes os.makedirs(self._bundle_path, mode=0o711) os.makedirs(self._runtime_root_path, mode=0o711) os.makedirs(self._logs_path, mode=0o711) os.makedirs(self._sandbox_shared_path, mode=0o777) os.makedirs(self._gotmp_dir, mode=0o711) os.chmod(self._sandbox_shared_path, mode=0o777, follow_symlinks=False) rootfs_path = os.path.join(self._tmp_dir, "rootfs") os.makedirs(rootfs_path, mode=0o755) if self._persistent_home_dir is not None: if not os.path.isdir(self._persistent_home_dir): raise self.SandboxException( f"Persistent home directory {self._persistent_home_dir} does not exist" ) oci_config["root"]["path"] = rootfs_path do_resource_limiting = True if not self._require_resource_limiting: try: self.check_cgroups() except self.EnvironmentNeedsSetupException: do_resource_limiting = False self._switcheroo = self._Switcheroo( libc=self._libc(), log_path=os.path.join(self._logs_path, "switcheroo.txt"), max_sandbox_ram_bytes=self._max_ram_bytes, do_resource_limiting=do_resource_limiting, ) try: self._switcheroo.do() except Exception as e: try: switcheroo_status = self._switcheroo._status() except Exception: raise e else: raise e.__class__(f"{e}; {switcheroo_status}") # Locate the interpreter to use. interpreter_path = sys.executable if self._language == self.LANGUAGE_BASH: interpreter_path = shutil.which("bash") if interpreter_path is None: raise RuntimeError("Interpreter not found") oci_config["mounts"].append( { "type": "bind", "source": interpreter_path, "destination": interpreter_path, "options": ["ro"], } ) # Populate rootfs. This is a multi-step process. # Create writable empty directories. for d in self.EMPTY_WRITABLE_DIRECTORIES: rootfs_subdir = os.path.join(rootfs_path, d.removeprefix(os.path.sep)) os.makedirs(rootfs_subdir, mode=0o755, exist_ok=True) oci_config["mounts"].append( { "type": "tmpfs", "destination": d, "options": [], } ) # Create read-only empty directories. for d in self.EMPTY_READ_ONLY_DIRECTORIES + [os.path.dirname(interpreter_path)]: rootfs_subdir = os.path.join(rootfs_path, d.removeprefix(os.path.sep)) os.makedirs(rootfs_subdir, mode=0o755, exist_ok=True) # Handle exposed host symlinks. These will show up as symlinks with the same # target path in the sandbox, so they do not expose the host's view of the # directory they point to. symlinks = set() for p in self.EXPOSED_SYSTEM_DIRECTORIES + self.EXPOSED_SYSTEM_FILES: if not os.path.islink(p): continue rootfs_subpath = os.path.join(rootfs_path, p.removeprefix(os.path.sep)) os.makedirs(os.path.dirname(rootfs_subpath), mode=0o755, exist_ok=True) os.symlink(src=os.readlink(p), dst=rootfs_subpath) symlinks.add(p) # Handle exposed host directories. for d in self.EXPOSED_SYSTEM_DIRECTORIES: if d in symlinks: continue # It is a symlink, so already handled. if not os.path.isdir(d): continue # The host does not have a directory at this path. rootfs_subdir = os.path.join(rootfs_path, d.removeprefix(os.path.sep)) os.makedirs(rootfs_subdir, mode=0o755, exist_ok=True) oci_config["mounts"].append( { "type": "bind", "source": d, "destination": d, "options": ["ro", "rbind"], } ) # Handle exposed host files. for f in self.EXPOSED_SYSTEM_FILES: if f in symlinks: continue # It is a symlink, so already handled. if not os.path.isfile(f): continue # The host does not have a file at this path. rootfs_subpath = os.path.join(rootfs_path, f.removeprefix(os.path.sep)) rootfs_subdir = os.path.dirname(rootfs_subpath) os.makedirs(rootfs_subdir, mode=0o755, exist_ok=True) oci_config["mounts"].append( { "type": "bind", "source": f, "destination": f, "options": ["ro"], } ) # Shared sandbox directory to propagate exit code and persistent files. oci_config["mounts"].append( { "type": "bind", "source": self._sandbox_shared_path, "destination": "/sandbox", "options": ["rw"], } ) if self._persistent_home_dir is not None: oci_config["mounts"].append( { "type": "bind", "source": self._persistent_home_dir, "destination": "/sandbox/persistent", "options": ["rw"], } ) # Sort mounts to ensure proper overlay order. oci_config["mounts"].sort(key=lambda m: m["destination"]) # Generate some /etc files that look normal. with open(os.path.join(rootfs_path, "etc/hostname"), "w") as hostname_f: hostname_f.write("sandbox\n") with open(os.path.join(rootfs_path, "etc/passwd"), "w") as passwd_f: passwd_f.write("user:x:1000:1000:user:/home/user:/bin/bash\n") # Generate command line to run in the sandbox. self._sandboxed_command = [ shutil.which("bash"), "-c", "; ".join( ( "echo OK > /sandbox/started", f"{interpreter_path} /dev/stdin", 'echo "$?" > /sandbox/.pre_exit_code || exit 1', "if [[ -d /sandbox/persistent ]]; then cp -rd --one-file-system /home/user/. /sandbox/persistent/ || exit 2; fi", "mv /sandbox/.pre_exit_code /sandbox/exit_code || exit 3", "exit 0", ) ), ] # Work around issue that gVisor does not preserve correct UID mappings when running as non-root user in the sandbox. # So map current user to 0:0, then create a new user namespace immediately before running command and remap to # correct UID/GID. oci_config["process"]["user"]["uid"] = 0 oci_config["process"]["user"]["gid"] = 0 oci_config["process"]["args"] = [ shutil.which("unshare"), "--map-user=1000", "--map-group=1000", ] + self._sandboxed_command # We are done. Write OCI config to bundle directory. with open(os.path.join(self._bundle_path, "config.json"), "w") as bundle_f: json.dump(oci_config, bundle_f, indent=2, sort_keys=True) def _run(self) -> subprocess.CompletedProcess: """ Spawn and wait for the sandbox. Runs in separate forked process. :return: A `CompletedProcess` object representing the return code and stdout/stderr of the code interpreter. :raises Sandbox.SandboxRuntimeException: If the sandbox failed to start or behaved incorrectly regardless of the code being evaluated. :raises Sandbox.ExecutionTimeoutError: If the code interpreter ran for longer than configured. :raises Sandbox.InterruptedExecutionError: If the code interpreter died without providing a return code; usually due to running over resource limits. :raises sandbox.CodeExecutionError: If the code interpreter failed to execute the given code. This does not represent a sandbox failure. """ try: self._setup_sandbox() network_mode = "host" if self._networking_allowed else "none" runsc_argv = [ self.get_runsc_path(), "--rootless=true", "--directfs=false", f"--network={network_mode}", "--ignore-cgroups=true", # We already took care of cgroups manually. f"--root={self._runtime_root_path}", f"--debug-log={self._logs_path}/", "run", f"--bundle={self._bundle_path}", "sandbox", ] runsc_env = os.environ.copy() runsc_env["TMPDIR"] = self._gotmp_dir started_marker_path = os.path.join(self._sandbox_shared_path, "started") resource_monitor_cancel = self._switcheroo.monitor_cgroup_resources() try: result = subprocess.run( runsc_argv, env=runsc_env, preexec_fn=self._switcheroo.move_process_to_sandbox_leaf_cgroup_lambda(), input=self._code + "\n", text=True, capture_output=True, timeout=self._max_runtime_seconds, check=True, ) except subprocess.TimeoutExpired as e: raise self.ExecutionTimeoutError( code=self._code, returncode=126, cmd=self._sandboxed_command, output=e.stdout, stderr=e.stderr, ) except subprocess.CalledProcessError as e: if os.path.isfile(started_marker_path): raise self.InterruptedExecutionError( code=self._code, returncode=127, cmd=self._sandboxed_command, output=e.stdout, stderr=e.stderr, ) logs = {} def process_log(filename, log_line): if self._debug or ( log_line and log_line[0] in "WEF" ): # Warning, Error, Fatal if filename not in logs: logs[filename] = [] logs[filename].append(log_line) self.debug_logs(process_log) stderr = e.stderr.strip() json_logs = json.dumps(logs) if self._debug: raise self.SandboxRuntimeException( f"Sandbox failed to start: {e}; stderr: {stderr}; logs: {json_logs}" ) raise self.SandboxRuntimeException( f"Sandbox failed to start: {e} (turn on debug mode to see more information); stderr: {stderr}; logs: {json_logs}" ) finally: resource_monitor_cancel() if not os.path.isfile(started_marker_path): raise self.SandboxRuntimeException( "Sandbox failed to start up properly" ) exit_code_path = os.path.join(self._sandbox_shared_path, "exit_code") if not os.path.isfile(exit_code_path): raise self.SandboxRuntimeException( "Sandbox failed to record an exit code" ) with open(exit_code_path, "r") as exit_code_f: exit_code_str = exit_code_f.read() try: exit_code = int(exit_code_str.strip()) except ValueError as e: raise self.SandboxRuntimeException( f"Sandbox recorded non-integer exit code: {e}" ) if exit_code != 0: raise self.CodeExecutionError( code=self._code, returncode=exit_code, cmd=self._sandboxed_command, output=result.stdout, stderr=result.stderr, ) return result finally: if self._switcheroo is not None: self._switcheroo.cleanup() def run(self) -> subprocess.CompletedProcess: """ Set up and run the sandbox in a separate process. :return: A `CompletedProcess` object representing the return code and stdout/stderr of the code interpreter. :raises FixableException: If an issue occurs but that can be fixed by the user. :raises Sandbox.SandboxRuntimeException: If the sandbox failed to start or behaved incorrectly regardless of the code being evaluated. :raises Sandbox.ExecutionTimeoutError: If the code interpreter ran for longer than configured. :raises Sandbox.InterruptedExecutionError: If the code interpreter died without providing a return code; usually due to running over resource limits. :raises sandbox.CodeExecutionError: If the code interpreter failed to execute the given code. This does not represent a sandbox failure. """ reexec_path = os.path.join(self._tmp_dir, "self.py") with open(reexec_path, "w") as reexec_f: reexec_f.write(self._SelfFile.contents()) new_env = os.environ.copy() new_env[self._MARKER_ENVIRONMENT_VARIABLE] = "1" data = json.dumps({"settings": self._settings}) try: result = subprocess.run( (sys.executable, reexec_path), env=new_env, input=data, text=True, capture_output=True, check=True, ) except subprocess.CalledProcessError as e: raise self.SandboxRuntimeException(f"{e} (stderr: {e.stderr})") else: if not result.stdout: raise self.SandboxRuntimeException( f"Subprocess interpreter did not produce any output (stderr: {result.stderr})" ) try: output = json.loads(result.stdout) except json.decoder.JSONDecodeError as e: raise self.SandboxRuntimeException( f"Subprocess interpreter produced invalid JSON (stdout: {result.stdout}): {e}" ) if "exception" in output: class_name = output["exception"]["name"] found_class = None for ex_class in ( self.PlatformNotSupportedException, self.SandboxRuntimeException, self.CodeExecutionError, self.ExecutionTimeoutError, self.InterruptedExecutionError, self.GVisorNotInstalledException, self.CorruptDownloadException, self.EnvironmentNeedsSetupException, self.ExecutionError, self.SandboxException, ): if ex_class.__name__ == class_name: found_class = ex_class break if found_class is None: exception_str = output["exception"]["str"] raise self.SandboxException(f"{class_name}: {exception_str}") raise found_class( *output["exception"]["args"], **output["exception"]["kwargs"] ) if "result" not in output: raise self.SandboxException( f"Invalid response from subprocess: {output}" ) return subprocess.CompletedProcess( args=output["result"]["args"], returncode=output["result"]["returncode"], stdout=base64.b64decode(output["result"]["stdout"]).decode( "utf-8", errors="replace" ), stderr=base64.b64decode(output["result"]["stderr"]).decode( "utf-8", errors="replace" ), ) def debug_logs(self, write_fn: typing.Callable[[str, str], typing.Any]): """ Write debug logs and other system information to the given function. May only be called after `run` returns, but may be called even when `run` fails. :param write_fn: A function that takes (filename, log line) as arguments. """ log_paths = [] all_logs = [] if os.path.isdir(self._logs_path): for log_filename in sorted(os.listdir(self._logs_path)): if not log_filename.endswith(".txt"): continue log_paths.append(os.path.join(self._logs_path, log_filename)) else: all_logs.append( ( "[meta]", f"Logs path {self._logs_path} does not exist or is not a directory.", ) ) log_paths.extend(self._EXTRA_DEBUG_LOG_PATHS) runsc_filenames = set() for log_path in log_paths: log_filename = os.path.basename(log_path) if log_filename.startswith("runsc.log.") and log_filename.endswith(".txt"): log_filename = "runsc." + self._LOG_FILENAME_TRUNCATE_RE.sub( "", log_filename ).removesuffix(".txt") if log_filename in runsc_filenames: runsc_filename_suffix = 2 while f"{log_filename}.{runsc_filename_suffix}" in runsc_filenames: runsc_filename_suffix += 1 log_filename = f"{log_filename}.{runsc_filename_suffix}" runsc_filenames.add(log_filename) try: with open(log_path, "rb") as log_f: for log_line in log_f: log_line = log_line.replace(b"\x00", b"\n").rstrip() if not log_line: continue for line in log_line.split(b"\n"): all_logs.append( (log_filename, line.decode("utf-8", errors="replace")) ) except Exception as e: all_logs.append((log_filename, f"Failed to open: {e}")) for cmd in self._EXTRA_DEBUG_LOG_COMMANDS: cmd_str = "`" + " ".join(cmd) + "`" try: result = subprocess.run(cmd, capture_output=True, timeout=1, check=True) except subprocess.CalledProcessError as e: all_logs.append( (cmd_str, f"Failed: {e} (stdout={e.stdout}, stderr={e.stderr})") ) except Exception as e: all_logs.append((cmd_str, f"Failed: {e}")) else: for line in result.stdout.replace(b"\x00", b"\n").split(b"\n"): line = line.rstrip() if not line: continue all_logs.append((cmd_str, line.decode("utf-8", errors="replace"))) for filename, log_entry in all_logs: write_fn(filename, log_entry) Sandbox.main() class UpdateCheck: """ Check for updates. """ RELEASES_URL = "https://github.com/EtiennePerot/safe-code-execution/releases.atom" USER_URL = "https://github.com/EtiennePerot/safe-code-execution/" ENABLED = True SELF_VERSION = None LAST_UPDATE_CHECK = None LAST_UPDATE_CACHE = None UPDATE_CHECK_INTERVAL = datetime.timedelta(days=3) VERSION_REGEX = re.compile(r"<title>\s*(v?\d+(?:\.\d+)+)\s*</title>") class VersionCheckError(Exception): pass @staticmethod def _parse_version(version_str): return tuple(int(c) for c in version_str.strip().removeprefix("v").split(".")) @staticmethod def _format_version(version): return "v" + ".".join(str(c) for c in version) @staticmethod def _compare(version_a, version_b): """ Returns -1 if version_a < version_b, 0 if equal, 1 if greater. """ for a, b in zip(version_a, version_b): if a < b: return -1 if a > b: return 1 return len @classmethod def disable(cls): cls.ENABLED = False @classmethod def init_from_frontmatter(cls, file_with_frontmatter): if not cls.ENABLED: return with open(file_with_frontmatter, "rb") as f: contents = f.read().decode("ascii").strip() if not contents.startswith('"""'): raise cls.VersionCheckError( f"Malformed file contents: {contents[:min(8, len(contents))]}[...]" ) contents = contents[len('"""') :].strip() version = None for line in contents.split("\n"): line = line.strip() if line == '"""': break if line.startswith("version:"): if version is not None: raise cls.VersionCheckError( f"Multiple 'version' lines found: {version} and {line}" ) version = line[len("version:") :].strip() if version is None: raise cls.VersionCheckError("Version metadata not found") cls.SELF_VERSION = cls._parse_version(version) @classmethod def _get_current_version(cls): assert ( cls.SELF_VERSION is not None ), "UpdateCheck.init_from_frontmatter must be called first." return cls.SELF_VERSION @classmethod def need_check(cls): if cls.LAST_UPDATE_CHECK is None: return True return ( datetime.datetime.now() - cls.LAST_UPDATE_CHECK >= cls.UPDATE_CHECK_INTERVAL ) @classmethod def _get_latest_version(cls): if not cls.need_check(): if type(cls.LAST_UPDATE_CACHE) is type(()): return cls.LAST_UPDATE_CACHE raise cls.LAST_UPDATE_CACHE try: try: releases_xml = urllib.request.urlopen(url=cls.RELEASES_URL).read() except urllib.error.HTTPError as e: cls.LAST_UPDATE_CACHE = cls.VersionCheckError( f"Failed to retrieve latest version: {e} (URL: {cls.RELEASES_URL})" ) raise cls.LAST_UPDATE_CACHE latest_version = None for match in cls.VERSION_REGEX.finditer(releases_xml.decode("utf-8")): version = cls._parse_version(match.group(1)) if latest_version is None or cls._compare(version, latest_version) == 1: latest_version = version if latest_version is None: cls.LAST_UPDATE_CACHE = cls.VersionCheckError( f"Failed to retrieve latest version: no release found (URL: {cls.RELEASES_URL})" ) raise cls.LAST_UPDATE_CACHE cls.LAST_UPDATE_CACHE = latest_version return latest_version finally: cls.LAST_UPDATE_CHECK = datetime.datetime.now() @classmethod def get_newer_version(cls) -> typing.Optional[str]: """ Check for the latest version and return it if newer than current. :raises VersionCheckError: If there was an error checking for version. :return: The latest version number if newer than current, else None. """ if not cls.ENABLED: return None try: current_version = cls._get_current_version() except cls.VersionCheckError as e: raise e.__class__(f"Checking current version: {e}") try: latest_version = cls._get_latest_version() except cls.VersionCheckError as e: raise e.__class__(f"Checking latest version: {e}") if cls._compare(current_version, latest_version) == -1: return cls._format_version(latest_version) return None UpdateCheck.init_from_frontmatter(os.path.abspath(__file__)) _SAMPLE_BASH_INSTRUCTIONS = ( "echo 'Hello from the sandbox!'", "date", "dmesg", "echo 'Bye from the sandbox!'", ) _SAMPLE_PYTHON_INSTRUCTIONS = ( "print('Hello from the sandbox!')", "import datetime, sys", "print('Current date and time:', datetime.datetime.now())", "sys.stdout.flush()", "import shutil, subprocess", "subprocess.run([shutil.which('dmesg')], check=True)", "print('Bye from the sandbox!')", ) def _do_self_tests(debug): user_storage_path = None def _want_generated_files(want_generated_files): def _verify(got_generated_files): try: for name in want_generated_files.keys(): if name not in got_generated_files: raise ValueError(f"File {name} is missing") for name in got_generated_files.keys(): if name not in want_generated_files: raise ValueError(f"Unexpected generated file: {name}") for name, predicates in want_generated_files.items(): if type(predicates) not in (type(()), type([])): predicates = (predicates,) for i, predicate in enumerate(predicates): try: result = predicate(got_generated_files[name]) if result is not None: raise ValueError(result) except ValueError as e: raise e.__class__( f"File {name} predicate #{i}: {e} (got: {got_generated_files[name]})" ) except ValueError as e: return e else: return None return _verify def _markdown_contains(substring): return lambda md: ( f"substring '{substring}' not found" if substring not in md else None ) def _markdown_does_not_contain(substring): return lambda md: ( f"substring '{substring}' unexpectedly found" if substring in md else None ) def _want_user_storage_num_files(num_files_predicate): if type(num_files_predicate) is type(42): want_num_files = num_files_predicate def num_files_predicate_fn(n): return n == want_num_files num_files_predicate = num_files_predicate_fn def _verify(): total_files = 0 for _, _, subfiles in os.walk(user_storage_path): total_files += len(list(f for f in subfiles if f != ".lock")) if not num_files_predicate(total_files): raise ValueError(f"User storage unexpectedly has {total_files} files") return _verify _self_tests = ( { "name": "simple_python", "language": "python", "code": _SAMPLE_PYTHON_INSTRUCTIONS, "debug": True, "status": "OK", }, { "name": "simple_bash", "language": "bash", "code": _SAMPLE_BASH_INSTRUCTIONS, "debug": True, "status": "OK", "post": _want_user_storage_num_files(0), }, { "name": "bad_syntax_python", "language": "python", "code": ("print('foo",), "debug": True, "status": "ERROR", }, { "name": "bad_syntax_bash", "language": "bash", "code": ("echo 'foo",), "debug": True, "status": "ERROR", }, { "name": "simple_file_creation", "language": "bash", "code": ("echo 'Hello world!' > today.txt",), "status": "OK", "generated_files": { "today.txt": ( _markdown_contains("today.txt"), _markdown_contains("Hello world!\n"), ) }, "post": _want_user_storage_num_files(1), }, { "name": "generated_file_rendering", "language": "bash", "code": ( "mkdir -p foo/bar", "dmesg > foo/bar/dmesg.txt", "touch empty.txt", "(echo '# Example markdown file'; echo 'This is markdown!') > markdown.md", "echo 'This file cannot `''`''` be inlined' > no_inline.txt", "echo 'Might be an image' > picture.png", "echo 'Might be music' > song.flac", "echo 'Might be a video' > movie.mkv", "echo 'Might be random data' > random_data", "echo 'Might be random data' > random_data.bin", ), "valves": { "WEB_ACCESSIBLE_DIRECTORY_URL": "http://myhost/URL/ROOT", }, "status": "OK", "generated_files": { "foo/bar/dmesg.txt": ( _markdown_contains("\U0001f4c4"), _markdown_contains("[foo/bar/dmesg.txt]"), _markdown_contains("(http://myhost/URL/ROOT/"), _markdown_contains("/foo/bar/dmesg.txt)"), _markdown_contains("gVisor"), _markdown_contains("```"), ), "empty.txt": ( _markdown_contains("\u2049"), _markdown_does_not_contain("\U0001f4c4"), _markdown_contains("`empty.txt`"), _markdown_contains("(empty)"), _markdown_does_not_contain("[empty.txt]"), _markdown_does_not_contain("(http://myhost/URL/ROOT"), ), "markdown.md": ( _markdown_contains("\U0001f4c3"), _markdown_does_not_contain("\U0001f4c4"), _markdown_contains( "\n> # Example markdown file\n> This is markdown!" ), _markdown_does_not_contain("```"), ), "no_inline.txt": ( _markdown_contains("\U0001f4c4"), _markdown_does_not_contain("This file cannot"), _markdown_does_not_contain("```"), ), "picture.png": ( _markdown_contains("\U0001f5bc"), _markdown_does_not_contain("Might be"), ), "song.flac": ( _markdown_contains("\U0001f3b5"), _markdown_does_not_contain("Might be"), ), "movie.mkv": ( _markdown_contains("\U0001f3ac"), _markdown_does_not_contain("Might be"), ), "random_data": ( _markdown_contains("\U0001f4be"), _markdown_does_not_contain("Might be"), ), "random_data.bin": ( _markdown_contains("\U0001f4be"), _markdown_does_not_contain("Might be"), ), }, "post": _want_user_storage_num_files(10), }, { "name": "generated_text_files_too_large_to_render", "language": "bash", "code": ( f"yes boop | head -{_Action._UserFile.MAX_INLINE_TEXT_LINES-1} > ok_lines.md", f"yes boop | head -{_Action._UserFile.MAX_INLINE_TEXT_LINES-1} > ok_lines.txt", f"yes boop | head -{_Action._UserFile.MAX_INLINE_TEXT_LINES+1} > too_many_lines.md", f"yes boop | head -{_Action._UserFile.MAX_INLINE_TEXT_LINES+1} > too_many_lines.txt", f"yes boop | tr '\\n' ' ' | head -c{_Action._UserFile.MAX_INLINE_TEXT_BYTES-1} > ok_bytes.txt", f"yes boop | tr '\\n' ' ' | head -c{_Action._UserFile.MAX_INLINE_TEXT_BYTES+1} > too_many_bytes.txt", ), "status": "OK", "generated_files": { "ok_lines.md": _markdown_contains("boop"), "ok_lines.txt": _markdown_contains("boop"), "too_many_lines.md": _markdown_does_not_contain("boop"), "too_many_lines.txt": _markdown_does_not_contain("boop"), "ok_bytes.txt": _markdown_contains("boop"), "too_many_bytes.txt": _markdown_does_not_contain("boop"), }, "post": _want_user_storage_num_files(16), }, { "name": "too_large_file_for_user_quota", "language": "bash", "code": (f"head -c{64 * 1024 * 1024} /dev/urandom > random_data.bin",), "valves": { "MAX_MEGABYTES_PER_USER": 32, "MAX_RAM_MEGABYTES": 2048, }, "status": "STORAGE_ERROR", "post": _want_user_storage_num_files(16), }, { "name": "too_many_files_generated_in_single_execution", "language": "bash", "code": ("for i in $(seq 1 128); do touch file$i.txt; done",), "valves": { "MAX_FILES_PER_EXECUTION": 64, }, "status": "STORAGE_ERROR", "post": _want_user_storage_num_files(16), }, { "name": "too_many_files_for_user_quota", "language": "bash", "code": ("for i in $(seq 1 128); do touch file$i.txt; done",), "valves": { "MAX_FILES_PER_EXECUTION": 256, "MAX_FILES_PER_USER": 64, }, "status": "STORAGE_ERROR", "post": _want_user_storage_num_files(16), }, { "name": "clear_out_existing_user_storage", "language": "bash", "code": ("for i in $(seq 1 128); do touch file$i.txt; done",), "valves": { "MAX_FILES_PER_EXECUTION": 256, # Covers enough quota for the last test's files to not be # deleted (6 of them), but all previous tests' files will. # We add +2 which covers 2 of the previous-previous tests' # files, but since files from one execution are either all # deleted or not, this should not make a difference, so the # extra +2 should not add to the expected number of files. "MAX_FILES_PER_USER": 128 + 6 + 2, }, "status": "OK", "generated_files": {f"file{i}.txt": () for i in range(1, 128 + 1)}, # Verify that only this test + the last test's files are here: "post": _want_user_storage_num_files(128 + 6), }, { "name": "long_running_code", "language": "python", "code": ( "import time", "time.sleep(15)", "print('Managed to sleep for 15 seconds.')", ), "valves": { "MAX_RUNTIME_SECONDS": 5, }, "status": "TIMEOUT", }, { "name": "ram_hog", "language": "python", "code": ( "import time", "f = open('/dev/urandom', 'rb')", "s = []", "for i in range(256): s.append(f.read(1024 * 1024))", "time.sleep(1)", "print('\\n'.join(line for line in open('/proc/self/status').read().split('\\n') if line.startswith('Vm')))", "print('Managed to hog', len(s), 'megabytes.')", ), "valves": { "MAX_RAM_MEGABYTES": 128, }, "status": "INTERRUPTED", }, ) def _print_output(obj): if obj.stdout: print(" \U0001f5e8 Output:", file=sys.stderr) for stdout_line in obj.stdout.split("\n"): print(f" {stdout_line}") if obj.stderr: print(" \U0001f41e Debug:", file=sys.stderr) for stderr_line in obj.stderr.split("\n"): print(f" {stderr_line}") success = True with tempfile.TemporaryDirectory(prefix="sandbox_self_test") as tmp_dir: user_storage_path = os.path.join(tmp_dir, "user_storage") os.makedirs(user_storage_path, mode=0o755) valve_name_prefix = ( _Action.Valves()._VALVE_OVERRIDE_ENVIRONMENT_VARIABLE_NAME_PREFIX ) for self_test in _self_tests: name = self_test["name"] language = self_test["language"] code = "\n".join(self_test["code"]) + "\n" want_status = self_test["status"] valves = self_test.get("valves", {}) want_generated_files = self_test.get("generated_files", {}) if type(want_generated_files) is type({}): want_generated_files = _want_generated_files(want_generated_files) post_tests = self_test.get("post", ()) if type(post_tests) not in (type(()), type([])): post_tests = (post_tests,) test_env = os.environ.copy() test_env[f"{valve_name_prefix}WEB_ACCESSIBLE_DIRECTORY_PATH"] = ( user_storage_path ) for valve_name, valve_value in valves.items(): test_env[f"{valve_name_prefix}{valve_name}"] = str(valve_value) test_argv = [ sys.executable, os.path.abspath(__file__), f"--language={language}", ] if debug or self_test.get("debug", False): test_argv.append("--debug") print(f"\u23f3 Running self-test: {name}", file=sys.stderr) try: result = subprocess.run( test_argv, env=test_env, input=code, text=True, capture_output=True, timeout=20, check=True, ) except subprocess.CalledProcessError as e: success = False _print_output(e) print( f"\u274c Self-test {name} failed: process failed: {e}", file=sys.stderr, ) except Exception as e: success = False exception_class = e.__class__ print( f"\u274c Self-test {name} failed: {exception_class}: {e}", file=sys.stderr, ) else: try: result_data = json.loads(result.stdout) except json.decoder.JSONDecodeError as e: _print_output(result) success = False print( f"\u274c Self-test {name} failed: JSON decoding failed: {e}; got: {result.stdout}", file=sys.stderr, ) else: if debug: _print_output(result) got_status = result_data["status"] got_generated_files = result_data.get("generated_files", {}) if got_status != want_status: success = False print( f"\u274c Self-test {name} failed: status was {got_status}, expected {want_status}", file=sys.stderr, ) elif want_generated_files(got_generated_files) is not None: success = False generated_files_error = want_generated_files( got_generated_files ) print( f"\u274c Self-test {name} failed: generated files are incorrect: {generated_files_error}", file=sys.stderr, ) else: post_test_failure = None for post_test in post_tests: try: post_test() except Exception as e: post_test_failure = e break if post_test_failure is not None: success = False print( f"\u274c Self-test {name} failed: post-test verification failed: {post_test_failure}", file=sys.stderr, ) else: print(f"\u2714 Self-test {name} passed.", file=sys.stderr) if success: print("\u2705 All function self-tests passed, good go to!", file=sys.stderr) sys.exit(0) else: print("\u2620 One or more function self-tests failed.", file=sys.stderr) sys.exit(1) assert False, "Unreachable" # Debug utility: Run code from stdin if running as a normal Python script. if __name__ == "__main__": parser = argparse.ArgumentParser( description="Run arbitrary code in a gVisor sandbox." ) parser.add_argument( "--language", choices=("python", "bash"), default="python", help="Language of the code to run.", ) parser.add_argument( "--use_sample_code", action="store_true", default=False, help="Run sample code for the given language; otherwise, read code from stdin.", ) parser.add_argument( "--self_test", action="store_true", default=False, help="Run series of self-tests.", ) parser.add_argument( "--debug", action="store_true", default=False, help="Enable debug mode." ) parser.add_argument( "--want_status", type=str, default="", help="If set, verify that the code evaluation status matches this or exit with error code.", ) args = parser.parse_args() if args.debug: os.environ[ _Action.Valves()._VALVE_OVERRIDE_ENVIRONMENT_VARIABLE_NAME_PREFIX + "DEBUG" ] = "true" if args.self_test: _do_self_tests(args.debug) if args.use_sample_code: if args.language == "bash": code = "\n".join(_SAMPLE_BASH_INSTRUCTIONS) + "\n" else: code = "\n".join(_SAMPLE_PYTHON_INSTRUCTIONS) + "\n" else: code = sys.stdin.read() async def _local_run(): def _dummy_emitter(event): if not args.want_status: print(f"Event: {event}", file=sys.stderr) action = Action() body = { "messages": [ { "role": "assistant", "content": f"```{args.language}\n{code}\n```\n", }, ], } output_str = await action.action(body=body, __event_emitter__=_dummy_emitter) if args.want_status: output = json.loads(output_str) got_status = output["status"] if got_status != args.want_status: raise RuntimeError( f"Code evaluation status is {got_status} but expected {args.want_status}" ) print( f"\u2705 Code evaluation status is {got_status} as expected.", file=sys.stderr, ) else: print(output_str) asyncio.run(_local_run())