test_sock_lowlevel.py

importsocket
importasyncio
importsys
importunittest

fromasyncioimportproactor_events
fromitertoolsimportcycle, islice
fromunittest.mockimportMock
fromtest.test_asyncioimportutilsastest_utils
fromtestimportsupport
fromtest.supportimportsocket_helper

ifsocket_helper.tcp_blackhole():
raiseunittest.SkipTest('Not relevant to ProactorEventLoop')


deftearDownModule():
asyncio.set_event_loop_policy(None)


classMyProto(asyncio.Protocol):
connected=None
done=None

def__init__(self, loop=None):
self.transport=None
self.state='INITIAL'
self.nbytes=0
ifloopisnotNone:
self.connected=loop.create_future()
self.done=loop.create_future()

def_assert_state(self, *expected):
ifself.statenotinexpected:
raiseAssertionError(f'state: {self.state!r}, expected: {expected!r}')

defconnection_made(self, transport):
self.transport=transport
self._assert_state('INITIAL')
self.state='CONNECTED'
ifself.connected:
self.connected.set_result(None)
transport.write(b'GET / HTTP/1.0\r\nHost: example.com\r\n\r\n')

defdata_received(self, data):
self._assert_state('CONNECTED')
self.nbytes+=len(data)

defeof_received(self):
self._assert_state('CONNECTED')
self.state='EOF'

defconnection_lost(self, exc):
self._assert_state('CONNECTED', 'EOF')
self.state='CLOSED'
ifself.done:
self.done.set_result(None)


classBaseSockTestsMixin:

defcreate_event_loop(self):
raiseNotImplementedError

defsetUp(self):
self.loop=self.create_event_loop()
self.set_event_loop(self.loop)
super().setUp()

deftearDown(self):
# just in case if we have transport close callbacks
ifnotself.loop.is_closed():
test_utils.run_briefly(self.loop)

self.doCleanups()
support.gc_collect()
super().tearDown()

def_basetest_sock_client_ops(self, httpd, sock):
ifnotisinstance(self.loop, proactor_events.BaseProactorEventLoop):
# in debug mode, socket operations must fail
# if the socket is not in blocking mode
self.loop.set_debug(True)
sock.setblocking(True)
withself.assertRaises(ValueError):
self.loop.run_until_complete(
self.loop.sock_connect(sock, httpd.address))
withself.assertRaises(ValueError):
self.loop.run_until_complete(
self.loop.sock_sendall(sock, b'GET / HTTP/1.0\r\n\r\n'))
withself.assertRaises(ValueError):
self.loop.run_until_complete(
self.loop.sock_recv(sock, 1024))
withself.assertRaises(ValueError):
self.loop.run_until_complete(
self.loop.sock_recv_into(sock, bytearray()))
withself.assertRaises(ValueError):
self.loop.run_until_complete(
self.loop.sock_accept(sock))

# test in non-blocking mode
sock.setblocking(False)
self.loop.run_until_complete(
self.loop.sock_connect(sock, httpd.address))
self.loop.run_until_complete(
self.loop.sock_sendall(sock, b'GET / HTTP/1.0\r\n\r\n'))
data=self.loop.run_until_complete(
self.loop.sock_recv(sock, 1024))
# consume data
self.loop.run_until_complete(
self.loop.sock_recv(sock, 1024))
sock.close()
self.assertTrue(data.startswith(b'HTTP/1.0 200 OK'))

def_basetest_sock_recv_into(self, httpd, sock):
# same as _basetest_sock_client_ops, but using sock_recv_into
sock.setblocking(False)
self.loop.run_until_complete(
self.loop.sock_connect(sock, httpd.address))
self.loop.run_until_complete(
self.loop.sock_sendall(sock, b'GET / HTTP/1.0\r\n\r\n'))
data=bytearray(1024)
withmemoryview(data) asbuf:
nbytes=self.loop.run_until_complete(
self.loop.sock_recv_into(sock, buf[:1024]))
# consume data
self.loop.run_until_complete(
self.loop.sock_recv_into(sock, buf[nbytes:]))
sock.close()
self.assertTrue(data.startswith(b'HTTP/1.0 200 OK'))

deftest_sock_client_ops(self):
withtest_utils.run_test_server() ashttpd:
sock=socket.socket()
self._basetest_sock_client_ops(httpd, sock)
sock=socket.socket()
self._basetest_sock_recv_into(httpd, sock)

asyncdef_basetest_sock_recv_racing(self, httpd, sock):
sock.setblocking(False)
awaitself.loop.sock_connect(sock, httpd.address)

task=asyncio.create_task(self.loop.sock_recv(sock, 1024))
awaitasyncio.sleep(0)
task.cancel()

asyncio.create_task(
self.loop.sock_sendall(sock, b'GET / HTTP/1.0\r\n\r\n'))
data=awaitself.loop.sock_recv(sock, 1024)
# consume data
awaitself.loop.sock_recv(sock, 1024)

self.assertTrue(data.startswith(b'HTTP/1.0 200 OK'))

asyncdef_basetest_sock_recv_into_racing(self, httpd, sock):
sock.setblocking(False)
awaitself.loop.sock_connect(sock, httpd.address)

data=bytearray(1024)
withmemoryview(data) asbuf:
task=asyncio.create_task(
self.loop.sock_recv_into(sock, buf[:1024]))
awaitasyncio.sleep(0)
task.cancel()

task=asyncio.create_task(
self.loop.sock_sendall(sock, b'GET / HTTP/1.0\r\n\r\n'))
nbytes=awaitself.loop.sock_recv_into(sock, buf[:1024])
# consume data
awaitself.loop.sock_recv_into(sock, buf[nbytes:])
self.assertTrue(data.startswith(b'HTTP/1.0 200 OK'))

awaittask

asyncdef_basetest_sock_send_racing(self, listener, sock):
listener.bind(('127.0.0.1', 0))
listener.listen(1)

# make connection
sock.setsockopt(socket.SOL_SOCKET, socket.SO_SNDBUF, 1024)
sock.setblocking(False)
task=asyncio.create_task(
self.loop.sock_connect(sock, listener.getsockname()))
awaitasyncio.sleep(0)
server=listener.accept()[0]
server.setblocking(False)

withserver:
awaittask

# fill the buffer until sending 5 chars would block
size=8192
whilesize>=4:
withself.assertRaises(BlockingIOError):
whileTrue:
sock.send(b' '*size)
size=int(size/2)

# cancel a blocked sock_sendall
task=asyncio.create_task(
self.loop.sock_sendall(sock, b'hello'))
awaitasyncio.sleep(0)
task.cancel()

# receive everything that is not a space
asyncdefrecv_all():
rv=b''
whileTrue:
buf=awaitself.loop.sock_recv(server, 8192)
ifnotbuf:
returnrv
rv+=buf.strip()
task=asyncio.create_task(recv_all())

# immediately make another sock_sendall call
awaitself.loop.sock_sendall(sock, b'world')
sock.shutdown(socket.SHUT_WR)
data=awaittask
# ProactorEventLoop could deliver hello, so endswith is necessary
self.assertTrue(data.endswith(b'world'))

# After the first connect attempt before the listener is ready,
# the socket needs time to "recover" to make the next connect call.
# On Linux, a second retry will do. On Windows, the waiting time is
# unpredictable; and on FreeBSD the socket may never come back
# because it's a loopback address. Here we'll just retry for a few
# times, and have to skip the test if it's not working. See also:
# https://stackoverflow.com/a/54437602/3316267
# https://lists.freebsd.org/pipermail/freebsd-current/2005-May/049876.html
asyncdef_basetest_sock_connect_racing(self, listener, sock):
listener.bind(('127.0.0.1', 0))
addr=listener.getsockname()
sock.setblocking(False)

task=asyncio.create_task(self.loop.sock_connect(sock, addr))
awaitasyncio.sleep(0)
task.cancel()

listener.listen(1)

skip_reason="Max retries reached"
foriinrange(128):
try:
awaitself.loop.sock_connect(sock, addr)
exceptConnectionRefusedErrorase:
skip_reason=e
exceptOSErrorase:
skip_reason=e

# Retry only for this error:
# [WinError 10022] An invalid argument was supplied
ifgetattr(e, 'winerror', 0) !=10022:
break
else:
# success
return

self.skipTest(skip_reason)

deftest_sock_client_racing(self):
withtest_utils.run_test_server() ashttpd:
sock=socket.socket()
withsock:
self.loop.run_until_complete(asyncio.wait_for(
self._basetest_sock_recv_racing(httpd, sock), 10))
sock=socket.socket()
withsock:
self.loop.run_until_complete(asyncio.wait_for(
self._basetest_sock_recv_into_racing(httpd, sock), 10))
listener=socket.socket()
sock=socket.socket()
withlistener, sock:
self.loop.run_until_complete(asyncio.wait_for(
self._basetest_sock_send_racing(listener, sock), 10))

deftest_sock_client_connect_racing(self):
listener=socket.socket()
sock=socket.socket()
withlistener, sock:
self.loop.run_until_complete(asyncio.wait_for(
self._basetest_sock_connect_racing(listener, sock), 10))

asyncdef_basetest_huge_content(self, address):
sock=socket.socket()
sock.setblocking(False)
DATA_SIZE=10_000_00

chunk=b'0123456789'* (DATA_SIZE//10)

awaitself.loop.sock_connect(sock, address)
awaitself.loop.sock_sendall(sock,
 (b'POST /loop HTTP/1.0\r\n'+
b'Content-Length: %d\r\n'%DATA_SIZE+
b'\r\n'))

task=asyncio.create_task(self.loop.sock_sendall(sock, chunk))

data=awaitself.loop.sock_recv(sock, DATA_SIZE)
# HTTP headers size is less than MTU,
# they are sent by the first packet always
self.assertTrue(data.startswith(b'HTTP/1.0 200 OK'))
whiledata.find(b'\r\n\r\n') ==-1:
data+=awaitself.loop.sock_recv(sock, DATA_SIZE)
# Strip headers
headers=data[:data.index(b'\r\n\r\n') +4]
data=data[len(headers):]

size=DATA_SIZE
checker=cycle(b'0123456789')

expected=bytes(islice(checker, len(data)))
self.assertEqual(data, expected)
size-=len(data)

whileTrue:
data=awaitself.loop.sock_recv(sock, DATA_SIZE)
ifnotdata:
break
expected=bytes(islice(checker, len(data)))
self.assertEqual(data, expected)
size-=len(data)
self.assertEqual(size, 0)

awaittask
sock.close()

deftest_huge_content(self):
withtest_utils.run_test_server() ashttpd:
self.loop.run_until_complete(
self._basetest_huge_content(httpd.address))

asyncdef_basetest_huge_content_recvinto(self, address):
sock=socket.socket()
sock.setblocking(False)
DATA_SIZE=10_000_00

chunk=b'0123456789'* (DATA_SIZE//10)

awaitself.loop.sock_connect(sock, address)
awaitself.loop.sock_sendall(sock,
 (b'POST /loop HTTP/1.0\r\n'+
b'Content-Length: %d\r\n'%DATA_SIZE+
b'\r\n'))

task=asyncio.create_task(self.loop.sock_sendall(sock, chunk))

array=bytearray(DATA_SIZE)
buf=memoryview(array)

nbytes=awaitself.loop.sock_recv_into(sock, buf)
data=bytes(buf[:nbytes])
# HTTP headers size is less than MTU,
# they are sent by the first packet always
self.assertTrue(data.startswith(b'HTTP/1.0 200 OK'))
whiledata.find(b'\r\n\r\n') ==-1:
nbytes=awaitself.loop.sock_recv_into(sock, buf)
data=bytes(buf[:nbytes])
# Strip headers
headers=data[:data.index(b'\r\n\r\n') +4]
data=data[len(headers):]

size=DATA_SIZE
checker=cycle(b'0123456789')

expected=bytes(islice(checker, len(data)))
self.assertEqual(data, expected)
size-=len(data)

whileTrue:
nbytes=awaitself.loop.sock_recv_into(sock, buf)
data=buf[:nbytes]
ifnotdata:
break
expected=bytes(islice(checker, len(data)))
self.assertEqual(data, expected)
size-=len(data)
self.assertEqual(size, 0)

awaittask
sock.close()

deftest_huge_content_recvinto(self):
withtest_utils.run_test_server() ashttpd:
self.loop.run_until_complete(
self._basetest_huge_content_recvinto(httpd.address))

asyncdef_basetest_datagram_recvfrom(self, server_address):
# Happy path, sock.sendto() returns immediately
data=b'\x01'*4096
withsocket.socket(socket.AF_INET, socket.SOCK_DGRAM) assock:
sock.setblocking(False)
awaitself.loop.sock_sendto(sock, data, server_address)
received_data, from_addr=awaitself.loop.sock_recvfrom(
sock, 4096)
self.assertEqual(received_data, data)
self.assertEqual(from_addr, server_address)

deftest_recvfrom(self):
withtest_utils.run_udp_echo_server() asserver_address:
self.loop.run_until_complete(
self._basetest_datagram_recvfrom(server_address))

asyncdef_basetest_datagram_recvfrom_into(self, server_address):
# Happy path, sock.sendto() returns immediately
withsocket.socket(socket.AF_INET, socket.SOCK_DGRAM) assock:
sock.setblocking(False)

buf=bytearray(4096)
data=b'\x01'*4096
awaitself.loop.sock_sendto(sock, data, server_address)
num_bytes, from_addr=awaitself.loop.sock_recvfrom_into(
sock, buf)
self.assertEqual(num_bytes, 4096)
self.assertEqual(buf, data)
self.assertEqual(from_addr, server_address)

buf=bytearray(8192)
awaitself.loop.sock_sendto(sock, data, server_address)
num_bytes, from_addr=awaitself.loop.sock_recvfrom_into(
sock, buf, 4096)
self.assertEqual(num_bytes, 4096)
self.assertEqual(buf[:4096], data[:4096])
self.assertEqual(from_addr, server_address)

deftest_recvfrom_into(self):
withtest_utils.run_udp_echo_server() asserver_address:
self.loop.run_until_complete(
self._basetest_datagram_recvfrom_into(server_address))

asyncdef_basetest_datagram_sendto_blocking(self, server_address):
# Sad path, sock.sendto() raises BlockingIOError
# This involves patching sock.sendto() to raise BlockingIOError but
# sendto() is not used by the proactor event loop
data=b'\x01'*4096
withsocket.socket(socket.AF_INET, socket.SOCK_DGRAM) assock:
sock.setblocking(False)
mock_sock=Mock(sock)
mock_sock.gettimeout=sock.gettimeout
mock_sock.sendto.configure_mock(side_effect=BlockingIOError)
mock_sock.fileno=sock.fileno
self.loop.call_soon(
lambda: setattr(mock_sock, 'sendto', sock.sendto)
 )
awaitself.loop.sock_sendto(mock_sock, data, server_address)

received_data, from_addr=awaitself.loop.sock_recvfrom(
sock, 4096)
self.assertEqual(received_data, data)
self.assertEqual(from_addr, server_address)

deftest_sendto_blocking(self):
ifsys.platform=='win32':
ifisinstance(self.loop, asyncio.ProactorEventLoop):
raiseunittest.SkipTest('Not relevant to ProactorEventLoop')

withtest_utils.run_udp_echo_server() asserver_address:
self.loop.run_until_complete(
self._basetest_datagram_sendto_blocking(server_address))

@socket_helper.skip_unless_bind_unix_socket
deftest_unix_sock_client_ops(self):
withtest_utils.run_test_unix_server() ashttpd:
sock=socket.socket(socket.AF_UNIX)
self._basetest_sock_client_ops(httpd, sock)
sock=socket.socket(socket.AF_UNIX)
self._basetest_sock_recv_into(httpd, sock)

deftest_sock_client_fail(self):
# Make sure that we will get an unused port
address=None
try:
s=socket.socket()
s.bind(('127.0.0.1', 0))
address=s.getsockname()
finally:
s.close()

sock=socket.socket()
sock.setblocking(False)
withself.assertRaises(ConnectionRefusedError):
self.loop.run_until_complete(
self.loop.sock_connect(sock, address))
sock.close()

deftest_sock_accept(self):
listener=socket.socket()
listener.setblocking(False)
listener.bind(('127.0.0.1', 0))
listener.listen(1)
client=socket.socket()
client.connect(listener.getsockname())

f=self.loop.sock_accept(listener)
conn, addr=self.loop.run_until_complete(f)
self.assertEqual(conn.gettimeout(), 0)
self.assertEqual(addr, client.getsockname())
self.assertEqual(client.getpeername(), listener.getsockname())
client.close()
conn.close()
listener.close()

deftest_cancel_sock_accept(self):
listener=socket.socket()
listener.setblocking(False)
listener.bind(('127.0.0.1', 0))
listener.listen(1)
sockaddr=listener.getsockname()
f=asyncio.wait_for(self.loop.sock_accept(listener), 0.1)
withself.assertRaises(asyncio.TimeoutError):
self.loop.run_until_complete(f)

listener.close()
client=socket.socket()
client.setblocking(False)
f=self.loop.sock_connect(client, sockaddr)
withself.assertRaises(ConnectionRefusedError):
self.loop.run_until_complete(f)

client.close()

deftest_create_connection_sock(self):
withtest_utils.run_test_server() ashttpd:
sock=None
infos=self.loop.run_until_complete(
self.loop.getaddrinfo(
*httpd.address, type=socket.SOCK_STREAM))
forfamily, type, proto, cname, addressininfos:
try:
sock=socket.socket(family=family, type=type, proto=proto)
sock.setblocking(False)
self.loop.run_until_complete(
self.loop.sock_connect(sock, address))
exceptBaseException:
pass
else:
break
else:
self.fail('Can not create socket.')

f=self.loop.create_connection(
lambda: MyProto(loop=self.loop), sock=sock)
tr, pr=self.loop.run_until_complete(f)
self.assertIsInstance(tr, asyncio.Transport)
self.assertIsInstance(pr, asyncio.Protocol)
self.loop.run_until_complete(pr.done)
self.assertGreater(pr.nbytes, 0)
tr.close()


ifsys.platform=='win32':

classSelectEventLoopTests(BaseSockTestsMixin,
test_utils.TestCase):

defcreate_event_loop(self):
returnasyncio.SelectorEventLoop()


classProactorEventLoopTests(BaseSockTestsMixin,
test_utils.TestCase):

defcreate_event_loop(self):
returnasyncio.ProactorEventLoop()


asyncdef_basetest_datagram_send_to_non_listening_address(self,
recvfrom):
# see:
# https://github.com/python/cpython/issues/91227
# https://github.com/python/cpython/issues/88906
# https://bugs.python.org/issue47071
# https://bugs.python.org/issue44743
# The Proactor event loop would fail to receive datagram messages
# after sending a message to an address that wasn't listening.

defcreate_socket():
sock=socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
sock.setblocking(False)
sock.bind(('127.0.0.1', 0))
returnsock

socket_1=create_socket()
addr_1=socket_1.getsockname()

socket_2=create_socket()
addr_2=socket_2.getsockname()

# creating and immediately closing this to try to get an address
# that is not listening
socket_3=create_socket()
addr_3=socket_3.getsockname()
socket_3.shutdown(socket.SHUT_RDWR)
socket_3.close()

socket_1_recv_task=self.loop.create_task(recvfrom(socket_1))
socket_2_recv_task=self.loop.create_task(recvfrom(socket_2))
awaitasyncio.sleep(0)

awaitself.loop.sock_sendto(socket_1, b'a', addr_2)
self.assertEqual(awaitsocket_2_recv_task, b'a')

awaitself.loop.sock_sendto(socket_2, b'b', addr_1)
self.assertEqual(awaitsocket_1_recv_task, b'b')
socket_1_recv_task=self.loop.create_task(recvfrom(socket_1))
awaitasyncio.sleep(0)

# this should send to an address that isn't listening
awaitself.loop.sock_sendto(socket_1, b'c', addr_3)
self.assertEqual(awaitsocket_1_recv_task, b'')
socket_1_recv_task=self.loop.create_task(recvfrom(socket_1))
awaitasyncio.sleep(0)

# socket 1 should still be able to receive messages after sending
# to an address that wasn't listening
socket_2.sendto(b'd', addr_1)
self.assertEqual(awaitsocket_1_recv_task, b'd')

socket_1.shutdown(socket.SHUT_RDWR)
socket_1.close()
socket_2.shutdown(socket.SHUT_RDWR)
socket_2.close()


deftest_datagram_send_to_non_listening_address_recvfrom(self):
asyncdefrecvfrom(socket):
data, _=awaitself.loop.sock_recvfrom(socket, 4096)
returndata

self.loop.run_until_complete(
self._basetest_datagram_send_to_non_listening_address(
recvfrom))


deftest_datagram_send_to_non_listening_address_recvfrom_into(self):
asyncdefrecvfrom_into(socket):
buf=bytearray(4096)
length, _=awaitself.loop.sock_recvfrom_into(socket, buf,
4096)
returnbuf[:length]

self.loop.run_until_complete(
self._basetest_datagram_send_to_non_listening_address(
recvfrom_into))

else:
importselectors

ifhasattr(selectors, 'KqueueSelector'):
classKqueueEventLoopTests(BaseSockTestsMixin,
test_utils.TestCase):

defcreate_event_loop(self):
returnasyncio.SelectorEventLoop(
selectors.KqueueSelector())

ifhasattr(selectors, 'EpollSelector'):
classEPollEventLoopTests(BaseSockTestsMixin,
test_utils.TestCase):

defcreate_event_loop(self):
returnasyncio.SelectorEventLoop(selectors.EpollSelector())

ifhasattr(selectors, 'PollSelector'):
classPollEventLoopTests(BaseSockTestsMixin,
test_utils.TestCase):

defcreate_event_loop(self):
returnasyncio.SelectorEventLoop(selectors.PollSelector())

# Should always exist.
classSelectEventLoopTests(BaseSockTestsMixin,
test_utils.TestCase):

defcreate_event_loop(self):
returnasyncio.SelectorEventLoop(selectors.SelectSelector())


if__name__=='__main__':
unittest.main()