PeerStore is now persistent.

1 files changed, 48 insertions, 11 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]