7 files changed, 42 insertions, 26 deletions

diff --git a/bittorrent.cabal b/bittorrent.cabal
index c5b04e42..e5fbc058 100644
--- a/bittorrent.cabal
+++ b/bittorrent.cabal
@@ -52,6 +52,7 @@ library
   if flag(testing)
      exposed-modules:  Network.BitTorrent.Exchange.Protocol
                      , Network.BitTorrent.Tracker.Protocol
+                     , Network.BitTorrent.DHT.Protocol
                      , System.IO.MMap.Fixed
   if !flag(testing)
      other-modules:    Network.BitTorrent.Exchange.Protocol
diff --git a/examples/Main.hs b/examples/Main.hs
index 1d3b711b..b8e3c11f 100644
--- a/examples/Main.hs
+++ b/examples/Main.hs
@@ -8,3 +8,6 @@ main = do
   [path]  <- getArgs
   torrent <- fromFile path
   print (contentLayout "./" (tInfo torrent))
+
+  withDefaultClient 3000 3001 $ \ client ->
+    addTorrent client $ TorrentLoc path "/tmp"
diff --git a/src/Data/Torrent.hs b/src/Data/Torrent.hs
index ae40bef4..3f555a5b 100644
--- a/src/Data/Torrent.hs
+++ b/src/Data/Torrent.hs
@@ -27,7 +27,7 @@
 module Data.Torrent
        ( -- * Torrent
          Torrent(..), ContentInfo(..), FileInfo(..)
-       , torrent, simpleTorrent
+       , mktorrent, simpleTorrent
        , torrentExt, isTorrentPath
        , fromFile
 
@@ -279,16 +279,16 @@ instance Hashable Torrent where
 -}
 
 -- | Smart constructor for 'Torrent' which compute info hash.
-torrent :: URI            -> ContentInfo
+mktorrent :: URI            -> ContentInfo
         -> Maybe [[URI]]  -> Maybe Text       -> Maybe ByteString
         -> Maybe Time     -> Maybe ByteString -> Maybe URI
         -> Maybe URI      -> Maybe ByteString
         -> Torrent
-torrent announce info = Torrent (hashlazy (BE.encoded info)) announce info
+mktorrent announce info = Torrent (hashlazy (BE.encoded info)) announce info
 
 -- | A simple torrent contains only required fields.
 simpleTorrent :: URI -> ContentInfo -> Torrent
-simpleTorrent announce info = torrent announce info
+simpleTorrent announce info = mktorrent announce info
                               Nothing Nothing Nothing
                               Nothing Nothing Nothing
                               Nothing Nothing
diff --git a/src/Network/BitTorrent.hs b/src/Network/BitTorrent.hs
index 8c3189d3..c166b1b1 100644
--- a/src/Network/BitTorrent.hs
+++ b/src/Network/BitTorrent.hs
@@ -46,8 +46,7 @@ import System.Torrent.Storage
 -- | Client session with default parameters. Use it for testing only.
 withDefaultClient :: PortNumber -> PortNumber -> (ClientSession -> IO ()) -> IO ()
 withDefaultClient dhtPort listPort action = do
-  withClientSession defaultThreadCount defaultExtensions $ \client -> do
-    action client
+  withClientSession defaultThreadCount defaultExtensions listPort dhtPort action
 
 {-----------------------------------------------------------------------
     Torrent management
diff --git a/src/Network/BitTorrent/Sessions.hs b/src/Network/BitTorrent/Sessions.hs
index 2e51fde6..9a1d0c6a 100644
--- a/src/Network/BitTorrent/Sessions.hs
+++ b/src/Network/BitTorrent/Sessions.hs
@@ -48,7 +48,7 @@ module Network.BitTorrent.Sessions
        , discover
        ) where
 
-import Prelude hiding (mapM_)
+import Prelude hiding (mapM_, elem)
 
 import Control.Applicative
 import Control.Concurrent
@@ -61,6 +61,7 @@ import Control.Monad.Trans
 import Data.IORef
 import Data.Map as M
 import Data.HashMap.Strict as HM
+import Data.Foldable as F
 import Data.Set as S
 
 import Data.Serialize hiding (get)
@@ -141,14 +142,13 @@ startDHT ClientSession {..} nodePort = withRunning peerListener failure start
 
 -- | 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.
-                  -> [Extension]      -- ^ Extensions allowed to use.
-                  -> IO ClientSession -- ^ Client with unique peer ID.
-openClientSession n exts = do
+openClientSession :: SessionCount -> [Extension] -> PortNumber -> PortNumber -> IO ClientSession
+openClientSession n exts listenerPort _ = do
   mgr <- Ev.new
   -- TODO kill this thread when leave client
   _   <- forkIO $ loop mgr
-  ClientSession
+
+  cs <- ClientSession
     <$> genPeerId
     <*> pure exts
     <*> newEmptyMVar
@@ -160,13 +160,21 @@ openClientSession n exts = do
     <*> newTVarIO (startProgress 0)
     <*> newTVarIO HM.empty
 
+  startListener cs listenerPort
+  return cs
+
 closeClientSession :: ClientSession -> IO ()
-closeClientSession ClientSession {..} =
-  stopService nodeListener `finally` stopService peerListener
--- TODO stop all swarm sessions
+closeClientSession ClientSession {..} = do
+  stopService nodeListener
+  stopService peerListener
+
+  sws <- readTVarIO swarmSessions
+  forM_ sws closeSwarmSession
 
-withClientSession :: SessionCount -> [Extension] -> (ClientSession -> IO ()) -> IO ()
-withClientSession c es = bracket (openClientSession c es) closeClientSession
+withClientSession :: SessionCount -> [Extension]
+                  -> PortNumber -> PortNumber
+                  -> (ClientSession -> IO ()) -> IO ()
+withClientSession c es l d = bracket (openClientSession c es l d) closeClientSession
 
 -- | Get current global progress of the client. This value is usually
 -- shown to a user.
@@ -222,6 +230,14 @@ discover swarm @ SwarmSession {..} action = {-# SCC discover #-} do
         initiatePeerSession swarm addr $ \conn ->
           runP2P conn action
 
+registerSwarmSession :: SwarmSession -> IO ()
+registerSwarmSession = undefined
+
+unregisterSwarmSession :: SwarmSession -> IO ()
+unregisterSwarmSession SwarmSession {..} =
+  atomically $ modifyTVar (swarmSessions clientSession) $
+    M.delete $ tInfoHash torrentMeta
+
 newSwarmSession :: Int -> Bitfield -> ClientSession -> Torrent
                 -> IO SwarmSession
 newSwarmSession n bf cs @ ClientSession {..} t @ Torrent {..}
@@ -246,11 +262,6 @@ closeSwarmSession se @ SwarmSession {..} = do
   -- TODO the order is important!
   closeStorage storage
 
-unregisterSwarmSession :: SwarmSession -> IO ()
-unregisterSwarmSession SwarmSession {..} =
-  atomically $ modifyTVar (swarmSessions clientSession) $
-  M.delete $ tInfoHash torrentMeta
-
 getSwarm :: ClientSession -> InfoHash -> IO SwarmSession
 getSwarm cs @ ClientSession {..} ih = do
   ss <- readTVarIO swarmSessions
diff --git a/src/Network/BitTorrent/Sessions/Types.lhs b/src/Network/BitTorrent/Sessions/Types.lhs
index f94dbfa6..3f9c6db1 100644
--- a/src/Network/BitTorrent/Sessions/Types.lhs
+++ b/src/Network/BitTorrent/Sessions/Types.lhs
@@ -73,10 +73,12 @@ Thread layout
 
 When client session created 2 new threads appear:
 
-  * DHT listener - replies to DHT requests;
-
   * Peer listener - accept new P2P connection initiated by other
-peers.
+peers;
+
+  * Tracker announcer - announce that the peer have this torrent.
+
+  * OPTIONAL: DHT listener - replies to DHT requests;
 
 When swarn session created 3 new threads appear:
 
diff --git a/tests/Main.hs b/tests/Main.hs
index adb50380..c0ef52db 100644
--- a/tests/Main.hs
+++ b/tests/Main.hs
@@ -118,7 +118,7 @@ instance Arbitrary ContentInfo where
     ]
 
 instance Arbitrary Torrent where
-  arbitrary = torrent <$> arbitrary
+  arbitrary = mktorrent <$> arbitrary
                  <*> arbitrary <*> arbitrary <*> arbitrary
                  <*> arbitrary <*> arbitrary <*> arbitrary
                  <*> arbitrary <*> pure Nothing <*> arbitrary