@ajksunkang-aios/kgraph-linux-x64 0.1.2 → 0.1.3

package/lib/site-packages/trio/_tests/test_highlevel_open_tcp_stream.py

@@ -0,0 +1,693 @@
+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
+
+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 = (pytest.RaisesExc(OSError, match="^sorry$"),) * 4
+    assert pytest.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)