4 files changed, 95 insertions, 27 deletions

diff --git a/examples/dhtd.hs b/examples/dhtd.hs
index 3df77190..19b45acb 100644
--- a/examples/dhtd.hs
+++ b/examples/dhtd.hs
@@ -18,6 +18,7 @@ import Data.Maybe
 import Data.String
 import qualified Data.ByteString as B (ByteString,writeFile,readFile)
         ; import Data.ByteString (ByteString)
+import qualified Data.ByteString.Char8 as B8
 import System.IO
 import System.IO.Error
 import Text.PrettyPrint.HughesPJClass
@@ -114,15 +115,16 @@ clientSession st signalQuit sock n h = do
     line <- map toLower . dropWhile isSpace <$> hGetLine h
     let cmd0 action = action >> clientSession st signalQuit sock n h
         cmd  action = cmd0 $ join $ runDHT st action
-    case line of
+        (c,args) = second (dropWhile isSpace) $ break isSpace line
+    case (c,args) of
 
-        "quit" -> hPutClient h "" >> hClose h
+        ("quit", _) -> hPutClient h "" >> hClose h
 
-        "stop" -> do hPutClient h "Terminating DHT Daemon."
-                     hClose h
-                     putMVar signalQuit ()
+        ("stop", _) -> do hPutClient h "Terminating DHT Daemon."
+                          hClose h
+                          putMVar signalQuit ()
 
-        "ls" -> cmd $ do
+        ("ls", _) -> cmd $ do
                     tbl <- getTable
                     t <- showTable
                     me <- myNodeIdAccordingTo (read "8.8.8.8:6881")
@@ -135,13 +137,13 @@ clientSession st signalQuit sock n h = do
                             , ("internet address", show ip)
                             , ("buckets", show $ R.shape tbl)]
                         ]
-        "external-ip" -> cmd $ do
+        ("external-ip", _) -> cmd $ do
                     ip <- routableAddress
                     return $ do
                     hPutClient h $ maybe "" (takeWhile (/=':') . show) ip
 
-        s | s=="swarms" || "swarms " `isPrefixOf` s -> cmd $ do
-                    let fltr = case dropWhile isSpace (drop 7 s) of
+        ("swarms", s) -> cmd $ do
+                    let fltr = case s of
                                     ('-':'v':cs) | all isSpace (take 1 cs)
                                       -> const True
                                     _ -> (\(h,c,n) -> c/=0 )
@@ -151,8 +153,8 @@ clientSession st signalQuit sock n h = do
                     return $ do
                     hPutClient h $ showReport r
 
-        s | "peers " `isPrefixOf` s -> cmd $ do
-                    let ih = fromString (drop 6 s)
+        ("peers ", s) -> cmd $ do
+                    let ih = fromString s
                     ps <- allPeers ih
                     return $ do
                     hPutClient h $ showReport $ map (((,) "") . show . pPrint) ps
@@ -168,6 +170,15 @@ main = do
                  (Lifted.ioError $ userError "unable to resolve bootstrap nodes")
     saved_nodes <- resume
 
+    peers'trial <- liftIO $ tryIOError $ B.readFile "bt-peers.dat"
+    saved_peers <-
+        either (const $ do liftIO $ putStrLn "Error reading bt-peers.dat"
+                           return Nothing)
+               (return . Just)
+            peers'trial
+
+    maybe (return ()) mergeSavedPeers saved_peers
+
     when (isJust saved_nodes) $ do
         b <- isBootstrapped
         tbl <- getTable
@@ -209,7 +220,9 @@ main = do
                ,("buckets", show $ R.shape tbl)
                ,("optBucketCount", show bc)
                ]
-    snapshot >>= liftIO . B.writeFile "dht-nodes.dat"
 
-    waitForSignal
+    waitForSignal -- Await unix socket to signal termination.
+
+    snapshot >>= liftIO . B.writeFile "dht-nodes.dat"
+    savePeerStore >>= liftIO . B.writeFile "bt-peers.dat"
 
diff --git a/src/Network/BitTorrent/DHT/ContactInfo.hs b/src/Network/BitTorrent/DHT/ContactInfo.hs
index 4302288c..117325fc 100644
--- a/src/Network/BitTorrent/DHT/ContactInfo.hs
+++ b/src/Network/BitTorrent/DHT/ContactInfo.hs
@@ -15,6 +15,8 @@ import Data.Serialize
 import Data.PSQueue as PSQ
 import Data.Time.Clock.POSIX
 import Data.ByteString (ByteString)
+import Data.Word
+import Network.Socket (SockAddr(..))
 
 import Data.Torrent
 import Network.BitTorrent.Address
@@ -122,6 +124,40 @@ data SwarmData ip = SwarmData
     , name :: Maybe ByteString
     }
 
+
+newtype SerializeAddress a = SerializeAddress { unserializeAddress :: a }
+
+instance Address a => Serialize (SerializeAddress a) where
+    get = SerializeAddress <$> do
+        c <- get
+        case (c::Word8) of
+            0x34 -> do ip4 <- get
+                       return $ fromJust $ fromAddr (ip4::IPv4)
+            0x36 -> do ip6 <- get
+                       return $ fromJust $ fromAddr (ip6::IPv6)
+            _ -> return $ error "cannot deserialize non-IP SerializeAddress"
+    put (SerializeAddress a)
+        | Just ip4 <- fromAddr a
+                    = put (0x34::Word8) >> put (ip4::IPv4)
+        | Just ip6 <- fromAddr a
+                    = put (0x36::Word8) >> put (ip6::IPv6)
+        | otherwise = return $ error "cannot serialize non-IP SerializeAddress"
+
+
+instance (Ord ip, Address ip) => Serialize (SwarmData ip) where
+    get = flip SwarmData <$> get
+                         <*> ( PSQ.fromList . L.map parseAddr <$> get )
+     where
+        parseAddr addr = (unserializeAddress <$> addr)
+                          :-> 0
+
+    put SwarmData{..} = do
+        put name
+        put $ L.map (\(addr :-> _) -> (SerializeAddress <$> addr))
+                -- XXX: should we serialize the timestamp?
+            $ PSQ.toList peers
+
+
 knownSwarms :: PeerStore ip -> [ (InfoHash, Int, Maybe ByteString) ]
 knownSwarms (PeerStore m) = L.map (\(ih,SwarmData q n) -> (ih,PSQ.size q,n)) $ HM.toList m
 
@@ -136,26 +172,27 @@ swarmInsert old new = SwarmData
     , name = name new <|> name old -- TODO: decodeUtf8' check
     }
 
+isSwarmOccupied SwarmData{..} = not $ PSQ.null peers
 
 -- | Empty store.
 instance Default (PeerStore a) where
   def = PeerStore HM.empty
   {-# INLINE def #-}
 
--- -- | Monoid under union operation.
--- instance Eq a => Monoid (PeerStore a) where
---   mempty  = def
---   {-# INLINE mempty #-}
--- 
---   mappend (PeerStore a) (PeerStore b) =
---     PeerStore (HM.unionWith L.union a b)
---   {-# INLINE mappend #-}
+-- | Monoid under union operation.
+instance Ord a => Monoid (PeerStore a) where
+  mempty  = def
+  {-# INLINE mempty #-}
+
+  mappend (PeerStore a) (PeerStore b) =
+    PeerStore (HM.unionWith swarmInsert a b)
+  {-# INLINE mappend #-}
 
 -- | Can be used to store peers between invocations of the client
 -- software.
-instance Serialize (PeerStore a) where
-  get = undefined
-  put = undefined
+instance (Ord a, Address a) => Serialize (PeerStore a) where
+  get = PeerStore . HM.fromList <$> get
+  put (PeerStore m) = put (L.filter (isSwarmOccupied . snd) $ HM.toList m)
 
 -- | Returns all peers associated with a given info hash.
 lookup :: Ord a => InfoHash -> PeerStore a -> [PeerAddr a]
diff --git a/src/Network/BitTorrent/DHT/Session.hs b/src/Network/BitTorrent/DHT/Session.hs
index bc9fda91..c08021c7 100644
--- a/src/Network/BitTorrent/DHT/Session.hs
+++ b/src/Network/BitTorrent/DHT/Session.hs
@@ -54,14 +54,16 @@ module Network.BitTorrent.DHT.Session
          -- ** Routing table
        , getTable
        , getClosest
-       , getSwarms
-       , allPeers
 
          -- ** Peer storage
        , insertPeer
        , getPeerList
        , insertTopic
        , deleteTopic
+       , getSwarms
+       , savePeerStore
+       , mergeSavedPeers
+       , allPeers
 
          -- ** Messaging
        , queryParallel
@@ -84,6 +86,7 @@ import Data.Fixed
 import Data.Hashable
 import Data.List as L
 import Data.Maybe
+import Data.Monoid
 import Data.Set as S
 import Data.Time
 import Network (PortNumber)
@@ -91,6 +94,7 @@ import System.Random (randomIO)
 import Data.Time.Clock.POSIX
 import Data.Text as Text
 import Text.PrettyPrint.HughesPJClass hiding ((<>),($$))
+import Data.Serialize as S
 
 
 import Data.Torrent as Torrent
@@ -409,6 +413,20 @@ getSwarms = do
     store <- asks contactInfo >>= liftIO . atomically . readTVar
     return $ P.knownSwarms store
 
+savePeerStore :: (Ord ip, Address ip) => DHT ip ByteString
+savePeerStore = do
+    var <- asks contactInfo
+    peers <- liftIO $ atomically $ readTVar var
+    return $ S.encode peers
+
+mergeSavedPeers :: (Ord ip, Address ip) => ByteString -> DHT ip ()
+mergeSavedPeers bs = do
+    var <- asks contactInfo
+    case S.decode bs of
+        Right newbies -> liftIO $ atomically $ modifyTVar' var (<> newbies)
+        Left _        -> return ()
+
+
 allPeers :: Ord ip => InfoHash -> DHT ip [ PeerAddr ip ]
 allPeers ih = do
     store <- asks contactInfo >>= liftIO . atomically . readTVar
diff --git a/src/Network/BitTorrent/Tracker/RPC/HTTP.hs b/src/Network/BitTorrent/Tracker/RPC/HTTP.hs
index de3fc5f5..bc52bddd 100644
--- a/src/Network/BitTorrent/Tracker/RPC/HTTP.hs
+++ b/src/Network/BitTorrent/Tracker/RPC/HTTP.hs
@@ -129,7 +129,7 @@ fillRequest Options {..} q r = r
 
 httpTracker :: BEncode a => Manager -> URI -> SimpleQuery -> IO a
 httpTracker Manager {..} uri q = packHttpException $ do
-  request  <- fillRequest options q <$> setUri def uri
+  request  <- fillRequest options q <$> setUri def {- http-client instance for Request -} uri
   response <- runResourceT $ httpLbs request httpMgr
   case BE.decode $ BL.toStrict $ responseBody response of
     Left  msg  -> throwIO (ParserFailure msg)