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(..), UdpTrackerResponse(..))
  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 announceResponse :: Word32 -> ReaderT UDPTrackerHandle IO UdpTrackerResponse
 175 announceResponse tid = do
 176   h <- ask
 177   resp <- liftIO $ recvResponse h
 178   case resp of
 179     (AnnounceResp tidr interval ss ls xs) ->
 180       if tidr == tid
 181       then do
 182         liftIO $ putStrLn "announce succeeded"
 183         return $ UdpTrackerResponse ls ss interval xs
 184       else
 185         return $ UdpTrackerResponse 0 0 0 []
 186     _ -> return $ UdpTrackerResponse 0 0 0 []
 187
 188 getIPPortPairs :: Get [Peer]
 189 getIPPortPairs = do
 190   empty <- isEmpty
 191   if empty
 192     then return []
 193     else do
 194     ip <- toIP <$> getByteString 4
 195     port <- toPort <$> getByteString 2
 196     ipportpairs <- getIPPortPairs
 197     return $ (Peer ip port) : ipportpairs
 198
 199 startSession :: String -> Port -> IO UDPTrackerHandle
 200 startSession host port = do
 201   s <- socket AF_INET Datagram defaultProtocol
 202   addrinfos <- getAddrInfo Nothing (Just host) (Just (show port))
 203   let (SockAddrInet p ip) = addrAddress $ head addrinfos
 204   putStrLn "connected to tracker"
 205   return UDPTrackerHandle { sock = s
 206                           , addr = (SockAddrInet (fromIntegral port) ip) }
 207
 208 closeSession :: UDPTrackerHandle -> IO ()
 209 closeSession (UDPTrackerHandle s _ _) = close s
 210
 211 trackerLoop :: String -> Port -> String -> ByteString -> FS.MsgChannel -> TState -> IO ()
 212 trackerLoop url sport peerId infohash fschan tstate = forever $ do
 213   st <- readMVar <$> FS.getStats fschan
 214   up <- fmap FS.bytesRead st
 215   down <- fmap FS.bytesWritten st
 216   handle <- startSession host port
 217   stats <- timeout (15*(10^6)) $ worker handle up down
 218   case stats of
 219     Nothing -> closeSession handle
 220     Just stats' -> do
 221       ps <- isEmptyMVar $ connectedPeers tstate
 222       if ps
 223         then
 224         putMVar (connectedPeers tstate) (peers stats')
 225         else
 226         void $ swapMVar (connectedPeers tstate) (peers stats')
 227       threadDelay $ fromIntegral (interval stats') * (10^6)
 228       return ()
 229   where
 230     port = getPort url
 231     host = getHostname url
 232     worker handle up down = flip runReaderT handle $ do
 233       t1 <- connectRequest
 234       cid <- connectResponse t1
 235       t2 <- announceRequest cid infohash peerId (fromIntegral up) (fromIntegral down) (fromIntegral (left tstate)) (fromIntegral sport)
 236       stats <- announceResponse t2
 237       return stats