""" Read and write chunks from files. Version 1.0.0. A file is divided into blocks, each of which has size L{BLOCK_SIZE} (except for the last block, which may be smaller). Blocks are encoded into chunks. One publishes the hash of the entire file. Clients who want to download the file first obtain the hash, then the clients can receive chunks in any order. Cryptographic hashing is used to verify each received chunk before writing to disk. Thus it is impossible to download corrupt data if one has the correct file hash. One obtains the hash of a complete file via L{CompleteChunkFile.file_hash}. One can read chunks from a complete file by the sequence operations of C{len()} and subscripting on a L{CompleteChunkFile} object. One can open an empty or partially downloaded file with L{PartialChunkFile}, and read and write chunks to this file. A chunk will fail to write if its contents and index are not consistent with the overall file hash passed to L{PartialChunkFile} when the partial chunk file was first created. The chunks have an overhead of less than 4% for files of size less than C{10**20} bytes. Benchmarks: - On a 3 GHz Pentium 3, it took 3.4 minutes to first make a L{CompleteChunkFile} object for a 4 GB file. Up to 10 MB of memory was used as the constructor ran. A metafile filename was passed to the constructor, and so the hash information was written to the metafile. The object used a negligible amount of memory after the constructor was finished. - Creation of L{CompleteChunkFile} objects in future runs of the program took negligible time, since the hash information was already stored in the metafile. @var BLOCK_SIZE: Size of a block. See L{BlockFile}. @var MAX_CHUNK_SIZE: Upper bound on the size of a chunk. See L{CompleteChunkFile}. free (adj.): unencumbered; not under the control of others Written by Connelly Barnes in 2005 and released into the public domain with no warranty of any kind, either expressed or implied. It probably won't make your computer catch on fire, or eat your children, but it might. Use at your own risk. """ import sha import os #import os.path from allmydata.util import bencode __all__ = ['CompleteChunkFile', 'PartialChunkFile'] __version__ = '1.0.0' BLOCK_SIZE = 65536 MAX_CHUNK_SIZE = BLOCK_SIZE + 4096 def hash(s): """ Cryptographic hash function used by this module. """ return sha.new(s).digest() def roundup_pow2(x): """ Round integer C{x} up to the nearest power of 2. """ ans = 1 while ans < x: ans *= 2 return ans class CompleteBinaryTreeMixin: """ Adds convenience methods to a complete binary tree. Assumes the total number of elements in the binary tree may be accessed via C{__len__}, and that each element can be retrieved using list subscripting. Tree is indexed like so:: 0 / \ 1 2 / \ / \ 3 4 5 6 / \ / \ / \ / \ 7 8 9 10 11 12 13 14 """ def parent(self, i): """ Index of the parent of C{i}. """ if i < 1 or (hasattr(self, '__len__') and i >= len(self)): raise IndexError('index out of range: ' + repr(i)) return (i - 1) // 2 def lchild(self, i): """ Index of the left child of C{i}. """ ans = 2 * i + 1 if i < 0 or (hasattr(self, '__len__') and ans >= len(self)): raise IndexError('index out of range: ' + repr(i)) return ans def rchild(self, i): """ Index of right child of C{i}. """ ans = 2 * i + 2 if i < 0 or (hasattr(self, '__len__') and ans >= len(self)): raise IndexError('index out of range: ' + repr(i)) return ans def sibling(self, i): """ Index of sibling of C{i}. """ parent = self.parent(i) if self.lchild(parent) == i: return self.rchild(parent) else: return self.lchild(parent) def needed(self, i): """ Return a list of nodes that are necessary for the hash chain. """ if i < 0 or i >= len(self): raise IndexError('index out of range: ' + repr(i)) needed = [] here = i while here != 0: needed.append(self.sibling(here)) here = self.parent(here) return needed class HashTree(CompleteBinaryTreeMixin, list): """ Compute Merkle hashes at any node in a complete binary tree. Tree is indexed like so:: 0 / \ 1 2 / \ / \ 3 4 5 6 / \ / \ / \ / \ 7 8 9 10 11 12 13 14 <- List passed to constructor. """ def __init__(self, L): """ Create complete binary tree from list of hash strings. The list is augmented by hashes so its length is a power of 2, and then this is used as the bottom row of the hash tree. The augmenting is done so that if the augmented element is at index C{i}, then its value is C{hash(bencode.bencode((i, '')))}. """ # Augment the list. start = len(L) end = roundup_pow2(len(L)) L = L + [None] * (end - start) for i in range(start, end): L[i] = hash(bencode.bencode((i, ''))) # Form each row of the tree. rows = [L] while len(rows[-1]) != 1: last = rows[-1] rows += [[hash(last[2*i] + last[2*i+1]) for i in xrange(len(last)//2)]] # Flatten the list of rows into a single list. rows.reverse() self[:] = sum(rows, []) class BlockFile: """ Reads and writes blocks of data to a binary file. It is assumed that the binary file does not change in size. @ivar file_name: Full path to file. @ivar file_size: Size of file in bytes. @ivar block_size: Size of each block. """ def __init__(self, file_name, mode, block_size, file_size=None): """ Initialize block reader or writer on given file name. If mode is 'r', the file must already exist and it is opened for reading only. If mode is 'w', the file will be created with size C{file_size} if it does not exist, and it is opened for reading and writing. Note that C{file_size} is ignored if the file already exists. """ self.mode = mode self.file_name = os.path.abspath(file_name) assert self.mode in ['r', 'w'] if mode == 'r': f = open(self.file_name, 'rb') f.close() # Create file if it doesn't exist. created = False if mode == 'w' and not os.path.exists(self.file_name): created = True buf = ' ' * 1024 f = open(self.file_name, 'wb') for i in xrange(file_size // len(buf)): f.write(buf) f.write(' ' * (file_size % len(buf))) f.close() self.file_size = os.stat(self.file_name).st_size if created: assert self.file_size == file_size self.block_size = block_size self.__block_count = self.file_size // self.block_size if self.file_size % self.block_size == 0: self.last_block_size = self.block_size else: self.last_block_size = self.file_size % self.block_size self.__block_count += 1 def __getitem__(self, i): """ Get block i. """ if i < 0 or i >= len(self): raise IndexError('block index out of range: ' + repr(i)) f = open(self.file_name, 'rb') try: f.seek(i * self.block_size) ans = f.read(self.block_size) finally: f.close() return ans def __setitem__(self, i, s): """ Set block i. """ if self.mode != 'w': raise ValueError('file opened for reading only') if i < 0 or i >= len(self): raise IndexError('block index out of range: ' + repr(i)) if i < len(self) - 1: if len(s) != self.block_size: raise ValueError('length of value must equal block_size') else: if len(s) != self.last_block_size: raise ValueError('length of value must equal last_block_size') f = open(self.file_name, 'rb+') try: f.seek(i * self.block_size) f.write(s) finally: f.close() def __len__(self): """ Get number of blocks. """ return int(self.__block_count) class MetaFile(CompleteBinaryTreeMixin): """ A L{HashTree} stored on disk, with a timestamp. The list of hashes can be accessed using subscripting and C{__len__}, in the same manner as for L{HashTree}. Note that the constructor takes the entire list associated with the L{HashTree}, not just the bottom row of the tree. @ivar meta_name: Full path to metafile. """ def __init__(self, meta_name, mode, L=None): """ Open an existing meta-file for reading or writing. If C{mode} is 'r', the meta-file must already exist and it is opened for reading only, and the list C{L} is ignored. If C{mode} is 'w', the file will be created if it does not exist (from the list of hashes given in C{L}), and it is opened for reading and writing. """ self.meta_name = os.path.abspath(meta_name) self.mode = mode assert self.mode in ['r', 'w'] # A timestamp is stored at index 0. The MetaFile instance # offsets all indices passed to __getitem__, __setitem__ by # this offset, and pretends it has length equal to # self.sublength. self.offset = 1 if self.mode == 'w': suggested_length = len(hash('')) * (len(L)+self.offset) else: suggested_length = None created = False if self.mode == 'w' and not os.path.exists(self.meta_name): created = True self.block_file = BlockFile(self.meta_name, self.mode, len(hash('')), suggested_length) self.sublength = len(self.block_file) - self.offset if created: for i in xrange(len(L)): self.block_file[i + self.offset] = L[i] def __getitem__(self, i): if i < 0 or i >= self.sublength: raise IndexError('bad meta-file block index') return self.block_file[i + self.offset] def __setitem__(self, i, value): if i < 0 or i >= self.sublength: raise IndexError('bad meta-file block index') self.block_file[i + self.offset] = value def __len__(self): return self.sublength def set_timestamp(self, file_name): """ Set meta file's timestamp equal to the timestamp for C{file_name}. """ st = os.stat(file_name) timestamp = bencode.bencode((st.st_size, st.st_mtime)) self.block_file[0] = sha.new(timestamp).digest() def check_timestamp(self, file_name): """ True if meta file's timestamp equals timestamp for C{file_name}. """ st = os.stat(file_name) timestamp = bencode.bencode((st.st_size, st.st_mtime)) return self.block_file[0] == sha.new(timestamp).digest() class CompleteChunkFile(BlockFile): """ Reads chunks from a fully-downloaded file. A chunk C{i} is created from block C{i}. Block C{i} is unencoded data read from the file by the L{BlockFile}. Chunk C{i} is an encoded string created from block C{i}. Chunks can be read using list subscripting. The total number of chunks (equals the total number of blocks) is given by L{__len__}. @ivar file_name: Full path to file. @ivar file_size: Size of file in bytes. @ivar file_hash: Hash of file. @ivar meta_name: Full path to metafile, or C{None}. @ivar tree: L{HashTree} or L{MetaFile} instance for the file. One can extract a hash from any node in the hash tree. """ def __init__(self, file_name, meta_name=None, callback=None): """ Initialize reader on the given file name. The entire file will be read and the hash will be computed from the file. This may take a long time, so C{callback()} is called frequently during this process. This allows you to reduce CPU usage if you wish. The C{meta_name} argument is optional. If it is specified, then the hashes for C{file_name} will be stored under the file C{meta_name}. If a C{CompleteChunkFile} is created on the same file and metafile in the future, then the hashes will not need to be recomputed and the constructor will return instantly. The metafile contains a file and date stamp, so that if the file stored in C{file_name} is modified, then the hashes will be recomputed. """ BlockFile.__init__(self, file_name, 'r', block_size=65536) # Whether we need to compute the hash tree compute_tree = False self.meta_name = meta_name if self.meta_name != None: self.meta_name = os.path.abspath(self.meta_name) self.meta = None if self.meta_name == None: compute_tree = True else: try: meta = MetaFile(self.meta_name, 'r') assert meta.check_timestamp(self.file_name) except (IOError, AssertionError): compute_tree = True # Compute the hash tree if needed. if compute_tree: chunk_hashes = [None] * len(self) for i in xrange(len(self)): triple = (self.file_size, i, BlockFile.__getitem__(self, i)) chunk_hashes[i] = hash(bencode.bencode(triple)) if callback: callback() self.tree = HashTree(chunk_hashes) del chunk_hashes # If a meta-file was given, make self.tree be a MetaFile instance. if self.meta_name != None: if compute_tree: # Did we compute the hash tree? Then store it to disk. self.tree = MetaFile(self.meta_name, 'w', self.tree) # Update its timestamp to be consistent with the file we # just hashed. self.tree.set_timestamp(self.file_name) else: # Read existing file from disk. self.tree = MetaFile(self.meta_name, 'r') self.file_hash = self.tree[0] def __getitem__(self, i): """ Get chunk C{i}. Raises C{ValueError} if the file's contents changed since the CompleteFileChunkReader was instantiated. """ return encode_chunk(BlockFile.__getitem__(self, i), i, self.file_size, self.tree) def encode_chunk(block, index, file_size, tree): """ Encode a chunk. Given a block at index C{index} in a file with size C{file_size}, and a L{HashTree} or L{MetaFile} instance C{tree}, computes and returns a chunk string for the given block. The C{tree} argument needs to have correct hashes only at certain indices. Check out the code for details. In any case, if a hash is wrong an exception will be raised. """ block_count = (len(tree) + 1) // 2 if index < 0 or index >= block_count: raise IndexError('block index out of range: ' + repr(index)) suffix = bencode.bencode((file_size, index, block)) current = len(tree) - block_count + index prefix = [] while current > 0: sibling = tree.sibling(current) prefix += [tree[current], tree[sibling]] current = tree.parent(current) prefix = ''.join(prefix) # Encode the chunk chunk = bencode.bencode((prefix, suffix)) # Check to make sure it decodes properly. decode_chunk(chunk, file_size, tree) return chunk def decode_chunk(chunk, file_size, tree): """ Decode a chunk. Given file with size C{file_size} and a L{HashTree} or L{MetaFile} instance C{tree}, return C{(index, block, tree_items)}. Here C{index} is the block index where string C{block} should be placed in the file. Also C{tree_items} is a dict mapping indices within the L{HashTree} or L{MetaFile} tree object associated with the given file to the corresponding hashes at those indices. These have been verified against the file's hash, so it is known that they are correct. Raises C{ValueError} if chunk verification fails. """ file_hash = tree[0] block_count = (len(tree) + 1) // 2 try: # Decode the chunk try: (prefix, suffix) = bencode.bdecode(chunk) except: raise AssertionError() assert isinstance(prefix, str) assert isinstance(suffix, str) # Verify the suffix against the hashes in the prefix. hash_len = len(hash('')) L = [prefix[hash_len*i:hash_len*(i+1)] for i in range(len(prefix)//hash_len)] L += [file_hash] assert L[0] == hash(suffix) branches = [] for i in range(0, len(L)-1, 2): if hash(L[i] + L[i+1]) == L[i+2]: branches += [0] elif hash(L[i+1] + L[i]) == L[i+2]: branches += [1] else: raise AssertionError() # Decode the suffix try: (claim_file_size, claim_index, block) = bencode.bdecode(suffix) except: raise AssertionError() assert isinstance(claim_file_size, int) or isinstance(claim_file_size, long) assert isinstance(claim_index, int) or isinstance(claim_index, long) assert isinstance(block, str) assert file_size == claim_file_size # Compute the index of the block, and check it. found_index = sum([branches[i]*2**i for i in range(len(branches))]) assert found_index == claim_index # Now fill in the tree_items dict. tree_items = {} current = (len(tree) - block_count) + found_index i = 0 while current > 0 and i + 1 < len(L): tree_items[current] = L[i] # Next item is our sibling. tree_items[tree.sibling(current)] = L[i+1] i += 2 current = tree.parent(current) return (found_index, block, tree_items) except AssertionError: raise ValueError('corrupt chunk') class PartialChunkFile(BlockFile): """ Reads and writes chunks to a partially downloaded file. @ivar file_name: Full path to file. @ivar file_size: Size of file in bytes. @ivar file_hash: Hash of file. @ivar meta_name: Full path to metafile. @ivar tree: L{MetaFile} instance for the file. The hashes in this hash tree are valid only for nodes that we have been sent hashes for. """ def __init__(self, file_name, meta_name, file_hash=None, file_size=None): """ Initialize reader/writer for the given file name and metafile name. If neither C{file_name} nor C{meta_file} exist, then both are created. The C{file_hash} and C{file_size} arguments are used to initialize the two files. If both C{file_name} and C{meta_file} exist, then the hash and file size arguments are ignored, and those values are instead read from the files. If one file exists and the other does not, an C{IOError} is raised. """ self.meta_name = os.path.abspath(meta_name) meta_exists = os.path.exists(self.meta_name) file_exists = os.path.exists(os.path.abspath(file_name)) BlockFile.__init__(self, os.path.abspath(file_name), 'w', BLOCK_SIZE, file_size) if file_exists and not meta_exists: raise IOError('metafile ' + repr(self.meta_name) + ' missing for file ' + repr(self.file_name)) if meta_exists and not file_exists: raise IOError('file ' + repr(self.file_name) + ' missing for metafile ' + repr(self.meta_name)) tree_count = 2 * roundup_pow2(len(self)) - 1 self.tree = MetaFile(self.meta_name, 'w', [hash('')] * tree_count) if not meta_exists and not file_exists: self.tree[0] = file_hash self.file_hash = self.tree[0] def __getitem__(self, i): """ Get chunk C{i}. Raises C{ValueError} if chunk has not yet been downloaded or is corrupted. """ return encode_chunk(BlockFile.__getitem__(self, i), i, self.file_size, self.tree) def __setitem__(self, i, chunk): """ Set chunk C{i}. Raises C{ValueError} if the chunk is invalid. """ (index, block, tree_items) = decode_chunk(chunk, self.file_size, self.tree) if index != i: raise ValueError('incorrect index for chunk') BlockFile.__setitem__(self, index, block) for (tree_index, tree_value) in tree_items.items(): self.tree[tree_index] = tree_value def test(filename1='temp-out', metaname1='temp-out.meta', filename2='temp-out2', metaname2='temp-out2.meta'): """ Unit tests. """ print 'Testing:' import random ntests = 100 max_file_size = 200000 # Test CompleteChunkFile. if os.path.exists(metaname1): os.remove(metaname1) for i in range(ntests): fsize = random.randrange(max_file_size) # Make some random string of size 'fsize' to go in the file. s = ''.join([sha.new(str(j)).digest() for j in range(fsize//20+1)]) assert len(s) >= fsize s = s[:fsize] f = open(filename1, 'wb') f.write(s) f.close() C = CompleteChunkFile(filename1) for j in range(len(C)): C[j] C = CompleteChunkFile(filename1, metaname1) for j in range(len(C)): C[j] C = CompleteChunkFile(filename1, metaname1) for j in range(len(C)): C[j] os.remove(metaname1) os.remove(filename1) print ' CompleteChunkFile: OK' # Test PartialChunkFile for i in range(ntests): fsize = random.randrange(max_file_size) # Make some random string of size 'fsize' to go in the file. s = ''.join([sha.new(str(j)).digest() for j in range(fsize//20+1)]) assert len(s) >= fsize s = s[:fsize] f = open(filename1, 'wb') f.write(s) f.close() C1 = CompleteChunkFile(filename1) if os.path.exists(filename2): os.remove(filename2) if os.path.exists(metaname2): os.remove(metaname2) C2 = PartialChunkFile(filename2, metaname2, C1.file_hash, C1.file_size) assert len(C1) == len(C2) assert C2.tree[0] == C1.tree[0] for j in range(len(C2)): try: C2[j] ok = False except ValueError: ok = True if not ok: raise AssertionError() for j in range(len(C2)//2): k = random.randrange(len(C2)) if len(C1) > 1: assert C1[k] != C1[(k+1)%len(C1)] try: C2[k] = C1[(k+1)%len(C1)] ok = False except ValueError: ok = True if not ok: raise AssertionError() C2[k] = C1[k] assert C2[k] == C1[k] for j in range(len(C2)): C2[j] = C1[j] assert C2[j] == C1[j] os.remove(filename1) os.remove(filename2) os.remove(metaname2) print ' PartialChunkFile: OK' if __name__ == '__main__': test()