6 files changed, 11 insertions, 10 deletions

diff --git a/src/Network/BitTorrent/Client.hs b/src/Network/BitTorrent/Client.hs
index d21b4d1e..b9a59f45 100644
--- a/src/Network/BitTorrent/Client.hs
+++ b/src/Network/BitTorrent/Client.hs
@@ -114,7 +114,7 @@ initClient opts @ Options {..} logFun = do
   let mkEx = Exchange.newManager (exchangeOptions pid opts) (connHandler tmap)
   (_, emgr) <- allocate mkEx Exchange.closeManager
 
-  let mkNode = DHT.newNode defaultHandlers def optNodeAddr logFun
+  let mkNode = DHT.newNode defaultHandlers def optNodeAddr logFun Nothing
   (_, node) <- allocate mkNode DHT.closeNode
 
   resourceMap <- getInternalState
diff --git a/src/Network/BitTorrent/DHT.hs b/src/Network/BitTorrent/DHT.hs
index 39b33478..7340b854 100644
--- a/src/Network/BitTorrent/DHT.hs
+++ b/src/Network/BitTorrent/DHT.hs
@@ -86,7 +86,7 @@ dht :: Address ip
     -> IO a        -- ^ result.
 dht opts addr action = do
   runStderrLoggingT $ LoggingT $ \ logger -> do
-    bracket (newNode defaultHandlers opts addr logger) closeNode $
+    bracket (newNode defaultHandlers opts addr logger Nothing) closeNode $
       \ node -> runDHT node action
 {-# INLINE dht #-}
 
diff --git a/src/Network/BitTorrent/DHT/Query.hs b/src/Network/BitTorrent/DHT/Query.hs
index ac53bd91..cb7d5c5f 100644
--- a/src/Network/BitTorrent/DHT/Query.hs
+++ b/src/Network/BitTorrent/DHT/Query.hs
@@ -160,7 +160,7 @@ search k action = do
     $(logWarnS) "search" "start query"
     responses <- lift $ queryParallel (action <$> batch)
     let (nodes, results) = partitionEithers responses
-    $(logWarnS) "search" "done query"
+    $(logWarnS) "search" ("done query more:" <> T.pack (show $ L.length nodes))
     leftover $ L.concat nodes
     mapM_ yield results
 
diff --git a/src/Network/BitTorrent/DHT/Session.hs b/src/Network/BitTorrent/DHT/Session.hs
index 9aa25866..38b3ed11 100644
--- a/src/Network/BitTorrent/DHT/Session.hs
+++ b/src/Network/BitTorrent/DHT/Session.hs
@@ -299,8 +299,9 @@ newNode :: Address ip
           -> Options          -- ^ various dht options;
           -> NodeAddr ip      -- ^ node address to bind;
           -> LogFun           -- ^
+          -> Maybe NodeId     -- ^ use this NodeId, if not given a new one is generated.
           -> IO (Node ip)     -- ^ a new DHT node running at given address.
-newNode hs opts naddr logger = do
+newNode hs opts naddr logger mbid = do
     s <- createInternalState
     runInternalState initNode s
       `onException` closeInternalState s
@@ -311,7 +312,7 @@ newNode hs opts naddr logger = do
       s      <- getInternalState
       (_, m) <- allocate (newManager rpcOpts nodeAddr hs) closeManager
       liftIO $ do
-        myId   <- genNodeId
+        myId   <- maybe genNodeId return mbid
         node   <- Node opts myId s m
                 <$> newMVar (nullTable myId (optBucketCount opts))
                 <*> newTVarIO def
diff --git a/tests/Network/BitTorrent/DHT/QuerySpec.hs b/tests/Network/BitTorrent/DHT/QuerySpec.hs
index 81c3b45b..6b3b5485 100644
--- a/tests/Network/BitTorrent/DHT/QuerySpec.hs
+++ b/tests/Network/BitTorrent/DHT/QuerySpec.hs
@@ -28,7 +28,7 @@ nullLogger _ _ _ _ = return ()
 
 simpleDHT :: [NodeHandler IPv4] -> DHT IPv4 a -> IO a
 simpleDHT hs m =
-  bracket (newNode hs def myAddr nullLogger) closeNode $ \ node ->
+  bracket (newNode hs def myAddr nullLogger Nothing) closeNode $ \ node ->
     runDHT node m
 
 getBootInfo :: IO (NodeInfo IPv4)
@@ -102,4 +102,4 @@ spec = parallel $ do
             sourceList [startNodes] $=
               search entryHash (getPeersQ entryHash) $=
                 CL.concat $$ CL.take 10
-          L.length peers `shouldBe` 10
\ No newline at end of file
+          L.length peers `shouldBe` 10
diff --git a/tests/Network/BitTorrent/DHT/SessionSpec.hs b/tests/Network/BitTorrent/DHT/SessionSpec.hs
index 8c536352..ef844b31 100644
--- a/tests/Network/BitTorrent/DHT/SessionSpec.hs
+++ b/tests/Network/BitTorrent/DHT/SessionSpec.hs
@@ -26,7 +26,7 @@ myAddr = "127.0.0.1:60000"
 
 simpleDHT :: DHT IPv4 a -> IO a
 simpleDHT m =
-  bracket (newNode defaultHandlers def myAddr nullLogger) closeNode $ \ node ->
+  bracket (newNode defaultHandlers def myAddr nullLogger Nothing) closeNode $ \ node ->
     runDHT node m
 
 isRight :: Either a b -> Bool
@@ -43,7 +43,7 @@ spec :: Spec
 spec = do
   describe "session" $ do
     it "is active until closeNode called" $ do
-      node <- newNode [] def myAddr nullLogger
+      node <- newNode [] def myAddr nullLogger Nothing
       runDHT node monadActive `shouldReturn` True
       runDHT node monadActive `shouldReturn` True
       closeNode node
@@ -106,4 +106,4 @@ spec = do
 
     describe "(<@>) operator" $ do
       it "" $
-         pending
\ No newline at end of file
+         pending