3 files changed, 74 insertions, 96 deletions

diff --git a/src/Network/BitTorrent.hs b/src/Network/BitTorrent.hs
index c68cceac..7c1e02e2 100644
--- a/src/Network/BitTorrent.hs
+++ b/src/Network/BitTorrent.hs
@@ -89,7 +89,6 @@ import Network.BitTorrent.Extension
 import Network.BitTorrent.Exchange
 import Network.BitTorrent.Exchange.Protocol
 import Network.BitTorrent.Tracker
-import Network.BitTorrent.Discovery
 
 import System.Torrent.Storage
 
@@ -100,8 +99,6 @@ import System.Torrent.Storage
 withDefaultClient :: PortNumber -> PortNumber -> (ClientSession -> IO ()) -> IO ()
 withDefaultClient dhtPort listPort action = do
   withClientSession defaultThreadCount defaultExtensions $ \client -> do
-    startListener client listPort
-    startDHT      client dhtPort
     action client
 
 {-----------------------------------------------------------------------
diff --git a/src/Network/BitTorrent/Discovery.hs b/src/Network/BitTorrent/Discovery.hs
deleted file mode 100644
index 8403461c..00000000
--- a/src/Network/BitTorrent/Discovery.hs
+++ /dev/null
@@ -1,59 +0,0 @@
-{-# LANGUAGE RecordWildCards   #-}
-{-# LANGUAGE OverloadedStrings #-}
-module Network.BitTorrent.Discovery
-       (discover, startListener, startDHT
-       ) where
-
-import Control.Monad
-import Control.Concurrent
-import Control.Exception
-import Network.Socket
-
-import Data.Torrent
-import Network.BitTorrent.Peer
-import Network.BitTorrent.Sessions
-import Network.BitTorrent.Exchange
-import Network.BitTorrent.Tracker
-import Network.BitTorrent.DHT
-
-
--- discover should hide tracker and DHT communication under the hood
--- thus we can obtain an unified interface
-
-discover :: SwarmSession -> P2P () -> IO ()
-discover swarm @ SwarmSession {..} action = {-# SCC discover #-} do
-  port <- listenerPort clientSession
-
-  let conn = TConnection {
-        tconnAnnounce = tAnnounce torrentMeta
-      , tconnInfoHash = tInfoHash torrentMeta
-      , tconnPeerId   = clientPeerId clientSession
-      , tconnPort     = port
-      }
-
-  progress <- getCurrentProgress clientSession
-
-  withTracker progress conn $ \tses -> do
-    forever $ do
-      addr <- getPeerAddr tses
-      forkThrottle swarm $ do
-        initiatePeerSession swarm addr $ \conn ->
-          runP2P conn action
-
-
-startListener :: ClientSession -> PortNumber -> IO ()
-startListener cs @ ClientSession {..} port =
-  startService peerListener port $ listener cs $ \conn @ (sock, PeerSession{..}) -> do
-      print "accepted"
-      let storage = error "storage"
-      runP2P conn (exchange storage)
-
-startDHT :: ClientSession -> PortNumber -> IO ()
-startDHT ClientSession {..} nodePort = withRunning peerListener failure start
-  where
-    start ClientService {..} = do
-      ses  <- newNodeSession servPort
-      startService nodeListener nodePort (dhtServer ses)
-
-    failure = throwIO $ userError msg
-    msg = "unable to start DHT server: peer listener is not running"
diff --git a/src/Network/BitTorrent/Sessions.hs b/src/Network/BitTorrent/Sessions.hs
index 31b30e43..da57ec3b 100644
--- a/src/Network/BitTorrent/Sessions.hs
+++ b/src/Network/BitTorrent/Sessions.hs
@@ -19,47 +19,44 @@ module Network.BitTorrent.Sessions
                        , clientPeerId, allowedExtensions
                        , nodeListener, peerListener
                        )
-        , withClientSession
-        , listenerPort, dhtPort
+       , withClientSession
+       , listenerPort, dhtPort
 
-        , ThreadCount
-        , defaultThreadCount
+       , ThreadCount
+       , defaultThreadCount
 
-        , TorrentLoc(..)
-        , registerTorrent
-        , unregisterTorrent
+       , TorrentLoc(..)
+       , registerTorrent
+       , unregisterTorrent
 
-        , getCurrentProgress
-        , getSwarmCount
-        , getPeerCount
+       , getCurrentProgress
+       , getSwarmCount
+       , getPeerCount
 
-          -- * Swarm
-        , SwarmSession( SwarmSession, torrentMeta, clientSession )
+         -- * Swarm
+       , SwarmSession( SwarmSession, torrentMeta, clientSession )
 
-        , SessionCount
-        , getSessionCount
+       , SessionCount
+       , getSessionCount
 
-        , newLeecher
-        , newSeeder
-        , getClientBitfield
+       , newLeecher
+       , newSeeder
+       , getClientBitfield
 
-          -- TODO hide this
-        , waitVacancy
-        , forkThrottle
+         -- * Peer
+       , PeerSession( PeerSession, connectedPeerAddr
+                    , swarmSession, enabledExtensions
+                    , sessionState
+                    )
+       , SessionState
+       , initiatePeerSession
+       , acceptPeerSession
+       , listener
 
-          -- * Peer
-        , PeerSession( PeerSession, connectedPeerAddr
-                     , swarmSession, enabledExtensions
-                     , sessionState
-                     )
-        , SessionState
-        , initiatePeerSession
-        , acceptPeerSession
-        , listener
-
-          -- * Timeouts
-        , updateIncoming, updateOutcoming
-        ) where
+         -- * Timeouts
+       , updateIncoming, updateOutcoming
+       , discover
+       ) where
 
 import Prelude hiding (mapM_)
 
@@ -90,10 +87,13 @@ import Data.Bitfield as BF
 import Data.Torrent
 import Network.BitTorrent.Extension
 import Network.BitTorrent.Peer
+import Network.BitTorrent.Sessions.Types
 import Network.BitTorrent.Exchange.Protocol as BT
 import Network.BitTorrent.Tracker.Protocol as BT
+import Network.BitTorrent.Tracker as BT
+import Network.BitTorrent.Exchange as BT
+import Network.BitTorrent.DHT as BT
 import System.Torrent.Storage
-import Network.BitTorrent.Sessions.Types
 
 {-----------------------------------------------------------------------
     Client Services
@@ -134,6 +134,23 @@ torrentPresence ClientSession {..} ih = do
     Client sessions
 -----------------------------------------------------------------------}
 
+startListener :: ClientSession -> PortNumber -> IO ()
+startListener cs @ ClientSession {..} port =
+  startService peerListener port $ listener cs $ \conn @ (sock, PeerSession{..}) -> do
+      print "accepted"
+      let storage = error "storage"
+      runP2P conn (exchange storage)
+
+startDHT :: ClientSession -> PortNumber -> IO ()
+startDHT ClientSession {..} nodePort = withRunning peerListener failure start
+  where
+    start ClientService {..} = do
+      ses  <- newNodeSession servPort
+      startService nodeListener nodePort (dhtServer ses)
+
+    failure = throwIO $ userError msg
+    msg = "unable to start DHT server: peer listener is not running"
+
 -- | Create a new client session. The data passed to this function are
 -- usually loaded from configuration file.
 openClientSession :: SessionCount     -- ^ Maximum count of active P2P Sessions.
@@ -194,6 +211,29 @@ defSeederConns = defaultUnchokeSlots
 defLeacherConns :: SessionCount
 defLeacherConns = defaultNumWant
 
+-- discovery should hide tracker and DHT communication under the hood
+-- thus we can obtain an unified interface
+
+discover :: SwarmSession -> P2P () -> IO ()
+discover swarm @ SwarmSession {..} action = {-# SCC discover #-} do
+  port <- listenerPort clientSession
+
+  let conn = TConnection {
+        tconnAnnounce = tAnnounce torrentMeta
+      , tconnInfoHash = tInfoHash torrentMeta
+      , tconnPeerId   = clientPeerId clientSession
+      , tconnPort     = port
+      }
+
+  progress <- getCurrentProgress clientSession
+
+  withTracker progress conn $ \tses -> do
+    forever $ do
+      addr <- getPeerAddr tses
+      forkThrottle swarm $ do
+        initiatePeerSession swarm addr $ \conn ->
+          runP2P conn action
+
 newSwarmSession :: Int -> Bitfield -> ClientSession -> Torrent
                 -> IO SwarmSession
 newSwarmSession n bf cs @ ClientSession {..} t @ Torrent {..}