Propagated MonadKRPC deletion to Network.BitTorrent.DHT.

2 files changed, 13 insertions, 7 deletions

diff --git a/src/Network/BitTorrent/DHT.hs b/src/Network/BitTorrent/DHT.hs
index 51d92127..d9328cea 100644
--- a/src/Network/BitTorrent/DHT.hs
+++ b/src/Network/BitTorrent/DHT.hs
@@ -70,8 +70,11 @@ import Network.BitTorrent.DHT.Query
 import Network.BitTorrent.DHT.Session
 import Network.DHT.Routing as T hiding (null)
 import qualified Data.Text as Text
+import Data.Typeable
 import Data.Monoid
 import Network.DatagramServer.Mainline (KMessageOf)
+import qualified Network.DatagramServer as KRPC (listen, Protocol(..))
+
 
 
 {-----------------------------------------------------------------------
@@ -97,8 +100,12 @@ dht :: (Ord ip, Address ip)
     -> IO a        -- ^ result.
 dht opts addr logfilter action = do
   runStderrLoggingT $ filterLogger logfilter $ LoggingT $ \ logger -> do
-    bracket (newNode defaultHandlers opts addr logger Nothing) closeNode $
-      \ node -> runDHT node action
+    bracket (newNode opts addr logger Nothing) closeNode $
+      \ node -> runDHT node $ do
+            hs <- defaultHandlers logger
+            m <- asks manager
+            liftIO $ KRPC.listen m hs (KRPC.Protocol Proxy Proxy)
+            action
 {-# INLINE dht #-}
 
 {-----------------------------------------------------------------------
diff --git a/src/Network/BitTorrent/DHT/Session.hs b/src/Network/BitTorrent/DHT/Session.hs
index d8665773..d4794038 100644
--- a/src/Network/BitTorrent/DHT/Session.hs
+++ b/src/Network/BitTorrent/DHT/Session.hs
@@ -377,8 +377,8 @@ locFromCS cs = case getCallStack cs of
 -- 'closeNode' function, otherwise socket or other scarce resources may
 -- leak.
 newNode :: Address ip
-          => [NodeHandler] -- ^ handlers to run on accepted queries;
-          -> Options       -- ^ various dht options;
+          => -- [NodeHandler] -- ^ handlers to run on accepted queries;
+             Options       -- ^ various dht options;
           -> NodeAddr ip   -- ^ node address to bind;
           -> LogFun        -- ^ invoked on log messages;
 #ifdef VERSION_bencoding
@@ -387,7 +387,7 @@ newNode :: Address ip
           -> Maybe (NodeId Tox.Message)     -- ^ use this NodeId, if not given a new one is generated.
 #endif
           -> IO (Node ip)     -- ^ a new DHT node running at given address.
-newNode hs opts naddr logger mbid = do
+newNode opts naddr logger mbid = do
     s <- createInternalState
     runInternalState initNode s
       `onException` closeInternalState s
@@ -396,7 +396,7 @@ newNode hs opts naddr logger mbid = do
     nodeAddr = toSockAddr naddr
     initNode = do
       s      <- getInternalState
-      (_, m) <- allocate (newManager rpcOpts (logt logger) nodeAddr hs) closeManager
+      (_, m) <- allocate (newManager rpcOpts (logt logger) nodeAddr []) closeManager
       liftIO $ do
         myId   <- maybe genNodeId return mbid
         node   <- Node opts myId s m
@@ -405,7 +405,6 @@ newNode hs opts naddr logger mbid = do
                 <*> newTVarIO S.empty
                 <*> (newTVarIO =<< nullSessionTokens)
                 <*> pure logger
-        KRPC.listen m (KRPC.Protocol Proxy Proxy)
         return node
 
 -- | Some resources like listener thread may live for