Give an actual name to MemoryIO on FreeBSD

Author: quantum5

dmoj/cptbox/utils.py

@@ -2,19 +2,65 @@
 import io
 import mmap
 import os
+from tempfile import NamedTemporaryFile
 from typing import Optional
 
 from dmoj.cptbox._cptbox import memory_fd_create, memory_fd_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):
+    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 to_bytes(self) -> bytes:
+        try:
+            with mmap.mmap(self.fileno(), 0, access=mmap.ACCESS_READ) as f:
+                return bytes(f)
+        except ValueError as e:
+            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)
+        _make_fd_readonly(self.fileno())
+
+    def close(self) -> None:
+        super().close()
+        os.unlink(self._name)
+
+    def to_path(self) -> str:
+        return self._name
+
+
+class MemoryIO(MmapableIO):
+    _name: Optional[str] = None
+
+    def __init__(self, *, prefill: Optional[bytes] = None, seal=False) -> None:
+        super().__init__(memory_fd_create(), prefill=prefill, seal=seal)
+
     def seal(self) -> None:
         fd = self.fileno()
         try:
@@ -26,20 +72,14 @@ def seal(self) -> None:
                 return
             raise
 
-        new_fd = os.open(f'/proc/self/fd/{fd}', os.O_RDONLY)
-        try:
-            os.dup2(new_fd, fd)
-        finally:
-            os.close(new_fd)
+        _make_fd_readonly(fd)
 
     def to_path(self) -> str:
         return f'/proc/{os.getpid()}/fd/{self.fileno()}'
 
-    def to_bytes(self) -> bytes:
-        try:
-            with mmap.mmap(self.fileno(), 0, access=mmap.ACCESS_READ) as f:
-                return bytes(f)
-        except ValueError as e:
-            if e.args[0] == 'cannot mmap an empty file':
-                return b''
-            raise
+
+# On FreeBSD and some other systems, if /proc/[pid]/fd doesn't exist,
+# then MemoryIO.to_path() will not work. We fall back to NamedFileIO
+# in that case.
+if not os.path.isdir(f'/proc/{os.getpid()}/fd'):
+    MemoryIO = NamedFileIO