(refactor) Use recvfrom from aiodnsresolver

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