AzurLaneAutoScript/module/smart_mgmt/happy_eyeballs.py

from __future__ import annotations

import math
import select
import socket
import threading
import time
from ipaddress import ip_address

from requests.adapters import HTTPAdapter
from urllib3 import HTTPSConnectionPool
from urllib3.connection import VerifiedHTTPSConnection
from urllib3.exceptions import ConnectTimeoutError, NewConnectionError
from urllib3.poolmanager import PoolKey, _default_key_normalizer
from urllib3.util.connection import _set_socket_options

# RFC 8305 Happy Eyeballs tuning.
# Inter-family: IPv6 first attempt goes out immediately, IPv4 follows after
# INTER_FAMILY_DELAY (in-flight capped at 2, one per family).
# The stagger drives invalid connections toward zero when IPv6 works.
INTER_FAMILY_DELAY = 0.25

DEFAULT_PER_ADDRESS_TIMEOUT = 2.0
DEFAULT_CONNECT_TIMEOUT = 5.0
DEFAULT_TOTAL_DEADLINE = 10.0
TOTAL_DEADLINE_MARGIN = 5.0


def _validate_timeout(value, name):
    if value is None:
        return None
    try:
        value = float(value)
    except (TypeError, ValueError):
        raise ValueError('%s must be a positive finite number or None' % name)
    if not math.isfinite(value) or value <= 0:
        raise ValueError('%s must be a positive finite number or None' % name)
    return value


def _resolve_timeouts(per_address_timeout, connect_timeout, total_deadline):
    """
    Apply defaults and validate the three connection timeouts.

    Missing parameters fall back to defaults.
    If total_deadline is missing but connect_timeout is given or less than
    connect_timeout, it becomes connect_timeout + TOTAL_DEADLINE_MARGIN.
    """
    per_address_timeout = _validate_timeout(per_address_timeout, 'per_address_timeout')
    if per_address_timeout is None:
        per_address_timeout = DEFAULT_PER_ADDRESS_TIMEOUT

    connect_given = connect_timeout is not None
    connect_timeout = _validate_timeout(connect_timeout, 'connect_timeout')
    if connect_timeout is None:
        connect_timeout = DEFAULT_CONNECT_TIMEOUT

    total_deadline = _validate_timeout(total_deadline, 'total_deadline')
    if total_deadline is None:
        if connect_given:
            total_deadline = connect_timeout + TOTAL_DEADLINE_MARGIN
        else:
            total_deadline = DEFAULT_TOTAL_DEADLINE
    elif total_deadline < connect_timeout:
        total_deadline = connect_timeout + TOTAL_DEADLINE_MARGIN

    return per_address_timeout, connect_timeout, total_deadline


def _is_public_address(ip_str: str) -> bool:
    """
    Rejects loopback/private/reserved/link-local addresses.
    """
    try:
        return ip_address(ip_str).is_global
    except ValueError:
        return False


def _start_connect(af, sa, socket_options, source_address):
    """
    Create a non-blocking socket and begin a non-blocking connect.

    Returns (sock, None) when the attempt is in flight, or (None, err)
    when the socket could not even be started (e.g. EAFNOSUPPORT).
    """
    sock = socket.socket(af, socket.SOCK_STREAM)
    try:
        sock.setblocking(False)
        if socket_options:
            _set_socket_options(sock, socket_options)
        if source_address:
            sock.bind(source_address)
        try:
            sock.connect(sa)
        except (BlockingIOError, InterruptedError):
            pass
        return sock, None
    except OSError as e:
        sock.close()
        return None, e


def _close_all(inflight):
    for sock in list(inflight.keys()):
        try:
            sock.close()
        except OSError:
            pass
    inflight.clear()


def _connect_loop(
        host,
        port,
        socket_options,
        source_address,
        cancel_event,
        per_address_timeout,
        connect_timeout
    ):
    """
    RFC 8305-inspired Happy Eyeballs connection loop.
    """
    infos = socket.getaddrinfo(host, port, socket.AF_UNSPEC, socket.SOCK_STREAM)
    # The caller may have given up (total_deadline) while getaddrinfo() was blocking;
    # bail out immediately instead of starting any connect.
    if cancel_event.is_set():
        raise socket.timeout('Happy eyeballs connect cancelled for %s:%s' % (host, port))
    v6, v4 = [], []

    for af, _, _, _, sa in infos:
        if not sa or not _is_public_address(sa[0]):
            continue
        if af == socket.AF_INET6:
            v6.append((af, sa))
        elif af == socket.AF_INET:
            v4.append((af, sa))

    if not v6 and not v4:
        raise socket.error('No usable public address for %s:%s' % (host, port))

    start = time.monotonic()
    deadline = start + connect_timeout
    # No stagger to wait out when there is no IPv6 candidate.
    v4_launch_time = start + (INTER_FAMILY_DELAY if v6 else 0.0)
    state = {'v6_idx': 0, 'v4_idx': 0, 'v4_launched': False}
    inflight = {}
    connect_error = None
    timed_out = False

    def launch(family):
        nonlocal connect_error
        candidates = v6 if family == 'v6' else v4
        index_key = family + '_idx'

        while state[index_key] < len(candidates):
            if cancel_event.is_set() or time.monotonic() >= deadline:
                return

            af, sa = candidates[state[index_key]]
            state[index_key] += 1

            sock, err = _start_connect(af, sa, socket_options, source_address)
            if sock is None:
                if err is not None:
                    connect_error = err
                continue

            inflight[sock] = {
                'family': family,
                'per_deadline': time.monotonic() + per_address_timeout,
            }
            return

    def close_and_advance(sock, family, err=None):
        nonlocal connect_error, timed_out
        if err is None:
            timed_out = True
        else:
            connect_error = err
        try:
            sock.close()
        except OSError:
            pass
        launch(family)

    if v6:
        launch('v6')

    while True:
        if cancel_event.is_set():
            _close_all(inflight)
            raise socket.timeout('Happy eyeballs connect cancelled for %s:%s' % (host, port))

        now = time.monotonic()
        if now >= deadline:
            break

        if v4 and not state['v4_launched'] and now >= v4_launch_time:
            state['v4_launched'] = True
            launch('v4')

        if not inflight:
            # Only waiting for the IPv4 stagger, or fully exhausted.
            if v4 and not state['v4_launched']:
                # Wait out the stagger window; Event.wait returns early if
                # cancelled (Windows select rejects empty fd lists).
                wait = max(0.0, v4_launch_time - time.monotonic())
                cancel_event.wait(wait)
                continue
            break

        next_times = [deadline]
        if v4 and not state['v4_launched']:
            next_times.append(v4_launch_time)
        for info in inflight.values():
            next_times.append(info['per_deadline'])
        select_timeout = max(0.0, min(next_times) - time.monotonic())

        try:
            _, writable, exceptional = select.select(
                [], list(inflight.keys()), list(inflight.keys()), select_timeout
            )
        except (OSError, ValueError):
            # select() failing is not "no events": the connect state machine
            # can no longer be trusted, close everything and abort.
            _close_all(inflight)
            raise socket.error(
                'select failed during Happy Eyeballs connect for %s:%s'
                % (host, port)
            )

        now = time.monotonic()
        if v4 and not state['v4_launched'] and now >= v4_launch_time:
            state['v4_launched'] = True
            launch('v4')

        winner = None
        for sock in writable:
            info = inflight.pop(sock, None)
            if info is None:
                continue
            try:
                err = sock.getsockopt(socket.SOL_SOCKET, socket.SO_ERROR)
            except OSError as e:
                close_and_advance(sock, info['family'], e)
                continue
            if err == 0:
                winner = sock
                break
            close_and_advance(sock, info['family'], socket.error(err))

        if winner is None:
            # Some platforms report refused connects via exceptfds rather
            # than writable; use the same SO_ERROR logic as writable.
            for sock in exceptional:
                info = inflight.pop(sock, None)
                if info is None:
                    continue
                try:
                    err = sock.getsockopt(socket.SOL_SOCKET, socket.SO_ERROR)
                except OSError as e:
                    close_and_advance(sock, info['family'], e)
                    continue
                if err == 0:
                    winner = sock
                    break
                close_and_advance(sock, info['family'], socket.error(err))

        if winner is not None:
            _close_all(inflight)
            winner.setblocking(True)
            return winner

        # Per-address short-timeout advancement for sockets not yet writable.
        for sock in list(inflight.keys()):
            info = inflight.get(sock)
            if info is None or now < info['per_deadline']:
                continue
            inflight.pop(sock, None)
            close_and_advance(sock, info['family'])

    exhausted = (
        not inflight
        and state['v6_idx'] >= len(v6)
        and (not v4 or (state['v4_launched'] and state['v4_idx'] >= len(v4)))
    )
    _close_all(inflight)
    if exhausted:
        if connect_error is not None:
            raise socket.error(
                'All connect attempts failed for %s:%s (last error: %s)'
                % (host, port, connect_error)
            )
        if timed_out:
            raise socket.timeout(
                'All connect attempts timed out for %s:%s'
                % (host, port)
            )
    raise socket.timeout('Happy eyeballs connect timeout exhausted for %s:%s' % (host, port))


def happy_eyeballs_create_connection(
        address,
        per_address_timeout=None,
        connect_timeout=None,
        total_deadline=None,
        source_address=None,
        socket_options=None,
    ):
    """
    Establish a TCP connection with RFC 8305-inspired Happy Eyeballs.

    The connect loop runs in a daemon worker thread so the blocking getaddrinfo()
    cannot stall the caller beyond total_deadline; on expiry the cancel event tells
    the worker to self-terminate and socket.timeout is raised.

    Args:
        per_address_timeout: window to wait for a single address;
            on expiry the address is abandoned and the next candidate of the same family is tried.
        connect_timeout: budget for the whole connect loop (getaddrinfo excluded);
            the loop gives up when exhausted.
        total_deadline: overall deadline including getaddrinfo;
            must be at least connect_timeout.

    Raises OSError subclasses (socket.timeout / socket.error) on failure
    so urllib3's exception translation in _new_conn applies unchanged.
    """
    host, port = address
    if host.startswith('[') and host.endswith(']'):
        host = host[1:-1]

    per_address_timeout, connect_timeout, total_deadline = _resolve_timeouts(
        per_address_timeout, connect_timeout, total_deadline
    )

    holder = {}
    cancel_event = threading.Event()

    def worker():
        try:
            sock = _connect_loop(
                host,
                port,
                socket_options,
                source_address,
                cancel_event,
                per_address_timeout,
                connect_timeout,
            )
            remaining = deadline - time.monotonic()
            if cancel_event.is_set() or remaining <= 0:
                sock.close()
                raise socket.timeout(
                    'Happy eyeballs total deadline exceeded: %s:%s'
                    % (host, port)
                )
            # Let the TLS handshake done by urllib3 after _new_conn()
            # inherit the remaining total_deadline budget.
            sock.settimeout(remaining)
            holder['sock'] = sock
        except Exception as e:
            holder['exc'] = e

    t = threading.Thread(target=worker, name='happy-eyeballs', daemon=True)
    deadline = time.monotonic() + total_deadline
    t.start()
    t.join(max(0.0, deadline - time.monotonic()))

    if t.is_alive():
        # Most likely stuck in uninterruptible getaddrinfo(). Signal cancel and
        # wait briefly; the thread will self-terminate once getaddrinfo() returns.
        cancel_event.set()
        t.join(0.5)
        # The worker may have completed just before the deadline and left a
        # connected socket behind; it must not survive a timed-out call.
        sock = holder.get('sock')
        if sock is not None:
            try:
                sock.close()
            except OSError:
                pass
        raise socket.timeout('Happy eyeballs total deadline exceeded: %s:%s' % (host, port))

    exc = holder.get('exc')
    if exc is not None:
        raise exc

    sock = holder.get('sock')
    if sock is None:
        raise socket.error('Happy eyeballs failed without a socket: %s:%s' % (host, port))

    return sock


_HE_TIMEOUT_KEYS = ('he_per_address_timeout', 'he_connect_timeout', 'he_total_deadline')


def _he_key_normalizer(context):
    """
    urllib3 1.22 pool key normalizer: drop the custom timeout keys before
    PoolKey construction. PoolKey only accepts its fixed _key_fields.
    """
    context = dict(context)
    for key in _HE_TIMEOUT_KEYS:
        context.pop(key, None)
    return _default_key_normalizer(PoolKey, context)


class HappyEyeballsHTTPSConnection(VerifiedHTTPSConnection):
    """HTTPS connection using Happy Eyeballs for TCP connection setup."""

    def __init__(self, *args, **kwargs):
        # Pop the custom timeouts before urllib3 sees them: they are not urllib3 parameters
        # and would be rejected by HTTPSConnection if left in **kwargs.
        # None values are defaulted later in _resolve_timeouts.
        self.he_per_address_timeout = kwargs.pop('he_per_address_timeout', None)
        self.he_connect_timeout = kwargs.pop('he_connect_timeout', None)
        self.he_total_deadline = kwargs.pop('he_total_deadline', None)
        super(HappyEyeballsHTTPSConnection, self).__init__(*args, **kwargs)

    def _new_conn(self):
        extra_kw = {}
        if self.source_address:
            extra_kw['source_address'] = self.source_address
        if self.socket_options:
            extra_kw['socket_options'] = self.socket_options
        try:
            conn = happy_eyeballs_create_connection(
                (self.host, self.port),
                self.he_per_address_timeout,
                self.he_connect_timeout,
                self.he_total_deadline,
                **extra_kw
            )
        except socket.timeout as e:
            raise ConnectTimeoutError(self, 'Connection to %s timed out. (happy eyeballs)' % self.host)
        except socket.error as e:
            raise NewConnectionError(self, 'Failed to establish a new connection: %s' % e)

        return conn


class HappyEyeballsHTTPSConnectionPool(HTTPSConnectionPool):
    ConnectionCls = HappyEyeballsHTTPSConnection


class HappyEyeballsHTTPAdapter(HTTPAdapter):
    """
    Mount onto a requests Session to route HTTPS through Happy Eyeballs.

    Only the HTTPS pool class is replaced; plain HTTP keeps urllib3 defaults.
    """

    def __init__(
            self,
            per_address_timeout=None,
            connect_timeout=None,
            total_deadline=None,
            *args,
            **kwargs
        ):
        # Timeouts pass through unresolved; defaults and validation happen
        # once in _resolve_timeouts at happy_eyeballs_create_connection.
        self.he_per_address_timeout = per_address_timeout
        self.he_connect_timeout = connect_timeout
        self.he_total_deadline = total_deadline

        super(HappyEyeballsHTTPAdapter, self).__init__(*args, **kwargs)

    def init_poolmanager(self, connections, maxsize, block=False, **pool_kwargs):
        super(HappyEyeballsHTTPAdapter, self).init_poolmanager(connections, maxsize, block=block, **pool_kwargs)
        scheme_map = dict(self.poolmanager.pool_classes_by_scheme)
        scheme_map['https'] = HappyEyeballsHTTPSConnectionPool
        self.poolmanager.pool_classes_by_scheme = scheme_map

        # Push the timeouts into PoolManager's shared pool kwargs so every
        # https pool created afterwards forwards them to its connections via
        # conn_kw (urllib3 1.22 _new_conn expands conn_kw into ConnectionCls).
        self.poolmanager.connection_pool_kw.update({
            'he_per_address_timeout': self.he_per_address_timeout,
            'he_connect_timeout': self.he_connect_timeout,
            'he_total_deadline': self.he_total_deadline,
        })
        # Replace only the https key normalizer
        key_scheme_map = dict(self.poolmanager.key_fn_by_scheme)
        key_scheme_map['https'] = _he_key_normalizer
        self.poolmanager.key_fn_by_scheme = key_scheme_map