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 (PortNumber)
  38 import Network.Socket (Socket, Family( AF_INET ), SocketType( Datagram ), defaultProtocol, SockAddr(..), socket, close, getAddrInfo, addrAddress, SockAddr(..))
  39 import Network.Socket.ByteString (sendTo, recvFrom)
  40 import System.Random (randomIO)
  41 import System.Timeout (timeout)
  42
  43 import FuncTorrent.PeerMsgs (Peer(..))
  44 import FuncTorrent.Tracker.Types (TrackerEventState(..), TState(..), UdpTrackerResponse(..))
  45 import FuncTorrent.Utils (IP, Port, toIP, toPort, getHostname, getPort)
  46 import qualified FuncTorrent.FileSystem as FS (MsgChannel, Stats(..), getStats)
  47
  48 -- UDP tracker: http://bittorrent.org/beps/bep_0015.html
  49 data Action = Connect
  50             | Announce
  51             | Scrape
  52             deriving (Show, Eq)
  53
  54 data UDPRequest = ConnectReq Word32
  55                 | AnnounceReq Word64 Word32 ByteString String Word64 Word64 Word64 TrackerEventState Word16
  56                 | ScrapeReq Integer Integer ByteString
  57                 deriving (Show, Eq)
  58
  59 data UDPResponse = ConnectResp Word32 Word64 -- transaction_id connection_id
  60                  | AnnounceResp Word32 Word32 Word32 Word32 [Peer] -- transaction_id interval leechers seeders [(ip, port)]
  61                  | ScrapeResp Integer Integer Integer Integer
  62                  | ErrorResp Integer String
  63                  deriving (Show, Eq)
  64
  65 data UDPTrackerHandle = UDPTrackerHandle { sock :: Socket
  66                                          , addr :: SockAddr
  67                                          , tid  :: Word32
  68                                          }
  69
  70 actionToInteger :: Action -> Integer
  71 actionToInteger Connect  = 0
  72 actionToInteger Announce = 1
  73 actionToInteger Scrape   = 2
  74
  75 intToAction :: Integer -> Action
  76 intToAction 0 = Connect
  77 intToAction 1 = Announce
  78 intToAction 2 = Scrape
  79
  80 eventToInteger :: TrackerEventState -> Integer
  81 eventToInteger None = 0
  82 eventToInteger Completed = 1
  83 eventToInteger Started = 2
  84 eventToInteger Stopped = 3
  85
  86 instance Binary UDPRequest where
  87   put (ConnectReq transId) = do
  88     putWord64be 0x41727101980
  89     putWord32be $ fromIntegral (actionToInteger Connect)
  90     putWord32be (fromIntegral transId)
  91   put (AnnounceReq connId transId infohash peerId down left up event port) = do
  92     putWord64be $ fromIntegral connId
  93     putWord32be $ fromIntegral (actionToInteger Announce)
  94     putWord32be $ fromIntegral transId
  95     putByteString infohash
  96     putByteString (BC.pack peerId)
  97     putWord64be (fromIntegral down)
  98     putWord64be (fromIntegral left)
  99     putWord64be (fromIntegral up)
 100     putWord32be $ fromIntegral (eventToInteger event)
 101     putWord32be 0
 102     putWord32be 0
 103     putWord32be 10
 104     putWord16be $ fromIntegral port
 105   put (ScrapeReq _ _ _) = undefined
 106   get = undefined
 107
 108 instance Binary UDPResponse where
 109   put = undefined
 110   get = do
 111     a <- getWord32be -- action
 112     case a of
 113       0 -> liftA2 ConnectResp (fromIntegral <$> getWord32be) (fromIntegral <$> getWord64be)
 114       1 -> do
 115         tid <- fromIntegral <$> getWord32be
 116         interval' <- fromIntegral <$> getWord32be
 117         l <- getWord32be -- leechers
 118         s <- getWord32be -- seeders
 119         ipportpairs <- getIPPortPairs -- [(ip, port)]
 120         return $ AnnounceResp tid interval' l s ipportpairs
 121       2 -> do
 122         tid <- fromIntegral <$> getWord32be
 123         _ <- getWord32be
 124         _ <- getWord32be
 125         _ <- getWord32be
 126         return $ ScrapeResp tid 0 0 0
 127       3 -> do -- error response
 128         tid <- fromIntegral <$> getWord32be
 129         bs  <- getByteString 4
 130         return $ ErrorResp tid $ BC.unpack bs
 131       _ -> error ("unknown response action type: " ++ show a)
 132
 133 sendRequest :: UDPTrackerHandle -> ByteString -> IO ()
 134 sendRequest h req = do
 135   n <- sendTo (sock h) req (addr h)
 136   -- sanity check with n?
 137   return ()
 138
 139 recvResponse :: UDPTrackerHandle -> IO UDPResponse
 140 recvResponse h = do
 141   (bs, saddr) <- recvFrom (sock h) (16*1024)
 142   return $ decode $ fromStrict bs
 143
 144 connectRequest :: ReaderT UDPTrackerHandle IO Word32
 145 connectRequest = do
 146   h <- ask
 147   tidi <- liftIO randomIO
 148   let pkt = encode $ ConnectReq tidi
 149   liftIO $ sendRequest h (toStrict pkt)
 150   return tidi
 151
 152 connectResponse :: Word32 -> ReaderT UDPTrackerHandle IO Word64
 153 connectResponse tid = do
 154   h <- ask
 155   resp <- liftIO $ recvResponse h
 156   -- check if nbytes is at least 16 bytes long
 157   case resp of
 158     (ConnectResp tidr cid) ->
 159       if tidr == tid
 160       then do
 161         liftIO $ putStrLn "connect succeeded"
 162         return cid
 163       else
 164         return 0
 165     _                      -> return 0
 166
 167 announceRequest :: Word64 -> ByteString -> String -> Word64 -> Word64 -> Word64 -> Word16 -> ReaderT UDPTrackerHandle IO Word32
 168 announceRequest cid infohash peerId up down left port = do
 169   h <- ask
 170   tidi <- liftIO randomIO
 171   let pkt = encode $ AnnounceReq cid tidi infohash peerId down left up None port
 172   liftIO $ sendRequest h (toStrict pkt)
 173   return tidi
 174
 175 announceResponse :: Word32 -> ReaderT UDPTrackerHandle IO UdpTrackerResponse
 176 announceResponse tid = do
 177   h <- ask
 178   resp <- liftIO $ recvResponse h
 179   case resp of
 180     (AnnounceResp tidr interval ss ls xs) ->
 181       if tidr == tid
 182       then do
 183         liftIO $ putStrLn "announce succeeded"
 184         return $ UdpTrackerResponse ls ss interval xs
 185       else
 186         return $ UdpTrackerResponse 0 0 0 []
 187     _ -> return $ UdpTrackerResponse 0 0 0 []
 188
 189 getIPPortPairs :: Get [Peer]
 190 getIPPortPairs = do
 191   empty <- isEmpty
 192   if empty
 193     then return []
 194     else do
 195     ip <- toIP <$> getByteString 4
 196     port <- toPort <$> getByteString 2
 197     ipportpairs <- getIPPortPairs
 198     return $ (Peer ip port) : ipportpairs
 199
 200 startSession :: String -> Port -> IO UDPTrackerHandle
 201 startSession host port = do
 202   s <- socket AF_INET Datagram defaultProtocol
 203   addrinfos <- getAddrInfo Nothing (Just host) (Just (show port))
 204   let (SockAddrInet p ip) = addrAddress $ head addrinfos
 205   putStrLn "connected to tracker"
 206   return UDPTrackerHandle { sock = s
 207                           , addr = (SockAddrInet (fromIntegral port) ip) }
 208
 209 closeSession :: UDPTrackerHandle -> IO ()
 210 closeSession (UDPTrackerHandle s _ _) = close s
 211
 212 trackerLoop :: String -> PortNumber -> String -> ByteString -> FS.MsgChannel -> TState -> IO ()
 213 trackerLoop url sport peerId infohash fschan tstate = forever $ do
 214   st <- readMVar <$> FS.getStats fschan
 215   up <- fmap FS.bytesRead st
 216   down <- fmap FS.bytesWritten st
 217   handle <- startSession host port
 218   stats <- timeout (15*(10^6)) $ worker handle up down
 219   case stats of
 220     Nothing -> closeSession handle
 221     Just stats' -> do
 222       ps <- isEmptyMVar $ connectedPeers tstate
 223       if ps
 224         then
 225         putMVar (connectedPeers tstate) (peers stats')
 226         else
 227         void $ swapMVar (connectedPeers tstate) (peers stats')
 228       threadDelay $ fromIntegral (interval stats') * (10^6)
 229       return ()
 230   where
 231     port = getPort url
 232     host = getHostname url
 233     worker handle up down = flip runReaderT handle $ do
 234       t1 <- connectRequest
 235       cid <- connectResponse t1
 236       t2 <- announceRequest cid infohash peerId (fromIntegral up) (fromIntegral down) (fromIntegral (left tstate)) (fromIntegral sport)
 237       stats <- announceResponse t2
 238       return stats