1 files changed, 52 insertions, 14 deletions

diff --git a/src/Network/BitTorrent/DHT.hs b/src/Network/BitTorrent/DHT.hs
index dbedf801..a5b4612f 100644
--- a/src/Network/BitTorrent/DHT.hs
+++ b/src/Network/BitTorrent/DHT.hs
@@ -57,17 +57,21 @@ import Control.Applicative
 import Control.Monad.Logger
 import Control.Monad.Reader
 import Control.Exception
-import Data.ByteString as BS
+import qualified Data.ByteString as BS
 import Data.Conduit as C
-import Data.Conduit.List as C
+import qualified Data.Conduit.List as C
 import Data.Serialize
 import Network.Socket
+import Text.PrettyPrint.HughesPJClass as PP (pPrint,render)
 
 import Data.Torrent
 import Network.BitTorrent.Address
 import Network.BitTorrent.DHT.Query
 import Network.BitTorrent.DHT.Session
-import Network.BitTorrent.DHT.Routing as T
+import Network.BitTorrent.DHT.Routing as T hiding (null)
+import qualified Data.Text as Text
+import Data.Monoid
+
 
 {-----------------------------------------------------------------------
 --  DHT types
@@ -156,15 +160,50 @@ resolveHostName NodeAddr {..} = do
 --
 --   This operation do block, use
 --   'Control.Concurrent.Async.Lifted.async' if needed.
---
 bootstrap :: Address ip => [NodeAddr ip] -> DHT ip ()
 bootstrap startNodes = do
   $(logInfoS) "bootstrap" "Start node bootstrapping"
   nid <- asks thisNodeId
-  -- TODO filter duplicated in startNodes list
-  -- TODO retransmissions for startNodes
-  aliveNodes <- queryParallel (pingQ <$> startNodes)
-  _ <- sourceList [aliveNodes] $= search nid (findNodeQ nid) $$ C.consume
+  let searchAll aliveNodes
+        = C.sourceList [aliveNodes] $= search nid (findNodeQ nid) $$ C.consume
+  -- Step 1: Use iterative searches to flesh out the table..
+  do knowns <- map (map $ nodeAddr . fst) . T.toList <$> getTable
+     alive_knowns <- queryParallel (pingQ <$> concat knowns)
+     nss <- searchAll alive_knowns
+     -- We only use the supplied bootstrap nodes when we don't know of any
+     -- others to try.
+     when (null nss) $ do
+        -- TODO filter duplicated in startNodes list
+        -- TODO retransmissions for startNodes
+        aliveNodes <- queryParallel (pingQ <$> startNodes)
+        _ <- searchAll aliveNodes
+        return ()
+  -- Step 2: Repeatedly refresh incomplete buckets until the table is full.
+  maxbuckets <- asks $ optBucketCount . options
+  flip fix 0 $ \loop icnt -> do
+      tbl <- getTable
+      let unfull = filter ((/=defaultBucketSize) . snd)
+          us = zip
+                -- is_last = True for the last bucket
+                (True:repeat False)
+                -- Only non-full buckets unless it is the last one and the
+                -- maximum number of buckets has not been reached.
+                $ case reverse $ zip [0..] $ T.shape tbl of
+                    p@(n,_):ps | n+1==maxbuckets -> unfull (p:ps)
+                    p:ps                         -> p:unfull ps
+                    []                           -> []
+      forM_ us $ \(is_last,(index,_)) -> do
+        sample <- liftIO $ genBucketSample nid (bucketRange index is_last)
+        $(logDebugS) "bootstrapping"
+            $ "BOOTSTRAP sample"
+              <> Text.pack (show (is_last,index,T.shape tbl))
+              <> " " <> Text.pack (render $ pPrint sample)
+        refreshNodes sample
+      $(logDebugS) "bootstrapping"
+        $ "BOOTSTRAP finished iteration "
+          <> Text.pack (show (icnt,T.shape tbl,us,defaultBucketSize))
+      when (not (null us) && icnt < div (3*maxbuckets) 2)
+        $ loop (succ icnt)
   $(logInfoS) "bootstrap" "Node bootstrapping finished"
 
 -- | Check if this node is already bootstrapped.
@@ -183,18 +222,17 @@ isBootstrapped = T.full <$> getTable
 --
 --   This is blocking operation, use
 --   'Control.Concurrent.Async.Lifted.async' if needed.
-restore :: Address ip => ByteString -> DHT ip ()
+restore :: Address ip => BS.ByteString -> DHT ip ()
 restore bs = do
-    tblvar <- asks routingTable
     case decode bs of
         Right tbl -> restoreTable tbl
-        Left _ -> return ()
+        Left e -> $(logWarnS) "restore" (Text.pack e)
 
 -- | Serialize current DHT session to byte string.
 --
 --   This is blocking operation, use
 -- 'Control.Concurrent.Async.Lifted.async' if needed.
-snapshot :: Address ip => DHT ip ByteString
+snapshot :: Address ip => DHT ip BS.ByteString
 snapshot = do
     tbl <- getTable
     return $ encode tbl
@@ -207,10 +245,10 @@ snapshot = do
 --
 --   This operation is incremental and do block.
 --
-lookup :: Address ip => InfoHash -> DHT ip `Source` [PeerAddr ip]
+lookup :: Address ip => InfoHash -> DHT ip `C.Source` [PeerAddr ip]
 lookup topic = do      -- TODO retry getClosest if bucket is empty
   closest <- lift $ getClosest topic
-  sourceList [closest] $= search topic (getPeersQ topic)
+  C.sourceList [closest] $= search topic (getPeersQ topic)
 
 -- TODO do not republish if the topic is already in announceSet