diff --git a/dnsrewriteproxy.py b/dnsrewriteproxy.py index e90642f..12f4224 100644 --- a/dnsrewriteproxy.py +++ b/dnsrewriteproxy.py @@ -22,6 +22,7 @@ from aiodnsresolver import ( ResourceRecord, pack, parse, + recvfrom, ) @@ -79,7 +80,7 @@ def DnsProxy( try: while True: - request_data, addr = await recvfrom(loop, sock, 512) + request_data, addr = await recvfrom(loop, [sock], 512) await upstream_queue.put((request_data, addr)) finally: # Finish upstream requests, which can add to to the downstream @@ -208,36 +209,6 @@ def error(query, rcode): ) -async def recvfrom(loop, sock, max_bytes): - try: - return sock.recvfrom(max_bytes) - except BlockingIOError: - pass - - def reader(): - try: - (data, addr) = sock.recvfrom(max_bytes) - except BlockingIOError: - pass - except BaseException as exception: - loop.remove_reader(fileno) - if not result.done(): - result.set_exception(exception) - else: - loop.remove_reader(fileno) - if not result.done(): - result.set_result((data, addr)) - - fileno = sock.fileno() - result = Future() - loop.add_reader(fileno, reader) - - try: - return await result - finally: - loop.remove_reader(fileno) - - async def sendto(loop, sock, data, addr): # In our cases, the UDP responses will always be 512 bytes or less. # Even if sendto sent some of the data, there is no way for the other