3 files changed, 27 insertions, 4 deletions

diff --git a/src/Network/BitTorrent/DHT.hs b/src/Network/BitTorrent/DHT.hs
index bb9be560..dbedf801 100644
--- a/src/Network/BitTorrent/DHT.hs
+++ b/src/Network/BitTorrent/DHT.hs
@@ -60,6 +60,7 @@ import Control.Exception
 import Data.ByteString as BS
 import Data.Conduit as C
 import Data.Conduit.List as C
+import Data.Serialize
 import Network.Socket
 
 import Data.Torrent
@@ -182,15 +183,21 @@ isBootstrapped = T.full <$> getTable
 --
 --   This is blocking operation, use
 --   'Control.Concurrent.Async.Lifted.async' if needed.
-restore :: ByteString -> IO (Node ip)
-restore = error "DHT.restore: not implemented"
+restore :: Address ip => ByteString -> DHT ip ()
+restore bs = do
+    tblvar <- asks routingTable
+    case decode bs of
+        Right tbl -> restoreTable tbl
+        Left _ -> return ()
 
 -- | Serialize current DHT session to byte string.
 --
 --   This is blocking operation, use
 -- 'Control.Concurrent.Async.Lifted.async' if needed.
-snapshot :: DHT ip ByteString
-snapshot = error "DHT.snapshot: not implemented"
+snapshot :: Address ip => DHT ip ByteString
+snapshot = do
+    tbl <- getTable
+    return $ encode tbl
 
 {-----------------------------------------------------------------------
 --  Operations
diff --git a/src/Network/BitTorrent/DHT/Query.hs b/src/Network/BitTorrent/DHT/Query.hs
index c5fcccb4..7f20ad6d 100644
--- a/src/Network/BitTorrent/DHT/Query.hs
+++ b/src/Network/BitTorrent/DHT/Query.hs
@@ -41,6 +41,7 @@ module Network.BitTorrent.DHT.Query
        , publish
 
        -- ** Routing table
+       , restoreTable
        , insertNode
        , refreshNodes
 
@@ -266,3 +267,10 @@ queryNode addr q = do
 q <@> addr = snd <$> queryNode addr q
 {-# INLINE (<@>) #-}
 
+restoreTable :: Address ip => Table ip -> DHT ip ()
+restoreTable tbl = do
+  tblvar <- asks routingTable
+  tbl0 <- liftIO $ takeMVar tblvar
+  mb <- routing $ merge tbl tbl0
+  maybe (return ()) (liftIO . putMVar tblvar) mb
+
diff --git a/src/Network/BitTorrent/DHT/Routing.hs b/src/Network/BitTorrent/DHT/Routing.hs
index fee52380..68edef56 100644
--- a/src/Network/BitTorrent/DHT/Routing.hs
+++ b/src/Network/BitTorrent/DHT/Routing.hs
@@ -46,6 +46,7 @@ module Network.BitTorrent.DHT.Routing
          -- * Construction
        , Network.BitTorrent.DHT.Routing.nullTable
        , Network.BitTorrent.DHT.Routing.insert
+       , Network.BitTorrent.DHT.Routing.merge
 
          -- * Conversion
        , Network.BitTorrent.DHT.Routing.TableEntry
@@ -457,3 +458,10 @@ toBucketList (One   b t) = b : toBucketList t
 
 toList :: Eq ip => Table ip -> [[TableEntry ip]]
 toList = L.map (L.map tableEntry . PSQ.toList) . toBucketList
+
+merge :: Eq ip => Table ip -> Table ip -> Routing ip (Table ip)
+merge a b = do
+    let ns = concatMap PSQ.toList $ toBucketList a
+    -- TODO: merge timestamps as well and let refresh take care of ping.
+    as <- filterM (needPing . nodeAddr . PSQ.key) ns
+    foldM (flip $ Network.BitTorrent.DHT.Routing.insert) b $ map PSQ.key as