1 files changed, 33 insertions, 28 deletions

diff --git a/src/Network/BitTorrent/DHT.hs b/src/Network/BitTorrent/DHT.hs
index 1c62a0e0..3867d182 100644
--- a/src/Network/BitTorrent/DHT.hs
+++ b/src/Network/BitTorrent/DHT.hs
@@ -33,7 +33,6 @@ module Network.BitTorrent.DHT
 
          -- * Initialization
        , snapshot
-       , restore
 
          -- * Operations
        , Network.BitTorrent.DHT.lookup
@@ -162,24 +161,39 @@ 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
+bootstrap :: Address ip => Maybe BS.ByteString -> [NodeAddr ip] -> DHT ip ()
+bootstrap mbs startNodes = do
+  restored <-
+    case decode <$> mbs of
+        Just (Right tbl) -> return (T.toList tbl)
+        Just (Left e)    -> do $(logWarnS) "restore" (Text.pack e)
+                               return []
+        Nothing          -> return []
+
   $(logInfoS) "bootstrap" "Start node bootstrapping"
-  nid <- asks thisNodeId
-  let searchAll aliveNodes
-        = C.sourceList [aliveNodes] $= search nid (findNodeQ nid) $$ C.consume
+  let searchAll aliveNodes = do
+        nid <- myNodeIdAccordingTo (error "FIXME")
+        C.sourceList [aliveNodes] $= search nid (findNodeQ nid) $$ C.consume
+  input_nodes <- (restored ++) . T.toList <$> getTable
   -- 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 ()
+  do let knowns = map (map $ nodeAddr . fst) input_nodes
+     -- Below, we reverse the nodes since the table serialization puts the
+     -- nearest nodes last and we want to choose a similar node id to bootstrap
+     -- faster.
+     (alive_knowns,_) <- unzip <$> queryParallel (pingQ <$> reverse (concat knowns))
+     b <- isBootstrapped
+     -- If our cached nodes are alive and our IP address did not change, it's possible
+     -- we are already bootsrapped, so no need to do any searches.
+     when (not b) $ do
+         nss <- searchAll $ take 2 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,_) <- unzip <$> queryParallel (pingQ <$> startNodes)
+            _ <- searchAll $ take 2 aliveNodes
+            return ()
   -- Step 2: Repeatedly refresh incomplete buckets until the table is full.
   maxbuckets <- asks $ optBucketCount . options
   flip fix 0 $ \loop icnt -> do
@@ -195,6 +209,7 @@ bootstrap startNodes = do
                     p:ps                         -> p:unfull ps
                     []                           -> []
       forM_ us $ \(is_last,(index,_)) -> do
+        nid <- myNodeIdAccordingTo (error "FIXME")
         sample <- liftIO $ genBucketSample nid (bucketRange index is_last)
         $(logDebugS) "bootstrapping"
             $ "BOOTSTRAP sample"
@@ -213,23 +228,13 @@ bootstrap startNodes = do
 --
 --   This operation do not block.
 --
-isBootstrapped :: DHT ip Bool
+isBootstrapped :: Eq ip => DHT ip Bool
 isBootstrapped = T.full <$> getTable
 
 {-----------------------------------------------------------------------
 -- Initialization
 -----------------------------------------------------------------------}
 
--- | Load previous session. (corrupted - exception/ignore ?)
---
---   This is blocking operation, use
---   'Control.Concurrent.Async.Lifted.async' if needed.
-restore :: Address ip => BS.ByteString -> DHT ip ()
-restore bs = do
-    case decode bs of
-        Right tbl -> restoreTable tbl
-        Left e -> $(logWarnS) "restore" (Text.pack e)
-
 -- | Serialize current DHT session to byte string.
 --
 --   This is blocking operation, use