9754258d46
In _PyAsyncIoStreamProtocol.write_loop(), memoryview objects pointing to C++ message memory were passed directly to transport.write(). Since transport.write() is non-blocking and only queues data for later transmission, the memoryview could reference freed memory after fulfill() was called. This caused message corruption when pipelining RPC calls with payloads larger than ~4000 bytes, as the C++ message memory would be freed before asyncio had a chance to transmit the data. The fix copies the data to Python bytes objects before passing to transport.write(), ensuring the data remains valid until asyncio transmits it. Includes regression test that verifies large payload integrity with both sequential and pipelined RPC calls.
113 lines
4.1 KiB
Python
113 lines
4.1 KiB
Python
"""
|
|
Regression test for use-after-free bug in async write with large payloads.
|
|
|
|
This test reproduces a bug where pipelining multiple RPC calls with large data
|
|
payloads (>~4000 bytes) causes memory corruption. The root cause was that
|
|
_PyAsyncIoStreamProtocol.write_loop() passed a memoryview pointing to C++
|
|
memory to transport.write(), then called fulfill() which freed the C++ memory.
|
|
Since transport.write() is non-blocking and buffers data asynchronously,
|
|
the data could be corrupted before being sent.
|
|
|
|
The fix is to copy the data to Python bytes before passing to transport.write().
|
|
|
|
See: https://github.com/capnproto/pycapnp/pull/392
|
|
"""
|
|
|
|
import pytest
|
|
import socket
|
|
|
|
import capnp
|
|
import test_capability_capnp
|
|
|
|
|
|
@pytest.fixture(autouse=True)
|
|
async def kj_loop():
|
|
async with capnp.kj_loop():
|
|
yield
|
|
|
|
|
|
class LargeResponseServer(test_capability_capnp.TestInterface.Server):
|
|
"""
|
|
Server that returns large text responses to trigger the use-after-free bug.
|
|
|
|
The bug manifests when response messages are >~4000 bytes.
|
|
"""
|
|
|
|
async def foo(self, i, j, **kwargs):
|
|
# Generate a large response string based on input
|
|
# The size is controlled by the input parameter 'i'
|
|
size = i
|
|
# Create a deterministic pattern that can be verified
|
|
pattern = "".join(chr(65 + (k % 26)) for k in range(size))
|
|
return pattern
|
|
|
|
|
|
async def test_large_response_sequential():
|
|
"""
|
|
Test that large RPC responses are received correctly when sent sequentially.
|
|
|
|
Tests various payload sizes including those >4000 bytes which trigger the bug.
|
|
"""
|
|
read_sock, write_sock = socket.socketpair()
|
|
read_stream = await capnp.AsyncIoStream.create_connection(sock=read_sock)
|
|
write_stream = await capnp.AsyncIoStream.create_connection(sock=write_sock)
|
|
|
|
_ = capnp.TwoPartyServer(write_stream, bootstrap=LargeResponseServer())
|
|
client = capnp.TwoPartyClient(read_stream)
|
|
cap = client.bootstrap().cast_as(test_capability_capnp.TestInterface)
|
|
|
|
# Test various sizes, including sizes that trigger the bug (>~4000 bytes)
|
|
test_sizes = [100, 1000, 4000, 5000, 8000]
|
|
|
|
for size in test_sizes:
|
|
response = await cap.foo(i=size, j=False)
|
|
|
|
# Verify the response has the correct length
|
|
assert (
|
|
len(response.x) == size
|
|
), f"Size mismatch for {size}: expected {size}, got {len(response.x)}"
|
|
|
|
# Verify the pattern is correct (not corrupted)
|
|
expected = "".join(chr(65 + (k % 26)) for k in range(size))
|
|
assert response.x == expected, (
|
|
f"Data corruption detected for {size} bytes payload! "
|
|
f"First 50 chars: expected '{expected[:50]}', got '{response.x[:50]}'"
|
|
)
|
|
|
|
|
|
async def test_large_response_pipelined():
|
|
"""
|
|
Test that pipelining multiple RPC calls with large responses works correctly.
|
|
|
|
This is a more aggressive test that sends multiple requests without awaiting,
|
|
then collects all results. This pattern is more likely to trigger the
|
|
use-after-free bug because multiple messages are queued in the write buffer.
|
|
"""
|
|
read_sock, write_sock = socket.socketpair()
|
|
read_stream = await capnp.AsyncIoStream.create_connection(sock=read_sock)
|
|
write_stream = await capnp.AsyncIoStream.create_connection(sock=write_sock)
|
|
|
|
_ = capnp.TwoPartyServer(write_stream, bootstrap=LargeResponseServer())
|
|
client = capnp.TwoPartyClient(read_stream)
|
|
cap = client.bootstrap().cast_as(test_capability_capnp.TestInterface)
|
|
|
|
# Test sizes that trigger the bug - send 3 pipelined requests
|
|
test_sizes = [5000, 6000, 8000]
|
|
|
|
# Send all requests without awaiting (pipelining)
|
|
promises = []
|
|
for size in test_sizes:
|
|
promise = cap.foo(i=size, j=False)
|
|
promises.append((size, promise))
|
|
|
|
# Now await all responses and verify
|
|
for size, promise in promises:
|
|
response = await promise
|
|
|
|
assert len(response.x) == size, f"Size mismatch for {size}"
|
|
|
|
expected = "".join(chr(65 + (k % 26)) for k in range(size))
|
|
assert (
|
|
response.x == expected
|
|
), f"Data corruption detected for {size} bytes payload!"
|