1 files changed, 26 insertions, 17 deletions

diff --git a/src/Network/BitTorrent/DHT/Query.hs b/src/Network/BitTorrent/DHT/Query.hs
index 4c980e22..87081d38 100644
--- a/src/Network/BitTorrent/DHT/Query.hs
+++ b/src/Network/BitTorrent/DHT/Query.hs
@@ -127,32 +127,37 @@ nodeHandler :: ( Address ip
                )
            => (NodeInfo KMessageOf ip () -> Maybe ReflectedIP -> IO ()) -> (NodeAddr ip -> IO (NodeId KMessageOf)) -> (Char -> String -> Text -> IO ()) -> QueryMethod KMessageOf -> (NodeAddr ip -> a -> IO b) -> NodeHandler
 -}
-nodeHandler ::
-    (Address addr, WireFormat raw msg, Pretty (NodeInfo dht addr u),
+nodeHandler :: forall raw dht addr u t q r.
+    (Address addr, WireFormat raw dht, Pretty (NodeInfo dht addr u),
      Default u,
-     IsString t, Functor msg,
+     IsString t, Functor dht,
+     KRPC dht (Query dht q) (Response dht r),
      SerializableTo raw (Response dht r),
      SerializableTo raw (Query dht q)) =>
     (NodeInfo dht addr u -> Maybe ReflectedIP -> IO ())
     -> (NodeAddr addr -> IO (NodeId dht))
     -> (Char -> t -> Text -> IO ())
-    -> QueryMethod msg
+    -> DHTData dht addr
+    -> QueryMethod dht
     -> (NodeAddr addr -> q -> IO r)
-    -> Handler IO msg raw
-nodeHandler insertNode myNodeIdAccordingTo logm method action = handler method $ \ sockAddr qry -> do
-  let remoteId  = queringNodeId qry
-      read_only = queryIsReadOnly qry
-      q         = queryParams qry
+    -> Handler IO dht raw
+nodeHandler insertNode myNodeIdAccordingTo logm dta method action = handler method $ \ sockAddr qry -> do
+  let remoteId = queringNodeId qry
+      qextra   = queryExtra qry
+      resptype = Proxy :: Proxy (Response dht r)
+      q        = queryParams qry
   case fromSockAddr sockAddr of
     Nothing    -> throwIO BadAddress
     Just naddr -> do
+      rextra <- liftIO $ makeResponseExtra dta qry resptype
       let ni = NodeInfo remoteId naddr def
       -- Do not route read-only nodes. (bep 43)
-      if read_only
-        then logm 'W' "nodeHandler" $ "READ-ONLY " <> T.pack (show $ pPrint ni)
-        else insertNode ni Nothing >> return () -- TODO need to block. why?
+      if fromRoutableNode qextra
+        then insertNode ni Nothing >> return () -- TODO need to block. why?
+        else logm 'W' "nodeHandler" $ "READ-ONLY " <> T.pack (show $ pPrint ni)
       Response
         <$> myNodeIdAccordingTo naddr
+        <*> pure rextra
         <*> action naddr q
 
 -- | Default 'Ping' handler.
@@ -208,11 +213,12 @@ kademliaHandlers :: forall raw dht u ip. (Eq ip, Ord ip, Address ip
 kademliaHandlers logger = do
     groknode <- insertNode1
     mynid <- myNodeIdAccordingTo1
+    dta <- asks dhtData
     let handler :: ( KRPC dht (Query dht a) (Response dht b)
                    , SerializableTo raw (Response dht b)
                    , SerializableTo raw (Query dht a)
                    ) => QueryMethod dht -> (NodeAddr ip -> a -> IO b) -> Handler IO dht raw
-        handler = nodeHandler groknode mynid (logt logger)
+        handler = nodeHandler groknode mynid (logt logger) dta
         dht = Proxy :: Proxy dht
     getclosest <- getClosest1
     return [ handler (namePing dht)      $ pingH dht
@@ -248,6 +254,7 @@ bthandlers getclosest dta =
 data MethodHandler raw dht ip =
         forall a b. ( SerializableTo raw (Response dht b)
                     , SerializableTo raw (Query dht a)
+                    , KRPC dht (Query dht a) (Response dht b)
                     ) => MethodHandler (QueryMethod dht) (NodeAddr ip -> a -> IO b)
 
 -- | Includes all default query handlers.
@@ -275,9 +282,9 @@ defaultHandlers :: forall raw dht u ip.
 defaultHandlers logger = do
     groknode <- insertNode1
     mynid <- myNodeIdAccordingTo1
-    let handler :: MethodHandler raw dht ip -> Handler IO dht raw
-        handler (MethodHandler name action) = nodeHandler groknode mynid (logt logger) name action
     dta <- asks dhtData
+    let handler :: MethodHandler raw dht ip -> Handler IO dht raw
+        handler (MethodHandler name action) = nodeHandler groknode mynid (logt logger) dta name action
     getclosest <- getClosest1
     hs <- kademliaHandlers logger
     return $ hs ++ L.map handler (dataHandlers (fmap (fmap (fmap (const ()))) . getclosest) dta)
@@ -664,10 +671,12 @@ queryNode' :: forall raw dht u a b ip.
     ) => NodeAddr ip -> a -> DHT raw dht u ip (NodeId dht, b, Maybe ReflectedIP)
 queryNode' addr q = do
   nid <- myNodeIdAccordingTo addr
+  dta <- asks dhtData
+  qextra <- liftIO $ makeQueryExtra dta (Proxy :: Proxy (Query dht q)) (Proxy :: Proxy (Response dht b))
   let read_only = False -- TODO: check for NAT issues. (BEP 43)
-  let KRPC.Method name = KRPC.method :: KRPC.Method dht (Query dht a) (Response dht b)
+  -- let KRPC.Method name = KRPC.method :: KRPC.Method dht (Query dht a) (Response dht b)
   mgr <- asks manager
-  (Response remoteId r, witnessed_ip) <- liftIO $ query' mgr name (toSockAddr addr) (Query nid read_only q)
+  (Response remoteId rextra r , witnessed_ip) <- liftIO $ query' mgr (toSockAddr addr) (Query nid qextra q)
   -- \$(logDebugS) "queryNode" $ "Witnessed address: " <> T.pack (show witnessed_ip)
   --                             <> " by " <> T.pack (show (toSockAddr addr))
   _ <- insertNode (NodeInfo remoteId addr def) witnessed_ip