diff --git a/include/pocketpy/interpreter/vm.h b/include/pocketpy/interpreter/vm.h index 92f47cec..81dee1d3 100644 --- a/include/pocketpy/interpreter/vm.h +++ b/include/pocketpy/interpreter/vm.h @@ -26,6 +26,9 @@ typedef struct VM { py_TValue builtins; // builtins module py_TValue main; // __main__ module + py_i64 max_steps; + py_i64 used_steps; + py_Callbacks callbacks; py_TValue ascii_literals[128+1]; @@ -136,4 +139,4 @@ py_Type pk_code__register(); py_TValue pk_builtins__register(); /* mappingproxy */ -void pk_mappingproxy__namedict(py_Ref out, py_Ref object); \ No newline at end of file +void pk_mappingproxy__namedict(py_Ref out, py_Ref object); diff --git a/include/pocketpy/pocketpy.h b/include/pocketpy/pocketpy.h index 30561501..267675b9 100644 --- a/include/pocketpy/pocketpy.h +++ b/include/pocketpy/pocketpy.h @@ -735,6 +735,7 @@ enum py_PredefinedTypes { tp_ImportError, tp_AssertionError, tp_KeyError, + tp_Timeout, /* linalg */ tp_vec2, tp_vec3, diff --git a/src/interpreter/ceval.c b/src/interpreter/ceval.c index 2b337614..dec2f112 100644 --- a/src/interpreter/ceval.c +++ b/src/interpreter/ceval.c @@ -109,6 +109,15 @@ FrameResult VM__run_top_frame(VM* self) { goto __ERROR; } + if (self->max_steps > 0) { + if (self->used_steps >= self->max_steps) { + py_exception(tp_Timeout, "execution timed out"); + goto __ERROR; + } + self->used_steps++; + } + + switch((Opcode)byte.op) { case OP_NO_OP: DISPATCH(); /*****************************************/ @@ -1440,4 +1449,4 @@ static bool stack_format_object(VM* self, c11_sv spec) { #undef POPX #undef SP #undef INSERT_THIRD -#undef vectorcall_opcall \ No newline at end of file +#undef vectorcall_opcall diff --git a/src/interpreter/vm.c b/src/interpreter/vm.c index 3c76f9fc..1b2fe381 100644 --- a/src/interpreter/vm.c +++ b/src/interpreter/vm.c @@ -176,6 +176,7 @@ void VM__ctor(VM* self) { INJECT_BUILTIN_EXC(ImportError, tp_Exception); INJECT_BUILTIN_EXC(AssertionError, tp_Exception); INJECT_BUILTIN_EXC(KeyError, tp_Exception); + INJECT_BUILTIN_EXC(Timeout, tp_Exception); #undef INJECT_BUILTIN_EXC #undef validate @@ -812,4 +813,4 @@ int py_replinput(char* buf, int max_size) { buf[size] = '\0'; return size; -} \ No newline at end of file +} diff --git a/src/public/modules.c b/src/public/modules.c index 9adcd670..8244e7fa 100644 --- a/src/public/modules.c +++ b/src/public/modules.c @@ -612,6 +612,20 @@ static bool builtins_exec(int argc, py_Ref argv) { return ok; } +static bool builtins_exec_jailed(int argc, py_Ref argv) { + VM* vm = pk_current_vm; + PY_CHECK_ARG_TYPE(0, tp_int); + int was_jailed = (vm->max_steps > 0); + if (!was_jailed) { + vm->max_steps = py_toint(py_arg(0)); + vm->used_steps = 0; + } + bool ok = _builtins_execdyn("exec", argc - 1, argv + 1, EXEC_MODE); + if (!was_jailed) vm->max_steps = 0; + py_newnone(py_retval()); + return ok; +} + static bool builtins_eval(int argc, py_Ref argv) { return _builtins_execdyn("eval", argc, argv, EVAL_MODE); } @@ -755,6 +769,7 @@ py_TValue pk_builtins__register() { py_bindfunc(builtins, "globals", builtins_globals); py_bindfunc(builtins, "locals", builtins_locals); py_bindfunc(builtins, "exec", builtins_exec); + py_bindfunc(builtins, "exec_jailed", builtins_exec_jailed); py_bindfunc(builtins, "eval", builtins_eval); py_bindfunc(builtins, "compile", builtins_compile); diff --git a/tests/66_eval.py b/tests/66_eval.py index 2a51c22a..dd41905f 100644 --- a/tests/66_eval.py +++ b/tests/66_eval.py @@ -74,3 +74,11 @@ exec(code, {'x': 42, 'res': res}) assert res == [42, 42] assert x == 33 +code = ''' +while True: + pass +''' +try: + exec_jailed(100000, code) +except Timeout: + pass