Pārlūkot izejas kodu

2009-11-23 Tatsuhiro Tsujikawa <t-tujikawa@users.sourceforge.net>

	Fixed ut_metadata data handling. Implemented
	UTMetadataDataExtensionMessage::doReceivedAction().  Initialize
	PeerStorage in HandshakeExtensionMessage::doReceivedAction() when
	metadata_size is received.
	* src/DefaultExtensionMessageFactory.cc
	* src/DefaultExtensionMessageFactory.h
	* src/HandshakeExtensionMessage.cc
	* src/HandshakeExtensionMessage.h
	* src/UTMetadataDataExtensionMessage.cc
	* src/UTMetadataDataExtensionMessage.h
	* src/UTMetadataRequestExtensionMessage.cc
	* src/UTMetadataRequestFactory.cc
	* src/UTMetadataRequestFactory.h
	* src/UTMetadataRequestTracker.cc
	* src/UTMetadataRequestTracker.h
	* test/DefaultExtensionMessageFactoryTest.cc
	* test/HandshakeExtensionMessageTest.cc
	* test/MockBtMessage.h
	* test/UTMetadataDataExtensionMessageTest.cc
	* test/UTMetadataRequestExtensionMessageTest.cc
	* test/UTMetadataRequestFactoryTest.cc
	* test/UTMetadataRequestTrackerTest.cc
	* test/extension_message_test_helper.h
Tatsuhiro Tsujikawa 16 gadi atpakaļ
vecāks
revīzija
c1730aeea9

+ 26 - 0
ChangeLog

@@ -1,3 +1,29 @@
+2009-11-23  Tatsuhiro Tsujikawa  <t-tujikawa@users.sourceforge.net>
+
+	Fixed ut_metadata data handling. Implemented
+	UTMetadataDataExtensionMessage::doReceivedAction().  Initialize
+	PeerStorage in HandshakeExtensionMessage::doReceivedAction() when
+	metadata_size is received.
+	* src/DefaultExtensionMessageFactory.cc
+	* src/DefaultExtensionMessageFactory.h
+	* src/HandshakeExtensionMessage.cc
+	* src/HandshakeExtensionMessage.h
+	* src/UTMetadataDataExtensionMessage.cc
+	* src/UTMetadataDataExtensionMessage.h
+	* src/UTMetadataRequestExtensionMessage.cc
+	* src/UTMetadataRequestFactory.cc
+	* src/UTMetadataRequestFactory.h
+	* src/UTMetadataRequestTracker.cc
+	* src/UTMetadataRequestTracker.h
+	* test/DefaultExtensionMessageFactoryTest.cc
+	* test/HandshakeExtensionMessageTest.cc
+	* test/MockBtMessage.h
+	* test/UTMetadataDataExtensionMessageTest.cc
+	* test/UTMetadataRequestExtensionMessageTest.cc
+	* test/UTMetadataRequestFactoryTest.cc
+	* test/UTMetadataRequestTrackerTest.cc
+	* test/extension_message_test_helper.h
+
 2009-11-23  Tatsuhiro Tsujikawa  <t-tujikawa@users.sourceforge.net>
 
 	Drop connection if ut_metadata reject message is received.

+ 13 - 19
src/DefaultExtensionMessageFactory.cc

@@ -50,6 +50,10 @@
 #include "UTMetadataRejectExtensionMessage.h"
 #include "message.h"
 #include "bencode.h"
+#include "PieceStorage.h"
+#include "UTMetadataRequestTracker.h"
+#include "BtRuntime.h"
+#include "RequestGroup.h"
 
 namespace aria2 {
 
@@ -93,16 +97,8 @@ DefaultExtensionMessageFactory::createMessage(const unsigned char* data, size_t
 	throw DL_ABORT_EX(StringFormat(MSG_TOO_SMALL_PAYLOAD_SIZE,
 				       "ut_metadata", length).str());
       }
-      std::string listdata;
-      listdata += 'l';
-      listdata += std::string(&data[1], &data[length]);
-      listdata += 'e';
-
-      const BDE& list = bencode::decode(listdata);
-      if(!list.isList() || list.empty()) {
-	throw DL_ABORT_EX("Bad ut_metadata");
-      }
-      const BDE& dict = list[0];
+      size_t end;
+      BDE dict = bencode::decode(data+1, length-1, end);
       if(!dict.isDict()) {
 	throw DL_ABORT_EX("Bad ut_metadata: dictionary not found");
       }
@@ -126,13 +122,9 @@ DefaultExtensionMessageFactory::createMessage(const unsigned char* data, size_t
 	return m;
       }
       case 1: {
-	if(list.size() != 2) {
+	if(end == length) {
 	  throw DL_ABORT_EX("Bad ut_metadata data: data not found");
 	}
-	const BDE& pieceData = list[1];
-	if(!pieceData.isString()) {
-	  throw DL_ABORT_EX("Bad ut_metadata data: data is not string");
-	}
 	const BDE& totalSize = dict["total_size"];
 	if(!totalSize.isInteger()) {
 	  throw DL_ABORT_EX("Bad ut_metadata data: total_size not found");
@@ -141,16 +133,18 @@ DefaultExtensionMessageFactory::createMessage(const unsigned char* data, size_t
 	  (new UTMetadataDataExtensionMessage(extensionMessageID));
 	m->setIndex(index.i());
 	m->setTotalSize(totalSize.i());
-	m->setData(pieceData.s());
-	// set tracker
-	// set piecestorage
+	m->setData(std::string(&data[1+end], &data[length]));
+	m->setUTMetadataRequestTracker(_tracker);
+	m->setPieceStorage(_dctx->getOwnerRequestGroup()->getPieceStorage());
+	m->setDownloadContext(_dctx);
+	m->setBtRuntime(_btRuntime);
 	return m;
       }
       case 2: {
 	SharedHandle<UTMetadataRejectExtensionMessage> m
 	  (new UTMetadataRejectExtensionMessage(extensionMessageID));
 	m->setIndex(index.i());
-	// set tracker if disconnecing peer on receive.
+	// No need to inject tracker because peer will be disconnected.
 	return m;
       }
       default:

+ 17 - 0
src/DefaultExtensionMessageFactory.h

@@ -46,6 +46,8 @@ class ExtensionMessageRegistry;
 class DownloadContext;
 class BtMessageFactory;
 class BtMessageDispatcher;
+class UTMetadataRequestTracker;
+class BtRuntime;
 
 class DefaultExtensionMessageFactory:public ExtensionMessageFactory {
 private:
@@ -57,10 +59,14 @@ private:
 
   SharedHandle<DownloadContext> _dctx;
 
+  SharedHandle<BtRuntime> _btRuntime;
+
   WeakHandle<BtMessageFactory> _messageFactory;
 
   WeakHandle<BtMessageDispatcher> _dispatcher;
 
+  WeakHandle<UTMetadataRequestTracker> _tracker;
+
   Logger* _logger;
 
 public:
@@ -99,6 +105,17 @@ public:
   {
     _dispatcher = disp;
   }
+  
+  void setUTMetadataRequestTracker
+  (const WeakHandle<UTMetadataRequestTracker>& tracker)
+  {
+    _tracker = tracker;
+  }
+
+  void setBtRuntime(const SharedHandle<BtRuntime>& btRuntime)
+  {
+    _btRuntime = btRuntime;
+  }
 };
 
 typedef SharedHandle<DefaultExtensionMessageFactory> DefaultExtensionMessageFactoryHandle;

+ 14 - 1
src/HandshakeExtensionMessage.cc

@@ -43,6 +43,8 @@
 #include "bencode.h"
 #include "DownloadContext.h"
 #include "bittorrent_helper.h"
+#include "RequestGroup.h"
+#include "PieceStorage.h"
 
 namespace aria2 {
 
@@ -110,8 +112,19 @@ void HandshakeExtensionMessage::doReceivedAction()
   }
   if(_metadataSize > 0) {
     BDE& attrs = _dctx->getAttribute(bittorrent::BITTORRENT);
-    if(!attrs.containsKey(bittorrent::METADATA_SIZE)) {
+    if(attrs.containsKey(bittorrent::METADATA_SIZE)) {
+      if(_metadataSize != (size_t)attrs[bittorrent::METADATA_SIZE].i()) {
+	throw DL_ABORT_EX("Wrong metadata_size. Which one is correct!?");
+      }
+    } else {
       attrs[bittorrent::METADATA_SIZE] = _metadataSize;
+      _dctx->getFirstFileEntry()->setLength(_metadataSize);
+      _dctx->markTotalLengthIsKnown();
+      _dctx->getOwnerRequestGroup()->initPieceStorage();
+      
+      SharedHandle<PieceStorage> pieceStorage =
+	_dctx->getOwnerRequestGroup()->getPieceStorage();
+      pieceStorage->setEndGamePieceNum(0);
     }
   }
 }

+ 2 - 4
src/HandshakeExtensionMessage.h

@@ -45,9 +45,7 @@ namespace aria2 {
 
 class Peer;
 class Logger;
-class HandshakeExtensionMessage;
 class DownloadContext;
-typedef SharedHandle<HandshakeExtensionMessage> HandshakeExtensionMessageHandle;
 
 class HandshakeExtensionMessage:public ExtensionMessage {
 private:
@@ -137,8 +135,8 @@ public:
 
   void setPeer(const SharedHandle<Peer>& peer);
 
-  static HandshakeExtensionMessageHandle create(const unsigned char* data,
-						size_t dataLength);
+  static SharedHandle<HandshakeExtensionMessage>
+  create(const unsigned char* data, size_t dataLength);
 };
 
 typedef SharedHandle<HandshakeExtensionMessage> HandshakeExtensionMessageHandle;

+ 43 - 15
src/UTMetadataDataExtensionMessage.cc

@@ -37,29 +37,30 @@
 #include "bencode.h"
 #include "util.h"
 #include "a2functional.h"
+#include "DownloadContext.h"
+#include "UTMetadataRequestTracker.h"
+#include "PieceStorage.h"
+#include "BtConstants.h"
+#include "MessageDigestHelper.h"
+#include "bittorrent_helper.h"
+#include "DiskAdaptor.h"
+#include "Piece.h"
+#include "BtRuntime.h"
+#include "LogFactory.h"
 
 namespace aria2 {
 
 UTMetadataDataExtensionMessage::UTMetadataDataExtensionMessage
-(uint8_t extensionMessageID):UTMetadataExtensionMessage(extensionMessageID) {}
+(uint8_t extensionMessageID):UTMetadataExtensionMessage(extensionMessageID),
+			     _logger(LogFactory::getInstance()) {}
 
 std::string UTMetadataDataExtensionMessage::getBencodedData()
 {
-  BDE list = BDE::list();
-
   BDE dict = BDE::dict();
   dict["msg_type"] = 1;
   dict["piece"] = _index;
   dict["total_size"] = _totalSize;
-
-  BDE data = _data;
-
-  list << dict;
-  list << data;
-
-  std::string encodedList = bencode::encode(list);
-  // Remove first 'l' and last 'e' and return.
-  return std::string(encodedList.begin()+1, encodedList.end()-1);
+  return bencode::encode(dict)+_data;
 }
 
 std::string UTMetadataDataExtensionMessage::toString() const
@@ -69,9 +70,36 @@ std::string UTMetadataDataExtensionMessage::toString() const
 
 void UTMetadataDataExtensionMessage::doReceivedAction()
 {
-  // Update tracker
-
-  // Write to pieceStorage
+  if(_tracker->tracks(_index)) {
+    _logger->debug("ut_metadata index=%lu found in tracking list",
+		  static_cast<unsigned long>(_index));
+    _tracker->remove(_index);
+    _pieceStorage->getDiskAdaptor()->writeData
+      (reinterpret_cast<const unsigned char*>(_data.c_str()), _data.size(),
+       _index*METADATA_PIECE_SIZE);
+    _pieceStorage->completePiece(_pieceStorage->getPiece(_index));
+    if(_pieceStorage->downloadFinished()) {
+      std::string metadata = util::toString(_pieceStorage->getDiskAdaptor());
+      unsigned char infoHash[INFO_HASH_LENGTH];
+      MessageDigestHelper::digest(infoHash, INFO_HASH_LENGTH,
+				  MessageDigestContext::SHA1,
+				  metadata.data(), metadata.size());
+      const BDE& attrs = _dctx->getAttribute(bittorrent::BITTORRENT);
+      if(std::string(&infoHash[0], &infoHash[INFO_HASH_LENGTH]) == 
+	 attrs[bittorrent::INFO_HASH].s()){
+	_logger->info("Got ut_metadata");
+	_btRuntime->setHalt(true);
+      } else {
+	_logger->info("Got wrong ut_metadata");
+	for(size_t i = 0; i < _dctx->getNumPieces(); ++i) {
+	  _pieceStorage->markPieceMissing(i);
+	}
+      }
+    }
+  } else {
+    _logger->debug("ut_metadata index=%lu is not tracked",
+		  static_cast<unsigned long>(_index));
+  }
 }
 
 } // namespace aria2

+ 37 - 0
src/UTMetadataDataExtensionMessage.h

@@ -39,11 +39,27 @@
 
 namespace aria2 {
 
+class DownloadContext;
+class PieceStorage;
+class UTMetadataRequestTracker;
+class BtRuntime;
+class Logger;
+
 class UTMetadataDataExtensionMessage:public UTMetadataExtensionMessage {
 private:
   size_t _totalSize;
 
   std::string _data;
+
+  SharedHandle<DownloadContext> _dctx;
+
+  SharedHandle<PieceStorage> _pieceStorage;
+
+  SharedHandle<BtRuntime> _btRuntime;
+
+  WeakHandle<UTMetadataRequestTracker> _tracker;
+
+  Logger* _logger;
 public:
   UTMetadataDataExtensionMessage(uint8_t extensionMessageID);
 
@@ -72,6 +88,27 @@ public:
   {
     return _data;
   }
+
+  void setPieceStorage(const SharedHandle<PieceStorage>& pieceStorage)
+  {
+    _pieceStorage = pieceStorage;
+  }
+
+  void setUTMetadataRequestTracker
+  (const WeakHandle<UTMetadataRequestTracker>& tracker)
+  {
+    _tracker = tracker;
+  }
+
+  void setDownloadContext(const SharedHandle<DownloadContext>& dctx)
+  {
+    _dctx = dctx;
+  }
+
+  void setBtRuntime(const SharedHandle<BtRuntime>& btRuntime)
+  {
+    _btRuntime = btRuntime;
+  }
 };
 
 } // namespace aria2

+ 2 - 0
src/UTMetadataRequestExtensionMessage.cc

@@ -48,6 +48,8 @@
 #include "BtConstants.h"
 #include "DownloadContext.h"
 #include "BtMessage.h"
+#include "PieceStorage.h"
+#include "BtRuntime.h"
 
 namespace aria2 {
 

+ 79 - 0
src/UTMetadataRequestFactory.cc

@@ -0,0 +1,79 @@
+/* <!-- copyright */
+/*
+ * aria2 - The high speed download utility
+ *
+ * Copyright (C) 2009 Tatsuhiro Tsujikawa
+ *
+ * This program is free software; you can redistribute it and/or modify
+ * it under the terms of the GNU General Public License as published by
+ * the Free Software Foundation; either version 2 of the License, or
+ * (at your option) any later version.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ * GNU General Public License for more details.
+ *
+ * You should have received a copy of the GNU General Public License
+ * along with this program; if not, write to the Free Software
+ * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301  USA
+ *
+ * In addition, as a special exception, the copyright holders give
+ * permission to link the code of portions of this program with the
+ * OpenSSL library under certain conditions as described in each
+ * individual source file, and distribute linked combinations
+ * including the two.
+ * You must obey the GNU General Public License in all respects
+ * for all of the code used other than OpenSSL.  If you modify
+ * file(s) with this exception, you may extend this exception to your
+ * version of the file(s), but you are not obligated to do so.  If you
+ * do not wish to do so, delete this exception statement from your
+ * version.  If you delete this exception statement from all source
+ * files in the program, then also delete it here.
+ */
+/* copyright --> */
+#include "UTMetadataRequestFactory.h"
+#include "PieceStorage.h"
+#include "DownloadContext.h"
+#include "Peer.h"
+#include "BtMessageDispatcher.h"
+#include "BtMessageFactory.h"
+#include "UTMetadataRequestExtensionMessage.h"
+#include "UTMetadataRequestTracker.h"
+#include "BtMessage.h"
+#include "LogFactory.h"
+
+namespace aria2 {
+
+UTMetadataRequestFactory::UTMetadataRequestFactory():
+  _logger(LogFactory::getInstance()) {}
+
+void UTMetadataRequestFactory::create
+(std::deque<SharedHandle<BtMessage> >& msgs, size_t num,
+ const SharedHandle<PieceStorage>& pieceStorage)
+{
+  for(size_t index = 0; index < _dctx->getNumPieces() && num; ++index) {
+    SharedHandle<Piece> p = pieceStorage->getMissingPiece(index);
+    if(p.isNull()) {
+      _logger->debug("ut_metadata piece %lu is used or already acquired.");
+      continue;
+    }
+    --num;
+    _logger->debug("Creating ut_metadata request index=%lu",
+		   static_cast<unsigned long>(index));
+    SharedHandle<UTMetadataRequestExtensionMessage> m
+      (new UTMetadataRequestExtensionMessage
+       (_peer->getExtensionMessageID("ut_metadata")));
+    m->setIndex(index);
+    m->setDownloadContext(_dctx);
+    m->setBtMessageDispatcher(_dispatcher);
+    m->setBtMessageFactory(_messageFactory);
+    m->setPeer(_peer);
+    
+    SharedHandle<BtMessage> msg = _messageFactory->createBtExtendedMessage(m);
+    msgs.push_back(msg);
+    _tracker->add(index);
+  }
+}
+
+} // namespace aria2

+ 105 - 0
src/UTMetadataRequestFactory.h

@@ -0,0 +1,105 @@
+/* <!-- copyright */
+/*
+ * aria2 - The high speed download utility
+ *
+ * Copyright (C) 2009 Tatsuhiro Tsujikawa
+ *
+ * This program is free software; you can redistribute it and/or modify
+ * it under the terms of the GNU General Public License as published by
+ * the Free Software Foundation; either version 2 of the License, or
+ * (at your option) any later version.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ * GNU General Public License for more details.
+ *
+ * You should have received a copy of the GNU General Public License
+ * along with this program; if not, write to the Free Software
+ * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301  USA
+ *
+ * In addition, as a special exception, the copyright holders give
+ * permission to link the code of portions of this program with the
+ * OpenSSL library under certain conditions as described in each
+ * individual source file, and distribute linked combinations
+ * including the two.
+ * You must obey the GNU General Public License in all respects
+ * for all of the code used other than OpenSSL.  If you modify
+ * file(s) with this exception, you may extend this exception to your
+ * version of the file(s), but you are not obligated to do so.  If you
+ * do not wish to do so, delete this exception statement from your
+ * version.  If you delete this exception statement from all source
+ * files in the program, then also delete it here.
+ */
+/* copyright --> */
+#ifndef _UT_METADATA_REQUEST_FACTORY_H_
+#define _UT_METADATA_REQUEST_FACTORY_H_
+
+#include "common.h"
+
+#include <deque>
+
+#include "SharedHandle.h"
+
+namespace aria2 {
+
+class PieceStorage;
+class DownloadContext;
+class Peer;
+class BtMessageDispatcher;
+class BtMessageFactory;
+class UTMetadataRequestTracker;
+class BtMessage;
+class Logger;
+
+class UTMetadataRequestFactory {
+private:
+  SharedHandle<DownloadContext> _dctx;
+
+  SharedHandle<Peer> _peer;
+
+  WeakHandle<BtMessageDispatcher> _dispatcher;
+
+  WeakHandle<BtMessageFactory> _messageFactory;
+
+  WeakHandle<UTMetadataRequestTracker> _tracker;
+
+  Logger* _logger;
+public:
+  UTMetadataRequestFactory();
+
+  // Creates at most num of ut_metadata request message and appends
+  // them to msgs. pieceStorage is used to identify missing piece.
+  void create(std::deque<SharedHandle<BtMessage> >& msgs, size_t num,
+	      const SharedHandle<PieceStorage>& pieceStorage);
+
+  void setDownloadContext(const SharedHandle<DownloadContext>& dctx)
+  {
+    _dctx = dctx;
+  }
+
+  void setBtMessageDispatcher(const WeakHandle<BtMessageDispatcher>& disp)
+  {
+    _dispatcher = disp;
+  }
+
+  void setBtMessageFactory(const WeakHandle<BtMessageFactory>& factory)
+  {
+    _messageFactory = factory;
+  }
+
+  void setPeer(const SharedHandle<Peer>& peer)
+  {
+    _peer = peer;
+  }
+
+  void setUTMetadataRequestTracker
+  (const WeakHandle<UTMetadataRequestTracker>& tracker)
+  {
+    _tracker = tracker;
+  }
+};
+
+} // namespace aria2
+
+#endif // _UT_METADATA_REQUEST_FACTORY_H_

+ 106 - 0
src/UTMetadataRequestTracker.cc

@@ -0,0 +1,106 @@
+/* <!-- copyright */
+/*
+ * aria2 - The high speed download utility
+ *
+ * Copyright (C) 2009 Tatsuhiro Tsujikawa
+ *
+ * This program is free software; you can redistribute it and/or modify
+ * it under the terms of the GNU General Public License as published by
+ * the Free Software Foundation; either version 2 of the License, or
+ * (at your option) any later version.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ * GNU General Public License for more details.
+ *
+ * You should have received a copy of the GNU General Public License
+ * along with this program; if not, write to the Free Software
+ * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301  USA
+ *
+ * In addition, as a special exception, the copyright holders give
+ * permission to link the code of portions of this program with the
+ * OpenSSL library under certain conditions as described in each
+ * individual source file, and distribute linked combinations
+ * including the two.
+ * You must obey the GNU General Public License in all respects
+ * for all of the code used other than OpenSSL.  If you modify
+ * file(s) with this exception, you may extend this exception to your
+ * version of the file(s), but you are not obligated to do so.  If you
+ * do not wish to do so, delete this exception statement from your
+ * version.  If you delete this exception statement from all source
+ * files in the program, then also delete it here.
+ */
+/* copyright --> */
+#include "UTMetadataRequestTracker.h"
+
+#include <algorithm>
+
+#include "LogFactory.h"
+
+namespace aria2 {
+
+UTMetadataRequestTracker::UTMetadataRequestTracker():
+  _logger(LogFactory::getInstance()) {}
+
+void UTMetadataRequestTracker::add(size_t index)
+{
+  _trackedRequests.push_back(RequestEntry(index));
+}
+
+bool UTMetadataRequestTracker::tracks(size_t index)
+{
+  return std::find(_trackedRequests.begin(), _trackedRequests.end(),
+		   RequestEntry(index)) != _trackedRequests.end();
+}
+
+void UTMetadataRequestTracker::remove(size_t index)
+{
+  std::vector<RequestEntry>::iterator i =
+    std::find(_trackedRequests.begin(), _trackedRequests.end(),
+	      RequestEntry(index));
+  if(i != _trackedRequests.end()) {
+    _trackedRequests.erase(i);
+  }
+}
+
+std::vector<size_t> UTMetadataRequestTracker::removeTimeoutEntry()
+{
+  std::vector<size_t> indexes;
+  const time_t TIMEOUT = 20;
+  for(std::vector<RequestEntry>::iterator i = _trackedRequests.begin();
+      i != _trackedRequests.end();) {
+    if((*i).elapsed(TIMEOUT)) {
+      LogFactory::getInstance()->debug
+	("ut_metadata request timeout. index=%lu",
+	 static_cast<unsigned long>((*i)._index));
+      indexes.push_back((*i)._index);
+      i = _trackedRequests.erase(i);
+    } else {
+      ++i;
+    }
+  }
+  return indexes;
+}
+
+size_t UTMetadataRequestTracker::avail() const
+{
+  const size_t MAX_OUTSTANDING_REQUEST = 1;
+  if(MAX_OUTSTANDING_REQUEST > count()) {
+    return MAX_OUTSTANDING_REQUEST-count();
+  } else {
+    return 0;
+  }
+}
+
+std::vector<size_t> UTMetadataRequestTracker::getAllTrackedIndex() const
+{
+  std::vector<size_t> indexes;
+  for(std::vector<RequestEntry>::const_iterator i = _trackedRequests.begin();
+      i != _trackedRequests.end(); ++i) {
+    indexes.push_back((*i)._index);
+  }
+  return indexes;
+}
+
+} // namespace aria2

+ 100 - 0
src/UTMetadataRequestTracker.h

@@ -0,0 +1,100 @@
+/* <!-- copyright */
+/*
+ * aria2 - The high speed download utility
+ *
+ * Copyright (C) 2009 Tatsuhiro Tsujikawa
+ *
+ * This program is free software; you can redistribute it and/or modify
+ * it under the terms of the GNU General Public License as published by
+ * the Free Software Foundation; either version 2 of the License, or
+ * (at your option) any later version.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ * GNU General Public License for more details.
+ *
+ * You should have received a copy of the GNU General Public License
+ * along with this program; if not, write to the Free Software
+ * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301  USA
+ *
+ * In addition, as a special exception, the copyright holders give
+ * permission to link the code of portions of this program with the
+ * OpenSSL library under certain conditions as described in each
+ * individual source file, and distribute linked combinations
+ * including the two.
+ * You must obey the GNU General Public License in all respects
+ * for all of the code used other than OpenSSL.  If you modify
+ * file(s) with this exception, you may extend this exception to your
+ * version of the file(s), but you are not obligated to do so.  If you
+ * do not wish to do so, delete this exception statement from your
+ * version.  If you delete this exception statement from all source
+ * files in the program, then also delete it here.
+ */
+/* copyright --> */
+#ifndef _UT_METADATA_REQUEST_TRACKER_H_
+#define _UT_METADATA_REQUEST_TRACKER_H_
+
+#include "common.h"
+
+#include <vector>
+
+#include "TimeA2.h"
+
+namespace aria2 {
+
+class Logger;
+
+class UTMetadataRequestTracker {
+private:
+  struct RequestEntry {
+    size_t _index;
+    Time _dispatchedTime;
+
+    RequestEntry(size_t index):_index(index) {}
+
+    bool elapsed(time_t t) const
+    {
+      return _dispatchedTime.elapsed(t);
+    }
+
+    bool operator==(const RequestEntry& e) const
+    {
+      return _index == e._index;
+    }
+  };
+
+  std::vector<RequestEntry> _trackedRequests;
+
+  Logger* _logger;
+public:
+  UTMetadataRequestTracker();
+
+  // Add request index to tracking list.
+  void add(size_t index);
+
+  // Returns true if request index is tracked.
+  bool tracks(size_t index);
+
+  // Remove index from tracking list.
+  void remove(size_t index);
+
+  // Returns all tracking indexes.
+  std::vector<size_t> getAllTrackedIndex() const;
+
+  // Removes request index which is timed out and returns their indexes.
+  std::vector<size_t> removeTimeoutEntry();
+
+  // Returns the number of tracking list.
+  size_t count() const
+  {
+    return _trackedRequests.size();
+  }
+
+  // Returns the number of additional index this tracker can track.
+  size_t avail() const;
+};
+
+} // namespace aria2
+
+#endif // _UT_METADATA_REQUEST_TRACKER_H_

+ 11 - 1
test/DefaultExtensionMessageFactoryTest.cc

@@ -21,6 +21,10 @@
 #include "UTMetadataRequestExtensionMessage.h"
 #include "UTMetadataDataExtensionMessage.h"
 #include "UTMetadataRejectExtensionMessage.h"
+#include "BtRuntime.h"
+#include "PieceStorage.h"
+#include "RequestGroup.h"
+#include "Option.h"
 
 namespace aria2 {
 
@@ -42,6 +46,7 @@ private:
   SharedHandle<MockBtMessageDispatcher> _dispatcher;
   SharedHandle<MockBtMessageFactory> _messageFactory;
   SharedHandle<DownloadContext> _dctx;
+  SharedHandle<RequestGroup> _requestGroup;
 public:
   void setUp()
   {
@@ -59,6 +64,11 @@ public:
 
     _dctx.reset(new DownloadContext());
 
+    SharedHandle<Option> option(new Option());
+    _requestGroup.reset(new RequestGroup(option));
+    _requestGroup->setDownloadContext(_dctx);
+    _dctx->setOwnerRequestGroup(_requestGroup.get());
+
     _factory.reset(new DefaultExtensionMessageFactory());
     _factory->setPeerStorage(_peerStorage);
     _factory->setPeer(_peer);
@@ -155,7 +165,7 @@ void DefaultExtensionMessageFactoryTest::testCreateMessage_UTMetadataRequest()
 void DefaultExtensionMessageFactoryTest::testCreateMessage_UTMetadataData()
 {
   std::string data = getExtensionMessageID("ut_metadata")+
-    "d8:msg_typei1e5:piecei1e10:total_sizei300ee10:0000000000";
+    "d8:msg_typei1e5:piecei1e10:total_sizei300ee0000000000";
   SharedHandle<UTMetadataDataExtensionMessage> m =
     createMessage<UTMetadataDataExtensionMessage>(data);
   CPPUNIT_ASSERT_EQUAL((size_t)1, m->getIndex());

+ 14 - 3
test/HandshakeExtensionMessageTest.cc

@@ -9,6 +9,8 @@
 #include "FileEntry.h"
 #include "DownloadContext.h"
 #include "bittorrent_helper.h"
+#include "Option.h"
+#include "RequestGroup.h"
 
 namespace aria2 {
 
@@ -88,9 +90,16 @@ void HandshakeExtensionMessageTest::testToString()
 
 void HandshakeExtensionMessageTest::testDoReceivedAction()
 {
-  SharedHandle<DownloadContext> ctx(new DownloadContext());
+  SharedHandle<DownloadContext> dctx
+    (new DownloadContext(METADATA_PIECE_SIZE, 0));
+  SharedHandle<Option> op(new Option());
+  RequestGroup rg(op);
+  rg.setDownloadContext(dctx);
+  dctx->setOwnerRequestGroup(&rg);
+
   BDE attrs = BDE::dict();
-  ctx->setAttribute(bittorrent::BITTORRENT, attrs);
+  dctx->setAttribute(bittorrent::BITTORRENT, attrs);
+  dctx->markTotalLengthIsUnknown();
 
   SharedHandle<Peer> peer(new Peer("192.168.0.1", 0));
   peer->allocateSessionResource(1024, 1024*1024);
@@ -101,7 +110,7 @@ void HandshakeExtensionMessageTest::testDoReceivedAction()
   msg.setExtension("a2_dht", 2);
   msg.setMetadataSize(1024);
   msg.setPeer(peer);
-  msg.setDownloadContext(ctx);
+  msg.setDownloadContext(dctx);
 
   msg.doReceivedAction();
 
@@ -109,6 +118,8 @@ void HandshakeExtensionMessageTest::testDoReceivedAction()
   CPPUNIT_ASSERT_EQUAL((uint8_t)1, peer->getExtensionMessageID("ut_pex"));
   CPPUNIT_ASSERT_EQUAL((uint8_t)2, peer->getExtensionMessageID("a2_dht"));
   CPPUNIT_ASSERT_EQUAL((int64_t)1024, attrs[bittorrent::METADATA_SIZE].i());
+  CPPUNIT_ASSERT_EQUAL((uint64_t)1024, dctx->getTotalLength());
+  CPPUNIT_ASSERT(dctx->knowsTotalLength());
 }
 
 void HandshakeExtensionMessageTest::testCreate()

+ 9 - 0
test/MockBtMessage.h

@@ -98,6 +98,15 @@ public:
 
 };
 
+template<typename T>
+class WrapBtMessage:public MockBtMessage {
+public:
+  SharedHandle<T> _m;
+  
+  WrapBtMessage(const SharedHandle<T>& m):_m(m) {}
+};
+
+
 } // namespace aria2
 
 #endif // _D_MOCK_BT_MESSAGE_H_

+ 58 - 1
test/UTMetadataDataExtensionMessageTest.cc

@@ -5,6 +5,17 @@
 #include <cppunit/extensions/HelperMacros.h>
 
 #include "BtConstants.h"
+#include "PieceStorage.h"
+#include "DownloadContext.h"
+#include "BtRuntime.h"
+#include "DirectDiskAdaptor.h"
+#include "ByteArrayDiskWriter.h"
+#include "BDE.h"
+#include "DownloadContext.h"
+#include "MockPieceStorage.h"
+#include "UTMetadataRequestTracker.h"
+#include "bittorrent_helper.h"
+#include "MessageDigestHelper.h"
 
 namespace aria2 {
 
@@ -41,7 +52,7 @@ void UTMetadataDataExtensionMessageTest::testGetBencodedData()
   msg.setTotalSize(data.size());
   msg.setData(data);
   CPPUNIT_ASSERT_EQUAL
-    (std::string("d8:msg_typei1e5:piecei1e10:total_sizei16384ee16384:")+data,
+    (std::string("d8:msg_typei1e5:piecei1e10:total_sizei16384ee")+data,
      msg.getBencodedData());
 }
 
@@ -55,6 +66,52 @@ void UTMetadataDataExtensionMessageTest::testToString()
 
 void UTMetadataDataExtensionMessageTest::testDoReceivedAction()
 {
+  SharedHandle<DirectDiskAdaptor> diskAdaptor(new DirectDiskAdaptor());
+  SharedHandle<ByteArrayDiskWriter> diskWriter(new ByteArrayDiskWriter());
+  diskAdaptor->setDiskWriter(diskWriter);
+  SharedHandle<MockPieceStorage> pieceStorage(new MockPieceStorage());
+  pieceStorage->setDiskAdaptor(diskAdaptor);
+  SharedHandle<BtRuntime> btRuntime(new BtRuntime());
+  SharedHandle<UTMetadataRequestTracker> tracker
+    (new UTMetadataRequestTracker());
+  SharedHandle<DownloadContext> dctx(new DownloadContext());
+  BDE attrs = BDE::dict();
+
+  std::string piece0 = std::string(METADATA_PIECE_SIZE, '0');
+  std::string piece1 = std::string(METADATA_PIECE_SIZE, '1');
+  std::string metadata = piece0+piece1;
+
+  unsigned char infoHash[INFO_HASH_LENGTH];
+  MessageDigestHelper::digest(infoHash, INFO_HASH_LENGTH,
+			      MessageDigestContext::SHA1,
+			      metadata.data(), metadata.size());
+  attrs[bittorrent::INFO_HASH] = std::string(&infoHash[0], &infoHash[20]);
+
+  dctx->setAttribute(bittorrent::BITTORRENT, attrs);
+
+  UTMetadataDataExtensionMessage m(1);
+  m.setPieceStorage(pieceStorage);
+  m.setBtRuntime(btRuntime);
+  m.setUTMetadataRequestTracker(tracker);
+  m.setDownloadContext(dctx);
+
+  m.setIndex(1);
+  m.setData(piece1);
+  
+  tracker->add(1);
+  m.doReceivedAction();
+  CPPUNIT_ASSERT(!tracker->tracks(1));
+
+  pieceStorage->setDownloadFinished(true);
+  // If piece is not tracked, it is ignored.
+  m.setIndex(0);
+  m.setData(piece0);
+  m.doReceivedAction();
+  CPPUNIT_ASSERT(!btRuntime->isHalt());
+
+  tracker->add(0);
+  m.doReceivedAction();
+  CPPUNIT_ASSERT(btRuntime->isHalt());
 }
 
 } // namespace aria2

+ 7 - 20
test/UTMetadataRequestExtensionMessageTest.cc

@@ -13,6 +13,9 @@
 #include "BtHandshakeMessage.h"
 #include "UTMetadataRejectExtensionMessage.h"
 #include "UTMetadataDataExtensionMessage.h"
+#include "PieceStorage.h"
+#include "BtRuntime.h"
+#include "extension_message_test_helper.h"
 
 namespace aria2 {
 
@@ -27,30 +30,14 @@ class UTMetadataRequestExtensionMessageTest:public CppUnit::TestFixture {
   CPPUNIT_TEST(testDoReceivedAction_data);
   CPPUNIT_TEST_SUITE_END();
 public:
-  class MockExtensionMessage:public MockBtMessage {
-  public:
-    SharedHandle<ExtensionMessage> _m;
-    
-    MockExtensionMessage(const SharedHandle<ExtensionMessage>& m):_m(m) {}
-  };
-
-  class MockBtMessageFactory2:public MockBtMessageFactory {
-  public:
-    virtual SharedHandle<BtMessage>
-    createBtExtendedMessage(const SharedHandle<ExtensionMessage>& extmsg)
-    {
-      return SharedHandle<BtMessage>(new MockExtensionMessage(extmsg));
-    }
-  };
-
   SharedHandle<DownloadContext> _dctx;
-  SharedHandle<MockBtMessageFactory2> _messageFactory;
+  SharedHandle<WrapExtBtMessageFactory> _messageFactory;
   SharedHandle<MockBtMessageDispatcher> _dispatcher;
   SharedHandle<Peer> _peer;
 
   void setUp()
   {
-    _messageFactory.reset(new MockBtMessageFactory2());
+    _messageFactory.reset(new WrapExtBtMessageFactory());
     _dispatcher.reset(new MockBtMessageDispatcher());
     _dctx.reset(new DownloadContext());
     BDE attrs = BDE::dict();
@@ -63,8 +50,8 @@ public:
   template<typename T>
   SharedHandle<T> getFirstDispatchedMessage()
   {
-    SharedHandle<MockExtensionMessage> wrapmsg =
-      dynamic_pointer_cast<MockExtensionMessage>
+    SharedHandle<WrapExtBtMessage> wrapmsg =
+      dynamic_pointer_cast<WrapExtBtMessage>
       (_dispatcher->messageQueue.front());
     
     SharedHandle<T> msg = dynamic_pointer_cast<T>(wrapmsg->_m);

+ 76 - 0
test/UTMetadataRequestFactoryTest.cc

@@ -0,0 +1,76 @@
+#include "UTMetadataRequestFactory.h"
+
+#include <vector>
+#include <deque>
+
+#include <cppunit/extensions/HelperMacros.h>
+
+#include "MockPieceStorage.h"
+#include "DownloadContext.h"
+#include "Peer.h"
+#include "BtMessage.h"
+#include "extension_message_test_helper.h"
+#include "BtHandshakeMessage.h"
+#include "ExtensionMessage.h"
+#include "UTMetadataRequestTracker.h"
+
+namespace aria2 {
+
+class UTMetadataRequestFactoryTest:public CppUnit::TestFixture {
+
+  CPPUNIT_TEST_SUITE(UTMetadataRequestFactoryTest);
+  CPPUNIT_TEST(testCreate);
+  CPPUNIT_TEST_SUITE_END();
+public:
+  void testCreate();
+
+  class MockPieceStorage2:public MockPieceStorage {
+  public:
+    std::set<size_t> missingIndexes;
+
+    virtual SharedHandle<Piece> getMissingPiece(size_t index)
+    {
+      if(missingIndexes.find(index) != missingIndexes.end()) {
+	return SharedHandle<Piece>(new Piece(index, 0));
+      } else {
+	return SharedHandle<Piece>();
+      }
+    }
+  };
+};
+
+
+CPPUNIT_TEST_SUITE_REGISTRATION(UTMetadataRequestFactoryTest);
+
+void UTMetadataRequestFactoryTest::testCreate()
+{
+  UTMetadataRequestFactory factory;
+  SharedHandle<DownloadContext> dctx
+    (new DownloadContext(METADATA_PIECE_SIZE, METADATA_PIECE_SIZE*2));
+  factory.setDownloadContext(dctx);
+  SharedHandle<MockPieceStorage2> ps(new MockPieceStorage2());
+  ps->missingIndexes.insert(0);
+  ps->missingIndexes.insert(1);
+  SharedHandle<WrapExtBtMessageFactory> messageFactory
+    (new WrapExtBtMessageFactory());
+  factory.setBtMessageFactory(messageFactory);
+  SharedHandle<Peer> peer(new Peer("peer", 6880));
+  peer->allocateSessionResource(0, 0);
+  factory.setPeer(peer);
+  SharedHandle<UTMetadataRequestTracker> tracker
+    (new UTMetadataRequestTracker());
+  factory.setUTMetadataRequestTracker(tracker);
+
+  std::deque<SharedHandle<BtMessage> > msgs;
+
+  factory.create(msgs, 1, ps);
+  CPPUNIT_ASSERT_EQUAL((size_t)1, msgs.size());
+
+  msgs.clear();
+
+  ps->missingIndexes.clear();
+  factory.create(msgs, 1, ps);
+  CPPUNIT_ASSERT_EQUAL((size_t)0, msgs.size());
+}
+
+} // namespace aria2

+ 72 - 0
test/UTMetadataRequestTrackerTest.cc

@@ -0,0 +1,72 @@
+#include "UTMetadataRequestTracker.h"
+
+#include <cppunit/extensions/HelperMacros.h>
+
+namespace aria2 {
+
+class UTMetadataRequestTrackerTest:public CppUnit::TestFixture {
+
+  CPPUNIT_TEST_SUITE(UTMetadataRequestTrackerTest);
+  CPPUNIT_TEST(testAdd);
+  CPPUNIT_TEST(testRemove);
+  CPPUNIT_TEST(testGetAllTrackedIndex);
+  CPPUNIT_TEST(testCount);
+  CPPUNIT_TEST(testAvail);
+  CPPUNIT_TEST_SUITE_END();
+public:
+  void testAdd();
+  void testRemove();
+  void testGetAllTrackedIndex();
+  void testCount();
+  void testAvail();
+};
+
+
+CPPUNIT_TEST_SUITE_REGISTRATION(UTMetadataRequestTrackerTest);
+
+void UTMetadataRequestTrackerTest::testAdd()
+{
+  UTMetadataRequestTracker tr;
+  tr.add(1);
+  CPPUNIT_ASSERT(tr.tracks(1));
+}
+
+void UTMetadataRequestTrackerTest::testRemove()
+{
+  UTMetadataRequestTracker tr;
+  tr.add(1);
+  tr.remove(1);
+  CPPUNIT_ASSERT(!tr.tracks(1));
+}
+
+void UTMetadataRequestTrackerTest::testGetAllTrackedIndex()
+{
+  UTMetadataRequestTracker tr;
+  tr.add(1);
+  tr.add(2);
+
+  std::vector<size_t> indexes = tr.getAllTrackedIndex();
+  CPPUNIT_ASSERT_EQUAL((size_t)2, indexes.size());
+  CPPUNIT_ASSERT_EQUAL((size_t)1, indexes[0]);
+  CPPUNIT_ASSERT_EQUAL((size_t)2, indexes[1]);
+}
+
+void UTMetadataRequestTrackerTest::testCount()
+{
+  UTMetadataRequestTracker tr;
+  tr.add(1);
+  tr.add(2);
+  CPPUNIT_ASSERT_EQUAL((size_t)2, tr.count());
+}
+
+void UTMetadataRequestTrackerTest::testAvail()
+{
+  UTMetadataRequestTracker tr;
+  CPPUNIT_ASSERT_EQUAL((size_t)1, tr.avail());
+  tr.add(1);
+  CPPUNIT_ASSERT_EQUAL((size_t)0, tr.avail());
+  tr.add(2);
+  CPPUNIT_ASSERT_EQUAL((size_t)0, tr.avail());
+}
+
+} // namespace aria2

+ 22 - 0
test/extension_message_test_helper.h

@@ -0,0 +1,22 @@
+#ifndef _D_EXTENSION_MESSAGE_TEST_HELPER_H_
+#define _D_EXTENSION_MESSAGE_TEST_HELPER_H_
+
+#include "MockBtMessage.h"
+#include "MockBtMessageFactory.h"
+
+namespace aria2 {
+
+typedef WrapBtMessage<ExtensionMessage> WrapExtBtMessage;
+
+class WrapExtBtMessageFactory:public MockBtMessageFactory {
+public:
+  virtual SharedHandle<BtMessage>
+  createBtExtendedMessage(const SharedHandle<ExtensionMessage>& extmsg)
+  {
+    return SharedHandle<BtMessage>(new WrapExtBtMessage(extmsg));
+  }
+};
+
+} // namespace aria2
+
+#endif // _D_EXTENSION_MESSAGE_TEST_HELPER_H_