Merge pull request #24 from python-discord/refactor/nsjail

Improve NsJail

13 files changed, 254 insertions, 148 deletions

diff --git a/.flake8 b/.flake8
index f8eec98..1a61894 100644
--- a/.flake8
+++ b/.flake8
@@ -13,6 +13,7 @@ ignore=
     D400,D401,D402,D405,D406,D407,D408,D409,D410,D411,D412,D413,D414
 exclude=
     __pycache__,.cache,
-    venv,.venv,
-    tests
+    venv,.venv
+per-file-ignores=tests/*:D1
 import-order-style=pycharm
+inline-quotes="
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
index 1d75342..d2737fe 100644
--- a/.pre-commit-config.yaml
+++ b/.pre-commit-config.yaml
@@ -1,5 +1,16 @@
 repos:
 -   repo: https://github.com/pre-commit/pre-commit-hooks
-    rev: v2.0.0
+    rev: v2.2.3
     hooks:
-    - id: flake8
\ No newline at end of file
+    - id: flake8
+      additional_dependencies: [
+        flake8-docstrings,
+        flake8-bugbear,
+        flake8-import-order,
+        flake8-tidy-imports,
+        flake8-todo,
+        flake8-string-format,
+        flake8-formatter-junit-xml,
+        flake8-quotes
+      ]
+
diff --git a/Pipfile b/Pipfile
index 72df21b..6b6aa24 100644
--- a/Pipfile
+++ b/Pipfile
@@ -21,6 +21,7 @@ flake8-tidy-imports = "*"
 flake8-todo = "*"
 flake8-string-format = "*"
 flake8-formatter-junit-xml = "*"
+flake8-quotes = "*"
 
 [requires]
 python_version = "3.7"
@@ -30,7 +31,7 @@ lint = "flake8"
 precommit = "pre-commit install"
 test = "pytest tests --cov . --cov-report term-missing -v"
 report = "pytest tests --cov . --cov-report=html"
-snekbox = "gunicorn -w 2 -b 0.0.0.0:8060 snekbox.api.app"
+snekbox = "gunicorn -w 2 -b 0.0.0.0:8060 --logger-class snekbox.GunicornLogger --access-logfile - snekbox.api.app"
 buildbox = "docker build -t pythondiscord/snekbox:latest -f docker/Dockerfile ."
 pushbox = "docker push pythondiscord/snekbox:latest"
 buildboxbase = "docker build -t pythondiscord/snekbox-base:latest -f docker/base.Dockerfile ."
diff --git a/Pipfile.lock b/Pipfile.lock
index e0c08a0..c09c916 100644
--- a/Pipfile.lock
+++ b/Pipfile.lock
@@ -1,7 +1,7 @@
 {
     "_meta": {
         "hash": {
-            "sha256": "8dda3bfb1f2f9109b882225e4c55e3a561e25a7d80845889b4ffe5ad250cc86e"
+            "sha256": "af53793d2c00001698021096041ed23e5de0a5553f974822e9ad7e58f70de4a9"
         },
         "pipfile-spec": 6,
         "requires": {
@@ -61,8 +61,7 @@
         },
         "pyrsistent": {
             "hashes": [
-                "sha256:16692ee739d42cf5e39cef8d27649a8c1fdb7aa99887098f1460057c5eb75c3a",
-                "sha256:1b304bab4d25fbe2b5bf32034d0d904bfc870f7f4ed9dccec6ae388978f0ef6f"
+                "sha256:16692ee739d42cf5e39cef8d27649a8c1fdb7aa99887098f1460057c5eb75c3a"
             ],
             "version": "==0.15.2"
         },
@@ -193,6 +192,13 @@
             ],
             "version": "==1.0.2"
         },
+        "flake8-quotes": {
+            "hashes": [
+                "sha256:10c9af6b472d4302a8e721c5260856c3f985c5c082b04841aefd2f808ac02038"
+            ],
+            "index": "pypi",
+            "version": "==2.0.1"
+        },
         "flake8-string-format": {
             "hashes": [
                 "sha256:68ea72a1a5b75e7018cae44d14f32473c798cf73d75cbaed86c6a9a907b770b2",
@@ -225,10 +231,10 @@
         },
         "importlib-metadata": {
             "hashes": [
-                "sha256:027cfc6524613de726789072f95d2e4cc64dd1dee8096d42d13f2ead5bd302f5",
-                "sha256:0d05199e1f0b1a8707a1b9c46476d4a49807fb56cb1b0737db1d37feb42fe31d"
+                "sha256:a9f185022cfa69e9ca5f7eabfd5a58b689894cb78a11e3c8c89398a8ccbb8e7f",
+                "sha256:df1403cd3aebeb2b1dcd3515ca062eecb5bd3ea7611f18cba81130c68707e879"
             ],
-            "version": "==0.15"
+            "version": "==0.17"
         },
         "junit-xml": {
             "hashes": [
@@ -257,6 +263,13 @@
             ],
             "version": "==1.3.3"
         },
+        "packaging": {
+            "hashes": [
+                "sha256:0c98a5d0be38ed775798ece1b9727178c4469d9c3b4ada66e8e6b7849f8732af",
+                "sha256:9e1cbf8c12b1f1ce0bb5344b8d7ecf66a6f8a6e91bcb0c84593ed6d3ab5c4ab3"
+            ],
+            "version": "==19.0"
+        },
         "pluggy": {
             "hashes": [
                 "sha256:0825a152ac059776623854c1543d65a4ad408eb3d33ee114dff91e57ec6ae6fc",
@@ -301,13 +314,20 @@
             ],
             "version": "==2.1.1"
         },
+        "pyparsing": {
+            "hashes": [
+                "sha256:1873c03321fc118f4e9746baf201ff990ceb915f433f23b395f5580d1840cb2a",
+                "sha256:9b6323ef4ab914af344ba97510e966d64ba91055d6b9afa6b30799340e89cc03"
+            ],
+            "version": "==2.4.0"
+        },
         "pytest": {
             "hashes": [
-                "sha256:1a8aa4fa958f8f451ac5441f3ac130d9fc86ea38780dd2715e6d5c5882700b24",
-                "sha256:b8bf138592384bd4e87338cb0f256bf5f615398a649d4bd83915f0e4047a5ca6"
+                "sha256:6032845e68a17a96e8da3088037f899b56357769a724122056265ca2ea1890ee",
+                "sha256:bea27a646a3d74cbbcf8d3d4a06b2dfc336baf3dc2cc85cf70ad0157e73e8322"
             ],
             "index": "pypi",
-            "version": "==4.5.0"
+            "version": "==4.6.2"
         },
         "pytest-cov": {
             "hashes": [
diff --git a/snekbox/__init__.py b/snekbox/__init__.py
index fc6070e..a48abd5 100644
--- a/snekbox/__init__.py
+++ b/snekbox/__init__.py
@@ -1,10 +1,34 @@
 import logging
+import os
 import sys
 
-logformat = logging.Formatter(fmt='[%(asctime)s] [%(process)s] [%(levelname)s] %(message)s',
-                              datefmt='%Y-%m-%d %H:%M:%S %z')
-log = logging.getLogger(__name__)
-log.setLevel(logging.DEBUG)
-console = logging.StreamHandler(sys.stdout)
-console.setFormatter(logformat)
-log.addHandler(console)
+from gunicorn import glogging
+
+DEBUG = os.environ.get("DEBUG", False)
+
+
+class GunicornLogger(glogging.Logger):
+    """Logger for Gunicorn with custom formatting and support for the DEBUG environment variable."""
+
+    error_fmt = "%(asctime)s | %(process)5s | %(name)30s | %(levelname)8s | %(message)s"
+    datefmt = "%Y-%m-%d %H:%M:%S"
+
+    def setup(self, cfg):
+        """Set up loggers and set error logger's level to DEBUG if the DEBUG env var is set."""
+        super().setup(cfg)
+
+        if DEBUG:
+            self.loglevel = logging.DEBUG
+        else:
+            self.loglevel = self.LOG_LEVELS.get(cfg.loglevel.lower(), logging.INFO)
+
+        self.error_log.setLevel(self.loglevel)
+
+
+log = logging.getLogger("snekbox")
+log.setLevel(logging.DEBUG if DEBUG else logging.INFO)
+log.propagate = True
+formatter = logging.Formatter(GunicornLogger.error_fmt, GunicornLogger.datefmt)
+handler = logging.StreamHandler(sys.stdout)
+handler.setFormatter(formatter)
+log.addHandler(handler)
diff --git a/snekbox/api/middleware/__init__.py b/snekbox/api/middleware/__init__.py
deleted file mode 100644
index ba97ca6..0000000
--- a/snekbox/api/middleware/__init__.py
+++ /dev/null
@@ -1,3 +0,0 @@
-from .logger import LoggingMiddleware
-
-__all__ = ("LoggingMiddleware",)
diff --git a/snekbox/api/middleware/logger.py b/snekbox/api/middleware/logger.py
deleted file mode 100644
index 393fc64..0000000
--- a/snekbox/api/middleware/logger.py
+++ /dev/null
@@ -1,11 +0,0 @@
-import logging
-
-log = logging.getLogger("snekbox.api")
-
-
-class LoggingMiddleware:
-    """Log basic information about responses."""
-
-    def process_response(self, req, resp, resource, req_succeeded):
-        """Log the method, route, and status of a response."""
-        log.info(f"{req.method} {req.relative_uri} {resp.status}")
diff --git a/snekbox/api/resources/eval.py b/snekbox/api/resources/eval.py
index b2f4260..4779557 100644
--- a/snekbox/api/resources/eval.py
+++ b/snekbox/api/resources/eval.py
@@ -36,7 +36,16 @@ class EvalResource:
     @validate(REQ_SCHEMA)
     def on_post(self, req, resp):
         """
-        Evaluate Python code and return the result.
+        Evaluate Python code and return stdout, stderr, and the return code.
+
+        The return codes mostly resemble those of a Unix shell. Some noteworthy cases:
+
+        - None
+            The NsJail process failed to launch
+        - 137 (SIGKILL)
+            Typically because NsJail killed the Python process due to time or memory constraints
+        - 255
+            NsJail encountered a fatal error
 
         Request body:
 
@@ -47,8 +56,9 @@ class EvalResource:
         Response format:
 
         >>> {
-        ...     "input": "print(1 + 1)",
-        ...     "output": "2\\n"
+        ...     "stdout": "2\\n",
+        ...     "stderr": "",
+        ...     "returncode": 0
         ... }
 
         Status codes:
@@ -63,12 +73,13 @@ class EvalResource:
         code = req.media["input"]
 
         try:
-            output = self.nsjail.python3(code)
+            result = self.nsjail.python3(code)
         except Exception:
             log.exception("An exception occurred while trying to process the request")
             raise falcon.HTTPInternalServerError
 
         resp.media = {
-            "input": code,
-            "output": output
+            "stdout": result.stdout,
+            "stderr": result.stderr,
+            "returncode": result.returncode
         }
diff --git a/snekbox/api/snekapi.py b/snekbox/api/snekapi.py
index 849e7d6..cb0356a 100644
--- a/snekbox/api/snekapi.py
+++ b/snekbox/api/snekapi.py
@@ -1,6 +1,5 @@
 import falcon
 
-from .middleware import LoggingMiddleware
 from .resources import EvalResource
 
 
@@ -22,6 +21,6 @@ class SnekAPI(falcon.API):
     """
 
     def __init__(self, *args, **kwargs):
-        super().__init__(middleware=[LoggingMiddleware()], *args, **kwargs)
+        super().__init__(*args, **kwargs)
 
         self.add_route("/eval", EvalResource())
diff --git a/snekbox/nsjail.py b/snekbox/nsjail.py
index 5c7d0f0..ff12ec4 100644
--- a/snekbox/nsjail.py
+++ b/snekbox/nsjail.py
@@ -1,94 +1,140 @@
-import os
+import logging
+import re
 import subprocess
 import sys
+import textwrap
+from pathlib import Path
+from tempfile import NamedTemporaryFile
+
+from snekbox import DEBUG
+
+log = logging.getLogger(__name__)
+
+# [level][timestamp][PID]? function_signature:line_no? message
+LOG_PATTERN = re.compile(
+    r"\[(?P<level>(I)|[WEF])\]\[.+?\](?(2)|(?P<func>\[\d+\] .+?:\d+ )) ?(?P<msg>.+)"
+)
+LOG_BLACKLIST = ("Process will be ",)
+
+# Explicitly define constants for NsJail's default values.
+CGROUP_PIDS_PARENT = Path("/sys/fs/cgroup/pids/NSJAIL")
+CGROUP_MEMORY_PARENT = Path("/sys/fs/cgroup/memory/NSJAIL")
+
+ENV = {
+    "PATH": (
+        "/snekbox/.venv/bin:/usr/local/bin:/usr/local/"
+        "sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin"
+    ),
+    "LANG": "en_US.UTF-8",
+    "PYTHON_VERSION": "3.7.3",
+    "PYTHON_PIP_VERSION": "19.0.3",
+    "PYTHONDONTWRITEBYTECODE": "1",
+}
 
 
 class NsJail:
-    """Core Snekbox functionality, providing safe execution of Python code."""
+    """
+    Core Snekbox functionality, providing safe execution of Python code.
+
+    NsJail configuration:
+
+    - Root directory is mounted as read-only
+    - Time limit of 2 seconds
+    - Maximum of 1 PID
+    - Maximum memory of 52428800 bytes
+    - Loopback interface is down
+    - procfs is disabled
+
+    Python configuration:
+
+    - Isolated mode
+        - Neither the script's directory nor the user's site packages are in sys.path
+        - All PYTHON* environment variables are ignored
+    - Import of the site module is disabled
+    """
 
-    def __init__(self,
-                 nsjail_binary='nsjail',
-                 python_binary=os.path.dirname(sys.executable) + os.sep + 'python3.7'):
+    def __init__(self, nsjail_binary="nsjail", python_binary=sys.executable):
         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'
-        ),
-        'LANG': 'en_US.UTF-8',
-        'PYTHON_VERSION': '3.7.3',
-        'PYTHON_PIP_VERSION': '19.0.3',
-        '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):
-        """
-        Execute Python 3 code in a isolated environment.
 
-        The value of ``cmd`` is passed using '-c' to a Python
-        interpreter that is started in a ``nsjail``, isolating it
-        from the rest of the system.
+        self._create_parent_cgroups()
 
-        Returns the output of executing the command (stdout) if
-        successful, or a error message if the execution failed.
+    @staticmethod
+    def _create_parent_cgroups(pids: Path = CGROUP_PIDS_PARENT, mem: Path = CGROUP_MEMORY_PARENT):
         """
-        args = [self.nsjail_binary, '-Mo',
-                '--rlimit_as', '700',
-                '--chroot', '/',
-                '-E', 'LANG=en_US.UTF-8',
-                '-R/usr', '-R/lib', '-R/lib64',
-                '--user', 'nobody',
-                '--group', 'nogroup',
-                '--time_limit', '2',
-                '--disable_proc',
-                '--iface_no_lo',
-                '--cgroup_pids_max=1',
-                '--cgroup_mem_max=52428800',
-                '--quiet', '--',
-                self.python_binary, '-ISq', '-c', cmd]
-        try:
-            proc = subprocess.Popen(args,
-                                    stdin=subprocess.PIPE,
-                                    stdout=subprocess.PIPE,
-                                    stderr=subprocess.PIPE,
-                                    env=self.env,
-                                    universal_newlines=True)
-        except ValueError:
-            return 'ValueError: embedded null byte'
-
-        stdout, stderr = proc.communicate()
-        if proc.returncode == 0:
-            output = stdout
-
-        elif proc.returncode == 1:
-            try:
-                filtered = []
-                for line in stderr.split('\n'):
-                    if not line.startswith('['):
-                        filtered.append(line)
-                output = '\n'.join(filtered)
-            except IndexError:
-                output = ''
+        Create the PIDs and memory cgroups which NsJail will use as its parent cgroups.
+
+        NsJail doesn't do this automatically because it requires privileges NsJail usually doesn't
+        have.
+        """
+        pids.mkdir(parents=True, exist_ok=True)
+        mem.mkdir(parents=True, exist_ok=True)
+
+    @staticmethod
+    def _parse_log(log_file):
+        """Parse and log NsJail's log messages."""
+        for line in log_file.read().decode("UTF-8").splitlines():
+            match = LOG_PATTERN.fullmatch(line)
+            if match is None:
+                log.warning(f"Failed to parse log line '{line}'")
+                continue
 
-        elif proc.returncode == 109:
-            return 'timed out or memory limit exceeded'
+            msg = match["msg"]
+            if not DEBUG and any(msg.startswith(s) for s in LOG_BLACKLIST):
+                # Skip blacklisted messages if not debugging.
+                continue
 
-        elif proc.returncode == 255:
-            return 'permission denied (root required)'
+            if DEBUG and match["func"]:
+                # Prepend PID, function signature, and line number if debugging.
+                msg = f"{match['func']}{msg}"
 
-        elif proc.returncode:
-            return f'unknown error, code: {proc.returncode}'
+            if match["level"] == "D":
+                log.debug(msg)
+            elif match["level"] == "I":
+                if DEBUG or msg.startswith("pid="):
+                    # Skip messages unrelated to process exit if not debugging.
+                    log.info(msg)
+            elif match["level"] == "W":
+                log.warning(msg)
+            else:
+                # Treat fatal as error.
+                log.error(msg)
+
+    def python3(self, code: str) -> subprocess.CompletedProcess:
+        """Execute Python 3 code in an isolated environment and return the completed process."""
+        with NamedTemporaryFile() as nsj_log:
+            args = (
+                self.nsjail_binary, "-Mo",
+                "--rlimit_as", "700",
+                "--chroot", "/",
+                "-E", "LANG=en_US.UTF-8",
+                "-R/usr", "-R/lib", "-R/lib64",
+                "--user", "nobody",
+                "--group", "nogroup",
+                "--time_limit", "2",
+                "--disable_proc",
+                "--iface_no_lo",
+                "--log", nsj_log.name,
+                "--cgroup_mem_max=52428800",
+                "--cgroup_mem_mount", str(CGROUP_MEMORY_PARENT.parent),
+                "--cgroup_mem_parent", CGROUP_MEMORY_PARENT.name,
+                "--cgroup_pids_max=1",
+                "--cgroup_pids_mount", str(CGROUP_PIDS_PARENT.parent),
+                "--cgroup_pids_parent", CGROUP_PIDS_PARENT.name,
+                "--",
+                self.python_binary, "-ISq", "-c", code
+            )
+
+            msg = "Executing code..."
+            if DEBUG:
+                msg = f"{msg[:-3]}:\n{textwrap.indent(code, '    ')}"
+            log.info(msg)
+
+            try:
+                result = subprocess.run(args, capture_output=True, env=ENV, text=True)
+            except ValueError:
+                return subprocess.CompletedProcess(args, None, "", "ValueError: embedded null byte")
 
-        else:
-            return 'unknown error, no error code'
+            self._parse_log(nsj_log)
 
-        return output
+        return result
diff --git a/tests/api/__init__.py b/tests/api/__init__.py
index fd4679a..dcee5b5 100644
--- a/tests/api/__init__.py
+++ b/tests/api/__init__.py
@@ -1,3 +1,4 @@
+from subprocess import CompletedProcess
 from unittest import mock
 
 from falcon import testing
@@ -11,7 +12,12 @@ class SnekAPITestCase(testing.TestCase):
 
         self.patcher = mock.patch("snekbox.api.resources.eval.NsJail", autospec=True)
         self.mock_nsjail = self.patcher.start()
-        self.mock_nsjail.return_value.python3.return_value = "test output"
+        self.mock_nsjail.return_value.python3.return_value = CompletedProcess(
+            args=[],
+            returncode=0,
+            stdout="output",
+            stderr="error"
+        )
         self.addCleanup(self.patcher.stop)
 
         self.app = SnekAPI()
diff --git a/tests/api/test_eval.py b/tests/api/test_eval.py
index a5b83fd..03f0e39 100644
--- a/tests/api/test_eval.py
+++ b/tests/api/test_eval.py
@@ -9,8 +9,9 @@ class TestEvalResource(SnekAPITestCase):
         result = self.simulate_post(self.PATH, json=body)
 
         self.assertEqual(result.status_code, 200)
-        self.assertEqual(body["input"], result.json["input"])
-        self.assertEqual("test output", result.json["output"])
+        self.assertEqual("output", result.json["stdout"])
+        self.assertEqual("error", result.json["stderr"])
+        self.assertEqual(0, result.json["returncode"])
 
     def test_post_invalid_schema_400(self):
         body = {"stuff": "foo"}
@@ -26,7 +27,7 @@ class TestEvalResource(SnekAPITestCase):
         self.assertEqual(expected, result.json)
 
     def test_post_invalid_content_type_415(self):
-        body = "{\"input\": \"foo\"}"
+        body = "{'input': 'foo'}"
         headers = {"Content-Type": "application/xml"}
         result = self.simulate_post(self.PATH, body=body, headers=headers)
 
diff --git a/tests/test_snekbox.py b/tests/test_snekbox.py
index c08178f..46319d6 100644
--- a/tests/test_snekbox.py
+++ b/tests/test_snekbox.py
@@ -7,44 +7,44 @@ nsjail = NsJail()
 
 class SnekTests(unittest.TestCase):
     def test_nsjail(self):
-        result = nsjail.python3('print("test")')
-        self.assertEquals(result.strip(), 'test')
+        result = nsjail.python3("print('test')")
+        self.assertEquals(result.strip(), "test")
 
     # def test_memory_error(self):
-    #     code = ('x = "*"\n'
-    #             'while True:\n'
-    #             '    x = x * 99\n')
+    #     code = ("x = "*"\n"
+    #             "while True:\n"
+    #             "    x = x * 99\n")
     #     result = nsjail.python3(code)
-    #     self.assertEquals(result.strip(), 'timed out or memory limit exceeded')
+    #     self.assertEquals(result.strip(), "timed out or memory limit exceeded")
 
     def test_timeout(self):
         code = (
-            'x = "*"\n'
-            'while True:\n'
-            '    try:\n'
-            '        x = x * 99\n'
-            '    except:\n'
-            '        continue\n'
+            "x = '*'\n"
+            "while True:\n"
+            "    try:\n"
+            "        x = x * 99\n"
+            "    except:\n"
+            "        continue\n"
         )
 
         result = nsjail.python3(code)
-        self.assertEquals(result.strip(), 'timed out or memory limit exceeded')
+        self.assertEquals(result.strip(), "timed out or memory limit exceeded")
 
     def test_kill(self):
-        code = ('import subprocess\n'
-                'print(subprocess.check_output("kill -9 6", shell=True).decode())')
+        code = ("import subprocess\n"
+                "print(subprocess.check_output('kill -9 6', shell=True).decode())")
         result = nsjail.python3(code)
-        if 'ModuleNotFoundError' in result.strip():
-            self.assertIn('ModuleNotFoundError', result.strip())
+        if "ModuleNotFoundError" in result.strip():
+            self.assertIn("ModuleNotFoundError", result.strip())
         else:
-            self.assertIn('(PIDs left: 0)', result.strip())
+            self.assertIn("(PIDs left: 0)", result.strip())
 
     def test_forkbomb(self):
-        code = ('import os\n'
-                'while 1:\n'
-                '    os.fork()')
+        code = ("import os\n"
+                "while 1:\n"
+                "    os.fork()")
         result = nsjail.python3(code)
-        self.assertIn('Resource temporarily unavailable', result.strip())
+        self.assertIn("Resource temporarily unavailable", result.strip())
 
     def test_juan_golf(self):  # in honour of Juan
         code = ("func = lambda: None\n"
@@ -53,4 +53,4 @@ class SnekTests(unittest.TestCase):
                 "exec(bytecode)")
 
         result = nsjail.python3(code)
-        self.assertEquals('unknown error, code: 111', result.strip())
+        self.assertEquals("unknown error, code: 111", result.strip())