src/FuncTorrent/Tracker/Udp.hs

   1 {-
   2  - Copyright (C) 2015-2016 Ramakrishnan Muthukrishnan <ram@rkrishnan.org>
   3  -
   4  - This file is part of FuncTorrent.
   5  -
   6  - FuncTorrent is free software; you can redistribute it and/or modify
   7  - it under the terms of the GNU General Public License as published by
   8  - the Free Software Foundation; either version 3 of the License, or
   9  - (at your option) any later version.
  10  -
  11  - FuncTorrent is distributed in the hope that it will be useful,
  12  - but WITHOUT ANY WARRANTY; without even the implied warranty of
  13  - MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
  14  - GNU General Public License for more details.
  15  -
  16  - You should have received a copy of the GNU General Public License
  17  - along with FuncTorrent; if not,  see <http://www.gnu.org/licenses/>
  18  -}
  19
  20 {-# LANGUAGE OverloadedStrings #-}
  21 module FuncTorrent.Tracker.Udp
  22        (trackerLoop
  23        ) where
  24
  25 import Control.Applicative (liftA2)
  26 import Control.Monad (liftM, forever, void)
  27 import Control.Concurrent (threadDelay)
  28 import Control.Concurrent.MVar (readMVar, putMVar, isEmptyMVar, swapMVar)
  29 import Control.Monad.Reader (ReaderT, runReaderT, ask, liftIO)
  30 import Data.Binary (Binary(..), encode, decode)
  31 import Data.Binary.Get (Get, isEmpty, getWord32be, getWord64be, getByteString)
  32 import Data.Binary.Put (putWord16be, putWord64be, putWord32be, putByteString)
  33 import Data.ByteString (ByteString)
  34 import qualified Data.ByteString.Char8 as BC
  35 import Data.ByteString.Lazy (fromStrict, toStrict)
  36 import Data.Word (Word16, Word32, Word64)
  37 import Network.Socket (Socket, Family( AF_INET ), SocketType( Datagram ), defaultProtocol, SockAddr(..), socket, close, getAddrInfo, addrAddress, SockAddr(..))
  38 import Network.Socket.ByteString (sendTo, recvFrom)
  39 import System.Random (randomIO)
  40 import System.Timeout (timeout)
  41
  42 import FuncTorrent.PeerMsgs (Peer(..))
  43 import FuncTorrent.Tracker.Types (TrackerEventState(..), TState(..))
  44 import FuncTorrent.Utils (IP, Port, toIP, toPort, getHostname, getPort)
  45 import qualified FuncTorrent.FileSystem as FS (MsgChannel, Stats(..), getStats)
  46
  47 -- UDP tracker: http://bittorrent.org/beps/bep_0015.html
  48 data Action = Connect
  49             | Announce
  50             | Scrape
  51             deriving (Show, Eq)
  52
  53 data UDPRequest = ConnectReq Word32
  54                 | AnnounceReq Word64 Word32 ByteString String Word64 Word64 Word64 TrackerEventState Word16
  55                 | ScrapeReq Integer Integer ByteString
  56                 deriving (Show, Eq)
  57
  58 data UDPResponse = ConnectResp Word32 Word64 -- transaction_id connection_id
  59                  | AnnounceResp Word32 Word32 Word32 Word32 [Peer] -- transaction_id interval leechers seeders [(ip, port)]
  60                  | ScrapeResp Integer Integer Integer Integer
  61                  | ErrorResp Integer String
  62                  deriving (Show, Eq)
  63
  64 data UDPTrackerHandle = UDPTrackerHandle { sock :: Socket
  65                                          , addr :: SockAddr
  66                                          , tid  :: Word32
  67                                          }
  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 event)
 100     putWord32be 0
 101     putWord32be 0
 102     putWord32be 10
 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 <$> getWord64be)
 113       1 -> do
 114         tid <- fromIntegral <$> getWord32be
 115         interval' <- fromIntegral <$> getWord32be
 116         l <- getWord32be -- leechers
 117         s <- getWord32be -- seeders
 118         ipportpairs <- getIPPortPairs -- [(ip, port)]
 119         return $ AnnounceResp tid interval' l s ipportpairs
 120       2 -> do
 121         tid <- fromIntegral <$> getWord32be
 122         _ <- getWord32be
 123         _ <- getWord32be
 124         _ <- getWord32be
 125         return $ ScrapeResp tid 0 0 0
 126       3 -> do -- error response
 127         tid <- fromIntegral <$> getWord32be
 128         bs  <- getByteString 4
 129         return $ ErrorResp tid $ BC.unpack bs
 130       _ -> error ("unknown response action type: " ++ show a)
 131
 132 sendRequest :: UDPTrackerHandle -> ByteString -> IO ()
 133 sendRequest h req = do
 134   n <- sendTo (sock h) req (addr h)
 135   -- sanity check with n?
 136   return ()
 137
 138 recvResponse :: UDPTrackerHandle -> IO UDPResponse
 139 recvResponse h = do
 140   (bs, saddr) <- recvFrom (sock h) (16*1024)
 141   return $ decode $ fromStrict bs
 142
 143 connectRequest :: ReaderT UDPTrackerHandle IO Word32
 144 connectRequest = do
 145   h <- ask
 146   tidi <- liftIO randomIO
 147   let pkt = encode $ ConnectReq tidi
 148   liftIO $ sendRequest h (toStrict pkt)
 149   return tidi
 150
 151 connectResponse :: Word32 -> ReaderT UDPTrackerHandle IO Word64
 152 connectResponse tid = do
 153   h <- ask
 154   resp <- liftIO $ recvResponse h
 155   -- check if nbytes is at least 16 bytes long
 156   case resp of
 157     (ConnectResp tidr cid) ->
 158       if tidr == tid
 159       then do
 160         liftIO $ putStrLn "connect succeeded"
 161         return cid
 162       else
 163         return 0
 164     _                      -> return 0
 165
 166 announceRequest :: Word64 -> ByteString -> String -> Word64 -> Word64 -> Word64 -> Word16 -> ReaderT UDPTrackerHandle IO Word32
 167 announceRequest cid infohash peerId up down left port = do
 168   h <- ask
 169   tidi <- liftIO randomIO
 170   let pkt = encode $ AnnounceReq cid tidi infohash peerId down left up None port
 171   liftIO $ sendRequest h (toStrict pkt)
 172   return tidi
 173
 174 data PeerStats = PeerStats { leechers :: Word32
 175                            , seeders :: Word32
 176                            , interval :: Word32
 177                            , peers :: [Peer]
 178                            } deriving (Show)
 179
 180 announceResponse :: Word32 -> ReaderT UDPTrackerHandle IO PeerStats
 181 announceResponse tid = do
 182   h <- ask
 183   resp <- liftIO $ recvResponse h
 184   case resp of
 185     (AnnounceResp tidr interval ss ls xs) ->
 186       if tidr == tid
 187       then do
 188         liftIO $ putStrLn "announce succeeded"
 189         return $ PeerStats ls ss interval xs
 190       else
 191         return $ PeerStats 0 0 0 []
 192     _ -> return $ PeerStats 0 0 0 []
 193
 194 getIPPortPairs :: Get [Peer]
 195 getIPPortPairs = do
 196   empty <- isEmpty
 197   if empty
 198     then return []
 199     else do
 200     ip <- toIP <$> getByteString 4
 201     port <- toPort <$> getByteString 2
 202     ipportpairs <- getIPPortPairs
 203     return $ (Peer ip port) : ipportpairs
 204
 205 startSession :: String -> Port -> IO UDPTrackerHandle
 206 startSession host port = do
 207   s <- socket AF_INET Datagram defaultProtocol
 208   addrinfos <- getAddrInfo Nothing (Just host) (Just (show port))
 209   let (SockAddrInet p ip) = addrAddress $ head addrinfos
 210   putStrLn "connected to tracker"
 211   return UDPTrackerHandle { sock = s
 212                           , addr = (SockAddrInet (fromIntegral port) ip) }
 213
 214 closeSession :: UDPTrackerHandle -> IO ()
 215 closeSession (UDPTrackerHandle s _ _) = close s
 216
 217 trackerLoop :: String -> Port -> String -> ByteString -> FS.MsgChannel -> TState -> IO ()
 218 trackerLoop url sport peerId infohash fschan tstate = forever $ do
 219   st <- readMVar <$> FS.getStats fschan
 220   up <- fmap FS.bytesRead st
 221   down <- fmap FS.bytesWritten st
 222   handle <- startSession host port
 223   stats <- timeout (15*(10^6)) $ worker handle up down
 224   case stats of
 225     Nothing -> closeSession handle
 226     Just stats' -> do
 227       ps <- isEmptyMVar $ connectedPeers tstate
 228       if ps
 229         then
 230         putMVar (connectedPeers tstate) (peers stats')
 231         else
 232         void $ swapMVar (connectedPeers tstate) (peers stats')
 233       threadDelay $ fromIntegral (interval stats') * (10^6)
 234       return ()
 235   where
 236     port = getPort url
 237     host = getHostname url
 238     worker handle up down = flip runReaderT handle $ do
 239       t1 <- connectRequest
 240       cid <- connectResponse t1
 241       t2 <- announceRequest cid infohash peerId (fromIntegral up) (fromIntegral down) (fromIntegral (left tstate)) (fromIntegral sport)
 242       stats <- announceResponse t2
 243       return stats