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
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)
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