1 files changed, 14 insertions, 6 deletions

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