1 files changed, 60 insertions, 8 deletions

diff --git a/src/Network/BitTorrent/DHT/Query.hs b/src/Network/BitTorrent/DHT/Query.hs
index e5d9bd5f..67dc4541 100644
--- a/src/Network/BitTorrent/DHT/Query.hs
+++ b/src/Network/BitTorrent/DHT/Query.hs
@@ -16,6 +16,8 @@
 {-# LANGUAGE TupleSections       #-}
 {-# LANGUAGE PartialTypeSignatures #-}
 {-# LANGUAGE GADTs #-}
+{-# LANGUAGE RankNTypes #-}
+{-# LANGUAGE MultiParamTypeClasses #-}
 module Network.BitTorrent.DHT.Query
        ( -- * Handler
          -- | To bind specific set of handlers you need to pass
@@ -25,6 +27,7 @@ module Network.BitTorrent.DHT.Query
        , getPeersH
        , announceH
        , defaultHandlers
+       , DataHandlers
 
          -- * Query
          -- ** Basic
@@ -113,6 +116,7 @@ import Data.Serialize
 import System.IO.Unsafe (unsafeInterleaveIO)
 import Data.String
 
+
 {-----------------------------------------------------------------------
 --  Handlers
 -----------------------------------------------------------------------}
@@ -215,20 +219,68 @@ kademliaHandlers logger = do
            , handler (nameFindNodes dht) $ findNodeH getclosest
            ]
 
+class DataHandlers raw dht where
+    dataHandlers ::
+        ( Ord ip , Hashable ip, Typeable ip, Serialize ip) =>
+         (NodeId dht  -> IO [NodeInfo dht ip ()])
+        -> DHTData dht ip
+        -> [MethodHandler raw dht ip]
+
+instance DataHandlers BValue KMessageOf where
+    dataHandlers = bthandlers
+
+bthandlers ::
+    ( Ord ip , Hashable ip, Typeable ip, Serialize ip) =>
+     (NodeId KMessageOf  -> IO [NodeInfo KMessageOf ip ()])
+    -> DHTData KMessageOf ip
+    -> [MethodHandler BValue KMessageOf ip]
+bthandlers getclosest dta =
+    [ MethodHandler "get_peers" $ getPeersH (getpeers dta) (sessionTokens dta)
+    , MethodHandler "announce_peer" $ announceH (contactInfo dta) (sessionTokens dta)
+    ]
+ where
+    getpeers dta ih = do
+      ps <- lookupPeers (contactInfo dta) ih
+      if L.null ps
+        then Left <$> getclosest (toNodeId ih)
+        else return (Right ps)
+
+data MethodHandler raw dht ip =
+        forall a b. ( SerializableTo raw (Response dht b)
+                    , SerializableTo raw (Query dht a)
+                    ) => MethodHandler (QueryMethod dht) (NodeAddr ip -> a -> IO b)
 
 -- | Includes all default query handlers.
-defaultHandlers :: forall ip. (Eq ip, Ord ip, Address ip) => LogFun -> DHT BValue KMessageOf () ip [NodeHandler]
+defaultHandlers :: forall raw dht u ip.
+                    ( Ord (TransactionID dht)
+                    , Ord (NodeId dht)
+                    , Show u
+                    , SerializableTo raw (Response dht (Ping dht))
+                    , SerializableTo raw (Query dht (Ping dht))
+                    , Show (QueryMethod dht)
+                    , Show (NodeId dht)
+                    , FiniteBits (NodeId dht)
+                    , Default u
+                    , Serialize (TransactionID dht)
+                    , WireFormat raw dht
+                    , Kademlia dht
+                    , KRPC (Query dht (Ping dht)) (Response dht (Ping dht))
+                    , Functor dht
+                    , Pretty (NodeInfo dht ip u)
+                    , KRPC (Query dht (FindNode dht ip)) (Response dht (NodeFound dht ip))
+                    , SerializableTo raw (Response dht (NodeFound dht ip))
+                    , SerializableTo raw (Query dht (FindNode dht ip))
+                    , Eq ip, Ord ip, Address ip, DataHandlers raw dht
+                    ) => LogFun -> DHT raw dht u ip [Handler IO dht raw]
 defaultHandlers logger = do
     groknode <- insertNode1
     mynid <- myNodeIdAccordingTo1
-    let handler :: KRPC (Query KMessageOf a) (Response KMessageOf b) => QueryMethod KMessageOf -> (NodeAddr ip -> a -> IO b) -> NodeHandler
-        handler = nodeHandler groknode mynid (logt logger)
-    toks <- asks sessionTokens
-    peers <- asks contactInfo
-    getpeers <- getPeerList1
+    let handler :: MethodHandler raw dht ip -> Handler IO dht raw
+        handler (MethodHandler name action) = nodeHandler groknode mynid (logt logger) name action
+    dta <- asks dhtData
+    getclosest <- getClosest1
     hs <- kademliaHandlers logger
-    return $ hs ++ [ handler "get_peers"     $ getPeersH getpeers toks
-                   , handler "announce_peer" $ announceH peers toks ]
+    return $ hs ++ L.map handler (dataHandlers (fmap (fmap (fmap (const ()))) . getclosest) dta)
 
 {-----------------------------------------------------------------------
 --  Basic queries