add memory limit check only

This commit is contained in:
virusdefender 2018-03-23 18:10:00 +08:00
parent 4060c41384
commit 66e6e32ebd
5 changed files with 30 additions and 8 deletions

View File

@ -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"]

View File

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

View File

@ -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 {

View File

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

View File

@ -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")