3 files changed, 28 insertions, 23 deletions

diff --git a/examples/dhtd.hs b/examples/dhtd.hs
index 7d3661e6..837cb210 100644
--- a/examples/dhtd.hs
+++ b/examples/dhtd.hs
@@ -195,6 +195,7 @@ data DHT = forall nid ni. ( Show     ni
     , dhtParseId       :: String -> Either String nid
     , dhtSearches      :: TVar (Map.Map (String,nid) (DHTSearch nid ni))
     , dhtFallbackNodes :: IO [ni]
+    , dhtBootstrap     :: [ni] -> [ni] -> IO ()
     }
 
 nodesFileName :: String -> String
@@ -1013,7 +1014,7 @@ main = do
         "" -> return (return (), Map.empty,return [],[])
         p  -> do
             addr <- getBindAddress p (ip6bt opts)
-            (bt,btR) <- Mainline.newClient swarms addr
+            (bt,btR,btBootstrap4, btBootstrap6) <- Mainline.newClient swarms addr
             quitBt <- forkListener "bt" (clientNet bt)
             mainlineSearches <- atomically $ newTVar Map.empty
             peerPort <- atomically $ newTVar 6881 -- BitTorrent client TCP port.
@@ -1083,6 +1084,9 @@ main = do
                                     [ {- TODO -}
                                     ]
                     , dhtSecretKey = return Nothing
+                    , dhtBootstrap = case wantip of
+                                        Want_IP4 -> btBootstrap4
+                                        Want_IP6 -> btBootstrap6
                     }
                 dhts = Map.fromList $
                         ("bt4", mainlineDHT Mainline.routing4 Want_IP4)
@@ -1107,11 +1111,11 @@ main = do
             addrTox <- getBindAddress toxport (ip6tox opts)
             hPutStrLn stderr $ "Supplied key: " ++ show (fmap (Tox.key2id . toPublic) (dhtkey opts))
             tox <- Tox.newTox keysdb addrTox (Just netCryptoSessionsState) (dhtkey opts)
-            quitTox <- Tox.forkTox tox
+            (quitTox, toxStrap4, toxStrap6) <- Tox.forkTox tox
 
             toxSearches <- atomically $ newTVar Map.empty
 
-            let toxDHT bkts = DHT
+            let toxDHT bkts wantip = DHT
                     { dhtBuckets = bkts (Tox.toxRouting tox)
                     , dhtPing    = Map.fromList
                                     [ ("ping", DHTPing
@@ -1262,11 +1266,14 @@ main = do
                                     [ {- TODO -}
                                     ]
                     , dhtSecretKey = return $ Just $ transportSecret (Tox.toxCryptoKeys tox)
+                    , dhtBootstrap = case wantip of
+                                        Want_IP4 -> toxStrap4
+                                        Want_IP6 -> toxStrap6
                     }
                 dhts = Map.fromList $
-                        ("tox4", toxDHT Tox.routing4)
+                        ("tox4", toxDHT Tox.routing4 Want_IP4)
                         : if ip6tox opts
-                            then [ ("tox6", toxDHT Tox.routing6) ]
+                            then [ ("tox6", toxDHT Tox.routing6 Want_IP6) ]
                             else []
                 ips :: IO [SockAddr]
                 ips = readExternals Tox.nodeAddr [ Tox.routing4 $ Tox.toxRouting tox
@@ -1309,7 +1316,8 @@ main = do
           $ \(netname, dht@DHT { dhtBuckets = bkts
                                , dhtQuery = qrys
                                , dhtPing = pings
-                               , dhtFallbackNodes = getBootstrapNodes }) -> do
+                               , dhtFallbackNodes = getBootstrapNodes
+                               , dhtBootstrap = bootstrap }) -> do
         btSaved <- loadNodes netname --  :: IO [Mainline.NodeInfo]
         putStrLn $ "Loaded "++show (length btSaved)++" nodes for "++netname++"."
         fallbackNodes <- getBootstrapNodes
@@ -1320,18 +1328,7 @@ main = do
                          $ Map.lookup "ping" pings
         fork $ do
             myThreadId >>= flip labelThread ("bootstrap."++netname)
-            case Map.lookup "node" qrys of
-                Just DHTQuery { qsearch = srch } -> do
-                    case eqT of
-                        Just witness -> let strapper = BucketRefresher
-                                                { refreshSearch = isNodesSearch witness srch
-                                                , refreshBuckets = bkts
-                                                , refreshPing = ping
-                                                }
-                                         in bootstrap strapper btSaved fallbackNodes
-                        _            -> error $ "Missing node-search for "++netname++"."
-                    saveNodes netname dht
-                Nothing -> return ()
+            bootstrap btSaved fallbackNodes
         return ()
 
     waitForSignal
diff --git a/src/Network/BitTorrent/MainlineDHT.hs b/src/Network/BitTorrent/MainlineDHT.hs
index 3e7a0eda..47d5b3e1 100644
--- a/src/Network/BitTorrent/MainlineDHT.hs
+++ b/src/Network/BitTorrent/MainlineDHT.hs
@@ -538,7 +538,12 @@ mkNodeInfo nid addr = NodeInfo
     , nodePort = fromMaybe 0 $ sockAddrPort addr
     }
 
-newClient :: SwarmsDatabase -> SockAddr -> IO (MainlineClient, Routing)
+newClient :: SwarmsDatabase -> SockAddr
+    -> IO ( MainlineClient
+          , Routing
+          , [NodeInfo] -> [NodeInfo] -> IO ()
+          , [NodeInfo] -> [NodeInfo] -> IO ()
+          )
 newClient swarms addr = do
     udp <- udpTransport addr
     nid <- NodeId <$> getRandomBytes 20
@@ -651,7 +656,7 @@ newClient swarms addr = do
     refresh_thread4 <- forkPollForRefresh $ refresher4 routing
     refresh_thread6 <- forkPollForRefresh $ refresher6 routing
 
-    return (client, routing)
+    return (client, routing, bootstrap (refresher4 routing), bootstrap (refresher6 routing))
 
 -- | Modifies a purely random 'NodeId' to one that is related to a given
 -- routable address in accordance with BEP 42.
diff --git a/src/Network/Tox.hs b/src/Network/Tox.hs
index 5f691ef3..37802e3c 100644
--- a/src/Network/Tox.hs
+++ b/src/Network/Tox.hs
@@ -88,7 +88,7 @@ import GHC.TypeLits
 import Crypto.Tox
 import Data.Word64Map                        (fitsInInt)
 import qualified Data.Word64Map              (empty)
-import Network.Kademlia.Bootstrap            (forkPollForRefresh)
+import Network.Kademlia.Bootstrap            (forkPollForRefresh, bootstrap)
 import Network.Tox.Crypto.Transport          (NetCrypto)
 import Network.Tox.Crypto.Handlers           (cryptoNetHandler, newSessionsState, defaultUnRecHook, defaultCryptoDataHooks, NetCryptoSessions(..))
 import qualified Network.Tox.DHT.Handlers    as DHT
@@ -354,7 +354,7 @@ newTox keydb addr mbSessionsState suppliedDHTKey = do
 onionTimeout :: Tox -> DHT.TransactionId -> Onion.OnionDestination RouteId -> STM (Onion.OnionDestination RouteId, Int)
 onionTimeout Tox { toxOnionRoutes = or } (DHT.TransactionId n8 _) od = lookupTimeout or n8 od
 
-forkTox :: Tox -> IO (IO ())
+forkTox :: Tox -> IO (IO (), [NodeInfo] -> [NodeInfo] -> IO (), [NodeInfo] -> [NodeInfo] -> IO ())
 forkTox tox = do
     _ <- forkListener "toxCrypto" (toxCrypto tox)
     _ <- forkListener "toxToRoute" (toxToRoute tox)
@@ -362,5 +362,8 @@ forkTox tox = do
     quit <- forkListener "toxDHT" (clientNet $ toxDHT tox)
     forkPollForRefresh (DHT.refresher4 $ toxRouting tox)
     forkPollForRefresh (DHT.refresher6 $ toxRouting tox)
-    return quit
+    return ( quit
+           , bootstrap (DHT.refresher4 $ toxRouting tox)
+           , bootstrap (DHT.refresher6 $ toxRouting tox)
+           )