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, getWord64be, 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 event)
  97     putWord32be 0
  98     putWord32be 0
  99     putWord32be 10
 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 <$> getWord64be)
 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   -- check if nbytes is at least 16 bytes long
 153   case resp of
 154     (ConnectResp tidr cid) ->
 155       if tidr == tid
 156       then do
 157         liftIO $ putStrLn "connect succeeded"
 158         return cid
 159       else
 160         return 0
 161     _                      -> return 0
 162
 163 announceRequest :: Word64 -> ByteString -> String -> Word64 -> Word64 -> Word64 -> Word16 -> ReaderT UDPTrackerHandle IO Word32
 164 announceRequest cid infohash peerId up down left port = do
 165   h <- ask
 166   tidi <- liftIO randomIO
 167   let pkt = encode $ AnnounceReq cid tidi infohash peerId down left up None port
 168   liftIO $ sendRequest h (toStrict pkt)
 169   return tidi
 170
 171 data PeerStats = PeerStats { leechers :: Word32
 172                            , seeders :: Word32
 173                            , peers :: [(IP, Port)]
 174                            } deriving (Show)
 175
 176 announceResponse :: Word32 -> ReaderT UDPTrackerHandle IO PeerStats
 177 announceResponse tid = do
 178   h <- ask
 179   resp <- liftIO $ recvResponse h
 180   case resp of
 181     (AnnounceResp tidr interval ss ls xs) ->
 182       if tidr == tid
 183       then do
 184         liftIO $ putStrLn "announce succeeded"
 185         return $ PeerStats ls ss xs
 186       else
 187         return $ PeerStats 0 0 []
 188     _ -> return $ PeerStats 0 0 []
 189
 190 getIPPortPairs :: Get [(IP, Port)]
 191 getIPPortPairs = do
 192   empty <- isEmpty
 193   if empty
 194     then return []
 195     else do
 196     ip <- toIP <$> getByteString 4
 197     port <- toPort <$> getByteString 2
 198     ipportpairs <- getIPPortPairs
 199     return $ (ip, port) : ipportpairs
 200
 201 startSession :: String -> Port -> IO UDPTrackerHandle
 202 startSession host port = do
 203   s <- socket AF_INET Datagram defaultProtocol
 204   addrinfos <- getAddrInfo Nothing (Just host) (Just (show port))
 205   let (SockAddrInet p ip) = addrAddress $ head addrinfos
 206   putStrLn "connected to tracker"
 207   return UDPTrackerHandle { sock = s
 208                           , addr = (SockAddrInet (fromIntegral port) ip) }
 209
 210 closeSession :: UDPTrackerHandle -> IO ()
 211 closeSession (UDPTrackerHandle s _ _) = close s
 212
 213 trackerLoop :: String -> Port -> String -> ByteString -> FS.MsgChannel -> TState -> IO ()
 214 trackerLoop url sport peerId infohash fschan tstate = do
 215   st' <- FS.getStats fschan
 216   st <- readMVar st'
 217   let up = FS.bytesRead st
 218       down = FS.bytesWritten st
 219       port = getPort url
 220       host = getHostname url
 221   putStrLn $ "host = " ++ (show host) ++ " port= " ++ (show port)
 222   handle <- startSession host port
 223   flip runReaderT handle $ do
 224     t1 <- connectRequest
 225     cid <- connectResponse t1
 226     t2 <- announceRequest cid infohash peerId (fromIntegral up) (fromIntegral down) (fromIntegral (left tstate)) (fromIntegral sport)
 227     stats <- announceResponse t2
 228     liftIO $ print stats
 229 --    _ <- threadDelay $