From 66e6e32ebd6e7d047a679ad9fdaadffd98c609dd Mon Sep 17 00:00:00 2001
From: virusdefender <virusdefender@outlook.com>
Date: Fri, 23 Mar 2018 18:10:00 +0800
Subject: [PATCH] add memory limit check only

---
 bindings/Python/_judger/__init__.py                |  5 +++--
 src/child.c                                        | 13 ++++++++-----
 src/main.c                                         |  9 ++++++++-
 src/runner.h                                       |  1 +
 tests/Python_and_core/testcase/integration/test.py | 10 ++++++++++
 5 files changed, 30 insertions(+), 8 deletions(-)

diff --git a/bindings/Python/_judger/__init__.py b/bindings/Python/_judger/__init__.py
index 267fbd2..7405016 100644
--- a/bindings/Python/_judger/__init__.py
+++ b/bindings/Python/_judger/__init__.py
@@ -40,11 +40,12 @@ def run(max_cpu_time,
         log_path,
         seccomp_rule_name,
         uid,
-        gid):
+        gid,
+        memory_limit_check_only=0):
     str_list_vars = ["args", "env"]
     int_vars = ["max_cpu_time", "max_real_time",
                 "max_memory", "max_stack", "max_output_size",
-                "max_process_number", "uid", "gid"]
+                "max_process_number", "uid", "gid", "memory_limit_check_only"]
     str_vars = ["exe_path", "input_path", "output_path", "error_path", "log_path"]
 
     proc_args = ["/usr/lib/judger/libjudger.so"]
diff --git a/src/child.c b/src/child.c
index 1451e42..8909e2d 100644
--- a/src/child.c
+++ b/src/child.c
@@ -48,11 +48,14 @@ void child_process(FILE *log_fp, struct config *_config) {
     }
 
     // set memory limit
-    if (_config->max_memory != UNLIMITED) {
-        struct rlimit max_memory;
-        max_memory.rlim_cur = max_memory.rlim_max = (rlim_t) (_config->max_memory) * 2;
-        if (setrlimit(RLIMIT_AS, &max_memory) != 0) {
-            CHILD_ERROR_EXIT(SETRLIMIT_FAILED);
+    // if memory_limit_check_only == 0, we only check memory usage number, because setrlimit(maxrss) will cause some crash issues
+    if (_config->memory_limit_check_only == 0) {
+        if (_config->max_memory != UNLIMITED) {
+            struct rlimit max_memory;
+            max_memory.rlim_cur = max_memory.rlim_max = (rlim_t) (_config->max_memory) * 2;
+            if (setrlimit(RLIMIT_AS, &max_memory) != 0) {
+                CHILD_ERROR_EXIT(SETRLIMIT_FAILED);
+            }
         }
     }
 
diff --git a/src/main.c b/src/main.c
index 7e938b1..c16d957 100644
--- a/src/main.c
+++ b/src/main.c
@@ -5,7 +5,7 @@
 #define STR_PLACE_HOLDER "<str>"
 
 struct arg_lit *verb, *help, *version;
-struct arg_int *max_cpu_time, *max_real_time, *max_memory, *max_stack,
+struct arg_int *max_cpu_time, *max_real_time, *max_memory, *max_stack, *memory_limit_check_only,
         *max_process_number, *max_output_size, *uid, *gid;
 struct arg_str *exe_path, *input_path, *output_path, *error_path, *args, *env, *log_path, *seccomp_rule_name;
 struct arg_end *end;
@@ -17,6 +17,7 @@ int main(int argc, char *argv[]) {
             max_cpu_time = arg_intn(NULL, "max_cpu_time", INT_PLACE_HOLDER, 0, 1, "Max CPU Time (ms)"),
             max_real_time = arg_intn(NULL, "max_real_time", INT_PLACE_HOLDER, 0, 1, "Max Real Time (ms)"),
             max_memory = arg_intn(NULL, "max_memory", INT_PLACE_HOLDER, 0, 1, "Max Memory (byte)"),
+            memory_limit_check_only = arg_intn(NULL, "memory_limit_check_only", INT_PLACE_HOLDER, 0, 1, "only check memory usage, do not setrlimit (default False)"),
             max_stack = arg_intn(NULL, "max_stack", INT_PLACE_HOLDER, 0, 1, "Max Stack (byte, default 16M)"),
             max_process_number = arg_intn(NULL, "max_process_number", INT_PLACE_HOLDER, 0, 1, "Max Process Number"),
             max_output_size = arg_intn(NULL, "max_output_size", INT_PLACE_HOLDER, 0, 1, "Max Output Size (byte)"),
@@ -83,6 +84,12 @@ int main(int argc, char *argv[]) {
         _config.max_memory = UNLIMITED;
     }
 
+    if (memory_limit_check_only->count > 0) {
+        _config.memory_limit_check_only = *memory_limit_check_only->ival == 0 ? 0 : 1;
+    } else {
+        _config.memory_limit_check_only = 0;
+    }
+
     if (max_stack->count > 0) {
         _config.max_stack = (long) *max_stack->ival;
     } else {
diff --git a/src/runner.h b/src/runner.h
index 4e7205d..bfefb44 100644
--- a/src/runner.h
+++ b/src/runner.h
@@ -46,6 +46,7 @@ struct config {
     long max_stack;
     int max_process_number;
     long max_output_size;
+    int memory_limit_check_only;
     char *exe_path;
     char *input_path;
     char *output_path;
diff --git a/tests/Python_and_core/testcase/integration/test.py b/tests/Python_and_core/testcase/integration/test.py
index 355775f..4a6f6d2 100644
--- a/tests/Python_and_core/testcase/integration/test.py
+++ b/tests/Python_and_core/testcase/integration/test.py
@@ -218,6 +218,16 @@ class IntegrationTest(base.BaseTestCase):
         self.assertTrue(result["memory"] < 3000000)
         del a
 
+    def test_memory_limit_check_only(self):
+        config = self.base_config
+        config["memory_limit_check_only"] = 1
+        config["max_memory"] = 64 * 1024 * 1024
+        config["exe_path"] = self._compile_c("memory2.c")
+        result = _judger.run(**config)
+        self.assertEqual(result["exit_code"], 0)
+        self.assertTrue(result["memory"] > 256 * 1024 * 1024)
+        self.assertEqual(result["result"], _judger.RESULT_MEMORY_LIMIT_EXCEEDED)
+
     def test_re1(self):
         config = self.base_config
         config["exe_path"] = self._compile_c("re1.c")