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

	Fixed the bug which causes segmentation fault with tellWaiting
	XML-RPC method when BitTorrent download is waiting.  The key of
	BtRegistry is changed from info hash to GID, because it is
	registered per RequestGroup, not info hash.
	* src/BtRegistry.cc
	* src/BtRegistry.h
	* src/BtSetup.cc
	* src/ConsoleStatCalc.cc
	* src/InitiatorMSEHandshakeCommand.cc
	* src/PeerInitiateConnectionCommand.cc
	* src/PeerInteractionCommand.cc
	* src/PeerInteractionCommand.h
	* src/PeerReceiveHandshakeCommand.cc
	* src/RequestGroup.cc
	* src/XmlRpcMethodImpl.cc
	* test/BtRegistryTest.cc
	* test/XmlRpcMethodTest.cc

11 files changed, 68 insertions, 74 deletions

diff --git a/src/BtRegistry.cc b/src/BtRegistry.cc
index 158da62b..2a94cafe 100644
--- a/src/BtRegistry.cc
+++ b/src/BtRegistry.cc
@@ -40,25 +40,40 @@
 #include "BtAnnounce.h"
 #include "BtRuntime.h"
 #include "BtProgressInfoFile.h"
+#include "bittorrent_helper.h"
 
 namespace aria2 {
 
 SharedHandle<DownloadContext>
+BtRegistry::getDownloadContext(int32_t gid) const
+{
+  return get(gid)._downloadContext;
+}
+
+SharedHandle<DownloadContext>
 BtRegistry::getDownloadContext(const std::string& infoHash) const
 {
-  return get(infoHash)._downloadContext;
+  SharedHandle<DownloadContext> dctx;
+  for(std::map<int32_t, BtObject>::const_iterator i = _pool.begin();
+      i != _pool.end(); ++i) {
+    const BDE& attrs =
+      (*i).second._downloadContext->getAttribute(bittorrent::BITTORRENT);
+    if(attrs[bittorrent::INFO_HASH].s() == infoHash) {
+      dctx = (*i).second._downloadContext;
+      break;
+    }
+  }
+  return dctx;
 }
 
-void BtRegistry::put(const std::string& infoHash, const BtObject& obj)
+void BtRegistry::put(int32_t gid, const BtObject& obj)
 {
-  remove(infoHash);
-  std::map<std::string, BtObject>::value_type p(infoHash, obj);
-  _pool.insert(p);
+  _pool[gid] = obj;
 }
 
-BtObject BtRegistry::get(const std::string& infoHash) const
+BtObject BtRegistry::get(int32_t gid) const
 {
-  std::map<std::string, BtObject>::const_iterator i = _pool.find(infoHash);
+  std::map<int32_t, BtObject>::const_iterator i = _pool.find(gid);
   if(i == _pool.end()) {
     return BtObject();
   } else {
@@ -66,9 +81,9 @@ BtObject BtRegistry::get(const std::string& infoHash) const
   }
 }
 
-bool BtRegistry::remove(const std::string& infoHash)
+bool BtRegistry::remove(int32_t gid)
 {
-  return _pool.erase(infoHash);
+  return _pool.erase(gid);
 }
 
 void BtRegistry::removeAll() {
diff --git a/src/BtRegistry.h b/src/BtRegistry.h
index 3b3858f8..3fc642ee 100644
--- a/src/BtRegistry.h
+++ b/src/BtRegistry.h
@@ -37,7 +37,6 @@
 
 #include "common.h"
 
-#include <string>
 #include <map>
 
 #include "SharedHandle.h"
@@ -87,19 +86,22 @@ struct BtObject {
 
 class BtRegistry {
 private:
-  std::map<std::string, BtObject> _pool;
+  std::map<int32_t, BtObject> _pool;
 public:
   SharedHandle<DownloadContext>
+  getDownloadContext(int32_t gid) const;
+
+  SharedHandle<DownloadContext>
   getDownloadContext(const std::string& infoHash) const;
 
-  void put(const std::string& infoHash, const BtObject& obj);
+  void put(int32_t gid, const BtObject& obj);
 
-  BtObject get(const std::string& infoHash) const;
+  BtObject get(int32_t gid) const;
 
   template<typename OutputIterator>
   OutputIterator getAllDownloadContext(OutputIterator dest)
   {
-    for(std::map<std::string, BtObject>::const_iterator i = _pool.begin();
+    for(std::map<int32_t, BtObject>::const_iterator i = _pool.begin();
 	i != _pool.end(); ++i) {
       *dest++ = (*i).second._downloadContext;
     }
@@ -108,7 +110,7 @@ public:
 
   void removeAll();
 
-  bool remove(const std::string& infoHash);
+  bool remove(int32_t gid);
 };
 
 } // namespace aria2
diff --git a/src/BtSetup.cc b/src/BtSetup.cc
index 2bd12719..1ebb7476 100644
--- a/src/BtSetup.cc
+++ b/src/BtSetup.cc
@@ -76,8 +76,7 @@ void BtSetup::setup(std::deque<Command*>& commands,
   const BDE& torrentAttrs =
     requestGroup->getDownloadContext()->getAttribute(bittorrent::BITTORRENT);
   bool metadataGetMode = !torrentAttrs.containsKey(bittorrent::METADATA);
-  BtObject btObject =
-    e->getBtRegistry()->get(torrentAttrs[bittorrent::INFO_HASH].s());
+  BtObject btObject = e->getBtRegistry()->get(requestGroup->getGID());
   SharedHandle<PieceStorage> pieceStorage = btObject._pieceStorage;
   SharedHandle<PeerStorage> peerStorage = btObject._peerStorage;
   SharedHandle<BtRuntime> btRuntime = btObject._btRuntime;
diff --git a/src/ConsoleStatCalc.cc b/src/ConsoleStatCalc.cc
index 25449ce5..10dc5140 100644
--- a/src/ConsoleStatCalc.cc
+++ b/src/ConsoleStatCalc.cc
@@ -116,10 +116,8 @@ static void printProgress
     << rg->getNumConnection();
 #ifdef ENABLE_BITTORRENT
   if(rg->getDownloadContext()->hasAttribute(bittorrent::BITTORRENT)) {
-    const BDE& torrentAttrs =
-      rg->getDownloadContext()->getAttribute(bittorrent::BITTORRENT);
     SharedHandle<PeerStorage> ps =
-      e->getBtRegistry()->get(torrentAttrs[bittorrent::INFO_HASH].s())._peerStorage;
+      e->getBtRegistry()->get(rg->getGID())._peerStorage;
     std::deque<SharedHandle<Peer> > peers;
     ps->getActivePeers(peers);
     o << " " << "SEED:"
diff --git a/src/InitiatorMSEHandshakeCommand.cc b/src/InitiatorMSEHandshakeCommand.cc
index 7afeab6e..87717751 100644
--- a/src/InitiatorMSEHandshakeCommand.cc
+++ b/src/InitiatorMSEHandshakeCommand.cc
@@ -152,11 +152,10 @@ bool InitiatorMSEHandshakeCommand::executeInternal() {
       }
       PeerInteractionCommand* c =
 	  new PeerInteractionCommand
-	(cuid, _requestGroup, peer, e, _btRuntime, _pieceStorage,
+	(cuid, _requestGroup, peer, e, _btRuntime, _pieceStorage, _peerStorage,
 	 socket,
 	 PeerInteractionCommand::INITIATOR_SEND_HANDSHAKE,
 	 peerConnection);
-      c->setPeerStorage(_peerStorage);
       e->commands.push_back(c);
       return true;
     }
diff --git a/src/PeerInitiateConnectionCommand.cc b/src/PeerInitiateConnectionCommand.cc
index bd16a354..a05cc4e0 100644
--- a/src/PeerInitiateConnectionCommand.cc
+++ b/src/PeerInitiateConnectionCommand.cc
@@ -89,9 +89,8 @@ bool PeerInitiateConnectionCommand::executeInternal() {
   } else {
     PeerInteractionCommand* command =
       new PeerInteractionCommand
-      (cuid, _requestGroup, peer, e, _btRuntime, _pieceStorage,
+      (cuid, _requestGroup, peer, e, _btRuntime, _pieceStorage, _peerStorage,
        socket, PeerInteractionCommand::INITIATOR_SEND_HANDSHAKE);
-    command->setPeerStorage(_peerStorage);
     e->commands.push_back(command);
   }
   return true;
diff --git a/src/PeerInteractionCommand.cc b/src/PeerInteractionCommand.cc
index 5cbf025f..5255fe0e 100644
--- a/src/PeerInteractionCommand.cc
+++ b/src/PeerInteractionCommand.cc
@@ -45,7 +45,6 @@
 #include "Socket.h"
 #include "Option.h"
 #include "DownloadContext.h"
-#include "BtRegistry.h"
 #include "Peer.h"
 #include "BtMessage.h"
 #include "BtRuntime.h"
@@ -83,6 +82,7 @@ PeerInteractionCommand::PeerInteractionCommand
  DownloadEngine* e,
  const SharedHandle<BtRuntime>& btRuntime,
  const SharedHandle<PieceStorage>& pieceStorage,
+ const SharedHandle<PeerStorage>& peerStorage,
  const SocketHandle& s,
  Seq sequence,
  const PeerConnectionHandle& passedPeerConnection)
@@ -90,6 +90,7 @@ PeerInteractionCommand::PeerInteractionCommand
    _requestGroup(requestGroup),
    _btRuntime(btRuntime),
    _pieceStorage(pieceStorage),
+   _peerStorage(peerStorage),
    sequence(sequence)
 {
   // TODO move following bunch of processing to separate method, like init()
@@ -101,9 +102,6 @@ PeerInteractionCommand::PeerInteractionCommand
 
   const BDE& torrentAttrs =
     _requestGroup->getDownloadContext()->getAttribute(bittorrent::BITTORRENT);
-  SharedHandle<BtRegistry> btRegistry = e->getBtRegistry();
-  SharedHandle<PeerStorage> peerStorage =
-    btRegistry->get(torrentAttrs[bittorrent::INFO_HASH].s())._peerStorage;
 
   bool metadataGetMode = !torrentAttrs.containsKey(bittorrent::METADATA);
 
@@ -354,12 +352,6 @@ bool PeerInteractionCommand::exitBeforeExecute()
   return _btRuntime->isHalt();
 }
 
-void PeerInteractionCommand::setPeerStorage
-(const SharedHandle<PeerStorage>& peerStorage)
-{
-  _peerStorage = peerStorage;
-}
-
 const SharedHandle<Option>& PeerInteractionCommand::getOption() const
 {
   return _requestGroup->getOption();
diff --git a/src/PeerInteractionCommand.h b/src/PeerInteractionCommand.h
index 894cda7d..8dba1272 100644
--- a/src/PeerInteractionCommand.h
+++ b/src/PeerInteractionCommand.h
@@ -80,14 +80,13 @@ public:
 			 DownloadEngine* e,
 			 const SharedHandle<BtRuntime>& btRuntime,
 			 const SharedHandle<PieceStorage>& pieceStorage,
+			 const SharedHandle<PeerStorage>& peerStorage,
 			 const SharedHandle<SocketCore>& s,
 			 Seq sequence,
 			 const SharedHandle<PeerConnection>& peerConnection =
 			 SharedHandle<PeerConnection>());
 
   virtual ~PeerInteractionCommand();
-
-  void setPeerStorage(const SharedHandle<PeerStorage>& peerStorage);
 };
 
 } // namespace aria2
diff --git a/src/PeerReceiveHandshakeCommand.cc b/src/PeerReceiveHandshakeCommand.cc
index 8052dda5..337cd5ec 100644
--- a/src/PeerReceiveHandshakeCommand.cc
+++ b/src/PeerReceiveHandshakeCommand.cc
@@ -93,17 +93,24 @@ bool PeerReceiveHandshakeCommand::executeInternal()
     // check info_hash
     std::string infoHash = std::string(&data[28], &data[28+INFO_HASH_LENGTH]);
 
-    BtObject btObject = e->getBtRegistry()->get(infoHash);
-    SharedHandle<DownloadContext> downloadContext = btObject._downloadContext;
+    SharedHandle<DownloadContext> downloadContext =
+      e->getBtRegistry()->getDownloadContext(infoHash);
+    if(downloadContext.isNull()) {
+      throw DL_ABORT_EX
+	(StringFormat("Unknown info hash %s",
+		      util::toHex(infoHash).c_str()).str());
+    }
+    BtObject btObject = e->getBtRegistry()->get
+      (downloadContext->getOwnerRequestGroup()->getGID());
     SharedHandle<BtRuntime> btRuntime = btObject._btRuntime;
     SharedHandle<PieceStorage> pieceStorage = btObject._pieceStorage;
     SharedHandle<PeerStorage> peerStorage = btObject._peerStorage;
-
-    if(downloadContext.isNull() || !btRuntime->ready()) {
+    if(!btRuntime->ready()) {
       throw DL_ABORT_EX
 	(StringFormat("Unknown info hash %s",
 		      util::toHex(infoHash).c_str()).str());
     }
+
     TransferStat tstat =
       downloadContext->getOwnerRequestGroup()->calculateStat();
     const unsigned int maxDownloadLimit =
@@ -130,10 +137,10 @@ bool PeerReceiveHandshakeCommand::executeInternal()
 	   e,
 	   btRuntime,
 	   pieceStorage,
+	   peerStorage,
 	   socket,
 	   PeerInteractionCommand::RECEIVER_WAIT_HANDSHAKE,
 	   _peerConnection);
-	command->setPeerStorage(peerStorage);
 	e->commands.push_back(command);
 	logger->debug(MSG_INCOMING_PEER_CONNECTION, cuid, peer->usedBy());
       }
diff --git a/src/RequestGroup.cc b/src/RequestGroup.cc
index a07bd8ff..7625799a 100644
--- a/src/RequestGroup.cc
+++ b/src/RequestGroup.cc
@@ -264,7 +264,7 @@ void RequestGroup::createInitialCommand
 	(_option->getAsInt(PREF_BT_TRACKER_INTERVAL));
       btAnnounce->shuffleAnnounce();
       
-      btRegistry->put(torrentAttrs[bittorrent::INFO_HASH].s(),
+      btRegistry->put(_gid,
 		      BtObject(_downloadContext,
 			       _pieceStorage,
 			       peerStorage,
@@ -774,22 +774,7 @@ void RequestGroup::releaseRuntimeResource(DownloadEngine* e)
 {
 #ifdef ENABLE_BITTORRENT
   if(_downloadContext->hasAttribute(bittorrent::BITTORRENT)) {
-    SharedHandle<BtRegistry> btRegistry = e->getBtRegistry();
-    const BDE& torrentAttrs =
-      _downloadContext->getAttribute(bittorrent::BITTORRENT);
-    const std::string& infoHash = torrentAttrs[bittorrent::INFO_HASH].s();
-    SharedHandle<DownloadContext> contextInReg =
-      btRegistry->getDownloadContext(infoHash);
-    // Make sure that the registered DownloadContext's GID is equal to
-    // _gid.  Even if createInitialCommand() throws exception without
-    // registering this DownloadContext, after that, this method is
-    // called. In this case, just finding DownloadContext using
-    // infoHash may detect another download's DownloadContext and
-    // deleting it from BtRegistry causes Segmentation Fault.
-    if(!contextInReg.isNull() &&
-       contextInReg->getOwnerRequestGroup()->getGID() == _gid) {
-      btRegistry->remove(infoHash);
-    }
+    e->getBtRegistry()->remove(_gid);
   }
 #endif // ENABLE_BITTORRENT
   if(!_pieceStorage.isNull()) {
diff --git a/src/XmlRpcMethodImpl.cc b/src/XmlRpcMethodImpl.cc
index 9f90d228..deba3c6a 100644
--- a/src/XmlRpcMethodImpl.cc
+++ b/src/XmlRpcMethodImpl.cc
@@ -290,17 +290,19 @@ static void gatherProgressCommon
 
 #ifdef ENABLE_BITTORRENT
 static void gatherProgressBitTorrent
-(BDE& entryDict, const BDE& torrentAttrs, const SharedHandle<BtRegistry>& btreg)
+(BDE& entryDict, const BDE& torrentAttrs, const BtObject& btObject)
 {
   const std::string& infoHash = torrentAttrs[bittorrent::INFO_HASH].s();
   entryDict["infoHash"] = util::toHex(infoHash);
 
-  SharedHandle<PeerStorage> peerStorage = btreg->get(infoHash)._peerStorage;
-  assert(!peerStorage.isNull());
+  if(!btObject.isNull()) {
+    SharedHandle<PeerStorage> peerStorage = btObject._peerStorage;
+    assert(!peerStorage.isNull());
 
-  std::deque<SharedHandle<Peer> > peers;
-  peerStorage->getActivePeers(peers);
-  entryDict["numSeeders"] = countSeeder(peers.begin(), peers.end());
+    std::deque<SharedHandle<Peer> > peers;
+    peerStorage->getActivePeers(peers);
+    entryDict["numSeeders"] = countSeeder(peers.begin(), peers.end());
+  }
 }
 
 static void gatherPeer(BDE& peers, const SharedHandle<PeerStorage>& ps)
@@ -335,8 +337,8 @@ static void gatherProgress
   if(group->getDownloadContext()->hasAttribute(bittorrent::BITTORRENT)) {
     const BDE& torrentAttrs =
       group->getDownloadContext()->getAttribute(bittorrent::BITTORRENT);
-    SharedHandle<BtRegistry> btreg = e->getBtRegistry();
-    gatherProgressBitTorrent(entryDict, torrentAttrs, btreg);
+    BtObject btObject = e->getBtRegistry()->get(group->getGID());
+    gatherProgressBitTorrent(entryDict, torrentAttrs, btObject);
   }
 #endif // ENABLE_BITTORRENT
 }
@@ -463,14 +465,11 @@ BDE GetPeersXmlRpcMethod::process
   }
   BDE peers = BDE::list();
   if(group->getDownloadContext()->hasAttribute(bittorrent::BITTORRENT)) {
-    SharedHandle<BtRegistry> btreg = e->getBtRegistry();
-    const BDE& torrentAttrs =
-      group->getDownloadContext()->getAttribute(bittorrent::BITTORRENT);
-    SharedHandle<PeerStorage> peerStorage =
-      btreg->get(torrentAttrs[bittorrent::INFO_HASH].s())._peerStorage;
-    assert(!peerStorage.isNull());
-    BDE entry = BDE::dict();
-    gatherPeer(peers, peerStorage);
+    BtObject btObject = e->getBtRegistry()->get(group->getGID());
+    if(!btObject.isNull()) {
+      assert(!btObject._peerStorage.isNull());
+      gatherPeer(peers, btObject._peerStorage);
+    }
   }
   return peers;
 }