3 files changed, 25 insertions, 19 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
diff --git a/src/Network/DatagramServer.hs b/src/Network/DatagramServer.hs
index f212ffdf..2140e2cd 100644
--- a/src/Network/DatagramServer.hs
+++ b/src/Network/DatagramServer.hs
@@ -196,7 +196,7 @@ data Manager h raw msg = Manager
   , listenerThread     :: !(MVar ThreadId)
   , transactionCounter :: {-# UNPACK #-} !TransactionCounter
   , pendingCalls       :: {-# UNPACK #-} !(PendingCalls msg raw)
-  , handlers           :: [Handler h msg raw]
+  , handlers           :: [Handler h msg raw] -- TODO delete this, it's not used
   , logMsg             :: Char -> String -> T.Text -> IO ()
   }
 
@@ -462,8 +462,8 @@ dispatchHandler :: ( Eq (QueryMethod msg)
                    , Show (QueryMethod msg)
                    , Serialize (TransactionID msg)
                    , Envelope msg
-                   ) => Manager IO raw msg -> QueryMethod msg -> msg raw -> SockAddr -> IO (KResult msg raw)
-dispatchHandler mgr@Manager{..} meth q addr = do
+                   ) => Manager IO raw msg -> [Handler IO msg raw] -> QueryMethod msg -> msg raw -> SockAddr -> IO (KResult msg raw)
+dispatchHandler mgr handlers meth q addr = do
   case L.lookup meth handlers of
     Nothing -> return $ Left $ KError MethodUnknown ("Unknown method " <> BC.pack (show meth)) (envelopeTransaction q)
     Just h  -> runHandler mgr meth h addr q
@@ -483,10 +483,10 @@ handleQuery :: ( WireFormat raw msg
                , Eq (QueryMethod msg)
                , Show (QueryMethod msg)
                , Serialize (TransactionID msg)
-               ) => Manager IO raw msg -> QueryMethod msg -> raw -> msg raw -> SockAddr -> IO ()
-handleQuery mgr@Manager{..} meth raw q addr = void $ fork $ do
+               ) => Manager IO raw msg -> [Handler IO msg raw] -> QueryMethod msg -> raw -> msg raw -> SockAddr -> IO ()
+handleQuery mgr@Manager{..} hs meth raw q addr = void $ fork $ do
   myThreadId >>= liftIO . flip labelThread "KRPC.handleQuery"
-  res <- dispatchHandler mgr meth q addr
+  res <- dispatchHandler mgr hs meth q addr
   let  res' = either buildError Just res
        ctx = error "TODO TOX ToxCipherContext 2 or () for Mainline"
        resbs = fmap (encodeHeaders ctx) res' :: Maybe BS.ByteString
@@ -520,8 +520,8 @@ listener :: forall raw msg.
             , Eq (QueryMethod msg)
             , Show (QueryMethod msg)
             , Serialize (TransactionID msg)
-            ) => Manager IO raw msg -> Protocol raw msg -> IO ()
-listener mgr@Manager{..} p = do
+            ) => Manager IO raw msg -> [Handler IO msg raw] -> Protocol raw msg -> IO ()
+listener mgr@Manager{..} hs p = do
   fix $ \again -> do
     let ctx = error "TODO TOX ToxCipherContext or () for Mainline"
     (bs, addr) <- liftIO $ do
@@ -532,7 +532,7 @@ listener mgr@Manager{..} p = do
                        return () -- Without transaction id, error message isn't very useful.
       Right (raw,m) ->
         case envelopeClass m of
-            Query meth  -> handleQuery mgr meth (raw::raw) m         addr
+            Query meth  -> handleQuery mgr hs meth (raw::raw) m         addr
             Response _  -> handleResponse mgr   raw        (Right m) addr
             Error e     -> handleResponse mgr   raw        (Left  e) addr
 
@@ -551,10 +551,10 @@ listen :: ( WireFormat raw msg
           , Eq (QueryMethod msg)
           , Show (QueryMethod msg)
           , Serialize (TransactionID msg)
-          ) => Manager IO raw msg -> Protocol raw msg -> IO ()
-listen mgr@Manager{..} p = do
+          ) => Manager IO raw msg -> [Handler IO msg raw] ->  Protocol raw msg -> IO ()
+listen mgr@Manager{..} hs p = do
   tid <- fork $ do
     myThreadId >>= liftIO . flip labelThread "KRPC.listen"
-    listener mgr p `Lifted.finally`
+    listener mgr hs p `Lifted.finally`
       liftIO (takeMVar listenerThread)
   liftIO $ putMVar listenerThread tid