From fd0262413e84c02f93d15eca1f87e1ba74cbe34e Mon Sep 17 00:00:00 2001
From: Kovid Goyal <kovid@kovidgoyal.net>
Date: Thu, 10 Mar 2022 06:30:03 +0530
Subject: [PATCH] Access to POSIX shared memory from Python

---
 kitty/data-types.c        |  20 ++++++++
 kitty/fast_data_types.pyi |   8 +++
 kitty/shm.py              | 100 ++++++++++++++++++++++++++++++++++++++
 3 files changed, 128 insertions(+)
 create mode 100644 kitty/shm.py

diff --git a/kitty/data-types.c b/kitty/data-types.c
index f2426a119..743f1afbc 100644
--- a/kitty/data-types.c
+++ b/kitty/data-types.c
@@ -156,6 +156,24 @@ close_tty(PyObject *self UNUSED, PyObject *args) {
 
 #undef TTY_ARGS
 
+static PyObject*
+py_shm_open(PyObject UNUSED *self, PyObject *args) {
+    char *name;
+    int flags, mode = 0600;
+    if (!PyArg_ParseTuple(args, "si|i", &name, &flags, &mode)) return NULL;
+    long fd = safe_shm_open(name, flags, mode);
+    if (fd < 0) return PyErr_SetFromErrnoWithFilenameObject(PyExc_OSError, PyTuple_GET_ITEM(args, 0));
+    return PyLong_FromLong(fd);
+}
+
+static PyObject*
+py_shm_unlink(PyObject UNUSED *self, PyObject *args) {
+    char *name;
+    if (!PyArg_ParseTuple(args, "s", &name)) return NULL;
+    if (shm_unlink(name) != 0) return PyErr_SetFromErrnoWithFilenameObject(PyExc_OSError, PyTuple_GET_ITEM(args, 0));
+    Py_RETURN_NONE;
+}
+
 static PyObject*
 wcwidth_wrap(PyObject UNUSED *self, PyObject *chr) {
     return PyLong_FromLong(wcwidth_std(PyLong_AsLong(chr)));
@@ -184,6 +202,8 @@ static PyMethodDef module_methods[] = {
     {"parse_bytes_dump", (PyCFunction)parse_bytes_dump, METH_VARARGS, ""},
     {"redirect_std_streams", (PyCFunction)redirect_std_streams, METH_VARARGS, ""},
     {"locale_is_valid", (PyCFunction)locale_is_valid, METH_VARARGS, ""},
+    {"shm_open", (PyCFunction)py_shm_open, METH_VARARGS, ""},
+    {"shm_unlink", (PyCFunction)py_shm_unlink, METH_VARARGS, ""},
 #ifdef __APPLE__
     METHODB(user_cache_dir, METH_NOARGS),
     METHODB(process_group_map, METH_NOARGS),
diff --git a/kitty/fast_data_types.pyi b/kitty/fast_data_types.pyi
index 9c9a97c84..a76129980 100644
--- a/kitty/fast_data_types.pyi
+++ b/kitty/fast_data_types.pyi
@@ -1361,3 +1361,11 @@ def set_os_window_title(os_window_id: int, title: str) -> None:
 
 def update_ime_position_for_window(window_id: int, force: bool = False, lost_focus: bool = False) -> bool:
     pass
+
+
+def shm_open(name: str, flags: int, mode: int = 0o600) -> int:
+    pass
+
+
+def shm_unlink(name: str) -> None:
+    pass
diff --git a/kitty/shm.py b/kitty/shm.py
new file mode 100644
index 000000000..6fafa9457
--- /dev/null
+++ b/kitty/shm.py
@@ -0,0 +1,100 @@
+#!/usr/bin/env python
+# License: GPLv3 Copyright: 2022, Kovid Goyal <kovid at kovidgoyal.net>
+
+# This is present in the python stdlib after version 3.7 but we need to support
+# 3.7 for another year, so sigh.
+
+import mmap
+import os
+import secrets
+from typing import Optional
+
+from kitty.fast_data_types import shm_open, shm_unlink
+
+
+def make_filename(safe_length: int = 14, prefix: str = '/ky-') -> str:
+    "Create a random filename for the shared memory object."
+    # number of random bytes to use for name
+    nbytes = (safe_length - len(prefix)) // 2
+    name = prefix + secrets.token_hex(nbytes)
+    return name
+
+
+class SharedMemory:
+
+    def __init__(self, name: Optional[str] = None, create: bool = False, size: int = 0, readonly: bool = False, mode: int = 0o600):
+        if not size >= 0:
+            raise ValueError("'size' must be a positive integer")
+        if create:
+            flags = os.O_CREAT | os.O_EXCL
+            if size <= 0:
+                raise ValueError("'size' must be > 0")
+        else:
+            flags = os.O_RDONLY if readonly else os.O_RDWR
+        if name is None and not flags & os.O_EXCL:
+            raise ValueError("'name' can only be None if create=True")
+
+        if name is None:
+            while True:
+                name = make_filename()
+                try:
+                    self._fd = shm_open(name, flags, mode)
+                except FileExistsError:
+                    continue
+                self._name = name
+                break
+        self._name = name
+        try:
+            if create and size:
+                os.ftruncate(self._fd, size)
+            stats = os.fstat(self._fd)
+            size = stats.st_size
+            self._mmap = mmap.mmap(self._fd, size)
+        except OSError:
+            self.unlink()
+            raise
+
+        self.size = size
+        self._buf: Optional[memoryview] = memoryview(self._mmap)
+
+    def __del__(self) -> None:
+        try:
+            self.close()
+        except OSError:
+            pass
+
+    @property
+    def name(self) -> str:
+        return self._name
+
+    @property
+    def buf(self) -> memoryview:
+        ans = self._buf
+        if ans is None:
+            raise RuntimeError('Cannot access the buffer of a closed shared memory object')
+        return ans
+
+    def __repr__(self) -> str:
+        return f'{self.__class__.__name__}({self.name!r}, size={self.size})'
+
+    def close(self) -> None:
+        """Closes access to the shared memory from this instance but does
+        not destroy the shared memory block."""
+        if self._buf is not None:
+            self._buf.release()
+            self._buf = None
+        if self._mmap is not None:
+            self._mmap.close()
+        if self._fd >= 0:
+            os.close(self._fd)
+            self._fd = -1
+
+    def unlink(self) -> None:
+        """Requests that the underlying shared memory block be destroyed.
+
+        In order to ensure proper cleanup of resources, unlink should be
+        called once (and only once) across all processes which have access
+        to the shared memory block."""
+        if self._name:
+            shm_unlink(self._name)
+            self._name = ''