1 files changed, 9 insertions, 3 deletions
diff --git a/src/DefaultExtensionMessageFactory.cc b/src/DefaultExtensionMessageFactory.cc
index 2295a70e..b8ccee7e 100644
--- a/src/DefaultExtensionMessageFactory.cc
+++ b/src/DefaultExtensionMessageFactory.cc
@@ -40,8 +40,8 @@
 #include "UTPexExtensionMessage.h"
 #include "LogFactory.h"
 #include "Logger.h"
-#include "BtRegistry.h"
 #include "StringFormat.h"
+#include "PeerStorage.h"
 
 namespace aria2 {
 
@@ -76,8 +76,8 @@ DefaultExtensionMessageFactory::createMessage(const unsigned char* data, size_t
     if(extensionName == "ut_pex") {
       // uTorrent compatible Peer-Exchange
       UTPexExtensionMessageHandle m =
-	UTPexExtensionMessage::create(_btContext, data, length);
-      m->setBtContext(_btContext);
+	UTPexExtensionMessage::create(data, length);
+      m->setPeerStorage(_peerStorage);
       return m;
     } else {
       throw DlAbortEx
@@ -92,6 +92,12 @@ void DefaultExtensionMessageFactory::setBtContext(const BtContextHandle& btConte
   _btContext = btContext;
 }
 
+void DefaultExtensionMessageFactory::setPeerStorage
+(const SharedHandle<PeerStorage>& peerStorage)
+{
+  _peerStorage = peerStorage;
+}
+
 void DefaultExtensionMessageFactory::setPeer(const PeerHandle& peer)
 {
   _peer = peer;