from __future__ import annotations
import socket
import sys
from socket import AddressFamily, SocketKind
from typing import TYPE_CHECKING
import attrs
import pytest
import trio
from trio._highlevel_open_tcp_stream import (
close_all,
format_host_port,
open_tcp_stream,
reorder_for_rfc_6555_section_5_4,
)
from trio.socket import AF_INET, AF_INET6, IPPROTO_TCP, SOCK_STREAM, SocketType
from trio.testing import Matcher, RaisesGroup
if TYPE_CHECKING:
from collections.abc import Sequence
from trio.testing import MockClock
if sys.version_info < (3, 11):
from exceptiongroup import BaseExceptionGroup
def test_close_all() -> None:
class CloseMe(SocketType):
closed = False
def close(self) -> None:
self.closed = True
class CloseKiller(SocketType):
def close(self) -> None:
raise OSError("os error text")
c: CloseMe = 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, match="os error text"):
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() -> None:
def fake4(
i: int,
) -> tuple[socket.AddressFamily, socket.SocketKind, int, str, tuple[str, int]]:
return (
AF_INET,
SOCK_STREAM,
IPPROTO_TCP,
"",
(f"10.0.0.{i}", 80),
)
def fake6(
i: int,
) -> tuple[socket.AddressFamily, socket.SocketKind, int, str, tuple[str, int]]:
return (AF_INET6, SOCK_STREAM, IPPROTO_TCP, "", (f"::{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() -> None:
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() -> None:
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() -> None:
with pytest.raises(ValueError, match=r"^host must be str or bytes, not None$"):
await open_tcp_stream(None, 80) # type: ignore[arg-type]
with pytest.raises(TypeError):
await open_tcp_stream("127.0.0.1", b"80") # type: ignore[arg-type]
def can_bind_127_0_0_2() -> bool:
with socket.socket() as s:
try:
s.bind(("127.0.0.2", 0))
except OSError:
return False
# s.getsockname() is typed as returning Any
return s.getsockname()[0] == "127.0.0.2" # type: ignore[no-any-return]
async def test_local_address_real() -> None:
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.
local_address = "127.0.0.2" if can_bind_127_0_0_2() else "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()
# accept returns tuple[SocketType, object], due to typeshed returning `Any`
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,
match=r"^all attempts to connect* to *127\.0\.0\.\d:\d+ failed$",
):
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
@attrs.define(eq=False, slots=False)
class FakeSocket(trio.socket.SocketType):
scenario: Scenario
_family: AddressFamily
_type: SocketKind
_proto: int
ip: str | int | None = None
port: str | int | None = None
succeeded: bool = False
closed: bool = False
failing: bool = False
@property
def type(self) -> SocketKind:
return self._type
@property
def family(self) -> AddressFamily: # pragma: no cover
return self._family
@property
def proto(self) -> int: # pragma: no cover
return self._proto
async def connect(self, sockaddr: tuple[str | int, str | int | None]) -> None:
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) -> None:
self.closed = True
# called when SocketStream is constructed
def setsockopt(self, *args: object, **kwargs: object) -> None:
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: int,
ip_list: Sequence[tuple[str, float, str]],
supported_families: set[AddressFamily],
) -> None:
# ip_list have to be unique
ip_order = [ip for (ip, _, _) in ip_list]
assert len(set(ip_order)) == len(ip_list)
ip_dict: dict[str | int, tuple[float, str]] = {}
for ip, delay, result in ip_list:
assert delay >= 0
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: dict[str | int, FakeSocket] = {}
self.connect_times: dict[str | int, float] = {}
def socket(
self,
family: AddressFamily | int | None = None,
type_: SocketKind | int | None = None,
proto: int | None = None,
) -> SocketType:
assert isinstance(family, AddressFamily)
assert isinstance(type_, SocketKind)
assert proto is not None
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: str) -> tuple[
AddressFamily,
SocketKind,
int,
str,
tuple[str, int, int, int] | tuple[str, int] | tuple[int, bytes],
]:
sockaddr: tuple[str, int] | tuple[str, int, int, int] | tuple[int, bytes]
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: bytes | None,
port: bytes | str | int | None,
family: int = -1,
type: int = -1,
proto: int = -1,
flags: int = -1,
) -> list[
tuple[
AddressFamily,
SocketKind,
int,
str,
tuple[str, int, int, int] | tuple[str, int] | tuple[int, bytes],
]
]:
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: tuple[str, int] | tuple[str, int, int, int],
flags: int,
) -> tuple[str, str]:
raise NotImplementedError
def check(self, succeeded: SocketType | None) -> None:
# 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: int,
# 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: Sequence[tuple[str, float, str]],
*,
# If False, AF_INET4/6 sockets error out on creation, before connect is
# even called.
ipv4_supported: bool = True,
ipv6_supported: bool = 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: tuple[type[BaseException], ...] | type[BaseException] = (),
happy_eyeballs_delay: float | None = 0.25,
local_address: str | None = None,
) -> tuple[SocketType, Scenario] | tuple[BaseException, Scenario]:
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,
happy_eyeballs_delay=happy_eyeballs_delay,
local_address=local_address,
)
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: MockClock) -> None:
sock, _scenario = await run_scenario(80, [("1.2.3.4", 0.123, "success")])
assert isinstance(sock, FakeSocket)
assert sock.ip == "1.2.3.4"
assert trio.current_time() == 0.123
async def test_one_host_slow_success(autojump_clock: MockClock) -> None:
sock, _scenario = await run_scenario(81, [("1.2.3.4", 100, "success")])
assert isinstance(sock, FakeSocket)
assert sock.ip == "1.2.3.4"
assert trio.current_time() == 100
async def test_one_host_quick_fail(autojump_clock: MockClock) -> None:
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: MockClock) -> None:
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: MockClock) -> None:
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: MockClock) -> None:
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 isinstance(sock, FakeSocket)
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: MockClock) -> None:
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 isinstance(sock, FakeSocket)
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: MockClock) -> None:
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 isinstance(sock, FakeSocket)
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_none_default(autojump_clock: MockClock) -> None:
"""Copy of test_basic_fallthrough, but specifying the delay =None"""
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=None,
)
assert isinstance(sock, FakeSocket)
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_custom_errors_expedite(autojump_clock: MockClock) -> None:
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 isinstance(sock, FakeSocket)
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: MockClock) -> None:
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)
subexceptions = (Matcher(OSError, match="^sorry$"),) * 4
assert RaisesGroup(
*subexceptions,
match="all attempts to connect to test.example.com:80 failed",
).matches(exc.__cause__)
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: MockClock) -> None:
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 isinstance(sock, FakeSocket)
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: MockClock) -> None:
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 isinstance(sock, FakeSocket)
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: MockClock) -> None:
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 isinstance(sock, FakeSocket)
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: MockClock) -> None:
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 isinstance(sock, FakeSocket)
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: MockClock) -> None:
exc, _scenario = await run_scenario(80, [], expect_error=OSError)
assert "no results found" in str(exc)
async def test_cancel(autojump_clock: MockClock) -> None:
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,
)
assert isinstance(exc, BaseException)
# 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=None)