1 files changed, 254 insertions, 0 deletions

diff --git a/dht/src/Network/BitTorrent/DHT/ContactInfo.hs b/dht/src/Network/BitTorrent/DHT/ContactInfo.hs
new file mode 100644
index 00000000..ec7e6658
--- /dev/null
+++ b/dht/src/Network/BitTorrent/DHT/ContactInfo.hs
@@ -0,0 +1,254 @@
+{-# LANGUAGE BangPatterns #-}
+module Network.BitTorrent.DHT.ContactInfo
+       ( PeerStore
+       , PeerAddr(..)
+       , Network.BitTorrent.DHT.ContactInfo.lookup
+       , Network.BitTorrent.DHT.ContactInfo.freshPeers
+       , Network.BitTorrent.DHT.ContactInfo.insertPeer
+       , deleteOlderThan
+       , knownSwarms
+       ) where
+
+import Control.Applicative
+import Data.Default
+import Data.List as L
+import Data.Maybe
+import Data.HashMap.Strict as HM
+import Data.Serialize
+import Data.Semigroup
+import Data.Wrapper.PSQ as PSQ
+import Data.Time.Clock.POSIX
+import Data.ByteString (ByteString)
+import Data.Word
+
+import Data.Torrent
+import Network.Address
+
+-- {-
+-- import Data.HashMap.Strict as HM
+--
+-- import Data.Torrent.InfoHash
+-- import Network.Address
+--
+-- -- increase prefix when table is too large
+-- -- decrease prefix when table is too small
+-- -- filter outdated peers
+--
+-- {-----------------------------------------------------------------------
+-- --  PeerSet
+-- -----------------------------------------------------------------------}
+--
+-- type PeerSet a = [(PeerAddr, NodeInfo a, Timestamp)]
+--
+-- -- compare PSQueue vs Ordered list
+--
+-- takeNewest :: PeerSet a -> [PeerAddr]
+-- takeNewest = undefined
+--
+-- dropOld :: Timestamp -> PeerSet a -> PeerSet a
+-- dropOld = undefined
+--
+-- insert :: PeerAddr -> Timestamp -> PeerSet a -> PeerSet a
+-- insert = undefined
+--
+-- type Mask = Int
+-- type Size = Int
+-- type Timestamp = Int
+--
+-- {-----------------------------------------------------------------------
+-- --  InfoHashMap
+-- -----------------------------------------------------------------------}
+--
+-- -- compare handwritten prefix tree versus IntMap
+--
+-- data Tree a
+--   = Nil
+--   | Tip !InfoHash !(PeerSet a)
+--   | Bin !InfoHash !Mask !Size !Timestamp (Tree a) (Tree a)
+--
+-- insertTree :: InfoHash -> a -> Tree a -> Tree a
+-- insertTree = undefined
+--
+-- type Prio = Int
+--
+-- --shrink :: ContactInfo ip -> Int
+-- shrink  Nil      = Nil
+-- shrink (Tip _ _) = undefined
+-- shrink (Bin _ _) = undefined
+--
+-- {-----------------------------------------------------------------------
+-- -- InfoHashMap
+-- -----------------------------------------------------------------------}
+--
+-- -- compare new design versus HashMap
+--
+-- data IntMap k p a
+-- type ContactInfo = Map InfoHash Timestamp (Set (PeerAddr IP) Timestamp)
+--
+-- data ContactInfo ip = PeerStore
+--   { maxSize    :: Int
+--   , prefixSize :: Int
+--   , thisNodeId :: NodeId
+--
+--   , count      :: Int  -- ^ Cached size of the 'peerSet'
+--   , peerSet    :: HashMap InfoHash [PeerAddr ip]
+--   }
+--
+-- size :: ContactInfo ip -> Int
+-- size = undefined
+--
+-- prefixSize :: ContactInfo ip -> Int
+-- prefixSize = undefined
+--
+-- lookup :: InfoHash -> ContactInfo ip -> [PeerAddr ip]
+-- lookup = undefined
+--
+-- insert :: InfoHash -> PeerAddr ip -> ContactInfo ip -> ContactInfo ip
+-- insert = undefined
+--
+-- -- | Limit in size.
+-- prune :: NodeId -> Int -> ContactInfo ip -> ContactInfo ip
+-- prune pref targetSize  Nil      = Nil
+-- prune pref targetSize (Tip _ _) = undefined
+--
+-- -- | Remove expired entries.
+-- splitGT :: Timestamp -> ContactInfo ip -> ContactInfo ip
+-- splitGT = undefined
+-- -}
+
+-- | Storage used to keep track a set of known peers in client,
+-- tracker or DHT sessions.
+newtype PeerStore = PeerStore (HashMap InfoHash SwarmData)
+
+type Timestamp = POSIXTime
+
+data SwarmData = SwarmData
+    { peers :: !(PSQ PeerAddr Timestamp)
+    , name :: !(Maybe ByteString)
+    }
+
+-- | This wrapper will serialize an ip address with a '4' or '6' prefix byte
+-- to indicate whether it is IPv4 or IPv6.
+--
+-- Note: it does not serialize port numbers.
+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 Serialize SwarmData where
+    get = flip SwarmData <$> get
+                         <*> ( PSQ.fromList . L.map parseAddr <$> get )
+     where
+        parseAddr (pid,addr,port) = PeerAddr { peerId = pid
+                                             , peerHost = unserializeAddress addr
+                                             , peerPort = port
+                                             }
+                                      :-> 0
+
+    put SwarmData{..} = do
+        put name
+        put $ L.map (\(addr :-> _) -> (peerId addr, SerializeAddress addr, peerPort addr))
+                -- XXX: should we serialize the timestamp?
+            $ PSQ.toList peers
+
+knownSwarms :: PeerStore -> [ (InfoHash, Int, Maybe ByteString) ]
+knownSwarms (PeerStore m) = L.map (\(ih,SwarmData q n) -> (ih,PSQ.size q,n)) $ HM.toList m
+
+swarmSingleton :: PeerAddr -> SwarmData
+swarmSingleton a = SwarmData
+    { peers = PSQ.singleton a 0
+    , name = Nothing }
+
+swarmInsert :: SwarmData -> SwarmData -> SwarmData
+swarmInsert new old = SwarmData
+    { peers = L.foldl' (\q (a :-> t) -> PSQ.insertWith newerTimeStamp a t q) (peers old) (PSQ.toList $ peers new)
+    , name = name new <|> name old -- TODO: decodeUtf8' check
+    }
+ where
+    newerTimeStamp newtime oldtime = if newtime > oldtime then newtime else oldtime
+
+isSwarmOccupied :: SwarmData -> Bool
+isSwarmOccupied SwarmData{..} = not $ PSQ.null peers
+
+-- | Empty store.
+instance Default (PeerStore) where
+  def = PeerStore HM.empty
+  {-# INLINE def #-}
+
+instance Semigroup PeerStore where
+  PeerStore a <> PeerStore b =
+    PeerStore (HM.unionWith swarmInsert a b)
+  {-# INLINE (<>) #-}
+
+-- | Monoid under union operation.
+instance Monoid PeerStore 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 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 :: InfoHash -> PeerStore -> [PeerAddr]
+lookup ih (PeerStore m) = maybe [] (PSQ.keys . peers) $ HM.lookup ih m
+
+batchSize :: Int
+batchSize = 64
+
+-- | Used in 'get_peers' DHT queries.
+freshPeers :: InfoHash -> Timestamp -> PeerStore -> ([PeerAddr], PeerStore)
+freshPeers ih tm (PeerStore m) = fromMaybe ([],PeerStore m) $ do
+    swarm <- HM.lookup ih m
+    let ps0 = take batchSize $ unfoldr (incomp minView) (peers swarm)
+        peers' = case reverse ps0 of
+                    (_,psq):_ -> psq
+                    _         -> peers swarm
+        ps = L.map (key . fst) ps0
+        m' = HM.insert ih swarm { peers = L.foldl' (\q p -> PSQ.insert p tm q) peers' ps } m
+    return $! m' `seq` (ps,PeerStore m')
+
+incomp :: (x -> Maybe (r,x)) -> x -> Maybe ((r,x),x)
+incomp !f !x = do
+   (result,x') <- f x
+   pure $! ( (result,x'), x' )
+
+-- | Used in 'announce_peer' DHT queries.
+insertPeer :: InfoHash -> Maybe ByteString -> PeerAddr -> PeerStore -> PeerStore
+insertPeer !ih !name !a !(PeerStore m) = seq a' $ PeerStore (HM.insertWith swarmInsert ih a' m)
+ where
+    a' = SwarmData { peers = PSQ.singleton a 0
+                   , name = name }
+
+deleteOlderThan :: POSIXTime -> PeerStore -> PeerStore
+deleteOlderThan cutoff (PeerStore m) = PeerStore $ HM.mapMaybe gc m
+    where
+        gc :: SwarmData -> Maybe SwarmData
+        gc swarms = fmap (\ps -> swarms { peers = ps }) $ gcPSQ (peers swarms)
+
+        gcPSQ :: PSQKey a => PSQ a Timestamp -> Maybe (PSQ a Timestamp)
+        gcPSQ ps = case minView ps of
+            Nothing                            -> Nothing
+            Just (_ :-> tm, ps') | tm < cutoff -> gcPSQ ps'
+            Just _                             -> Just ps