Refactor BtBitfieldMessage

2 files changed, 13 insertions, 24 deletions

diff --git a/src/BtBitfieldMessage.cc b/src/BtBitfieldMessage.cc
index e53cdf85..eaeb610e 100644
--- a/src/BtBitfieldMessage.cc
+++ b/src/BtBitfieldMessage.cc
@@ -48,16 +48,12 @@ namespace aria2 {
 
 const char BtBitfieldMessage::NAME[] = "bitfield";
 
-BtBitfieldMessage::BtBitfieldMessage()
-    : SimpleBtMessage(ID, NAME), bitfieldLength_(0)
-{
-}
+BtBitfieldMessage::BtBitfieldMessage() : SimpleBtMessage(ID, NAME) {}
 
 BtBitfieldMessage::BtBitfieldMessage(const unsigned char* bitfield,
                                      size_t bitfieldLength)
-    : SimpleBtMessage(ID, NAME), bitfieldLength_(0)
+    : SimpleBtMessage(ID, NAME), bitfield_(bitfield, bitfield + bitfieldLength)
 {
-  setBitfield(bitfield, bitfieldLength);
 }
 
 BtBitfieldMessage::~BtBitfieldMessage() {}
@@ -65,13 +61,7 @@ BtBitfieldMessage::~BtBitfieldMessage() {}
 void BtBitfieldMessage::setBitfield(const unsigned char* bitfield,
                                     size_t bitfieldLength)
 {
-  if (bitfield_.get() == bitfield) {
-    return;
-  }
-
-  bitfieldLength_ = bitfieldLength;
-  bitfield_ = make_unique<unsigned char[]>(bitfieldLength_);
-  memcpy(bitfield_.get(), bitfield, bitfieldLength_);
+  bitfield_.assign(bitfield, bitfield + bitfieldLength);
 }
 
 std::unique_ptr<BtBitfieldMessage>
@@ -89,9 +79,9 @@ void BtBitfieldMessage::doReceivedAction()
   if (isMetadataGetMode()) {
     return;
   }
-  getPieceStorage()->updatePieceStats(bitfield_.get(), bitfieldLength_,
+  getPieceStorage()->updatePieceStats(bitfield_.data(), bitfield_.size(),
                                       getPeer()->getBitfield());
-  getPeer()->setBitfield(bitfield_.get(), bitfieldLength_);
+  getPeer()->setBitfield(bitfield_.data(), bitfield_.size());
   if (getPeer()->isSeeder() && getPieceStorage()->downloadFinished()) {
     throw DL_ABORT_EX(MSG_GOOD_BYE_SEEDER);
   }
@@ -105,19 +95,19 @@ std::vector<unsigned char> BtBitfieldMessage::createMessage()
    * bitfield --- bitfield, bitfieldLength bytes
    * total: 5+bitfieldLength bytes
    */
-  const size_t msgLength = 5 + bitfieldLength_;
+  const size_t msgLength = 5 + bitfield_.size();
   auto msg = std::vector<unsigned char>(msgLength);
   bittorrent::createPeerMessageString(msg.data(), msgLength,
-                                      1 + bitfieldLength_, ID);
-  memcpy(msg.data() + 5, bitfield_.get(), bitfieldLength_);
+                                      1 + bitfield_.size(), ID);
+  std::copy(std::begin(bitfield_), std::end(bitfield_), std::begin(msg) + 5);
   return msg;
 }
 
 std::string BtBitfieldMessage::toString() const
 {
   std::string s = NAME;
-  s += " ";
-  s += util::toHex(bitfield_.get(), bitfieldLength_);
+  s += ' ';
+  s += util::toHex(bitfield_.data(), bitfield_.size());
   return s;
 }
 
diff --git a/src/BtBitfieldMessage.h b/src/BtBitfieldMessage.h
index 0015d439..4d492b83 100644
--- a/src/BtBitfieldMessage.h
+++ b/src/BtBitfieldMessage.h
@@ -41,8 +41,7 @@ namespace aria2 {
 
 class BtBitfieldMessage : public SimpleBtMessage {
 private:
-  std::unique_ptr<unsigned char[]> bitfield_;
-  size_t bitfieldLength_;
+  std::vector<unsigned char> bitfield_;
 
 public:
   BtBitfieldMessage();
@@ -57,9 +56,9 @@ public:
 
   void setBitfield(const unsigned char* bitfield, size_t bitfieldLength);
 
-  const unsigned char* getBitfield() const { return bitfield_.get(); }
+  const unsigned char* getBitfield() const { return bitfield_.data(); }
 
-  size_t getBitfieldLength() const { return bitfieldLength_; }
+  size_t getBitfieldLength() const { return bitfield_.size(); }
 
   static std::unique_ptr<BtBitfieldMessage> create(const unsigned char* data,
                                                    size_t dataLength);