From aa34a52993f971b5e8abeb620df768c470d98308 Mon Sep 17 00:00:00 2001
From: messense <messense@icloud.com>
Date: Thu, 20 Dec 2018 11:19:56 +0800
Subject: [PATCH] Add timeout and max_memory limit support

---
 .../extension/mini_racer_extension.cc         | 76 ++++++++++++++++---
 py_mini_racer/py_mini_racer.py                | 40 ++++++++--
 tests/test_eval.py                            | 21 +++++
 3 files changed, 121 insertions(+), 16 deletions(-)

diff --git a/py_mini_racer/extension/mini_racer_extension.cc b/py_mini_racer/extension/mini_racer_extension.cc
index ccbc5c97..c6929451 100644
--- a/py_mini_racer/extension/mini_racer_extension.cc
+++ b/py_mini_racer/extension/mini_racer_extension.cc
@@ -35,6 +35,8 @@ enum BinaryTypes {
 
     type_execute_exception = 200,
     type_parse_exception   = 201,
+    type_oom_exception = 202,
+    type_timeout_exception = 203,
 };
 
 /* This is a generic store for arbitrary JSON like values.
@@ -63,6 +65,8 @@ void BinaryValueFree(BinaryValue *v) {
     switch(v->type) {
     case type_execute_exception:
     case type_parse_exception:
+    case type_oom_exception:
+    case type_timeout_exception:
     case type_str_utf8:
         free(v->str_val);
         break;
@@ -119,6 +123,7 @@ struct EvalResult {
     bool parsed;
     bool executed;
     bool terminated;
+    bool timed_out;
     Persistent<Value>* value;
     Persistent<Value>* message;
     Persistent<Value>* backtrace;
@@ -144,10 +149,31 @@ typedef struct {
     Local<String>* eval;
     useconds_t timeout;
     EvalResult* result;
+    size_t max_memory;
 } EvalParams;
 
+enum IsolateFlags {
+    MEM_SOFTLIMIT_VALUE,
+    MEM_SOFTLIMIT_REACHED,
+};
+
 static Platform* current_platform = NULL;
 
+static void gc_callback(Isolate *isolate, GCType type, GCCallbackFlags flags) {
+    if((bool)isolate->GetData(MEM_SOFTLIMIT_REACHED)) return;
+
+    size_t softlimit = *(size_t*) isolate->GetData(MEM_SOFTLIMIT_VALUE);
+
+    HeapStatistics stats;
+    isolate->GetHeapStatistics(&stats);
+    size_t used = stats.used_heap_size();
+
+    if(used > softlimit) {
+        isolate->SetData(MEM_SOFTLIMIT_REACHED, (void*)true);
+        isolate->TerminateExecution();
+    }
+}
+
 static void init_v8() {
     if (current_platform == NULL) {
         V8::InitializeICU();
@@ -161,6 +187,7 @@ static void* breaker(void *d) {
   EvalParams* data = (EvalParams*)d;
   usleep(data->timeout*1000);
   pthread_setcancelstate(PTHREAD_CANCEL_DISABLE, NULL);
+  data->result->timed_out = true;
   data->context_info->isolate->TerminateExecution();
   return NULL;
 }
@@ -178,10 +205,16 @@ static void* nogvl_context_eval(void* arg) {
 
     Context::Scope context_scope(context);
 
+    // Memory softlimit
+    isolate->SetData(MEM_SOFTLIMIT_VALUE, (void*)false);
+    // Memory softlimit hit flag
+    isolate->SetData(MEM_SOFTLIMIT_REACHED, (void*)false);
+
     MaybeLocal<Script> parsed_script = Script::Compile(context, *eval_params->eval);
     result->parsed = !parsed_script.IsEmpty();
     result->executed = false;
     result->terminated = false;
+    result->timed_out = false;
     result->value = NULL;
 
     if (!result->parsed) {
@@ -191,10 +224,16 @@ static void* nogvl_context_eval(void* arg) {
 
         pthread_t breaker_thread = 0;
 
+        // timeout limit
         auto timeout = eval_params->timeout;
         if (timeout > 0) {
             pthread_create(&breaker_thread, NULL, breaker, (void*)eval_params);
         }
+        // memory limit
+        if (eval_params->max_memory > 0) {
+            isolate->SetData(MEM_SOFTLIMIT_VALUE, &eval_params->max_memory);
+            isolate->AddGCEpilogueCallback(gc_callback);
+        }
 
         MaybeLocal<Value> maybe_value = parsed_script.ToLocalChecked()->Run(context);
 
@@ -204,7 +243,6 @@ static void* nogvl_context_eval(void* arg) {
         }
 
         result->executed = !maybe_value.IsEmpty();
-
         if (!result->executed) {
             if (trycatch.HasCaught()) {
                 if (!trycatch.Exception()->IsNull()) {
@@ -213,7 +251,12 @@ static void* nogvl_context_eval(void* arg) {
                 } else if(trycatch.HasTerminated()) {
                     result->terminated = true;
                     result->message = new Persistent<Value>();
-                    Local<String> tmp = String::NewFromUtf8(isolate, "JavaScript was terminated (either by timeout or explicitly)");
+                    Local<String> tmp;
+                    if (result->timed_out) {
+                        tmp = String::NewFromUtf8(isolate, "JavaScript was terminated by timeout");
+                    } else {
+                        tmp = String::NewFromUtf8(isolate, "JavaScript was terminated");
+                    }
                     result->message->Reset(isolate, tmp);
                 }
 
@@ -497,7 +540,8 @@ ContextInfo *MiniRacer_init_context()
 
 static BinaryValue* MiniRacer_eval_context_unsafe(
         ContextInfo *context_info,
-        char *utf_str, int   str_len)
+        char *utf_str, int str_len,
+        unsigned long timeout, size_t max_memory)
 {
     EvalParams eval_params;
     EvalResult eval_result{};
@@ -529,9 +573,13 @@ static BinaryValue* MiniRacer_eval_context_unsafe(
         eval_params.eval = &eval;
         eval_params.result = &eval_result;
         eval_params.timeout = 0;
-
-        // FIXME - we should allow setting a timeout here
-        //     eval_params.timeout = (useconds_t)NUM2LONG(timeout);
+        eval_params.max_memory = 0;
+        if (timeout > 0) {
+            eval_params.timeout = (useconds_t)timeout;
+        }
+        if (max_memory > 0) {
+            eval_params.max_memory = max_memory;
+        }
 
         nogvl_context_eval(&eval_params);
 
@@ -571,9 +619,19 @@ static BinaryValue* MiniRacer_eval_context_unsafe(
 
     else if (!eval_result.executed) {
         result = xalloc(result);
-        result->type = type_execute_exception;
         result->str_val = nullptr;
 
+        bool mem_softlimit_reached = (bool)context_info->isolate->GetData(MEM_SOFTLIMIT_REACHED);
+        if (mem_softlimit_reached) {
+            result->type = type_oom_exception;
+        } else {
+            if (eval_result.timed_out) {
+                result->type = type_timeout_exception;
+            } else {
+                result->type = type_execute_exception;
+            }
+        }
+
         if (bmessage && bmessage->type == type_str_utf8 &&
                 bbacktrace && bbacktrace->type == type_str_utf8) {
             // +1 for \n, +1 for NUL terminator
@@ -639,8 +697,8 @@ class BufferOutputStream: public OutputStream {
 
 extern "C" {
 
-BinaryValue* mr_eval_context(ContextInfo *context_info, char *str, int len) {
-    BinaryValue *res = MiniRacer_eval_context_unsafe(context_info, str, len);
+BinaryValue* mr_eval_context(ContextInfo *context_info, char *str, int len, unsigned long timeout, size_t max_memory) {
+    BinaryValue *res = MiniRacer_eval_context_unsafe(context_info, str, len, timeout, max_memory);
     return res;
 }
 
diff --git a/py_mini_racer/py_mini_racer.py b/py_mini_racer/py_mini_racer.py
index 588f7436..4d4e40f3 100644
--- a/py_mini_racer/py_mini_racer.py
+++ b/py_mini_racer/py_mini_racer.py
@@ -35,6 +35,14 @@ class JSEvalException(MiniRacerBaseException):
     """ JS could not be executed """
     pass
 
+class JSOOMException(JSEvalException):
+    """ JS execution out of memory """
+    pass
+
+class JSTimeoutException(JSEvalException):
+    """ JS execution timed out """
+    pass
+
 class JSConversionException(MiniRacerBaseException):
     """ type could not be converted """
     pass
@@ -88,7 +96,9 @@ def _fetch_ext_handle():
     _ext_handle.mr_eval_context.argtypes = [
         ctypes.c_void_p,
         ctypes.c_char_p,
-        ctypes.c_int]
+        ctypes.c_int,
+        ctypes.c_ulong,
+        ctypes.c_size_t]
     _ext_handle.mr_eval_context.restype = ctypes.POINTER(PythonValue)
 
     _ext_handle.mr_free_value.argtypes = [ctypes.c_void_p]
@@ -123,13 +133,13 @@ def free(self, res):
 
         self.ext.mr_free_value(res)
 
-    def execute(self, js_str):
+    def execute(self, js_str, timeout=0, max_memory=0):
         """ Exec the given JS value """
 
         wrapped = "(function(){return (%s)})()" % js_str
-        return self.eval(wrapped)
+        return self.eval(wrapped, timeout, max_memory)
 
-    def eval(self, js_str):
+    def eval(self, js_str, timeout=0, max_memory=0):
         """ Eval the JavaScript string """
 
         if is_unicode(js_str):
@@ -137,9 +147,14 @@ def eval(self, js_str):
         else:
             bytes_val = js_str
 
+        res = None
         self.lock.acquire()
         try:
-            res = self.ext.mr_eval_context(self.ctx, bytes_val, len(bytes_val))
+            res = self.ext.mr_eval_context(self.ctx,
+                                           bytes_val,
+                                           len(bytes_val),
+                                           ctypes.c_ulong(timeout),
+                                           ctypes.c_size_t(max_memory))
 
             if bool(res) is False:
                 raise JSConversionException()
@@ -147,7 +162,8 @@ def eval(self, js_str):
             return python_value
         finally:
             self.lock.release()
-            self.free(res)
+            if res is not None:
+                self.free(res)
 
     def call(self, identifier, *args, **kwargs):
         """ Call the named function with provided arguments
@@ -156,10 +172,12 @@ def call(self, identifier, *args, **kwargs):
         """
 
         encoder = kwargs.get('encoder', None)
+        timeout = kwargs.get('timeout', 0)
+        max_memory = kwargs.get('max_memory', 0)
 
         json_args = json.dumps(args, separators=(',', ':'), cls=encoder)
         js = "{identifier}.apply(this, {json_args})"
-        return self.eval(js.format(identifier=identifier, json_args=json_args))
+        return self.eval(js.format(identifier=identifier, json_args=json_args), timeout, max_memory)
 
     def heap_stats(self):
         """ Return heap statistics """
@@ -210,6 +228,8 @@ class PythonTypes(object):
 
     execute_exception = 200
     parse_exception = 201
+    oom_exception = 202
+    timeout_exception = 203
 
 
 class PythonValue(ctypes.Structure):
@@ -273,6 +293,12 @@ def to_python(self):
         elif self.type == PythonTypes.execute_exception:
             msg = ctypes.c_char_p(self.value).value
             raise JSEvalException(msg.decode('utf-8', errors='replace'))
+        elif self.type == PythonTypes.oom_exception:
+            msg = ctypes.c_char_p(self.value).value
+            raise JSOOMException(msg)
+        elif self.type == PythonTypes.timeout_exception:
+            msg = ctypes.c_char_p(self.value).value
+            raise JSTimeoutException(msg)
         elif self.type == PythonTypes.date:
             timestamp = self._double_value()
             # JS timestamp are milliseconds, in python we are in seconds
diff --git a/tests/test_eval.py b/tests/test_eval.py
index 4a55deee..010bfa2e 100644
--- a/tests/test_eval.py
+++ b/tests/test_eval.py
@@ -3,6 +3,7 @@
 
 """ Basic JS types tests """
 
+import time
 import unittest
 import six
 
@@ -75,6 +76,26 @@ def test_null_byte(self):
         result = context.eval(in_val)
         self.assertEqual(result, s)
 
+    def test_timeout(self):
+        timeout_ms = 100
+        with self.assertRaises(py_mini_racer.JSTimeoutException):
+            start_time = time.clock()
+            self.mr.eval('while(1) { }', timeout=timeout_ms)
+            duration = time.clock() - start_time
+            assert timeout_ms <= duration * 1000 <= timeout_ms + 10
+
+    def test_max_memory(self):
+        with self.assertRaises(py_mini_racer.JSOOMException):
+            self.mr.eval('''let s = 1000;
+                var a = new Array(s);
+                a.fill(0);
+                while(true) {
+                    s *= 1.1;
+                    let n = new Array(Math.floor(s));
+                    n.fill(0);
+                    a = a.concat(n);
+                }''', max_memory=200000000)
+
 
 if __name__ == '__main__':
     import sys