Fix send_all byte offsets for typed memoryviews

newsfragments/3543.bugfix.rst

@@ -0,0 +1,4 @@
+Fixed silent data loss in ``trio.SocketStream.send_all`` and
+``trio.lowlevel.FdStream.send_all`` when a partial write occurs while sending a
+contiguous memoryview with multi-byte elements. Multidimensional and
+zero-dimensional views are now also handled in byte units.

src/trio/_highlevel_socket.py

@@ -112,16 +112,18 @@ async def send_all(self, data: bytes | bytearray | memoryview) -> None:
         with self._send_conflict_detector:
             with _translate_socket_errors_to_stream_errors():
                 with memoryview(data) as data:
-                    if not data:
+                    if not data.nbytes:
                         if self.socket.fileno() == -1:
                             raise trio.ClosedResourceError("socket was already closed")
                         await trio.lowlevel.checkpoint()
                         return
-                    total_sent = 0
-                    while total_sent < len(data):
-                        with data[total_sent:] as remaining:
-                            sent = await self.socket.send(remaining)
-                        total_sent += sent
+                    # send() returns byte counts, so offsets must also be in bytes.
+                    with data.cast("B") as data:
+                        total_sent = 0
+                        while total_sent < len(data):
+                            with data[total_sent:] as remaining:
+                                sent = await self.socket.send(remaining)
+                            total_sent += sent
 
     async def wait_send_all_might_not_block(self) -> None:
         with self._send_conflict_detector:

src/trio/_tests/test_highlevel_socket.py

@@ -9,6 +9,7 @@
 
 from .. import _core, socket as tsocket
 from .._highlevel_socket import *
+from .._socket import _SocketType
 from ..testing import (
     assert_checkpoints,
     check_half_closeable_stream,
@@ -119,6 +120,45 @@ async def receiver() -> None:
         assert await b.receive_some(10) == b""
 
 
+@pytest.mark.parametrize("shape", [(16,), (4, 4), ()])
+async def test_SocketStream_send_all_memoryview(
+    monkeypatch: pytest.MonkeyPatch,
+    shape: tuple[int, ...],
+) -> None:
+    original_send = _SocketType.send
+
+    async def partial_send(
+        self: _SocketType,
+        data: bytes | bytearray | memoryview,
+        flags: int = 0,
+    ) -> int:
+        # Use a real socket, but force partial writes independently of OS buffers.
+        with memoryview(data).cast("B") as view:
+            return await original_send(self, view[:3], flags)
+
+    monkeypatch.setattr(_SocketType, "send", partial_send)
+    payload = bytes(range(32)) if shape else b"ab"
+    left, right = tsocket.socketpair()
+    async with SocketStream(left) as sender, SocketStream(right) as receiver:
+        with memoryview(payload).cast("H", shape=shape) as data:
+            await sender.send_all(data)
+        await sender.send_eof()
+        received = bytearray()
+        async for chunk in receiver:
+            received.extend(chunk)
+        assert received == payload
+
+
+async def test_SocketStream_send_all_empty_noncontiguous_memoryview() -> None:
+    left, right = tsocket.socketpair()
+    async with SocketStream(left) as sender, SocketStream(right) as receiver:
+        with memoryview(b"abc")[0:0:2] as data:
+            with assert_checkpoints():
+                await sender.send_all(data)
+        await sender.send_eof()
+        assert await receiver.receive_some() == b""
+
+
 async def fill_stream(s: SocketStream) -> None:
     async def sender() -> None:
         while True:

src/trio/_tests/test_unix_pipes.py

@@ -10,7 +10,7 @@
 
 from .. import _core
 from .._core._tests.tutil import gc_collect_harder, skip_if_fbsd_pipes_broken
-from ..testing import check_one_way_stream, wait_all_tasks_blocked
+from ..testing import assert_checkpoints, check_one_way_stream, wait_all_tasks_blocked
 
 if TYPE_CHECKING:
     from .._file_io import _HasFileNo
@@ -66,6 +66,40 @@ async def test_send_pipe() -> None:
         os.close(r)
 
 
+@pytest.mark.parametrize("shape", [(16,), (4, 4), ()])
+async def test_send_pipe_memoryview(
+    monkeypatch: pytest.MonkeyPatch,
+    shape: tuple[int, ...],
+) -> None:
+    original_write = os.write
+
+    def partial_write(fd: int, data: bytes | bytearray | memoryview) -> int:
+        with memoryview(data).cast("B") as view:
+            return original_write(fd, view[:3])
+
+    monkeypatch.setattr(os, "write", partial_write)
+    payload = bytes(range(32)) if shape else b"ab"
+    send, receive = await make_pipe()
+    async with send, receive:
+        with memoryview(payload).cast("H", shape=shape) as data:
+            await send.send_all(data)
+        await send.aclose()
+        received = bytearray()
+        async for chunk in receive:
+            received.extend(chunk)
+        assert received == payload
+
+
+async def test_send_pipe_empty_noncontiguous_memoryview() -> None:
+    send, receive = await make_pipe()
+    async with send, receive:
+        with memoryview(b"abc")[0:0:2] as data:
+            with assert_checkpoints():
+                await send.send_all(data)
+        await send.aclose()
+        assert await receive.receive_some() == b""
+
+
 async def test_receive_pipe() -> None:
     r, w = os.pipe()
     async with FdStream(r) as recv:

src/trio/_unix_pipes.py

@@ -134,23 +134,27 @@ async def send_all(self, data: bytes | bytearray | memoryview) -> None:
             if self._fd_holder.closed:
                 raise trio.ClosedResourceError("file was already closed")
             await trio.lowlevel.checkpoint()
-            length = len(data)
             # adapted from the SocketStream code
-            with memoryview(data) as view:
-                sent = 0
-                while sent < length:
-                    with view[sent:] as remaining:
-                        try:
-                            sent += os.write(self._fd_holder.fd, remaining)
-                        except BlockingIOError:
-                            await trio.lowlevel.wait_writable(self._fd_holder.fd)
-                        except OSError as e:
-                            if e.errno == errno.EBADF:
-                                raise trio.ClosedResourceError(
-                                    "file was already closed",
-                                ) from None
-                            else:
-                                raise trio.BrokenResourceError from e
+            with memoryview(data) as raw_view:
+                length = raw_view.nbytes
+                if not length:
+                    return
+                # os.write() returns byte counts, regardless of the input format.
+                with raw_view.cast("B") as view:
+                    sent = 0
+                    while sent < length:
+                        with view[sent:] as remaining:
+                            try:
+                                sent += os.write(self._fd_holder.fd, remaining)
+                            except BlockingIOError:
+                                await trio.lowlevel.wait_writable(self._fd_holder.fd)
+                            except OSError as e:
+                                if e.errno == errno.EBADF:
+                                    raise trio.ClosedResourceError(
+                                        "file was already closed",
+                                    ) from None
+                                else:
+                                    raise trio.BrokenResourceError from e
 
     async def wait_send_all_might_not_block(self) -> None:
         with self._send_conflict_detector: