Track down the closed memfd.

Previously, I had been trying to close the memfd in _shutdown_shm, and
ignoring EBADF.  It turns out that XCB will close the memfd when you
send it to the X server.

I think this was one potential cause of the issues I saw in
test_thread_safety: the two threads would be reallocated each others'
fds, leading to thread A closing an fd that thread B was using,
thinking that it was thread A's memfd.

Fix so that the memfd is only explicitly closed in an error situation.

Author: jholveck

src/mss/linux/xcb.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from ctypes import Structure, _Pointer, c_int, c_uint8, c_uint16, c_uint32
+from ctypes import _Pointer, c_int
 
 from . import xcbgen
 

src/mss/linux/xcbhelpers.py

@@ -230,10 +230,9 @@ def __str__(self) -> str:
         ext_desc = f"\n  Extension:  {details['extension']}" if "extension" in details else ""
         msg += (
             f"\nX Error of failed request:  {error_desc}"
-            f"\n  Major opcode of failed request:  {major_desc}"
-            f"{ext_desc}"
-            f"\n  Minor opcode of failed request:  {minor_desc}"
-            f"\n  Resource id in failed request:  {details['resource_id']}"
+            f"\n  Major opcode of failed request:  {major_desc}{ext_desc}"
+            + (f"\n  Minor opcode of failed request:  {minor_desc}" if details["minor_code"] != 0 else "")
+            + f"\n  Resource id in failed request:  {details['resource_id']}"
             f"\n  Serial number of failed request:  {details['full_sequence']}"
         )
         return msg

src/mss/linux/xshmgetimage.py

@@ -1,9 +1,9 @@
 from __future__ import annotations
 
-import errno
+import enum
 import os
 from mmap import PROT_READ, mmap  # type: ignore[attr-defined]
-from typing import TYPE_CHECKING, Any, Literal
+from typing import TYPE_CHECKING, Any
 
 from mss.exception import ScreenShotError
 from mss.linux import xcb
@@ -16,6 +16,12 @@
     from mss.screenshot import ScreenShot
 
 
+class ShmStatus(enum.Enum):
+    UNKNOWN = enum.auto()
+    AVAILABLE = enum.auto()
+    UNAVAILABLE = enum.auto()
+
+
 class MSS(MSSXCBBase):
     """Multiple ScreenShots implementation for GNU/Linux.
 
@@ -31,65 +37,78 @@ def __init__(self, /, **kwargs: Any) -> None:
         self._buf: mmap | None = None
         self._shmseg: xcb.ShmSeg | None = None
 
-        # _shm_works is True if at least one screenshot has been taken, False if it's known to fail, and None until
-        # then.
-        self._shm_works: bool | None = self._setup_shm()
+        self.shm_status: ShmStatus = self._setup_shm()
+        self.shm_failed_reason: str | None = None
 
     def _shm_report_issue(self, msg: str, *args: Any) -> None:
         """Debugging hook for troubleshooting MIT-SHM issues.
 
         This will be called whenever MIT-SHM is disabled.  The optional
         arguments are not well-defined; exceptions are common.
         """
-        print(msg, args)
+        full_msg = msg
+        if args:
+            full_msg += " | " + ", ".join(str(arg) for arg in args)
+        self.shm_failed_reason = full_msg
 
-    def _setup_shm(self) -> Literal[False] | None:
+    def _setup_shm(self) -> ShmStatus:  # noqa: PLR0911
         assert self.conn is not None  # noqa: S101
 
-        shm_ext_data = xcb.get_extension_data(self.conn, LIB.shm_id)
-        if not shm_ext_data.present:
-            self._shm_report_issue("MIT-SHM extension not present")
-            return False
-
-        # We use the FD-based version of ShmGetImage, so we require the extension to be at least 1.3.
-        shm_version_data = xcb.shm_query_version(self.conn)
-        shm_version = (shm_version_data.major_version, shm_version_data.minor_version)
-        if shm_version < (1, 2):
-            self._shm_report_issue("MIT-SHM version too old", shm_version)
-            return False
-
-        # We allocate something large enough for the root, so we don't have to reallocate each time the window is
-        # resized.
-        # TODO(jholveck): Check in _grab_impl that we're not going to exceed this size.  That can happen if the
-        # root is resized.
-        size = self.pref_screen.width_in_pixels * self.pref_screen.height_in_pixels * 4
-
         try:
-            self._memfd = os.memfd_create("mss-shm-buf", flags=os.MFD_CLOEXEC)  # type: ignore[attr-defined]
-        except OSError as e:
-            self._shm_report_issue("memfd_create failed", e)
-            self._shutdown_shm()
-            return False
-        os.ftruncate(self._memfd, size)
+            shm_ext_data = xcb.get_extension_data(self.conn, LIB.shm_id)
+            if not shm_ext_data.present:
+                self._shm_report_issue("MIT-SHM extension not present")
+                return ShmStatus.UNAVAILABLE
+
+            # We use the FD-based version of ShmGetImage, so we require the extension to be at least 1.2.
+            shm_version_data = xcb.shm_query_version(self.conn)
+            shm_version = (shm_version_data.major_version, shm_version_data.minor_version)
+            if shm_version < (1, 2):
+                self._shm_report_issue("MIT-SHM version too old", shm_version)
+                return ShmStatus.UNAVAILABLE
+
+            # We allocate something large enough for the root, so we don't have to reallocate each time the window is
+            # resized.
+            self._bufsize = self.pref_screen.width_in_pixels * self.pref_screen.height_in_pixels * 4
+
+            if not hasattr(os, "memfd_create"):
+                self._shm_report_issue("os.memfd_create not available")
+                return ShmStatus.UNAVAILABLE
+            try:
+                self._memfd = os.memfd_create("mss-shm-buf", flags=os.MFD_CLOEXEC)  # type: ignore[attr-defined]
+            except OSError as e:
+                self._shm_report_issue("memfd_create failed", e)
+                self._shutdown_shm()
+                return ShmStatus.UNAVAILABLE
+            os.ftruncate(self._memfd, self._bufsize)
 
-        try:
-            self._buf = mmap(self._memfd, size, prot=PROT_READ)  # type: ignore[call-arg]
-        except OSError as e:
-            self._shm_report_issue("mmap failed", e)
-            self._shutdown_shm()
-            return False
+            try:
+                self._buf = mmap(self._memfd, self._bufsize, prot=PROT_READ)  # type: ignore[call-arg]
+            except OSError as e:
+                self._shm_report_issue("mmap failed", e)
+                self._shutdown_shm()
+                return ShmStatus.UNAVAILABLE
 
-        self._shmseg = xcb.ShmSeg(xcb.generate_id(self.conn).value)
-        try:
-            # This will normally be what raises an exception if you're on a remote connection.  I previously thought
-            # the server deferred that until the GetImage call, but I had not been properly checking the status here.
-            xcb.shm_attach_fd(self.conn, self._shmseg, self._memfd, read_only=False)
-        except xcb.XError as e:
-            self._shm_report_issue("Cannot attach MIT-SHM segment", e)
+            self._shmseg = xcb.ShmSeg(xcb.generate_id(self.conn).value)
+            try:
+                # This will normally be what raises an exception if you're on a remote connection.  (I previously
+                # thought the server deferred that until the GetImage call, but I had not been properly checking the
+                # status.)
+                # This will close _memfd, on success or on failure.
+                try:
+                    xcb.shm_attach_fd(self.conn, self._shmseg, self._memfd, read_only=False)
+                finally:
+                    self._memfd = None
+            except xcb.XError as e:
+                self._shm_report_issue("Cannot attach MIT-SHM segment", e)
+                self._shutdown_shm()
+                return ShmStatus.UNAVAILABLE
+
+        except Exception:
             self._shutdown_shm()
-            return False
+            raise
 
-        return None
+        return ShmStatus.UNKNOWN
 
     def close(self) -> None:
         self._shutdown_shm()
@@ -103,15 +122,7 @@ def _shutdown_shm(self) -> None:
             self._buf.close()
             self._buf = None
         if self._memfd is not None:
-            # TODO(jholveck): For some reason, at this point, self._memfd is no longer valid.  If I try to close it,
-            # I get EBADF, even if I try to close it before closing the mmap.  The theories I have about this involve
-            # the mmap object taking control, but it doesn't make sense that I could still use shm_attach_fd in that
-            # case.  I need to investigate before releasing.
-            try:
-                os.close(self._memfd)
-            except OSError as e:
-                if e.errno != errno.EBADF:
-                    raise
+            os.close(self._memfd)
             self._memfd = None
 
     def _grab_impl_xshmgetimage(self, monitor: Monitor) -> ScreenShot:
@@ -121,6 +132,16 @@ def _grab_impl_xshmgetimage(self, monitor: Monitor) -> ScreenShot:
         assert self._buf is not None  # noqa: S101
         assert self._shmseg is not None  # noqa: S101
 
+        required_size = monitor["width"] * monitor["height"] * 4
+        if required_size > self._bufsize:
+            # This is temporary.  The permanent fix will depend on how
+            # issue https://github.com/BoboTiG/python-mss/issues/432 is resolved.
+            msg = (
+                "Requested capture size exceeds the allocated buffer. If you have resized the screen, "
+                "please recreate your MSS object."
+            )
+            raise ScreenShotError(msg)
+
         img_reply = xcb.shm_get_image(
             self.conn,
             self.drawable.value,
@@ -154,19 +175,19 @@ def _grab_impl_xshmgetimage(self, monitor: Monitor) -> ScreenShot:
 
     def _grab_impl(self, monitor: Monitor) -> ScreenShot:
         """Retrieve all pixels from a monitor. Pixels have to be RGBX."""
-        if self._shm_works == False:  # noqa: E712
+        if self.shm_status == ShmStatus.UNAVAILABLE:
             return super()._grab_impl_xgetimage(monitor)
 
         try:
             rv = self._grab_impl_xshmgetimage(monitor)
         except XProtoError as e:
-            if self._shm_works is not None:
+            if self.shm_status != ShmStatus.UNKNOWN:
                 raise
             self._shm_report_issue("MIT-SHM GetImage failed", e)
-            self._shm_works = False
+            self.shm_status = ShmStatus.UNAVAILABLE
             self._shutdown_shm()
             rv = super()._grab_impl_xgetimage(monitor)
         else:
-            self._shm_works = True
+            self.shm_status = ShmStatus.AVAILABLE
 
         return rv

src/tests/test_gnu_linux.py

@@ -310,3 +310,41 @@ def test_with_cursor_failure(display: str) -> None:
             pytest.raises(ScreenShotError),
         ):
             sct.grab(sct.monitors[1])
+
+
+def test_shm_available() -> None:
+    """Verify that the xshmgetimage backend doesn't always fallback.
+
+    Since this backend does an automatic fallback for certain types of
+    anticipated issues, that could cause some failures to be masked.
+    Ensure this isn't happening.
+    """
+    with (
+        pyvirtualdisplay.Display(size=(WIDTH, HEIGHT), color_depth=DEPTH) as vdisplay,
+        mss.mss(display=vdisplay.new_display_var, backend="xshmgetimage") as sct,
+    ):
+        assert isinstance(sct, mss.linux.xshmgetimage.MSS)  # For Mypy
+        # The status currently isn't established as final until a grab succeeds.
+        sct.grab(sct.monitors[0])
+        assert sct.shm_status == mss.linux.xshmgetimage.ShmStatus.AVAILABLE
+
+
+def test_shm_fallback() -> None:
+    """Verify that the xshmgetimage backend falls back if MIT-SHM fails.
+
+    The most common case when a fallback is needed is with a TCP
+    connection, such as the one used with ssh relaying.  By using
+    DISPLAY=localhost:99 instead of DISPLAY=:99, we connect over TCP
+    instead of a local-domain socket.  This is sufficient to prevent
+    MIT-SHM from completing its setup: the extension is available, but
+    won't be able to attach a shared memory segment.
+    """
+    with (
+        pyvirtualdisplay.Display(size=(WIDTH, HEIGHT), color_depth=DEPTH, extra_args=["-listen", "tcp"]) as vdisplay,
+        mss.mss(display=f"localhost{vdisplay.new_display_var}", backend="xshmgetimage") as sct,
+    ):
+        assert isinstance(sct, mss.linux.xshmgetimage.MSS)  # For Mypy
+        # Ensure that the grab call completes without exception.
+        sct.grab(sct.monitors[0])
+        # Ensure that it really did have to fall back; otherwise, we'd need to change how we test this case.
+        assert sct.shm_status == mss.linux.xshmgetimage.ShmStatus.UNAVAILABLE

src/tests/test_implementation.py

@@ -245,10 +245,6 @@ def test_grab_with_tuple_percents(mss_impl: Callable[..., MSSBase]) -> None:
 def test_thread_safety(backend: str) -> None:
     """Regression test for issue #169."""
 
-    # This currently breaks on xshmgetimage.  Resource exhaustion?  Something leaking?
-    if backend == "xshmgetimage":
-        pytest.xfail()
-
     def record(check: dict) -> None:
         """Record for one second."""
         start_time = time.time()