n8n/packages/@n8n/task-runner-python/src/task_executor.py
Emilia 95a81ed219
feat(core): Add N8N_RUNNERS_ALLOW_TRANSITIVE_IMPORTS for the Python task runner (#33130)
Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
2026-06-30 08:20:12 +00:00

673 lines
21 KiB
Python

import ast
import builtins
import collections
import importlib
import io
import json
import logging
import multiprocessing
import os
import sys
import textwrap
import traceback
from typing import cast
from src.errors import (
TaskCancelledError,
TaskKilledError,
TaskResultMissingError,
TaskResultReadError,
TaskRuntimeError,
TaskTimeoutError,
TaskSubprocessFailedError,
SecurityViolationError,
)
from src._sandbox_callables import _SafePrint, _GuardedImport, _SafeFormat
from src.format_validation import find_blocked_format_tokens
from src.import_validation import validate_module_import
from src.config.security_config import SecurityConfig
from src.message_types.broker import NodeMode, Items, Query
from src.message_types.pipe import (
PipeResultMessage,
PipeErrorMessage,
TaskErrorInfo,
PrintArgs,
)
from src.pipe_reader import PipeReader
from src.constants import (
EXECUTOR_CIRCULAR_REFERENCE_KEY,
EXECUTOR_SAFE_FORMAT_KEY,
EXECUTOR_USER_OUTPUT_KEY,
EXECUTOR_ALL_ITEMS_FILENAME,
EXECUTOR_PER_ITEM_FILENAME,
ERROR_DANGEROUS_STRING_PATTERN,
SIGTERM_EXIT_CODE,
SIGKILL_EXIT_CODE,
PIPE_MSG_PREFIX_LENGTH,
FORMAT_METHOD_NAMES,
)
from multiprocessing.context import ForkServerProcess
from multiprocessing.connection import Connection
logger = logging.getLogger(__name__)
MULTIPROCESSING_CONTEXT = multiprocessing.get_context("forkserver")
MAX_PRINT_ARGS_ALLOWED = 100
# Captured at module load before any allowlist guards are installed, so a
# fresh wrapper always delegates to the real implementations rather than
# stacking on top of a previously-installed wrapper.
_PRISTINE_IMPORT_MODULE = importlib.import_module
_PRISTINE_DUNDER_IMPORT = importlib.__import__
type PipeConnection = Connection
class FormatGuardTransformer(ast.NodeTransformer):
def visit_Call(self, node: ast.Call) -> ast.AST:
self.generic_visit(node)
if (
isinstance(node.func, ast.Attribute)
and node.func.attr in FORMAT_METHOD_NAMES
):
replacement = ast.Call(
func=ast.Name(id=EXECUTOR_SAFE_FORMAT_KEY, ctx=ast.Load()),
args=[
ast.Constant(value=node.func.attr),
node.func.value,
*node.args,
],
keywords=node.keywords,
)
return ast.copy_location(replacement, node)
return node
def _validate_format_template(template: str) -> None:
token = next(find_blocked_format_tokens(template), None)
if token is not None:
raise SecurityViolationError(
description=ERROR_DANGEROUS_STRING_PATTERN.format(attr=token),
)
def _validate_field_expression(expr: str) -> None:
# Wrap as a complete template so the existing parser can scan it.
_validate_format_template("{" + expr + "}")
_TEMPLATE_METHODS = frozenset({"format", "format_map", "vformat"})
_FIELD_METHODS = frozenset({"get_field"})
def _resolve_template_arg(method_name: str, receiver, args):
"""Return ``(template, field)`` for the call: at most one element is
non-``None``. Normalises bound and unbound call forms so the same
validation runs for both ``"tpl".format(...)``, ``str.format("tpl", ...)``,
``Formatter().format("tpl", ...)``, and ``Formatter.format(f, "tpl", ...)``.
The UserString short-circuit applies only to template methods; for
``get_field`` the field expression lives in ``args``, not in the
receiver's stored data.
"""
is_template = method_name in _TEMPLATE_METHODS
is_field = method_name in _FIELD_METHODS
if not is_template and not is_field:
return (None, None)
# ``str.format``/``UserString.format``: receiver itself is the template.
if is_template:
if isinstance(receiver, str):
return (receiver, None)
if isinstance(receiver, collections.UserString):
return (receiver.data, None)
if isinstance(receiver, type):
if issubclass(receiver, str):
# Unbound ``str.format(template, ...)``.
candidate = args[0] if args else None
elif issubclass(receiver, collections.UserString) and is_template:
# Unbound ``UserString.format(self, ...)`` — args[0] is the instance.
inst = args[0] if args else None
candidate = inst.data if isinstance(inst, collections.UserString) else None
else:
# Unbound on an arbitrary class (e.g. ``Formatter``): args[0] is
# the instance, the template/field lives at args[1].
candidate = args[1] if len(args) >= 2 else None
else:
# Bound on a non-str instance: ``args[0]`` is the template/field.
candidate = args[0] if args else None
if not isinstance(candidate, str):
return (None, None)
if is_field:
return (None, candidate)
return (candidate, None)
def _safe_format_impl(method_name: str, receiver, /, *args, **kwargs):
template, field = _resolve_template_arg(method_name, receiver, args)
if template is not None:
_validate_format_template(template)
if field is not None:
_validate_field_expression(field)
return getattr(receiver, method_name)(*args, **kwargs)
# Injected into user globals as a hardened callable so its ``__globals__``
# (which carries this module's sensitive namespace) is not reachable from user
# code; the implementation is held on a denied slot.
_safe_format = _SafeFormat(_safe_format_impl)
class TaskExecutor:
"""Responsible for executing Python code tasks in isolated subprocesses."""
@staticmethod
def create_process(
code: str,
node_mode: NodeMode,
items: Items,
security_config: SecurityConfig,
query: Query = None,
) -> tuple[ForkServerProcess, PipeConnection, PipeConnection]:
"""Create a subprocess for executing a Python code task and a pipe for communication."""
fn = (
TaskExecutor._all_items
if node_mode == "all_items"
else TaskExecutor._per_item
)
# thread in runner process reads, subprocess writes
read_conn, write_conn = MULTIPROCESSING_CONTEXT.Pipe(duplex=False)
process = MULTIPROCESSING_CONTEXT.Process(
target=fn,
args=(
code,
items,
write_conn,
security_config,
query,
),
)
return process, read_conn, write_conn
@staticmethod
def execute_process(
process: ForkServerProcess,
read_conn: PipeConnection,
write_conn: PipeConnection,
task_timeout: int,
continue_on_fail: bool,
) -> tuple[Items, PrintArgs, int]:
"""Execute a subprocess for a Python code task."""
print_args: PrintArgs = []
pipe_reader = PipeReader(read_conn.fileno(), read_conn)
pipe_reader.start()
try:
try:
process.start()
except Exception as e:
raise TaskSubprocessFailedError(-1, e)
finally:
write_conn.close()
process.join(timeout=task_timeout)
if process.is_alive():
TaskExecutor.stop_process(process)
raise TaskTimeoutError(task_timeout)
if process.exitcode == SIGTERM_EXIT_CODE:
raise TaskCancelledError()
if process.exitcode == SIGKILL_EXIT_CODE:
raise TaskKilledError()
if process.exitcode != 0:
assert process.exitcode is not None
raise TaskSubprocessFailedError(process.exitcode)
pipe_reader.join(timeout=task_timeout)
if pipe_reader.is_alive():
try:
read_conn.close()
except Exception:
pass
raise TaskResultReadError(
TimeoutError(f"Pipe reader timed out after {task_timeout}s")
)
if pipe_reader.error:
raise TaskResultReadError(pipe_reader.error)
if pipe_reader.pipe_message is None:
raise TaskResultMissingError()
returned = pipe_reader.pipe_message
if "error" in returned:
error_msg = cast(PipeErrorMessage, returned)
raise TaskRuntimeError(error_msg["error"])
if "result" not in returned:
raise TaskResultMissingError()
result_msg = cast(PipeResultMessage, returned)
result = result_msg["result"]
print_args = result_msg.get("print_args", [])
assert pipe_reader.message_size is not None
result_size_bytes = pipe_reader.message_size
return result, print_args, result_size_bytes
except Exception as e:
if continue_on_fail:
return [{"json": {"error": str(e)}}], print_args, 0
raise
@staticmethod
def stop_process(process: ForkServerProcess | None):
"""Stop a running subprocess, gracefully else force-killing."""
if process is None or not process.is_alive():
return
try:
process.terminate()
process.join(timeout=1) # 1s grace period
if process.is_alive():
process.kill()
process.join()
except (ProcessLookupError, ConnectionError, BrokenPipeError):
# subprocess is dead or unreachable
pass
@staticmethod
def _all_items(
raw_code: str,
items: Items,
write_conn,
security_config: SecurityConfig,
query: Query = None,
):
"""Execute a Python code task in all-items mode."""
if security_config.runner_env_deny:
os.environ.clear()
TaskExecutor._sanitize_sys_modules(security_config)
TaskExecutor._harden_importlib(security_config)
print_args: PrintArgs = []
sys.stderr = stderr_capture = io.StringIO()
try:
compiled_code = TaskExecutor._compile_user_code(
raw_code, EXECUTOR_ALL_ITEMS_FILENAME
)
globals = {
"__builtins__": TaskExecutor._filter_builtins(security_config),
"_items": items,
"_query": query,
"print": TaskExecutor._create_custom_print(print_args),
EXECUTOR_SAFE_FORMAT_KEY: _safe_format,
}
exec(compiled_code, globals)
result = cast(Items, globals[EXECUTOR_USER_OUTPUT_KEY])
TaskExecutor._put_result(write_conn.fileno(), result, print_args)
except BaseException as e:
TaskExecutor._put_error(
write_conn.fileno(), e, stderr_capture.getvalue(), print_args
)
@staticmethod
def _per_item(
raw_code: str,
items: Items,
write_conn,
security_config: SecurityConfig,
_query: Query = None, # unused, only to keep signatures consistent across modes
):
"""Execute a Python code task in per-item mode."""
if security_config.runner_env_deny:
os.environ.clear()
TaskExecutor._sanitize_sys_modules(security_config)
TaskExecutor._harden_importlib(security_config)
print_args: PrintArgs = []
sys.stderr = stderr_capture = io.StringIO()
try:
compiled_code = TaskExecutor._compile_user_code(
raw_code, EXECUTOR_PER_ITEM_FILENAME
)
filtered_builtins = TaskExecutor._filter_builtins(security_config)
custom_print = TaskExecutor._create_custom_print(print_args)
result: Items = []
for index, item in enumerate(items):
globals = {
"__builtins__": filtered_builtins,
"_item": item,
"print": custom_print,
EXECUTOR_SAFE_FORMAT_KEY: _safe_format,
}
exec(compiled_code, globals)
user_output = globals[EXECUTOR_USER_OUTPUT_KEY]
if user_output is None:
continue
json_data = TaskExecutor._extract_json_data_per_item(user_output)
output_item = {"json": json_data, "pairedItem": {"item": index}}
if isinstance(user_output, dict) and "binary" in user_output:
output_item["binary"] = user_output["binary"]
result.append(output_item)
TaskExecutor._put_result(write_conn.fileno(), result, print_args)
except BaseException as e:
TaskExecutor._put_error(
write_conn.fileno(), e, stderr_capture.getvalue(), print_args
)
@staticmethod
def _wrap_code(raw_code: str) -> str:
indented_code = textwrap.indent(raw_code, " ")
return f"def _user_function():\n{indented_code}\n\n{EXECUTOR_USER_OUTPUT_KEY} = _user_function()"
@staticmethod
def _compile_user_code(raw_code: str, filename: str):
wrapped_code = TaskExecutor._wrap_code(raw_code)
tree = ast.parse(wrapped_code, filename, "exec")
tree = FormatGuardTransformer().visit(tree)
ast.fix_missing_locations(tree)
return compile(tree, filename, "exec")
@staticmethod
def _extract_json_data_per_item(user_output):
if not isinstance(user_output, dict):
return user_output
if "json" in user_output:
return user_output["json"]
if "binary" in user_output:
return {k: v for k, v in user_output.items() if k != "binary"}
return user_output
@staticmethod
def _put_result(write_fd: int, result: Items, print_args: PrintArgs):
message: PipeResultMessage = {
"result": result,
"print_args": TaskExecutor._truncate_print_args(print_args),
}
data = json.dumps(message, default=str, ensure_ascii=False).encode("utf-8")
length_bytes = len(data).to_bytes(PIPE_MSG_PREFIX_LENGTH, "big")
try:
TaskExecutor._write_bytes(write_fd, length_bytes)
TaskExecutor._write_bytes(write_fd, data)
finally:
try:
os.close(write_fd)
except Exception:
pass
@staticmethod
def _put_error(
write_fd: int,
e: BaseException,
stderr: str = "",
print_args: PrintArgs | None = None,
):
if print_args is None:
print_args = []
task_error_info: TaskErrorInfo = {
"message": f"Process exited with code {e.code}"
if isinstance(e, SystemExit)
else str(e),
"description": getattr(e, "description", ""),
"stack": traceback.format_exc(),
"stderr": stderr,
}
message: PipeErrorMessage = {
"error": task_error_info,
"print_args": TaskExecutor._truncate_print_args(print_args),
}
data = json.dumps(message, default=str, ensure_ascii=False).encode("utf-8")
length_bytes = len(data).to_bytes(PIPE_MSG_PREFIX_LENGTH, "big")
try:
TaskExecutor._write_bytes(write_fd, length_bytes)
TaskExecutor._write_bytes(write_fd, data)
finally:
try:
os.close(write_fd)
except Exception:
pass
# ========== print() ==========
@staticmethod
def _create_custom_print(print_args: PrintArgs):
return _SafePrint(print_args, TaskExecutor._format_print_args)
@staticmethod
def _format_print_args(*args) -> list[str]:
"""
Takes the args passed to a `print()` call in user code and converts them
to string representations suitable for display in a browser console.
Expects all args to be serializable.
"""
formatted = []
for arg in args:
if isinstance(arg, str):
formatted.append(f"'{arg}'")
elif arg is None or isinstance(arg, (int, float, bool)):
formatted.append(str(arg))
elif isinstance(arg, dict) and EXECUTOR_CIRCULAR_REFERENCE_KEY in arg:
formatted.append(f"[Circular {arg.get('__type__', 'Object')}]")
else:
formatted.append(json.dumps(arg, default=str, ensure_ascii=False))
return formatted
@staticmethod
def _truncate_print_args(print_args: PrintArgs) -> PrintArgs:
"""Truncate print_args to prevent pipe buffer overflow."""
if not print_args or len(print_args) <= MAX_PRINT_ARGS_ALLOWED:
return print_args
truncated = print_args[:MAX_PRINT_ARGS_ALLOWED]
truncated.append(
[
f"[Output truncated - {len(print_args) - MAX_PRINT_ARGS_ALLOWED} more print statements]"
]
)
return truncated
# ========== security ==========
@staticmethod
def _filter_builtins(security_config: SecurityConfig):
"""Get __builtins__ with denied ones removed."""
if len(security_config.builtins_deny) == 0:
filtered = dict(__builtins__)
else:
filtered = {
k: v
for k, v in __builtins__.items()
if k not in security_config.builtins_deny
}
filtered["__import__"] = TaskExecutor._create_safe_import(security_config)
class _ImmutableBuiltins:
__slots__ = ()
def __getitem__(self, key):
return filtered[key]
def __contains__(self, key):
return key in filtered
def __iter__(self):
return iter(filtered)
def __len__(self):
return len(filtered)
def keys(self):
return filtered.keys()
def values(self):
return filtered.values()
def items(self):
return filtered.items()
def get(self, key, default=None):
return filtered.get(key, default)
def __getattr__(self, name):
try:
return filtered[name]
except KeyError:
raise AttributeError(name) from None
def __setattr__(self, name, value):
raise AttributeError("read-only")
def __delattr__(self, name):
raise AttributeError("read-only")
def __repr__(self):
return f"ImmutableBuiltins({len(filtered)} keys)"
return _ImmutableBuiltins()
@staticmethod
def _sanitize_sys_modules(security_config: SecurityConfig):
safe_modules = {
"builtins",
"__main__",
"sys",
"traceback",
"linecache",
"importlib",
"importlib.machinery",
}
if "*" in security_config.stdlib_allow:
safe_modules.update(sys.stdlib_module_names)
else:
safe_modules.update(security_config.stdlib_allow)
if "*" in security_config.external_allow:
safe_modules.update(
name
for name in sys.modules.keys()
if name not in sys.stdlib_module_names
)
else:
safe_modules.update(security_config.external_allow)
# keep modules marked as safe and submodules of those
safe_prefixes = [safe + "." for safe in safe_modules]
modules_to_remove = [
name
for name in sys.modules.keys()
if name not in safe_modules
and not any(name.startswith(prefix) for prefix in safe_prefixes)
]
for module_name in modules_to_remove:
del sys.modules[module_name]
@staticmethod
def _create_safe_import(security_config: SecurityConfig):
return _GuardedImport(
security_config, validate_module_import, builtins.__import__
)
@staticmethod
def _harden_importlib(security_config: SecurityConfig) -> None:
"""Route ``importlib``'s import entry points through the same
allowlist used for ``import`` statements in user code. The
replacement callables are class instances from
``src._sandbox_callables`` whose introspection-deny set blocks user
code from reaching back through ``__globals__`` etc."""
setattr(
importlib,
"import_module",
_GuardedImport(
security_config,
validate_module_import,
_PRISTINE_IMPORT_MODULE,
trust_eligible=True,
),
)
setattr(
importlib,
"__import__",
_GuardedImport(
security_config,
validate_module_import,
_PRISTINE_DUNDER_IMPORT,
trust_eligible=True,
),
)
# ========== pipe I/O ==========
@staticmethod
def _write_bytes(fd: int, data: bytes):
total_written = 0
while total_written < len(data):
written = os.write(fd, data[total_written:])
if written == 0:
raise OSError("Write failed")
total_written += written