From ba6d0a8a10af687393134fc1e9662100ce67df52 Mon Sep 17 00:00:00 2001 From: ionite34 Date: Sat, 19 Nov 2022 21:10:20 -0500 Subject: Implement files request form --- tests/test_integration.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) (limited to 'tests/test_integration.py') diff --git a/tests/test_integration.py b/tests/test_integration.py index 7c5db2b..eba5e60 100644 --- a/tests/test_integration.py +++ b/tests/test_integration.py @@ -7,7 +7,7 @@ from tests.gunicorn_utils import run_gunicorn def run_code_in_snekbox(code: str) -> tuple[str, int]: - body = {"input": code} + body = {"args": ["-c", code]} json_data = json.dumps(body).encode("utf-8") req = urllib.request.Request("http://localhost:8060/eval") -- cgit v1.2.3 From 33203063804966baf5ecc4c532bd3ae4e1daee11 Mon Sep 17 00:00:00 2001 From: ionite34 Date: Tue, 22 Nov 2022 15:40:57 -0500 Subject: Add file send integration tests --- tests/test_integration.py | 57 ++++++++++++++++++++++++++++++++++++++++++++--- 1 file changed, 54 insertions(+), 3 deletions(-) (limited to 'tests/test_integration.py') diff --git a/tests/test_integration.py b/tests/test_integration.py index eba5e60..086abab 100644 --- a/tests/test_integration.py +++ b/tests/test_integration.py @@ -1,14 +1,25 @@ import json import unittest import urllib.request +from base64 import b64encode from multiprocessing.dummy import Pool +from textwrap import dedent from tests.gunicorn_utils import run_gunicorn -def run_code_in_snekbox(code: str) -> tuple[str, int]: +def b64encode_code(data: str): + data = dedent(data).strip() + return b64encode(data.encode()).decode("ascii") + + +def snekbox_run_code(code: str) -> tuple[str, int]: body = {"args": ["-c", code]} - json_data = json.dumps(body).encode("utf-8") + return snekbox_request(body) + + +def snekbox_request(content: dict) -> tuple[str, int]: + json_data = json.dumps(content).encode("utf-8") req = urllib.request.Request("http://localhost:8060/eval") req.add_header("Content-Type", "application/json; charset=utf-8") @@ -34,9 +45,49 @@ class IntegrationTests(unittest.TestCase): args = [code] * processes with Pool(processes) as p: - results = p.map(run_code_in_snekbox, args) + results = p.map(snekbox_run_code, args) responses, statuses = zip(*results) self.assertTrue(all(status == 200 for status in statuses)) self.assertTrue(all(json.loads(response)["returncode"] == 0 for response in responses)) + + def test_files_send_receive(self): + """Test sending and receiving files to snekbox.""" + with run_gunicorn(): + request = { + "args": ["main.py"], + "files": [ + { + "path": "main.py", + "content": b64encode_code( + """ + from mod import lib + print(lib.var) + + with open('output.txt', 'w') as f: + f.write('file write test') + """ + ), + }, + {"path": "mod/__init__.py"}, + {"path": "mod/lib.py", "content": b64encode_code("var = 'hello'")}, + ], + } + + expected = { + "stdout": "hello\n", + "returncode": 0, + "files": [ + { + "path": "output.txt", + "size": len("file write test"), + "content": b64encode_code("file write test"), + } + ], + } + + response, status = snekbox_request(request) + + self.assertEqual(200, status) + self.assertEqual(expected, json.loads(response)) -- cgit v1.2.3 From 78b4b6af18a40db3162aad56eb726a26c5a74e8c Mon Sep 17 00:00:00 2001 From: ionite34 Date: Thu, 24 Nov 2022 10:32:32 +0800 Subject: Refactor output files in `output` dir --- snekbox/memfs.py | 16 +++++++++++----- snekbox/nsjail.py | 18 +++++++++++++++--- snekbox/snekio.py | 13 ++++++++++--- tests/test_integration.py | 21 +++++++++++++++------ tests/test_nsjail.py | 6 +++--- 5 files changed, 54 insertions(+), 20 deletions(-) (limited to 'tests/test_integration.py') diff --git a/snekbox/memfs.py b/snekbox/memfs.py index 16f5cfb..a39b690 100644 --- a/snekbox/memfs.py +++ b/snekbox/memfs.py @@ -62,7 +62,12 @@ class MemFS: @property def home(self) -> Path: """Path to home directory.""" - return Path(self.path, "home") + return self.path / "home" + + @property + def output(self) -> Path: + """Path to output directory.""" + return self.home / "output" def __enter__(self) -> MemFS: """Mounts a new tempfs, returns self.""" @@ -79,7 +84,8 @@ class MemFS: else: raise RuntimeError("Failed to generate a unique tempdir name in 10 attempts") - self.mkdir("home") + self.mkdir(self.home) + self.mkdir(self.output) return self def __exit__( @@ -98,17 +104,17 @@ class MemFS: return folder def attachments( - self, max_count: int, pattern: str = "output*" + self, max_count: int, pattern: str = "**/*" ) -> Generator[FileAttachment, None, None]: """Return a list of attachments in the tempdir.""" count = 0 - for file in self.home.glob(pattern): + for file in self.output.rglob(pattern): if count > max_count: log.info(f"Max attachments {max_count} reached, skipping remaining files") break if file.is_file(): count += 1 - yield FileAttachment.from_path(file) + yield FileAttachment.from_path(file, relative_to=self.output) def cleanup(self) -> None: """Unmounts tmpfs.""" diff --git a/snekbox/nsjail.py b/snekbox/nsjail.py index 88af3cf..4538287 100644 --- a/snekbox/nsjail.py +++ b/snekbox/nsjail.py @@ -53,8 +53,21 @@ class NsJail: memfs_instance_size: int = 48 * 1024 * 1024, files_limit: int = 100, files_timeout: float = 15, - files_pattern: str = "output*", + files_pattern: str = "**/*", ): + """ + Initialize NsJail. + + Args: + nsjail_path: Path to the NsJail binary. + config_path: Path to the NsJail configuration file. + max_output_size: Maximum size of the output in bytes. + read_chunk_size: Size of the read buffer in bytes. + memfs_instance_size: Size of the tmpfs instance in bytes. + files_limit: Maximum number of files to parse for attach. + files_timeout: Maximum time in seconds to wait for files to be written / read. + files_pattern: Pattern to match files to attach. + """ self.nsjail_path = nsjail_path self.config_path = config_path self.max_output_size = max_output_size @@ -178,7 +191,6 @@ class NsJail: ) with NamedTemporaryFile() as nsj_log, MemFS(self.memfs_instance_size) as fs: - # Add the temp dir to be mounted as cwd nsjail_args = ( # Set fslimit to unlimited, cannot be set in cfg # due to upstream protobuf parsing issue @@ -204,7 +216,7 @@ class NsJail: *iter_lstrip(py_args), ] - # Write files if any + # Write provided files if any for file in files: file.save_to(fs.home) log.info(f"Created file at {(fs.home / file.path)!r}.") diff --git a/snekbox/snekio.py b/snekbox/snekio.py index bc71501..7852301 100644 --- a/snekbox/snekio.py +++ b/snekbox/snekio.py @@ -57,9 +57,16 @@ class FileAttachment(Generic[T]): return cls(path, content) @classmethod - def from_path(cls, file: Path) -> FileAttachment[bytes]: - """Create an attachment from a file path.""" - return cls(file.name, file.read_bytes()) + def from_path(cls, file: Path, relative_to: Path | None = None) -> FileAttachment[bytes]: + """ + Create an attachment from a file path. + + Args: + file: The file to attach. + relative_to: The root for the path name. + """ + path = file.relative_to(relative_to) if relative_to else file + return cls(str(path), file.read_bytes()) @property def size(self) -> int: diff --git a/tests/test_integration.py b/tests/test_integration.py index 086abab..ba0d9b5 100644 --- a/tests/test_integration.py +++ b/tests/test_integration.py @@ -62,11 +62,15 @@ class IntegrationTests(unittest.TestCase): "path": "main.py", "content": b64encode_code( """ + from pathlib import Path from mod import lib print(lib.var) - with open('output.txt', 'w') as f: - f.write('file write test') + with open('output/test.txt', 'w') as f: + f.write('test 1') + + Path('output/dir').mkdir() + Path('output/dir/test2.txt').write_text('test 2') """ ), }, @@ -80,10 +84,15 @@ class IntegrationTests(unittest.TestCase): "returncode": 0, "files": [ { - "path": "output.txt", - "size": len("file write test"), - "content": b64encode_code("file write test"), - } + "path": "dir/test2.txt", + "size": len("test 2"), + "content": b64encode_code("test 2"), + }, + { + "path": "test.txt", + "size": len("test 1"), + "content": b64encode_code("test 1"), + }, ], } diff --git a/tests/test_nsjail.py b/tests/test_nsjail.py index da2afea..d63180d 100644 --- a/tests/test_nsjail.py +++ b/tests/test_nsjail.py @@ -190,16 +190,16 @@ class NsJailTests(unittest.TestCase): data = "a" * 1024 size = 32 * 1024 * 1024 - with open("src", "w") as f: + with open("output/file", "w") as f: for _ in range((size // 1024) - 5): f.write(data) for i in range(100): - os.symlink("src", f"output{i}") + os.symlink("file", f"output/file{i}") """ ).strip() - nsjail = NsJail(memfs_instance_size=48 * 1024 * 1024, files_timeout=1) + nsjail = NsJail(memfs_instance_size=32 * 1024 * 1024, files_timeout=1) result = nsjail.python3(["-c", code]) self.assertEqual(result.returncode, None) self.assertEqual( -- cgit v1.2.3 From 94fe1fb29102dce384b776c1ccd1823e37f5287f Mon Sep 17 00:00:00 2001 From: ionite34 Date: Mon, 28 Nov 2022 12:20:35 +0800 Subject: Add input/args integration test --- tests/test_integration.py | 13 +++++++++++++ 1 file changed, 13 insertions(+) (limited to 'tests/test_integration.py') diff --git a/tests/test_integration.py b/tests/test_integration.py index ba0d9b5..5aa1b43 100644 --- a/tests/test_integration.py +++ b/tests/test_integration.py @@ -52,6 +52,19 @@ class IntegrationTests(unittest.TestCase): self.assertTrue(all(status == 200 for status in statuses)) self.assertTrue(all(json.loads(response)["returncode"] == 0 for response in responses)) + def test_eval(self): + """Test normal eval requests without files.""" + with run_gunicorn(): + cases = [ + ({"input": "print('Hello')"}, "Hello"), + ({"args": ["-c", "print('Hello')"]}, "Hello"), + ] + for body, expected in cases: + with self.subTest(body=body): + response, status = snekbox_request(body) + self.assertEqual(status, 200) + self.assertEqual(json.loads(response)["stdout"], expected) + def test_files_send_receive(self): """Test sending and receiving files to snekbox.""" with run_gunicorn(): -- cgit v1.2.3 From ad7082e6a5c4462e5064fbf3848bf61ba5c22fdf Mon Sep 17 00:00:00 2001 From: ionite34 Date: Mon, 28 Nov 2022 12:24:59 +0800 Subject: Fix input append to args --- snekbox/api/resources/eval.py | 1 + tests/test_integration.py | 4 ++-- 2 files changed, 3 insertions(+), 2 deletions(-) (limited to 'tests/test_integration.py') diff --git a/snekbox/api/resources/eval.py b/snekbox/api/resources/eval.py index c3e553a..28b739d 100644 --- a/snekbox/api/resources/eval.py +++ b/snekbox/api/resources/eval.py @@ -122,6 +122,7 @@ class EvalResource: # If `input` is supplied, default `args` to `-c` if "input" in body: body.setdefault("args", ["-c"]) + body["args"].append(body["input"]) try: result = self.nsjail.python3( py_args=body["args"], diff --git a/tests/test_integration.py b/tests/test_integration.py index 5aa1b43..01fa657 100644 --- a/tests/test_integration.py +++ b/tests/test_integration.py @@ -56,8 +56,8 @@ class IntegrationTests(unittest.TestCase): """Test normal eval requests without files.""" with run_gunicorn(): cases = [ - ({"input": "print('Hello')"}, "Hello"), - ({"args": ["-c", "print('Hello')"]}, "Hello"), + ({"input": "print('Hello')"}, "Hello\n"), + ({"args": ["-c", "print('abc12')"]}, "abc12\n"), ] for body, expected in cases: with self.subTest(body=body): -- cgit v1.2.3 From 9a2dc79a20feb3a0470a5096455fb826acb1f1b9 Mon Sep 17 00:00:00 2001 From: Ionite Date: Thu, 2 Mar 2023 16:56:23 -0500 Subject: Update unit tests for home output directory --- tests/test_integration.py | 6 +++--- tests/test_nsjail.py | 6 +++--- 2 files changed, 6 insertions(+), 6 deletions(-) (limited to 'tests/test_integration.py') diff --git a/tests/test_integration.py b/tests/test_integration.py index 01fa657..91b01e6 100644 --- a/tests/test_integration.py +++ b/tests/test_integration.py @@ -79,11 +79,11 @@ class IntegrationTests(unittest.TestCase): from mod import lib print(lib.var) - with open('output/test.txt', 'w') as f: + with open('test.txt', 'w') as f: f.write('test 1') - Path('output/dir').mkdir() - Path('output/dir/test2.txt').write_text('test 2') + Path('dir').mkdir() + Path('dir/test2.txt').write_text('test 2') """ ), }, diff --git a/tests/test_nsjail.py b/tests/test_nsjail.py index 71405c7..138c68a 100644 --- a/tests/test_nsjail.py +++ b/tests/test_nsjail.py @@ -191,12 +191,12 @@ class NsJailTests(unittest.TestCase): data = "a" * 1024 size = 32 * 1024 * 1024 - with open("output/file", "w") as f: + with open("file", "w") as f: for _ in range((size // 1024) - 5): f.write(data) for i in range(100): - os.symlink("file", f"output/file{i}") + os.symlink("file", f"file{i}") """ ).strip() @@ -210,7 +210,7 @@ class NsJailTests(unittest.TestCase): def test_file_write_error(self): """Test errors during file write.""" - result = self.nsjail.python3([""], [FileAttachment("output", "hello".encode())]) + result = self.nsjail.python3([""], [FileAttachment("../dev", "hello".encode())]) self.assertEqual(result.returncode, None) self.assertEqual(result.stdout, "IsADirectoryError: Failed to create file 'output'.") -- cgit v1.2.3