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, close, getAddrInfo, addrAddress, SockAddr(..))
  37 import Network.Socket.ByteString (sendTo, recvFrom)
  38 import System.Random (randomIO)
  39
  40 import FuncTorrent.Tracker.Types (TrackerEventState(..), TState(..))
  41 import FuncTorrent.Utils (IP, Port, toIP, toPort, getHostname, getPort)
  42 import qualified FuncTorrent.FileSystem as FS (MsgChannel, Stats(..), getStats)
  43
  44 -- UDP tracker: http://bittorrent.org/beps/bep_0015.html
  45 data Action = Connect
  46             | Announce
  47             | Scrape
  48             deriving (Show, Eq)
  49
  50 data UDPRequest = ConnectReq Word32
  51                 | AnnounceReq Word64 Word32 ByteString String Word64 Word64 Word64 TrackerEventState Word16
  52                 | ScrapeReq Integer Integer ByteString
  53                 deriving (Show, Eq)
  54
  55 data UDPResponse = ConnectResp Word32 Word64 -- transaction_id connection_id
  56                  | AnnounceResp Word32 Word32 Word32 Word32 [(IP, Port)] -- transaction_id interval leechers seeders [(ip, port)]
  57                  | ScrapeResp Integer Integer Integer Integer
  58                  | ErrorResp Integer String
  59                  deriving (Show, Eq)
  60
  61 data UDPTrackerHandle = UDPTrackerHandle { sock :: Socket
  62                                          , addr :: SockAddr
  63                                          , tid  :: Word32
  64                                          }
  65
  66 actionToInteger :: Action -> Integer
  67 actionToInteger Connect  = 0
  68 actionToInteger Announce = 1
  69 actionToInteger Scrape   = 2
  70
  71 intToAction :: Integer -> Action
  72 intToAction 0 = Connect
  73 intToAction 1 = Announce
  74 intToAction 2 = Scrape
  75
  76 eventToInteger :: TrackerEventState -> Integer
  77 eventToInteger None = 0
  78 eventToInteger Completed = 1
  79 eventToInteger Started = 2
  80 eventToInteger Stopped = 3
  81
  82 instance Binary UDPRequest where
  83   put (ConnectReq transId) = do
  84     putWord64be 0x41727101980
  85     putWord32be $ fromIntegral (actionToInteger Connect)
  86     putWord32be (fromIntegral transId)
  87   put (AnnounceReq connId transId infohash peerId down left up event port) = do
  88     putWord64be $ fromIntegral connId
  89     putWord32be $ fromIntegral (actionToInteger Announce)
  90     putWord32be $ fromIntegral transId
  91     putByteString infohash
  92     putByteString (BC.pack peerId)
  93     putWord64be (fromIntegral down)
  94     putWord64be (fromIntegral left)
  95     putWord64be (fromIntegral up)
  96     putWord32be $ fromIntegral (eventToInteger None)
  97     putWord32be 0
  98     -- key is optional, we will not send it for now
  99     putWord32be $ fromIntegral (-1)
 100     putWord16be $ fromIntegral port
 101   put (ScrapeReq _ _ _) = undefined
 102   get = undefined
 103
 104 instance Binary UDPResponse where
 105   put = undefined
 106   get = do
 107     a <- getWord32be -- action
 108     case a of
 109       0 -> liftA2 ConnectResp (fromIntegral <$> getWord32be) (fromIntegral <$> getWord32be)
 110       1 -> do
 111         tid <- fromIntegral <$> getWord32be
 112         interval' <- fromIntegral <$> getWord32be
 113         l <- getWord32be -- leechers
 114         s <- getWord32be -- seeders
 115         ipportpairs <- getIPPortPairs -- [(ip, port)]
 116         return $ AnnounceResp tid interval' l s ipportpairs
 117       2 -> do
 118         tid <- fromIntegral <$> getWord32be
 119         _ <- getWord32be
 120         _ <- getWord32be
 121         _ <- getWord32be
 122         return $ ScrapeResp tid 0 0 0
 123       3 -> do -- error response
 124         tid <- fromIntegral <$> getWord32be
 125         bs  <- getByteString 4
 126         return $ ErrorResp tid $ BC.unpack bs
 127       _ -> error ("unknown response action type: " ++ show a)
 128
 129 sendRequest :: UDPTrackerHandle -> ByteString -> IO ()
 130 sendRequest h req = do
 131   n <- sendTo (sock h) req (addr h)
 132   -- sanity check with n?
 133   return ()
 134
 135 recvResponse :: UDPTrackerHandle -> IO UDPResponse
 136 recvResponse h = do
 137   (bs, saddr) <- recvFrom (sock h) (16*1024)
 138   return $ decode $ fromStrict bs
 139
 140 connectRequest :: ReaderT UDPTrackerHandle IO Word32
 141 connectRequest = do
 142   h <- ask
 143   tidi <- liftIO randomIO
 144   let pkt = encode $ ConnectReq tidi
 145   liftIO $ sendRequest h (toStrict pkt)
 146   return tidi
 147
 148 connectResponse :: Word32 -> ReaderT UDPTrackerHandle IO Word64
 149 connectResponse tid = do
 150   h <- ask
 151   resp <- liftIO $ recvResponse h
 152   liftIO $ print resp
 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 -> String -> Word64 -> Word64 -> Word64 -> Word16 -> ReaderT UDPTrackerHandle IO Word32
 165 announceRequest cid infohash peerId 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 peerId 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   putStrLn "connected to tracker"
 209   return $ UDPTrackerHandle { sock = s
 210                             , addr = (SockAddrInet (fromIntegral port) ip) }
 211
 212 closeSession :: UDPTrackerHandle -> IO ()
 213 closeSession (UDPTrackerHandle s _ _) = close s
 214
 215 trackerLoop :: String -> Port -> String -> ByteString -> FS.MsgChannel -> TState -> IO ()
 216 trackerLoop url sport peerId infohash fschan tstate = do
 217   st' <- FS.getStats fschan
 218   st <- readMVar st'
 219   let up = FS.bytesRead st
 220       down = FS.bytesWritten st
 221       port = getPort url
 222       host = getHostname url
 223   putStrLn $ "host = " ++ (show host) ++ " port= " ++ (show port)
 224   handle <- startSession host port
 225   flip runReaderT handle $ do
 226     t1 <- connectRequest
 227     cid <- connectResponse t1
 228     liftIO $ print cid
 229     t2 <- announceRequest cid infohash peerId (fromIntegral up) (fromIntegral down) (fromIntegral (left tstate)) (fromIntegral sport)
 230     stats <- announceResponse t2
 231     liftIO $ print stats