From 95b19be7cd65d9263e7edb7ab63d3067fe4ec8c5 Mon Sep 17 00:00:00 2001
From: Ramakrishnan Muthukrishnan <ram@rkrishnan.org>
Date: Mon, 13 Jun 2016 20:17:32 +0530
Subject: [PATCH] UDP Tracker: connect + announce. Does not work

---
 src/FuncTorrent/Tracker.hs       | 47 ++++++++++++----
 src/FuncTorrent/Tracker/Types.hs |  2 +-
 src/FuncTorrent/Tracker/Udp.hs   | 95 +++++++++++++++++++++++++-------
 3 files changed, 113 insertions(+), 31 deletions(-)

diff --git a/src/FuncTorrent/Tracker.hs b/src/FuncTorrent/Tracker.hs
index 815d081..5e059f7 100644
--- a/src/FuncTorrent/Tracker.hs
+++ b/src/FuncTorrent/Tracker.hs
@@ -29,20 +29,49 @@ import Control.Concurrent.Chan (Chan, newChan, readChan, writeChan)
 import Control.Concurrent.MVar (newEmptyMVar, putMVar, readMVar)
 import Control.Monad.State (StateT, liftIO, get, runStateT)
 import Control.Monad (forever)
-import Data.ByteString.Char8 (ByteString)
+import Data.ByteString.Char8 (ByteString, pack, unpack)
 import Data.List (isPrefixOf)
 import Network (PortNumber)
 
-import FuncTorrent.Tracker.Http (trackerLoop)
+import qualified FuncTorrent.Tracker.Http as HT (trackerLoop)
+import qualified FuncTorrent.Tracker.Udp as UT (trackerLoop)
 import FuncTorrent.Tracker.Types (TState(..), TrackerEventState(..), TrackerProtocol(..), TrackerMsg(..))
+import FuncTorrent.Utils (Port, toPort)
 import qualified FuncTorrent.FileSystem as FS (MsgChannel)
 import FuncTorrent.Peer (Peer)
 
 type MsgChannel = Chan TrackerMsg
 
+data TrackerUrl = TrackerUrl { protocol :: TrackerProtocol
+                             , host :: String
+                             , port :: Port
+                             , path :: String
+                             }
+
 newTracker :: IO MsgChannel
 newTracker = newChan
 
+parseUrl :: String -> TrackerUrl
+parseUrl url = TrackerUrl proto host port path
+  where proto = getTrackerType url
+        host = getHostname url
+        port = getPort url
+        path = getPath url
+
+getTrackerType :: String -> TrackerProtocol
+getTrackerType url | isPrefixOf "http://" url = Http
+                   | isPrefixOf "udp://" url  = Udp
+                   | otherwise                = UnknownProtocol
+
+getHostname :: String -> String
+getHostname url = takeWhile (/= ':') $ drop 2 $ dropWhile (/= '/') url
+
+getPort :: String -> Port
+getPort url = toPort . pack $ takeWhile (/= '/') $ drop 1 $ dropWhile (/= ':') $ drop 2 $ dropWhile (/= '/') url
+
+getPath :: String -> String
+getPath url = dropWhile (/= '/') $ dropWhile (/= ':') $ drop 1 $ dropWhile (/= ':') url
+
 runTracker :: MsgChannel -> FS.MsgChannel -> ByteString -> PortNumber
            -> String -> [String] -> Integer -> IO ()
 runTracker msgChannel fsChan infohash port peerId announceList sz = do
@@ -51,20 +80,18 @@ runTracker msgChannel fsChan infohash port peerId announceList sz = do
                              , connectedPeers = ps
                              , left = sz }
       turl = head announceList
+      host = getHostname turl
   case getTrackerType turl of
     Http -> do
-      _ <- forkIO $ trackerLoop turl port peerId infohash fsChan initialTState
+      _ <- forkIO $ HT.trackerLoop host port peerId infohash fsChan initialTState
       runStateT (msgHandler msgChannel) initialTState
       return ()
-    _ -> do
+    Udp -> do
+      _ <- forkIO $ UT.trackerLoop host (fromIntegral port) peerId infohash fsChan initialTState
+      return ()
+    _ ->
       error "Tracker Protocol unimplemented"
 
-getTrackerType :: String -> TrackerProtocol
-getTrackerType url | isPrefixOf "http://" url = Http
-                   | isPrefixOf "udp://" url  = Udp
-                   | otherwise                = UnknownProtocol
-
-
 msgHandler :: MsgChannel -> StateT TState IO ()
 msgHandler c = forever $ do
   st <- get
diff --git a/src/FuncTorrent/Tracker/Types.hs b/src/FuncTorrent/Tracker/Types.hs
index c79fcef..2165b08 100644
--- a/src/FuncTorrent/Tracker/Types.hs
+++ b/src/FuncTorrent/Tracker/Types.hs
@@ -39,7 +39,7 @@ data TrackerProtocol = Http
 data TrackerEventState = None
                        | Started
                        | Completed
-                       | Error ByteString
+                       | Stopped
                        deriving (Show, Eq)
 
 data TrackerMsg = GetStatusMsg TrackerEventState
diff --git a/src/FuncTorrent/Tracker/Udp.hs b/src/FuncTorrent/Tracker/Udp.hs
index e24944a..4e541a9 100644
--- a/src/FuncTorrent/Tracker/Udp.hs
+++ b/src/FuncTorrent/Tracker/Udp.hs
@@ -19,10 +19,12 @@
 
 {-# LANGUAGE OverloadedStrings #-}
 module FuncTorrent.Tracker.Udp
-       (
+       (trackerLoop
        ) where
 
 import Control.Applicative (liftA2)
+import Control.Monad (liftM)
+import Control.Concurrent.MVar (readMVar)
 import Control.Monad.Reader (ReaderT, runReaderT, ask, liftIO)
 import Data.Binary (Binary(..), encode, decode)
 import Data.Binary.Get (Get, isEmpty, getWord32be, getByteString)
@@ -30,13 +32,15 @@ import Data.Binary.Put (putWord16be, putWord64be, putWord32be, putByteString)
 import Data.ByteString (ByteString)
 import qualified Data.ByteString.Char8 as BC
 import Data.ByteString.Lazy (fromStrict, toStrict)
-import Data.Word (Word32, Word64)
-import Network.Socket (Socket, Family( AF_INET ), SocketType( Datagram ), defaultProtocol, SockAddr(..), socket, inet_addr, close)
+import Data.Word (Word16, Word32, Word64)
+import Network.Socket (Socket, Family( AF_INET ), SocketType( Datagram ), defaultProtocol, SockAddr(..), socket, inet_addr, close, getAddrInfo, addrAddress, SockAddr(..))
 import Network.Socket.ByteString (sendTo, recvFrom)
 import System.Random (randomIO)
 
 import FuncTorrent.Tracker.Types (TrackerEventState(..))
 import FuncTorrent.Utils (IP, Port, toIP, toPort)
+import qualified FuncTorrent.FileSystem as FS (MsgChannel, Stats(..), getStats)
+import FuncTorrent.Tracker.Types(TState(..))
 
 -- UDP tracker: http://bittorrent.org/beps/bep_0015.html
 data Action = Connect
@@ -45,12 +49,12 @@ data Action = Connect
             deriving (Show, Eq)
 
 data UDPRequest = ConnectReq Word32
-                | AnnounceReq Integer Integer ByteString String Integer Integer Integer TrackerEventState Integer
+                | AnnounceReq Word64 Word32 ByteString String Word64 Word64 Word64 TrackerEventState Word16
                 | ScrapeReq Integer Integer ByteString
                 deriving (Show, Eq)
 
 data UDPResponse = ConnectResp Word32 Word64 -- transaction_id connection_id
-                 | AnnounceResp Integer Integer Integer Integer [(IP, Port)] -- transaction_id interval leechers seeders [(ip, port)]
+                 | AnnounceResp Word32 Word32 Word32 Word32 [(IP, Port)] -- transaction_id interval leechers seeders [(ip, port)]
                  | ScrapeResp Integer Integer Integer Integer
                  | ErrorResp Integer String
                  deriving (Show, Eq)
@@ -74,6 +78,7 @@ eventToInteger :: TrackerEventState -> Integer
 eventToInteger None = 0
 eventToInteger Completed = 1
 eventToInteger Started = 2
+eventToInteger Stopped = 3
 
 instance Binary UDPRequest where
   put (ConnectReq transId) = do
@@ -106,10 +111,10 @@ instance Binary UDPResponse where
       1 -> do
         tid <- fromIntegral <$> getWord32be
         interval' <- fromIntegral <$> getWord32be
-        _ <- getWord32be -- leechers
-        _ <- getWord32be -- seeders
+        l <- getWord32be -- leechers
+        s <- getWord32be -- seeders
         ipportpairs <- getIPPortPairs -- [(ip, port)]
-        return $ AnnounceResp tid interval' 0 0 ipportpairs
+        return $ AnnounceResp tid interval' l s ipportpairs
       2 -> do
         tid <- fromIntegral <$> getWord32be
         _ <- getWord32be
@@ -130,23 +135,59 @@ sendRequest h req = do
 
 recvResponse :: UDPTrackerHandle -> IO UDPResponse
 recvResponse h = do
-  (bs, saddr) <- recvFrom (sock h) 32
+  (bs, saddr) <- recvFrom (sock h) (16*1024)
   return $ decode $ fromStrict bs
 
-connectRequest :: ReaderT UDPTrackerHandle IO ()
+connectRequest :: ReaderT UDPTrackerHandle IO Word32
 connectRequest = do
   h <- ask
-  let pkt = encode $ ConnectReq (tid h)
+  tidi <- liftIO randomIO
+  let pkt = encode $ ConnectReq tidi
   liftIO $ sendRequest h (toStrict pkt)
+  return tidi
 
-connectResponse :: Word32 -> ReaderT UDPTrackerHandle IO Bool
-connectResponse itid = do
+connectResponse :: Word32 -> ReaderT UDPTrackerHandle IO Word64
+connectResponse tid = do
   h <- ask
   resp <- liftIO $ recvResponse h
   -- check if nbytes is at least 16 bytes long
   case resp of
-    (ConnectResp tid cid) -> return $ tid == itid
-    _                     -> return False
+    (ConnectResp tidr cid) ->
+      if tidr == tid
+      then do
+        liftIO $ putStrLn "connect succeeded"
+        return cid
+      else
+        return 0
+    _                      -> return 0
+
+announceRequest :: Word64 -> ByteString -> Word64 -> Word64 -> Word64 -> Word16 -> ReaderT UDPTrackerHandle IO Word32
+announceRequest cid infohash up down left port = do
+  h <- ask
+  tidi <- liftIO randomIO
+  -- connId transId infohash peerId down left up event port)
+  let pkt = encode $ AnnounceReq cid tidi infohash "foo" down left up None port
+  liftIO $ sendRequest h (toStrict pkt)
+  return tidi
+
+data PeerStats = PeerStats { leechers :: Word32
+                           , seeders :: Word32
+                           , peers :: [(IP, Port)]
+                           } deriving (Show)
+
+announceResponse :: Word32 -> ReaderT UDPTrackerHandle IO PeerStats
+announceResponse tid = do
+  h <- ask
+  resp <- liftIO $ recvResponse h
+  case resp of
+    (AnnounceResp tidr interval ss ls xs) ->
+      if tidr == tid
+      then do
+        liftIO $ putStrLn "announce succeeded"
+        return $ PeerStats ls ss xs
+      else
+        return $ PeerStats 0 0 []
+    _ -> return $ PeerStats 0 0 []
 
 getIPPortPairs :: Get [(IP, Port)]
 getIPPortPairs = do
@@ -159,15 +200,29 @@ getIPPortPairs = do
     ipportpairs <- getIPPortPairs
     return $ (ip, port) : ipportpairs
 
-startSession :: IP -> Port -> IO UDPTrackerHandle
-startSession ip port = do
+startSession :: String -> Port -> IO UDPTrackerHandle
+startSession host port = do
   s <- socket AF_INET Datagram defaultProtocol
-  hostAddr <- inet_addr ip
+  addrinfos <- getAddrInfo Nothing (Just host) (Just (show port))
+  let (SockAddrInet p ip) = addrAddress $ head addrinfos
+  hostAddr <- inet_addr (show ip)
   putStrLn "connected to tracker"
-  r <- randomIO
   return $ UDPTrackerHandle { sock = s
-                            , tid = r
                             , addr = (SockAddrInet (fromIntegral port) hostAddr) }
   
 closeSession :: UDPTrackerHandle -> IO ()
 closeSession (UDPTrackerHandle s _ _) = close s
+
+trackerLoop :: String -> Port -> String -> ByteString -> FS.MsgChannel -> TState -> IO ()
+trackerLoop host port peerId infohash fschan tstate = do
+  st' <- FS.getStats fschan
+  st <- readMVar st'
+  let up = FS.bytesRead st
+      down = FS.bytesWritten st
+  handle <- startSession host 2710
+  flip runReaderT handle $ do
+    t1 <- connectRequest
+    cid <- connectResponse t1
+    t2 <- announceRequest cid infohash (fromIntegral up) (fromIntegral down) (fromIntegral (left tstate)) (fromIntegral port)
+    stats <- announceResponse t2
+    liftIO $ print stats
-- 
2.45.2