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 (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 data UDPRequest = ConnectReq Integer
  60                 | AnnounceReq Integer Integer ByteString String Integer Integer Integer TrackerEventState Integer
  61                 | ScrapeReq Integer Integer ByteString
  62                 deriving (Show, Eq)
  63
  64 data UDPResponse = ConnectResp Integer Integer
  65                  | AnnounceResp Integer Integer Integer Integer Integer Integer
  66                  | ScrapeResp Integer Integer Integer Integer
  67                  deriving (Show, Eq)
  68
  69 actionToInteger :: Action -> Integer
  70 actionToInteger Connect  = 0
  71 actionToInteger Announce = 1
  72 actionToInteger Scrape   = 2
  73
  74 intToAction :: Integer -> Action
  75 intToAction 0 = Connect
  76 intToAction 1 = Announce
  77 intToAction 2 = Scrape
  78
  79 eventToInteger :: TrackerEventState -> Integer
  80 eventToInteger None = 0
  81 eventToInteger Completed = 1
  82 eventToInteger Started = 2
  83 eventToInteger Stopped = 3
  84
  85 instance Binary UDPRequest where
  86   put (ConnectReq transId) = do
  87     putWord64be 0x41727101980
  88     putWord32be $ fromIntegral (actionToInteger Connect)
  89     putWord32be (fromIntegral transId)
  90   put (AnnounceReq connId transId infohash peerId down left up event port) = do
  91     putWord64be $ fromIntegral connId
  92     putWord32be $ fromIntegral (actionToInteger Announce)
  93     putWord32be $ fromIntegral transId
  94     putByteString infohash
  95     putByteString (BC.pack peerId)
  96     putWord64be (fromIntegral down)
  97     putWord64be (fromIntegral left)
  98     putWord64be (fromIntegral up)
  99     putWord32be $ fromIntegral (eventToInteger None)
 100     putWord32be 0
 101     -- key is optional, we will not send it for now
 102     putWord32be $ fromIntegral (-1)
 103     putWord16be $ fromIntegral port
 104   put (ScrapeReq _ _ _) = undefined
 105   get = undefined
 106
 107 instance Binary UDPResponse where
 108   put = undefined
 109   get = do
 110     a <- getWord32be -- action
 111     case a of
 112       0 -> liftA2 ConnectResp (fromIntegral <$> getWord32be) (fromIntegral <$> getWord32be)
 113       1 -> do
 114         tid <- fromIntegral <$> getWord32be
 115         interval' <- fromIntegral <$> getWord32be
 116         _ <- getWord32be -- leechers
 117         _ <- getWord32be -- seeders
 118         _ <- getWord32be -- ip
 119         _ <- getWord16be -- port
 120         return $ AnnounceResp tid interval' 0 0 0 0
 121       2 -> do
 122         tid <- fromIntegral <$> getWord32be
 123         _ <- getWord32be
 124         _ <- getWord32be
 125         _ <- getWord32be
 126         return $ ScrapeResp tid 0 0 0
 127       _ -> error ("unknown response action type: " ++ show a)
 128
 129 initialTrackerState :: Integer -> IO TState
 130 initialTrackerState sz = do
 131   ps <- newEmptyMVar
 132   up <- newMVar 0
 133   down <- newMVar 0
 134   return $ TState { currentState = None
 135                   , connectedPeers = ps
 136                   , uploaded = up
 137                   , downloaded = down
 138                   , left = sz }
 139
 140 -- | Deserialize tracker response
 141 mkTrackerResponse :: BVal -> Either ByteString TrackerResponse
 142 mkTrackerResponse resp =
 143     case lookup "failure reason" body of
 144       Just (Bstr err) -> Left err
 145       Just _ -> Left "Unknown failure"
 146       Nothing ->
 147           let (Just (Bint i)) = lookup "interval" body
 148               (Just (Bstr peersBS)) = lookup "peers" body
 149               pl = map makePeer (splitN 6 peersBS)
 150           in Right TrackerResponse {
 151                    interval = i
 152                  , peers = pl
 153                  , complete = Nothing
 154                  , incomplete = Nothing
 155                  }
 156     where
 157       (Bdict body) = resp
 158
 159       toInt :: String -> Integer
 160       toInt = read
 161
 162       toPort :: ByteString -> Integer
 163       toPort = read . ("0x" ++) . unpack . B16.encode
 164
 165       toIP :: ByteString -> String
 166       toIP = Data.List.intercalate "." .
 167              map (show . toInt . ("0x" ++) . unpack) .
 168                  splitN 2 . B16.encode
 169
 170       makePeer :: ByteString -> Peer
 171       makePeer peer = Peer "" (toIP ip') (toPort port')
 172           where (ip', port') = splitAt 4 peer
 173
 174 --- | URL encode hash as per RFC1738
 175 --- TODO: Add tests
 176 --- REVIEW: Why is this not written in terms of `Network.HTTP.Base.urlEncode` or
 177 --- equivalent library function?
 178 urlEncodeHash :: ByteString -> String
 179 urlEncodeHash bs = concatMap (encode' . unpack) (splitN 2 bs)
 180   where encode' b@[c1, c2] = let c =  chr (read ("0x" ++ b))
 181                             in escape c c1 c2
 182         encode' _ = ""
 183         escape i c1 c2 | i `elem` nonSpecialChars = [i]
 184                        | otherwise = "%" ++ [c1] ++ [c2]
 185
 186         nonSpecialChars = ['A'..'Z'] ++ ['a'..'z'] ++ ['0'..'9'] ++ "-_.~"
 187
 188 -- | Make arguments that should be posted to tracker.
 189 -- This is a separate pure function for testability.
 190 mkArgs :: PortNumber -> String -> Integer -> Integer -> Metainfo -> [(String, ByteString)]
 191 mkArgs port peer_id up down m =
 192   let fileSize = lengthInBytes $ info m
 193       bytesLeft = fileSize - down
 194   in
 195     [("info_hash", pack . urlEncodeHash . B16.encode . infoHash $ m),
 196      ("peer_id", pack . urlEncode $ peer_id),
 197      ("port", pack $ show port),
 198      ("uploaded", pack $ show up),
 199      ("downloaded", pack $ show down),
 200      ("left", pack $ show bytesLeft),
 201      ("compact", "1"),
 202      ("event", "started")]
 203
 204 trackerLoop :: PortNumber -> String -> Metainfo -> TState -> IO ByteString
 205 trackerLoop port peerId m st = do
 206   up <- readMVar $ uploaded st
 207   down <- readMVar $ downloaded st
 208   resp <- sendGetRequest (head . announceList $ m) $ mkArgs port peerId up down m
 209   case decode resp of
 210     Left e -> return $ pack (show e)
 211     Right trackerInfo ->
 212       case mkTrackerResponse trackerInfo of
 213         Left e -> return e
 214         Right tresp -> do
 215           _ <- threadDelay $ fromIntegral (interval tresp)
 216           _ <- putMVar (connectedPeers st) (peers tresp)
 217           trackerLoop port peerId m st
 218