Bittorrent udp tracker server transport.

3 files changed, 109 insertions, 16 deletions

diff --git a/dht/src/Network/BitTorrent/Tracker/Message.hs b/dht/src/Network/BitTorrent/Tracker/Message.hs
index e9d12006..f44c7700 100644
--- a/dht/src/Network/BitTorrent/Tracker/Message.hs
+++ b/dht/src/Network/BitTorrent/Tracker/Message.hs
@@ -102,6 +102,7 @@ module Network.BitTorrent.Tracker.Message
 
 import Control.Applicative
 import Control.Monad
+import Crypto.Random
 import Data.BEncode as BE hiding (Result)
 import Data.BEncode.BDict as BE
 import Data.ByteString as BS
@@ -810,8 +811,11 @@ newtype TransactionId = TransactionId Word32
 instance Show TransactionId where
   showsPrec _ (TransactionId tid) = showString "0x" <> showHex tid
 
-genTransactionId :: IO TransactionId
-genTransactionId = (TransactionId . fromIntegral) <$> genToken
+genTransactionId :: DRG g => g -> (TransactionId, g)
+genTransactionId g = let (ws, g') = randomBytesGenerate 4 g
+                         Right w = S.runGet S.getWord32be ws
+                     in (TransactionId w, g')
+
 
 data Request
   = Connect
diff --git a/dht/src/Network/BitTorrent/Tracker/Transport.hs b/dht/src/Network/BitTorrent/Tracker/Transport.hs
new file mode 100644
index 00000000..60cb832d
--- /dev/null
+++ b/dht/src/Network/BitTorrent/Tracker/Transport.hs
@@ -0,0 +1,76 @@
+{-# LANGUAGE TupleSections #-}
+module Network.BitTorrent.Tracker.Transport where
+
+import Control.Concurrent.STM
+import Crypto.Random
+import Data.ByteString (ByteString)
+import Data.Functor.Contravariant
+import Data.IntMap.Strict (IntMap)
+import qualified Data.Serialize as S
+import Network.Socket
+
+import Data.QueryResponse.Table
+import Data.TableMethods
+import Network.Address
+import Network.BitTorrent.Tracker.Message
+import Network.QueryResponse as QR
+
+parseTracker :: ByteString -> SockAddr -> Either String (Transaction Request, SockAddr)
+parseTracker bs saddr = fmap (, saddr) $ S.decode bs
+
+encodeTracker :: Transaction Response -> SockAddr -> (ByteString, SockAddr)
+encodeTracker resp saddr = (S.encode resp,saddr)
+
+type Handler = MethodHandlerA String TransactionId SockAddr (Transaction Request) (Transaction Response)
+
+type Client = QR.ClientA String MessageId TransactionId SockAddr (Transaction Request) (Transaction Response)
+
+handlers :: MessageId -> Maybe Handler
+handlers ConnectId  = Nothing
+handlers AnnounceId = Nothing
+handlers ScrapeId   = Nothing
+handlers ErrorId    = Nothing
+handlers _          = Nothing
+
+methodOfRequest :: Request -> MessageId
+methodOfRequest Connect  {} = ConnectId
+methodOfRequest Announce {} = AnnounceId
+methodOfRequest Scrape   {} = ScrapeId
+
+classify :: Transaction Request -> MessageClass String MessageId TransactionId SockAddr (Transaction Request)
+classify qry = IsQuery (methodOfRequest $ request qry) (transIdQ qry)
+
+dispatch :: DRG g =>
+            DispatchMethodsA
+                        (g, IntMap (Maybe (Transaction Response) -> IO ()))
+                        String
+                        MessageId
+                        TransactionId
+                        SockAddr
+                        (Transaction Request)
+                        (Transaction Response)
+dispatch = DispatchMethods
+    { classifyInbound = classify
+    , lookupHandler   = handlers
+    , tableMethods    = transactionMethods (contramap fromEnum intMapMethods) genTransactionId
+    }
+
+newClient :: ErrorReporter SockAddr (Transaction Request) MessageId TransactionId String
+             -> TransportA String SockAddr (Transaction Request) (Transaction Response)
+             -> IO (ClientA String
+                            MessageId
+                            TransactionId
+                            SockAddr
+                            (Transaction Request)
+                            (Transaction Response))
+newClient err net = do
+    drg <- drgNew
+    state <- newTVarIO (drg, mempty)
+    return QR.Client
+        { clientNet           = net
+        , clientDispatcher    = dispatch
+        , clientErrorReporter = err
+        , clientPending       = state
+        , clientAddress       = const $ return localhost4
+        , clientResponseId    = return
+        }
diff --git a/dht/src/Network/QueryResponse.hs b/dht/src/Network/QueryResponse.hs
index 44039ee0..877c7ab6 100644
--- a/dht/src/Network/QueryResponse.hs
+++ b/dht/src/Network/QueryResponse.hs
@@ -271,12 +271,16 @@ data MessageClass err meth tid addr x
     | IsUnknown err                                       -- ^ None of the above.
 
 -- | Handler for an inbound query of type /x/ from an address of type _addr_.
-data MethodHandler err tid addr x = forall a b. MethodHandler
+type MethodHandler err tid addr x = MethodHandlerA  err tid addr x x
+
+-- | Handler for an inbound query of type /x/ with outbound response of type
+-- /y/ to an address of type /addr/.
+data MethodHandlerA err tid addr x y = forall a b. MethodHandler
     { -- | Parse the query into a more specific type for this method.
       methodParse :: x -> Either err a
       -- | Serialize the response for transmission, given a context /ctx/ and the origin
       -- and destination addresses.
-    , methodSerialize :: tid -> addr -> addr -> b -> x
+    , methodSerialize :: tid -> addr -> addr -> b -> y
       -- | Fully typed action to perform upon the query.  The remote origin
       -- address of the query is provided to the handler.
     , methodAction :: addr -> a -> IO b
@@ -290,49 +294,56 @@ data MethodHandler err tid addr x = forall a b. MethodHandler
     , noreplyAction :: addr -> a -> IO ()
     }
 
+
 -- | To dispatch responses to our outbound queries, we require three
 -- primitives.  See the 'transactionMethods' function to create these
 -- primitives out of a lookup table and a generator for transaction ids.
 --
 -- The type variable /d/ is used to represent the current state of the
 -- transaction generator and the table of pending transactions.
-data TransactionMethods d qid addr x = TransactionMethods
+data TransactionMethods d qid addr y = TransactionMethods
     {
       -- | Before a query is sent, this function stores an 'MVar' to which the
       -- response will be written too.  The returned /qid/ is a transaction id
       -- that can be used to forget the 'MVar' if the remote peer is not
       -- responding.
       dispatchRegister :: POSIXTime -- time of expiry
-                          -> (Maybe x -> IO ()) -- callback upon response (or timeout)
+                          -> (Maybe y -> IO ()) -- callback upon response (or timeout)
                           -> addr
                           -> d
                           -> STM (qid, d)
-      -- | This method is invoked when an incoming packet /x/ indicates it is
+      -- | This method is invoked when an incoming packet /y/ indicates it is
       -- a response to the transaction with id /qid/.  The returned IO action
       -- will write the packet to the correct 'MVar' thus completing the
       -- dispatch.
-    , dispatchResponse :: qid -> x -> d -> STM (d, IO ())
+    , dispatchResponse :: qid -> y -> d -> STM (d, IO ())
       -- | When a timeout interval elapses, this method is called to remove the
       -- transaction from the table.
     , dispatchCancel :: qid -> d -> STM d
     }
 
 -- | A set of methods necessary for dispatching incoming packets.
-data DispatchMethods tbl err meth tid addr x = DispatchMethods
+type DispatchMethods tbl err meth tid addr x = DispatchMethodsA tbl err meth tid addr x x
+
+-- | A set of methods necessary for dispatching incoming packets.
+data DispatchMethodsA tbl err meth tid addr x y = DispatchMethods
     { -- | Classify an inbound packet as a query or response.
       classifyInbound :: x -> MessageClass err meth tid addr x
       -- | Lookup the handler for a inbound query.
-    , lookupHandler :: meth -> Maybe (MethodHandler err tid addr x)
+    , lookupHandler :: meth -> Maybe (MethodHandlerA err tid addr x y)
       -- | Methods for handling incoming responses.
-    , tableMethods :: TransactionMethods tbl tid addr x
+    , tableMethods :: TransactionMethods tbl tid addr y
     }
 
 -- | All inputs required to implement a query\/response client.
-data Client err meth tid addr x = forall tbl. Client
+type Client err meth tid addr x = ClientA err meth tid addr x x
+
+-- | All inputs required to implement a query\/response client.
+data ClientA err meth tid addr x y = forall tbl. Client
     { -- | The 'Transport' used to dispatch and receive packets.
-      clientNet :: Transport err addr x
+      clientNet :: TransportA err addr x y
       -- | Methods for handling inbound packets.
-    , clientDispatcher :: DispatchMethods tbl err meth tid addr x
+    , clientDispatcher :: DispatchMethodsA tbl err meth tid addr x y
       -- | Methods for reporting various conditions.
     , clientErrorReporter :: ErrorReporter addr x meth tid err
       -- | State necessary for routing inbound responses and assigning unique
@@ -355,7 +366,7 @@ data Client err meth tid addr x = forall tbl. Client
 -- | These four parameters are required to implement an outgoing query.  A
 -- peer-to-peer algorithm will define a 'MethodSerializer' for every 'MethodHandler' that
 -- might be returned by 'lookupHandler'.
-data MethodSerializer tid addr x meth a b = MethodSerializer
+data MethodSerializerA tid addr x y meth a b = MethodSerializer
     { -- | Returns the microseconds to wait for a response to this query being
       -- sent to the given address.  The /addr/ may also be modified to add
       -- routing information.
@@ -369,9 +380,11 @@ data MethodSerializer tid addr x meth a b = MethodSerializer
       -- auxiliary notations on all outgoing packets.
     , wrapQuery :: tid -> addr -> addr -> a -> x
       -- | Parse an inbound packet /x/ into a response /b/ for this query.
-    , unwrapResponse :: x -> b
+    , unwrapResponse :: y -> b
     }
 
+type MethodSerializer tid addr x meth a b = MethodSerializerA tid addr x x meth a b
+
 microsecondsDiff :: Int -> POSIXTime
 microsecondsDiff us = fromIntegral us / 1000000