diff --git a/src/network/bmobject.py b/src/network/bmobject.py index 12b997d7..c6684f3c 100644 --- a/src/network/bmobject.py +++ b/src/network/bmobject.py @@ -36,6 +36,12 @@ class BMObjectUnwantedStreamError(Exception): errorCodes = ("Object in unwanted stream") +class BMObjectInvalidStreamError(Exception): + """Exception indicating the object is in a stream + outside of specification.""" + errorCodes = ("Object in invalid stream") + + class BMObjectInvalidError(Exception): """The object's data does not match object specification.""" errorCodes = ("Invalid object") @@ -106,6 +112,9 @@ class BMObject(object): # pylint: disable=too-many-instance-attributes 'The streamNumber %i isn\'t one we are interested in.', self.streamNumber) raise BMObjectUnwantedStreamError() + if self.streamNumber < protocol.MIN_VALID_STREAM \ + or self.streamNumber > protocol.MAX_VALID_STREAM: + raise BMObjectInvalidStreamError() def checkAlreadyHave(self): """ diff --git a/src/network/bmproto.py b/src/network/bmproto.py index 009c471f..ee43523d 100644 --- a/src/network/bmproto.py +++ b/src/network/bmproto.py @@ -22,7 +22,8 @@ from network.advanceddispatcher import AdvancedDispatcher from network.bmobject import ( BMObject, BMObjectAlreadyHaveError, BMObjectExpiredError, BMObjectInsufficientPOWError, BMObjectInvalidDataError, - BMObjectInvalidError, BMObjectUnwantedStreamError + BMObjectInvalidError, BMObjectUnwantedStreamError, + BMObjectInvalidStreamError ) from network.constants import ( ADDRESS_ALIVE, MAX_MESSAGE_SIZE, MAX_OBJECT_COUNT, @@ -409,6 +410,10 @@ class BMProto(AdvancedDispatcher, ObjectTracker): self.object.inventoryHash, acceptmismatch) if not acceptmismatch: raise + except BMObjectInvalidStreamError: + BMProto.stopDownloadingObject( + self.object.inventoryHash) + raise try: self.object.checkObjectByType() diff --git a/src/protocol.py b/src/protocol.py index 291d4ebd..4cab2dd6 100644 --- a/src/protocol.py +++ b/src/protocol.py @@ -90,6 +90,11 @@ def isBitSetWithinBitfield(fourByteString, n): x, = unpack('>L', fourByteString) return x & 2**n != 0 +# Streams + + +MIN_VALID_STREAM = 1 +MAX_VALID_STREAM = 2**63 - 1 # IP addresses