Fix synchro

Author: NoahStapp

pymongo/synchronous/auth.py

@@ -174,13 +174,20 @@ def _auth_key(nonce: str, username: str, password: str) -> str:
     return md5hash.hexdigest()
 
 
-def _canonicalize_hostname(hostname: str) -> str:
+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]
 
+    # For forward just to resolve the cname as dns.lookup() will not return it.
+    if option == "forward":
+        return canonname.lower()
+
     try:
         name = socket.getnameinfo(sockaddr, socket.NI_NAMEREQD)
     except socket.gaierror:
@@ -202,9 +209,8 @@ def _authenticate_gssapi(credentials: MongoCredential, conn: Connection) -> None
         props = credentials.mechanism_properties
         # Starting here and continuing through the while loop below - establish
         # the security context. See RFC 4752, Section 3.1, first paragraph.
-        host = conn.address[0]
-        if props.canonicalize_host_name:
-            host = _canonicalize_hostname(host)
+        host = props.service_host or conn.address[0]
+        host = _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/synchronous/encryption.py

@@ -19,6 +19,7 @@
 import contextlib
 import enum
 import socket
+import time as time  # noqa: PLC0414 # needed in sync version
 import uuid
 import weakref
 from copy import deepcopy
@@ -67,7 +68,7 @@
     EncryptedCollectionError,
     EncryptionError,
     InvalidOperation,
-    PyMongoError,
+    NetworkTimeout,
     ServerSelectionTimeoutError,
 )
 from pymongo.network_layer import BLOCKING_IO_ERRORS, sendall
@@ -80,14 +81,21 @@
 from pymongo.synchronous.cursor import Cursor
 from pymongo.synchronous.database import Database
 from pymongo.synchronous.mongo_client import MongoClient
-from pymongo.synchronous.pool import _configured_socket, _raise_connection_failure
+from pymongo.synchronous.pool import (
+    _configured_socket,
+    _get_timeout_details,
+    _raise_connection_failure,
+)
 from pymongo.typings import _DocumentType, _DocumentTypeArg
 from pymongo.uri_parser import parse_host
 from pymongo.write_concern import WriteConcern
 
 if TYPE_CHECKING:
     from pymongocrypt.mongocrypt import MongoCryptKmsContext
 
+    from pymongo.pyopenssl_context import _sslConn
+    from pymongo.typings import _Address
+
 
 _IS_SYNC = True
 
@@ -103,6 +111,13 @@
 _KEY_VAULT_OPTS = CodecOptions(document_class=RawBSONDocument)
 
 
+def _connect_kms(address: _Address, opts: PoolOptions) -> Union[socket.socket, _sslConn]:
+    try:
+        return _configured_socket(address, opts)
+    except Exception as exc:
+        _raise_connection_failure(address, exc, timeout_details=_get_timeout_details(opts))
+
+
 @contextlib.contextmanager
 def _wrap_encryption_errors() -> Iterator[None]:
     """Context manager to wrap encryption related errors."""
@@ -166,18 +181,22 @@ def kms_request(self, kms_context: MongoCryptKmsContext) -> None:
                 None,  # crlfile
                 False,  # allow_invalid_certificates
                 False,  # allow_invalid_hostnames
-                False,
-            )  # disable_ocsp_endpoint_check
+                False,  # disable_ocsp_endpoint_check
+            )
         # CSOT: set timeout for socket creation.
         connect_timeout = max(_csot.clamp_remaining(_KMS_CONNECT_TIMEOUT), 0.001)
         opts = PoolOptions(
             connect_timeout=connect_timeout,
             socket_timeout=connect_timeout,
             ssl_context=ctx,
         )
-        host, port = parse_host(endpoint, _HTTPS_PORT)
+        address = parse_host(endpoint, _HTTPS_PORT)
+        sleep_u = kms_context.usleep
+        if sleep_u:
+            sleep_sec = float(sleep_u) / 1e6
+            time.sleep(sleep_sec)
         try:
-            conn = _configured_socket((host, port), opts)
+            conn = _connect_kms(address, opts)
             try:
                 sendall(conn, message)
                 while kms_context.bytes_needed > 0:
@@ -194,20 +213,29 @@ def kms_request(self, kms_context: MongoCryptKmsContext) -> None:
                     if not data:
                         raise OSError("KMS connection closed")
                     kms_context.feed(data)
-            # Async raises an OSError instead of returning empty bytes
-            except OSError as err:
-                raise OSError("KMS connection closed") from err
-            except BLOCKING_IO_ERRORS:
-                raise socket.timeout("timed out") from None
+            except MongoCryptError:
+                raise  # Propagate MongoCryptError errors directly.
+            except Exception as exc:
+                # Wrap I/O errors in PyMongo exceptions.
+                if isinstance(exc, BLOCKING_IO_ERRORS):
+                    exc = socket.timeout("timed out")
+                _raise_connection_failure(address, exc, timeout_details=_get_timeout_details(opts))
             finally:
                 conn.close()
-        except (PyMongoError, MongoCryptError):
-            raise  # Propagate pymongo errors directly.
-        except asyncio.CancelledError:
-            raise
-        except Exception as error:
-            # Wrap I/O errors in PyMongo exceptions.
-            _raise_connection_failure((host, port), error)
+        except MongoCryptError:
+            raise  # Propagate MongoCryptError errors directly.
+        except Exception as exc:
+            remaining = _csot.remaining()
+            if isinstance(exc, NetworkTimeout) or (remaining is not None and remaining <= 0):
+                raise
+            # Mark this attempt as failed and defer to libmongocrypt to retry.
+            try:
+                kms_context.fail()
+            except MongoCryptError as final_err:
+                exc = MongoCryptError(
+                    f"{final_err}, last attempt failed with: {exc}", final_err.code
+                )
+                raise exc from final_err
 
     def collection_info(self, database: str, filter: bytes) -> Optional[bytes]:
         """Get the collection info for a namespace.

pymongo/synchronous/network.py

@@ -17,6 +17,7 @@
 
 import datetime
 import logging
+import time
 from typing import (
     TYPE_CHECKING,
     Any,
@@ -30,16 +31,20 @@
 
 from bson import _decode_all_selective
 from pymongo import _csot, helpers_shared, message
-from pymongo.compression_support import _NO_COMPRESSION
+from pymongo.common import MAX_MESSAGE_SIZE
+from pymongo.compression_support import _NO_COMPRESSION, decompress
 from pymongo.errors import (
     NotPrimaryError,
     OperationFailure,
+    ProtocolError,
 )
 from pymongo.logger import _COMMAND_LOGGER, _CommandStatusMessage, _debug_log
-from pymongo.message import _OpMsg
+from pymongo.message import _UNPACK_REPLY, _OpMsg, _OpReply
 from pymongo.monitoring import _is_speculative_authenticate
 from pymongo.network_layer import (
-    receive_message,
+    _UNPACK_COMPRESSION_HEADER,
+    _UNPACK_HEADER,
+    receive_data,
     sendall,
 )
 
@@ -51,15 +56,15 @@
     from pymongo.read_preferences import _ServerMode
     from pymongo.synchronous.client_session import ClientSession
     from pymongo.synchronous.mongo_client import MongoClient
-    from pymongo.synchronous.pool import ConnectionProtocol
+    from pymongo.synchronous.pool import Connection
     from pymongo.typings import _Address, _CollationIn, _DocumentOut, _DocumentType
     from pymongo.write_concern import WriteConcern
 
 _IS_SYNC = True
 
 
 def command(
-    conn: ConnectionProtocol,
+    conn: Connection,
     dbname: str,
     spec: MutableMapping[str, Any],
     is_mongos: bool,
@@ -189,7 +194,7 @@ def command(
         )
 
     try:
-        sendall(conn, msg)
+        sendall(conn.conn, msg)
         if use_op_msg and unacknowledged:
             # Unacknowledged, fake a successful command response.
             reply = None
@@ -292,3 +297,45 @@ def command(
         )
 
     return response_doc  # type: ignore[return-value]
+
+
+def receive_message(
+    conn: Connection, request_id: Optional[int], max_message_size: int = MAX_MESSAGE_SIZE
+) -> Union[_OpReply, _OpMsg]:
+    """Receive a raw BSON message or raise socket.error."""
+    if _csot.get_timeout():
+        deadline = _csot.get_deadline()
+    else:
+        timeout = conn.conn.gettimeout()
+        if timeout:
+            deadline = time.monotonic() + timeout
+        else:
+            deadline = None
+    # Ignore the response's request id.
+    length, _, response_to, op_code = _UNPACK_HEADER(receive_data(conn, 16, deadline))
+    # No request_id for exhaust cursor "getMore".
+    if request_id is not None:
+        if request_id != response_to:
+            raise ProtocolError(f"Got response id {response_to!r} but expected {request_id!r}")
+    if length <= 16:
+        raise ProtocolError(
+            f"Message length ({length!r}) not longer than standard message header size (16)"
+        )
+    if length > max_message_size:
+        raise ProtocolError(
+            f"Message length ({length!r}) is larger than server max "
+            f"message size ({max_message_size!r})"
+        )
+    if op_code == 2012:
+        op_code, _, compressor_id = _UNPACK_COMPRESSION_HEADER(receive_data(conn, 9, deadline))
+        data = decompress(receive_data(conn, length - 25, deadline), compressor_id)
+    else:
+        data = receive_data(conn, length - 16, deadline)
+
+    try:
+        unpack_reply = _UNPACK_REPLY[op_code]
+    except KeyError:
+        raise ProtocolError(
+            f"Got opcode {op_code!r} but expected {_UNPACK_REPLY.keys()!r}"
+        ) from None
+    return unpack_reply(data)