feat: add resource limit

err.h

@@ -8,5 +8,7 @@
 #define ERR_SECCOMP_LOAD    13
 #define ERR_UNSETENV        14
 #define ERR_NO_RULE_FOUND   15
+#define ERR_RLIMIT_ENV      16
+#define ERR_RLIMIT_SET      17
 
 #endif // WOJ_SANDBOX_ERR_H

library.c

@@ -1,3 +1,4 @@
+#include "resource.h"
 #include "rules/lang.h"
 #include "sandbox.h"
 #include "utils/log.h"
@@ -5,5 +6,6 @@
 static __attribute__((constructor)) void inject(void) {
     LOG_INFO("Setting up...");
     register_lang_c();
+    setup_rlimit();
     setup_seccomp();
 }

resource.c

@@ -0,0 +1,58 @@
+#include "resource.h"
+#include "err.h"
+#include "utils/log.h"
+
+#include <stdlib.h>
+#include <sys/resource.h>
+
+#define UNSET_ENV(name)                                               \
+    do {                                                              \
+        if (unsetenv(name)) {                                         \
+            LOG_ERR("Failed to unset environment variable %s", name); \
+            exit(ERR_UNSETENV);                                       \
+        }                                                             \
+    } while (0)
+
+#define SET_LIMIT(resource, limit, name, unit)                         \
+    do {                                                               \
+        if (limit) {                                                   \
+            LOG_INFO("Setting " name " limit to %ld " unit, limit);    \
+            if (setrlimit(resource, &(struct rlimit){limit, limit})) { \
+                LOG_ERR("Failed to set " name " limit");               \
+                exit(ERR_RLIMIT_SET);                                  \
+            }                                                          \
+        }                                                              \
+    } while (0)
+
+void setup_rlimit(void) {
+    LOG_INFO("Setting resource limit");
+
+    char *mem_limit_str   = getenv(LIMIT_MEMORY); // in mb
+    char *nproc_limit_str = getenv(LIMIT_NPROC);
+    char *time_limit_str  = getenv(LIMIT_TIME); // in ms
+
+    long mem_limit = 0, nproc_limit = 0, time_limit = 0;
+
+    if (mem_limit_str) {
+        // convert to bytes and double the memory limit
+        mem_limit = strtol(mem_limit_str, NULL, 10) * 1024 * 1024 * 2;
+    }
+
+    if (nproc_limit_str) {
+        nproc_limit = strtol(nproc_limit_str, NULL, 10);
+    }
+
+    if (time_limit_str) {
+        // add 2 seconds to avoid time limit exceeded
+        long limit = strtol(time_limit_str, NULL, 10);
+        time_limit = limit ? (limit / 1000 + 2) : 0;
+    }
+
+    SET_LIMIT(RLIMIT_AS, mem_limit, "memory", "bytes");
+    SET_LIMIT(RLIMIT_NPROC, nproc_limit, "nproc", "processes"); // per user
+    SET_LIMIT(RLIMIT_CPU, time_limit, "time", "seconds"); // except blocked time
+
+    UNSET_ENV(LIMIT_MEMORY);
+    UNSET_ENV(LIMIT_NPROC);
+    UNSET_ENV(LIMIT_TIME);
+}

resource.h

@@ -0,0 +1,11 @@
+#ifndef WOJ_SANDBOX_RESOURCE_H
+#define WOJ_SANDBOX_RESOURCE_H
+
+// Configuration Environment Variables
+#define LIMIT_MEMORY "LIMIT_MEMORY"
+#define LIMIT_TIME   "LIMIT_TIME"
+#define LIMIT_NPROC  "LIMIT_NPROC"
+
+void setup_rlimit(void);
+
+#endif // WOJ_SANDBOX_RESOURCE_H

rules/lang_c.c

@@ -5,8 +5,11 @@
 #include <seccomp.h>
 
 void setup_lang_c(scmp_filter_ctx ctx) {
-    int white[]   = {SCMP_SYS(read),  SCMP_SYS(write),      SCMP_SYS(getpid),
+    int white[] = {
-                     SCMP_SYS(futex), SCMP_SYS(exit_group), SCMP_SYS(newfstatat)};
+        SCMP_SYS(read),  SCMP_SYS(write),      SCMP_SYS(lseek),
+        SCMP_SYS(mmap),  SCMP_SYS(munmap),     SCMP_SYS(getpid),
+        SCMP_SYS(futex), SCMP_SYS(newfstatat), SCMP_SYS(exit_group),
+    };
     int white_len = sizeof(white) / sizeof(white[0]);
 
     for (int i = 0; i < white_len; i++) {

rules/rules.c

@@ -12,6 +12,7 @@ void register_rule(struct rule *rule) { list_add(&rule->list, &seccomp_rules); }
 void setup_rule(char *name, scmp_filter_ctx ctx) {
     struct list_head *current;
     struct rule      *rule;
+
     list_for_each(current, &seccomp_rules) {
         rule = list_entry(current, struct rule, list);
         if (strcmp(rule->name, name) == 0) {
@@ -19,6 +20,7 @@ void setup_rule(char *name, scmp_filter_ctx ctx) {
             return;
         }
     }
+
     LOG_ERR("No rule found for %s", name);
     dump_rules();
     exit(ERR_NO_RULE_FOUND);

sandbox.c

@@ -27,6 +27,15 @@ void add_syscall_name(const char *syscall_name, scmp_filter_ctx ctx,
     add_syscall_nr(syscall_nr, ctx, action);
 }
 
+#define UNSET_ENV(name)                                               \
+    do {                                                              \
+        if (unsetenv(name)) {                                         \
+            LOG_ERR("Failed to unset environment variable %s", name); \
+            seccomp_release(ctx);                                     \
+            exit(ERR_UNSETENV);                                       \
+        }                                                             \
+    } while (0)
+
 void setup_seccomp(void) {
     LOG_INFO("Setting seccomp rules...");
 
@@ -51,21 +60,10 @@ void setup_seccomp(void) {
         exit(ERR_SECCOMP_INIT);
     }
 
-    setup_rule(template, ctx);
+    if (template) setup_rule(template, ctx);
 
-#define UNSETENV(name)                                                \
+    UNSET_ENV(SANDBOX_TEMPLATE);
-    do {                                                              \
+    UNSET_ENV(SANDBOX_ACTION);
-        if (unsetenv(name)) {                                         \
-            LOG_ERR("Failed to unset environment variable %s", name); \
-            seccomp_release(ctx);                                     \
-            exit(ERR_UNSETENV);                                       \
-        }                                                             \
-    } while (0)
-
-    UNSETENV(SANDBOX_TEMPLATE);
-    UNSETENV(SANDBOX_ACTION);
-
-#undef UNSETENV
 
     if (seccomp_load(ctx)) {
         LOG_ERR("Failed to load seccomp context");

test.c

@@ -1,9 +1,31 @@
 #include "utils/log.h"
+#include <stdlib.h>
 
 int main() {
-    char buf[128];
+    LOG_INFO("Testing Memory Limit");
-    scanf("%s", buf);
+    void *p = malloc(sizeof(int) * 1024 * 1024 * 10);
-    printf("Hello %s\n", buf);
+    if (!p) {
+        LOG_ERR("malloc failed");
+    }
 
+    LOG_INFO("Testing NPROC Limit");
+    pid_t pid = fork();
+    if (pid == -1) {
+        perror("fork failed");
+    } else if (pid == 0) {
+        LOG_INFO("Child process");
+        exit(0);
+    } else {
+        LOG_INFO("Parent process");
+    }
+
+    LOG_INFO("Testing Time Limit 1");
+    sleep(5);
+
+    LOG_INFO("Testing Time Limit 2");
+    for (volatile int i = 0; i != -1; i++)
+        ;
+
+    LOG_INFO("Exiting...");
     return 0;
 }