PSQ instead of list for peer set. Also: dhtd "swarms" command.

5 files changed, 77 insertions, 27 deletions

diff --git a/src/Network/BitTorrent/DHT/ContactInfo.hs b/src/Network/BitTorrent/DHT/ContactInfo.hs
index d7c92e35..979dbb62 100644
--- a/src/Network/BitTorrent/DHT/ContactInfo.hs
+++ b/src/Network/BitTorrent/DHT/ContactInfo.hs
@@ -1,14 +1,19 @@
 module Network.BitTorrent.DHT.ContactInfo
        ( PeerStore
        , Network.BitTorrent.DHT.ContactInfo.lookup
-       , Network.BitTorrent.DHT.ContactInfo.insert
+       , 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.Torrent
 import Network.BitTorrent.Address
@@ -107,21 +112,43 @@ import Network.BitTorrent.Address
 
 -- | Storage used to keep track a set of known peers in client,
 -- tracker or DHT sessions.
-newtype PeerStore ip = PeerStore (HashMap InfoHash [PeerAddr ip])
+newtype PeerStore ip = PeerStore (HashMap InfoHash (SwarmData ip))
+
+type Timestamp = POSIXTime
+
+data SwarmData ip = SwarmData
+    { peers :: PSQ (PeerAddr ip) Timestamp
+    , name :: Maybe ByteString
+    }
+
+knownSwarms :: PeerStore ip -> [ (InfoHash, Int, Maybe ByteString) ]
+knownSwarms (PeerStore m) = L.map (\(ih,SwarmData q n) -> (ih,PSQ.size q,n)) $ HM.toList m
+
+swarmSinglton :: Ord ip => PeerAddr ip -> SwarmData ip
+swarmSinglton 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
+    }
+
 
 -- | 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 Eq a => Monoid (PeerStore a) where
+--   mempty  = def
+--   {-# INLINE mempty #-}
+-- 
+--   mappend (PeerStore a) (PeerStore b) =
+--     PeerStore (HM.unionWith L.union a b)
+--   {-# INLINE mappend #-}
 
 -- | Can be used to store peers between invocations of the client
 -- software.
@@ -130,9 +157,12 @@ instance Serialize (PeerStore a) where
   put = undefined
 
 -- | Used in 'get_peers' DHT queries.
-lookup :: InfoHash -> PeerStore a -> [PeerAddr a]
-lookup ih (PeerStore m) = fromMaybe [] $ HM.lookup ih m
+lookup :: Ord a => InfoHash -> PeerStore a -> [PeerAddr a]
+lookup ih (PeerStore m) = maybe [] (PSQ.keys . peers) $ HM.lookup ih m
 
 -- | Used in 'announce_peer' DHT queries.
-insert :: Eq a => InfoHash -> PeerAddr a -> PeerStore a -> PeerStore a
-insert ih a (PeerStore m) = PeerStore (HM.insertWith L.union ih [a] m)
+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 }
diff --git a/src/Network/BitTorrent/DHT/Message.hs b/src/Network/BitTorrent/DHT/Message.hs
index 9d66741f..d31cce82 100644
--- a/src/Network/BitTorrent/DHT/Message.hs
+++ b/src/Network/BitTorrent/DHT/Message.hs
@@ -86,6 +86,7 @@ module Network.BitTorrent.DHT.Message
 import Control.Applicative
 import Data.BEncode as BE
 import Data.BEncode.BDict
+import Data.ByteString (ByteString)
 import Data.List as L
 import Data.Monoid
 import Data.Serialize as S
@@ -251,6 +252,9 @@ peers_key = "values"
 token_key :: BKey
 token_key = "token"
 
+name_key :: BKey
+name_key = "name"
+
 instance (Typeable ip, Serialize ip) => BEncode (GotPeers ip) where
   toBEncode GotPeers {..} = toDict $
     case peers of
@@ -295,11 +299,17 @@ data Announce = Announce
     -- | infohash of the torrent;
   , topic    :: InfoHash
 
+    -- | some clients announce the friendly name of the torrent here.
+  , announcedName :: Maybe ByteString
+
     -- | the port /this/ peer is listening;
   , port     :: PortNumber
 
+    -- TODO: optional boolean "seed" key
+
     -- | received in response to a previous get_peers query.
   , sessionToken :: Token
+
   } deriving (Show, Eq, Typeable)
 
 port_key :: BKey
@@ -312,6 +322,7 @@ instance BEncode Announce where
   toBEncode Announce {..} = toDict $
        implied_port_key .=? flagField impliedPort
     .: info_hash_key    .=! topic
+    .: name_key         .=? announcedName
     .: port_key         .=! port
     .: token_key        .=! sessionToken
     .: endDict
@@ -321,6 +332,7 @@ instance BEncode Announce where
   fromBEncode = fromDict $ do
     Announce <$> (boolField <$> optional (field (req implied_port_key)))
              <*>! info_hash_key
+             <*>? name_key
              <*>! port_key
              <*>! token_key
     where
diff --git a/src/Network/BitTorrent/DHT/Query.hs b/src/Network/BitTorrent/DHT/Query.hs
index 44083d81..0bec867d 100644
--- a/src/Network/BitTorrent/DHT/Query.hs
+++ b/src/Network/BitTorrent/DHT/Query.hs
@@ -105,7 +105,7 @@ findNodeH = nodeHandler $ \ _ (FindNode nid) -> do
   NodeFound <$> getClosest nid
 
 -- | Default 'GetPeers' handler.
-getPeersH :: Address ip => NodeHandler ip
+getPeersH :: Ord ip => Address ip => NodeHandler ip
 getPeersH = nodeHandler $ \ naddr (GetPeers ih) -> do
   ps <- getPeerList ih
   tok <- grantToken naddr
@@ -113,19 +113,19 @@ getPeersH = nodeHandler $ \ naddr (GetPeers ih) -> do
   return $ GotPeers ps tok
 
 -- | Default 'Announce' handler.
-announceH :: Address ip => NodeHandler ip
+announceH :: Ord ip => Address ip => NodeHandler ip
 announceH = nodeHandler $ \ naddr @ NodeAddr {..} (Announce {..}) -> do
   valid <- checkToken naddr sessionToken
   unless valid $ do
     throwIO $ InvalidParameter "token"
 
   let annPort  = if impliedPort then nodePort else port
-  let peerAddr = PeerAddr Nothing nodeHost annPort
-  insertPeer topic peerAddr
+      peerAddr = PeerAddr Nothing nodeHost annPort
+  insertPeer topic announcedName peerAddr
   return Announced
 
 -- | Includes all default query handlers.
-defaultHandlers :: Address ip => [NodeHandler ip]
+defaultHandlers :: Ord ip => Address ip => [NodeHandler ip]
 defaultHandlers = [pingH, findNodeH, getPeersH, announceH]
 
 {-----------------------------------------------------------------------
@@ -168,7 +168,7 @@ announceQ ih p NodeInfo {..} = do
       | False     -> undefined -- TODO check if we can announce
       | otherwise -> return (Left ns)
     Right _ -> do -- TODO *probably* add to peer cache
-      Announced <- Announce False ih p grantedToken <@> nodeAddr
+      Announced <- Announce False ih Nothing p grantedToken <@> nodeAddr
       return (Right [nodeAddr])
 
 {-----------------------------------------------------------------------
diff --git a/src/Network/BitTorrent/DHT/Routing.hs b/src/Network/BitTorrent/DHT/Routing.hs
index 38207be5..8a6849a1 100644
--- a/src/Network/BitTorrent/DHT/Routing.hs
+++ b/src/Network/BitTorrent/DHT/Routing.hs
@@ -460,7 +460,7 @@ compatibleNodeId tbl = genBucketSample prefix br
  where
     br = bucketRange (L.length (shape tbl) - 1) True
     bs = BS.pack $ take nodeIdSize $ tablePrefix tbl ++ repeat 0
-    prefix = NodeId bs
+    prefix = asNodeId bs
 
 tablePrefix :: Table ip -> [Word8]
 tablePrefix = map (packByte . take 8 . (++repeat False))
diff --git a/src/Network/BitTorrent/DHT/Session.hs b/src/Network/BitTorrent/DHT/Session.hs
index b85e97fa..339b18eb 100644
--- a/src/Network/BitTorrent/DHT/Session.hs
+++ b/src/Network/BitTorrent/DHT/Session.hs
@@ -53,6 +53,7 @@ module Network.BitTorrent.DHT.Session
          -- ** Routing table
        , getTable
        , getClosest
+       , getSwarms
 
          -- ** Peer storage
        , insertPeer
@@ -74,6 +75,7 @@ import Control.Monad.Logger
 import Control.Monad.Reader
 import Control.Monad.Trans.Control
 import Control.Monad.Trans.Resource
+import Data.ByteString
 import Data.Conduit.Lazy
 import Data.Default
 import Data.Fixed
@@ -89,7 +91,8 @@ import Data.Torrent as Torrent
 import Network.KRPC as KRPC hiding (Options, def)
 import qualified Network.KRPC as KRPC (def)
 import Network.BitTorrent.Address
-import Network.BitTorrent.DHT.ContactInfo as P
+import Network.BitTorrent.DHT.ContactInfo (PeerStore)
+import qualified Network.BitTorrent.DHT.ContactInfo as P
 import Network.BitTorrent.DHT.Message
 import Network.BitTorrent.DHT.Routing as R
 import Network.BitTorrent.DHT.Token   as T
@@ -395,6 +398,11 @@ getTable = do
   let nil = nullTable myId (optBucketCount opts)
   liftIO (maybe nil R.myBuckets <$> atomically (readTVar var))
 
+getSwarms :: Ord ip => DHT ip [ (InfoHash, Int, Maybe ByteString) ]
+getSwarms = do
+    store <- asks contactInfo >>= liftIO . atomically . readTVar
+    return $ P.knownSwarms store
+
 -- | Find a set of closest nodes from routing table of this node. (in
 -- no particular order)
 --
@@ -416,14 +424,14 @@ refreshContacts =
 
 
 -- | Insert peer to peer store. Used to handle announce requests.
-insertPeer :: Eq ip => InfoHash -> PeerAddr ip -> DHT ip ()
-insertPeer ih addr = do
+insertPeer :: Ord ip => InfoHash -> Maybe ByteString -> PeerAddr ip -> DHT ip ()
+insertPeer ih name addr = do
   refreshContacts
   var <- asks contactInfo
-  liftIO $ atomically $ modifyTVar' var (P.insert ih addr)
+  liftIO $ atomically $ modifyTVar' var (P.insertPeer ih name addr)
 
 -- | Get peer set for specific swarm.
-lookupPeers :: InfoHash -> DHT ip [PeerAddr ip]
+lookupPeers :: Ord ip => InfoHash -> DHT ip [PeerAddr ip]
 lookupPeers ih = do
   refreshContacts
   var <- asks contactInfo
@@ -433,7 +441,7 @@ lookupPeers ih = do
 --
 --   This operation use 'getClosest' as failback so it may block.
 --
-getPeerList :: Eq ip => InfoHash -> DHT ip (PeerList ip)
+getPeerList :: Ord ip => InfoHash -> DHT ip (PeerList ip)
 getPeerList ih = do
   ps <- lookupPeers ih
   if L.null ps