module Network.BitTorrent.DHT.ContactInfo ( PeerStore , Network.BitTorrent.DHT.ContactInfo.lookup , Network.BitTorrent.DHT.ContactInfo.freshPeers , Network.BitTorrent.DHT.ContactInfo.insertPeer , 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.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 -- {- -- import Data.HashMap.Strict as HM -- -- import Data.Torrent.InfoHash -- import Network.BitTorrent.Address -- -- -- increase prefix when table is too large -- -- decrease prefix when table is too small -- -- filter outdated peers -- -- {----------------------------------------------------------------------- -- -- PeerSet -- -----------------------------------------------------------------------} -- -- type PeerSet a = [(PeerAddr a, NodeInfo a, Timestamp)] -- -- -- compare PSQueue vs Ordered list -- -- takeNewest :: PeerSet a -> [PeerAddr a] -- takeNewest = undefined -- -- dropOld :: Timestamp -> PeerSet a -> PeerSet a -- dropOld = undefined -- -- insert :: PeerAddr a -> 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 ip = PeerStore (HashMap InfoHash (SwarmData ip)) type Timestamp = POSIXTime data SwarmData ip = SwarmData { peers :: PSQ (PeerAddr ip) Timestamp , 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 swarmSingleton :: Ord ip => PeerAddr ip -> SwarmData ip swarmSingleton a = SwarmData { peers = PSQ.singleton a 0 , name = Nothing } swarmInsert :: Ord ip => SwarmData ip -> SwarmData ip -> SwarmData ip swarmInsert old new = SwarmData { peers = L.foldl' (\q (a :-> t) -> PSQ.insertWith (\p _ -> p) a t q) (peers old) (PSQ.toList $ peers new) , 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 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 (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] lookup ih (PeerStore m) = maybe [] (PSQ.keys . peers) $ HM.lookup ih m batchSize = 64 -- | Used in 'get_peers' DHT queries. freshPeers :: Ord a => InfoHash -> Timestamp -> PeerStore a -> ([PeerAddr a], PeerStore a) freshPeers ih tm (PeerStore m) = (ps, PeerStore m') where swarm = fromMaybe (SwarmData PSQ.empty Nothing) $ HM.lookup ih m 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 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 :: Ord a => InfoHash -> Maybe ByteString -> PeerAddr a -> PeerStore a -> PeerStore a insertPeer ih name a (PeerStore m) = PeerStore (HM.insertWith swarmInsert ih a' m) where a' = SwarmData { peers = PSQ.singleton a 0 , name = name }