2 files changed, 35 insertions, 31 deletions

diff --git a/src/Network/BitTorrent/Tracker.hs b/src/Network/BitTorrent/Tracker.hs
index e4bcf2cd..2319a551 100644
--- a/src/Network/BitTorrent/Tracker.hs
+++ b/src/Network/BitTorrent/Tracker.hs
@@ -8,6 +8,7 @@
 --   This module provides high level API for peer->tracker
 --   communication.
 --
+{-# LANGUAGE RecordWildCards #-}
 module Network.BitTorrent.Tracker
        ( module Network.BitTorrent.Tracker.Scrape
 
@@ -155,28 +156,26 @@ withTracker :: Progress -> TConnection -> (TSession -> IO a) -> IO a
 withTracker initProgress conn action = bracket start end (action . fst)
   where
     start = do
-      res <- sendRequest (startedReq conn initProgress)
-      case res of
-        Left err -> ioError (userError err)
-        Right (Failure err) -> ioError (userError (show err))
-        Right resp -> do
-          se  <- newSession initProgress (respInterval resp) (respPeers resp)
-          tid <- forkIO (syncSession se)
-          return (se, tid)
-
-
-    syncSession se = do
-      waitInterval se
-      pr    <- getProgress se
-      eresp <- sendRequest (regularReq defaultNumWant conn pr)
-      case eresp of
-        Right (OK { respInterval = i, respPeers = ps }) -> do
-          writeIORef (seInterval se) i
-          atomically $ writeTVar (sePeers se) ps
-        _             -> return ()
-      syncSession se
-
+      resp <- askTracker (startedReq conn initProgress)
+      se   <- newSession 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 (regularReq defaultNumWant conn pr)
+        case resp of
+          Right (OK {..}) -> do
+            writeIORef seInterval respInterval
+            atomically $ writeTVar sePeers respPeers
+          _ -> return ()
+      where
+        isIOException :: IOException -> Maybe IOException
+        isIOException = return
 
     end (se, tid) = do
       killThread tid
-      getProgress se >>= sendRequest . stoppedReq conn
+      pr <- getProgress se
+      askTracker $ stoppedReq conn pr
diff --git a/src/Network/BitTorrent/Tracker/Protocol.hs b/src/Network/BitTorrent/Tracker/Protocol.hs
index 0371a187..af48e3e9 100644
--- a/src/Network/BitTorrent/Tracker/Protocol.hs
+++ b/src/Network/BitTorrent/Tracker/Protocol.hs
@@ -24,7 +24,7 @@ module Network.BitTorrent.Tracker.Protocol
        ( module Network.BitTorrent.Tracker.Scrape
 
        , Event(..), TRequest(..), TResponse(..)
-       , sendRequest
+       , askTracker
 
          -- * Defaults
        , defaultPorts, defaultNumWant
@@ -220,16 +220,21 @@ defaultPorts = [6881..6889]
 defaultNumWant :: Int
 defaultNumWant = 25
 
--- | TODO rename to ask for peers
+-- | 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.
 --
-sendRequest :: TRequest -> IO (Result TResponse)
-sendRequest req = do
-  let r = mkHTTPRequest (encodeRequest req)
-
-  rawResp  <- simpleHTTP r
-  respBody <- getResponseBody rawResp
-  return (decoded respBody)
+askTracker :: TRequest -> IO TResponse
+askTracker req = do
+    let r = mkHTTPRequest (encodeRequest req)
 
+    rawResp  <- simpleHTTP r
+    respBody <- getResponseBody rawResp
+    checkResult $ decoded respBody
   where
     mkHTTPRequest :: URI -> Request ByteString
     mkHTTPRequest uri = Request uri GET [] ""
+
+    checkResult (Left err)            = ioError (userError err)
+    checkResult (Right (Failure err)) = ioError (userError (show err))
+    checkResult (Right resp)          = return resp