Expose UDP tracker specific message types

Those can be used to implement UDP tracker server by third party
libraries or projects.

1 files changed, 162 insertions, 1 deletions

diff --git a/src/Network/BitTorrent/Tracker/Message.hs b/src/Network/BitTorrent/Tracker/Message.hs
index edafdaba..8131ecf0 100644
--- a/src/Network/BitTorrent/Tracker/Message.hs
+++ b/src/Network/BitTorrent/Tracker/Message.hs
@@ -25,6 +25,7 @@
 {-# LANGUAGE DeriveDataTypeable         #-}
 {-# LANGUAGE DeriveFunctor              #-}
 {-# LANGUAGE ScopedTypeVariables        #-}
+{-# LANGUAGE TypeFamilies               #-}
 {-# OPTIONS -fno-warn-orphans           #-}
 module Network.BitTorrent.Tracker.Message
        ( -- * Announce
@@ -76,6 +77,21 @@ module Network.BitTorrent.Tracker.Message
 
          -- ** Extra
        , queryToSimpleQuery
+
+         -- * UDP specific
+         -- ** Connection
+       , ConnectionId
+       , initialConnectionId
+
+         -- ** Messages
+       , Request (..)
+       , Response (..)
+       , responseName
+
+         -- ** Transaction
+       , genTransactionId
+       , TransactionId
+       , Transaction  (..)
        )
        where
 
@@ -104,7 +120,9 @@ import Network
 import Network.HTTP.Types.QueryLike
 import Network.HTTP.Types.URI hiding (urlEncode)
 import Network.HTTP.Types.Status
-import Network.Socket
+import Network.Socket hiding (Connected)
+import Numeric
+import System.Entropy
 import Text.Read (readMaybe)
 
 import Data.Torrent.InfoHash
@@ -765,3 +783,146 @@ scrapeType = "text/plain"
 --
 parseFailureStatus :: ParamParseFailure -> Status
 parseFailureStatus = mkStatus <$> parseFailureCode <*> parseFailureMessage
+
+{-----------------------------------------------------------------------
+--  UDP specific message types
+-----------------------------------------------------------------------}
+
+genToken :: IO Word64
+genToken = do
+    bs <- getEntropy 8
+    either err return $ runGet getWord64be bs
+  where
+    err = error "genToken: impossible happen"
+
+-- | Connection Id is used for entire tracker session.
+newtype ConnectionId  = ConnectionId Word64
+  deriving (Eq, Serialize)
+
+instance Show ConnectionId where
+  showsPrec _ (ConnectionId cid) = showString "0x" <> showHex cid
+
+initialConnectionId :: ConnectionId
+initialConnectionId =  ConnectionId 0x41727101980
+
+-- | Transaction Id is used within a UDP RPC.
+newtype TransactionId = TransactionId Word32
+  deriving (Eq, Ord, Enum, Bounded, Serialize)
+
+instance Show TransactionId where
+  showsPrec _ (TransactionId tid) = showString "0x" <> showHex tid
+
+genTransactionId :: IO TransactionId
+genTransactionId = (TransactionId . fromIntegral) <$> genToken
+
+data Request
+  = Connect
+  | Announce  AnnounceQuery
+  | Scrape    ScrapeQuery
+    deriving Show
+
+data Response
+  = Connected ConnectionId
+  | Announced AnnounceInfo
+  | Scraped   [ScrapeEntry]
+  | Failed    Text
+    deriving Show
+
+responseName :: Response -> String
+responseName (Connected _) = "connected"
+responseName (Announced _) = "announced"
+responseName (Scraped   _) = "scraped"
+responseName (Failed    _) = "failed"
+
+data family Transaction a
+data instance Transaction Request  = TransactionQ
+    { connIdQ  :: {-# UNPACK #-} !ConnectionId
+    , transIdQ :: {-# UNPACK #-} !TransactionId
+    , request  :: !Request
+    } deriving Show
+data instance Transaction Response = TransactionR
+    { transIdR :: {-# UNPACK #-} !TransactionId
+    , response :: !Response
+    } deriving Show
+
+-- TODO newtype
+newtype MessageId = MessageId Word32
+                    deriving (Show, Eq, Num, Serialize)
+
+connectId, announceId, scrapeId, errorId :: MessageId
+connectId  = 0
+announceId = 1
+scrapeId   = 2
+errorId    = 3
+
+instance Serialize (Transaction Request) where
+  put TransactionQ {..} = do
+    case request of
+      Connect        -> do
+        put initialConnectionId
+        put connectId
+        put transIdQ
+
+      Announce ann -> do
+        put connIdQ
+        put announceId
+        put transIdQ
+        put ann
+
+      Scrape   hashes -> do
+        put connIdQ
+        put scrapeId
+        put transIdQ
+        forM_ hashes put
+
+  get = do
+      cid <- get
+      mid <- get
+      TransactionQ cid <$> S.get <*> getBody mid
+    where
+      getBody :: MessageId -> S.Get Request
+      getBody msgId
+        | msgId == connectId  = pure Connect
+        | msgId == announceId = Announce <$> get
+        | msgId == scrapeId   = Scrape   <$> many get
+        |       otherwise     = fail errMsg
+        where
+          errMsg = "unknown request: " ++ show msgId
+
+instance Serialize (Transaction Response) where
+  put TransactionR {..} = do
+    case response of
+      Connected conn -> do
+        put connectId
+        put transIdR
+        put conn
+
+      Announced info -> do
+        put announceId
+        put transIdR
+        put info
+
+      Scraped infos -> do
+        put scrapeId
+        put transIdR
+        forM_ infos put
+
+      Failed info -> do
+        put errorId
+        put transIdR
+        put (encodeUtf8 info)
+
+
+  get = do
+      mid <- get
+      TransactionR <$> get <*> getBody mid
+    where
+      getBody :: MessageId -> S.Get Response
+      getBody msgId
+        | msgId == connectId  = Connected <$> get
+        | msgId == announceId = Announced <$> get
+        | msgId == scrapeId   = Scraped   <$> many get
+        | msgId == errorId    = (Failed . decodeUtf8) <$> get
+        |       otherwise     = fail msg
+        where
+          msg = "unknown response: " ++ show msgId