Disable memory swapping and add a memory limit test

If memory swapping was enabled locally, the memory test would fail.
Explicitly disabling swapping also removes reliance on the assumption
that it'll be disabled in production.

* Add a constant for the maximum memory
* Simplify the timeout test; it'd otherwise first run out of memory now

1 files changed, 8 insertions, 1 deletions

diff --git a/scripts/.profile b/scripts/.profile
index 415e4f6..bff260d 100644
--- a/scripts/.profile
+++ b/scripts/.profile
@@ -1,12 +1,19 @@
 nsjpy() {
+    local MEM_MAX=52428800
+
+    # All arguments except the last are considered to be for NsJail, not Python.
     local nsj_args=""
     while [ "$#" -gt 1 ]; do
         nsj_args="${nsj_args:+${nsj_args} }$1"
         shift
     done
 
+    # Set up cgroups and disable memory swapping.
     mkdir -p /sys/fs/cgroup/pids/NSJAIL
     mkdir -p /sys/fs/cgroup/memory/NSJAIL
+    echo "${MEM_MAX}" > /sys/fs/cgroup/memory/NSJAIL/memory.limit_in_bytes
+    echo "${MEM_MAX}" > /sys/fs/cgroup/memory/NSJAIL/memory.memsw.limit_in_bytes
+
     nsjail \
         -Mo \
         --rlimit_as 700 \
@@ -19,7 +26,7 @@ nsjpy() {
         --disable_proc \
         --iface_no_lo \
         --cgroup_pids_max=1 \
-        --cgroup_mem_max=52428800 \
+        --cgroup_mem_max="${MEM_MAX}" \
         $nsj_args -- \
         /snekbox/.venv/bin/python3 -Iq -c "$@"
 }