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