Skip to content

Commit 0855c8a

Browse files
cbbm142John Pastore
authored andcommitted
Preserve write buffer when raw.write() returns None on blocked socket
Add a check in cheroot.makefile.BufferedWriter._flush_unlocked to prevent clearing of the write buffer when raw.write() returns None due to a blocked stream. Also limits each write call to SOCK_WRITE_BLOCKSIZE bytes, which was defined but not previously used in this method.
1 parent 2ffb0ba commit 0855c8a

2 files changed

Lines changed: 42 additions & 2 deletions

File tree

cheroot/makefile.py

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -26,13 +26,17 @@ def write(self, b):
2626
def _flush_unlocked(self):
2727
self._checkClosed('flush of closed file')
2828
while self._write_buf:
29+
n = None
2930
try:
3031
# ssl sockets only except 'bytes', not bytearrays
3132
# so perhaps we should conditionally wrap this for perf?
32-
n = self.raw.write(bytes(self._write_buf))
33+
n = self.raw.write(
34+
bytes(self._write_buf[:SOCK_WRITE_BLOCKSIZE]),
35+
)
3336
except io.BlockingIOError as e:
3437
n = e.characters_written
35-
del self._write_buf[:n]
38+
if n:
39+
del self._write_buf[:n]
3640

3741

3842
class StreamReader(io.BufferedReader):

cheroot/test/test_makefile.py

Lines changed: 36 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -51,3 +51,39 @@ def test_bytes_written():
5151
wfile = makefile.MakeFile(sock, 'w')
5252
wfile.write(b'bar')
5353
assert wfile.bytes_written == 3
54+
55+
56+
def test_flush_preserves_data_when_raw_write_returns_none():
57+
"""_flush_unlocked() must not treat None from raw.write() as a byte count.
58+
59+
io.RawIOBase.write() returns None when a non-blocking socket cannot accept
60+
data. del self._write_buf[:None] is equivalent to del self._write_buf[:]
61+
which silently clears the entire buffer, truncating the response without
62+
raising an exception.
63+
"""
64+
data = b'x' * 32768 # larger than SOCK_WRITE_BLOCKSIZE to stress the loop
65+
66+
sock = MockSocket()
67+
wfile = makefile.MakeFile(sock, 'w')
68+
wfile._write_buf.extend(data)
69+
70+
written = bytearray()
71+
call_count = 0
72+
73+
def mock_raw_write(b):
74+
nonlocal call_count
75+
call_count += 1
76+
if call_count == 1:
77+
return (
78+
None # simulates socket returning None on first blocked write
79+
)
80+
written.extend(b)
81+
return len(b)
82+
83+
wfile.raw.write = mock_raw_write
84+
wfile._flush_unlocked()
85+
86+
assert bytes(written) == data, (
87+
f'Expected {len(data)} bytes but only {len(written)} reached raw.write(): '
88+
'buffer was silently discarded when raw.write() returned None'
89+
)

0 commit comments

Comments
 (0)