2
0
Эх сурвалжийг харах

Merge pull request #934 from ThomasWaldmann/pass-meta

pass meta-data around, fixes #765
TW 9 жил өмнө
parent
commit
87e211f5c6

+ 31 - 23
borg/archive.py

@@ -18,7 +18,7 @@ from io import BytesIO
 from . import xattr
 from .compress import Compressor, COMPR_BUFFER
 from .constants import *  # NOQA
-from .helpers import Error, uid2user, user2uid, gid2group, group2gid, \
+from .helpers import Chunk, Error, uid2user, user2uid, gid2group, group2gid, \
     parse_timestamp, to_localtime, format_time, format_timedelta, \
     Manifest, Statistics, decode_dict, make_path_safe, StableDict, int_to_bigint, bigint_to_int, \
     ProgressIndicatorPercent, ChunkIteratorFileWrapper, remove_surrogates, log_multi, \
@@ -45,7 +45,7 @@ class DownloadPipeline:
 
     def unpack_many(self, ids, filter=None, preload=False):
         unpacker = msgpack.Unpacker(use_list=False)
-        for data in self.fetch_many(ids):
+        for _, data in self.fetch_many(ids):
             unpacker.feed(data)
             items = [decode_dict(item, ITEM_TEXT_KEYS) for item in unpacker]
             if filter:
@@ -87,7 +87,7 @@ class ChunkBuffer:
         if self.buffer.tell() == 0:
             return
         self.buffer.seek(0)
-        chunks = list(bytes(s) for s in self.chunker.chunkify(self.buffer))
+        chunks = list(Chunk(bytes(s)) for s in self.chunker.chunkify(self.buffer))
         self.buffer.seek(0)
         self.buffer.truncate(0)
         # Leave the last partial chunk in the buffer unless flush is True
@@ -109,7 +109,7 @@ class CacheChunkBuffer(ChunkBuffer):
         self.stats = stats
 
     def write_chunk(self, chunk):
-        id_, _, _ = self.cache.add_chunk(self.key.id_hash(chunk), chunk, self.stats)
+        id_, _, _ = self.cache.add_chunk(self.key.id_hash(chunk.data), chunk, self.stats)
         return id_
 
 
@@ -166,7 +166,7 @@ class Archive:
             self.zeros = b'\0' * (1 << chunker_params[1])
 
     def _load_meta(self, id):
-        data = self.key.decrypt(id, self.repository.get(id))
+        _, data = self.key.decrypt(id, self.repository.get(id))
         metadata = msgpack.unpackb(data)
         if metadata[b'version'] != 1:
             raise Exception('Unknown archive metadata version')
@@ -263,7 +263,7 @@ Number of files: {0.stats.nfiles}'''.format(
         metadata.update(additional_metadata or {})
         data = msgpack.packb(StableDict(metadata), unicode_errors='surrogateescape')
         self.id = self.key.id_hash(data)
-        self.cache.add_chunk(self.id, data, self.stats)
+        self.cache.add_chunk(self.id, Chunk(data), self.stats)
         self.manifest.archives[name] = {'id': self.id, 'time': metadata['time']}
         self.manifest.write()
         self.repository.commit()
@@ -287,7 +287,8 @@ Number of files: {0.stats.nfiles}'''.format(
         add(self.id)
         for id, chunk in zip(self.metadata[b'items'], self.repository.get_many(self.metadata[b'items'])):
             add(id)
-            unpacker.feed(self.key.decrypt(id, chunk))
+            _, data = self.key.decrypt(id, chunk)
+            unpacker.feed(data)
             for item in unpacker:
                 if b'chunks' in item:
                     stats.nfiles += 1
@@ -354,7 +355,7 @@ Number of files: {0.stats.nfiles}'''.format(
                 # Extract chunks, since the item which had the chunks was not extracted
             with open(path, 'wb') as fd:
                 ids = [c.id for c in item[b'chunks']]
-                for data in self.pipeline.fetch_many(ids, is_preloaded=True):
+                for _, data in self.pipeline.fetch_many(ids, is_preloaded=True):
                     if sparse and self.zeros.startswith(data):
                         # all-zero chunk: create a hole in a sparse file
                         fd.seek(len(data), 1)
@@ -449,7 +450,7 @@ Number of files: {0.stats.nfiles}'''.format(
         metadata[key] = value
         data = msgpack.packb(metadata, unicode_errors='surrogateescape')
         new_id = self.key.id_hash(data)
-        self.cache.add_chunk(new_id, data, self.stats)
+        self.cache.add_chunk(new_id, Chunk(data), self.stats)
         self.manifest.archives[self.name] = {'id': new_id, 'time': metadata[b'time']}
         self.cache.chunk_decref(self.id, self.stats)
         self.id = new_id
@@ -469,7 +470,8 @@ Number of files: {0.stats.nfiles}'''.format(
         for (i, (items_id, data)) in enumerate(zip(items_ids, self.repository.get_many(items_ids))):
             if progress:
                 pi.show(i)
-            unpacker.feed(self.key.decrypt(items_id, data))
+            _, data = self.key.decrypt(items_id, data)
+            unpacker.feed(data)
             self.cache.chunk_decref(items_id, stats)
             for item in unpacker:
                 if b'chunks' in item:
@@ -531,8 +533,8 @@ Number of files: {0.stats.nfiles}'''.format(
         uid, gid = 0, 0
         fd = sys.stdin.buffer  # binary
         chunks = []
-        for chunk in self.chunker.chunkify(fd):
-            chunks.append(cache.add_chunk(self.key.id_hash(chunk), chunk, self.stats))
+        for data in self.chunker.chunkify(fd):
+            chunks.append(cache.add_chunk(self.key.id_hash(data), Chunk(data), self.stats))
         self.stats.nfiles += 1
         t = int_to_bigint(int(time.time()) * 1000000000)
         item = {
@@ -588,8 +590,8 @@ Number of files: {0.stats.nfiles}'''.format(
             fh = Archive._open_rb(path)
             with os.fdopen(fh, 'rb') as fd:
                 chunks = []
-                for chunk in self.chunker.chunkify(fd, fh):
-                    chunks.append(cache.add_chunk(self.key.id_hash(chunk), chunk, self.stats))
+                for data in self.chunker.chunkify(fd, fh):
+                    chunks.append(cache.add_chunk(self.key.id_hash(data), Chunk(data), self.stats))
                     if self.show_progress:
                         self.stats.show_progress(item=item, dt=0.2)
             cache.memorize_file(path_hash, st, [c.id for c in chunks])
@@ -735,7 +737,7 @@ class ArchiveChecker:
         manifest = Manifest(self.key, self.repository)
         for chunk_id, _ in self.chunks.iteritems():
             cdata = self.repository.get(chunk_id)
-            data = self.key.decrypt(chunk_id, cdata)
+            _, data = self.key.decrypt(chunk_id, cdata)
             # Some basic sanity checks of the payload before feeding it into msgpack
             if len(data) < 2 or ((data[0] & 0xf0) != 0x80) or ((data[1] & 0xe0) != 0xa0):
                 continue
@@ -766,9 +768,9 @@ class ArchiveChecker:
                 self.possibly_superseded.add(id_)
 
         def add_callback(chunk):
-            id_ = self.key.id_hash(chunk)
+            id_ = self.key.id_hash(chunk.data)
             cdata = self.key.encrypt(chunk)
-            add_reference(id_, len(chunk), len(cdata), cdata)
+            add_reference(id_, len(chunk.data), len(cdata), cdata)
             return id_
 
         def add_reference(id_, size, csize, cdata=None):
@@ -795,7 +797,7 @@ class ArchiveChecker:
                     self.error_found = True
                     data = bytes(size)
                     chunk_id = self.key.id_hash(data)
-                    cdata = self.key.encrypt(data)
+                    cdata = self.key.encrypt(Chunk(data))
                     csize = len(cdata)
                     add_reference(chunk_id, size, csize, cdata)
                 else:
@@ -835,7 +837,8 @@ class ArchiveChecker:
                 if state > 0:
                     unpacker.resync()
                 for chunk_id, cdata in zip(items, repository.get_many(items)):
-                    unpacker.feed(self.key.decrypt(chunk_id, cdata))
+                    _, data = self.key.decrypt(chunk_id, cdata)
+                    unpacker.feed(data)
                     try:
                         for item in unpacker:
                             if isinstance(item, dict):
@@ -872,7 +875,7 @@ class ArchiveChecker:
                     continue
                 mark_as_possibly_superseded(archive_id)
                 cdata = self.repository.get(archive_id)
-                data = self.key.decrypt(archive_id, cdata)
+                _, data = self.key.decrypt(archive_id, cdata)
                 archive = StableDict(msgpack.unpackb(data))
                 if archive[b'version'] != 1:
                     raise Exception('Unknown archive metadata version')
@@ -890,7 +893,7 @@ class ArchiveChecker:
                 archive[b'items'] = items_buffer.chunks
                 data = msgpack.packb(archive, unicode_errors='surrogateescape')
                 new_archive_id = self.key.id_hash(data)
-                cdata = self.key.encrypt(data)
+                cdata = self.key.encrypt(Chunk(data))
                 add_reference(new_archive_id, len(data), len(cdata), cdata)
                 info[b'id'] = new_archive_id
 
@@ -1045,7 +1048,7 @@ class ArchiveRecreater:
         chunk_iterator = self.create_chunk_iterator(archive, target, item)
         consume(chunk_iterator, len(new_chunks))
         for chunk in chunk_iterator:
-            chunk_id = self.key.id_hash(chunk)
+            chunk_id = self.key.id_hash(chunk.data)
             if chunk_id in self.seen_chunks:
                 new_chunks.append(self.cache.chunk_incref(chunk_id, target.stats))
             else:
@@ -1076,7 +1079,12 @@ class ArchiveRecreater:
             # The target.chunker will read the file contents through ChunkIteratorFileWrapper chunk-by-chunk
             # (does not load the entire file into memory)
             file = ChunkIteratorFileWrapper(chunk_iterator)
-            chunk_iterator = target.chunker.chunkify(file)
+
+            def _chunk_iterator():
+                for data in target.chunker.chunkify(file):
+                    yield Chunk(data)
+
+            chunk_iterator = _chunk_iterator()
         return chunk_iterator
 
     def process_partial_chunks(self, target):

+ 3 - 3
borg/archiver.py

@@ -121,14 +121,14 @@ class Archiver:
                 a = next(chunks1, end)
                 if a is end:
                     return not blen - bi and next(chunks2, end) is end
-                a = memoryview(a)
+                a = memoryview(a.data)
                 alen = len(a)
                 ai = 0
             if not blen - bi:
                 b = next(chunks2, end)
                 if b is end:
                     return not alen - ai and next(chunks1, end) is end
-                b = memoryview(b)
+                b = memoryview(b.data)
                 blen = len(b)
                 bi = 0
             slicelen = min(alen - ai, blen - bi)
@@ -868,7 +868,7 @@ class Archiver:
         """dump (decrypted, decompressed) archive items metadata (not: data)"""
         archive = Archive(repository, key, manifest, args.location.archive)
         for i, item_id in enumerate(archive.metadata[b'items']):
-            data = key.decrypt(item_id, repository.get(item_id))
+            _, data = key.decrypt(item_id, repository.get(item_id))
             filename = '%06d_%s.items' % (i, hexlify(item_id).decode('ascii'))
             print('Dumping', filename)
             with open(filename, 'wb') as fd:

+ 5 - 5
borg/cache.py

@@ -279,7 +279,7 @@ Chunk index:    {0.total_unique_chunks:20d} {0.total_chunks:20d}"""
         def fetch_and_build_idx(archive_id, repository, key):
             chunk_idx = ChunkIndex()
             cdata = repository.get(archive_id)
-            data = key.decrypt(archive_id, cdata)
+            _, data = key.decrypt(archive_id, cdata)
             chunk_idx.add(archive_id, 1, len(data), len(cdata))
             archive = msgpack.unpackb(data)
             if archive[b'version'] != 1:
@@ -287,7 +287,7 @@ Chunk index:    {0.total_unique_chunks:20d} {0.total_chunks:20d}"""
             decode_dict(archive, (b'name',))
             unpacker = msgpack.Unpacker()
             for item_id, chunk in zip(archive[b'items'], repository.get_many(archive[b'items'])):
-                data = key.decrypt(item_id, chunk)
+                _, data = key.decrypt(item_id, chunk)
                 chunk_idx.add(item_id, 1, len(data), len(chunk))
                 unpacker.feed(data)
                 for item in unpacker:
@@ -368,14 +368,14 @@ Chunk index:    {0.total_unique_chunks:20d} {0.total_chunks:20d}"""
             self.do_cache = os.path.isdir(archive_path)
             self.chunks = create_master_idx(self.chunks)
 
-    def add_chunk(self, id, data, stats, overwrite=False):
+    def add_chunk(self, id, chunk, stats, overwrite=False):
         if not self.txn_active:
             self.begin_txn()
-        size = len(data)
+        size = len(chunk.data)
         refcount = self.seen_chunk(id, size)
         if refcount and not overwrite:
             return self.chunk_incref(id, stats)
-        data = self.key.encrypt(data)
+        data = self.key.encrypt(chunk)
         csize = len(data)
         self.repository.put(id, data, wait=False)
         self.chunks.add(id, 1, size, csize)

+ 3 - 3
borg/fuse.py

@@ -73,7 +73,7 @@ class FuseOperations(llfuse.Operations):
         """
         unpacker = msgpack.Unpacker()
         for key, chunk in zip(archive.metadata[b'items'], self.repository.get_many(archive.metadata[b'items'])):
-            data = self.key.decrypt(key, chunk)
+            _, data = self.key.decrypt(key, chunk)
             unpacker.feed(data)
             for item in unpacker:
                 segments = prefix + os.fsencode(os.path.normpath(item[b'path'])).split(b'/')
@@ -229,8 +229,8 @@ class FuseOperations(llfuse.Operations):
                 offset -= s
                 continue
             n = min(size, s - offset)
-            chunk = self.key.decrypt(id, self.repository.get(id))
-            parts.append(chunk[offset:offset + n])
+            _, data = self.key.decrypt(id, self.repository.get(id))
+            parts.append(data[offset:offset + n])
             offset = 0
             size -= n
             if not size:

+ 14 - 5
borg/helpers.py

@@ -37,6 +37,14 @@ import msgpack.fallback
 import socket
 
 
+# meta dict, data bytes
+_Chunk = namedtuple('_Chunk', 'meta data')
+
+
+def Chunk(data, **meta):
+    return _Chunk(meta, data)
+
+
 class Error(Exception):
     """Error base class"""
 
@@ -93,7 +101,7 @@ class Manifest:
         if not key:
             key = key_factory(repository, cdata)
         manifest = cls(key, repository)
-        data = key.decrypt(None, cdata)
+        _, data = key.decrypt(None, cdata)
         manifest.id = key.id_hash(data)
         m = msgpack.unpackb(data)
         if not m.get(b'version') == 1:
@@ -114,7 +122,7 @@ class Manifest:
             'config': self.config,
         }))
         self.id = self.key.id_hash(data)
-        self.repository.put(self.MANIFEST_ID, self.key.encrypt(data))
+        self.repository.put(self.MANIFEST_ID, self.key.encrypt(Chunk(data)))
 
     def list_archive_infos(self, sort_by=None, reverse=False):
         # inexpensive Archive.list_archives replacement if we just need .name, .id, .ts
@@ -1260,8 +1268,8 @@ class ItemFormatter:
         if b'chunks' not in item:
             return ""
         hash = hashlib.new(hash_function)
-        for chunk in self.archive.pipeline.fetch_many([c.id for c in item[b'chunks']]):
-            hash.update(chunk)
+        for _, data in self.archive.pipeline.fetch_many([c.id for c in item[b'chunks']]):
+            hash.update(data)
         return hash.hexdigest()
 
     def format_time(self, key, item):
@@ -1284,7 +1292,8 @@ class ChunkIteratorFileWrapper:
         remaining = len(self.chunk) - self.chunk_offset
         if not remaining:
             try:
-                self.chunk = memoryview(next(self.chunk_iterator))
+                chunk = next(self.chunk_iterator)
+                self.chunk = memoryview(chunk.data)
             except StopIteration:
                 self.exhausted = True
                 return 0  # EOF

+ 8 - 7
borg/key.py

@@ -7,7 +7,7 @@ import textwrap
 from hmac import compare_digest
 from hashlib import sha256, pbkdf2_hmac
 
-from .helpers import IntegrityError, get_keys_dir, Error, yes
+from .helpers import Chunk, IntegrityError, get_keys_dir, Error, yes
 from .logger import create_logger
 logger = create_logger()
 
@@ -77,7 +77,7 @@ class KeyBase:
         """Return HMAC hash using the "id" HMAC key
         """
 
-    def encrypt(self, data):
+    def encrypt(self, chunk):
         pass
 
     def decrypt(self, id, data):
@@ -101,7 +101,8 @@ class PlaintextKey(KeyBase):
     def id_hash(self, data):
         return sha256(data).digest()
 
-    def encrypt(self, data):
+    def encrypt(self, chunk):
+        meta, data = chunk
         return b''.join([self.TYPE_STR, self.compressor.compress(data)])
 
     def decrypt(self, id, data):
@@ -110,7 +111,7 @@ class PlaintextKey(KeyBase):
         data = self.compressor.decompress(memoryview(data)[1:])
         if id and sha256(data).digest() != id:
             raise IntegrityError('Chunk id verification failed')
-        return data
+        return Chunk(data)
 
 
 class AESKeyBase(KeyBase):
@@ -133,8 +134,8 @@ class AESKeyBase(KeyBase):
         """
         return hmac_sha256(self.id_key, data)
 
-    def encrypt(self, data):
-        data = self.compressor.compress(data)
+    def encrypt(self, chunk):
+        data = self.compressor.compress(chunk.data)
         self.enc_cipher.reset()
         data = b''.join((self.enc_cipher.iv[8:], self.enc_cipher.encrypt(data)))
         hmac = hmac_sha256(self.enc_hmac_key, data)
@@ -156,7 +157,7 @@ class AESKeyBase(KeyBase):
             hmac_computed = hmac_sha256(self.id_key, data)
             if not compare_digest(hmac_computed, hmac_given):
                 raise IntegrityError('Chunk id verification failed')
-        return data
+        return Chunk(data)
 
     def extract_nonce(self, payload):
         if not (payload[0] == self.TYPE or

+ 3 - 3
borg/testsuite/archive.py

@@ -14,9 +14,9 @@ class MockCache:
     def __init__(self):
         self.objects = {}
 
-    def add_chunk(self, id, data, stats=None):
-        self.objects[id] = data
-        return id, len(data), len(data)
+    def add_chunk(self, id, chunk, stats=None):
+        self.objects[id] = chunk.data
+        return id, len(chunk.data), len(chunk.data)
 
 
 class ArchiveTimestampTestCase(BaseTestCase):

+ 5 - 3
borg/testsuite/archiver.py

@@ -23,7 +23,7 @@ from ..archiver import Archiver
 from ..cache import Cache
 from ..constants import *  # NOQA
 from ..crypto import bytes_to_long, num_aes_blocks
-from ..helpers import Manifest
+from ..helpers import Chunk, Manifest, EXIT_SUCCESS, EXIT_WARNING, EXIT_ERROR
 from ..key import KeyfileKeyBase
 from ..remote import RemoteRepository, PathNotAllowed
 from ..repository import Repository
@@ -1677,8 +1677,10 @@ def test_get_args():
 
 def test_compare_chunk_contents():
     def ccc(a, b):
-        compare1 = Archiver.compare_chunk_contents(iter(a), iter(b))
-        compare2 = Archiver.compare_chunk_contents(iter(b), iter(a))
+        chunks_a = [Chunk(data) for data in a]
+        chunks_b = [Chunk(data) for data in b]
+        compare1 = Archiver.compare_chunk_contents(iter(chunks_a), iter(chunks_b))
+        compare2 = Archiver.compare_chunk_contents(iter(chunks_b), iter(chunks_a))
         assert compare1 == compare2
         return compare1
     assert ccc([

+ 2 - 2
borg/testsuite/helpers.py

@@ -13,7 +13,7 @@ import time
 from ..helpers import Location, format_file_size, format_timedelta, make_path_safe, \
     prune_within, prune_split, get_cache_dir, get_keys_dir, Statistics, is_slow_msgpack, \
     yes, TRUISH, FALSISH, DEFAULTISH, \
-    StableDict, int_to_bigint, bigint_to_int, parse_timestamp, CompressionSpec, ChunkerParams, \
+    StableDict, int_to_bigint, bigint_to_int, parse_timestamp, CompressionSpec, ChunkerParams, Chunk, \
     ProgressIndicatorPercent, ProgressIndicatorEndless, load_excludes, parse_pattern, \
     PatternMatcher, RegexPattern, PathPrefixPattern, FnmatchPattern, ShellPattern, partial_format, ChunkIteratorFileWrapper
 from . import BaseTestCase, environment_variable, FakeInputs
@@ -902,7 +902,7 @@ def test_partial_format():
 
 
 def test_chunk_file_wrapper():
-    cfw = ChunkIteratorFileWrapper(iter([b'abc', b'def']))
+    cfw = ChunkIteratorFileWrapper(iter([Chunk(b'abc'), Chunk(b'def')]))
     assert cfw.read(2) == b'ab'
     assert cfw.read(50) == b'cdef'
     assert cfw.exhausted

+ 14 - 14
borg/testsuite/key.py

@@ -6,7 +6,7 @@ from binascii import hexlify, unhexlify
 
 from ..crypto import bytes_to_long, num_aes_blocks
 from ..key import PlaintextKey, PassphraseKey, KeyfileKey
-from ..helpers import Location
+from ..helpers import Location, Chunk
 from . import BaseTestCase
 
 
@@ -47,17 +47,17 @@ class KeyTestCase(BaseTestCase):
 
     def test_plaintext(self):
         key = PlaintextKey.create(None, None)
-        data = b'foo'
-        self.assert_equal(hexlify(key.id_hash(data)), b'2c26b46b68ffc68ff99b453c1d30413413422d706483bfa0f98a5e886266e7ae')
-        self.assert_equal(data, key.decrypt(key.id_hash(data), key.encrypt(data)))
+        chunk = Chunk(b'foo')
+        self.assert_equal(hexlify(key.id_hash(chunk.data)), b'2c26b46b68ffc68ff99b453c1d30413413422d706483bfa0f98a5e886266e7ae')
+        self.assert_equal(chunk, key.decrypt(key.id_hash(chunk.data), key.encrypt(chunk)))
 
     def test_keyfile(self):
         os.environ['BORG_PASSPHRASE'] = 'test'
         key = KeyfileKey.create(self.MockRepository(), self.MockArgs())
         self.assert_equal(bytes_to_long(key.enc_cipher.iv, 8), 0)
-        manifest = key.encrypt(b'XXX')
+        manifest = key.encrypt(Chunk(b'XXX'))
         self.assert_equal(key.extract_nonce(manifest), 0)
-        manifest2 = key.encrypt(b'XXX')
+        manifest2 = key.encrypt(Chunk(b'XXX'))
         self.assert_not_equal(manifest, manifest2)
         self.assert_equal(key.decrypt(None, manifest), key.decrypt(None, manifest2))
         self.assert_equal(key.extract_nonce(manifest2), 1)
@@ -67,15 +67,15 @@ class KeyTestCase(BaseTestCase):
         # Key data sanity check
         self.assert_equal(len(set([key2.id_key, key2.enc_key, key2.enc_hmac_key])), 3)
         self.assert_equal(key2.chunk_seed == 0, False)
-        data = b'foo'
-        self.assert_equal(data, key2.decrypt(key.id_hash(data), key.encrypt(data)))
+        chunk = Chunk(b'foo')
+        self.assert_equal(chunk, key2.decrypt(key.id_hash(chunk.data), key.encrypt(chunk)))
 
     def test_keyfile2(self):
         with open(os.path.join(os.environ['BORG_KEYS_DIR'], 'keyfile'), 'w') as fd:
             fd.write(self.keyfile2_key_file)
         os.environ['BORG_PASSPHRASE'] = 'passphrase'
         key = KeyfileKey.detect(self.MockRepository(), self.keyfile2_cdata)
-        self.assert_equal(key.decrypt(self.keyfile2_id, self.keyfile2_cdata), b'payload')
+        self.assert_equal(key.decrypt(self.keyfile2_id, self.keyfile2_cdata).data, b'payload')
 
     def test_passphrase(self):
         os.environ['BORG_PASSPHRASE'] = 'test'
@@ -85,9 +85,9 @@ class KeyTestCase(BaseTestCase):
         self.assert_equal(hexlify(key.enc_hmac_key), b'b885a05d329a086627412a6142aaeb9f6c54ab7950f996dd65587251f6bc0901')
         self.assert_equal(hexlify(key.enc_key), b'2ff3654c6daf7381dbbe718d2b20b4f1ea1e34caa6cc65f6bb3ac376b93fed2a')
         self.assert_equal(key.chunk_seed, -775740477)
-        manifest = key.encrypt(b'XXX')
+        manifest = key.encrypt(Chunk(b'XXX'))
         self.assert_equal(key.extract_nonce(manifest), 0)
-        manifest2 = key.encrypt(b'XXX')
+        manifest2 = key.encrypt(Chunk(b'XXX'))
         self.assert_not_equal(manifest, manifest2)
         self.assert_equal(key.decrypt(None, manifest), key.decrypt(None, manifest2))
         self.assert_equal(key.extract_nonce(manifest2), 1)
@@ -98,6 +98,6 @@ class KeyTestCase(BaseTestCase):
         self.assert_equal(key.enc_hmac_key, key2.enc_hmac_key)
         self.assert_equal(key.enc_key, key2.enc_key)
         self.assert_equal(key.chunk_seed, key2.chunk_seed)
-        data = b'foo'
-        self.assert_equal(hexlify(key.id_hash(data)), b'818217cf07d37efad3860766dcdf1d21e401650fed2d76ed1d797d3aae925990')
-        self.assert_equal(data, key2.decrypt(key2.id_hash(data), key.encrypt(data)))
+        chunk = Chunk(b'foo')
+        self.assert_equal(hexlify(key.id_hash(chunk.data)), b'818217cf07d37efad3860766dcdf1d21e401650fed2d76ed1d797d3aae925990')
+        self.assert_equal(chunk, key2.decrypt(key2.id_hash(chunk.data), key.encrypt(chunk)))