1 files changed, 52 insertions, 0 deletions

diff --git a/kad/src/Network/Kademlia/Persistence.hs b/kad/src/Network/Kademlia/Persistence.hs
new file mode 100644
index 00000000..32ec169d
--- /dev/null
+++ b/kad/src/Network/Kademlia/Persistence.hs
@@ -0,0 +1,52 @@
+{-# LANGUAGE NamedFieldPuns            #-}
+{-# LANGUAGE OverloadedStrings         #-}
+module Network.Kademlia.Persistence where
+
+import Network.Kademlia.CommonAPI
+import Network.Kademlia.Routing as R
+
+import Control.Concurrent.STM
+import qualified Data.Aeson           as J
+         ;import Data.Aeson           as J (FromJSON)
+import qualified Data.ByteString.Lazy as L
+import qualified Data.HashMap.Strict  as HashMap
+import Data.List
+import qualified Data.Vector          as V
+import System.IO.Error
+
+saveNodes :: String -> DHT -> IO ()
+saveNodes netname DHT{dhtBuckets} = do
+    bkts <- atomically $ readTVar dhtBuckets
+    let ns = map fst $ concat $ R.toList bkts
+        bs = J.encode ns
+        fname = nodesFileName netname
+    L.writeFile fname bs
+
+loadNodes :: FromJSON ni => String -> IO [ni]
+loadNodes netname = do
+    let fname = nodesFileName netname
+    attempt <- tryIOError $ do
+        J.decode <$> L.readFile fname
+         >>= maybe (ioError $ userError "Nothing") return
+    either (const $ fallbackLoad fname) return attempt
+
+nodesFileName :: String -> String
+nodesFileName netname = netname ++ "-nodes.json"
+
+fallbackLoad :: FromJSON t => FilePath -> IO [t]
+fallbackLoad fname = do
+    attempt <- tryIOError $ do
+        J.decode <$> L.readFile fname
+         >>= maybe (ioError $ userError "Nothing") return
+    let go r = do
+            let m = HashMap.lookup "nodes" (r :: J.Object)
+                ns0 = case m of Just (J.Array v) -> V.toList v
+                                Nothing        -> []
+                ns1 = zip (map J.fromJSON ns0) ns0
+                issuc (J.Error _,_) = False
+                issuc _             = True
+                (ss,fs) = partition issuc ns1
+                ns = map (\(J.Success n,_) -> n) ss
+            mapM_ print (map snd fs)  >> return ns
+    either (const $ return []) go attempt
+