1 files changed, 192 insertions, 54 deletions

diff --git a/src/Network/BitTorrent/Tracker/UDP.hs b/src/Network/BitTorrent/Tracker/UDP.hs
index 13e1298b..23c6cacc 100644
--- a/src/Network/BitTorrent/Tracker/UDP.hs
+++ b/src/Network/BitTorrent/Tracker/UDP.hs
@@ -6,50 +6,93 @@
 --   Portability :  portable
 --
 --   This module implement low-level UDP tracker protocol.
---   For more info see: http://www.bittorrent.org/beps/bep_0015.html
+--   For more info see:
+--   <http://www.bittorrent.org/beps/bep_0015.html>
 --
 {-# LANGUAGE RecordWildCards            #-}
 {-# LANGUAGE FlexibleInstances          #-}
 {-# LANGUAGE GeneralizedNewtypeDeriving #-}
 module Network.BitTorrent.Tracker.UDP
-       ( Request(..), Response(..)
+       ( UDPTracker
+       , initialTracker
+       , putTracker
+       , connectUDP
+       , freshConnection
        ) where
 
 import Control.Applicative
+import Control.Exception
 import Control.Monad
+import Data.ByteString (ByteString)
+import Data.IORef
+import Data.List as L
+import Data.Maybe
+import Data.Monoid
 import Data.Serialize
-import Data.Word
 import Data.Text
 import Data.Text.Encoding
+import Data.Time
+import Data.Word
+import Text.Read (readMaybe)
 import Network.Socket hiding (Connected)
 import Network.Socket.ByteString as BS
+import Network.URI
+import System.Entropy
+import Numeric
 
 import Data.Torrent.Metainfo ()
 import Network.BitTorrent.Tracker.Protocol
 
+{-----------------------------------------------------------------------
+  Tokens
+-----------------------------------------------------------------------}
 
+genToken :: IO Word64
+genToken = do
+    bs <- getEntropy 8
+    either err return $ runGet getWord64be bs
+  where
+    err = error "genToken: impossible happen"
+
+-- TODO rename
 -- | Connection Id is used for entire tracker session.
-newtype ConnId  = ConnId  { getConnId  :: Word64 }
-                  deriving (Show, Eq, Serialize)
+newtype ConnId  = ConnId Word64
+                  deriving (Eq, Serialize)
 
--- | Transaction Id is used for within UDP RPC.
-newtype TransId = TransId { getTransId :: Word32 }
-                  deriving (Show, Eq, Serialize)
+instance Show ConnId where
+  showsPrec _ (ConnId cid) = showString "0x" <> showHex cid
 
-genTransactionId :: IO TransId
-genTransactionId = return (TransId 0)
+genConnectionId :: IO ConnId
+genConnectionId = ConnId <$> genToken
 
 initialConnectionId :: ConnId
-initialConnectionId = ConnId 0
+initialConnectionId = ConnId 0x41727101980
+
+-- TODO rename
+-- | Transaction Id is used within a UDP RPC.
+newtype TransId = TransId Word32
+                  deriving (Eq, Serialize)
+
+instance Show TransId where
+  showsPrec _ (TransId tid) = showString "0x" <> showHex tid
+
+genTransactionId :: IO TransId
+genTransactionId = (TransId . fromIntegral) <$> genToken
+
+{-----------------------------------------------------------------------
+  Transactions
+-----------------------------------------------------------------------}
 
 data Request  = Connect
               | Announce  AnnounceQuery
               | Scrape    ScrapeQuery
+                deriving Show
 
-data Response = Connected
+data Response = Connected ConnId
               | Announced AnnounceInfo
               | Scraped   [ScrapeInfo]
               | Failed    Text
+                deriving Show
 
 -- TODO rename to message?
 data Transaction a = Transaction
@@ -70,7 +113,7 @@ instance Serialize (Transaction Request) where
   put Transaction {..} = do
     case body of
       Connect        -> do
-        put connId
+        put initialConnectionId
         put connectId
         put transId
 
@@ -82,15 +125,15 @@ instance Serialize (Transaction Request) where
 
       Scrape   hashes -> do
         put connId
-        put announceId
+        put scrapeId
         put transId
         forM_ hashes put
 
   get = do
     cid <- get
-    rid <- getWord32be
+    mid <- getWord32be
     tid <- get
-    bod <- getBody rid
+    bod <- getBody mid
 
     return $ Transaction {
         connId  = cid
@@ -100,7 +143,7 @@ instance Serialize (Transaction Request) where
     where
       getBody :: MessageId -> Get Request
       getBody msgId
-        | msgId == connectId  = return Connect
+        | msgId == connectId  = pure Connect
         | msgId == announceId = Announce <$> get
         | msgId == scrapeId   = Scrape   <$> many get
         |       otherwise     = fail "unknown message id"
@@ -108,45 +151,42 @@ instance Serialize (Transaction Request) where
 instance Serialize (Transaction Response) where
   put Transaction {..} = do
     case body of
-      Connected -> do
-        put connId
+      Connected conn -> do
         put connectId
         put transId
+        put conn
 
       Announced info -> do
-        put connId
         put announceId
         put transId
         put info
 
       Scraped infos -> do
-        put connId
         put scrapeId
         put transId
         forM_ infos put
 
       Failed info -> do
-        put connId
         put errorId
         put transId
         put (encodeUtf8 info)
 
 
   get = do
-    cid <- get
-    rid <- getWord32be
+--    cid <- get
+    mid <- getWord32be
     tid <- get
-    bod <- getBody rid
+    bod <- getBody mid
 
     return $ Transaction {
-        connId  = cid
+        connId  = initialConnectionId -- TODO
       , transId = tid
       , body    = bod
       }
     where
       getBody :: MessageId -> Get Response
       getBody msgId
-        | msgId == connectId  = return $ Connected
+        | msgId == connectId  = Connected <$> get
         | msgId == announceId = Announced <$> get
         | msgId == scrapeId   = Scraped   <$> many get
         | msgId == errorId    = do
@@ -154,38 +194,136 @@ instance Serialize (Transaction Response) where
           case decodeUtf8' bs of
             Left ex   -> fail (show ex)
             Right msg -> return $ Failed msg
-        |      otherwise      = fail "unknown message id"
+        |  otherwise  = fail $ "unknown message id: " ++ show msgId
 
-maxPacketSize :: Int
-maxPacketSize = 98 -- announce request packet
+{-----------------------------------------------------------------------
+  Connection
+-----------------------------------------------------------------------}
 
-call :: Request -> IO Response
-call request = do
-  tid <- genTransactionId
-  let trans = Transaction initialConnectionId tid request
-
-  let addr = error "TODO"
-  sock <- socket AF_INET Datagram defaultProtocol
-  BS.sendAllTo sock (encode trans) addr
-  (resp, addr') <- BS.recvFrom sock 4096
-  if addr' /= addr
-    then error "address mismatch"
-    else case decode resp of
-      Left msg -> error msg
-      Right (Transaction {..}) -> do
-        if tid /= transId
-          then error "transaction id mismatch"
-          else return body
+connectionLifetime :: NominalDiffTime
+connectionLifetime = 60
+
+connectionLifetimeServer :: NominalDiffTime
+connectionLifetimeServer = 120
 
 data Connection = Connection
+    { connectionId        :: ConnId
+    , connectionTimestamp :: UTCTime
+    } deriving Show
+
+initialConnection :: IO Connection
+initialConnection = Connection initialConnectionId <$> getCurrentTime
 
-type URI = ()
+isExpired :: Connection -> IO Bool
+isExpired Connection {..} = do
+  currentTime <- getCurrentTime
+  let timeDiff = diffUTCTime currentTime connectionTimestamp
+  return $ timeDiff > connectionLifetime
 
-connectTracker :: URI -> IO Connection
-connectTracker = undefined
+{-----------------------------------------------------------------------
+  RPC
+-----------------------------------------------------------------------}
 
-announceTracker :: Connection -> AnnounceQuery -> IO AnnounceInfo
-announceTracker = undefined
+maxPacketSize :: Int
+maxPacketSize = 98 -- announce request packet
 
-scrape :: Connection -> ScrapeQuery -> IO [ScrapeInfo]
-scrape = undefined
\ No newline at end of file
+setPort :: PortNumber -> SockAddr -> SockAddr
+setPort p (SockAddrInet  _ h)     = SockAddrInet  p h
+setPort p (SockAddrInet6 _ f h s) = SockAddrInet6 p f h s
+setPort _  addr = addr
+
+getTrackerAddr :: URI -> IO SockAddr
+getTrackerAddr URI { uriAuthority = Just (URIAuth {..}) } = do
+  infos <- getAddrInfo Nothing (Just uriRegName) Nothing
+  let port = fromMaybe 0 (readMaybe (L.drop 1 uriPort) :: Maybe Int)
+  case infos of
+    AddrInfo {..} : _ -> return $ setPort (fromIntegral port) addrAddress
+    _                 -> fail "getTrackerAddr: unable to lookup host addr"
+getTrackerAddr _       = fail "getTrackerAddr: hostname unknown"
+
+call :: SockAddr -> ByteString -> IO ByteString
+call addr arg = bracket open close rpc
+  where
+    open = socket AF_INET Datagram defaultProtocol
+    rpc sock = do
+      BS.sendAllTo sock arg addr
+      (res, addr') <- BS.recvFrom sock maxPacketSize
+      unless (addr' == addr) $ do
+        throwIO $ userError "address mismatch"
+      return res
+
+-- TODO retransmissions
+-- TODO blocking
+data UDPTracker = UDPTracker
+    { trackerURI        :: URI
+    , trackerConnection :: IORef Connection
+    }
+
+updateConnection :: ConnId -> UDPTracker -> IO ()
+updateConnection cid UDPTracker {..} = do
+  newConnection <- Connection cid <$> getCurrentTime
+  writeIORef trackerConnection newConnection
+
+getConnectionId :: UDPTracker -> IO ConnId
+getConnectionId UDPTracker {..}
+  = connectionId <$> readIORef trackerConnection
+
+putTracker :: UDPTracker -> IO ()
+putTracker UDPTracker {..} = do
+  print trackerURI
+  print =<< readIORef trackerConnection
+
+transaction :: UDPTracker -> Request -> IO (Transaction Response)
+transaction tracker @ UDPTracker {..} request = do
+  cid <- getConnectionId tracker
+  tid <- genTransactionId
+  let trans = Transaction cid tid request
+
+  addr <- getTrackerAddr trackerURI
+  res  <- call addr (encode trans)
+  case decode res of
+    Right (response @ Transaction {..})
+      | tid == transId -> return response
+      |   otherwise    -> throwIO $ userError "transaction id mismatch"
+    Left msg           -> throwIO $ userError msg
+
+connectUDP :: UDPTracker -> IO ConnId
+connectUDP tracker = do
+  Transaction _ tid resp <- transaction tracker Connect
+  case resp of
+    Connected cid -> return cid
+
+initialTracker :: URI -> IO UDPTracker
+initialTracker uri = do
+  tracker <- UDPTracker uri <$> (newIORef =<< initialConnection)
+  connId  <- connectUDP tracker
+  updateConnection connId tracker
+  return tracker
+
+freshConnection :: UDPTracker -> IO ()
+freshConnection tracker @ UDPTracker {..} = do
+  conn    <- readIORef trackerConnection
+  expired <- isExpired conn
+  when expired $ do
+    connId <- connectUDP tracker
+    updateConnection connId tracker
+
+{-
+
+announceUDP :: UDPTracker -> AnnounceQuery -> IO AnnounceInfo
+announceUDP t query = do
+  Transaction tid cid resp <- call transaction (Announce query)
+  case resp of
+    Announced info -> return info
+    _              -> fail "response type mismatch"
+
+scrapeUDP :: UDPTracker -> ScrapeQuery -> IO Scrape
+scrapeUDP UDPTracker {..} query = do
+  resp <- call trackerURI $ Scrape query
+  case resp of
+    Scraped scrape -> return undefined
+
+instance Tracker UDPTracker where
+  announce = announceUDP
+  scrape_  = scrapeUDP
+-}
\ No newline at end of file