diff --git a/src/BtPieceMessage.cc b/src/BtPieceMessage.cc index b0362d72..04025f48 100644 --- a/src/BtPieceMessage.cc +++ b/src/BtPieceMessage.cc @@ -157,7 +157,7 @@ void BtPieceMessage::doReceivedAction() size_t BtPieceMessage::MESSAGE_HEADER_LENGTH = 13; -unsigned char* BtPieceMessage::createMessageHeader() +void BtPieceMessage::createMessageHeader(unsigned char* msgHeader) const { /** * len --- 9+blockLength, 4bytes @@ -166,12 +166,10 @@ unsigned char* BtPieceMessage::createMessageHeader() * begin --- begin, 4bytes * total: 13bytes */ - unsigned char* msgHeader = new unsigned char[MESSAGE_HEADER_LENGTH]; bittorrent::createPeerMessageString(msgHeader, MESSAGE_HEADER_LENGTH, 9+blockLength_, ID); bittorrent::setIntParam(&msgHeader[5], index_); bittorrent::setIntParam(&msgHeader[9], begin_); - return msgHeader; } size_t BtPieceMessage::getMessageHeaderLength() @@ -181,13 +179,19 @@ size_t BtPieceMessage::getMessageHeaderLength() namespace { struct PieceSendUpdate : public ProgressUpdate { - PieceSendUpdate(const SharedHandle& peer) - : peer(peer) {} + PieceSendUpdate(const SharedHandle& peer, size_t headerLength) + : peer(peer), headerLength(headerLength) {} virtual void update(size_t length, bool complete) { + if(headerLength > 0) { + size_t m = std::min(headerLength, length); + headerLength -= m; + length -= m; + } peer->updateUploadLength(length); } SharedHandle peer; + size_t headerLength; }; } // namespace @@ -201,8 +205,6 @@ void BtPieceMessage::send() getPeer()->getIPAddress().c_str(), getPeer()->getPort(), toString().c_str())); - getPeerConnection()->pushBytes(createMessageHeader(), - getMessageHeaderLength()); int64_t pieceDataOffset = static_cast(index_)*downloadContext_->getPieceLength()+begin_; pushPieceData(pieceDataOffset, blockLength_); @@ -211,14 +213,18 @@ void BtPieceMessage::send() void BtPieceMessage::pushPieceData(int64_t offset, int32_t length) const { assert(length <= 16*1024); - array_ptr buf(new unsigned char[length]); + array_ptr buf + (new unsigned char[length+MESSAGE_HEADER_LENGTH]); + createMessageHeader(buf); ssize_t r; - r = getPieceStorage()->getDiskAdaptor()->readData(buf, length, offset); + r = getPieceStorage()->getDiskAdaptor()->readData(buf+MESSAGE_HEADER_LENGTH, + length, offset); if(r == length) { unsigned char* dbuf = buf; buf.reset(0); - getPeerConnection()->pushBytes(dbuf, length, - new PieceSendUpdate(getPeer())); + getPeerConnection()->pushBytes(dbuf, length+MESSAGE_HEADER_LENGTH, + new PieceSendUpdate(getPeer(), + MESSAGE_HEADER_LENGTH)); // To avoid upload rate overflow, we update the length here at // once. downloadContext_->updateUploadLength(length); diff --git a/src/BtPieceMessage.h b/src/BtPieceMessage.h index 26e01ea5..09fef7dc 100644 --- a/src/BtPieceMessage.h +++ b/src/BtPieceMessage.h @@ -96,7 +96,7 @@ public: virtual void doReceivedAction(); - unsigned char* createMessageHeader(); + void createMessageHeader(unsigned char* msgHeader) const; size_t getMessageHeaderLength();