nsjail workaround, fix forkbomb, add forkbomb unit test

4 files changed, 26 insertions, 19 deletions

diff --git a/.travis.yml b/.travis.yml
index a91ecba..1c053fb 100644
--- a/.travis.yml
+++ b/.travis.yml
@@ -19,7 +19,6 @@ before_install:
 
 env:
   global:
-    - PYTHONEXECUTABLE=/home/travis/build/discord-python/snekbox/.venv/bin/python3
     - PIPENV_VENV_IN_PROJECT=1
     - PIPENV_IGNORE_VIRTUALENVS=1
     - PIPENV_NOSPIN=1
diff --git a/README.md b/README.md
index 45fa887..d0b565d 100644
--- a/README.md
+++ b/README.md
@@ -57,6 +57,10 @@ chmod +x /usr/bin/nsjail
 give nsjail a test run
 
 ```bash
+# This is a workaround because nsjail can't create the directories automatically
+sudo mkdir -p /sys/fs/cgroup/pids/NSJAIL \
+  && mkdir -p /sys/fs/cgroup/memory/NSJAIL
+
 nsjail -Mo \
 --rlimit_as 700 \
 --chroot / \
@@ -67,11 +71,13 @@ nsjail -Mo \
 --time_limit 2 \
 --disable_proc \
 --iface_no_lo \
+--cgroup_pids_max=1 \
+--cgroup_mem_max=52428800 \
 --quiet -- \
 python3.6 -ISq -c "print('test')"
 ```
 
-> if it fails, try without the `--cgroup_pids_max=1`
+> if it fails, try without the `--cgroup_pids_max=1` and `--cgroup_mem_max=52428800`
 
 ## Development environment
 
diff --git a/snekbox.py b/snekbox.py
index 60bd06d..5a11bf2 100644
--- a/snekbox.py
+++ b/snekbox.py
@@ -17,6 +17,7 @@ class Snekbox(object):
 
         self.nsjail_binary = nsjail_binary
         self.python_binary = python_binary
+        self.nsjail_workaround()
 
     env = {
         'PATH': '/snekbox/.venv/bin:/usr/local/bin:/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin',
@@ -26,6 +27,12 @@ class Snekbox(object):
         'PYTHONDONTWRITEBYTECODE': '1',
     }
 
+    def nsjail_workaround(self):
+        dirs = ['/sys/fs/cgroup/pids/NSJAIL', '/sys/fs/cgroup/memory/NSJAIL']
+        for d in dirs:
+            if not os.path.exists(d):
+                os.makedirs(d)
+
     def python3(self, cmd):
         args = [self.nsjail_binary, '-Mo',
                 '--rlimit_as', '700',
@@ -37,7 +44,8 @@ class Snekbox(object):
                 '--time_limit', '2',
                 '--disable_proc',
                 '--iface_no_lo',
-                # '--cgroup_pids_max=1',  # This doesn't work :(
+                '--cgroup_pids_max=1',
+                '--cgroup_mem_max=52428800'
                 '--quiet', '--',
                 self.python_binary, '-ISq', '-c', cmd]
 
@@ -82,24 +90,10 @@ class Snekbox(object):
                     exchange=snekid)
         exit(0)
 
-    def stopwatch(self, process):
-        log.debug(f'3 second timer started for process {process.pid}')
-        for _ in range(3):
-            time.sleep(1)
-            if not process.is_alive():
-                log.debug(f'Clean exit on process {process.pid}')
-                exit(0)
-
-        process.terminate()
-        log.debug(f'Terminated process {process.pid} forcefully')
-
     def message_handler(self, ch, method, properties, body, thread_ws=None):
         p = multiprocessing.Process(target=self.execute, args=(body,))
         p.daemon = True
         p.start()
-        t = threading.Thread(target=self.stopwatch, args=(p,))
-        t.daemon = True
-        t.start()
 
         ch.basic_ack(delivery_tag=method.delivery_tag)
 
diff --git a/tests/test_snekbox.py b/tests/test_snekbox.py
index addcf5d..5f065aa 100644
--- a/tests/test_snekbox.py
+++ b/tests/test_snekbox.py
@@ -23,7 +23,7 @@ class SnekTests(unittest.TestCase):
                 '    x = x * 99\n')
 
         result = snek.python3(code)
-        self.assertEquals(result.strip(), 'MemoryError')
+        self.assertEquals(result.strip(), 'timed out or memory limit exceeded')
 
     def test_timeout(self):
         code = ('x = "*"\n'
@@ -43,7 +43,15 @@ class SnekTests(unittest.TestCase):
         if 'ModuleNotFoundError' in result.strip():
             self.assertIn('ModuleNotFoundError', result.strip())
         else:
-            self.assertIn('returned non-zero exit status 1.', result.strip())
+            self.assertIn('(PIDs left: 0)', result.strip())
+
+    def test_forkbomb(self):
+        code = ('import os\n'
+                'while 1:\n'
+                '    os.fork()')
+        result = snek.python3(code)
+
+        self.assertIn('(PIDs left: 0)', result.strip())
 
 
 class RMQTests(unittest.TestCase):