PYTHON-5021 - Fix usages of getaddrinfo to be non-blocking

Author: NoahStapp

pymongo/asynchronous/auth.py

@@ -15,6 +15,7 @@
 """Authentication helpers."""
 from __future__ import annotations
 
+import asyncio
 import functools
 import hashlib
 import hmac
@@ -177,15 +178,28 @@ def _auth_key(nonce: str, username: str, password: str) -> str:
     return md5hash.hexdigest()
 
 
-def _canonicalize_hostname(hostname: str, option: str | bool) -> str:
+async def _canonicalize_hostname(hostname: str, option: str | bool) -> str:
     """Canonicalize hostname following MIT-krb5 behavior."""
     # https://github.com/krb5/krb5/blob/d406afa363554097ac48646a29249c04f498c88e/src/util/k5test.py#L505-L520
     if option in [False, "none"]:
         return hostname
 
-    af, socktype, proto, canonname, sockaddr = socket.getaddrinfo(
-        hostname, None, 0, 0, socket.IPPROTO_TCP, socket.AI_CANONNAME
-    )[0]
+    if not _IS_SYNC:
+        loop = asyncio.get_event_loop()
+        af, socktype, proto, canonname, sockaddr = (
+            await loop.getaddrinfo(
+                hostname,
+                None,
+                family=0,
+                type=0,
+                proto=socket.IPPROTO_TCP,
+                flags=socket.AI_CANONNAME,
+            )
+        )[0]  # type: ignore[index]
+    else:
+        af, socktype, proto, canonname, sockaddr = socket.getaddrinfo(
+            hostname, None, 0, 0, socket.IPPROTO_TCP, socket.AI_CANONNAME
+        )[0]
 
     # For forward just to resolve the cname as dns.lookup() will not return it.
     if option == "forward":
@@ -213,7 +227,7 @@ async def _authenticate_gssapi(credentials: MongoCredential, conn: AsyncConnecti
         # Starting here and continuing through the while loop below - establish
         # the security context. See RFC 4752, Section 3.1, first paragraph.
         host = props.service_host or conn.address[0]
-        host = _canonicalize_hostname(host, props.canonicalize_host_name)
+        host = await _canonicalize_hostname(host, props.canonicalize_host_name)
         service = props.service_name + "@" + host
         if props.service_realm is not None:
             service = service + "@" + props.service_realm

pymongo/asynchronous/pool.py

@@ -783,7 +783,7 @@ def __repr__(self) -> str:
         )
 
 
-def _create_connection(address: _Address, options: PoolOptions) -> socket.socket:
+async def _create_connection(address: _Address, options: PoolOptions) -> socket.socket:
     """Given (host, port) and PoolOptions, connect and return a socket object.
 
     Can raise socket.error.
@@ -814,7 +814,14 @@ def _create_connection(address: _Address, options: PoolOptions) -> socket.socket
         family = socket.AF_UNSPEC
 
     err = None
-    for res in socket.getaddrinfo(host, port, family, socket.SOCK_STREAM):
+    if not _IS_SYNC:
+        loop = asyncio.get_event_loop()
+        results = await loop.getaddrinfo(  # type: ignore[assignment]
+            host, port, family=family, type=socket.SOCK_STREAM
+        )
+    else:
+        results = socket.getaddrinfo(host, port, family, socket.SOCK_STREAM)  # type: ignore[assignment]
+    for res in results:  # type: ignore[attr-defined]
         af, socktype, proto, dummy, sa = res
         # SOCK_CLOEXEC was new in CPython 3.2, and only available on a limited
         # number of platforms (newer Linux and *BSD). Starting with CPython 3.4
@@ -863,7 +870,7 @@ async def _configured_socket(
 
     Sets socket's SSL and timeout options.
     """
-    sock = _create_connection(address, options)
+    sock = await _create_connection(address, options)
     ssl_context = options._ssl_context
 
     if ssl_context is None:

pymongo/synchronous/auth.py

@@ -15,6 +15,7 @@
 """Authentication helpers."""
 from __future__ import annotations
 
+import asyncio
 import functools
 import hashlib
 import hmac
@@ -180,9 +181,22 @@ def _canonicalize_hostname(hostname: str, option: str | bool) -> str:
     if option in [False, "none"]:
         return hostname
 
-    af, socktype, proto, canonname, sockaddr = socket.getaddrinfo(
-        hostname, None, 0, 0, socket.IPPROTO_TCP, socket.AI_CANONNAME
-    )[0]
+    if not _IS_SYNC:
+        loop = asyncio.get_event_loop()
+        af, socktype, proto, canonname, sockaddr = (
+            loop.getaddrinfo(
+                hostname,
+                None,
+                family=0,
+                type=0,
+                proto=socket.IPPROTO_TCP,
+                flags=socket.AI_CANONNAME,
+            )
+        )[0]  # type: ignore[index]
+    else:
+        af, socktype, proto, canonname, sockaddr = socket.getaddrinfo(
+            hostname, None, 0, 0, socket.IPPROTO_TCP, socket.AI_CANONNAME
+        )[0]
 
     # For forward just to resolve the cname as dns.lookup() will not return it.
     if option == "forward":

pymongo/synchronous/pool.py

@@ -812,7 +812,14 @@ def _create_connection(address: _Address, options: PoolOptions) -> socket.socket
         family = socket.AF_UNSPEC
 
     err = None
-    for res in socket.getaddrinfo(host, port, family, socket.SOCK_STREAM):
+    if not _IS_SYNC:
+        loop = asyncio.get_event_loop()
+        results = loop.getaddrinfo(  # type: ignore[assignment]
+            host, port, family=family, type=socket.SOCK_STREAM
+        )
+    else:
+        results = socket.getaddrinfo(host, port, family, socket.SOCK_STREAM)  # type: ignore[assignment]
+    for res in results:  # type: ignore[attr-defined]
         af, socktype, proto, dummy, sa = res
         # SOCK_CLOEXEC was new in CPython 3.2, and only available on a limited
         # number of platforms (newer Linux and *BSD). Starting with CPython 3.4

test/asynchronous/test_auth.py

@@ -275,10 +275,10 @@ async def test_gssapi_threaded(self):
     async def test_gssapi_canonicalize_host_name(self):
         # Test the low level method.
         assert GSSAPI_HOST is not None
-        result = _canonicalize_hostname(GSSAPI_HOST, "forward")
+        result = await _canonicalize_hostname(GSSAPI_HOST, "forward")
         if "compute-1.amazonaws.com" not in result:
             self.assertEqual(result, GSSAPI_HOST)
-        result = _canonicalize_hostname(GSSAPI_HOST, "forwardAndReverse")
+        result = await _canonicalize_hostname(GSSAPI_HOST, "forwardAndReverse")
         self.assertEqual(result, GSSAPI_HOST)
 
         # Use the equivalent named CANONICALIZE_HOST_NAME.