Example #1
0
    def __init__(self,
                 sock_fam: socket.AddressFamily,
                 icmp_proto: int,
                 src: IPAddressType,
                 send_limiter: TokenBucket,
                 *,
                 loop: asyncio.AbstractEventLoop = None) -> None:
        tcp_sock = socket.socket(sock_fam, socket.SOCK_RAW, socket.IPPROTO_TCP)
        icmp_sock = socket.socket(sock_fam, socket.SOCK_RAW, icmp_proto)
        for sock in (tcp_sock, icmp_sock):
            sock.setblocking(False)
            sock.setsockopt(socket.SOL_SOCKET, _SO_RCVBUFFORCE, 8_388_608)
            sock.bind((str(src), 0))

        if loop is None:
            loop = asyncio.get_event_loop()
        loop.add_reader(tcp_sock.fileno(), self._handle_read)
        loop.add_reader(icmp_sock.fileno(), self._handle_icmp_read)
        loop.add_writer(tcp_sock.fileno(), self._handle_write)

        self._sock = tcp_sock
        self._icmp_sock = icmp_sock
        self._src_addr: IPAddressType = src
        self._test_map: Dict[Tuple[int, bytes, int],
                             "BaseTest[IPAddressType]"] = {}
        self._send_queue: Deque[OutgoingPacket[IPAddressType]] = deque()
        self._send_limiter = send_limiter
        self._sent_rsts: Dict[Tuple[int, bytes, int], int] = {}
        self._loop = loop
Example #2
0
    def _sendfile_cb(self, fut: 'asyncio.Future[None]', out_fd: int,
                     in_fd: int, offset: int, count: int,
                     loop: asyncio.AbstractEventLoop,
                     registered: bool) -> None:
        if registered:
            loop.remove_writer(out_fd)
        if fut.cancelled():
            return

        try:
            n = os.sendfile(out_fd, in_fd, offset, count)
            if n == 0:  # EOF reached
                n = count
        except (BlockingIOError, InterruptedError):
            n = 0
        except Exception as exc:
            set_exception(fut, exc)
            return

        if n < count:
            loop.add_writer(out_fd, self._sendfile_cb, fut, out_fd, in_fd,
                            offset + n, count - n, loop, True)
        else:
            set_result(fut, None)
Example #3
0
async def wait_for_writer(fd: Union[int, io.BufferedWriter], loop: asyncio.AbstractEventLoop) -> None:
    f: asyncio.Future = asyncio.Future()
    loop.add_writer(fd, f.set_result, None)
    f.add_done_callback(lambda f: loop.remove_writer(fd))
    await f