import pytest
|
import sys
|
import socket
|
|
import attr
|
|
import trio
|
from trio.socket import AF_INET, AF_INET6, SOCK_STREAM, IPPROTO_TCP
|
from trio._highlevel_open_tcp_stream import (
|
reorder_for_rfc_6555_section_5_4,
|
close_all,
|
open_tcp_stream,
|
format_host_port,
|
)
|
|
if sys.version_info < (3, 11):
|
from exceptiongroup import BaseExceptionGroup
|
|
|
def test_close_all():
|
class CloseMe:
|
closed = False
|
|
def close(self):
|
self.closed = True
|
|
class CloseKiller:
|
def close(self):
|
raise OSError
|
|
c = CloseMe()
|
with close_all() as to_close:
|
to_close.add(c)
|
assert c.closed
|
|
c = CloseMe()
|
with pytest.raises(RuntimeError):
|
with close_all() as to_close:
|
to_close.add(c)
|
raise RuntimeError
|
assert c.closed
|
|
c = CloseMe()
|
with pytest.raises(OSError):
|
with close_all() as to_close:
|
to_close.add(CloseKiller())
|
to_close.add(c)
|
assert c.closed
|
|
|
def test_reorder_for_rfc_6555_section_5_4():
|
def fake4(i):
|
return (
|
AF_INET,
|
SOCK_STREAM,
|
IPPROTO_TCP,
|
"",
|
("10.0.0.{}".format(i), 80),
|
)
|
|
def fake6(i):
|
return (AF_INET6, SOCK_STREAM, IPPROTO_TCP, "", ("::{}".format(i), 80))
|
|
for fake in fake4, fake6:
|
# No effect on homogeneous lists
|
targets = [fake(0), fake(1), fake(2)]
|
reorder_for_rfc_6555_section_5_4(targets)
|
assert targets == [fake(0), fake(1), fake(2)]
|
|
# Single item lists also OK
|
targets = [fake(0)]
|
reorder_for_rfc_6555_section_5_4(targets)
|
assert targets == [fake(0)]
|
|
# If the list starts out with different families in positions 0 and 1,
|
# then it's left alone
|
orig = [fake4(0), fake6(0), fake4(1), fake6(1)]
|
targets = list(orig)
|
reorder_for_rfc_6555_section_5_4(targets)
|
assert targets == orig
|
|
# If not, it's reordered
|
targets = [fake4(0), fake4(1), fake4(2), fake6(0), fake6(1)]
|
reorder_for_rfc_6555_section_5_4(targets)
|
assert targets == [fake4(0), fake6(0), fake4(1), fake4(2), fake6(1)]
|
|
|
def test_format_host_port():
|
assert format_host_port("127.0.0.1", 80) == "127.0.0.1:80"
|
assert format_host_port(b"127.0.0.1", 80) == "127.0.0.1:80"
|
assert format_host_port("example.com", 443) == "example.com:443"
|
assert format_host_port(b"example.com", 443) == "example.com:443"
|
assert format_host_port("::1", "http") == "[::1]:http"
|
assert format_host_port(b"::1", "http") == "[::1]:http"
|
|
|
# Make sure we can connect to localhost using real kernel sockets
|
async def test_open_tcp_stream_real_socket_smoketest():
|
listen_sock = trio.socket.socket()
|
await listen_sock.bind(("127.0.0.1", 0))
|
_, listen_port = listen_sock.getsockname()
|
listen_sock.listen(1)
|
client_stream = await open_tcp_stream("127.0.0.1", listen_port)
|
server_sock, _ = await listen_sock.accept()
|
await client_stream.send_all(b"x")
|
assert await server_sock.recv(1) == b"x"
|
await client_stream.aclose()
|
server_sock.close()
|
|
listen_sock.close()
|
|
|
async def test_open_tcp_stream_input_validation():
|
with pytest.raises(ValueError):
|
await open_tcp_stream(None, 80)
|
with pytest.raises(TypeError):
|
await open_tcp_stream("127.0.0.1", b"80")
|
|
|
def can_bind_127_0_0_2():
|
with socket.socket() as s:
|
try:
|
s.bind(("127.0.0.2", 0))
|
except OSError:
|
return False
|
return s.getsockname()[0] == "127.0.0.2"
|
|
|
async def test_local_address_real():
|
with trio.socket.socket() as listener:
|
await listener.bind(("127.0.0.1", 0))
|
listener.listen()
|
|
# It's hard to test local_address properly, because you need multiple
|
# local addresses that you can bind to. Fortunately, on most Linux
|
# systems, you can bind to any 127.*.*.* address, and they all go
|
# through the loopback interface. So we can use a non-standard
|
# loopback address. On other systems, the only address we know for
|
# certain we have is 127.0.0.1, so we can't really test local_address=
|
# properly -- passing local_address=127.0.0.1 is indistinguishable
|
# from not passing local_address= at all. But, we can still do a smoke
|
# test to make sure the local_address= code doesn't crash.
|
if can_bind_127_0_0_2():
|
local_address = "127.0.0.2"
|
else:
|
local_address = "127.0.0.1"
|
|
async with await open_tcp_stream(
|
*listener.getsockname(), local_address=local_address
|
) as client_stream:
|
assert client_stream.socket.getsockname()[0] == local_address
|
if hasattr(trio.socket, "IP_BIND_ADDRESS_NO_PORT"):
|
assert client_stream.socket.getsockopt(
|
trio.socket.IPPROTO_IP, trio.socket.IP_BIND_ADDRESS_NO_PORT
|
)
|
|
server_sock, remote_addr = await listener.accept()
|
await client_stream.aclose()
|
server_sock.close()
|
assert remote_addr[0] == local_address
|
|
# Trying to connect to an ipv4 address with the ipv6 wildcard
|
# local_address should fail
|
with pytest.raises(OSError):
|
await open_tcp_stream(*listener.getsockname(), local_address="::")
|
|
# But the ipv4 wildcard address should work
|
async with await open_tcp_stream(
|
*listener.getsockname(), local_address="0.0.0.0"
|
) as client_stream:
|
server_sock, remote_addr = await listener.accept()
|
server_sock.close()
|
assert remote_addr == client_stream.socket.getsockname()
|
|
|
# Now, thorough tests using fake sockets
|
|
|
@attr.s(eq=False)
|
class FakeSocket(trio.socket.SocketType):
|
scenario = attr.ib()
|
family = attr.ib()
|
type = attr.ib()
|
proto = attr.ib()
|
|
ip = attr.ib(default=None)
|
port = attr.ib(default=None)
|
succeeded = attr.ib(default=False)
|
closed = attr.ib(default=False)
|
failing = attr.ib(default=False)
|
|
async def connect(self, sockaddr):
|
self.ip = sockaddr[0]
|
self.port = sockaddr[1]
|
assert self.ip not in self.scenario.sockets
|
self.scenario.sockets[self.ip] = self
|
self.scenario.connect_times[self.ip] = trio.current_time()
|
delay, result = self.scenario.ip_dict[self.ip]
|
await trio.sleep(delay)
|
if result == "error":
|
raise OSError("sorry")
|
if result == "postconnect_fail":
|
self.failing = True
|
self.succeeded = True
|
|
def close(self):
|
self.closed = True
|
|
# called when SocketStream is constructed
|
def setsockopt(self, *args, **kwargs):
|
if self.failing:
|
# raise something that isn't OSError as SocketStream
|
# ignores those
|
raise KeyboardInterrupt
|
|
|
class Scenario(trio.abc.SocketFactory, trio.abc.HostnameResolver):
|
def __init__(self, port, ip_list, supported_families):
|
# ip_list have to be unique
|
ip_order = [ip for (ip, _, _) in ip_list]
|
assert len(set(ip_order)) == len(ip_list)
|
ip_dict = {}
|
for ip, delay, result in ip_list:
|
assert 0 <= delay
|
assert result in ["error", "success", "postconnect_fail"]
|
ip_dict[ip] = (delay, result)
|
|
self.port = port
|
self.ip_order = ip_order
|
self.ip_dict = ip_dict
|
self.supported_families = supported_families
|
self.socket_count = 0
|
self.sockets = {}
|
self.connect_times = {}
|
|
def socket(self, family, type, proto):
|
if family not in self.supported_families:
|
raise OSError("pretending not to support this family")
|
self.socket_count += 1
|
return FakeSocket(self, family, type, proto)
|
|
def _ip_to_gai_entry(self, ip):
|
if ":" in ip:
|
family = trio.socket.AF_INET6
|
sockaddr = (ip, self.port, 0, 0)
|
else:
|
family = trio.socket.AF_INET
|
sockaddr = (ip, self.port)
|
return (family, SOCK_STREAM, IPPROTO_TCP, "", sockaddr)
|
|
async def getaddrinfo(self, host, port, family, type, proto, flags):
|
assert host == b"test.example.com"
|
assert port == self.port
|
assert family == trio.socket.AF_UNSPEC
|
assert type == trio.socket.SOCK_STREAM
|
assert proto == 0
|
assert flags == 0
|
return [self._ip_to_gai_entry(ip) for ip in self.ip_order]
|
|
async def getnameinfo(self, sockaddr, flags): # pragma: no cover
|
raise NotImplementedError
|
|
def check(self, succeeded):
|
# sockets only go into self.sockets when connect is called; make sure
|
# all the sockets that were created did in fact go in there.
|
assert self.socket_count == len(self.sockets)
|
|
for ip, socket in self.sockets.items():
|
assert ip in self.ip_dict
|
if socket is not succeeded:
|
assert socket.closed
|
assert socket.port == self.port
|
|
|
async def run_scenario(
|
# The port to connect to
|
port,
|
# A list of
|
# (ip, delay, result)
|
# tuples, where delay is in seconds and result is "success" or "error"
|
# The ip's will be returned from getaddrinfo in this order, and then
|
# connect() calls to them will have the given result.
|
ip_list,
|
*,
|
# If False, AF_INET4/6 sockets error out on creation, before connect is
|
# even called.
|
ipv4_supported=True,
|
ipv6_supported=True,
|
# Normally, we return (winning_sock, scenario object)
|
# If this is True, we require there to be an exception, and return
|
# (exception, scenario object)
|
expect_error=(),
|
**kwargs,
|
):
|
supported_families = set()
|
if ipv4_supported:
|
supported_families.add(trio.socket.AF_INET)
|
if ipv6_supported:
|
supported_families.add(trio.socket.AF_INET6)
|
scenario = Scenario(port, ip_list, supported_families)
|
trio.socket.set_custom_hostname_resolver(scenario)
|
trio.socket.set_custom_socket_factory(scenario)
|
|
try:
|
stream = await open_tcp_stream("test.example.com", port, **kwargs)
|
assert expect_error == ()
|
scenario.check(stream.socket)
|
return (stream.socket, scenario)
|
except AssertionError: # pragma: no cover
|
raise
|
except expect_error as exc:
|
scenario.check(None)
|
return (exc, scenario)
|
|
|
async def test_one_host_quick_success(autojump_clock):
|
sock, scenario = await run_scenario(80, [("1.2.3.4", 0.123, "success")])
|
assert sock.ip == "1.2.3.4"
|
assert trio.current_time() == 0.123
|
|
|
async def test_one_host_slow_success(autojump_clock):
|
sock, scenario = await run_scenario(81, [("1.2.3.4", 100, "success")])
|
assert sock.ip == "1.2.3.4"
|
assert trio.current_time() == 100
|
|
|
async def test_one_host_quick_fail(autojump_clock):
|
exc, scenario = await run_scenario(
|
82, [("1.2.3.4", 0.123, "error")], expect_error=OSError
|
)
|
assert isinstance(exc, OSError)
|
assert trio.current_time() == 0.123
|
|
|
async def test_one_host_slow_fail(autojump_clock):
|
exc, scenario = await run_scenario(
|
83, [("1.2.3.4", 100, "error")], expect_error=OSError
|
)
|
assert isinstance(exc, OSError)
|
assert trio.current_time() == 100
|
|
|
async def test_one_host_failed_after_connect(autojump_clock):
|
exc, scenario = await run_scenario(
|
83, [("1.2.3.4", 1, "postconnect_fail")], expect_error=KeyboardInterrupt
|
)
|
assert isinstance(exc, KeyboardInterrupt)
|
|
|
# With the default 0.250 second delay, the third attempt will win
|
async def test_basic_fallthrough(autojump_clock):
|
sock, scenario = await run_scenario(
|
80,
|
[
|
("1.1.1.1", 1, "success"),
|
("2.2.2.2", 1, "success"),
|
("3.3.3.3", 0.2, "success"),
|
],
|
)
|
assert sock.ip == "3.3.3.3"
|
# current time is default time + default time + connection time
|
assert trio.current_time() == (0.250 + 0.250 + 0.2)
|
assert scenario.connect_times == {
|
"1.1.1.1": 0,
|
"2.2.2.2": 0.250,
|
"3.3.3.3": 0.500,
|
}
|
|
|
async def test_early_success(autojump_clock):
|
sock, scenario = await run_scenario(
|
80,
|
[
|
("1.1.1.1", 1, "success"),
|
("2.2.2.2", 0.1, "success"),
|
("3.3.3.3", 0.2, "success"),
|
],
|
)
|
assert sock.ip == "2.2.2.2"
|
assert trio.current_time() == (0.250 + 0.1)
|
assert scenario.connect_times == {
|
"1.1.1.1": 0,
|
"2.2.2.2": 0.250,
|
# 3.3.3.3 was never even started
|
}
|
|
|
# With a 0.450 second delay, the first attempt will win
|
async def test_custom_delay(autojump_clock):
|
sock, scenario = await run_scenario(
|
80,
|
[
|
("1.1.1.1", 1, "success"),
|
("2.2.2.2", 1, "success"),
|
("3.3.3.3", 0.2, "success"),
|
],
|
happy_eyeballs_delay=0.450,
|
)
|
assert sock.ip == "1.1.1.1"
|
assert trio.current_time() == 1
|
assert scenario.connect_times == {
|
"1.1.1.1": 0,
|
"2.2.2.2": 0.450,
|
"3.3.3.3": 0.900,
|
}
|
|
|
async def test_custom_errors_expedite(autojump_clock):
|
sock, scenario = await run_scenario(
|
80,
|
[
|
("1.1.1.1", 0.1, "error"),
|
("2.2.2.2", 0.2, "error"),
|
("3.3.3.3", 10, "success"),
|
# .25 is the default timeout
|
("4.4.4.4", 0.25, "success"),
|
],
|
)
|
assert sock.ip == "4.4.4.4"
|
assert trio.current_time() == (0.1 + 0.2 + 0.25 + 0.25)
|
assert scenario.connect_times == {
|
"1.1.1.1": 0,
|
"2.2.2.2": 0.1,
|
"3.3.3.3": 0.1 + 0.2,
|
"4.4.4.4": 0.1 + 0.2 + 0.25,
|
}
|
|
|
async def test_all_fail(autojump_clock):
|
exc, scenario = await run_scenario(
|
80,
|
[
|
("1.1.1.1", 0.1, "error"),
|
("2.2.2.2", 0.2, "error"),
|
("3.3.3.3", 10, "error"),
|
("4.4.4.4", 0.250, "error"),
|
],
|
expect_error=OSError,
|
)
|
assert isinstance(exc, OSError)
|
assert isinstance(exc.__cause__, BaseExceptionGroup)
|
assert len(exc.__cause__.exceptions) == 4
|
assert trio.current_time() == (0.1 + 0.2 + 10)
|
assert scenario.connect_times == {
|
"1.1.1.1": 0,
|
"2.2.2.2": 0.1,
|
"3.3.3.3": 0.1 + 0.2,
|
"4.4.4.4": 0.1 + 0.2 + 0.25,
|
}
|
|
|
async def test_multi_success(autojump_clock):
|
sock, scenario = await run_scenario(
|
80,
|
[
|
("1.1.1.1", 0.5, "error"),
|
("2.2.2.2", 10, "success"),
|
("3.3.3.3", 10 - 1, "success"),
|
("4.4.4.4", 10 - 2, "success"),
|
("5.5.5.5", 0.5, "error"),
|
],
|
happy_eyeballs_delay=1,
|
)
|
assert not scenario.sockets["1.1.1.1"].succeeded
|
assert (
|
scenario.sockets["2.2.2.2"].succeeded
|
or scenario.sockets["3.3.3.3"].succeeded
|
or scenario.sockets["4.4.4.4"].succeeded
|
)
|
assert not scenario.sockets["5.5.5.5"].succeeded
|
assert sock.ip in ["2.2.2.2", "3.3.3.3", "4.4.4.4"]
|
assert trio.current_time() == (0.5 + 10)
|
assert scenario.connect_times == {
|
"1.1.1.1": 0,
|
"2.2.2.2": 0.5,
|
"3.3.3.3": 1.5,
|
"4.4.4.4": 2.5,
|
"5.5.5.5": 3.5,
|
}
|
|
|
async def test_does_reorder(autojump_clock):
|
sock, scenario = await run_scenario(
|
80,
|
[
|
("1.1.1.1", 10, "error"),
|
# This would win if we tried it first...
|
("2.2.2.2", 1, "success"),
|
# But in fact we try this first, because of section 5.4
|
("::3", 0.5, "success"),
|
],
|
happy_eyeballs_delay=1,
|
)
|
assert sock.ip == "::3"
|
assert trio.current_time() == 1 + 0.5
|
assert scenario.connect_times == {
|
"1.1.1.1": 0,
|
"::3": 1,
|
}
|
|
|
async def test_handles_no_ipv4(autojump_clock):
|
sock, scenario = await run_scenario(
|
80,
|
# Here the ipv6 addresses fail at socket creation time, so the connect
|
# configuration doesn't matter
|
[
|
("::1", 10, "success"),
|
("2.2.2.2", 0, "success"),
|
("::3", 0.1, "success"),
|
("4.4.4.4", 0, "success"),
|
],
|
happy_eyeballs_delay=1,
|
ipv4_supported=False,
|
)
|
assert sock.ip == "::3"
|
assert trio.current_time() == 1 + 0.1
|
assert scenario.connect_times == {
|
"::1": 0,
|
"::3": 1.0,
|
}
|
|
|
async def test_handles_no_ipv6(autojump_clock):
|
sock, scenario = await run_scenario(
|
80,
|
# Here the ipv6 addresses fail at socket creation time, so the connect
|
# configuration doesn't matter
|
[
|
("::1", 0, "success"),
|
("2.2.2.2", 10, "success"),
|
("::3", 0, "success"),
|
("4.4.4.4", 0.1, "success"),
|
],
|
happy_eyeballs_delay=1,
|
ipv6_supported=False,
|
)
|
assert sock.ip == "4.4.4.4"
|
assert trio.current_time() == 1 + 0.1
|
assert scenario.connect_times == {
|
"2.2.2.2": 0,
|
"4.4.4.4": 1.0,
|
}
|
|
|
async def test_no_hosts(autojump_clock):
|
exc, scenario = await run_scenario(80, [], expect_error=OSError)
|
assert "no results found" in str(exc)
|
|
|
async def test_cancel(autojump_clock):
|
with trio.move_on_after(5) as cancel_scope:
|
exc, scenario = await run_scenario(
|
80,
|
[
|
("1.1.1.1", 10, "success"),
|
("2.2.2.2", 10, "success"),
|
("3.3.3.3", 10, "success"),
|
("4.4.4.4", 10, "success"),
|
],
|
expect_error=BaseExceptionGroup,
|
)
|
# What comes out should be 1 or more Cancelled errors that all belong
|
# to this cancel_scope; this is the easiest way to check that
|
raise exc
|
assert cancel_scope.cancelled_caught
|
|
assert trio.current_time() == 5
|
|
# This should have been called already, but just to make sure, since the
|
# exception-handling logic in run_scenario is a bit complicated and the
|
# main thing we care about here is that all the sockets were cleaned up.
|
scenario.check(succeeded=False)
|
