PYTHON-1029 - Assert messageLength <= maxMessageSizeBytes

Author: behackett

pymongo/network.py

@@ -31,6 +31,7 @@
     _HAS_POLL = False
 
 from pymongo import helpers, message
+from pymongo.common import MAX_MESSAGE_SIZE
 from pymongo.errors import AutoReconnect, NotMasterError, OperationFailure
 from pymongo.read_concern import DEFAULT_READ_CONCERN
 
@@ -115,7 +116,8 @@ def command(sock, dbname, spec, slave_ok, is_mongos,
     return response_doc
 
 
-def receive_message(sock, operation, request_id):
+def receive_message(
+        sock, operation, request_id, max_message_size=MAX_MESSAGE_SIZE):
     """Receive a raw BSON message or raise socket.error."""
     header = _receive_data_on_socket(sock, 16)
     length = _UNPACK_INT(header[:4])[0]
@@ -133,6 +135,10 @@ def receive_message(sock, operation, request_id):
     assert length > 16, ("wire protocol error: message length is shorter"
                          " than standard message header: %r" % (length,))
 
+    assert length <= max_message_size, (
+        "wire protocol error: message length (%r) is larger than server max "
+        "message size (%r)" % (length, max_message_size))
+
     return _receive_data_on_socket(sock, length - 16)
 
 

pymongo/pool.py

@@ -20,6 +20,7 @@
 from bson import DEFAULT_CODEC_OPTIONS
 from bson.py3compat import u, itervalues
 from pymongo import auth, helpers, thread_util
+from pymongo.common import MAX_MESSAGE_SIZE
 from pymongo.errors import (AutoReconnect,
                             ConnectionFailure,
                             ConfigurationError,
@@ -165,7 +166,8 @@ def __init__(self, sock, pool, ismaster, address):
         self.is_writable = ismaster.is_writable if ismaster else None
         self.max_wire_version = ismaster.max_wire_version if ismaster else None
         self.max_bson_size = ismaster.max_bson_size if ismaster else None
-        self.max_message_size = ismaster.max_message_size if ismaster else None
+        self.max_message_size = (
+            ismaster.max_message_size if ismaster else MAX_MESSAGE_SIZE)
         self.max_write_batch_size = (
             ismaster.max_write_batch_size if ismaster else None)
 
@@ -238,7 +240,8 @@ def receive_message(self, operation, request_id):
         If any exception is raised, the socket is closed.
         """
         try:
-            return receive_message(self.sock, operation, request_id)
+            return receive_message(
+                self.sock, operation, request_id, self.max_message_size)
         except BaseException as error:
             self._raise_connection_failure(error)