diff --git a/src/BtHaveMessage.cc b/src/BtHaveMessage.cc index 50b40f39..92ae8f8c 100644 --- a/src/BtHaveMessage.cc +++ b/src/BtHaveMessage.cc @@ -55,10 +55,13 @@ void BtHaveMessage::doReceivedAction() if(isMetadataGetMode()) { return; } - getPeer()->updateBitfield(getIndex(), 1); - getPieceStorage()->addPieceStats(getIndex()); - if(getPeer()->isSeeder() && getPieceStorage()->downloadFinished()) { - throw DL_ABORT_EX(MSG_GOOD_BYE_SEEDER); + size_t index = getIndex(); + if(!getPeer()->hasPiece(index)) { + getPeer()->updateBitfield(index, 1); + getPieceStorage()->addPieceStats(index); + if(getPeer()->isSeeder() && getPieceStorage()->downloadFinished()) { + throw DL_ABORT_EX(MSG_GOOD_BYE_SEEDER); + } } } diff --git a/src/BtPieceMessage.cc b/src/BtPieceMessage.cc index 04025f48..12ffa2f0 100644 --- a/src/BtPieceMessage.cc +++ b/src/BtPieceMessage.cc @@ -179,8 +179,10 @@ size_t BtPieceMessage::getMessageHeaderLength() namespace { struct PieceSendUpdate : public ProgressUpdate { - PieceSendUpdate(const SharedHandle& peer, size_t headerLength) - : peer(peer), headerLength(headerLength) {} + PieceSendUpdate(const SharedHandle& peer, size_t headerLength, + size_t index, const SharedHandle& pieceStorage) + : peer(peer), headerLength(headerLength), index(index), + pieceStorage(pieceStorage) {} virtual void update(size_t length, bool complete) { if(headerLength > 0) { @@ -189,9 +191,17 @@ struct PieceSendUpdate : public ProgressUpdate { length -= m; } peer->updateUploadLength(length); + if(complete && !peer->hasPiece(index)) { + // Update peer's bitfield because peer may not send HAVE message + // to us. + peer->updateBitfield(index, 1); + pieceStorage->addPieceStats(index); + } } SharedHandle peer; size_t headerLength; + size_t index; + SharedHandle pieceStorage; }; } // namespace @@ -224,7 +234,9 @@ void BtPieceMessage::pushPieceData(int64_t offset, int32_t length) const buf.reset(0); getPeerConnection()->pushBytes(dbuf, length+MESSAGE_HEADER_LENGTH, new PieceSendUpdate(getPeer(), - MESSAGE_HEADER_LENGTH)); + MESSAGE_HEADER_LENGTH, + index_, + getPieceStorage())); // To avoid upload rate overflow, we update the length here at // once. downloadContext_->updateUploadLength(length); diff --git a/test/BtHaveMessageTest.cc b/test/BtHaveMessageTest.cc index a0cd5a77..5c1c7e2a 100644 --- a/test/BtHaveMessageTest.cc +++ b/test/BtHaveMessageTest.cc @@ -115,6 +115,7 @@ void BtHaveMessageTest::testDoReceivedAction_goodByeSeeder() msg.doReceivedAction(); pieceStorage->setDownloadFinished(true); + peer->updateBitfield(1, 0); try { msg.doReceivedAction(); CPPUNIT_FAIL("exception must be thrown.");