Skip to content

Commit

Permalink
Strip protocol header from pickles where protocol > 1.
Browse files Browse the repository at this point in the history
  • Loading branch information
tseaver committed Jun 7, 2013
1 parent 9542630 commit ee4df92
Show file tree
Hide file tree
Showing 2 changed files with 42 additions and 8 deletions.
47 changes: 39 additions & 8 deletions src/ZODB/tests/testUtils.py
Expand Up @@ -41,23 +41,23 @@ class TestUtils(unittest.TestCase):
for i in range(NUM)]
all = small + large

def checkLongToStringToLong(self):
def test_LongToStringToLong(self):
for num in self.all:
s = p64(num)
n = U64(s)
self.assertEqual(num, n, "U64() failed")
n2 = u64(s)
self.assertEqual(num, n2, "u64() failed")

def checkKnownConstants(self):
def test_KnownConstants(self):
self.assertEqual(b"\000\000\000\000\000\000\000\001", p64(1))
self.assertEqual(b"\000\000\000\001\000\000\000\000", p64(1<<32))
self.assertEqual(u64(b"\000\000\000\000\000\000\000\001"), 1)
self.assertEqual(U64(b"\000\000\000\000\000\000\000\001"), 1)
self.assertEqual(u64(b"\000\000\000\001\000\000\000\000"), 1<<32)
self.assertEqual(U64(b"\000\000\000\001\000\000\000\000"), 1<<32)

def checkPersistentIdHandlesDescriptor(self):
def test_PersistentIdHandlesDescriptor(self):
from ZODB.serialize import ObjectWriter
class P(Persistent):
pass
Expand All @@ -70,7 +70,7 @@ class P(Persistent):
# deduce the class path from a pickle, instead of actually loading
# the pickle (and so also trying to import application module and
# class objects, which isn't a good idea on a ZEO server when avoidable).
def checkConflictErrorDoesntImport(self):
def test_ConflictErrorDoesntImport(self):
from ZODB.serialize import ObjectWriter
from ZODB.POSException import ConflictError
from ZODB.tests.MinPO import MinPO
Expand Down Expand Up @@ -98,9 +98,40 @@ def checkConflictErrorDoesntImport(self):
else:
self.fail("expected ConflictError, but no exception raised")

def test_get_pickle_metadata_w_protocol_0_class_pickle(self):
from ZODB.utils import get_pickle_metadata
from ZODB._compat import dumps
pickle = dumps(ExampleClass, protocol=0)
self.assertEqual(get_pickle_metadata(pickle),
(__name__, ExampleClass.__name__))

def test_get_pickle_metadata_w_protocol_1_class_pickle(self):
from ZODB.utils import get_pickle_metadata
from ZODB._compat import dumps
pickle = dumps(ExampleClass, protocol=1)
self.assertEqual(get_pickle_metadata(pickle),
(__name__, ExampleClass.__name__))

def test_get_pickle_metadata_w_protocol_2_class_pickle(self):
from ZODB.utils import get_pickle_metadata
from ZODB._compat import dumps
pickle = dumps(ExampleClass, protocol=2)
self.assertEqual(get_pickle_metadata(pickle),
(__name__, ExampleClass.__name__))

def test_get_pickle_metadata_w_protocol_3_class_pickle(self):
from ZODB.utils import get_pickle_metadata
from ZODB._compat import dumps
pickle = dumps(ExampleClass, protocol=3)
self.assertEqual(get_pickle_metadata(pickle),
(__name__, ExampleClass.__name__))


class ExampleClass(object):
pass

def test_suite():
suite = unittest.TestSuite()
suite.addTest(unittest.makeSuite(TestUtils, 'check'))
suite.addTest(doctest.DocFileSuite('../utils.txt', checker=checker))
return suite
return unittest.TestSuite((
unittest.makeSuite(TestUtils),
doctest.DocFileSuite('../utils.txt', checker=checker),
))
3 changes: 3 additions & 0 deletions src/ZODB/utils.py
Expand Up @@ -217,6 +217,9 @@ def get_pickle_metadata(data):
# ZODB's data records contain two pickles. The first is the class
# of the object, the second is the object. We're only trying to
# pick apart the first here, to extract the module and class names.
if data[0] == 0x80: # protocol marker, protocol > 1
protocol = data[1]
data = data[2:]
if data.startswith(b'(c'): # pickle MARK GLOBAL opcode sequence
global_prefix = 2
elif data.startswith(b'c'): # pickle GLOBAL opcode
Expand Down

0 comments on commit ee4df92

Please sign in to comment.