From 11987749fc6e6d3e53ea737d46d5ab13a16faeb8 Mon Sep 17 00:00:00 2001 From: James Crayne Date: Sat, 28 Sep 2019 13:43:29 -0400 Subject: Factor out some new libraries word64-map: Data.Word64Map network-addr: Network.Address tox-crypto: Crypto.Tox lifted-concurrent: Control.Concurrent.Lifted.Instrument Control.Concurrent.Async.Lifted.Instrument psq-wrap: Data.Wrapper.PSQInt Data.Wrapper.PSQ minmax-psq: Data.MinMaxPSQ tasks: Control.Concurrent.Tasks kad: Network.Kademlia Network.Kademlia.Bootstrap Network.Kademlia.Routing Network.Kademlia.CommonAPI Network.Kademlia.Persistence Network.Kademlia.Search --- dht/src/Network/BitTorrent/DHT/ContactInfo.hs | 254 ++++++++++++++++++++++++++ 1 file changed, 254 insertions(+) create mode 100644 dht/src/Network/BitTorrent/DHT/ContactInfo.hs (limited to 'dht/src/Network/BitTorrent/DHT/ContactInfo.hs') 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 -- cgit v1.2.3