src/FuncTorrent/Tracker.hs

   1 {-# LANGUAGE OverloadedStrings #-}
   2 module FuncTorrent.Tracker
   3     (TState(..),
   4      initialTrackerState,
   5      trackerLoop,
   6     ) where
   7
   8 import Prelude hiding (lookup, splitAt)
   9
  10 import Control.Applicative (liftA2)
  11 import Control.Concurrent (threadDelay)
  12 import Control.Concurrent.MVar (MVar, newEmptyMVar, newMVar, readMVar, putMVar)
  13 import Data.Binary (Binary(..))
  14 import Data.Binary.Put (putWord16be, putWord32be, putWord64be, putByteString)
  15 import Data.Binary.Get (Get, isEmpty, getByteString, getWord16be, getWord32be)
  16 import Data.ByteString (ByteString)
  17 import Data.ByteString.Char8 as BC (pack, unpack, splitAt)
  18 import Data.Char (chr)
  19 import Data.List (intercalate)
  20 import Data.Map as M (lookup)
  21 import Network (PortNumber)
  22 import Network.HTTP.Base (urlEncode)
  23 import qualified Data.ByteString.Base16 as B16 (encode)
  24
  25 import FuncTorrent.Bencode (BVal(..), decode)
  26 import FuncTorrent.Metainfo (Info(..), Metainfo(..))
  27 import FuncTorrent.Network (sendGetRequest)
  28 import FuncTorrent.Peer (Peer(..))
  29 import FuncTorrent.Utils (splitN)
  30
  31 -- | Tracker response
  32 data TrackerResponse = TrackerResponse {
  33   interval :: Integer
  34   , peers :: [Peer]
  35   , complete :: Maybe Integer
  36   , incomplete :: Maybe Integer
  37   } deriving (Show, Eq)
  38
  39 data TrackerEventState = None
  40                        | Started
  41                        | Stopped
  42                        | Completed
  43                        deriving (Show, Eq)
  44
  45 data TState = TState {
  46     uploaded :: MVar Integer
  47   , downloaded :: MVar Integer
  48   , left :: Integer
  49   , currentState :: TrackerEventState
  50   , connectedPeers :: MVar [Peer]
  51   }
  52
  53 -- UDP tracker: http://bittorrent.org/beps/bep_0015.html
  54 data Action = Connect
  55             | Announce
  56             | Scrape
  57             deriving (Show, Eq)
  58
  59 type IP = String
  60 type Port = Integer
  61
  62 data UDPRequest = ConnectReq Integer
  63                 | AnnounceReq Integer Integer ByteString String Integer Integer Integer TrackerEventState Integer
  64                 | ScrapeReq Integer Integer ByteString
  65                 deriving (Show, Eq)
  66
  67 data UDPResponse = ConnectResp Integer Integer -- transaction_id connection_id
  68                  | AnnounceResp Integer Integer Integer Integer [(IP, Port)] -- transaction_id interval leechers seeders [(ip, port)]
  69                  | ScrapeResp Integer Integer Integer Integer
  70                  deriving (Show, Eq)
  71
  72 actionToInteger :: Action -> Integer
  73 actionToInteger Connect  = 0
  74 actionToInteger Announce = 1
  75 actionToInteger Scrape   = 2
  76
  77 intToAction :: Integer -> Action
  78 intToAction 0 = Connect
  79 intToAction 1 = Announce
  80 intToAction 2 = Scrape
  81
  82 eventToInteger :: TrackerEventState -> Integer
  83 eventToInteger None = 0
  84 eventToInteger Completed = 1
  85 eventToInteger Started = 2
  86 eventToInteger Stopped = 3
  87
  88 instance Binary UDPRequest where
  89   put (ConnectReq transId) = do
  90     putWord64be 0x41727101980
  91     putWord32be $ fromIntegral (actionToInteger Connect)
  92     putWord32be (fromIntegral transId)
  93   put (AnnounceReq connId transId infohash peerId down left up event port) = do
  94     putWord64be $ fromIntegral connId
  95     putWord32be $ fromIntegral (actionToInteger Announce)
  96     putWord32be $ fromIntegral transId
  97     putByteString infohash
  98     putByteString (BC.pack peerId)
  99     putWord64be (fromIntegral down)
 100     putWord64be (fromIntegral left)
 101     putWord64be (fromIntegral up)
 102     putWord32be $ fromIntegral (eventToInteger None)
 103     putWord32be 0
 104     -- key is optional, we will not send it for now
 105     putWord32be $ fromIntegral (-1)
 106     putWord16be $ fromIntegral port
 107   put (ScrapeReq _ _ _) = undefined
 108   get = undefined
 109
 110 instance Binary UDPResponse where
 111   put = undefined
 112   get = do
 113     a <- getWord32be -- action
 114     case a of
 115       0 -> liftA2 ConnectResp (fromIntegral <$> getWord32be) (fromIntegral <$> getWord32be)
 116       1 -> do
 117         tid <- fromIntegral <$> getWord32be
 118         interval' <- fromIntegral <$> getWord32be
 119         _ <- getWord32be -- leechers
 120         _ <- getWord32be -- seeders
 121         ipportpairs <- getIPPortPairs -- [(ip, port)]
 122         return $ AnnounceResp tid interval' 0 0 ipportpairs
 123       2 -> do
 124         tid <- fromIntegral <$> getWord32be
 125         _ <- getWord32be
 126         _ <- getWord32be
 127         _ <- getWord32be
 128         return $ ScrapeResp tid 0 0 0
 129       _ -> error ("unknown response action type: " ++ show a)
 130
 131 getIPPortPairs :: Get [(IP, Port)]
 132 getIPPortPairs = do
 133   empty <- isEmpty
 134   if empty
 135     then return []
 136     else do
 137     ip <- toIP <$> getByteString 6
 138     port <- toPort <$> getByteString 2
 139     ipportpairs <- getIPPortPairs
 140     return $ (ip, port) : ipportpairs
 141
 142 initialTrackerState :: Integer -> IO TState
 143 initialTrackerState sz = do
 144   ps <- newEmptyMVar
 145   up <- newMVar 0
 146   down <- newMVar 0
 147   return $ TState { currentState = None
 148                   , connectedPeers = ps
 149                   , uploaded = up
 150                   , downloaded = down
 151                   , left = sz }
 152
 153 -- | Deserialize tracker response
 154 mkTrackerResponse :: BVal -> Either ByteString TrackerResponse
 155 mkTrackerResponse resp =
 156     case lookup "failure reason" body of
 157       Just (Bstr err) -> Left err
 158       Just _ -> Left "Unknown failure"
 159       Nothing ->
 160           let (Just (Bint i)) = lookup "interval" body
 161               (Just (Bstr peersBS)) = lookup "peers" body
 162               pl = map makePeer (splitN 6 peersBS)
 163           in Right TrackerResponse {
 164                    interval = i
 165                  , peers = pl
 166                  , complete = Nothing
 167                  , incomplete = Nothing
 168                  }
 169     where
 170       (Bdict body) = resp
 171
 172 toInt :: String -> Integer
 173 toInt = read
 174
 175 makePeer :: ByteString -> Peer
 176 makePeer peer = Peer "" (toIP ip') (toPort port')
 177   where (ip', port') = splitAt 4 peer
 178
 179 toPort :: ByteString -> Port
 180 toPort = read . ("0x" ++) . unpack . B16.encode
 181
 182 toIP :: ByteString -> IP
 183 toIP = Data.List.intercalate "." .
 184        map (show . toInt . ("0x" ++) . unpack) .
 185        splitN 2 . B16.encode
 186
 187 --- | URL encode hash as per RFC1738
 188 --- TODO: Add tests
 189 --- REVIEW: Why is this not written in terms of `Network.HTTP.Base.urlEncode` or
 190 --- equivalent library function?
 191 urlEncodeHash :: ByteString -> String
 192 urlEncodeHash bs = concatMap (encode' . unpack) (splitN 2 bs)
 193   where encode' b@[c1, c2] = let c =  chr (read ("0x" ++ b))
 194                             in escape c c1 c2
 195         encode' _ = ""
 196         escape i c1 c2 | i `elem` nonSpecialChars = [i]
 197                        | otherwise = "%" ++ [c1] ++ [c2]
 198
 199         nonSpecialChars = ['A'..'Z'] ++ ['a'..'z'] ++ ['0'..'9'] ++ "-_.~"
 200
 201 -- | Make arguments that should be posted to tracker.
 202 -- This is a separate pure function for testability.
 203 mkArgs :: PortNumber -> String -> Integer -> Integer -> Metainfo -> [(String, ByteString)]
 204 mkArgs port peer_id up down m =
 205   let fileSize = lengthInBytes $ info m
 206       bytesLeft = fileSize - down
 207   in
 208     [("info_hash", pack . urlEncodeHash . B16.encode . infoHash $ m),
 209      ("peer_id", pack . urlEncode $ peer_id),
 210      ("port", pack $ show port),
 211      ("uploaded", pack $ show up),
 212      ("downloaded", pack $ show down),
 213      ("left", pack $ show bytesLeft),
 214      ("compact", "1"),
 215      ("event", "started")]
 216
 217 trackerLoop :: PortNumber -> String -> Metainfo -> TState -> IO ByteString
 218 trackerLoop port peerId m st = do
 219   up <- readMVar $ uploaded st
 220   down <- readMVar $ downloaded st
 221   resp <- sendGetRequest (head . announceList $ m) $ mkArgs port peerId up down m
 222   case decode resp of
 223     Left e -> return $ pack (show e)
 224     Right trackerInfo ->
 225       case mkTrackerResponse trackerInfo of
 226         Left e -> return e
 227         Right tresp -> do
 228           _ <- threadDelay $ fromIntegral (interval tresp)
 229           _ <- putMVar (connectedPeers st) (peers tresp)
 230           trackerLoop port peerId m st
 231