add jailed exec which raises timeout after a number of instructions have been interpreted

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];

include/pocketpy/pocketpy.h

@@ -735,6 +735,7 @@ enum py_PredefinedTypes {
     tp_ImportError,
     tp_AssertionError,
     tp_KeyError,
+    tp_Timeout,
     /* linalg */
     tp_vec2,
     tp_vec3,

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();
             /*****************************************/

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

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);
 

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