WIP: tracker protocol

2 files changed, 75 insertions, 43 deletions

diff --git a/dht/src/Network/BitTorrent/Tracker/Transport.hs b/dht/src/Network/BitTorrent/Tracker/Transport.hs
index 60cb832d..5d225a7e 100644
--- a/dht/src/Network/BitTorrent/Tracker/Transport.hs
+++ b/dht/src/Network/BitTorrent/Tracker/Transport.hs
@@ -1,5 +1,10 @@
 {-# LANGUAGE TupleSections #-}
-module Network.BitTorrent.Tracker.Transport where
+module Network.BitTorrent.Tracker.Transport
+    ( implementTracker
+    , Callbacks(..)
+    , parseTracker
+    , encodeTracker
+    ) where
 
 import Control.Concurrent.STM
 import Crypto.Random
@@ -7,7 +12,7 @@ import Data.ByteString (ByteString)
 import Data.Functor.Contravariant
 import Data.IntMap.Strict (IntMap)
 import qualified Data.Serialize as S
-import Network.Socket
+import Network.Socket (SockAddr)
 
 import Data.QueryResponse.Table
 import Data.TableMethods
@@ -15,9 +20,11 @@ import Network.Address
 import Network.BitTorrent.Tracker.Message
 import Network.QueryResponse as QR
 
+-- | For use with 'layerTransport'.
 parseTracker :: ByteString -> SockAddr -> Either String (Transaction Request, SockAddr)
 parseTracker bs saddr = fmap (, saddr) $ S.decode bs
 
+-- | For use with 'layerTransport'.
 encodeTracker :: Transaction Response -> SockAddr -> (ByteString, SockAddr)
 encodeTracker resp saddr = (S.encode resp,saddr)
 
@@ -25,12 +32,28 @@ type Handler = MethodHandlerA String TransactionId SockAddr (Transaction Request
 
 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
+
+-- | Handlers for inbound tracker requests.  It is safe to throw 'DropQuery'
+-- from these, which should be done if the ConnectionId for an address does not
+-- match the one issued by the 'callOnConnect' method.
+data Callbacks = Callbacks
+    { callOnConnect  :: SockAddr -> IO ConnectionId
+    , callOnAnnounce :: SockAddr -> ConnectionId -> AnnounceQuery -> IO AnnounceInfo
+    , callOnScrape   :: SockAddr -> ConnectionId -> ScrapeQuery -> IO [ScrapeEntry]
+    }
+
+handlers :: Callbacks -> MessageId -> Maybe Handler
+handlers c ConnectId  = Just $ MethodHandler (const $ Right ())
+                                             (\qid from to conid -> TransactionR qid $ Connected conid)
+                                             (\from () -> callOnConnect c from)
+handlers c AnnounceId = Just $ MethodHandler (\TransactionQ{connIdQ=cid,request=Announce qry} -> Right (cid,qry))
+                                             (\qid from to info -> TransactionR qid $ Announced info)
+                                             (uncurry . callOnAnnounce c)
+handlers c ScrapeId   = Just $ MethodHandler (\TransactionQ{connIdQ=cid,request=Scrape qry} -> Right (cid,qry))
+                                             (\qid from to entries -> TransactionR qid $ Scraped entries)
+                                             (uncurry . callOnScrape c)
+handlers c ErrorId    = Nothing
+handlers c _          = Nothing
 
 methodOfRequest :: Request -> MessageId
 methodOfRequest Connect  {} = ConnectId
@@ -41,36 +64,34 @@ classify :: Transaction Request -> MessageClass String MessageId TransactionId S
 classify qry = IsQuery (methodOfRequest $ request qry) (transIdQ qry)
 
 dispatch :: DRG g =>
-            DispatchMethodsA
-                        (g, IntMap (Maybe (Transaction Response) -> IO ()))
+            Callbacks
+            -> DispatchMethodsA
+                        (g, IntMap (Maybe (Transaction Request) -> IO ()))
                         String
                         MessageId
                         TransactionId
                         SockAddr
                         (Transaction Request)
                         (Transaction Response)
-dispatch = DispatchMethods
+dispatch c = DispatchMethods
     { classifyInbound = classify
-    , lookupHandler   = handlers
+    , lookupHandler   = handlers c
     , 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
+implementTracker :: ErrorReporter SockAddr (Transaction Request) MessageId TransactionId String
+                     -> TransportA String SockAddr (Transaction Request) (Transaction Response)
+                     -> Callbacks
+                     -> IO (TransportA String SockAddr (Transaction Request) (Transaction Response))
+implementTracker err net c = do
     drg <- drgNew
     state <- newTVarIO (drg, mempty)
-    return QR.Client
-        { clientNet           = net
-        , clientDispatcher    = dispatch
-        , clientErrorReporter = err
-        , clientPending       = state
-        , clientAddress       = const $ return localhost4
-        , clientResponseId    = return
-        }
+    let client = QR.Client
+                    { clientNet           = net
+                    , clientDispatcher    = dispatch c
+                    , clientErrorReporter = err
+                    , clientPending       = state
+                    , clientAddress       = const $ return localhost4
+                    , clientResponseId    = return
+                    }
+    return $ addHandler (\err -> return ()) (handleMessage client) net
diff --git a/dht/src/Network/QueryResponse.hs b/dht/src/Network/QueryResponse.hs
index 877c7ab6..20e7ecf0 100644
--- a/dht/src/Network/QueryResponse.hs
+++ b/dht/src/Network/QueryResponse.hs
@@ -165,7 +165,7 @@ partitionTransport parse encodex tr =
 -- Note: If you add a handler to one of the branches before applying a
 -- 'mergeTransports' combinator, then this handler may not block or return
 -- Nothing.
-addHandler :: (err -> IO ()) -> (addr -> x -> IO (Maybe (x -> x))) -> Transport err addr x -> Transport err addr x
+addHandler :: (err -> IO ()) -> (addr -> x -> IO (Maybe (x -> x))) -> TransportA err addr x y -> TransportA err addr x y
 addHandler onParseError f tr = tr
     { awaitMessage = \kont -> fix $ \eat -> awaitMessage tr $ \case
         Arrival addr x -> f addr x >>= maybe (join $ atomically eat) (kont . Arrival addr . ($ x))
@@ -283,6 +283,8 @@ data MethodHandlerA err tid addr x y = forall a b. MethodHandler
     , 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.
+      --
+      -- TODO: Allow queries to be ignored?
     , methodAction :: addr -> a -> IO b
     }
       -- | See also 'IsUnsolicited' which likely makes this constructor unnecessary.
@@ -301,22 +303,22 @@ data MethodHandlerA err tid addr x y = forall a b. MethodHandler
 --
 -- 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 y = TransactionMethods
+data TransactionMethods d qid addr x = 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 y -> IO ()) -- callback upon response (or timeout)
+                          -> (Maybe x -> IO ()) -- callback upon response (or timeout)
                           -> addr
                           -> d
                           -> STM (qid, d)
-      -- | This method is invoked when an incoming packet /y/ indicates it is
+      -- | This method is invoked when an incoming packet /x/ 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 -> y -> d -> STM (d, IO ())
+    , dispatchResponse :: qid -> x -> 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
@@ -332,7 +334,7 @@ data DispatchMethodsA tbl err meth tid addr x y = DispatchMethods
       -- | Lookup the handler for a inbound query.
     , lookupHandler :: meth -> Maybe (MethodHandlerA err tid addr x y)
       -- | Methods for handling incoming responses.
-    , tableMethods :: TransactionMethods tbl tid addr y
+    , tableMethods :: TransactionMethods tbl tid addr x
     }
 
 -- | All inputs required to implement a query\/response client.
@@ -466,17 +468,26 @@ contramapAddr f (MethodHandler p s a)
 contramapAddr f (NoReply p a)
     = NoReply p (\addr arg -> a (f addr) arg)
 
+-- | Query handlers can throw this to ignore a query instead of responding to
+-- it.
+data DropQuery = DropQuery
+ deriving Show
+
+instance Exception DropQuery
+
 -- | Attempt to invoke a 'MethodHandler' upon a given inbound query.  If the
 -- parse is successful, the returned IO action will construct our reply if
 -- there is one.  Otherwise, a parse err is returned.
-dispatchQuery :: MethodHandler err tid addr x -- ^ Handler to invoke.
-                 -> tid                       -- ^ The transaction id for this query\/response session.
-                 -> addr                      -- ^ Our own address, to which the query was sent.
-                 -> x                         -- ^ The query packet.
-                 -> addr                      -- ^ The origin address of the query.
-                 -> Either err (IO (Maybe x))
+dispatchQuery :: MethodHandlerA err tid addr x y -- ^ Handler to invoke.
+                 -> tid                          -- ^ The transaction id for this query\/response session.
+                 -> addr                         -- ^ Our own address, to which the query was sent.
+                 -> x                            -- ^ The query packet.
+                 -> addr                         -- ^ The origin address of the query.
+                 -> Either err (IO (Maybe y))
 dispatchQuery (MethodHandler unwrapQ wrapR f) tid self x addr =
-    fmap (\a -> Just . wrapR tid self addr <$> f addr a) $ unwrapQ x
+    fmap (\a -> catch (Just . wrapR tid self addr <$> f addr a)
+                      (\DropQuery -> return Nothing))
+        $ unwrapQ x
 dispatchQuery (NoReply unwrapQ f) tid self x addr =
     fmap (\a -> f addr a >> return Nothing) $ unwrapQ x
 
@@ -515,7 +526,7 @@ transactionMethods methods generate = transactionMethods' id id methods generate
 -- in a forked thread that loops until 'awaitMessage' returns 'Nothing' or
 -- throws an exception.
 handleMessage ::
-    Client err meth tid addr x
+    ClientA err meth tid addr x y
     -> addr
     -> x
     -> IO (Maybe (x -> x))