2 files changed, 68 insertions, 13 deletions

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