Skip to content

Commit 69de9ea

Browse files
julianz-julianz-
committed
improved test for moving buffer
1 parent d95dcf5 commit 69de9ea

File tree

1 file changed

+215
-80
lines changed

1 file changed

+215
-80
lines changed

‎tests/test_ssl.py‎

Lines changed: 215 additions & 80 deletions
Original file line numberDiff line numberDiff line change
@@ -3131,6 +3131,195 @@ def test_wantWriteError(self) -> None:
31313131

31323132
# XXX want_read
31333133

3134+
def _fill_client_buffer(self, client_socket: socket) -> None:
3135+
"""
3136+
Attempts to fill the client's raw send buffer until
3137+
EWOULDBLOCK is hit.
3138+
"""
3139+
print("--- Phase 1: Filling client socket buffer ---")
3140+
for msg in [b"x" * 65536, b"x" * 16]:
3141+
for _ in range(
3142+
1024 * 1024 * 64
3143+
): # Large loop count to ensure buffer fill
3144+
try:
3145+
client_socket.send(msg)
3146+
time.sleep(0.01)
3147+
except OSError as e:
3148+
if e.errno == EWOULDBLOCK:
3149+
print("Client socket buffer filled (EWOULDBLOCK hit).")
3150+
return # Buffer successfully filled, exit function
3151+
raise # pragma: no cover # Re-raise unexpected OSErrors
3152+
else: # If the inner loop completes without hitting EWOULDBLOCK
3153+
pytest.fail(
3154+
"Failed to fill socket buffer, cannot test bad write error"
3155+
) # pragma: no cover
3156+
3157+
def _attempt_want_write_error(self, client: Connection) -> int:
3158+
"""
3159+
Attempts to send application data over SSL to trigger WantWriteError.
3160+
Returns (True, successful_size) if triggered,
3161+
otherwise calls pytest.fail.
3162+
"""
3163+
print("--- Phase 2: Attempting to trigger WantWriteError ---")
3164+
test_sizes = [
3165+
64,
3166+
128,
3167+
256,
3168+
512,
3169+
1024,
3170+
2048,
3171+
4096,
3172+
8192,
3173+
16384,
3174+
32768,
3175+
65536,
3176+
]
3177+
initial_want_write_triggered = False
3178+
successful_size = -1
3179+
3180+
for size in test_sizes:
3181+
msg2 = b"Y" * size
3182+
try:
3183+
client.send(msg2)
3184+
# Continue loop to try larger sizes until an error is hit
3185+
except SSL.WantWriteError:
3186+
print(
3187+
f"Got WantWriteError with message size {size} "
3188+
"(this is what we want)."
3189+
)
3190+
initial_want_write_triggered = True
3191+
successful_size = size * 2 # double it to be really sure
3192+
break # Exit loop as desired error was triggered
3193+
3194+
if not initial_want_write_triggered:
3195+
pytest.fail(
3196+
"Could not induce WantWriteError with any message size."
3197+
) # pragma: no cover
3198+
3199+
return successful_size
3200+
3201+
def _drain_server_buffers(
3202+
self, server: Connection, server_socket: socket
3203+
) -> None:
3204+
"""Reads from server SSL and raw sockets to drain any pending data."""
3205+
print("--- Phase 3: Draining server buffers ---")
3206+
total_read = 0
3207+
read_chunks = []
3208+
3209+
try:
3210+
# First, try to read any SSL data that might be available
3211+
try:
3212+
server.recv(65536)
3213+
except (SSL.WantReadError, SSL.Error) as ssl_error:
3214+
print(f"No SSL data available or SSL error: {ssl_error}")
3215+
3216+
# Now read raw data from the underlying server socket to
3217+
# drain buffer
3218+
server_socket.setblocking(False) # Ensure non-blocking
3219+
consecutive_empty_reads = 0
3220+
3221+
while (
3222+
total_read < 1024 * 1024
3223+
): # Read up to 1MB or until no more data
3224+
try:
3225+
data = server_socket.recv(65536) # Read raw data
3226+
read_chunks.append(data)
3227+
total_read += len(data)
3228+
print(
3229+
f"Read {len(data)} bytes of raw data from "
3230+
f"server socket (total: {total_read})."
3231+
)
3232+
time.sleep(0.01) # Small delay between reads
3233+
3234+
except OSError as e:
3235+
if e.errno == EWOULDBLOCK:
3236+
consecutive_empty_reads += 1
3237+
if consecutive_empty_reads >= 5:
3238+
print(
3239+
"No more raw data available from server "
3240+
"socket after retries."
3241+
)
3242+
break
3243+
print(
3244+
"No data available, waiting... "
3245+
f"(attempt {consecutive_empty_reads})."
3246+
)
3247+
time.sleep(0.1) # Wait longer when buffer is empty
3248+
continue
3249+
else: # pragma: no cover
3250+
print(f"OSError while reading from server socket: {e}")
3251+
break
3252+
3253+
print(
3254+
f"Finished reading from server. Total bytes read: {total_read}"
3255+
)
3256+
print("Allowing network buffers to settle...")
3257+
time.sleep(0.1)
3258+
3259+
except Exception as read_exception: # pragma: no cover
3260+
print(f"Exception while reading from server: {read_exception}")
3261+
3262+
def _perform_moving_buffer_test(
3263+
self, client: Connection, successful_size: int
3264+
) -> bool:
3265+
"""
3266+
Attempts a retry write with a moving buffer and checks for
3267+
'bad write retry' error.
3268+
Returns True if 'bad write retry' occurs, False otherwise.
3269+
"""
3270+
print("--- Phase 4: Performing moving buffer retry test ---")
3271+
assert successful_size > -1, (
3272+
"successful_size must be greater than -1 for a WantWriteError "
3273+
"to be triggered"
3274+
)
3275+
msg3 = b"Z" * successful_size
3276+
3277+
print(
3278+
"Attempting retry with different buffer "
3279+
f"(same size {successful_size})."
3280+
)
3281+
3282+
try:
3283+
client.send(msg3)
3284+
print(f"Retry succeeded with {successful_size} bytes written.")
3285+
return False # Retry succeeded
3286+
except SSL.Error as e:
3287+
reason = get_ssl_error_reason(e)
3288+
if reason == "bad write retry":
3289+
print(f"Got expected SSL error: {e} ({reason}).")
3290+
return True # Bad write retry
3291+
else:
3292+
pytest.fail(
3293+
f"Retry failed with unexpected SSL error: {e} ({reason})."
3294+
) # pragma: no cover
3295+
# If any other exception occurs, it will propagate up
3296+
3297+
def _shutdown_connections(
3298+
self,
3299+
client: Connection,
3300+
server: Connection,
3301+
client_socket: socket,
3302+
server_socket: socket,
3303+
) -> None:
3304+
"""Helper to safely shut down SSL connections and close sockets."""
3305+
print("--- Cleanup: Shutting down connections ---")
3306+
try:
3307+
if client:
3308+
client.shutdown()
3309+
except Exception as e:
3310+
print(f"Error during client SSL shutdown: {e}")
3311+
try:
3312+
if server:
3313+
server.shutdown()
3314+
except Exception as e:
3315+
print(f"Error during server SSL shutdown: {e}")
3316+
finally:
3317+
if client_socket:
3318+
client_socket.close()
3319+
if server_socket:
3320+
server_socket.close()
3321+
print("Connections closed.")
3322+
31343323
def _badwriteretry(self, mode: int) -> bool:
31353324
"""
31363325
Tries to force a "bad write retry" error over an SSL connection
@@ -3141,78 +3330,26 @@ def _badwriteretry(self, mode: int) -> bool:
31413330
create_ssl_nonblocking_connection(mode)
31423331
)
31433332
result = False # Default return value
3144-
written = 0
31453333

31463334
try:
3147-
# Fill up the client's raw send buffer so the SSL connection
3148-
# won't be able to write anything. Start by sending larger chunks
3149-
# and continue by writing smaller chunks so we can be sure we
3150-
# completely fill the buffer.
3151-
for msg in [b"x" * 65536, b"x" * 16]:
3152-
for i in range(1024 * 1024 * 64):
3153-
try:
3154-
written = client_socket.send(msg)
3155-
print(f"Sent {written} bytes to fill buffer")
3156-
except OSError as e:
3157-
if e.errno == EWOULDBLOCK:
3158-
break
3159-
raise
3160-
else:
3161-
pytest.fail(
3162-
"Failed to fill socket buffer, cannot test \
3163-
bad write error"
3164-
)
3165-
3166-
# Now, attempt to send application data over the established
3167-
# SSL connection. Since the underlying raw socket's buffer is full,
3168-
# this should cause a WantWriteError.
3169-
msg2 = b"Y" * 65536
3170-
3171-
try:
3172-
written = client.send(msg2)
3173-
except SSL.WantWriteError:
3174-
try:
3175-
# After a WantWriteError if the connection has partially
3176-
# written the last buffer it will expect a retry write.
3177-
# This next write should fail but for two different reasons
3178-
# depending on whether SSL_MODE_ACCEPT_MOVING_WRITE_BUFFER
3179-
# was set
3180-
msg3 = b"Z" * 65536
3181-
written = client.send(msg3)
3182-
pytest.fail("Retry succeeded unexpectedly")
3183-
except SSL.Error as e:
3184-
reason = get_ssl_error_reason(e)
3185-
if reason == "bad write retry":
3186-
# Got SSL error on retry (expected if not using \
3187-
# SSL_MODE_ACCEPT_MOVING_WRITE_BUFFER)
3188-
result = True
3189-
else:
3190-
# when using SSL_MODE_ACCEPT_MOVING_WRITE_BUFFER
3191-
# we expect this to fail for a WantWriteError
3192-
result = False
3193-
3194-
except SSL.Error as e:
3195-
reason = get_ssl_error_reason(e)
3196-
pytest.fail(f"Got unexpected SSL error on retry: {e} {reason}")
3197-
except Exception as e:
3198-
pytest.fail(f"Unexpected exception during send: {e}")
3199-
3335+
# --- Main Test Flow ---
3336+
self._fill_client_buffer(client_socket)
3337+
3338+
successful_size = self._attempt_want_write_error(client)
3339+
3340+
# if WantWriteError was not triggered the test fails in
3341+
# _attempt_want_write_error().
3342+
# proceed with draining and retry
3343+
self._drain_server_buffers(server, server_socket)
3344+
result = self._perform_moving_buffer_test(client, successful_size)
3345+
except Exception as e: # pragma: no cover
3346+
pytest.fail(f"Unexpected exception during test: {e}.")
32003347
finally:
3201-
# Cleanup: shut down SSL connections and close raw sockets
3202-
try:
3203-
if client:
3204-
client.shutdown()
3205-
if server:
3206-
server.shutdown()
3207-
except Exception as e:
3208-
print(f"Error during SSL shutdown: {e}")
3209-
finally:
3210-
if client_socket:
3211-
client_socket.close()
3212-
if server_socket:
3213-
server_socket.close()
3214-
3215-
return result # Return the result after cleanup
3348+
self._shutdown_connections(
3349+
client, server, client_socket, server_socket
3350+
)
3351+
3352+
return result
32163353

32173354
def test_moving_write_buffer_should_pass(self) -> None:
32183355
"""
@@ -3227,11 +3364,10 @@ def test_moving_write_buffer_should_pass(self) -> None:
32273364
)
32283365
result = self._badwriteretry(mode)
32293366

3230-
if result:
3231-
pytest.fail(
3232-
"Using SSL_MODE_ACCEPT_MOVING_WRITE_BUFFER failed to \
3233-
prevent bad write retry"
3234-
)
3367+
assert result is False, (
3368+
"Using SSL_MODE_ACCEPT_MOVING_WRITE_BUFFER failed to prevent bad "
3369+
"write retry. A bad write retry occurred when it should not have."
3370+
)
32353371

32363372
def test_moving_write_buffer_should_fail(self) -> None:
32373373
"""
@@ -3244,12 +3380,11 @@ def test_moving_write_buffer_should_fail(self) -> None:
32443380
mode = 0
32453381
result = self._badwriteretry(mode)
32463382

3247-
if not result:
3248-
pytest.fail(
3249-
"Use of a moving buffer without \
3250-
SSL_MODE_ACCEPT_MOVING_WRITE_BUFFER should trigger \
3251-
a bad write retry error"
3252-
)
3383+
assert result is True, (
3384+
"Use of a moving buffer without "
3385+
"SSL_MODE_ACCEPT_MOVING_WRITE_BUFFER should trigger "
3386+
"a bad write retry error."
3387+
)
32533388

32543389
def test_get_finished_before_connect(self) -> None:
32553390
"""

0 commit comments

Comments
 (0)