From 52948c75289eb1569e7ae22f2bc47ac2bdd03350 Mon Sep 17 00:00:00 2001 From: Quantum Date: Sun, 30 Jan 2022 15:55:04 -0500 Subject: [PATCH] Implement graceful memfd fallback for FreeBSD --- dmoj/cptbox/_cptbox.pyi | 4 +- dmoj/cptbox/_cptbox.pyx | 12 ++--- dmoj/cptbox/helper.cpp | 11 ++-- dmoj/cptbox/helper.h | 4 +- dmoj/cptbox/utils.py | 113 ++++++++++++++++++++++++++++++++-------- dmoj/problem.py | 10 ++-- 6 files changed, 111 insertions(+), 43 deletions(-) diff --git a/dmoj/cptbox/_cptbox.pyi b/dmoj/cptbox/_cptbox.pyi index d7124f890..362aad3dd 100644 --- a/dmoj/cptbox/_cptbox.pyi +++ b/dmoj/cptbox/_cptbox.pyi @@ -100,8 +100,8 @@ AT_FDCWD: int bsd_get_proc_cwd: Callable[[int], str] bsd_get_proc_fdno: Callable[[int, int], str] -memory_fd_create: Callable[[], int] -memory_fd_seal: Callable[[int], None] +memfd_create: Callable[[], int] +memfd_seal: Callable[[int], None] class BufferProxy: def _get_real_buffer(self): ... diff --git a/dmoj/cptbox/_cptbox.pyx b/dmoj/cptbox/_cptbox.pyx index 03c9937d6..4ee3a80b3 100644 --- a/dmoj/cptbox/_cptbox.pyx +++ b/dmoj/cptbox/_cptbox.pyx @@ -134,8 +134,8 @@ cdef extern from 'helper.h' nogil: PTBOX_SPAWN_FAIL_EXECVE PTBOX_SPAWN_FAIL_SETAFFINITY - int _memory_fd_create "memory_fd_create"() - int _memory_fd_seal "memory_fd_seal"(int fd) + int cptbox_memfd_create() + int cptbox_memfd_seal(int fd) cdef extern from 'fcntl.h' nogil: @@ -215,14 +215,14 @@ def bsd_get_proc_fdno(pid_t pid, int fd): free(buf) return res -def memory_fd_create(): - cdef int fd = _memory_fd_create() +def memfd_create(): + cdef int fd = cptbox_memfd_create() if fd < 0: PyErr_SetFromErrno(OSError) return fd -def memory_fd_seal(int fd): - cdef int result = _memory_fd_seal(fd) +def memfd_seal(int fd): + cdef int result = cptbox_memfd_seal(fd) if result == -1: PyErr_SetFromErrno(OSError) diff --git a/dmoj/cptbox/helper.cpp b/dmoj/cptbox/helper.cpp index cc3ac89f9..fd2bf45c4 100644 --- a/dmoj/cptbox/helper.cpp +++ b/dmoj/cptbox/helper.cpp @@ -328,19 +328,16 @@ char *bsd_get_proc_fdno(pid_t pid, int fdno) { return bsd_get_proc_fd(pid, 0, fdno); } -int memory_fd_create(void) { +int cptbox_memfd_create(void) { #ifdef __FreeBSD__ - char filename[] = "/tmp/cptbox-memoryfd-XXXXXXXX"; - int fd = mkstemp(filename); - if (fd >= 0) - unlink(filename); - return fd; + errno = ENOSYS; + return -1; #else return memfd_create("cptbox memory_fd", MFD_ALLOW_SEALING); #endif } -int memory_fd_seal(int fd) { +int cptbox_memfd_seal(int fd) { #ifdef __FreeBSD__ errno = ENOSYS; return -1; diff --git a/dmoj/cptbox/helper.h b/dmoj/cptbox/helper.h index a22e06ec6..ec96f9392 100644 --- a/dmoj/cptbox/helper.h +++ b/dmoj/cptbox/helper.h @@ -35,7 +35,7 @@ int cptbox_child_run(const struct child_config *config); char *bsd_get_proc_cwd(pid_t pid); char *bsd_get_proc_fdno(pid_t pid, int fdno); -int memory_fd_create(void); -int memory_fd_seal(int fd); +int cptbox_memfd_create(void); +int cptbox_memfd_seal(int fd); #endif diff --git a/dmoj/cptbox/utils.py b/dmoj/cptbox/utils.py index 323271639..133a0ae60 100644 --- a/dmoj/cptbox/utils.py +++ b/dmoj/cptbox/utils.py @@ -1,39 +1,42 @@ -import errno import io import mmap import os +from abc import ABCMeta, abstractmethod +from tempfile import NamedTemporaryFile, TemporaryFile from typing import Optional -from dmoj.cptbox._cptbox import memory_fd_create, memory_fd_seal +from dmoj.cptbox._cptbox import memfd_create, memfd_seal -class MemoryIO(io.FileIO): - def __init__(self, prefill: Optional[bytes] = None, seal=False) -> None: - super().__init__(memory_fd_create(), 'r+') +def _make_fd_readonly(fd): + new_fd = os.open(f'/proc/self/fd/{fd}', os.O_RDONLY) + try: + os.dup2(new_fd, fd) + finally: + os.close(new_fd) + + +class MmapableIO(io.FileIO, metaclass=ABCMeta): + def __init__(self, fd, *, prefill: Optional[bytes] = None, seal=False) -> None: + super().__init__(fd, 'r+') + if prefill: self.write(prefill) if seal: self.seal() - def seal(self) -> None: - fd = self.fileno() - try: - memory_fd_seal(fd) - except OSError as e: - if e.errno == errno.ENOSYS: - # FreeBSD - self.seek(0, os.SEEK_SET) - return - raise + @classmethod + @abstractmethod + def usable_with_name(cls) -> bool: + ... - new_fd = os.open(f'/proc/self/fd/{fd}', os.O_RDONLY) - try: - os.dup2(new_fd, fd) - finally: - os.close(new_fd) + @abstractmethod + def seal(self) -> None: + ... + @abstractmethod def to_path(self) -> str: - return f'/proc/{os.getpid()}/fd/{self.fileno()}' + ... def to_bytes(self) -> bytes: try: @@ -43,3 +46,71 @@ def to_bytes(self) -> bytes: if e.args[0] == 'cannot mmap an empty file': return b'' raise + + +class NamedFileIO(MmapableIO): + _name: str + + def __init__(self, *, prefill: Optional[bytes] = None, seal=False) -> None: + with NamedTemporaryFile(delete=False) as f: + self._name = f.name + super().__init__(os.dup(f.fileno()), prefill=prefill, seal=seal) + + def seal(self) -> None: + self.seek(0, os.SEEK_SET) + + def close(self) -> None: + super().close() + os.unlink(self._name) + + def to_path(self) -> str: + return self._name + + @classmethod + def usable_with_name(cls): + return True + + +class UnnamedFileIO(MmapableIO): + def __init__(self, *, prefill: Optional[bytes] = None, seal=False) -> None: + with TemporaryFile() as f: + super().__init__(os.dup(f.fileno()), prefill=prefill, seal=seal) + + def seal(self) -> None: + self.seek(0, os.SEEK_SET) + _make_fd_readonly(self.fileno()) + + def to_path(self) -> str: + return f'/proc/{os.getpid()}/fd/{self.fileno()}' + + @classmethod + def usable_with_name(cls): + with cls() as f: + return os.path.exists(f.to_path()) + + +class MemfdIO(MmapableIO): + def __init__(self, *, prefill: Optional[bytes] = None, seal=False) -> None: + super().__init__(memfd_create(), prefill=prefill, seal=seal) + + def seal(self) -> None: + fd = self.fileno() + memfd_seal(fd) + _make_fd_readonly(fd) + + def to_path(self) -> str: + return f'/proc/{os.getpid()}/fd/{self.fileno()}' + + @classmethod + def usable_with_name(cls): + try: + with cls() as f: + return os.path.exists(f.to_path()) + except OSError: + return False + + +# Try to use memfd if possible, otherwise fallback to unlinked temporary files +# (UnnamedFileIO). On FreeBSD and some other systems, /proc/[pid]/fd doesn't +# exist, so to_path() will not work. We fall back to NamedFileIO in that case. +MemoryIO = next((i for i in (MemfdIO, UnnamedFileIO, NamedFileIO) if i.usable_with_name())) diff --git a/dmoj/problem.py b/dmoj/problem.py index e84e931b5..407e136da 100644 --- a/dmoj/problem.py +++ b/dmoj/problem.py @@ -30,7 +30,7 @@ from dmoj import checkers from dmoj.checkers import Checker from dmoj.config import ConfigNode, InvalidInitException -from dmoj.cptbox.utils import MemoryIO +from dmoj.cptbox.utils import MemoryIO, MmapableIO from dmoj.judgeenv import env, get_problem_root from dmoj.utils.helper_files import compile_with_auxiliary_files, parse_helper_file_error from dmoj.utils.module import load_module_from_file @@ -275,7 +275,7 @@ def open(self, key: str): return self.archive.open(zipinfo) raise KeyError('file "%s" could not be found in "%s"' % (key, self.problem_root_dir)) - def as_fd(self, key: str, normalize: bool = False) -> MemoryIO: + def as_fd(self, key: str, normalize: bool = False) -> MmapableIO: memory = MemoryIO() with self.open(key) as f: if normalize: @@ -344,7 +344,7 @@ class TestCase(BaseTestCase): batch: int output_prefix_length: int has_binary_data: bool - _input_data_fd: Optional[MemoryIO] + _input_data_fd: Optional[MmapableIO] _generated: Optional[Tuple[bytes, bytes]] def __init__(self, count: int, batch_no: int, config: ConfigNode, problem: Problem): @@ -451,14 +451,14 @@ def _run_generator(self, gen: Union[str, ConfigNode], args: Optional[Iterable[st def input_data(self) -> bytes: return self.input_data_fd().to_bytes() - def input_data_fd(self) -> MemoryIO: + def input_data_fd(self) -> MmapableIO: if self._input_data_fd: return self._input_data_fd result = self._input_data_fd = self._make_input_data_fd() return result - def _make_input_data_fd(self) -> MemoryIO: + def _make_input_data_fd(self) -> MmapableIO: gen = self.config.generator # don't try running the generator if we specify an output file explicitly,