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