Skip to content

Commit 740fc3c

Browse files
committed
PYTHON-1033, PYTHON-1034 - Make protocol asserts exceptions
This change introduces a new exception, pymongo.errors.ProtocolError, and uses it in place of asserts related to the wire protocol.
1 parent 1413bb0 commit 740fc3c

File tree

3 files changed

+23
-14
lines changed

3 files changed

+23
-14
lines changed

pymongo/errors.py

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,10 @@ class PyMongoError(Exception):
2626
"""Base class for all PyMongo exceptions."""
2727

2828

29+
class ProtocolError(PyMongoError):
30+
"""Raised for failures related to the wire protocol."""
31+
32+
2933
class ConnectionFailure(PyMongoError):
3034
"""Raised when a connection to the database cannot be made or is lost."""
3135

pymongo/helpers.py

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,7 @@
2828
ExecutionTimeout,
2929
NotMasterError,
3030
OperationFailure,
31+
ProtocolError,
3132
WriteError,
3233
WriteConcernError,
3334
WTimeoutError)
@@ -106,7 +107,8 @@ def _unpack_response(response, cursor_id=None, codec_options=CodecOptions()):
106107
response_flag = struct.unpack("<i", response[:4])[0]
107108
if response_flag & 1:
108109
# Shouldn't get this response if we aren't doing a getMore
109-
assert cursor_id is not None
110+
if cursor_id is None:
111+
raise ProtocolError("No cursor id for getMore operation")
110112

111113
# Fake a getMore command response. OP_GET_MORE provides no document.
112114
msg = "Cursor not found, cursor id: %d" % (cursor_id,)

pymongo/network.py

Lines changed: 16 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -32,7 +32,10 @@
3232

3333
from pymongo import helpers, message
3434
from pymongo.common import MAX_MESSAGE_SIZE
35-
from pymongo.errors import AutoReconnect, NotMasterError, OperationFailure
35+
from pymongo.errors import (AutoReconnect,
36+
NotMasterError,
37+
OperationFailure,
38+
ProtocolError)
3639
from pymongo.read_concern import DEFAULT_READ_CONCERN
3740

3841
_UNPACK_INT = struct.Struct("<i").unpack
@@ -121,21 +124,21 @@ def receive_message(
121124
length = _UNPACK_INT(header[:4])[0]
122125

123126
actual_op = _UNPACK_INT(header[12:])[0]
124-
assert operation == actual_op, ("wire protocol error: "
125-
"unknown opcode %r" % (actual_op,))
127+
if operation != actual_op:
128+
raise ProtocolError("Got opcode %r but expected "
129+
"%r" % (actual_op, operation))
126130
# No request_id for exhaust cursor "getMore".
127131
if request_id is not None:
128132
response_id = _UNPACK_INT(header[8:12])[0]
129-
assert request_id == response_id, (
130-
"wire protocol error: got response id %r but expected %r"
131-
% (response_id, request_id))
132-
133-
assert length > 16, ("wire protocol error: message length is shorter"
134-
" than standard message header: %r" % (length,))
135-
136-
assert length <= max_message_size, (
137-
"wire protocol error: message length (%r) is larger than server max "
138-
"message size (%r)" % (length, max_message_size))
133+
if request_id != response_id:
134+
raise ProtocolError("Got response id %r but expected "
135+
"%r" % (response_id, request_id))
136+
if length <= 16:
137+
raise ProtocolError("Message length (%r) not longer than standard "
138+
"message header size (16)" % (length,))
139+
if length > max_message_size:
140+
raise ProtocolError("Message length (%r) is larger than server max "
141+
"message size (%r)" % (length, max_message_size))
139142

140143
return _receive_data_on_socket(sock, length - 16)
141144

0 commit comments

Comments
 (0)