THRIFT-3532 Add configurable string and container read size limit to Python protocols
This closes #787
diff --git a/lib/py/src/protocol/TProtocol.py b/lib/py/src/protocol/TProtocol.py
index 450e0fa..d9aa2e8 100644
--- a/lib/py/src/protocol/TProtocol.py
+++ b/lib/py/src/protocol/TProtocol.py
@@ -18,6 +18,7 @@
#
from thrift.Thrift import TException, TType, TFrozenDict
+from thrift.transport.TTransport import TTransportException
from ..compat import binary_to_str, str_to_binary
import six
@@ -48,6 +49,15 @@
def __init__(self, trans):
self.trans = trans
+ @staticmethod
+ def _check_length(limit, length):
+ if length < 0:
+ raise TTransportException(TTransportException.NEGATIVE_SIZE,
+ 'Negative length: %d' % length)
+ if limit is not None and length > limit:
+ raise TTransportException(TTransportException.SIZE_LIMIT,
+ 'Length exceeded max allowed: %d' % limit)
+
def writeMessageBegin(self, name, ttype, seqid):
pass