From a6b7bd03421b09c87c48c56440be35060d53d1fd Mon Sep 17 00:00:00 2001
From: Tatsuhiro Tsujikawa <tatsuhiro.t@gmail.com>
Date: Tue, 15 Jul 2014 23:36:10 +0900
Subject: [PATCH] Show bitfield for unknown length download in aria2.tellStatus
 RPC method

Generally, bitfield is not available for download whose total length
is unknown.  We create bitfield when download is completed (usually
connection EOF) so that we can use it to show additional info in RPC
aria2.tellStatus response.  Specifically, bitfield is now shown.  And
completedLength under files key (or completedLength in
aria2.getFiles() response) is correctly shown.
---
 src/UnknownLengthPieceStorage.cc | 39 ++++++++++++++++++++++++++++----
 src/UnknownLengthPieceStorage.h  | 16 ++++++-------
 2 files changed, 43 insertions(+), 12 deletions(-)

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<Piece>& piec
     totalLength_ = piece_->getLength();
     diskAdaptor_->setTotalLength(totalLength_);
     piece_.reset();
+
+    createBitfield();
   }
 }
 
@@ -219,10 +222,17 @@ std::shared_ptr<DiskAdaptor> 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<BitfieldMan>(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<BitfieldMan> bitfield_;
+
   bool downloadFinished_;
 
   std::shared_ptr<Piece> piece_;
+
+  void createBitfield();
+
 public:
   UnknownLengthPieceStorage(const std::shared_ptr<DownloadContext>& 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
   {