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