Add BEP12: Multitracker Protocol Extension

3 files changed, 203 insertions, 3 deletions

diff --git a/bittorrent.cabal b/bittorrent.cabal
index f8e598c1..4b2431ce 100644
--- a/bittorrent.cabal
+++ b/bittorrent.cabal
@@ -74,7 +74,7 @@ library
 --                     Network.BitTorrent.Exchange.Session
                        Network.BitTorrent.Exchange.Status
                        Network.BitTorrent.Exchange.Wire
---                     Network.BitTorrent.Tracker
+                       Network.BitTorrent.Tracker.List
                        Network.BitTorrent.Tracker.Message
                        Network.BitTorrent.Tracker.RPC
                        Network.BitTorrent.Tracker.RPC.HTTP
@@ -153,12 +153,15 @@ library
                      , wai          >= 1.4 && < 2.0
                      , iproute
 
+                       -- RNG/PRNG
+                     , entropy        >= 0.2
+                     , random         >= 1.0.0.2
+                     , random-shuffle >= 0.0.0.4
+
                        -- System
                      , directory   >= 1.2
-                     , entropy     >= 0.2
                      , filepath    >= 1.3
                      , mmap        >= 0.5
-                     , random      >= 1.0.0.2
 
   ghc-options:         -Wall
   ghc-prof-options:
diff --git a/src/Network/BitTorrent/Tracker/List.hs b/src/Network/BitTorrent/Tracker/List.hs
new file mode 100644
index 00000000..b6371f1d
--- /dev/null
+++ b/src/Network/BitTorrent/Tracker/List.hs
@@ -0,0 +1,153 @@
+-- |
+--   Copyright   :  (c) Sam Truzjan 2014
+--   License     :  BSD
+--   Maintainer  :  pxqr.sta@gmail.com
+--   Stability   :  experimental
+--   Portability :  portable
+--
+--   Multitracker Metadata Extension support.
+--
+--   For more info see: <http://www.bittorrent.org/beps/bep_0012.html>
+--
+module Network.BitTorrent.Tracker.List
+       ( -- * Tracker list
+         TrackerList
+
+         -- * Construction
+       , trackerList
+       , shuffleTiers
+
+         -- * Traversals
+       , traverseAll
+       , traverseTiers
+       ) where
+
+import Prelude hiding (mapM, foldr)
+import Control.Applicative
+import Control.Exception
+import Data.List as L (elem, any, filter, null)
+import Data.Foldable
+import Data.Traversable
+import Network.URI
+import System.Random.Shuffle
+
+import Data.Torrent
+import Network.BitTorrent.Tracker.RPC as RPC
+
+{-----------------------------------------------------------------------
+--  Tracker list datatype
+-----------------------------------------------------------------------}
+
+type Tier a = [a]
+
+-- | Tracker list is either a single tracker or list of tiers. All
+-- trackers in each tier must be checked before the client goes on to
+-- the next tier.
+data TrackerList a
+  = Announce a           -- ^ torrent file 'announce'      field only
+  | TierList [Tier a]    -- ^ torrent file 'announce-list' field only
+    deriving (Show, Eq)
+
+instance Functor TrackerList where
+  fmap f (Announce a) = Announce (f a)
+  fmap f (TierList a) = TierList (fmap (fmap f) a)
+
+instance Foldable TrackerList where
+  foldr f z (Announce a ) = f a z
+  foldr f z (TierList xs) = foldr (flip (foldr f)) z xs
+
+instance Traversable TrackerList where
+  traverse f (Announce a ) = Announce <$> f a
+  traverse f (TierList xs) = TierList <$> traverse (traverse f) xs
+
+{-----------------------------------------------------------------------
+--  List extraction
+-----------------------------------------------------------------------}
+-- BEP12 do not expose any restrictions for the content of
+-- 'announce-list' key - there are some /bad/ cases can happen with
+-- poorly designed or even malicious torrent creation software.
+--
+-- Bad case #1: announce-list is present, but empty.
+--
+--   { tAnnounce     = Just "http://a.com"
+--   , tAnnounceList = Just [[]]
+--   }
+--
+-- Bad case #2: announce uri do not present in announce list.
+--
+--   { tAnnounce     = Just "http://a.com"
+--   , tAnnounceList = Just [["udp://a.com"]]
+--   }
+--
+--  The addBackup function solves both problems by adding announce uri
+--  as backup tier.
+--
+addBackup :: [[URI]] -> URI -> [[URI]]
+addBackup tiers bkp
+  | L.any (L.elem bkp) tiers = tiers
+  |         otherwise        = tiers ++ [[bkp]]
+
+fixList :: Maybe [[URI]] -> Maybe URI -> Maybe [[URI]]
+fixList mxss mx = do
+  xss <- mxss
+  let xss' = L.filter (not . L.null) xss
+  return $ maybe xss' (addBackup xss') mx
+
+-- | Extract set of trackers from torrent file. The 'tAnnounce' key is
+-- only ignored if the 'tAnnounceList' key is present.
+trackerList :: Torrent -> Maybe (TrackerList URI)
+trackerList Torrent {..} = TierList <$> (tAnnounceList `fixList` tAnnounce)
+                       <|> Announce <$>  tAnnounce
+
+-- | Shuffle /order of trackers/ in each tier, preserving original
+-- /order of tiers/. This can help to balance the load between the
+-- trackers.
+shuffleTiers :: TrackerList a -> IO (TrackerList a)
+shuffleTiers (Announce a ) = return (Announce a)
+shuffleTiers (TierList xs) = TierList <$> mapM shuffleM xs
+
+{-----------------------------------------------------------------------
+--  Special traversals (suppressed RPC exceptions)
+-----------------------------------------------------------------------}
+
+catchRPC :: IO a -> IO a -> IO a
+catchRPC a b = catch a (f b)
+  where
+    f :: a -> RpcException -> a
+    f = const
+
+throwRPC :: String -> IO a
+throwRPC = throwIO . GenericException
+
+-- | Like 'traverse' but ignore 'RpcExceptions'.
+traverseAll :: (a -> IO a) -> TrackerList a -> IO (TrackerList a)
+traverseAll action = traverse (action $?)
+  where
+    f $? x = catchRPC (f x) (return x)
+
+-- | Like 'traverse' but put working trackers to the head of tiers.
+-- This can help to avoid exceessive requests to not available
+-- trackers at each reannounce.
+traverseTiers :: (a -> IO a) -> TrackerList a -> IO (TrackerList a)
+traverseTiers action (Announce a)     = Announce <$> action a
+traverseTiers action (TierList tiers) = TierList <$> goTiers (goTier []) tiers
+  where
+    goTiers _ []       = throwRPC "traverseTiers: no tiers"
+    goTiers f (x : xs) = catchRPC shortcut failback
+      where
+        shortcut = do
+          x' <- f x
+          return (x' : xs)
+
+        failback = do
+          xs' <- goTiers f xs
+          return (x : xs')
+
+    goTier _      []       = throwRPC "traverseTiers: no trackers in tier"
+    goTier failed (a : as) = catchRPC shortcut failback
+      where
+        shortcut = do
+          a' <- action a
+          return (a' : as ++ failed) -- failed trackers at the end
+
+        failback = goTier (a : failed) as
diff --git a/tests/Network/BitTorrent/Tracker/ListSpec.hs b/tests/Network/BitTorrent/Tracker/ListSpec.hs
new file mode 100644
index 00000000..8decd3c9
--- /dev/null
+++ b/tests/Network/BitTorrent/Tracker/ListSpec.hs
@@ -0,0 +1,44 @@
+module Network.BitTorrent.Tracker.ListSpec (spec) where
+import Control.Exception
+import Data.Default
+import Data.Foldable as F
+import Data.List as L
+import Data.Maybe
+import Network.URI
+import Test.Hspec
+
+import Data.Torrent
+import Network.BitTorrent.Tracker.List
+import Network.BitTorrent.Tracker.RPC
+
+
+uris :: [URI]
+uris = fmap (fromJust . parseURI . renderURI) [1..10 :: Int]
+  where
+    renderURI n = "http://" ++ show n ++ ".org"
+
+list :: TrackerList URI
+list = fromJust $ trackerList def { tAnnounceList = Just [uris] }
+
+spec :: Spec
+spec = do
+  describe "TrackerList" $ do
+    it "trackerList is not empty" $ do
+      pending
+
+    it "shuffleTiers (may fail with very small probability)" $ do
+      list' <- shuffleTiers list
+      list' `shouldSatisfy` (/= list)
+
+    it "traverseAll" $ do
+      xs <- traverseAll (\ uri -> if uri == L.last uris
+               then throwIO (GenericException "")
+               else return uri { uriScheme = "udp://" }) list
+      let udps = F.sum $ fmap (fromEnum . ("udp://" ==) . uriScheme) xs
+      udps `shouldBe` pred (L.length uris)
+
+    it "traverseTiers" $ do
+      xs' <- traverseTiers (\ uri -> if uri == L.last uris then return uri
+                              else throwIO (GenericException "")) list
+
+      return ()