storagewrapper.py

bittorrent source by python. please enjoy

    ## out of order compatability
    ############################################################################
    def _initalloc(self, pos, piece):
        assert self.rplaces[pos] < 0
        assert self.places[piece] == NO_PLACE
        p = self.piece_size * pos
        length = self._piecelen(pos)
        self.places[piece] = pos
        self.rplaces[pos] = piece

    def _move_piece(self, oldpos, newpos):
        assert self.rplaces[newpos] < 0
        assert self.rplaces[oldpos] >= 0
        df = self._storage_read(oldpos, self._piecelen(newpos))
        yield df
        data = df.getResult()
        df = self._storage_write(newpos, data)
        yield df
        df.getResult()
        piece = self.rplaces[oldpos]
        self.places[piece] = newpos
        self.rplaces[oldpos] = ALLOCATED
        self.rplaces[newpos] = piece
        if not self.have[piece]:
            return
        data = buffer(data, 0, self._piecelen(piece))
        if sha(data).digest() != self.hashes[piece]:
            raise BTFailure(_("data corrupted on disk - "
                              "maybe you have two copies running?"))
    ############################################################################



    def get_piece_range_for_filename(self, filename):
        begin, end = self.storage.get_byte_range_for_filename(filename)
        begin = int(begin / self.piece_size)
        end = int(end / self.piece_size)
        return begin, end

    def _waspre(self, piece, piece_len=None):
        if piece_len is None:
            piece_len = self._piecelen(piece)
        return self.storage.was_preallocated(piece * self.piece_size, piece_len)

    def _piecelen(self, piece):
        if piece < self.numpieces - 1:
            return self.piece_size
        else:
            return self.total_length - piece * self.piece_size

    def get_total_length(self):
        """Returns the total length of the torrent in bytes."""
        return self.total_length

    def get_num_pieces(self):
        """Returns the total number of pieces in this torrent."""
        return self.numpieces

    def get_amount_left(self):
        """Returns the number of bytes left to download."""
        return self.amount_left

    def do_I_have_anything(self):
        return self.amount_left < self.total_length

    def get_have_list(self):
        return self.have.tostring()

    def do_I_have(self, index):
        return self.have[index]

    def _block_piece(self, index, df):
        self.blocking_pieces[index] = df
        df.addCallback(lambda x: self.blocking_pieces.pop(index))
        return df

    def write(self, index, begin, piece, source):
        df = launch_coroutine(_wrap_task(self.add_task),
                              self._write,
                              index, begin, piece, source)
        return df

    def _write(self, index, begin, piece, source):

        if index in self.blocking_pieces:
            df = self.blocking_pieces[index]
            yield df
            df.getResult()

        if self.places[index] < 0:
            # since old versions of BT wrote out-of-order, we could
            # come across a piece which is misplaced. move it to the
            # correct place.
            if self.rplaces[index] >= 0:
                new_pos = self.rplaces[index]
                df = launch_coroutine(_wrap_task(self.add_task),
                                      self._move_piece, index, new_pos)
                yield self._block_piece(index, df)
                df.getResult()

            self._initalloc(index, index)
            
        df = self.datapig.got_piece(index, begin, piece, source)
        if df is not None:
            yield df
            df.getResult()

        df = self._storage_write(self.places[index], piece, offset=begin)
        yield df
        df.getResult()
        r = (begin, len(piece))
        self.active_requests[index].remove(r)

        hashcheck = False
        if not self.want_requests(index) and len(self.active_requests[index]) == 0:
            hashcheck = True
            df = self.hashcheck_piece(self.places[index])
            yield df
            passed = df.getResult()
            length = self._piecelen(index)
            del self.inactive_requests[index]
            del self.active_requests[index]
            self.full_pieces.remove(index)
            if passed:
                self.have[index] = True
                self.have_set.add(index)
                # In this function we perform a lookup undownloaded[filename]
                # which results in a KeyError.  --Dave HEREDAVE
                self.storage.downloaded(index * self.piece_size, length)
                self.amount_left -= length

                self.datapig.finished_piece(index)
            else: # hashcheck fail
                self.data_flunked(length, index)
                self.amount_inactive += length

                self.datapig.failed_piece(index)
        self.fastresume_dirty = True
        yield hashcheck

    def read(self, index, begin, length):
        df = launch_coroutine(_wrap_task(self.add_task),
                              self._read, index, begin, length)
        return df

    def _read(self, index, begin, length):
        if not self.have[index]:
            yield None

        if index in self.blocking_pieces:
            df = self.blocking_pieces[index]
            yield df
            df.getResult()

        if index not in self.checked_pieces:
            df = self.hashcheck_piece(self.places[index])
            yield df
            passed = df.getResult()
            if not passed:
                # TODO: this case should cause a total file hash check and
                # reconnect when done.
                raise BTFailure, _("told file complete on start-up, but piece failed hash check")
        if begin + length > self._piecelen(index):
            yield None
        df = self._storage_read(self.places[index], length, offset=begin)
        yield df
        data = df.getResult()
        yield data

    def _storage_read(self, index, amount, offset=0):
        return self.storage.read(index * self.piece_size + offset, amount)

    def _storage_write(self, index, data, offset=0):
        return self.storage.write(index * self.piece_size + offset, data)

    ## partials
    ############################################################################
    def _check_partial(self, pos, partials, data):
        index = None
        missing = False
        if self.partial_mark is None:
            self.partial_mark = ("BitTorrent - this part has not been downloaded " +
                                 "yet." + self.infohash +
                                 struct.pack('>i', self.config['download_slice_size']))
        marklen = len(self.partial_mark)+4
        for i in xrange(0, len(data) - marklen,
                        self.config['download_slice_size']):
            if data[i:i+marklen-4] == self.partial_mark:
                ind = struct.unpack('>i', data[i+marklen-4:i+marklen])[0]
                if index is None:
                    index = ind
                    parts = []
                if ind >= self.numpieces or ind != index:
                    return
                parts.append(i)
            else:
                missing = True
        if index is not None and missing:
            i += self.config['download_slice_size']
            if i < len(data):
                parts.append(i)
            partials[index] = (pos, parts)

    def _make_pending(self, index, parts):
        length = self._piecelen(index)
        x = 0
        request_size = self.config['download_slice_size']
        for x in xrange(0, length, request_size):
            partlen = min(request_size, length - x)
            if x not in parts:
                self.amount_left_with_partials -= partlen
    ############################################################################


    ## request manager stuff
    ############################################################################
    def want_requests(self, index):
        # blah, this is dumb.
        if self.have[index]:
            return False
        # if all requests are pending, we'll have a blank list
        if (index in self.inactive_requests and
            len(self.inactive_requests[index]) == 0):
            return False
        return True

    def iter_want(self):
        for index in self.have_set.iterneg(0, self.numpieces):
            # if all requests are pending, we'll have a blank list
            if (index in self.inactive_requests and
                len(self.inactive_requests[index]) == 0):
                continue
            yield index

    def new_request(self, index, full=False):
        # returns (begin, length)
        if index not in self.inactive_requests:
            self._make_inactive(index)
        rs = self.inactive_requests[index]
        if full:
            s = SparseSet()
            while rs:
                r = rs.pop()
                s.add(r[0], r[0] + r[1])
            r = s.largest_range()
            s.remove(*r)
            for b, e in s.iterrange():
                l = e - b
                rs.extend(self._break_up(b, l))
        else:
            # why min? do we want all the blocks in order?
            r = min(rs)
            rs.remove(r)
        if len(rs) == 0:
            self.full_pieces.add(index)
        self.amount_inactive -= r[1]
        assert self.amount_inactive >= 0, 'Amount inactive: %d' % self.amount_inactive
        if self.amount_inactive == 0:
            self.endgame = True
        self.active_requests[index].append(r)
        return r

    def _break_up(self, begin, length):
        l = []
        x = 0
        request_size = self.config['download_slice_size']
        while x + request_size < length:
            l.append((begin + x, request_size))
            x += request_size
        l.append((begin + x, length - x))
        return l

    def _make_inactive(self, index):
        length = self._piecelen(index)
        self.inactive_requests[index] = self._break_up(0, length)
        self.active_requests[index] = []

    def request_lost(self, index, begin, length):
        if len(self.inactive_requests[index]) == 0:
            self.full_pieces.remove(index)
        self.amount_inactive += length
        r = (begin, length)
        self.active_requests[index].remove(r)
        l = self._break_up(begin, length)
        for r in l:
            self.inactive_requests[index].append(r)
    ############################################################################