diff --git a/mocket/io.py b/mocket/io.py index e815e0e..29e370d 100644 --- a/mocket/io.py +++ b/mocket/io.py @@ -3,9 +3,6 @@ from __future__ import annotations import io -import os - -from mocket.mocket import Mocket class MocketSocketIO(io.BytesIO): @@ -21,7 +18,7 @@ def __init__(self, address: tuple) -> None: super().__init__() def write(self, content: bytes) -> int: - """Write content to the buffer and the pipe if available. + """Write content to the in-memory buffer. Args: content: Bytes to write @@ -29,9 +26,4 @@ def write(self, content: bytes) -> int: Returns: Number of bytes written """ - super().write(content) - - _, w_fd = Mocket.get_pair(self._address) - if w_fd: - os.write(w_fd, content) - return len(content) + return super().write(content) diff --git a/mocket/mocket.py b/mocket/mocket.py index 75ae628..9baa2ac 100644 --- a/mocket/mocket.py +++ b/mocket/mocket.py @@ -22,7 +22,10 @@ class Mocket: """Singleton class managing all mock socket operations and entries.""" + _socket_ios: ClassVar[dict[Address, Any]] = {} _socket_pairs: ClassVar[dict[Address, tuple[int, int]]] = {} + _pipe_uses_data: ClassVar[dict[Address, bool]] = {} + _pending_readables: ClassVar[dict[Address, int]] = {} _address: ClassVar[Address | tuple[None, None]] = (None, None) _entries: ClassVar[dict[Address, list[MocketEntry]]] = collections.defaultdict(list) _requests: ClassVar[list] = [] @@ -90,6 +93,39 @@ def set_pair(cls, address: Address, pair: tuple[int, int]) -> None: """ cls._socket_pairs[address] = pair + @classmethod + def get_io(cls, address: Address) -> Any: + """Get the shared socket I/O buffer for an address, if any.""" + return cls._socket_ios.get(address) + + @classmethod + def set_io(cls, address: Address, socket_io: Any) -> None: + """Store the shared socket I/O buffer for an address.""" + cls._socket_ios[address] = socket_io + + @classmethod + def pipe_uses_data(cls, address: Address) -> bool: + """Return whether the pipe for an address carries mirrored response bytes.""" + return cls._pipe_uses_data.get(address, False) + + @classmethod + def set_pipe_uses_data(cls, address: Address, uses_data: bool) -> None: + """Store whether the pipe for an address carries mirrored response bytes.""" + cls._pipe_uses_data[address] = uses_data + + @classmethod + def get_pending_readables(cls, address: Address) -> int: + """Get the number of readiness bytes currently queued for an address.""" + return cls._pending_readables.get(address, 0) + + @classmethod + def set_pending_readables(cls, address: Address, count: int) -> None: + """Store the number of readiness bytes queued for an address.""" + if count > 0: + cls._pending_readables[address] = count + else: + cls._pending_readables.pop(address, None) + @classmethod def register(cls, *entries: MocketEntry) -> None: """Register mock entries with Mocket. @@ -132,10 +168,15 @@ def collect(cls, data: Any) -> None: @classmethod def reset(cls) -> None: """Reset all Mocket state and clean up file descriptors.""" + for socket_io in cls._socket_ios.values(): + socket_io.close() for r_fd, w_fd in cls._socket_pairs.values(): os.close(r_fd) os.close(w_fd) + cls._socket_ios = {} cls._socket_pairs = {} + cls._pipe_uses_data = {} + cls._pending_readables = {} cls._entries = collections.defaultdict(list) cls._requests = [] cls._record_storage = None diff --git a/mocket/socket.py b/mocket/socket.py index bd79528..a23887e 100644 --- a/mocket/socket.py +++ b/mocket/socket.py @@ -200,8 +200,14 @@ def proto(self) -> int: @property def io(self) -> MocketSocketIO: """Get or create the socket I/O object.""" - if self._io is None: - self._io = MocketSocketIO((self._host, self._port)) + if self._io is None or getattr(self._io, "closed", False): + address = self._address_key() + self._io = Mocket.get_io(address) + if self._io is not None and getattr(self._io, "closed", False): + self._io = None + if self._io is None: + self._io = MocketSocketIO(address) + Mocket.set_io(address, self._io) return self._io def fileno(self) -> int: @@ -214,9 +220,76 @@ def fileno(self) -> int: r_fd, _ = Mocket.get_pair(address) if not r_fd: r_fd, w_fd = os.pipe() + os.set_blocking(r_fd, False) + os.set_blocking(w_fd, False) Mocket.set_pair(address, (r_fd, w_fd)) + if self._io is not None and self._buffered_bytes(): + if Mocket.pipe_uses_data(address): + self._mirror_buffer_to_pipe() + else: + self._sync_readable_pipe() return r_fd + def _address_key(self) -> Address: + """Return the current socket address tuple.""" + return self._host, self._port + + def _buffered_bytes(self) -> int: + """Return the number of unread bytes buffered in the socket I/O.""" + return len(self.io.getvalue()) - self.io.tell() + + def _clear_readable_pipe(self) -> None: + """Drain any stale readiness bytes from the pipe for this socket.""" + address = self._address_key() + r_fd, _ = Mocket.get_pair(address) + if not r_fd: + return + + while True: + try: + if not os.read(r_fd, self._buflen): + break + except BlockingIOError: + break + + Mocket.set_pending_readables(address, 0) + + def _mirror_buffer_to_pipe(self) -> None: + """Mirror unread response bytes into the pipe for small payloads.""" + address = self._address_key() + _, w_fd = Mocket.get_pair(address) + if not w_fd: + return + + unread = self.io.getvalue()[self.io.tell() :] + if unread: + os.write(w_fd, unread) + + def _sync_readable_pipe(self) -> None: + """Keep the readiness pipe in sync with the unread buffer size. + + The pipe is used only to wake selector-based async clients. Response bytes + remain in the in-memory buffer so large payloads do not block on OS pipe + capacity. + """ + address = self._address_key() + _, w_fd = Mocket.get_pair(address) + if not w_fd: + return + + pending = Mocket.get_pending_readables(address) + desired = min(self._buffered_bytes(), self._buflen) + if desired <= pending: + return + + try: + written = os.write(w_fd, b"\0" * (desired - pending)) + except BlockingIOError: + written = 0 + + if written: + Mocket.set_pending_readables(address, pending + written) + def gettimeout(self) -> float | None: """Get the socket timeout. @@ -375,10 +448,20 @@ def sendall( response = self.true_sendall(data, *args, **kwargs) if response is not None: + address = self._address_key() + # Ensure the address pipe exists before deciding whether to mirror + # response bytes or only publish readiness signals. + self.fileno() self.io.seek(0) + self._clear_readable_pipe() self.io.write(response) self.io.truncate() self.io.seek(0) + Mocket.set_pipe_uses_data(address, len(response) <= self._buflen) + if Mocket.pipe_uses_data(address): + self._mirror_buffer_to_pipe() + else: + self._sync_readable_pipe() def sendmsg( self, @@ -537,11 +620,32 @@ def recv(self, buffersize: int, flags: int | None = None) -> bytes: Raises: BlockingIOError: If socket is non-blocking and no data available """ - r_fd, _ = Mocket.get_pair((self._host, self._port)) - if r_fd: - return os.read(r_fd, buffersize) + if buffersize is None: + buffersize = self._buflen + + address = self._address_key() + r_fd, _ = Mocket.get_pair(address) + if r_fd and Mocket.pipe_uses_data(address): + try: + pipe_data = os.read(r_fd, buffersize) + except BlockingIOError: + pipe_data = b"" + if pipe_data: + # Keep in-memory buffer position in sync with bytes drained from the pipe. + self.io.seek(self.io.tell() + len(pipe_data)) + return pipe_data + + pending = Mocket.get_pending_readables(address) + if r_fd and self._buffered_bytes() and pending == 0: + self._sync_readable_pipe() + pending = Mocket.get_pending_readables(address) + data = self.io.read(buffersize) if data: + if r_fd and pending: + drained = os.read(r_fd, min(len(data), pending)) + Mocket.set_pending_readables(address, pending - len(drained)) + self._sync_readable_pipe() return data # used by Redis mock exc = BlockingIOError() diff --git a/mocket/ssl/context.py b/mocket/ssl/context.py index aeaab6b..be688cc 100644 --- a/mocket/ssl/context.py +++ b/mocket/ssl/context.py @@ -4,6 +4,7 @@ from typing import Any +from mocket.mocket import Mocket from mocket.socket import MocketSocket from mocket.ssl.socket import MocketSSLSocket @@ -111,7 +112,23 @@ def wrap_bio( MocketSSLSocket instance """ ssl_obj = MocketSSLSocket() - ssl_obj._host = server_hostname + if isinstance(server_hostname, bytes): + hostname = server_hostname.decode("utf-8", errors="replace") + else: + hostname = server_hostname + + current_address = Mocket._address + if isinstance(current_address, tuple) and len(current_address) == 2: + current_host, current_port = current_address + else: + current_host, current_port = None, None + + effective_host = hostname if hostname is not None else current_host + ssl_obj._host = effective_host + if current_port is not None: + ssl_obj._port = current_port + if effective_host is not None and current_port is not None: + ssl_obj._address = (effective_host, current_port) return ssl_obj diff --git a/mocket/ssl/socket.py b/mocket/ssl/socket.py index 94984fc..baf7c10 100644 --- a/mocket/ssl/socket.py +++ b/mocket/ssl/socket.py @@ -27,6 +27,9 @@ def __init__(self, *args: Any, **kwargs: Any) -> None: self._did_handshake: bool = False self._sent_non_empty_bytes: bool = False + self._has_written: bool = False + self._ssl_pending: bytes = b"" + self._ssl_pending_pos: int = 0 self._original_socket: MocketSocket = self def read(self, buffersize: int | None = None) -> bytes: @@ -38,13 +41,28 @@ def read(self, buffersize: int | None = None) -> bytes: Returns: Bytes read from the socket - Raises: - ssl.SSLWantReadError: If handshake not completed and no data """ - rv = self.io.read(buffersize) + if self._ssl_pending_pos < len(self._ssl_pending): + if buffersize is None: + rv = self._ssl_pending[self._ssl_pending_pos :] + self._ssl_pending_pos = len(self._ssl_pending) + else: + end = self._ssl_pending_pos + buffersize + rv = self._ssl_pending[self._ssl_pending_pos : end] + self._ssl_pending_pos = min(end, len(self._ssl_pending)) + else: + rv = b"" if rv: self._sent_non_empty_bytes = True - if self._did_handshake and not self._sent_non_empty_bytes: + + # asyncio SSL transports probe reads before writing request bytes. + # Keep that non-blocking behavior, but once writes happened we must + # return empty bytes instead of surfacing SSLWantReadError. + if ( + self._did_handshake + and not self._sent_non_empty_bytes + and not self._has_written + ): raise ssl.SSLWantReadError("The operation did not complete (read)") return rv @@ -57,7 +75,14 @@ def write(self, data: bytes) -> int | None: Returns: Number of bytes written """ - return self.send(encode_to_bytes(data)) + self._has_written = self._has_written or bool(data) + bytes_sent = self.send(encode_to_bytes(data)) + + # Keep a private read buffer for SSL protocol consumers so response + # parsing does not depend on shared socket I/O cursor state. + self._ssl_pending = self.io.getvalue() + self._ssl_pending_pos = 0 + return bytes_sent def do_handshake(self) -> None: """Perform SSL handshake (mock implementation).""" @@ -72,8 +97,13 @@ def getpeercert(self, binary_form: bool = False) -> _PeerCertRetDictType: Returns: Mock certificate dictionary """ - if not (self._host and self._port): - self._address = self._host, self._port = Mocket._address + if self._host is None or self._port is None: + current_host, current_port = Mocket._address + if self._host is None: + self._host = current_host + if self._port is None: + self._port = current_port + self._address = (self._host, self._port) now = datetime.now() shift = now + timedelta(days=30 * 12) @@ -156,5 +186,8 @@ def _create( ssl_socket._io = sock._io ssl_socket._entry = sock._entry + ssl_socket._has_written = getattr(sock, "_has_written", False) + ssl_socket._ssl_pending = getattr(sock, "_ssl_pending", b"") + ssl_socket._ssl_pending_pos = getattr(sock, "_ssl_pending_pos", 0) return ssl_socket diff --git a/tests/test_socket.py b/tests/test_socket.py index 68e71ae..31e0d63 100644 --- a/tests/test_socket.py +++ b/tests/test_socket.py @@ -1,11 +1,16 @@ +import os import socket +import ssl import struct from unittest.mock import MagicMock import pytest -from mocket import Mocket, MocketEntry, mocketize +from mocket import Mocket, MocketEntry, Mocketizer, mocketize +from mocket.mockhttp import Entry from mocket.socket import MocketSocket +from mocket.ssl.context import MocketSSLContext +from mocket.ssl.socket import MocketSSLSocket @pytest.mark.parametrize("blocking", (False, True)) @@ -119,6 +124,59 @@ def test_getsockopt(): assert result == socket.SOCK_STREAM +@pytest.mark.parametrize( + ("server_hostname", "expected_host"), + [ + (b"httpbin.local", "httpbin.local"), + ("httpbin.local", "httpbin.local"), + (None, "httpbin.local"), + ("", ""), + (b"mocket-\xff.local", "mocket-�.local"), + ], +) +def test_wrap_bio_uses_current_mocket_address( + monkeypatch, server_hostname, expected_host +): + monkeypatch.setattr(Mocket, "_address", ("httpbin.local", 443)) + ssl_obj = MocketSSLContext().wrap_bio( + incoming=None, + outgoing=None, + server_hostname=server_hostname, + ) + + assert ssl_obj._host == expected_host + assert ssl_obj._port == 443 + assert ssl_obj._address == (expected_host, 443) + + +def test_wrap_bio_preserves_empty_server_hostname_on_getpeercert(monkeypatch): + monkeypatch.setattr(Mocket, "_address", ("httpbin.local", 443)) + ssl_obj = MocketSSLContext().wrap_bio( + incoming=None, + outgoing=None, + server_hostname="", + ) + + ssl_obj.getpeercert() + + assert ssl_obj._host == "" + assert ssl_obj._address == ("", 443) + + +def test_getpeercert_does_not_overwrite_empty_host_when_port_missing(monkeypatch): + monkeypatch.setattr(Mocket, "_address", ("httpbin.local", 443)) + ssl_obj = MocketSSLSocket() + ssl_obj._host = "" + ssl_obj._port = None + ssl_obj._address = ("", None) + + ssl_obj.getpeercert() + + assert ssl_obj._host == "" + assert ssl_obj._port == 443 + assert ssl_obj._address == ("", 443) + + def test_recvfrom_into(): sock = MocketSocket(socket.AF_INET, socket.SOCK_STREAM) test_data = b"abc123" @@ -143,7 +201,193 @@ def test_setsockopt_with_optlen(): sock = MocketSocket(socket.AF_INET, socket.SOCK_STREAM) sock._true_socket = MagicMock() linger_value = struct.pack("ii", 1, 5) - sock.setsockopt(socket.SOL_SOCKET, socket.SO_LINGER, linger_value, len(linger_value)) + sock.setsockopt( + socket.SOL_SOCKET, socket.SO_LINGER, linger_value, len(linger_value) + ) sock._true_socket.setsockopt.assert_called_once_with( socket.SOL_SOCKET, socket.SO_LINGER, linger_value, len(linger_value) ) + + +def test_ssl_read_empty_after_handshake_returns_empty_bytes(): + """After handshake, empty SSL reads should not raise SSLWantReadError.""" + sock = MocketSSLSocket() + sock._io = type("MockIO", (), {"read": lambda self, n: b""})() + sock._did_handshake = True + sock._has_written = True + + assert sock.read(1024) == b"" + + +def test_ssl_read_empty_after_handshake_before_write_raises_want_read(): + """After handshake but before writes, empty reads should signal WANT_READ.""" + sock = MocketSSLSocket() + sock._io = type("MockIO", (), {"read": lambda self, n: b""})() + sock._did_handshake = True + sock._has_written = False + + with pytest.raises(ssl.SSLWantReadError): + sock.read(1024) + + +def test_ssl_ciper_returns_mock_tuple(): + """Exercise the SSL mock cipher tuple branch for coverage.""" + sock = MocketSSLSocket() + assert sock.ciper() == ("ADH", "AES256", "SHA") + + +def test_ssl_getpeercert_uses_mocket_address_when_unset(monkeypatch): + """Cover getpeercert fallback when host/port are not yet assigned.""" + monkeypatch.setattr(Mocket, "_address", ("example.local", 443)) + sock = MocketSSLSocket() + + cert = sock.getpeercert() + + assert sock._host == "example.local" + assert sock._port == 443 + assert sock._address == ("example.local", 443) + assert cert["subjectAltName"][1] == ("DNS", "example.local") + + +# --------------------------------------------------------------------------- +# New pipe-mechanism tests added to maintain coverage after the large-response +# deadlock fix. +# --------------------------------------------------------------------------- + + +def test_mocket_shared_io_same_address(): + """Two MocketSocket instances for the same address share one I/O buffer.""" + addr = ("localhost", 9001) + with Mocketizer(): + s1 = MocketSocket() + s1.connect(addr) + s2 = MocketSocket() + s2.connect(addr) + # Accessing .io lazily creates and registers the shared buffer + assert s1.io is s2.io + + +def test_mocket_shared_io_recreated_after_close(): + """If the shared buffer is closed, accessing .io creates a fresh one.""" + addr = ("localhost", 9002) + with Mocketizer(): + s = MocketSocket() + s.connect(addr) + old_io = s.io + old_io.close() + # Force re-evaluation by clearing the cached reference + s._io = None + new_io = s.io + assert not new_io.closed + assert new_io is not old_io + + +def test_mocket_pipe_uses_data_false_for_large_response(): + """Responses larger than _buflen use readiness-only signaling, not data in pipe.""" + addr = ("localhost", 9003) + large_body = "x" * 70_000 + Entry.single_register( + method=Entry.GET, + uri="http://localhost:9003/", + body=large_body, + ) + with Mocketizer(): + s = MocketSocket() + s.connect(addr) + request = b"GET / HTTP/1.1\r\nHost: localhost\r\n\r\n" + s.sendall(request) + assert not Mocket.pipe_uses_data(addr), ( + "large response must use readiness-only pipe" + ) + + +@mocketize +def test_large_response_recv_drains_readiness_sentinel(): + """recv() for a large response drains readiness bytes and re-syncs the pipe. + + This exercises the _sync_readable_pipe() re-sync call (socket.py line 638) + that runs after draining sentinel bytes from a readiness-only pipe. + """ + addr = ("localhost", 9004) + large_body = "y" * 70_000 + Entry.single_register( + method=Entry.GET, + uri="http://localhost:9004/", + body=large_body, + ) + s = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + s.connect(addr) + # Calling fileno() forces pipe creation so the readiness path is active. + s.fileno() + s.sendall(b"GET / HTTP/1.1\r\nHost: localhost\r\n\r\n") + + # The response is large → pipe_uses_data=False → readiness sentinels are used. + # Collecting the full body exercises the drain + re-sync branch on line 638. + received = b"" + while True: + try: + chunk = s.recv(65536) + if not chunk: + break + received += chunk + except BlockingIOError: + break + s.close() + assert large_body.encode() in received + + +def test_mocket_get_set_io(): + """Mocket.get_io / set_io round-trip.""" + from mocket.io import MocketSocketIO + + addr = ("localhost", 9005) + buf = MocketSocketIO(addr) + Mocket.set_io(addr, buf) + assert Mocket.get_io(addr) is buf + + +def test_mocket_pipe_uses_data_flag(): + """Mocket.pipe_uses_data / set_pipe_uses_data round-trip.""" + addr = ("localhost", 9006) + assert not Mocket.pipe_uses_data(addr) + Mocket.set_pipe_uses_data(addr, True) + assert Mocket.pipe_uses_data(addr) + Mocket.set_pipe_uses_data(addr, False) + assert not Mocket.pipe_uses_data(addr) + + +def test_mocket_reset_clears_shared_ios(): + """Mocket.reset() clears the shared I/O buffer registry.""" + from mocket.io import MocketSocketIO + + addr = ("localhost", 9007) + Mocket.set_io(addr, MocketSocketIO(addr)) + Mocket.reset() + assert Mocket.get_io(addr) is None + + +def test_mirror_buffer_to_pipe_writes_data(): + """_mirror_buffer_to_pipe writes the unread buffer content into the pipe.""" + addr = ("localhost", 9008) + s = MocketSocket() + s.connect(addr) + + # Simulate a small response already buffered (no sendall needed) + response = b"hello pipe" + s.io.seek(0) + s.io.write(response) + s.io.seek(0) + + # Create the pipe manually + r_fd, w_fd = os.pipe() + os.set_blocking(r_fd, False) + os.set_blocking(w_fd, False) + Mocket.set_pair(addr, (r_fd, w_fd)) + + s._mirror_buffer_to_pipe() + pipe_bytes = os.read(r_fd, 64) + os.close(r_fd) + os.close(w_fd) + Mocket._socket_pairs.pop(addr, None) + + assert pipe_bytes == response