Fix gzip decompression bomb vulnerability (CWE-409)

json_tricks/decoders.py

@@ -309,14 +309,20 @@ def _bin_str_to_ndarray(data, order, shape, np_type_name, data_endianness):
 
 	assert order in [None, 'C'], 'specifying different memory order is not (yet) supported ' \
 								 'for binary numpy format (got order = {})'.format(order)
+	np_type = dtype(np_type_name)
 	if data.startswith('b64.gz:'):
 		data = standard_b64decode(data[7:])
-		data = gzip_decompress(data)
+		# the encoder writes exactly size * itemsize bytes, so anything beyond that is corrupt or hostile
+		expected_bytes = np_type.itemsize
+		for dimension in shape:
+			expected_bytes *= dimension
+		if expected_bytes < 0:
+			raise ValueError('numpy array has invalid shape {}'.format(shape))
+		data = gzip_decompress(data, max_size=expected_bytes)
 	elif data.startswith('b64:'):
 		data = standard_b64decode(data[4:])
 	else:
 		raise ValueError('found numpy array buffer, but did not understand header; supported: b64 or b64.gz')
-	np_type = dtype(np_type_name)
 	if data_endianness == sys.byteorder:
 		pass
 	if data_endianness == 'little':
@@ -347,6 +353,12 @@ def _lists_of_obj_to_ndarray(data, order, shape, dtype):
 	From nested list of objects (that aren't native numpy numbers) to ndarray.
 	"""
 	from numpy import empty, ndindex
+	# the declared shape must be backed by real data, or it alone would size the allocation
+	level = [data]
+	for size in shape:
+		if any(not isinstance(node, (list, tuple)) or len(node) != size for node in level):
+			raise ValueError('nested data does not match declared shape {}'.format(shape))
+		level = [item for node in level for item in node]
 	arr = empty(shape, dtype=dtype, order=order)
 	dec_data = data
 	for indx in ndindex(arr.shape):

json_tricks/utils.py

@@ -199,12 +199,21 @@ def gzip_compress(data, compresslevel):
 	return buf.getvalue()
 
 
-def gzip_decompress(data):
+def gzip_decompress(data, max_size=None):
 	"""
 	Do gzip decompression, without the timestamp. Just like gzip.decompress, but that's py3.2+.
+
+	:param max_size: If given, decompress at most this many bytes, and raise if the stream
+			holds more. Pass it when the size is known in advance, so that a corrupt or hostile
+			stream cannot expand without bound (CWE-409).
 	"""
-	with gzip.GzipFile(fileobj=io.BytesIO(data)) as f:
-		return f.read()
+	with gzip.GzipFile(fileobj=io.BytesIO(data)) as fh:
+		if max_size is None:
+			return fh.read()
+		result = fh.read(max_size + 1)
+	if len(result) > max_size:
+		raise ValueError('gzip stream holds more than the expected {} bytes'.format(max_size))
+	return result
 
 
 is_py3 = (version[:2] == '3.')

tests/test_np.py

@@ -1,13 +1,14 @@
 #!/usr/bin/env python
 # -*- coding: utf-8 -*-
 
+from base64 import standard_b64encode
 from copy import deepcopy
 from os.path import join
 from tempfile import mkdtemp
 import sys
 from warnings import catch_warnings, simplefilter
 
-from pytest import warns
+from pytest import raises, warns
 from numpy import arange, ones, array, array_equal, finfo, iinfo, pi
 from numpy import int8, int16, int32, int64, uint8, uint16, uint32, uint64, \
 	float16, float32, float64, complex64, complex128, zeros, ndindex
@@ -17,7 +18,7 @@
 from json_tricks import numpy_encode
 from json_tricks.np import dump, dumps, load, loads
 from json_tricks.np_utils import encode_scalars_inplace
-from json_tricks.utils import JsonTricksDeprecation, gzip_decompress
+from json_tricks.utils import JsonTricksDeprecation, gzip_compress, gzip_decompress
 from .test_bare import cls_instance
 from .test_class import MyTestCls
 
@@ -212,6 +213,25 @@ def test_dtype_object():
 	assert array_equal(back, arr)
 
 
+def test_dtype_object_shape_must_match_data():
+	# a declared shape must never size the allocation on its own
+	with raises(ValueError):
+		loads('{"__ndarray__": [], "dtype": "object", "shape": [100000000]}')
+	# only the first branch is nested deeply, so the declared 2**26 elements do not exist
+	node = 0
+	for _ in range(26):
+		node = [node, 0]
+	with raises(ValueError):
+		loads(dumps({'__ndarray__': node, 'dtype': 'object', 'shape': [2] * 26}))
+
+
+def test_compact_gzip_surplus_data_rejected():
+	# 1MB of zeros behind a shape that claims only 40
+	payload = standard_b64encode(gzip_compress(b'\x00' * (1 << 20), 9)).decode('ascii')
+	with raises(ValueError):
+		loads('{"__ndarray__": "b64.gz:%s", "dtype": "float32", "shape": [10]}' % payload)
+
+
 def test_compact_mode_unspecified():
 	# Other tests may have raised deprecation warning, so reset the cache here
 	numpy_encode._warned_compact = False