6 files changed, 473 insertions, 382 deletions

diff --git a/bittorrent.cabal b/bittorrent.cabal
index 22c2794a..fc19fef9 100644
--- a/bittorrent.cabal
+++ b/bittorrent.cabal
@@ -44,9 +44,10 @@ library
                      , Network.BitTorrent.Tracker
                      , Network.BitTorrent.Exchange
                      , Network.BitTorrent.DHT
+                     , Network.BitTorrent.Discovery
                      , System.Torrent.Storage
-                       Network.BitTorrent.Internal
-  other-modules:
+
+  other-modules:       Network.BitTorrent.Internal
   if flag(testing)
      exposed-modules:  Network.BitTorrent.Exchange.Protocol
                      , Network.BitTorrent.Tracker.Protocol
@@ -54,6 +55,7 @@ library
   if !flag(testing)
      other-modules:    Network.BitTorrent.Exchange.Protocol
                      , Network.BitTorrent.Tracker.Protocol
+                     , Network.BitTorrent.DHT.Protocol
                      , System.IO.MMap.Fixed
 
   build-depends:       base       == 4.*
diff --git a/src/Network/BitTorrent.hs b/src/Network/BitTorrent.hs
index 3602dd7e..06df77dd 100644
--- a/src/Network/BitTorrent.hs
+++ b/src/Network/BitTorrent.hs
@@ -17,8 +17,7 @@ module Network.BitTorrent
          -- ** Client
        , ClientSession( clientPeerId, allowedExtensions )
 
-       , newClient
-       , defaultClient
+       , withDefaultClient
 
        , Progress(..)
        , getCurrentProgress
@@ -91,36 +90,25 @@ import Network
 import Data.Bitfield as BF
 import Data.Torrent
 import Network.BitTorrent.Internal
+import Network.BitTorrent.Peer
+import Network.BitTorrent.Extension
 import Network.BitTorrent.Exchange
 import Network.BitTorrent.Exchange.Protocol
 import Network.BitTorrent.Tracker
-import Network.BitTorrent.Extension
-import Network.BitTorrent.Peer
+import Network.BitTorrent.Discovery
 
 import System.Torrent.Storage
 
+-- TODO remove fork from Network.BitTorrent.Exchange
+-- TODO make all forks in Internal.
 
 -- | Client session with default parameters. Use it for testing only.
-defaultClient :: IO ClientSession
-defaultClient = newClient defaultThreadCount defaultExtensions
-
--- 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
-  let conn = TConnection (tAnnounce torrentMeta)
-                         (tInfoHash torrentMeta)
-                         (clientPeerId clientSession)
-                         (listenerPort clientSession)
-
-  progress <- getCurrentProgress clientSession
-
-  withTracker progress conn $ \tses -> do
-    forever $ do
-      addr <- getPeerAddr tses
-      spawnP2P swarm addr $ do
-        action
+withDefaultClient :: PortNumber -> PortNumber -> (ClientSession -> IO ()) -> IO ()
+withDefaultClient dhtPort listPort action = do
+  withClientSession defaultThreadCount defaultExtensions $ \client -> do
+    startListener client listPort
+    startDHT      client dhtPort
+    action client
 
 {-----------------------------------------------------------------------
     Torrent management
diff --git a/src/Network/BitTorrent/DHT.hs b/src/Network/BitTorrent/DHT.hs
index f3c993c3..b0aac002 100644
--- a/src/Network/BitTorrent/DHT.hs
+++ b/src/Network/BitTorrent/DHT.hs
@@ -1,338 +1,6 @@
-{-# LANGUAGE OverloadedStrings #-}
-{-# LANGUAGE RecordWildCards   #-}
 module Network.BitTorrent.DHT
-       (
-         newNodeSession
-
-         -- * Tracker
-       , ping
-       , findNode
-       , getPeers
-       , announcePeer
-
-         -- * Server
+       ( newNodeSession
        , dhtServer
        ) where
 
-import Control.Applicative
-import Control.Concurrent
-import Control.Concurrent.STM
-import Control.Monad
-import Control.Exception
-import Data.ByteString
-import Data.Serialize as S
-import Data.Function
-import Data.Ord
-import Data.Maybe
-import Data.List as L
-import Data.Map as M
-import Data.HashMap.Strict as HM
-import Network
-import Network.Socket
-import System.Entropy
-
-import Remote.KRPC
-import Remote.KRPC.Protocol
-import Data.BEncode
-import Data.Torrent
-import Network.BitTorrent.Peer
-
-{-----------------------------------------------------------------------
-    Node
------------------------------------------------------------------------}
-
-type NodeId = ByteString
-
--- WARN is the 'system' random suitable for this?
--- | Generate random NodeID used for the entire session.
---   Distribution of ID's should be as uniform as possible.
---
-genNodeId :: IO NodeId
-genNodeId = getEntropy 20
-
-instance Serialize PortNumber where
-  get = fromIntegral <$> getWord16be
-  put = putWord16be . fromIntegral
-
-
-data NodeAddr = NodeAddr {
-    nodeIP   :: {-# UNPACK #-} !HostAddress
-  , nodePort :: {-# UNPACK #-} !PortNumber
-  } deriving (Show, Eq)
-
-instance Serialize NodeAddr where
-  get = NodeAddr <$> getWord32be <*> get
-  put NodeAddr {..} = do
-    putWord32be nodeIP
-    put         nodePort
-
-
-data NodeInfo = NodeInfo {
-    nodeID   :: !NodeId
-  , nodeAddr :: !NodeAddr
-  } deriving (Show, Eq)
-
-instance Serialize NodeInfo where
-  get = NodeInfo <$> getByteString 20 <*> get
-  put NodeInfo {..} = put nodeID >> put nodeAddr
-
-type CompactInfo = ByteString
-
-decodeCompact :: CompactInfo -> [NodeInfo]
-decodeCompact = either (const []) id . S.runGet (many get)
-
-encodeCompact :: [NodeId] -> CompactInfo
-encodeCompact = S.runPut . mapM_ put
-
-decodePeerList :: [BEncode] -> [PeerAddr]
-decodePeerList = undefined
-
-encodePeerList :: [PeerAddr] -> [BEncode]
-encodePeerList = undefined
-
-type Distance = NodeId
-
-
-{-----------------------------------------------------------------------
-    Tokens
------------------------------------------------------------------------}
-
-type Secret = Int
-
-genSecret :: IO Secret
-genSecret = undefined
-
--- | Instead of periodically loop over the all nodes in the routing
---   table with some given interval (or some other tricky method
---   e.g. using timeouts) we can just update tokens on demand - if no
---   one asks for a token then the token _should_ not change at all.
---
-type Token = ByteString
-
-genToken :: NodeAddr -> Secret -> Token
-genToken = return undefined
-
-defaultToken :: Token
-defaultToken = "0xdeadbeef"
-
-{-----------------------------------------------------------------------
-    Routing table
------------------------------------------------------------------------}
-
-type ContactInfo  = HashMap InfoHash [PeerAddr]
-
-insertPeer :: InfoHash -> PeerAddr -> ContactInfo -> ContactInfo
-insertPeer ih addr = HM.insertWith (++) ih [addr]
-
-lookupPeers :: InfoHash -> ContactInfo -> [PeerAddr]
-lookupPeers ih = fromMaybe [] . HM.lookup ih
-
--- TODO use more compact routing table
-type RoutingTable = HashMap NodeId NodeAddr
-
-insertNode :: NodeId -> NodeAddr -> RoutingTable -> RoutingTable
-insertNode = HM.insert
-
-type Alpha = Int
-
-defaultAlpha :: Alpha
-defaultAlpha = 8
-
--- TODO
-kclosest :: Int -> NodeId -> RoutingTable -> [NodeId]
-kclosest = undefined
-
-{-----------------------------------------------------------------------
-    Node session
------------------------------------------------------------------------}
-
-data NodeSession = NodeSession {
-    nodeId       :: !NodeId
-  , routingTable :: !(TVar RoutingTable)
-  , contactInfo  :: !(TVar ContactInfo)
-  , alpha        :: !Alpha
-  , listenerPort :: !PortNumber
-  }
-
-instance Eq NodeSession where
-  (==) = (==) `on` nodeId
-
-instance Ord NodeSession where
-  compare = comparing nodeId
-
-newNodeSession :: PortNumber -> IO NodeSession
-newNodeSession lport
-  = NodeSession
-    <$> genNodeId
-    <*> newTVarIO HM.empty
-    <*> newTVarIO HM.empty
-    <*> pure defaultAlpha
-    <*> pure lport
-
-assignToken :: NodeSession -> NodeId -> IO Token
-assignToken _ _ = return ""
-
--- TODO
-checkToken :: NodeId -> Token -> NodeSession -> IO Bool
-checkToken nid token _ = return True
-
-updateTimestamp :: NodeSession -> NodeId -> IO ()
-updateTimestamp = error "updateTimestamp"
-
-updateToken :: NodeSession -> NodeId -> Token -> IO ()
-updateToken _ _ _ = error "updateToken"
-
-{-----------------------------------------------------------------------
-    DHT Queries
------------------------------------------------------------------------}
-
-pingM :: Method NodeId NodeId
-pingM = method "ping" ["id"] ["id"]
-
-findNodeM :: Method (NodeId, NodeId) (NodeId, CompactInfo)
-findNodeM = method "find_node" ["id", "target"] ["id", "nodes"]
-
--- | Lookup peers by a torrent infohash. This method might return
--- different kind of responses depending on the routing table of
--- queried node:
---
---   * If quieried node contains a peer list for the given infohash
---   then the node should return the list in a "value" key. Note that
---   list is encoded as compact peer address, not a compact node info.
---   The result of 'get_peers' method have the following scheme:
---
---      > { "id"     : "dht_server_node_id"
---      > , "token"  : "assigned_token"
---      > , "values" : ["_IP_PO", "_ip_po"]
---      > }
---
---   * If quieried node does not contain a list of peers associated
---   with the given infohash, then node should return
---
---      > { "id"    : "dht_server_node_id"
---      > , "token" : "assigned_token"
---      > , "nodes" : "compact_nodes_info"
---      > }
---
---   The resulting dictionaries might differ only in a values\/nodes
---   keys.
---
-getPeersM :: Method (NodeId, InfoHash) BEncode
-getPeersM = method "get_peers" ["id", "info_hash"] []
-
--- | Used to announce that the peer, controlling the quering node is
--- downloading a torrent on a port.
-announcePeerM :: Method (NodeId, InfoHash, PortNumber, Token) NodeId
-announcePeerM = method "announce_peer" ["id", "info_hash", "port", "token"] ["id"]
-
-{-----------------------------------------------------------------------
-    DHT Tracker
------------------------------------------------------------------------}
--- TODO: update node timestamp on each successful call
-
--- | Note that tracker side query functions could throw RPCException.
-type DHT a b  = NodeSession -> NodeAddr -> a -> IO b
-
-ping :: DHT () ()
-ping NodeSession {..} addr @ NodeAddr {..} () = do
-  nid <- call (nodeIP, nodePort) pingM nodeId
-  atomically $ modifyTVar' routingTable $ HM.insert nid addr
-
-findNode :: DHT NodeId [NodeInfo]
-findNode ses @ NodeSession {..} NodeAddr {..} qnid = do
-  (nid, info) <- call (nodeIP, nodePort) findNodeM (nodeId, qnid)
-  updateTimestamp ses nid
-  return (decodeCompact info)
-
-getPeers :: DHT InfoHash (Either [NodeInfo] [PeerAddr])
-getPeers ses @ NodeSession {..} NodeAddr {..} ih = do
-    resp <- call (nodeIP, nodePort) getPeersM (nodeId, ih)
-    (nid, tok, res) <- extrResp resp
-    updateTimestamp ses nid
-    updateToken ses nid tok
-    return res
-  where
-    extrResp (BDict d)
-      | Just (BString nid   ) <- M.lookup "id"     d
-      , Just (BString tok   ) <- M.lookup "token"  d
-      , Just (BList   values) <- M.lookup "values" d
-      = return $ (nid, tok, Right $ decodePeerList values)
-
-      | Just (BString nid   ) <- M.lookup "id"     d
-      , Just (BString tok   ) <- M.lookup "token"  d
-      , Just (BString nodes)  <- M.lookup "nodes"  d
-      = return (nid, tok, Left $ decodeCompact nodes)
-
-    extrResp  _ = throw $ RPCException msg
-      where msg = ProtocolError "unable to extract getPeers resp"
-
--- remove token from signature, handle the all token stuff by NodeSession
-
--- | Note that before ever calling this method you should call the
--- getPeerList.
-announcePeer :: DHT (InfoHash, Token) NodeId
-announcePeer ses @ NodeSession {..} NodeAddr {..} (ih, tok) = do
-  nid <- call (nodeIP, nodePort) announcePeerM (nodeId, ih, listenerPort, tok)
-  updateTimestamp ses nid
-  return nid
-
-{-----------------------------------------------------------------------
-    DHT Server
------------------------------------------------------------------------}
--- TODO: update node timestamp on each successful call
--- NOTE: ensure all server operations run in O(1)
-
-type ServerHandler a b = NodeSession -> NodeAddr -> a -> IO b
-
-pingS :: ServerHandler NodeId NodeId
-pingS NodeSession {..} addr nid = do
-  atomically $ modifyTVar' routingTable $ insertNode nid addr
-  return nodeId
-
-findNodeS :: ServerHandler (NodeId, NodeId) (NodeId, CompactInfo)
-findNodeS ses @ NodeSession {..} _ (nid, qnid) = do
-  updateTimestamp ses nid
-  rt <- atomically $ readTVar routingTable
-  return (nodeId, encodeCompact $ kclosest alpha qnid rt)
-
-getPeersS :: ServerHandler (NodeId, InfoHash) BEncode
-getPeersS ses @ NodeSession {..} _ (nid, ih) = do
-    updateTimestamp ses nid
-    mkResp <$> assignToken ses nid <*> findPeers
-  where
-    findPeers = do
-      list <- lookupPeers ih <$> readTVarIO contactInfo
-      if not (L.null list)
-        then return $ Right list
-        else do
-          rt    <- readTVarIO routingTable
-          let nodes = kclosest alpha (getInfoHash ih) rt
-          return $ Left nodes
-
-    mkDict tok res = [("id",BString nodeId), ("token", BString tok), res]
-    mkResult (Left  nodes ) = ("nodes",  BString $ encodeCompact  nodes)
-    mkResult (Right values) = ("values", BList   $ encodePeerList values)
-    mkResp tok = BDict . M.fromList . mkDict tok . mkResult
-
-announcePeerS :: ServerHandler (NodeId, InfoHash, PortNumber, Token) NodeId
-announcePeerS ses @ NodeSession {..} NodeAddr {..} (nid, ih, port, token) = do
-  updateTimestamp ses nid
-  registered <- checkToken nid token ses
-  when registered $ do
-    atomically $ do
-      let peerAddr = PeerAddr Nothing nodeIP port
-      modifyTVar contactInfo $ insertPeer ih peerAddr
-  return nodeId
-
-dhtTracker :: NodeSession -> InfoHash -> Chan PeerAddr -> IO ()
-dhtTracker  = undefined
-
-dhtServer :: NodeSession -> PortNumber -> IO ()
-dhtServer s p = server p methods
-  where
-    methods =
-      [ pingM         ==> pingS         s undefined
-      , findNodeM     ==> findNodeS     s undefined
-      , getPeersM     ==> getPeersS     s undefined
-      , announcePeerM ==> announcePeerS s undefined
-      ]
\ No newline at end of file
+import Network.BitTorrent.DHT.Protocol
\ No newline at end of file
diff --git a/src/Network/BitTorrent/DHT/Protocol.hs b/src/Network/BitTorrent/DHT/Protocol.hs
new file mode 100644
index 00000000..5267a916
--- /dev/null
+++ b/src/Network/BitTorrent/DHT/Protocol.hs
@@ -0,0 +1,339 @@
+{-# LANGUAGE OverloadedStrings #-}
+{-# LANGUAGE RecordWildCards   #-}
+module Network.BitTorrent.DHT.Protocol
+       (
+         newNodeSession
+
+         -- * Tracker
+       , ping
+       , findNode
+       , getPeers
+       , announcePeer
+
+         -- * Server
+       , dhtServer
+       ) where
+
+import Control.Applicative
+import Control.Concurrent
+import Control.Concurrent.STM
+import Control.Monad
+import Control.Exception
+import Data.ByteString
+import Data.Serialize as S
+import Data.Function
+import Data.Ord
+import Data.Maybe
+import Data.List as L
+import Data.Map as M
+import Data.HashMap.Strict as HM
+import Network
+import Network.Socket
+import System.Entropy
+
+import Remote.KRPC
+import Remote.KRPC.Protocol
+import Data.BEncode
+import Data.Torrent
+import Network.BitTorrent.Peer
+
+{-----------------------------------------------------------------------
+    Node
+-----------------------------------------------------------------------}
+
+type NodeId = ByteString
+
+-- WARN is the 'system' random suitable for this?
+-- | Generate random NodeID used for the entire session.
+--   Distribution of ID's should be as uniform as possible.
+--
+genNodeId :: IO NodeId
+genNodeId = getEntropy 20
+
+instance Serialize PortNumber where
+  get = fromIntegral <$> getWord16be
+  put = putWord16be . fromIntegral
+
+
+data NodeAddr = NodeAddr {
+    nodeIP   :: {-# UNPACK #-} !HostAddress
+  , nodePort :: {-# UNPACK #-} !PortNumber
+  } deriving (Show, Eq)
+
+instance Serialize NodeAddr where
+  get = NodeAddr <$> getWord32be <*> get
+  put NodeAddr {..} = do
+    putWord32be nodeIP
+    put         nodePort
+
+
+data NodeInfo = NodeInfo {
+    nodeID   :: !NodeId
+  , nodeAddr :: !NodeAddr
+  } deriving (Show, Eq)
+
+instance Serialize NodeInfo where
+  get = NodeInfo <$> getByteString 20 <*> get
+  put NodeInfo {..} = put nodeID >> put nodeAddr
+
+type CompactInfo = ByteString
+
+decodeCompact :: CompactInfo -> [NodeInfo]
+decodeCompact = either (const []) id . S.runGet (many get)
+
+encodeCompact :: [NodeId] -> CompactInfo
+encodeCompact = S.runPut . mapM_ put
+
+decodePeerList :: [BEncode] -> [PeerAddr]
+decodePeerList = undefined
+
+encodePeerList :: [PeerAddr] -> [BEncode]
+encodePeerList = undefined
+
+type Distance = NodeId
+
+{-----------------------------------------------------------------------
+    Tokens
+-----------------------------------------------------------------------}
+
+type Secret = Int
+
+genSecret :: IO Secret
+genSecret = error "secret"
+
+-- | Instead of periodically loop over the all nodes in the routing
+--   table with some given interval (or some other tricky method
+--   e.g. using timeouts) we can just update tokens on demand - if no
+--   one asks for a token then the token _should_ not change at all.
+--
+type Token = ByteString
+
+defaultToken :: Token
+defaultToken = "0xdeadbeef"
+
+genToken :: NodeAddr -> Secret -> Token
+genToken _ _ = defaultToken
+
+{-----------------------------------------------------------------------
+    Routing table
+-----------------------------------------------------------------------}
+
+type ContactInfo  = HashMap InfoHash [PeerAddr]
+
+insertPeer :: InfoHash -> PeerAddr -> ContactInfo -> ContactInfo
+insertPeer ih addr = HM.insertWith (++) ih [addr]
+
+lookupPeers :: InfoHash -> ContactInfo -> [PeerAddr]
+lookupPeers ih = fromMaybe [] . HM.lookup ih
+
+-- TODO use more compact routing table
+type RoutingTable = HashMap NodeId NodeAddr
+
+insertNode :: NodeId -> NodeAddr -> RoutingTable -> RoutingTable
+insertNode = HM.insert
+
+type Alpha = Int
+
+defaultAlpha :: Alpha
+defaultAlpha = 8
+
+-- TODO
+kclosest :: Int -> NodeId -> RoutingTable -> [NodeId]
+kclosest = undefined
+
+{-----------------------------------------------------------------------
+    Node session
+-----------------------------------------------------------------------}
+
+data NodeSession = NodeSession {
+    nodeId        :: !NodeId
+  , routingTable  :: !(TVar RoutingTable)
+  , contactInfo   :: !(TVar ContactInfo)
+--  , currentSecret :: !(TVar Secret)
+--  , secretTimestamp :: !(TVar Timestamp)
+  , alpha         :: !Alpha
+  , listenerPort  :: !PortNumber
+  }
+
+instance Eq NodeSession where
+  (==) = (==) `on` nodeId
+
+instance Ord NodeSession where
+  compare = comparing nodeId
+
+newNodeSession :: PortNumber -> IO NodeSession
+newNodeSession lport
+  = NodeSession
+    <$> genNodeId
+    <*> newTVarIO HM.empty
+    <*> newTVarIO HM.empty
+    <*> pure defaultAlpha
+    <*> pure lport
+
+assignToken :: NodeSession -> NodeId -> IO Token
+assignToken _ _ = return ""
+
+-- TODO
+checkToken :: NodeId -> Token -> NodeSession -> IO Bool
+checkToken nid token _ = return True
+
+updateTimestamp :: NodeSession -> NodeId -> IO ()
+updateTimestamp = error "updateTimestamp"
+
+updateToken :: NodeSession -> NodeId -> Token -> IO ()
+updateToken _ _ _ = error "updateToken"
+
+{-----------------------------------------------------------------------
+    DHT Queries
+-----------------------------------------------------------------------}
+
+pingM :: Method NodeId NodeId
+pingM = method "ping" ["id"] ["id"]
+
+findNodeM :: Method (NodeId, NodeId) (NodeId, CompactInfo)
+findNodeM = method "find_node" ["id", "target"] ["id", "nodes"]
+
+-- | Lookup peers by a torrent infohash. This method might return
+-- different kind of responses depending on the routing table of
+-- queried node:
+--
+--   * If quieried node contains a peer list for the given infohash
+--   then the node should return the list in a "value" key. Note that
+--   list is encoded as compact peer address, not a compact node info.
+--   The result of 'get_peers' method have the following scheme:
+--
+--      > { "id"     : "dht_server_node_id"
+--      > , "token"  : "assigned_token"
+--      > , "values" : ["_IP_PO", "_ip_po"]
+--      > }
+--
+--   * If quieried node does not contain a list of peers associated
+--   with the given infohash, then node should return
+--
+--      > { "id"    : "dht_server_node_id"
+--      > , "token" : "assigned_token"
+--      > , "nodes" : "compact_nodes_info"
+--      > }
+--
+--   The resulting dictionaries might differ only in a values\/nodes
+--   keys.
+--
+getPeersM :: Method (NodeId, InfoHash) BEncode
+getPeersM = method "get_peers" ["id", "info_hash"] []
+
+-- | Used to announce that the peer, controlling the quering node is
+-- downloading a torrent on a port.
+announcePeerM :: Method (NodeId, InfoHash, PortNumber, Token) NodeId
+announcePeerM = method "announce_peer" ["id", "info_hash", "port", "token"] ["id"]
+
+{-----------------------------------------------------------------------
+    DHT Tracker
+-----------------------------------------------------------------------}
+-- TODO: update node timestamp on each successful call
+
+-- | Note that tracker side query functions could throw RPCException.
+type DHT a b  = NodeSession -> NodeAddr -> a -> IO b
+
+ping :: DHT () ()
+ping NodeSession {..} addr @ NodeAddr {..} () = do
+  nid <- call (nodeIP, nodePort) pingM nodeId
+  atomically $ modifyTVar' routingTable $ HM.insert nid addr
+
+findNode :: DHT NodeId [NodeInfo]
+findNode ses @ NodeSession {..} NodeAddr {..} qnid = do
+  (nid, info) <- call (nodeIP, nodePort) findNodeM (nodeId, qnid)
+  updateTimestamp ses nid
+  return (decodeCompact info)
+
+getPeers :: DHT InfoHash (Either [NodeInfo] [PeerAddr])
+getPeers ses @ NodeSession {..} NodeAddr {..} ih = do
+    resp <- call (nodeIP, nodePort) getPeersM (nodeId, ih)
+    (nid, tok, res) <- extrResp resp
+    updateTimestamp ses nid
+    updateToken ses nid tok
+    return res
+  where
+    extrResp (BDict d)
+      | Just (BString nid   ) <- M.lookup "id"     d
+      , Just (BString tok   ) <- M.lookup "token"  d
+      , Just (BList   values) <- M.lookup "values" d
+      = return $ (nid, tok, Right $ decodePeerList values)
+
+      | Just (BString nid   ) <- M.lookup "id"     d
+      , Just (BString tok   ) <- M.lookup "token"  d
+      , Just (BString nodes)  <- M.lookup "nodes"  d
+      = return (nid, tok, Left $ decodeCompact nodes)
+
+    extrResp  _ = throw $ RPCException msg
+      where msg = ProtocolError "unable to extract getPeers resp"
+
+-- remove token from signature, handle the all token stuff by NodeSession
+
+-- | Note that before ever calling this method you should call the
+-- getPeerList.
+announcePeer :: DHT (InfoHash, Token) NodeId
+announcePeer ses @ NodeSession {..} NodeAddr {..} (ih, tok) = do
+  nid <- call (nodeIP, nodePort) announcePeerM (nodeId, ih, listenerPort, tok)
+  updateTimestamp ses nid
+  return nid
+
+{-----------------------------------------------------------------------
+    DHT Server
+-----------------------------------------------------------------------}
+-- TODO: update node timestamp on each successful call
+-- NOTE: ensure all server operations run in O(1)
+
+type ServerHandler a b = NodeSession -> NodeAddr -> a -> IO b
+
+pingS :: ServerHandler NodeId NodeId
+pingS NodeSession {..} addr nid = do
+  atomically $ modifyTVar' routingTable $ insertNode nid addr
+  return nodeId
+
+findNodeS :: ServerHandler (NodeId, NodeId) (NodeId, CompactInfo)
+findNodeS ses @ NodeSession {..} _ (nid, qnid) = do
+  updateTimestamp ses nid
+  rt <- atomically $ readTVar routingTable
+  return (nodeId, encodeCompact $ kclosest alpha qnid rt)
+
+getPeersS :: ServerHandler (NodeId, InfoHash) BEncode
+getPeersS ses @ NodeSession {..} _ (nid, ih) = do
+    updateTimestamp ses nid
+    mkResp <$> assignToken ses nid <*> findPeers
+  where
+    findPeers = do
+      list <- lookupPeers ih <$> readTVarIO contactInfo
+      if not (L.null list)
+        then return $ Right list
+        else do
+          rt    <- readTVarIO routingTable
+          let nodes = kclosest alpha (getInfoHash ih) rt
+          return $ Left nodes
+
+    mkDict tok res = [("id",BString nodeId), ("token", BString tok), res]
+    mkResult (Left  nodes ) = ("nodes",  BString $ encodeCompact  nodes)
+    mkResult (Right values) = ("values", BList   $ encodePeerList values)
+    mkResp tok = BDict . M.fromList . mkDict tok . mkResult
+
+announcePeerS :: ServerHandler (NodeId, InfoHash, PortNumber, Token) NodeId
+announcePeerS ses @ NodeSession {..} NodeAddr {..} (nid, ih, port, token) = do
+  updateTimestamp ses nid
+  registered <- checkToken nid token ses
+  when registered $ do
+    atomically $ do
+      let peerAddr = PeerAddr Nothing nodeIP port
+      modifyTVar contactInfo $ insertPeer ih peerAddr
+  return nodeId
+
+dhtTracker :: NodeSession -> InfoHash -> Chan PeerAddr -> IO ()
+dhtTracker  = undefined
+
+dhtServer :: NodeSession -> PortNumber -> IO ()
+dhtServer s p = server p methods
+  where
+    methods =
+      [ pingM         ==> pingS         s undefined
+      , findNodeM     ==> findNodeS     s undefined
+      , getPeersM     ==> getPeersS     s undefined
+      , announcePeerM ==> announcePeerS s undefined
+      ]
\ No newline at end of file
diff --git a/src/Network/BitTorrent/Discovery.hs b/src/Network/BitTorrent/Discovery.hs
new file mode 100644
index 00000000..770ae818
--- /dev/null
+++ b/src/Network/BitTorrent/Discovery.hs
@@ -0,0 +1,56 @@
+{-# 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.Internal
+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
+      spawnP2P swarm addr $ do
+        action
+
+startListener :: ClientSession -> PortNumber -> IO ()
+startListener cs @ ClientSession {..} port =
+  putMVar peerListener =<< startService port (listener cs (error "listener"))
+
+startDHT :: ClientSession -> PortNumber -> IO ()
+startDHT ClientSession {..} nodePort = do
+    maybe failure start =<< tryTakeMVar peerListener
+  where
+    start ClientService {..} = do
+      ses  <- newNodeSession servPort
+      serv <- startService nodePort (dhtServer ses)
+      putMVar nodeListener serv
+
+    failure = throwIO $ userError msg
+    msg = "unable to start DHT server: peer listener is not running"
diff --git a/src/Network/BitTorrent/Internal.lhs b/src/Network/BitTorrent/Internal.lhs
index 15606c57..8461a841 100644
--- a/src/Network/BitTorrent/Internal.lhs
+++ b/src/Network/BitTorrent/Internal.lhs
@@ -23,10 +23,18 @@
 >        ( -- * Progress
 >          Progress(..), startProgress
 >
+>        , ClientService(..)
+>
 >          -- * Client
->        , ClientSession (clientPeerId, allowedExtensions)
+>        , ClientSession ( ClientSession
+>                        , clientPeerId, allowedExtensions
+>                        , nodeListener, peerListener
+>                        )
+>        , withClientSession
 >        , listenerPort, dhtPort
 >
+>        , startService
+>
 >        , ThreadCount
 >        , defaultThreadCount
 >
@@ -34,8 +42,6 @@
 >        , registerTorrent
 >        , unregisterTorrent
 >
->        , newClient
->
 >        , getCurrentProgress
 >        , getSwarmCount
 >        , getPeerCount
@@ -64,6 +70,7 @@
 >        , SessionState
 >        , initiatePeerSession
 >        , acceptPeerSession
+>        , listener
 >
 >          -- ** Broadcasting
 >        , available
@@ -89,7 +96,7 @@
 > import Control.Concurrent.STM
 > import Control.Concurrent.MSem as MSem
 > import Control.Lens
-> import Control.Monad (when)
+> import Control.Monad (when, forever)
 > import Control.Exception
 > import Control.Monad.Trans
 
@@ -106,7 +113,7 @@
 > import Data.Serialize hiding (get)
 > import Text.PrettyPrint
 
-> import Network
+> import Network hiding (accept)
 > import Network.Socket
 > import Network.Socket.ByteString
 
@@ -118,6 +125,7 @@
 > import Network.BitTorrent.Peer
 > import Network.BitTorrent.Exchange.Protocol as BT
 > import Network.BitTorrent.Tracker.Protocol as BT
+> import Network.BitTorrent.DHT.Protocol as BT
 
 Progress
 ------------------------------------------------------------------------
@@ -195,7 +203,7 @@ Peer session is one always forked thread.
 When client\/swarm\/peer session gets closed kill the corresponding
 threads, but flush data to disc. (for e.g. storage block map)
 
-So for e.g., in order to obtain our first block we need to run at
+So for e.g., in order to obtain our first block we need to spawn at
 least 7 threads: main thread, 2 client session threads, 3 swarm session
 threads and PeerSession thread.
 
@@ -296,10 +304,8 @@ so we can abstract out into ClientService:
 >   , servThread :: !ThreadId
 >   } deriving Show
 
-startService :: PortNumber -> IO a -> IO ClientService
-startService p m = forkIO $ handle $ m p
-  where
-    handle :: IOError -> IO ()
+> startService :: PortNumber -> (PortNumber -> IO ()) -> IO ClientService
+> startService port m = ClientService port <$> forkIO (m port)
 
 > stopService :: ClientService -> IO ()
 > stopService ClientService {..} = killThread servThread
@@ -339,8 +345,8 @@ and different enabled extensions at the same time.
 >     -- 'PeerSession'.
 >   , allowedExtensions :: [Extension]
 
->   , peerListener      :: !ClientService
->   , nodeListener      :: !ClientService
+>   , peerListener      :: !(MVar ClientService)
+>   , nodeListener      :: !(MVar ClientService)
 
 >     -- | Semaphor used to bound number of active P2P sessions.
 >   , activeThreads     :: !(MSem ThreadCount)
@@ -407,20 +413,20 @@ Retrieving client info
 
 > -- | Create a new client session. The data passed to this function are
 > -- usually loaded from configuration file.
-> newClient :: SessionCount     -- ^ Maximum count of active P2P Sessions.
+> openClientSession :: SessionCount     -- ^ Maximum count of active P2P Sessions.
 >           -> [Extension]      -- ^ Extensions allowed to use.
 >           -> IO ClientSession -- ^ Client with unique peer ID.
 
-> newClient n exts = do
+> openClientSession n exts = do
 >   mgr <- Ev.new
 >   -- TODO kill this thread when leave client
 >   _   <- forkIO $ loop mgr
-
+>
 >   ClientSession
 >     <$> genPeerId
 >     <*> pure exts
->     <*> pure (ClientService 10 undefined) -- TODO
->     <*> pure (ClientService 20 undefined) -- TODO
+>     <*> newEmptyMVar
+>     <*> newEmptyMVar
 >     <*> MSem.new n
 >     <*> pure n
 >     <*> newTVarIO M.empty
@@ -428,11 +434,21 @@ Retrieving client info
 >     <*> newTVarIO (startProgress 0)
 >     <*> newTVarIO HM.empty
 
-> listenerPort :: ClientSession -> PortNumber
-> listenerPort = servPort . peerListener
+> closeClientSession :: ClientSession -> IO ()
+> closeClientSession ClientSession {..} =
+>     maybeStop (tryTakeMVar peerListener) `finally`
+>     maybeStop (tryTakeMVar nodeListener)
+>   where
+>     maybeStop m = maybe (return ()) stopService =<< m
+
+> withClientSession :: SessionCount -> [Extension] -> (ClientSession -> IO ()) -> IO ()
+> withClientSession c es = bracket (openClientSession c es) closeClientSession
 
-> dhtPort :: ClientSession -> PortNumber
-> dhtPort = servPort . nodeListener
+> listenerPort :: ClientSession -> IO PortNumber
+> listenerPort ClientSession {..} = servPort <$> readMVar peerListener
+
+> dhtPort :: ClientSession -> IO PortNumber
+> dhtPort ClientSession {..} = servPort <$> readMVar nodeListener
 
 Swarm sessions
 ------------------------------------------------------------------------
@@ -736,8 +752,10 @@ TODO: use STM semaphore
 > sendClientStatus (sock, PeerSession {..}) = do
 >   cbf <- readTVarIO $ clientBitfield $ swarmSession
 >   sendAll sock $ encode $ Bitfield cbf
+>
+>   port <- dhtPort $ clientSession swarmSession
 >   when (ExtDHT `elem` enabledExtensions) $ do
->     sendAll sock $ encode $ Port $ dhtPort $ clientSession swarmSession
+>     sendAll sock $ encode $ Port port
 
 Exchange action depends on session and socket, whereas session depends
 on socket:
@@ -786,6 +804,26 @@ Used the a peer want to connect to the client.
 >       sendClientStatus (sock, ps)
 >       return ps
 
+
+> listener :: ClientSession -> Exchange -> PortNumber -> IO ()
+> listener cs action serverPort = bracket openListener close loop
+>   where
+>    loop sock = forever $ handle isIOError $ do
+>      (conn, addr) <- accept sock
+>      case addr of
+>        SockAddrInet port host -> do
+>          acceptPeerSession cs (PeerAddr Nothing host port) conn action
+>        _                      -> return ()
+>
+>    isIOError :: IOError -> IO ()
+>    isIOError _ = return ()
+>
+>    openListener  = do
+>      sock <- socket AF_INET Stream defaultProtocol
+>      bindSocket sock (SockAddrInet serverPort 0)
+>      listen sock 1
+>      return sock
+
 Broadcasting: Have, Cancel, Bitfield, SuggestPiece
 ------------------------------------------------------------------------