fix asyncio bugs

asyncio_bug_fixes.py

# a collection of bugfix backports to earlier asyncio, use asyncio_bug_fixes.patch_loop() to patch it
import asyncio
import asyncio.base_events
import functools
import socket
import ipaddress
from asyncio import futures


# Linux's sock.type is a bitmask that can include extra info about socket.
_SOCKET_TYPE_MASK = 0
if hasattr(socket, 'SOCK_NONBLOCK'):
    _SOCKET_TYPE_MASK |= socket.SOCK_NONBLOCK
if hasattr(socket, 'SOCK_CLOEXEC'):
    _SOCKET_TYPE_MASK |= socket.SOCK_CLOEXEC


@functools.lru_cache(maxsize=1024, typed=True)
def _ipaddr_info(host, port, family, type, proto):
    # Try to skip getaddrinfo if "host" is already an IP. Since getaddrinfo
    # blocks on an exclusive lock on some platforms, users might handle name
    # resolution in their own code and pass in resolved IPs.
    if proto not in {0, socket.IPPROTO_TCP, socket.IPPROTO_UDP} or host is None:
        return None

    type &= ~_SOCKET_TYPE_MASK
    if type == socket.SOCK_STREAM:
        proto = socket.IPPROTO_TCP
    elif type == socket.SOCK_DGRAM:
        proto = socket.IPPROTO_UDP
    else:
        return None

    if port in {None, '', b''}:
        port = 0
    elif isinstance(port, (bytes, str)):
        port = int(port)

    if hasattr(socket, 'inet_pton'):
        if family == socket.AF_UNSPEC:
            afs = [socket.AF_INET, socket.AF_INET6]
        else:
            afs = [family]

        for af in afs:
            # Linux's inet_pton doesn't accept an IPv6 zone index after host,
            # like '::1%lo0', so strip it. If we happen to make an invalid
            # address look valid, we fail later in sock.connect or sock.bind.
            try:
                if af == socket.AF_INET6:
                    socket.inet_pton(af, host.partition('%')[0])
                else:
                    socket.inet_pton(af, host)
                return af, type, proto, '', (host, port)
            except OSError:
                pass

        # "host" is not an IP address.
        return None

    # No inet_pton. (On Windows it's only available since Python 3.4.)
    # Even though getaddrinfo with AI_NUMERICHOST would be non-blocking, it
    # still requires a lock on some platforms, and waiting for that lock could
    # block the event loop. Use ipaddress instead, it's just text parsing.
    try:
        addr = ipaddress.IPv4Address(host)
    except ValueError:
        try:
            addr = ipaddress.IPv6Address(host.partition('%')[0])
        except ValueError:
            return None

    af = socket.AF_INET if addr.version == 4 else socket.AF_INET6
    if family not in (socket.AF_UNSPEC, af):
        # "host" is wrong IP version for "family".
        return None

    return af, type, proto, '', (host, port)


def _fixed_getaddrinfo(loop, host, port, *,
                       family=0, type=0, proto=0, flags=0):
    info = _ipaddr_info(host, port, family, type, proto)
    if info is not None:
        fut = futures.Future(loop=loop)
        fut.set_result([info])
        return fut
    elif loop._debug:
        return loop.run_in_executor(None, loop._getaddrinfo_debug,
                                    host, port, family, type, proto, flags)
    else:
        return loop.run_in_executor(None, socket.getaddrinfo,
                                    host, port, family, type, proto, flags)


def patch_loop(loop=None):
    if not loop:
        loop = asyncio.get_event_loop()

    # implement https://github.com/python/asyncio/pull/302 if it's not already available
    if '_ipaddr_info' not in dir(asyncio.base_events):
        loop.getaddrinfo = lambda *args, **kwargs: _fixed_getaddrinfo(loop, *args, **kwargs)

    # return the patched loop if possible
    return loop