THRIFT-3532 Add configurable string and container read size limit to Python protocols
This closes #787
diff --git a/lib/py/src/protocol/TCompactProtocol.py b/lib/py/src/protocol/TCompactProtocol.py
index 6023066..3d9c0e6 100644
--- a/lib/py/src/protocol/TCompactProtocol.py
+++ b/lib/py/src/protocol/TCompactProtocol.py
@@ -126,7 +126,9 @@
TYPE_BITS = 0x07
TYPE_SHIFT_AMOUNT = 5
- def __init__(self, trans):
+ def __init__(self, trans,
+ string_length_limit=None,
+ container_length_limit=None):
TProtocolBase.__init__(self, trans)
self.state = CLEAR
self.__last_fid = 0
@@ -134,6 +136,14 @@
self.__bool_value = None
self.__structs = []
self.__containers = []
+ self.string_length_limit = string_length_limit
+ self.container_length_limit = container_length_limit
+
+ def _check_string_length(self, length):
+ self._check_length(self.string_length_limit, length)
+
+ def _check_container_length(self, length):
+ self._check_length(self.container_length_limit, length)
def __writeVarint(self, n):
writeVarint(self.trans, n)
@@ -344,6 +354,7 @@
type = self.__getTType(size_type)
if size == 15:
size = self.__readSize()
+ self._check_container_length(size)
self.__containers.append(self.state)
self.state = CONTAINER_READ
return type, size
@@ -353,6 +364,7 @@
def readMapBegin(self):
assert self.state in (VALUE_READ, CONTAINER_READ), self.state
size = self.__readSize()
+ self._check_container_length(size)
types = 0
if size > 0:
types = self.__readUByte()
@@ -391,8 +403,9 @@
return val
def __readBinary(self):
- len = self.__readSize()
- return self.trans.readAll(len)
+ size = self.__readSize()
+ self._check_string_length(size)
+ return self.trans.readAll(size)
readBinary = reader(__readBinary)
def __getTType(self, byte):
@@ -400,8 +413,13 @@
class TCompactProtocolFactory(object):
- def __init__(self):
- pass
+ def __init__(self,
+ string_length_limit=None,
+ container_length_limit=None):
+ self.string_length_limit = string_length_limit
+ self.container_length_limit = container_length_limit
def getProtocol(self, trans):
- return TCompactProtocol(trans)
+ return TCompactProtocol(trans,
+ self.string_length_limit,
+ self.container_length_limit)