4 files changed, 97 insertions, 103 deletions

diff --git a/src/Network/BitTorrent/Tracker.hs b/src/Network/BitTorrent/Tracker.hs
index c707cedd..0501f428 100644
--- a/src/Network/BitTorrent/Tracker.hs
+++ b/src/Network/BitTorrent/Tracker.hs
@@ -41,7 +41,7 @@ import Network.URI
 
 import Data.Torrent.Metainfo
 import Network.BitTorrent.Peer
-import Network.BitTorrent.Tracker.Protocol
+import Network.BitTorrent.Tracker.Protocol as Tracker
 import Network.BitTorrent.Tracker.HTTP
 
 {-----------------------------------------------------------------------
@@ -83,9 +83,8 @@ genericReq ses pr =   AnnounceQuery {
 --   'startedReq'. It includes necessary 'Started' event field.
 --
 startedReq :: TConnection -> Progress -> AnnounceQuery
-startedReq ses pr = (genericReq ses pr) {
-    reqIP         = Nothing
-  , reqNumWant    = Just defaultNumWant
+startedReq ses pr = (genericReq ses pr)
+  { reqNumWant    = Just defaultNumWant
   , reqEvent      = Just Started
   }
 
@@ -94,9 +93,8 @@ startedReq ses pr = (genericReq ses pr) {
 --   so new peers could connect to the client.
 --
 regularReq :: Int -> TConnection -> Progress -> AnnounceQuery
-regularReq numWant ses pr = (genericReq ses pr) {
-    reqIP         = Nothing
-  , reqNumWant    = Just numWant
+regularReq numWant ses pr = (genericReq ses pr)
+  { reqNumWant    = Just numWant
   , reqEvent      = Nothing
   }
 
@@ -104,9 +102,8 @@ regularReq numWant ses pr = (genericReq ses pr) {
 -- gracefully.
 --
 stoppedReq :: TConnection -> Progress -> AnnounceQuery
-stoppedReq ses pr = (genericReq ses pr) {
-    reqIP         = Nothing
-  , reqNumWant    = Nothing
+stoppedReq ses pr = (genericReq ses pr)
+  { reqNumWant    = Nothing
   , reqEvent      = Just Stopped
   }
 
@@ -115,9 +112,8 @@ stoppedReq ses pr = (genericReq ses pr) {
 -- complete.
 --
 completedReq :: TConnection -> Progress -> AnnounceQuery
-completedReq ses pr = (genericReq ses pr) {
-    reqIP         = Nothing
-  , reqNumWant    = Nothing
+completedReq ses pr = (genericReq ses pr)
+  { reqNumWant    = Nothing
   , reqEvent      = Just Completed
   }
 
@@ -153,6 +149,7 @@ data TSession = TSession {
     seProgress   :: TVar Progress
   , seInterval   :: IORef TimeInterval
   , sePeers      :: BoundedChan PeerAddr
+  , seTracker    :: HTTPTracker
   }
 
 type PeerCount = Int
@@ -167,8 +164,9 @@ getProgress :: TSession -> IO Progress
 getProgress = readTVarIO . seProgress
 
 newSession :: PeerCount -> Progress -> TimeInterval -> [PeerAddr]
+           -> HTTPTracker
            -> IO TSession
-newSession chanSize pr i ps
+newSession chanSize pr i ps tr
   | chanSize < 1
   = throwIO $ userError "size of chan should be more that 1"
 
@@ -183,6 +181,7 @@ newSession chanSize pr i ps
     TSession <$> newTVarIO pr
              <*> newIORef i
              <*> pure chan
+             <*> pure tr
 
 waitInterval :: TSession -> IO ()
 waitInterval TSession {..} = do
@@ -191,39 +190,45 @@ waitInterval TSession {..} = do
   where
     sec = 1000 * 1000 :: Int
 
+announceLoop :: IO (BoundedChan PeerAddr)
+announceLoop = undefined
+
+openSession :: Progress -> TConnection -> IO TSession
+openSession initProgress conn = do
+  t    <- Tracker.connect (tconnAnnounce conn)
+  resp <- Tracker.announce t (startedReq conn initProgress)
+  newSession defaultChanSize initProgress
+          (respInterval resp) (respPeers resp) t
+
+closeSession :: TConnection -> TSession -> IO ()
+closeSession conn se @ TSession {..} = do
+  pr <- getProgress se
+  Tracker.announce seTracker (stoppedReq conn pr)
+  return ()
+
+syncSession :: TConnection -> TSession -> IO ()
+syncSession conn se @ TSession {..} = forever $ do
+  waitInterval se
+  pr   <- getProgress se
+  resp <- tryJust isIOException $ do
+   Tracker.announce seTracker (regularReq defaultNumWant conn pr)
+  case resp of
+    Left _ -> return ()
+    Right (AnnounceInfo {..}) -> do
+      writeIORef seInterval respInterval
+
+      -- we rely on the fact that union on lists is not
+      -- commutative: this implements the heuristic "old peers
+      -- in head"
+      old <- BC.getChanContents sePeers
+      let combined = L.union old respPeers
+      BC.writeList2Chan sePeers combined
+ where
+    isIOException :: IOException -> Maybe IOException
+    isIOException = return
+
 withTracker :: Progress -> TConnection -> (TSession -> IO a) -> IO a
-withTracker initProgress conn action = bracket start end (action . fst)
-  where
-    start = do
-      resp <- askTracker (tconnAnnounce conn) (startedReq conn initProgress)
-      se   <- newSession defaultChanSize initProgress
-                         (respInterval resp) (respPeers resp)
-
-      tid  <- forkIO (syncSession se)
-      return (se, tid)
-
-    syncSession se @ TSession {..} = forever $ do
-        waitInterval se
-        pr   <- getProgress se
-        resp <- tryJust isIOException $ do
-                    askTracker (tconnAnnounce conn) (regularReq defaultNumWant conn pr)
-        case resp of
-          Right (AnnounceInfo {..}) -> do
-            writeIORef seInterval respInterval
-
-            -- we rely on the fact that union on lists is not
-            -- commutative: this implements the heuristic "old peers
-            -- in head"
-            old <- BC.getChanContents sePeers
-            let combined = L.union old respPeers
-            BC.writeList2Chan sePeers combined
-
-          _ -> return ()
-      where
-        isIOException :: IOException -> Maybe IOException
-        isIOException = return
-
-    end (se, tid) = do
-      killThread tid
-      pr <- getProgress se
-      leaveTracker (tconnAnnounce conn) (stoppedReq conn pr)
+withTracker initProgress conn
+    = bracket
+      (openSession initProgress conn)
+      (closeSession conn)
diff --git a/src/Network/BitTorrent/Tracker/HTTP.hs b/src/Network/BitTorrent/Tracker/HTTP.hs
index 8d3a6412..ce517b34 100644
--- a/src/Network/BitTorrent/Tracker/HTTP.hs
+++ b/src/Network/BitTorrent/Tracker/HTTP.hs
@@ -14,17 +14,17 @@
 --   <https://wiki.theory.org/BitTorrentSpecification#Tracker_HTTP.2FHTTPS_Protocol>
 --
 module Network.BitTorrent.Tracker.HTTP
-       ( askTracker, leaveTracker
-       , scrapeURL, scrape, scrapeOne
+       ( HTTPTracker
+
+         -- * Extra
+       , scrapeURL
        ) where
 
-import Control.Applicative
-import Control.Monad
+import Control.Exception
 import Data.BEncode
 import Data.ByteString as B
 import Data.ByteString.Char8 as BC
 import Data.List as L
-import Data.Map as M
 import Data.Monoid
 import Data.URLEncoded as URL
 import Network.URI
@@ -34,37 +34,34 @@ import Data.Torrent.Metainfo hiding (announce)
 import Network.BitTorrent.Tracker.Protocol
 
 
-data HTTPTracker = HTTPTracker URI
+data HTTPTracker = HTTPTracker
+  { announceURI :: URI
+  } deriving Show
 
-instance Tracker URI where
-  announce = askTracker
-  scrape_  uri ihs = do
-    e <- scrape uri ihs
-    case e of
-      Left str  -> error str
-      Right si  -> return si
+instance Tracker HTTPTracker where
+  connect  = return . HTTPTracker
+  announce = announceHTTP
+  scrape   = scrapeHTTP
 
 {-----------------------------------------------------------------------
   Announce
 -----------------------------------------------------------------------}
 
 encodeRequest :: URI -> AnnounceQuery -> URI
-encodeRequest announce req = URL.urlEncode req
-                    `addToURI`      announce
+encodeRequest announceURI req = URL.urlEncode req
+                    `addToURI`      announceURI
                     `addHashToURI`  reqInfoHash req
 
 mkGET :: URI -> Request ByteString
 mkGET uri = Request uri GET [] ""
 
--- TODO rename to something like "announceBlahBlah"
-
 -- | Send request and receive response from the tracker specified in
 -- announce list. This function throws 'IOException' if it couldn't
 -- send request or receive response or decode response.
 --
-askTracker :: URI -> AnnounceQuery -> IO AnnounceInfo
-askTracker announce req = do
-    let r = mkGET (encodeRequest announce req)
+announceHTTP :: HTTPTracker -> AnnounceQuery -> IO AnnounceInfo
+announceHTTP HTTPTracker {..} req = do
+    let r = mkGET (encodeRequest announceURI req)
 
     rawResp  <- simpleHTTP r
     respBody <- getResponseBody rawResp
@@ -76,13 +73,6 @@ askTracker announce req = do
       = ioError $ userError $ show err ++ " in tracker response"
     checkResult (Right resp)          = return resp
 
--- | The same as the 'askTracker' but ignore response. Used in
--- conjunction with 'Stopped'.
-leaveTracker :: URI -> AnnounceQuery -> IO ()
-leaveTracker announce req = do
-  let r = mkGET (encodeRequest announce req)
-  void $ simpleHTTP r >>= getResponseBody
-
 {-----------------------------------------------------------------------
   Scrape
 -----------------------------------------------------------------------}
@@ -113,25 +103,15 @@ scrapeURL uri ihs = do
 --   all available torrents.
 --   Note that the 'URI' should be /announce/ URI, not /scrape/ URI.
 --
-scrape :: URI                -- ^ Announce 'URI'.
-       -> [InfoHash]         -- ^ Torrents to be scrapped.
-       -> IO (Result Scrape) -- ^ 'ScrapeInfo' for each torrent.
-scrape announce ihs
-  | Just uri<- scrapeURL announce ihs = do
+scrapeHTTP :: HTTPTracker -- ^ Announce 'URI'.
+           -> [InfoHash]  -- ^ Torrents to be scrapped.
+           -> IO Scrape   -- ^ 'ScrapeInfo' for each torrent.
+scrapeHTTP HTTPTracker {..} ihs
+  | Just uri <- scrapeURL announceURI ihs = do
     rawResp  <- simpleHTTP (Request uri GET [] "")
     respBody <- getResponseBody rawResp
-    return (decoded (BC.pack respBody))
+    case decoded (BC.pack respBody) of
+      Left  e -> throwIO $ userError $ e ++ " in scrape response"
+      Right r -> return r
 
-  | otherwise = return (Left "Tracker do not support scraping")
-
--- | More particular version of 'scrape', just for one torrent.
---
-scrapeOne :: URI                     -- ^ Announce 'URI'
-          -> InfoHash                -- ^ Hash of the torrent info.
-          -> IO (Result ScrapeInfo)  -- ^ 'ScrapeInfo' for the torrent.
-scrapeOne uri ih = extract <$> scrape uri [ih]
-  where
-    extract (Right m)
-      | Just s <- M.lookup ih m = Right s
-      | otherwise = Left "unable to find info hash in response dict"
-    extract (Left e) = Left e
+  | otherwise = throwIO $ userError "Tracker do not support scraping"
diff --git a/src/Network/BitTorrent/Tracker/Protocol.hs b/src/Network/BitTorrent/Tracker/Protocol.hs
index 3f264aed..965f3480 100644
--- a/src/Network/BitTorrent/Tracker/Protocol.hs
+++ b/src/Network/BitTorrent/Tracker/Protocol.hs
@@ -39,12 +39,13 @@ module Network.BitTorrent.Tracker.Protocol
 
          -- * TODO
        , Tracker(..)
+       , scrapeOne
        )
        where
 
 import Control.Applicative
+import Control.Exception
 import Control.Monad
-
 import Data.Aeson.TH
 import Data.Char as Char
 import Data.Map  as M
@@ -58,12 +59,11 @@ import Data.Text.Encoding
 import Data.Serialize hiding (Result)
 import Data.URLEncoded as URL
 import Data.Torrent.Metainfo
-
 import Network
+import Network.URI
 import Network.Socket
 
 import Network.BitTorrent.Peer
-import Network.BitTorrent.Sessions.Types
 
 {-----------------------------------------------------------------------
   Announce messages
@@ -365,5 +365,13 @@ instance Serialize ScrapeInfo where
 
 -- | Set of tracker RPCs.
 class Tracker s where
+  connect  :: URI -> IO s
   announce :: s -> AnnounceQuery -> IO AnnounceInfo
-  scrape_  :: s -> ScrapeQuery   -> IO Scrape
+  scrape   :: s -> ScrapeQuery   -> IO Scrape
+
+-- | More particular version of 'scrape', just for one torrent.
+--
+scrapeOne :: Tracker t => t -> InfoHash -> IO ScrapeInfo
+scrapeOne uri ih = scrape uri [ih] >>= maybe err return . M.lookup ih
+  where
+    err = throwIO $ userError "unable to find info hash in response dict"
diff --git a/src/Network/BitTorrent/Tracker/UDP.hs b/src/Network/BitTorrent/Tracker/UDP.hs
index e5475a23..dc1b4897 100644
--- a/src/Network/BitTorrent/Tracker/UDP.hs
+++ b/src/Network/BitTorrent/Tracker/UDP.hs
@@ -15,7 +15,6 @@
 {-# LANGUAGE TypeFamilies               #-}
 module Network.BitTorrent.Tracker.UDP
        ( UDPTracker
-       , initialTracker
 
          -- * Debug
        , putTracker
@@ -316,8 +315,8 @@ scrapeUDP tracker scr = do
   freshConnection tracker
   resp <- transaction tracker (Scrape scr)
   case resp of
-    Scraped scrape -> return $ M.fromList $ L.zip scr scrape
-    _              -> fail "scrape: response type mismatch"
+    Scraped info -> return $ M.fromList $ L.zip scr info
+    _            -> fail "scrape: response type mismatch"
 
 {-----------------------------------------------------------------------
   Retransmission
@@ -342,6 +341,8 @@ retransmission action = go minTimeout
         maybe (go (2 * curTimeout)) return r
 
 {----------------------------------------------------------------------}
+
 instance Tracker UDPTracker where
+  connect    = initialTracker
   announce t = retransmission . announceUDP t
-  scrape_  t = retransmission . scrapeUDP   t
+  scrape   t = retransmission . scrapeUDP   t