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 20
 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   print $ BC.length req
 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   liftIO $ print resp
 154   -- check if nbytes is at least 16 bytes long
 155   case resp of
 156     (ConnectResp tidr cid) ->
 157       if tidr == tid
 158       then do
 159         liftIO $ putStrLn "connect succeeded"
 160         return cid
 161       else
 162         return 0
 163     _                      -> return 0
 164
 165 announceRequest :: Word64 -> ByteString -> String -> Word64 -> Word64 -> Word64 -> Word16 -> ReaderT UDPTrackerHandle IO Word32
 166 announceRequest cid infohash peerId up down left port = do
 167   h <- ask
 168   tidi <- liftIO randomIO
 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 "connected: connect id"
 229     liftIO $ print cid
 230     t2 <- announceRequest cid infohash peerId (fromIntegral up) (fromIntegral down) (fromIntegral (left tstate)) (fromIntegral sport)
 231     liftIO $ print "announce request"
 232     liftIO $ print t2
 233     liftIO $ print "waiting for announce response"
 234     stats <- announceResponse t2
 235     liftIO $ print stats