Skip to content

Commit 0ec2e35

Browse files
jfolzjfolz
committed
fix problems associated with packing memoryviews
fix wrong length when packing multibyte memoryviews in fallback add tests for memoryviews of different types and sizes and check contents of packed data
1 parent ceb9635 commit 0ec2e35

3 files changed

Lines changed: 164 additions & 18 deletions

File tree

‎msgpack/_packer.pyx‎

Lines changed: 12 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -39,6 +39,7 @@ cdef extern from "pack.h":
3939
int msgpack_pack_ext(msgpack_packer* pk, char typecode, size_t l)
4040

4141
cdef int DEFAULT_RECURSE_LIMIT=511
42+
cdef size_t ITEM_LIMIT = (2**32)-1
4243

4344

4445
cdef class Packer(object):
@@ -178,7 +179,7 @@ cdef class Packer(object):
178179
ret = msgpack_pack_double(&self.pk, dval)
179180
elif PyBytes_CheckExact(o) if strict_types else PyBytes_Check(o):
180181
L = len(o)
181-
if L > (2**32)-1:
182+
if L > ITEM_LIMIT:
182183
raise PackValueError("bytes is too large")
183184
rawval = o
184185
ret = msgpack_pack_bin(&self.pk, L)
@@ -189,7 +190,7 @@ cdef class Packer(object):
189190
raise TypeError("Can't encode unicode string: no encoding is specified")
190191
o = PyUnicode_AsEncodedString(o, self.encoding, self.unicode_errors)
191192
L = len(o)
192-
if L > (2**32)-1:
193+
if L > ITEM_LIMIT:
193194
raise PackValueError("unicode string is too large")
194195
rawval = o
195196
ret = msgpack_pack_raw(&self.pk, L)
@@ -198,7 +199,7 @@ cdef class Packer(object):
198199
elif PyDict_CheckExact(o):
199200
d = <dict>o
200201
L = len(d)
201-
if L > (2**32)-1:
202+
if L > ITEM_LIMIT:
202203
raise PackValueError("dict is too large")
203204
ret = msgpack_pack_map(&self.pk, L)
204205
if ret == 0:
@@ -209,7 +210,7 @@ cdef class Packer(object):
209210
if ret != 0: break
210211
elif not strict_types and PyDict_Check(o):
211212
L = len(o)
212-
if L > (2**32)-1:
213+
if L > ITEM_LIMIT:
213214
raise PackValueError("dict is too large")
214215
ret = msgpack_pack_map(&self.pk, L)
215216
if ret == 0:
@@ -223,13 +224,13 @@ cdef class Packer(object):
223224
longval = o.code
224225
rawval = o.data
225226
L = len(o.data)
226-
if L > (2**32)-1:
227+
if L > ITEM_LIMIT:
227228
raise PackValueError("EXT data is too large")
228229
ret = msgpack_pack_ext(&self.pk, longval, L)
229230
ret = msgpack_pack_raw_body(&self.pk, rawval, L)
230231
elif PyList_CheckExact(o) if strict_types else (PyTuple_Check(o) or PyList_Check(o)):
231232
L = len(o)
232-
if L > (2**32)-1:
233+
if L > ITEM_LIMIT:
233234
raise PackValueError("list is too large")
234235
ret = msgpack_pack_array(&self.pk, L)
235236
if ret == 0:
@@ -240,7 +241,7 @@ cdef class Packer(object):
240241
if PyObject_GetBuffer(o, &view, PyBUF_SIMPLE) != 0:
241242
raise PackValueError("could not get buffer for memoryview")
242243
L = view.len
243-
if L > (2**32)-1:
244+
if L > ITEM_LIMIT:
244245
PyBuffer_Release(&view);
245246
raise PackValueError("memoryview is too large")
246247
ret = msgpack_pack_bin(&self.pk, L)
@@ -271,8 +272,8 @@ cdef class Packer(object):
271272
msgpack_pack_ext(&self.pk, typecode, len(data))
272273
msgpack_pack_raw_body(&self.pk, data, len(data))
273274

274-
def pack_array_header(self, long long size):
275-
if size > (2**32-1):
275+
def pack_array_header(self, size_t size):
276+
if size > ITEM_LIMIT:
276277
raise PackValueError
277278
cdef int ret = msgpack_pack_array(&self.pk, size)
278279
if ret == -1:
@@ -284,8 +285,8 @@ cdef class Packer(object):
284285
self.pk.length = 0
285286
return buf
286287

287-
def pack_map_header(self, long long size):
288-
if size > (2**32-1):
288+
def pack_map_header(self, size_t size):
289+
if size > ITEM_LIMIT:
289290
raise PackValueError
290291
cdef int ret = msgpack_pack_map(&self.pk, size)
291292
if ret == -1:

‎msgpack/fallback.py‎

Lines changed: 24 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -685,7 +685,7 @@ def _pack(self, obj, nest_limit=DEFAULT_RECURSE_LIMIT,
685685
default_used = True
686686
continue
687687
raise PackOverflowError("Integer value out of range")
688-
if self._use_bin_type and check(obj, (bytes, memoryview)):
688+
if self._use_bin_type and check(obj, bytes):
689689
n = len(obj)
690690
if n <= 0xff:
691691
self._buffer.write(struct.pack('>BB', 0xc4, n))
@@ -696,7 +696,7 @@ def _pack(self, obj, nest_limit=DEFAULT_RECURSE_LIMIT,
696696
else:
697697
raise PackValueError("Bytes is too large")
698698
return self._buffer.write(obj)
699-
if check(obj, (Unicode, bytes, memoryview)):
699+
if check(obj, (Unicode, bytes)):
700700
if check(obj, Unicode):
701701
if self._encoding is None:
702702
raise TypeError(
@@ -715,6 +715,28 @@ def _pack(self, obj, nest_limit=DEFAULT_RECURSE_LIMIT,
715715
else:
716716
raise PackValueError("String is too large")
717717
return self._buffer.write(obj)
718+
if check(obj, memoryview):
719+
n = len(obj) * obj.itemsize
720+
if self._use_bin_type:
721+
if n <= 0xff:
722+
self._buffer.write(struct.pack('>BB', 0xc4, n))
723+
elif n <= 0xffff:
724+
self._buffer.write(struct.pack(">BH", 0xc5, n))
725+
elif n <= 0xffffffff:
726+
self._buffer.write(struct.pack(">BI", 0xc6, n))
727+
else:
728+
raise PackValueError("memoryview is too large")
729+
return self._buffer.write(obj)
730+
else:
731+
if n <= 0x1f:
732+
self._buffer.write(struct.pack('B', 0xa0 + n))
733+
elif n <= 0xffff:
734+
self._buffer.write(struct.pack(">BH", 0xda, n))
735+
elif n <= 0xffffffff:
736+
self._buffer.write(struct.pack(">BI", 0xdb, n))
737+
else:
738+
raise PackValueError("memoryview is too large")
739+
return self._buffer.write(obj)
718740
if check(obj, float):
719741
if self._use_float:
720742
return self._buffer.write(struct.pack(">Bf", 0xca, obj))

‎test/test_memoryview.py‎

Lines changed: 128 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -2,11 +2,134 @@
22
# coding: utf-8
33

44

5+
from array import array
56
from msgpack import packb, unpackb
7+
import sys
68

79

8-
def test_pack_memoryview():
9-
data = bytearray(range(256))
10-
view = memoryview(data)
11-
unpacked = unpackb(packb(view))
12-
assert data == unpacked
10+
# For Python < 3:
11+
# - array type only supports old buffer interface
12+
# - array.frombytes is not available, must use deprecated array.fromstring
13+
if sys.version_info[0] < 3:
14+
def __memoryview(obj):
15+
return memoryview(buffer(obj))
16+
17+
def __make_array(f, data):
18+
a = array(f)
19+
a.fromstring(data)
20+
return a
21+
22+
def __get_data(a):
23+
return a.tostring()
24+
else:
25+
__memoryview = memoryview
26+
27+
def __make_array(f, data):
28+
a = array(f)
29+
a.frombytes(data)
30+
return a
31+
32+
def __get_data(a):
33+
return a.tobytes()
34+
35+
36+
def __run_test(format, nbytes, expected_header, expected_prefix, use_bin_type):
37+
# create a new array
38+
original_array = array(format)
39+
original_array.fromlist([255] * (nbytes // original_array.itemsize))
40+
original_data = __get_data(original_array)
41+
view = __memoryview(original_array)
42+
43+
# pack, unpack, and reconstruct array
44+
packed = packb(view, use_bin_type=use_bin_type)
45+
unpacked = unpackb(packed)
46+
reconstructed_array = __make_array(format, unpacked)
47+
48+
# check that we got the right amount of data
49+
assert len(original_data) == nbytes
50+
# check packed header
51+
assert packed[:1] == expected_header
52+
# check packed length prefix, if any
53+
assert packed[1:1+len(expected_prefix)] == expected_prefix
54+
# check packed data
55+
assert packed[1+len(expected_prefix):] == original_data
56+
# check array unpacked correctly
57+
assert original_array == reconstructed_array
58+
59+
60+
# -----------
61+
# test fixstr
62+
# -----------
63+
64+
65+
def test_memoryview_byte_fixstr():
66+
__run_test('B', 31, b'\xbf', b'', False)
67+
68+
69+
def test_memoryview_float_fixstr():
70+
__run_test('f', 28, b'\xbc', b'', False)
71+
72+
73+
# ----------
74+
# test str16
75+
# ----------
76+
77+
78+
def test_memoryview_byte_str16():
79+
__run_test('B', 2**8, b'\xda', b'\x01\x00', False)
80+
81+
82+
def test_memoryview_float_str16():
83+
__run_test('f', 2**8, b'\xda', b'\x01\x00', False)
84+
85+
86+
# ----------
87+
# test str32
88+
# ----------
89+
90+
91+
def test_memoryview_byte_str32():
92+
__run_test('B', 2**16, b'\xdb', b'\x00\x01\x00\x00', False)
93+
94+
95+
def test_memoryview_float_str32():
96+
__run_test('f', 2**16, b'\xdb', b'\x00\x01\x00\x00', False)
97+
98+
99+
# ---------
100+
# test bin8
101+
# ---------
102+
103+
104+
def test_memoryview_byte_bin8():
105+
__run_test('B', 1, b'\xc4', b'\x01', True)
106+
107+
108+
def test_memoryview_float_bin8():
109+
__run_test('f', 4, b'\xc4', b'\x04', True)
110+
111+
112+
# ----------
113+
# test bin16
114+
# ----------
115+
116+
117+
def test_memoryview_byte_bin16():
118+
__run_test('B', 2**8, b'\xc5', b'\x01\x00', True)
119+
120+
121+
def test_memoryview_float_bin16():
122+
__run_test('f', 2**8, b'\xc5', b'\x01\x00', True)
123+
124+
125+
# ----------
126+
# test bin32
127+
# ----------
128+
129+
130+
def test_memoryview_byte_bin32():
131+
__run_test('B', 2**16, b'\xc6', b'\x00\x01\x00\x00', True)
132+
133+
134+
def test_memoryview_float_bin32():
135+
__run_test('f', 2**16, b'\xc6', b'\x00\x01\x00\x00', True)

0 commit comments

Comments
 (0)