THRIFT-2757: py - Added integer limits checks in compact and JSON protocols

Client: Python
Patch: Bret Curtis

This closes: #234

3 files changed, 39 insertions, 13 deletions

diff --git a/lib/py/src/protocol/TCompactProtocol.py b/lib/py/src/protocol/TCompactProtocol.py
index 79deda8fd..7054ab0de 100644
--- a/lib/py/src/protocol/TCompactProtocol.py
+++ b/lib/py/src/protocol/TCompactProtocol.py
@@ -45,6 +45,7 @@ reader = make_helper(VALUE_READ, CONTAINER_READ)
 
 
 def makeZigZag(n, bits):
+  checkIntegerLimits(n, bits)
   return (n << 1) ^ (n >> (bits - 1))
 
 
diff --git a/lib/py/src/protocol/TJSONProtocol.py b/lib/py/src/protocol/TJSONProtocol.py
index 3048197d4..9c1877ba6 100644
--- a/lib/py/src/protocol/TJSONProtocol.py
+++ b/lib/py/src/protocol/TJSONProtocol.py
@@ -17,7 +17,8 @@
 # under the License.
 #
 
-from TProtocol import TType, TProtocolBase, TProtocolException
+from TProtocol import TType, TProtocolBase, TProtocolException, \
+    checkIntegerLimits
 import base64
 import json
 import math
@@ -449,12 +450,21 @@ class TJSONProtocol(TJSONProtocolBase):
   def writeBool(self, boolean):
     self.writeJSONNumber(1 if boolean is True else 0)
 
-  def writeInteger(self, integer):
-    self.writeJSONNumber(integer)
-  writeByte = writeInteger
-  writeI16 = writeInteger
-  writeI32 = writeInteger
-  writeI64 = writeInteger
+  def writeByte(self, byte):
+    checkIntegerLimits(byte, 8)
+    self.writeJSONNumber(byte)
+
+  def writeI16(self, i16):
+    checkIntegerLimits(i16, 16)
+    self.writeJSONNumber(i16)
+
+  def writeI32(self, i32):
+    checkIntegerLimits(i32, 32)
+    self.writeJSONNumber(i32)
+
+  def writeI64(self, i64):
+    checkIntegerLimits(i64, 64)
+    self.writeJSONNumber(i64)
 
   def writeDouble(self, dbl):
     self.writeJSONNumber(dbl)
@@ -524,12 +534,21 @@ class TSimpleJSONProtocol(TJSONProtocolBase):
     writeSetBegin = _writeCollectionBegin
     writeSetEnd = _writeCollectionEnd
 
-    def writeInteger(self, integer):
-        self.writeJSONNumber(integer)
-    writeByte = writeInteger
-    writeI16 = writeInteger
-    writeI32 = writeInteger
-    writeI64 = writeInteger
+    def writeByte(self, byte):
+        checkIntegerLimits(byte, 8)
+        self.writeJSONNumber(byte)
+
+    def writeI16(self, i16):
+        checkIntegerLimits(i16, 16)
+        self.writeJSONNumber(i16)
+
+    def writeI32(self, i32):
+        checkIntegerLimits(i32, 32)
+        self.writeJSONNumber(i32)
+
+    def writeI64(self, i64):
+        checkIntegerLimits(i64, 64)
+        self.writeJSONNumber(i64)
     
     def writeBool(self, boolean):
         self.writeJSONNumber(1 if boolean is True else 0)
diff --git a/lib/py/src/protocol/TProtocol.py b/lib/py/src/protocol/TProtocol.py
index 40e05b10d..bd69067f8 100644
--- a/lib/py/src/protocol/TProtocol.py
+++ b/lib/py/src/protocol/TProtocol.py
@@ -402,6 +402,12 @@ class TProtocolBase:
     else:
       writer(val)
 
+def checkIntegerLimits(i, bits):
+    lo = -(2 ** (bits - 1))
+    hi = 2 ** (bits - 1) - 1
+    if not lo <= i <= hi:
+        raise TProtocolException(TProtocolException.INVALID_DATA,
+                                 "i%d value: %d is outside range: [%d, %d]" % (bits, i, lo, hi))
 
 class TProtocolFactory:
   def getProtocol(self, trans):