1 files changed, 454 insertions, 0 deletions

diff --git a/bittorrent/src/Network/BitTorrent/Tracker/RPC/UDP.hs b/bittorrent/src/Network/BitTorrent/Tracker/RPC/UDP.hs
new file mode 100644
index 00000000..31b6b870
--- /dev/null
+++ b/bittorrent/src/Network/BitTorrent/Tracker/RPC/UDP.hs
@@ -0,0 +1,454 @@
+-- |
+--   Copyright   :  (c) Sam Truzjan 2013-2014
+--   License     :  BSD3
+--   Maintainer  :  pxqr.sta@gmail.com
+--   Stability   :  provisional
+--   Portability :  portable
+--
+--   This module implement UDP tracker protocol.
+--
+--   For protocol details and uri scheme see:
+--   <http://www.bittorrent.org/beps/bep_0015.html>,
+--   <https://www.iana.org/assignments/uri-schemes/prov/udp>
+--
+{-# LANGUAGE RecordWildCards            #-}
+{-# LANGUAGE FlexibleInstances          #-}
+{-# LANGUAGE GeneralizedNewtypeDeriving #-}
+{-# LANGUAGE DeriveDataTypeable         #-}
+module Network.BitTorrent.Tracker.RPC.UDP
+       ( -- * Manager
+         Options (..)
+       , Manager
+       , newManager
+       , closeManager
+       , withManager
+
+         -- * RPC
+       , RpcException (..)
+       , announce
+       , scrape
+       ) where
+
+import Control.Applicative
+import Control.Concurrent
+import Control.Exception
+import Control.Monad
+import Data.Default
+import Data.IORef
+import Data.List as L
+import Data.Map  as M
+import Data.Maybe
+import Data.Serialize
+import Data.Text as T
+import Data.Time
+import Data.Time.Clock.POSIX
+import Data.Traversable
+import Data.Typeable
+import Text.Read (readMaybe)
+import Network.Socket hiding (Connected, connect, listen)
+import Network.Socket.ByteString as BS
+import Network.URI
+import System.Timeout
+
+import Network.BitTorrent.Tracker.Message
+
+{-----------------------------------------------------------------------
+--  Options
+-----------------------------------------------------------------------}
+
+-- | 'System.Timeout.timeout' specific.
+sec :: Int
+sec = 1000000
+
+-- | See <http://www.bittorrent.org/beps/bep_0015.html#time-outs>
+defMinTimeout :: Int
+defMinTimeout = 15
+
+-- | See <http://www.bittorrent.org/beps/bep_0015.html#time-outs>
+defMaxTimeout :: Int
+defMaxTimeout = 15 * 2 ^ (8 :: Int)
+
+-- | See: <http://www.bittorrent.org/beps/bep_0015.html#time-outs>
+defMultiplier :: Int
+defMultiplier = 2
+
+-- TODO why 98?
+defMaxPacketSize :: Int
+defMaxPacketSize = 98
+
+-- | Manager configuration.
+data Options = Options
+  { -- | Max size of a /response/ packet.
+    --
+    --   'optMaxPacketSize' /must/ be a positive value.
+    --
+    optMaxPacketSize :: {-# UNPACK #-} !Int
+
+    -- | Starting timeout interval in seconds. If a response is not
+    -- received after 'optMinTimeout' then 'Manager' repeat RPC with
+    -- timeout interval multiplied by 'optMultiplier' and so on until
+    -- timeout interval reach 'optMaxTimeout'.
+    --
+    --   'optMinTimeout' /must/ be a positive value.
+    --
+  , optMinTimeout    :: {-# UNPACK #-} !Int
+
+    -- | Final timeout interval in seconds. After 'optMaxTimeout'
+    -- reached and tracker still not responding both 'announce' and
+    -- 'scrape' functions will throw 'TimeoutExpired' exception.
+    --
+    --   'optMaxTimeout' /must/ be greater than 'optMinTimeout'.
+    --
+  , optMaxTimeout    :: {-# UNPACK #-} !Int
+
+    -- | 'optMultiplier' /must/ be a positive value.
+  , optMultiplier    :: {-# UNPACK #-} !Int
+  } deriving (Show, Eq)
+
+-- | Options suitable for bittorrent client.
+instance Default Options where
+  def = Options
+    { optMaxPacketSize = defMaxPacketSize
+    , optMinTimeout    = defMinTimeout
+    , optMaxTimeout    = defMaxTimeout
+    , optMultiplier    = defMultiplier
+    }
+
+checkOptions :: Options -> IO ()
+checkOptions Options {..} = do
+  unless (optMaxPacketSize > 0) $ do
+    throwIO $ userError "optMaxPacketSize must be positive"
+
+  unless (optMinTimeout > 0) $ do
+    throwIO $ userError "optMinTimeout must be positive"
+
+  unless (optMaxTimeout > 0) $ do
+    throwIO $ userError "optMaxTimeout must be positive"
+
+  unless (optMultiplier > 0) $ do
+    throwIO $ userError "optMultiplier must be positive"
+
+  unless (optMaxTimeout > optMinTimeout) $ do
+    throwIO $ userError "optMaxTimeout must be greater than optMinTimeout"
+
+
+{-----------------------------------------------------------------------
+--  Manager state
+-----------------------------------------------------------------------}
+
+type ConnectionCache     = Map SockAddr Connection
+
+type PendingResponse     = MVar (Either RpcException Response)
+type PendingTransactions = Map TransactionId PendingResponse
+type PendingQueries      = Map SockAddr      PendingTransactions
+
+-- | UDP tracker manager.
+data Manager = Manager
+  { options         :: !Options
+  , sock            :: !Socket
+--  , dnsCache        :: !(IORef (Map URI SockAddr))
+  , connectionCache :: !(IORef ConnectionCache)
+  , pendingResps    :: !(MVar  PendingQueries)
+  , listenerThread  :: !(MVar ThreadId)
+  }
+
+initManager :: Options -> IO Manager
+initManager opts = Manager opts
+  <$> socket AF_INET Datagram defaultProtocol
+  <*> newIORef M.empty
+  <*> newMVar  M.empty
+  <*> newEmptyMVar
+
+unblockAll :: PendingQueries -> IO ()
+unblockAll m = traverse (traverse unblockCall) m >> return ()
+  where
+    unblockCall ares = putMVar ares (Left ManagerClosed)
+
+resetState :: Manager -> IO ()
+resetState Manager {..} = do
+    writeIORef          connectionCache err
+    m    <- swapMVar    pendingResps    err
+    unblockAll m
+    mtid <- tryTakeMVar listenerThread
+    case mtid of
+      Nothing -> return () -- thread killed by 'closeManager'
+      Just _  -> return () -- thread killed by exception from 'listen'
+    return ()
+  where
+    err = error "UDP tracker manager closed"
+
+-- | This function will throw 'IOException' on invalid 'Options'.
+newManager :: Options -> IO Manager
+newManager opts = do
+  checkOptions opts
+  mgr <- initManager opts
+  tid <- forkIO (listen mgr `finally` resetState mgr)
+  putMVar (listenerThread mgr) tid
+  return mgr
+
+-- | Unblock all RPCs by throwing 'ManagerClosed' exception. No rpc
+-- calls should be performed after manager becomes closed.
+closeManager :: Manager -> IO ()
+closeManager Manager {..} = do
+  close sock
+  mtid <- tryTakeMVar listenerThread
+  case mtid of
+    Nothing  -> return ()
+    Just tid -> killThread tid
+
+-- | Normally you need to use 'Control.Monad.Trans.Resource.allocate'.
+withManager :: Options -> (Manager -> IO a) -> IO a
+withManager opts = bracket (newManager opts) closeManager
+
+{-----------------------------------------------------------------------
+--  Exceptions
+-----------------------------------------------------------------------}
+
+data RpcException
+    -- | Unable to lookup hostname;
+  = HostUnknown
+
+    -- | Unable to lookup hostname;
+  | HostLookupFailed
+
+    -- | Expecting 'udp:', but some other scheme provided.
+  | UnrecognizedScheme String
+
+    -- | Tracker exists but not responding for specific number of seconds.
+  | TimeoutExpired Int
+
+    -- | Tracker responded with unexpected message type.
+  | UnexpectedResponse
+    { expectedMsg :: String
+    , actualMsg   :: String
+    }
+
+    -- | RPC succeed, but tracker responded with error code.
+  | QueryFailed Text
+
+    -- | RPC manager closed while waiting for response.
+  | ManagerClosed
+    deriving (Eq, Show, Typeable)
+
+instance Exception RpcException
+
+{-----------------------------------------------------------------------
+--  Host Addr resolution
+-----------------------------------------------------------------------}
+
+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
+
+resolveURI :: URI -> IO SockAddr
+resolveURI 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
+    _                 -> throwIO HostLookupFailed
+resolveURI _       = throwIO HostUnknown
+
+-- TODO caching?
+getTrackerAddr :: Manager -> URI -> IO SockAddr
+getTrackerAddr _ uri
+  | uriScheme uri == "udp:" = resolveURI uri
+  |        otherwise        = throwIO (UnrecognizedScheme (uriScheme uri))
+
+{-----------------------------------------------------------------------
+  Connection
+-----------------------------------------------------------------------}
+
+connectionLifetime :: NominalDiffTime
+connectionLifetime = 60
+
+data Connection = Connection
+  { connectionId        :: ConnectionId
+  , connectionTimestamp :: UTCTime
+  } deriving Show
+
+-- placeholder for the first 'connect'
+initialConnection :: Connection
+initialConnection = Connection initialConnectionId (posixSecondsToUTCTime 0)
+
+establishedConnection :: ConnectionId -> IO Connection
+establishedConnection cid = Connection cid <$> getCurrentTime
+
+isExpired :: Connection -> IO Bool
+isExpired Connection {..} = do
+  currentTime <- getCurrentTime
+  let timeDiff = diffUTCTime currentTime connectionTimestamp
+  return $ timeDiff > connectionLifetime
+
+{-----------------------------------------------------------------------
+--  Transactions
+-----------------------------------------------------------------------}
+
+-- | Sometimes 'genTransactionId' may return already used transaction
+-- id. We use a good entropy source but the issue /still/ (with very
+-- small probabality) may happen. If the collision happen then this
+-- function tries to find nearest unused slot, otherwise pending
+-- transactions table is full.
+firstUnused :: SockAddr -> TransactionId -> PendingQueries -> TransactionId
+firstUnused addr rid m = do
+  case M.splitLookup rid <$> M.lookup addr m of
+    Nothing                -> rid
+    Just (_ , Nothing, _ ) -> rid
+    Just (lt, Just _ , gt) ->
+      case backwardHole (keys lt) rid <|> forwardHole rid (keys gt) of
+        Nothing  -> error "firstUnused: table is full" -- impossible
+        Just tid -> tid
+  where
+    forwardHole a []
+      | a == maxBound = Nothing
+      |   otherwise   = Just (succ a)
+    forwardHole a (b : xs)
+      | succ a == b   = forwardHole b xs
+      |   otherwise   = Just (succ a)
+
+    backwardHole [] a
+      | a == minBound = Nothing
+      |   otherwise   = Just (pred a)
+    backwardHole (b : xs) a
+      | b == pred a   = backwardHole xs b
+      |   otherwise   = Just (pred a)
+
+register :: SockAddr -> TransactionId -> PendingResponse
+         -> PendingQueries -> PendingQueries
+register addr tid ares = M.alter insertId addr
+  where
+    insertId Nothing  = Just (M.singleton tid ares)
+    insertId (Just m) = Just (M.insert tid ares m)
+
+unregister :: SockAddr -> TransactionId
+           -> PendingQueries -> PendingQueries
+unregister addr tid = M.update deleteId addr
+  where
+    deleteId m
+      | M.null m' = Nothing
+      | otherwise = Just m'
+      where
+        m' = M.delete tid m
+
+-- | Generate a new unused transaction id and register as pending.
+allocTransaction :: Manager -> SockAddr -> PendingResponse -> IO TransactionId
+allocTransaction Manager {..} addr ares =
+  modifyMVar pendingResps $ \ m -> do
+    rndId <- genTransactionId
+    let tid = firstUnused addr rndId m
+    return (register addr tid ares m, tid)
+
+-- | Wake up blocked thread and return response back.
+commitTransaction :: Manager -> SockAddr -> TransactionId -> Response -> IO ()
+commitTransaction Manager {..} addr tid resp =
+  modifyMVarMasked_ pendingResps $ \ m -> do
+    case M.lookup tid =<< M.lookup addr m of
+      Nothing   -> return m -- tracker responded after 'cancelTransaction' fired
+      Just ares -> do
+        putMVar ares (Right resp)
+        return $ unregister addr tid m
+
+-- | Abort transaction forcefully.
+cancelTransaction :: Manager -> SockAddr -> TransactionId -> IO ()
+cancelTransaction Manager {..} addr tid =
+  modifyMVarMasked_ pendingResps $ \m ->
+    return $ unregister addr tid m
+
+-- | Handle responses from trackers.
+listen :: Manager -> IO ()
+listen mgr @ Manager {..} = do
+  forever $ do
+    (bs, addr) <- BS.recvFrom sock (optMaxPacketSize options)
+    case decode bs of
+      Left  _                   -> return () -- parser failed, ignoring
+      Right (TransactionR {..}) -> commitTransaction mgr addr transIdR response
+
+-- | Perform RPC transaction. If the action interrupted transaction
+-- will be aborted.
+transaction :: Manager -> SockAddr -> Connection -> Request -> IO Response
+transaction mgr @ Manager {..} addr conn request = do
+    ares <- newEmptyMVar
+    tid  <- allocTransaction mgr addr ares
+    performTransaction tid ares
+      `onException` cancelTransaction mgr addr tid
+  where
+    performTransaction tid ares = do
+      let trans = TransactionQ (connectionId conn) tid request
+      BS.sendAllTo sock (encode trans) addr
+      takeMVar ares >>= either throwIO return
+
+{-----------------------------------------------------------------------
+--  Connection cache
+-----------------------------------------------------------------------}
+
+connect :: Manager -> SockAddr -> Connection -> IO ConnectionId
+connect m addr conn = do
+  resp <- transaction m addr conn Connect
+  case resp of
+    Connected cid -> return cid
+    Failed    msg -> throwIO $ QueryFailed msg
+    _ -> throwIO $ UnexpectedResponse "connected" (responseName resp)
+
+newConnection :: Manager -> SockAddr -> IO Connection
+newConnection m addr = do
+  connId  <- connect m addr initialConnection
+  establishedConnection connId
+
+refreshConnection :: Manager -> SockAddr -> Connection -> IO Connection
+refreshConnection mgr addr conn = do
+  expired <- isExpired conn
+  if expired
+    then do
+      connId <- connect mgr addr conn
+      establishedConnection connId
+    else do
+      return conn
+
+withCache :: Manager -> SockAddr
+          -> (Maybe Connection -> IO Connection) -> IO Connection
+withCache mgr addr action = do
+  cache <- readIORef (connectionCache mgr)
+  conn  <- action (M.lookup addr cache)
+  writeIORef (connectionCache mgr) (M.insert addr conn cache)
+  return conn
+
+getConnection :: Manager -> SockAddr -> IO Connection
+getConnection mgr addr = withCache mgr addr $
+  maybe (newConnection mgr addr) (refreshConnection mgr addr)
+
+{-----------------------------------------------------------------------
+--  RPC
+-----------------------------------------------------------------------}
+
+retransmission :: Options -> IO a -> IO a
+retransmission Options {..} action = go optMinTimeout
+  where
+    go curTimeout
+      | curTimeout > optMaxTimeout = throwIO $ TimeoutExpired curTimeout
+      |         otherwise          = do
+        r <- timeout (curTimeout * sec) action
+        maybe (go (optMultiplier * curTimeout)) return r
+
+queryTracker :: Manager -> URI -> Request -> IO Response
+queryTracker mgr uri req = do
+  addr <- getTrackerAddr mgr uri
+  retransmission (options mgr) $ do
+    conn <- getConnection  mgr addr
+    transaction mgr addr conn req
+
+-- | This function can throw 'RpcException'.
+announce :: Manager -> URI -> AnnounceQuery -> IO AnnounceInfo
+announce mgr uri q = do
+  resp <- queryTracker mgr uri (Announce q)
+  case resp of
+    Announced info -> return info
+    _ -> throwIO $ UnexpectedResponse "announce" (responseName resp)
+
+-- | This function can throw 'RpcException'.
+scrape :: Manager -> URI -> ScrapeQuery -> IO ScrapeInfo
+scrape mgr uri ihs = do
+  resp <- queryTracker mgr uri (Scrape ihs)
+  case resp of
+    Scraped info -> return $ L.zip ihs info
+    _ -> throwIO $ UnexpectedResponse "scrape" (responseName resp)