Skip to content

Commit aba1797

Browse files
authored
Fix HTTPS httpx mocking in wrap_bio() plus regressions occurred (#324)
* fix httpx https wrap_bio address handling * refine ssl wrap_bio regression fix * clarify ssl wrap_bio address assignment * harden ssl wrap_bio hostname decoding * simplify ssl wrap_bio fallback logic * handle empty ssl wrap_bio hostnames explicitly * rename ssl wrap_bio hostname locals * preserve ssl wrap_bio port state consistently * Prevent deadlock when mocking large async HTTP responses. * Fix SSL host fallback check for empty server_hostname * Preserve empty SSL host when filling missing peer cert fields * Add more tests for coverage and preventing regressions.
1 parent c5daccb commit aba1797

6 files changed

Lines changed: 456 additions & 25 deletions

File tree

mocket/io.py

Lines changed: 2 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -3,9 +3,6 @@
33
from __future__ import annotations
44

55
import io
6-
import os
7-
8-
from mocket.mocket import Mocket
96

107

118
class MocketSocketIO(io.BytesIO):
@@ -21,17 +18,12 @@ def __init__(self, address: tuple) -> None:
2118
super().__init__()
2219

2320
def write(self, content: bytes) -> int:
24-
"""Write content to the buffer and the pipe if available.
21+
"""Write content to the in-memory buffer.
2522
2623
Args:
2724
content: Bytes to write
2825
2926
Returns:
3027
Number of bytes written
3128
"""
32-
super().write(content)
33-
34-
_, w_fd = Mocket.get_pair(self._address)
35-
if w_fd:
36-
os.write(w_fd, content)
37-
return len(content)
29+
return super().write(content)

mocket/mocket.py

Lines changed: 41 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -22,7 +22,10 @@
2222
class Mocket:
2323
"""Singleton class managing all mock socket operations and entries."""
2424

25+
_socket_ios: ClassVar[dict[Address, Any]] = {}
2526
_socket_pairs: ClassVar[dict[Address, tuple[int, int]]] = {}
27+
_pipe_uses_data: ClassVar[dict[Address, bool]] = {}
28+
_pending_readables: ClassVar[dict[Address, int]] = {}
2629
_address: ClassVar[Address | tuple[None, None]] = (None, None)
2730
_entries: ClassVar[dict[Address, list[MocketEntry]]] = collections.defaultdict(list)
2831
_requests: ClassVar[list] = []
@@ -90,6 +93,39 @@ def set_pair(cls, address: Address, pair: tuple[int, int]) -> None:
9093
"""
9194
cls._socket_pairs[address] = pair
9295

96+
@classmethod
97+
def get_io(cls, address: Address) -> Any:
98+
"""Get the shared socket I/O buffer for an address, if any."""
99+
return cls._socket_ios.get(address)
100+
101+
@classmethod
102+
def set_io(cls, address: Address, socket_io: Any) -> None:
103+
"""Store the shared socket I/O buffer for an address."""
104+
cls._socket_ios[address] = socket_io
105+
106+
@classmethod
107+
def pipe_uses_data(cls, address: Address) -> bool:
108+
"""Return whether the pipe for an address carries mirrored response bytes."""
109+
return cls._pipe_uses_data.get(address, False)
110+
111+
@classmethod
112+
def set_pipe_uses_data(cls, address: Address, uses_data: bool) -> None:
113+
"""Store whether the pipe for an address carries mirrored response bytes."""
114+
cls._pipe_uses_data[address] = uses_data
115+
116+
@classmethod
117+
def get_pending_readables(cls, address: Address) -> int:
118+
"""Get the number of readiness bytes currently queued for an address."""
119+
return cls._pending_readables.get(address, 0)
120+
121+
@classmethod
122+
def set_pending_readables(cls, address: Address, count: int) -> None:
123+
"""Store the number of readiness bytes queued for an address."""
124+
if count > 0:
125+
cls._pending_readables[address] = count
126+
else:
127+
cls._pending_readables.pop(address, None)
128+
93129
@classmethod
94130
def register(cls, *entries: MocketEntry) -> None:
95131
"""Register mock entries with Mocket.
@@ -132,10 +168,15 @@ def collect(cls, data: Any) -> None:
132168
@classmethod
133169
def reset(cls) -> None:
134170
"""Reset all Mocket state and clean up file descriptors."""
171+
for socket_io in cls._socket_ios.values():
172+
socket_io.close()
135173
for r_fd, w_fd in cls._socket_pairs.values():
136174
os.close(r_fd)
137175
os.close(w_fd)
176+
cls._socket_ios = {}
138177
cls._socket_pairs = {}
178+
cls._pipe_uses_data = {}
179+
cls._pending_readables = {}
139180
cls._entries = collections.defaultdict(list)
140181
cls._requests = []
141182
cls._record_storage = None

mocket/socket.py

Lines changed: 109 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -200,8 +200,14 @@ def proto(self) -> int:
200200
@property
201201
def io(self) -> MocketSocketIO:
202202
"""Get or create the socket I/O object."""
203-
if self._io is None:
204-
self._io = MocketSocketIO((self._host, self._port))
203+
if self._io is None or getattr(self._io, "closed", False):
204+
address = self._address_key()
205+
self._io = Mocket.get_io(address)
206+
if self._io is not None and getattr(self._io, "closed", False):
207+
self._io = None
208+
if self._io is None:
209+
self._io = MocketSocketIO(address)
210+
Mocket.set_io(address, self._io)
205211
return self._io
206212

207213
def fileno(self) -> int:
@@ -214,9 +220,76 @@ def fileno(self) -> int:
214220
r_fd, _ = Mocket.get_pair(address)
215221
if not r_fd:
216222
r_fd, w_fd = os.pipe()
223+
os.set_blocking(r_fd, False)
224+
os.set_blocking(w_fd, False)
217225
Mocket.set_pair(address, (r_fd, w_fd))
226+
if self._io is not None and self._buffered_bytes():
227+
if Mocket.pipe_uses_data(address):
228+
self._mirror_buffer_to_pipe()
229+
else:
230+
self._sync_readable_pipe()
218231
return r_fd
219232

233+
def _address_key(self) -> Address:
234+
"""Return the current socket address tuple."""
235+
return self._host, self._port
236+
237+
def _buffered_bytes(self) -> int:
238+
"""Return the number of unread bytes buffered in the socket I/O."""
239+
return len(self.io.getvalue()) - self.io.tell()
240+
241+
def _clear_readable_pipe(self) -> None:
242+
"""Drain any stale readiness bytes from the pipe for this socket."""
243+
address = self._address_key()
244+
r_fd, _ = Mocket.get_pair(address)
245+
if not r_fd:
246+
return
247+
248+
while True:
249+
try:
250+
if not os.read(r_fd, self._buflen):
251+
break
252+
except BlockingIOError:
253+
break
254+
255+
Mocket.set_pending_readables(address, 0)
256+
257+
def _mirror_buffer_to_pipe(self) -> None:
258+
"""Mirror unread response bytes into the pipe for small payloads."""
259+
address = self._address_key()
260+
_, w_fd = Mocket.get_pair(address)
261+
if not w_fd:
262+
return
263+
264+
unread = self.io.getvalue()[self.io.tell() :]
265+
if unread:
266+
os.write(w_fd, unread)
267+
268+
def _sync_readable_pipe(self) -> None:
269+
"""Keep the readiness pipe in sync with the unread buffer size.
270+
271+
The pipe is used only to wake selector-based async clients. Response bytes
272+
remain in the in-memory buffer so large payloads do not block on OS pipe
273+
capacity.
274+
"""
275+
address = self._address_key()
276+
_, w_fd = Mocket.get_pair(address)
277+
if not w_fd:
278+
return
279+
280+
pending = Mocket.get_pending_readables(address)
281+
desired = min(self._buffered_bytes(), self._buflen)
282+
if desired <= pending:
283+
return
284+
285+
try:
286+
written = os.write(w_fd, b"\0" * (desired - pending))
287+
except BlockingIOError:
288+
written = 0
289+
290+
if written:
291+
Mocket.set_pending_readables(address, pending + written)
292+
220293
def gettimeout(self) -> float | None:
221294
"""Get the socket timeout.
222295
@@ -375,10 +448,20 @@ def sendall(
375448
response = self.true_sendall(data, *args, **kwargs)
376449

377450
if response is not None:
451+
address = self._address_key()
452+
# Ensure the address pipe exists before deciding whether to mirror
453+
# response bytes or only publish readiness signals.
454+
self.fileno()
378455
self.io.seek(0)
456+
self._clear_readable_pipe()
379457
self.io.write(response)
380458
self.io.truncate()
381459
self.io.seek(0)
460+
Mocket.set_pipe_uses_data(address, len(response) <= self._buflen)
461+
if Mocket.pipe_uses_data(address):
462+
self._mirror_buffer_to_pipe()
463+
else:
464+
self._sync_readable_pipe()
382465

383466
def sendmsg(
384467
self,
@@ -537,11 +620,32 @@ def recv(self, buffersize: int, flags: int | None = None) -> bytes:
537620
Raises:
538621
BlockingIOError: If socket is non-blocking and no data available
539622
"""
540-
r_fd, _ = Mocket.get_pair((self._host, self._port))
541-
if r_fd:
542-
return os.read(r_fd, buffersize)
623+
if buffersize is None:
624+
buffersize = self._buflen
625+
626+
address = self._address_key()
627+
r_fd, _ = Mocket.get_pair(address)
628+
if r_fd and Mocket.pipe_uses_data(address):
629+
try:
630+
pipe_data = os.read(r_fd, buffersize)
631+
except BlockingIOError:
632+
pipe_data = b""
633+
if pipe_data:
634+
# Keep in-memory buffer position in sync with bytes drained from the pipe.
635+
self.io.seek(self.io.tell() + len(pipe_data))
636+
return pipe_data
637+
638+
pending = Mocket.get_pending_readables(address)
639+
if r_fd and self._buffered_bytes() and pending == 0:
640+
self._sync_readable_pipe()
641+
pending = Mocket.get_pending_readables(address)
642+
543643
data = self.io.read(buffersize)
544644
if data:
645+
if r_fd and pending:
646+
drained = os.read(r_fd, min(len(data), pending))
647+
Mocket.set_pending_readables(address, pending - len(drained))
648+
self._sync_readable_pipe()
545649
return data
546650
# used by Redis mock
547651
exc = BlockingIOError()

mocket/ssl/context.py

Lines changed: 18 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@
44

55
from typing import Any
66

7+
from mocket.mocket import Mocket
78
from mocket.socket import MocketSocket
89
from mocket.ssl.socket import MocketSSLSocket
910

@@ -111,7 +112,23 @@ def wrap_bio(
111112
MocketSSLSocket instance
112113
"""
113114
ssl_obj = MocketSSLSocket()
114-
ssl_obj._host = server_hostname
115+
if isinstance(server_hostname, bytes):
116+
hostname = server_hostname.decode("utf-8", errors="replace")
117+
else:
118+
hostname = server_hostname
119+
120+
current_address = Mocket._address
121+
if isinstance(current_address, tuple) and len(current_address) == 2:
122+
current_host, current_port = current_address
123+
else:
124+
current_host, current_port = None, None
125+
126+
effective_host = hostname if hostname is not None else current_host
127+
ssl_obj._host = effective_host
128+
if current_port is not None:
129+
ssl_obj._port = current_port
130+
if effective_host is not None and current_port is not None:
131+
ssl_obj._address = (effective_host, current_port)
115132
return ssl_obj
116133

117134

mocket/ssl/socket.py

Lines changed: 40 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,9 @@ def __init__(self, *args: Any, **kwargs: Any) -> None:
2727

2828
self._did_handshake: bool = False
2929
self._sent_non_empty_bytes: bool = False
30+
self._has_written: bool = False
31+
self._ssl_pending: bytes = b""
32+
self._ssl_pending_pos: int = 0
3033
self._original_socket: MocketSocket = self
3134

3235
def read(self, buffersize: int | None = None) -> bytes:
@@ -38,13 +41,28 @@ def read(self, buffersize: int | None = None) -> bytes:
3841
Returns:
3942
Bytes read from the socket
4043
41-
Raises:
42-
ssl.SSLWantReadError: If handshake not completed and no data
4344
"""
44-
rv = self.io.read(buffersize)
45+
if self._ssl_pending_pos < len(self._ssl_pending):
46+
if buffersize is None:
47+
rv = self._ssl_pending[self._ssl_pending_pos :]
48+
self._ssl_pending_pos = len(self._ssl_pending)
49+
else:
50+
end = self._ssl_pending_pos + buffersize
51+
rv = self._ssl_pending[self._ssl_pending_pos : end]
52+
self._ssl_pending_pos = min(end, len(self._ssl_pending))
53+
else:
54+
rv = b""
4555
if rv:
4656
self._sent_non_empty_bytes = True
47-
if self._did_handshake and not self._sent_non_empty_bytes:
57+
58+
# asyncio SSL transports probe reads before writing request bytes.
59+
# Keep that non-blocking behavior, but once writes happened we must
60+
# return empty bytes instead of surfacing SSLWantReadError.
61+
if (
62+
self._did_handshake
63+
and not self._sent_non_empty_bytes
64+
and not self._has_written
65+
):
4866
raise ssl.SSLWantReadError("The operation did not complete (read)")
4967
return rv
5068

@@ -57,7 +75,14 @@ def write(self, data: bytes) -> int | None:
5775
Returns:
5876
Number of bytes written
5977
"""
60-
return self.send(encode_to_bytes(data))
78+
self._has_written = self._has_written or bool(data)
79+
bytes_sent = self.send(encode_to_bytes(data))
80+
81+
# Keep a private read buffer for SSL protocol consumers so response
82+
# parsing does not depend on shared socket I/O cursor state.
83+
self._ssl_pending = self.io.getvalue()
84+
self._ssl_pending_pos = 0
85+
return bytes_sent
6186

6287
def do_handshake(self) -> None:
6388
"""Perform SSL handshake (mock implementation)."""
@@ -72,8 +97,13 @@ def getpeercert(self, binary_form: bool = False) -> _PeerCertRetDictType:
7297
Returns:
7398
Mock certificate dictionary
7499
"""
75-
if not (self._host and self._port):
76-
self._address = self._host, self._port = Mocket._address
100+
if self._host is None or self._port is None:
101+
current_host, current_port = Mocket._address
102+
if self._host is None:
103+
self._host = current_host
104+
if self._port is None:
105+
self._port = current_port
106+
self._address = (self._host, self._port)
77107

78108
now = datetime.now()
79109
shift = now + timedelta(days=30 * 12)
@@ -156,5 +186,8 @@ def _create(
156186

157187
ssl_socket._io = sock._io
158188
ssl_socket._entry = sock._entry
189+
ssl_socket._has_written = getattr(sock, "_has_written", False)
190+
ssl_socket._ssl_pending = getattr(sock, "_ssl_pending", b"")
191+
ssl_socket._ssl_pending_pos = getattr(sock, "_ssl_pending_pos", 0)
159192

160193
return ssl_socket

0 commit comments

Comments
 (0)