path: root/src/Network/BitTorrent/DHT/ContactInfo.hs

{-# LANGUAGE BangPatterns #-}
module Network.BitTorrent.DHT.ContactInfo
       ( PeerStore
       , PeerAddr(..)
       , 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.Wrapper.PSQ as PSQ
import Data.Time.Clock.POSIX
import Data.ByteString (ByteString)
import Data.Word
import Network.Socket (SockAddr(..))

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"


{- XXX: What happened to the ports?
instance Serialize SwarmData 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 -> [ (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 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) where
  def = PeerStore HM.empty
  {-# INLINE def #-}

-- | 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 = 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 $! (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 }