Added startup nodes to CommonAPI search query.

2 files changed, 25 insertions, 11 deletions

diff --git a/dht/examples/dhtd.hs b/dht/examples/dhtd.hs
index f9dc777d..3095e7b4 100644
--- a/dht/examples/dhtd.hs
+++ b/dht/examples/dhtd.hs
@@ -295,12 +295,11 @@ forkSearch ::
     -> nid
     -> DHTQuery nid ni
     -> TVar (Map.Map (String,nid) (DHTSearch nid ni))
-    -> TVar (BucketList ni)
+    -> [ni]
     -> ThreadId
     -> TVar (Maybe (IO ()))
     -> STM ()
-forkSearch method nid DHTQuery{qsearch,qshowTok,qshowR} dhtSearches dhtBuckets  tid kvar = do
-    ns <- R.kclosest (searchSpace qsearch) (searchK qsearch) nid <$> readTVar dhtBuckets
+forkSearch method nid DHTQuery{qsearch,qshowTok,qshowR} dhtSearches ns tid kvar = do
     st <- newSearch qsearch nid ns
     results <- newTVar Set.empty
     let storeResult r = modifyTVar' results (Set.insert (qshowR r))
@@ -1046,10 +1045,11 @@ clientSession s@Session{..} sock cnum h = do
                                 -- STM action decides not to add a new search.  This is so that
                                 -- I can store the ThreadId into new DHTSearch structure.
                                 tid <- fork $ join $ atomically (readTVar kvar >>= maybe retry return)
+                                ns <- case qry of DHTQuery{qbootNodes} -> qbootNodes nid
                                 join $ atomically $ do
                                     schs <- readTVar dhtSearches
                                     case Map.lookup (method,nid) schs of
-                                        Nothing  -> do forkSearch method nid qry dhtSearches (refreshBuckets dhtBuckets) tid kvar
+                                        Nothing  -> do forkSearch method nid qry dhtSearches ns tid kvar
                                                        return $ presentSearches
                                         Just sch -> do writeTVar kvar (Just $ return ())
                                                        return $ reportSearchResults method h sch
@@ -1503,8 +1503,8 @@ initTox runio opts ssvar keysdb mbxmpp invc = case porttox opts of
 
             tcpSearches <- atomically $ newTVar Map.empty
 
-            let toxDHT bkts wantip = DHT
-                    { dhtBuckets = bkts (Tox.toxRouting tox)
+            let toxDHT bkts wantip = let toxBkts = bkts (Tox.toxRouting tox) in DHT
+                    { dhtBuckets = toxBkts
                     , dhtPing    = Map.fromList
                                     [ ("ping", DHTPing
                                             { pingQuery = noArgPing $ fmap (bool Nothing (Just ())) . Tox.pingUDP (Tox.toxDHT tox)
@@ -1521,7 +1521,7 @@ initTox runio opts ssvar keysdb mbxmpp invc = case porttox opts of
                                             , pingShowResult = show
                                             })]
                     , dhtQuery   = Map.fromList
-                                    [ ("node", DHTQuery
+                                    [ ("node", fix $ \q -> DHTQuery
                                         { qsearch = Tox.nodeSearch (Tox.toxDHT tox)
                                                                    (Tox.nodesOfInterest $ Tox.toxRouting tox)
                                         , qhandler = (\ni -> fmap Tox.unwrapNodes
@@ -1529,8 +1529,9 @@ initTox runio opts ssvar keysdb mbxmpp invc = case porttox opts of
                                                                 . Tox.GetNodes)
                                         , qshowR = show -- NodeInfo
                                         , qshowTok = (const Nothing)
+                                        , qbootNodes = genericBootNodes (refreshBuckets toxBkts) q
                                         })
-                                    , ("toxid", DHTQuery
+                                    , ("toxid", fix $ \q -> DHTQuery
                                         { qsearch = Tox.toxQSearch tox
                                         , qhandler = -- qhandler :: ni -> nid -> IO ([ni], [r], tok)
                                                          (\ni nid ->
@@ -1544,6 +1545,7 @@ initTox runio opts ssvar keysdb mbxmpp invc = case porttox opts of
                                                                                   (Tox.AnnounceRequest zeros32 nid Tox.zeroID))
                                         , qshowR  = show -- Rendezvous
                                         , qshowTok = Just . show -- Nonce32
+                                        , qbootNodes = genericBootNodes (refreshBuckets toxBkts) q
                                         })
                                     ]
                     , dhtParseId  = readEither :: String -> Either String Tox.NodeId
@@ -1690,7 +1692,7 @@ initTox runio opts ssvar keysdb mbxmpp invc = case porttox opts of
                                     { pingQuery = noArgPing $ TCP.tcpPing (TCP.tcpClient tcpclient)
                                     , pingShowResult = show
                                     }
-                    , dhtQuery = Map.singleton "node" DHTQuery
+                    , dhtQuery = Map.singleton "node" $ fix $ \q -> DHTQuery
                                     { qsearch = TCP.nodeSearch tcpprober tcpclient
                                     , qhandler = \ni nid -> do
                                         ns <- R.kclosest (searchSpace $ TCP.nodeSearch tcpprober tcpclient)
@@ -1700,6 +1702,7 @@ initTox runio opts ssvar keysdb mbxmpp invc = case porttox opts of
                                         return (ns,ns,Just ())
                                     , qshowR = show -- TCP.NodeInfo
                                     , qshowTok = (const Nothing)
+                                    , qbootNodes = genericBootNodes (refreshBuckets tcpRefresher) q
                                     }
                     , dhtAnnouncables = Map.empty
                     , dhtParseId = readEither :: String -> Either String Tox.NodeId
@@ -1808,19 +1811,20 @@ main = do
                                         , pingShowResult = show
                                         }
                     , dhtQuery   = Map.fromList
-                                    [ ("node", DHTQuery
+                                    [ ("node", fix $ \q -> DHTQuery
                                         { qsearch = (Mainline.nodeSearch bt)
                                         , qhandler = (\ni -> fmap Mainline.unwrapNodes
                                                                 . Mainline.findNodeH btR ni
                                                                 . flip Mainline.FindNode (Just Want_Both))
                                         , qshowR = show
                                         , qshowTok = (const Nothing)
+                                        , qbootNodes = genericBootNodes (refreshBuckets $ bkts btR) q
                                         })
                                     -- Maybe (sr :~: pr, stok :~: ptok, sni :~: pni )
                                     -- sr = InfoHash
                                     -- stok = Token
                                     -- sni = NodeInfo
-                                    , ("peer", DHTQuery
+                                    , ("peer", fix $ \q -> DHTQuery
                                         { qsearch = (Mainline.peerSearch bt)
                                         , qhandler = (\ni -> fmap Mainline.unwrapPeers
                                                                 . Mainline.getPeersH btR swarms ni
@@ -1828,6 +1832,7 @@ main = do
                                                                 . (read . show)) -- TODO: InfoHash -> NodeId
                                         , qshowR = (show . pPrint)
                                         , qshowTok = (Just . show)
+                                        , qbootNodes = genericBootNodes (refreshBuckets $ bkts btR) q
                                         })
                                     ]
                     , dhtParseId  = readEither :: String -> Either String Mainline.NodeId
diff --git a/kad/src/Network/Kademlia/CommonAPI.hs b/kad/src/Network/Kademlia/CommonAPI.hs
index bcbfe9d8..4de7909d 100644
--- a/kad/src/Network/Kademlia/CommonAPI.hs
+++ b/kad/src/Network/Kademlia/CommonAPI.hs
@@ -1,4 +1,5 @@
 {-# LANGUAGE ExistentialQuantification #-}
+{-# LANGUAGE NamedFieldPuns            #-}
 module Network.Kademlia.CommonAPI
     ( module Network.Kademlia.CommonAPI
     , refreshBuckets
@@ -58,8 +59,16 @@ data DHTQuery nid ni = forall addr r tok qk.
     , qhandler    :: ni -> nid -> IO ([ni], [r], Maybe tok) -- ^ Invoked on local node, when there is no query destination.
     , qshowR      :: r -> String
     , qshowTok    :: tok -> Maybe String
+    , qbootNodes  :: nid -> IO [ni]
     }
 
+-- Can be used to initialize qbootNodes like this:
+-- fix \q -> DHTQuery { ... , qbootNodes = genericBootNodes bkts q }
+genericBootNodes :: Ord nid => TVar (BucketList ni) -> DHTQuery nid ni -> nid -> IO [ni]
+genericBootNodes dhtBuckets DHTQuery{qsearch} nid =
+    atomically $ R.kclosest (searchSpace qsearch) (searchK qsearch) nid
+                 <$> readTVar dhtBuckets
+
 data DHTAnnouncable nid = forall dta tok ni r.
                         ( Show r
                         , Typeable dta -- information being announced