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     putWord32be 0
  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   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   -- connId transId infohash peerId down left up event port)
 170   let pkt = encode $ AnnounceReq cid tidi infohash peerId down left up None port
 171   liftIO $ sendRequest h (toStrict pkt)
 172   return tidi
 173
 174 data PeerStats = PeerStats { leechers :: Word32
 175                            , seeders :: Word32
 176                            , peers :: [(IP, Port)]
 177                            } deriving (Show)
 178
 179 announceResponse :: Word32 -> ReaderT UDPTrackerHandle IO PeerStats
 180 announceResponse tid = do
 181   h <- ask
 182   resp <- liftIO $ recvResponse h
 183   case resp of
 184     (AnnounceResp tidr interval ss ls xs) ->
 185       if tidr == tid
 186       then do
 187         liftIO $ putStrLn "announce succeeded"
 188         return $ PeerStats ls ss xs
 189       else
 190         return $ PeerStats 0 0 []
 191     _ -> return $ PeerStats 0 0 []
 192
 193 getIPPortPairs :: Get [(IP, Port)]
 194 getIPPortPairs = do
 195   empty <- isEmpty
 196   if empty
 197     then return []
 198     else do
 199     ip <- toIP <$> getByteString 6
 200     port <- toPort <$> getByteString 2
 201     ipportpairs <- getIPPortPairs
 202     return $ (ip, port) : ipportpairs
 203
 204 startSession :: String -> Port -> IO UDPTrackerHandle
 205 startSession host port = do
 206   s <- socket AF_INET Datagram defaultProtocol
 207   addrinfos <- getAddrInfo Nothing (Just host) (Just (show port))
 208   let (SockAddrInet p ip) = addrAddress $ head addrinfos
 209   putStrLn "connected to tracker"
 210   return UDPTrackerHandle { sock = s
 211                           , addr = (SockAddrInet (fromIntegral port) ip) }
 212
 213 closeSession :: UDPTrackerHandle -> IO ()
 214 closeSession (UDPTrackerHandle s _ _) = close s
 215
 216 trackerLoop :: String -> Port -> String -> ByteString -> FS.MsgChannel -> TState -> IO ()
 217 trackerLoop url sport 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       port = getPort url
 223       host = getHostname url
 224   putStrLn $ "host = " ++ (show host) ++ " port= " ++ (show port)
 225   handle <- startSession host port
 226   flip runReaderT handle $ do
 227     t1 <- connectRequest
 228     cid <- connectResponse t1
 229     liftIO $ print "connect response"
 230     liftIO $ print cid
 231     t2 <- announceRequest cid infohash peerId (fromIntegral up) (fromIntegral down) (fromIntegral (left tstate)) (fromIntegral sport)
 232     liftIO $ print "announce request"
 233     liftIO $ print t2
 234     stats <- announceResponse t2
 235     liftIO $ print stats