path: root/BitTorrent/StorageWrapper.py

# The contents of this file are subject to the BitTorrent Open Source License
# Version 1.1 (the License).  You may not copy or use this file, in either
# source code or executable form, except in compliance with the License.  You
# may obtain a copy of the License at http://www.bittorrent.com/license/.
#
# Software distributed under the License is distributed on an AS IS basis,
# WITHOUT WARRANTY OF ANY KIND, either express or implied.  See the License
# for the specific language governing rights and limitations under the
# License.

# Written by Bram Cohen

from __future__ import division

from sha import sha
from array import array
from binascii import b2a_hex

from BitTorrent.bitfield import Bitfield
from BitTorrent import BTFailure, INFO, WARNING, ERROR, CRITICAL

def toint(s):
    return int(b2a_hex(s), 16)

def tobinary(i):
    return (chr(i >> 24) + chr((i >> 16) & 0xFF) +
        chr((i >> 8) & 0xFF) + chr(i & 0xFF))

NO_PLACE = -1

ALLOCATED = -1
UNALLOCATED = -2
FASTRESUME_PARTIAL = -3

class StorageWrapper(object):

    def __init__(self, storage, config, hashes, piece_size, finished,
            statusfunc, flag, data_flunked, infohash, errorfunc, resumefile):
        self.numpieces = len(hashes)
        self.storage = storage
        self.config = config
        check_hashes = config['check_hashes']
        self.hashes = hashes
        self.piece_size = piece_size
        self.data_flunked = data_flunked
        self.errorfunc = errorfunc
        self.total_length = storage.get_total_length()
        self.amount_left = self.total_length
        self.partial_mark = "BitTorrent - this part has not been "+\
                            "downloaded yet."+infohash+\
                            tobinary(config['download_slice_size'])
        if self.total_length <= piece_size * (self.numpieces - 1):
            raise BTFailure, _("bad data in responsefile - total too small")
        if self.total_length > piece_size * self.numpieces:
            raise BTFailure, _("bad data in responsefile - total too big")
        self.finished = finished
        self.numactive = array('H', [0] * self.numpieces)
        self.inactive_requests = [1] * self.numpieces
        self.amount_inactive = self.total_length
        self.endgame = False
        self.have = Bitfield(self.numpieces)
        self.waschecked = Bitfield(self.numpieces)
        if self.numpieces < 32768:
            typecode = 'h'
        else:
            typecode = 'l'
        self.places = array(typecode, [NO_PLACE] * self.numpieces)
        if not check_hashes:
            self.rplaces = array(typecode, range(self.numpieces))
            fastresume = True
        else:
            self.rplaces = self._load_fastresume(resumefile, typecode)
            if self.rplaces is not None:
                fastresume = True
            else:
                self.rplaces = array(typecode, [UNALLOCATED] * self.numpieces)
                fastresume = False
        self.holepos = 0
        self.stat_numfound = 0
        self.stat_numflunked = 0
        self.stat_numdownloaded = 0
        self.stat_active = {}
        self.stat_new = {}
        self.stat_dirty = {}
        self.download_history = {}
        self.failed_pieces = {}

        if self.numpieces == 0:
            return
        targets = {}
        total = 0
        if not fastresume:
            for i in xrange(self.numpieces):
                if self._waspre(i):
                    self.rplaces[i] = ALLOCATED
                    total += 1
                else:
                    targets[hashes[i]] = i
        if total and check_hashes:
            statusfunc(_("checking existing file"), 0)
        def markgot(piece, pos):
            if self.have[piece]:
                if piece != pos:
                    return
                self.rplaces[self.places[pos]] = ALLOCATED
                self.places[pos] = self.rplaces[pos] = pos
                return
            self.places[piece] = pos
            self.rplaces[pos] = piece
            self.have[piece] = True
            self.amount_left -= self._piecelen(piece)
            self.amount_inactive -= self._piecelen(piece)
            self.inactive_requests[piece] = None
            if not fastresume:
                self.waschecked[piece] = True
            self.stat_numfound += 1
        lastlen = self._piecelen(self.numpieces - 1)
        partials = {}
        for i in xrange(self.numpieces):
            if not self._waspre(i):
                if self.rplaces[i] != UNALLOCATED:
                    raise BTFailure(_("--check_hashes 0 or fastresume info "
                                      "doesn't match file state (missing data)"))
                continue
            elif fastresume:
                t = self.rplaces[i]
                if t >= 0:
                    markgot(t, i)
                    continue
                if t == UNALLOCATED:
                    raise BTFailure(_("Bad fastresume info (files contain more "
                                      "data)"))
                if t == ALLOCATED:
                    continue
                if t!= FASTRESUME_PARTIAL:
                    raise BTFailure(_("Bad fastresume info (illegal value)"))
                data = self.storage.read(self.piece_size * i,
                                         self._piecelen(i))
                self._check_partial(i, partials, data)
                self.rplaces[i] = ALLOCATED
            else:
                data = self.storage.read(piece_size * i, self._piecelen(i))
                sh = sha(buffer(data, 0, lastlen))
                sp = sh.digest()
                sh.update(buffer(data, lastlen))
                s = sh.digest()
                if s == hashes[i]:
                    markgot(i, i)
                elif s in targets and self._piecelen(i) == self._piecelen(targets[s]):
                    markgot(targets[s], i)
                elif not self.have[self.numpieces - 1] and sp == hashes[-1] and (i == self.numpieces - 1 or not self._waspre(self.numpieces - 1)):
                    markgot(self.numpieces - 1, i)
                else:
                    self._check_partial(i, partials, data)
                statusfunc(fractionDone = 1 - self.amount_left /
                           self.total_length)
            if flag.isSet():
                return
        self.amount_left_with_partials = self.amount_left
        for piece in partials:
            if self.places[piece] < 0:
                pos = partials[piece][0]
                self.places[piece] = pos
                self.rplaces[pos] = piece
                self._make_partial(piece, partials[piece][1])
        for i in xrange(self.numpieces):
            if self.rplaces[i] != UNALLOCATED:
                self.storage.allocated(piece_size * i, self._piecelen(i))
            if self.have[i]:
                self.storage.downloaded(piece_size * i, self._piecelen(i))

    def _waspre(self, piece):
        return self.storage.was_preallocated(piece * self.piece_size, self._piecelen(piece))

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

    def _check_partial(self, pos, partials, data):
        index = None
        missing = False
        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 = toint(data[i+marklen-4:i+marklen])
                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_partial(self, index, parts):
        length = self._piecelen(index)
        l = []
        self.inactive_requests[index] = l
        x = 0
        self.amount_left_with_partials -= self._piecelen(index)
        self.download_history[index] = {}
        request_size = self.config['download_slice_size']
        for x in xrange(0, self._piecelen(index), request_size):
            partlen = min(request_size, length - x)
            if x in parts:
                l.append((x, partlen))
                self.amount_left_with_partials += partlen
            else:
                self.amount_inactive -= partlen
                self.download_history[index][x] = None
        self.stat_dirty[index] = 1

    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)
        if self.rplaces[pos] == UNALLOCATED:
            self.storage.allocated(p, length)
        self.places[piece] = pos
        self.rplaces[pos] = piece
        # "if self.rplaces[pos] != ALLOCATED:" to skip extra mark writes
        mark = self.partial_mark + tobinary(piece)
        mark += chr(0xff) * (self.config['download_slice_size'] - len(mark))
        mark *= (length - 1) // len(mark) + 1
        self.storage.write(p, buffer(mark, 0, length))

    def _move_piece(self, oldpos, newpos):
        assert self.rplaces[newpos] < 0
        assert self.rplaces[oldpos] >= 0
        data = self.storage.read(self.piece_size * oldpos,
                                 self._piecelen(newpos))
        self.storage.write(self.piece_size * newpos, data)
        if self.rplaces[newpos] == UNALLOCATED:
            self.storage.allocated(self.piece_size * newpos, len(data))
        piece = self.rplaces[oldpos]
        self.places[piece] = newpos
        self.rplaces[oldpos] = ALLOCATED
        self.rplaces[newpos] = piece
        if not self.have[piece]:
            return
        data = data[:self._piecelen(piece)]
        if sha(data).digest() != self.hashes[piece]:
            raise BTFailure(_("data corrupted on disk - "
                              "maybe you have two copies running?"))

    def _get_free_place(self):
        while self.rplaces[self.holepos] >= 0:
            self.holepos += 1
        return self.holepos

    def get_amount_left(self):
        return self.amount_left

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

    def _make_inactive(self, index):
        length = self._piecelen(index)
        l = []
        x = 0
        request_size = self.config['download_slice_size']
        while x + request_size < length:
            l.append((x, request_size))
            x += request_size
        l.append((x, length - x))
        self.inactive_requests[index] = l

    def _load_fastresume(self, resumefile, typecode):
        if resumefile is not None:
            try:
                r = array(typecode)
                r.fromfile(resumefile, self.numpieces)
                return r
            except Exception, e:
                self.errorfunc(WARNING, _("Couldn't read fastresume data: ") +
                               str(e))
        return None

    def write_fastresume(self, resumefile):
        for i in xrange(self.numpieces):
            if self.rplaces[i] >= 0 and not self.have[self.rplaces[i]]:
                self.rplaces[i] = FASTRESUME_PARTIAL
        self.rplaces.tofile(resumefile)

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

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

    def do_I_have_requests(self, index):
        return not not self.inactive_requests[index]

    def new_request(self, index):
        # returns (begin, length)
        if self.inactive_requests[index] == 1:
            self._make_inactive(index)
        self.numactive[index] += 1
        self.stat_active[index] = 1
        if index not in self.stat_dirty:
            self.stat_new[index] = 1
        rs = self.inactive_requests[index]
        r = min(rs)
        rs.remove(r)
        self.amount_inactive -= r[1]
        if self.amount_inactive == 0:
            self.endgame = True
        return r

    def piece_came_in(self, index, begin, piece, source = None):
        if self.places[index] < 0:
            if self.rplaces[index] == ALLOCATED:
                self._initalloc(index, index)
            else:
                n = self._get_free_place()
                if self.places[n] >= 0:
                    oldpos = self.places[n]
                    self._move_piece(oldpos, n)
                    n = oldpos
                if self.rplaces[index] < 0 or index == n:
                    self._initalloc(n, index)
                else:
                    self._move_piece(index, n)
                    self._initalloc(index, index)

        if index in self.failed_pieces:
            old = self.storage.read(self.places[index] * self.piece_size +
                                    begin, len(piece))
            if old != piece:
                self.failed_pieces[index][self.download_history[index][begin]]\
                    = None
        self.download_history.setdefault(index, {})
        self.download_history[index][begin] = source

        self.storage.write(self.places[index] * self.piece_size + begin, piece)
        self.stat_dirty[index] = 1
        self.numactive[index] -= 1
        if self.numactive[index] == 0:
            del self.stat_active[index]
        if index in self.stat_new:
            del self.stat_new[index]
        if not self.inactive_requests[index] and not self.numactive[index]:
            del self.stat_dirty[index]
            if sha(self.storage.read(self.piece_size * self.places[index], self._piecelen(index))).digest() == self.hashes[index]:
                self.have[index] = True
                self.storage.downloaded(index * self.piece_size,
                                        self._piecelen(index))
                self.inactive_requests[index] = None
                self.waschecked[index] = True
                self.amount_left -= self._piecelen(index)
                self.stat_numdownloaded += 1
                for d in self.download_history[index].itervalues():
                    if d is not None:
                        d.good(index)
                del self.download_history[index]
                if index in self.failed_pieces:
                    for d in self.failed_pieces[index]:
                        if d is not None:
                            d.bad(index)
                    del self.failed_pieces[index]
                if self.amount_left == 0:
                    self.finished()
            else:
                self.data_flunked(self._piecelen(index), index)
                self.inactive_requests[index] = 1
                self.amount_inactive += self._piecelen(index)
                self.stat_numflunked += 1

                self.failed_pieces[index] = {}
                allsenders = {}
                for d in self.download_history[index].itervalues():
                    allsenders[d] = None
                if len(allsenders) == 1:
                    culprit = allsenders.keys()[0]
                    if culprit is not None:
                        culprit.bad(index, bump = True)
                    del self.failed_pieces[index] # found the culprit already
                return False
        return True

    def request_lost(self, index, begin, length):
        self.inactive_requests[index].append((begin, length))
        self.amount_inactive += length
        self.numactive[index] -= 1
        if not self.numactive[index] and index in self.stat_active:
            del self.stat_active[index]
            if index in self.stat_new:
                del self.stat_new[index]

    def get_piece(self, index, begin, length):
        if not self.have[index]:
            return None
        if not self.waschecked[index]:
            if sha(self.storage.read(self.piece_size * self.places[index], self._piecelen(index))).digest() != self.hashes[index]:
                raise BTFailure, _("told file complete on start-up, but piece failed hash check")
            self.waschecked[index] = True
        if begin + length > self._piecelen(index):
            return None
        return self.storage.read(self.piece_size * self.places[index] + begin, length)