diff --git a/src/UnknownLengthPieceStorage.cc b/src/UnknownLengthPieceStorage.cc index 2a50f2b5..e067c01b 100644 --- a/src/UnknownLengthPieceStorage.cc +++ b/src/UnknownLengthPieceStorage.cc @@ -43,6 +43,7 @@ #include "DownloadContext.h" #include "Piece.h" #include "FileEntry.h" +#include "BitfieldMan.h" namespace aria2 { @@ -182,6 +183,8 @@ void UnknownLengthPieceStorage::completePiece(const std::shared_ptr& piec totalLength_ = piece_->getLength(); diskAdaptor_->setTotalLength(totalLength_); piece_.reset(); + + createBitfield(); } } @@ -219,10 +222,17 @@ std::shared_ptr UnknownLengthPieceStorage::getDiskAdaptor() int32_t UnknownLengthPieceStorage::getPieceLength(size_t index) { - if(index == 0) { - return totalLength_; - } else { - return 0; + // TODO Basically, PieceStorage::getPieceLength() is only used by + // BitTorrent, and it does not use UnknownLengthPieceStorage. + abort(); +} + +void UnknownLengthPieceStorage::createBitfield() +{ + if(totalLength_ > 0) { + bitfield_ = make_unique(downloadContext_->getPieceLength(), + totalLength_); + bitfield_->setAllBit(); } } @@ -232,6 +242,9 @@ void UnknownLengthPieceStorage::markAllPiecesDone() totalLength_ = piece_->getLength(); piece_.reset(); } + + createBitfield(); + downloadFinished_ = true; } @@ -257,4 +270,22 @@ void UnknownLengthPieceStorage::setDiskWriterFactory diskWriterFactory_ = diskWriterFactory; } +const unsigned char* UnknownLengthPieceStorage::getBitfield() +{ + if(bitfield_) { + return bitfield_->getBitfield(); + } + + return nullptr; +} + +size_t UnknownLengthPieceStorage::getBitfieldLength() +{ + if(bitfield_) { + return bitfield_->getBitfieldLength(); + } + + return 0; +} + } // namespace aria2 diff --git a/src/UnknownLengthPieceStorage.h b/src/UnknownLengthPieceStorage.h index 5aad9667..d39f3cae 100644 --- a/src/UnknownLengthPieceStorage.h +++ b/src/UnknownLengthPieceStorage.h @@ -43,6 +43,7 @@ class Option; class DownloadContext; class DiskWriterFactory; class DirectDiskAdaptor; +class BitfieldMan; class UnknownLengthPieceStorage:public PieceStorage { private: @@ -54,9 +55,14 @@ private: int64_t totalLength_; + std::unique_ptr bitfield_; + bool downloadFinished_; std::shared_ptr piece_; + + void createBitfield(); + public: UnknownLengthPieceStorage(const std::shared_ptr& downloadContext); @@ -203,18 +209,12 @@ public: */ virtual void initStorage() CXX11_OVERRIDE; - virtual const unsigned char* getBitfield() CXX11_OVERRIDE - { - return nullptr; - } + virtual const unsigned char* getBitfield() CXX11_OVERRIDE; virtual void setBitfield(const unsigned char* bitfield, size_t bitfieldLength) CXX11_OVERRIDE {} - virtual size_t getBitfieldLength() CXX11_OVERRIDE - { - return 0; - } + virtual size_t getBitfieldLength() CXX11_OVERRIDE; virtual bool isSelectiveDownloadingMode() CXX11_OVERRIDE {