7 files changed, 64 insertions, 35 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
diff --git a/src/Network/DHT/Mainline.hs b/src/Network/DHT/Mainline.hs
index 6ef6d450..e5517a3a 100644
--- a/src/Network/DHT/Mainline.hs
+++ b/src/Network/DHT/Mainline.hs
@@ -155,7 +155,7 @@ read_only_key = "ro"
 instance BEncode a => BEncode (Query KMessageOf a) where
   toBEncode Query {..} = toDict $
       BDict.union ( node_id_key .=! queringNodeId
-                   .: read_only_key .=? bool Nothing (Just (1 :: Integer)) queryIsReadOnly
+                   .: read_only_key .=? bool Nothing (Just (1 :: Integer)) (queryIsReadOnly queryExtra)
                    .: endDict)
                   (dict (toBEncode queryParams))
     where
@@ -164,7 +164,7 @@ instance BEncode a => BEncode (Query KMessageOf a) where
 
   fromBEncode v = do
     Query <$> fromDict (field (req node_id_key)) v
-          <*> fromDict (fromMaybe False <$>? read_only_key) v
+          <*> fromDict (IsReadOnlyQuery . fromMaybe False <$>? read_only_key) v
           <*> fromBEncode v
 #else
 data Query a = Query a
@@ -174,11 +174,11 @@ data Query a = Query a
 instance BEncode a => BEncode (Response KMessageOf a) where
   toBEncode = toBEncode . toQuery
     where
-      toQuery (Response nid a) = Query nid False a
+      toQuery (Response nid MainlineResponseData a) = Query nid (IsReadOnlyQuery False) a
 
   fromBEncode b = fromQuery <$> fromBEncode b
     where
-      fromQuery (Query nid _ a) = Response nid a
+      fromQuery (Query nid _ a) = Response nid MainlineResponseData a
 #else
 data Response KMessageOf a = Response KMessageOf a
 #endif
diff --git a/src/Network/DHT/Types.hs b/src/Network/DHT/Types.hs
index 73a7be65..bd2825fb 100644
--- a/src/Network/DHT/Types.hs
+++ b/src/Network/DHT/Types.hs
@@ -27,22 +27,23 @@ data TableParameters msg ip u = TableParameters
 -- of the querying node.
 data Query dht a = Query
   { queringNodeId   :: NodeId dht -- ^ node id of /quering/ node;
-  , queryIsReadOnly :: Bool   -- ^ node is read-only as per BEP 43
+  , queryExtra :: QueryExtra dht -- , queryIsReadOnly :: Bool   -- node is read-only as per BEP 43
   , queryParams     :: a      -- ^ query parameters.
   } deriving (Typeable,Generic)
 
-deriving instance (Eq (NodeId dht), Eq a ) => Eq (Query dht a)
-deriving instance (Show (NodeId dht), Show a ) => Show (Query dht a)
+deriving instance (Eq (NodeId dht), Eq (QueryExtra dht), Eq a ) => Eq (Query dht a)
+deriving instance (Show (NodeId dht), Show (QueryExtra dht), Show a ) => Show (Query dht a)
 
 -- | All responses have an \"id\" key and value containing the node ID
 -- of the responding node.
 data Response dht a = Response
   { queredNodeId :: NodeId dht -- ^ node id of /quered/ node;
+  , responseExtra :: ResponseExtra dht
   , responseVals :: a      -- ^ query result.
   } deriving (Typeable,Generic)
 
-deriving instance (Eq (NodeId dht), Eq a ) => Eq (Response dht a)
-deriving instance (Show (NodeId dht), Show a ) => Show (Response dht a)
+deriving instance (Eq (NodeId dht), Eq (ResponseExtra dht), Eq a ) => Eq (Response dht a)
+deriving instance (Show (NodeId dht), Show (ResponseExtra dht), Show a ) => Show (Response dht a)
 
 -- | The most basic query is a ping. Ping query is used to check if a
 -- quered node is still alive.
diff --git a/src/Network/DatagramServer.hs b/src/Network/DatagramServer.hs
index e004eef3..0e0297a0 100644
--- a/src/Network/DatagramServer.hs
+++ b/src/Network/DatagramServer.hs
@@ -329,8 +329,8 @@ query :: forall h a b raw msg.
     , SerializableTo raw a
     , WireFormat raw msg
     , KRPC msg a b
-    ) => Manager raw msg -> QueryMethod msg -> SockAddr -> a -> IO b
-query mgr meth addr params = queryK mgr meth addr params (\_ x _ -> x)
+    ) => Manager raw msg -> SockAddr -> a -> IO b
+query mgr addr params = queryK mgr addr params (\_ x _ -> x)
 
 -- | Like 'query' but possibly returns your externally routable IP address.
 query' :: forall h a b raw msg.
@@ -340,8 +340,8 @@ query' :: forall h a b raw msg.
     , Serialize (TransactionID msg)
     , SerializableTo raw a , WireFormat raw msg
     , KRPC msg a b
-    ) => Manager raw msg -> QueryMethod msg -> SockAddr -> a -> IO (b , Maybe ReflectedIP)
-query' mgr meth addr params = queryK mgr meth addr params (const (,))
+    ) => Manager raw msg -> SockAddr -> a -> IO (b , Maybe ReflectedIP)
+query' mgr addr params = queryK mgr addr params (const (,))
 
 -- | Enqueue a query, but give us the complete BEncoded content sent by the
 -- remote Node.  This is useful for handling extensions that this library does
@@ -354,8 +354,8 @@ queryRaw :: forall h a b raw msg.
     , SerializableTo raw a
     , WireFormat raw msg
     , KRPC msg a b
-    ) => Manager raw msg -> QueryMethod msg -> SockAddr -> a -> IO (b , raw)
-queryRaw mgr meth addr params = queryK mgr meth addr params (\raw x _ -> (x,raw))
+    ) => Manager raw msg -> SockAddr -> a -> IO (b , raw)
+queryRaw mgr addr params = queryK mgr addr params (\raw x _ -> (x,raw))
 
 queryK :: forall h a b x raw msg.
     ( SerializableTo raw b
@@ -366,10 +366,10 @@ queryK :: forall h a b x raw msg.
     , Serialize (TransactionID msg)
     , KRPC msg a b
     ) =>
-    Manager raw msg -> QueryMethod msg -> SockAddr -> a -> (raw -> b -> Maybe ReflectedIP -> x) -> IO x
-queryK mgr@Manager{..} meth addr params kont = do
+    Manager raw msg -> SockAddr -> a -> (raw -> b -> Maybe ReflectedIP -> x) -> IO x
+queryK mgr@Manager{..} addr params kont = do
   tid <- liftIO $ genTransactionId transactionCounter
-  -- let queryMethod = method :: Method a b
+  let Method meth = method :: Method msg a b
   let signature = querySignature meth tid addr
   logMsg 'D' "query.sending" signature
 
diff --git a/src/Network/DatagramServer/Mainline.hs b/src/Network/DatagramServer/Mainline.hs
index 87825ce0..7cf96e44 100644
--- a/src/Network/DatagramServer/Mainline.hs
+++ b/src/Network/DatagramServer/Mainline.hs
@@ -49,6 +49,8 @@ module Network.DatagramServer.Mainline
        , KMessageOf (..)
        , KMessage
        , KQueryArgs
+       , QueryExtra(..)
+       , ResponseExtra(..)
 
        , NodeId(..)
        , nodeIdSize
@@ -277,6 +279,12 @@ instance Envelope KMessageOf where
     newtype NodeId KMessageOf = NodeId Word160
       deriving (Show, Eq, Ord, Typeable, Bits, FiniteBits)
 
+    newtype QueryExtra KMessageOf = IsReadOnlyQuery { queryIsReadOnly :: Bool }
+      deriving (Show, Eq, Ord, Typeable)
+
+    data ResponseExtra KMessageOf = MainlineResponseData
+      deriving (Show, Eq, Ord, Typeable)
+
     envelopePayload (Q q) = queryArgs q
     envelopePayload (R r) = respVals r
     envelopePayload (E _) = error "TODO: messagePayload for KError"
@@ -296,6 +304,9 @@ instance Envelope KMessageOf where
 
     uniqueTransactionId cnt = return $ TID $ Char8.pack (show cnt)
 
+    fromRoutableNode (IsReadOnlyQuery b) = not b
+
+
 instance WireFormat BValue KMessageOf where
     type SerializableTo BValue = BEncode
     type CipherContext BValue KMessageOf = ()
diff --git a/src/Network/DatagramServer/Types.hs b/src/Network/DatagramServer/Types.hs
index e2d56dc0..13f79afb 100644
--- a/src/Network/DatagramServer/Types.hs
+++ b/src/Network/DatagramServer/Types.hs
@@ -94,6 +94,8 @@ class Envelope envelope where
     data TransactionID envelope
     type QueryMethod envelope
     data NodeId envelope
+    data QueryExtra envelope
+    data ResponseExtra envelope
 
     envelopePayload     :: envelope a -> a
     envelopeTransaction :: envelope a -> TransactionID envelope
@@ -119,6 +121,9 @@ class Envelope envelope where
 
     uniqueTransactionId :: Int -> IO (TransactionID envelope)
 
+    fromRoutableNode :: QueryExtra envelope -> Bool
+    fromRoutableNode _ = True
+
 -- | In Kademlia, the distance metric is XOR and the result is
 -- interpreted as an unsigned integer.
 newtype NodeDistance nodeid = NodeDistance nodeid
diff --git a/src/Network/KRPC/Method.hs b/src/Network/KRPC/Method.hs
index 2033f808..a31380cc 100644
--- a/src/Network/KRPC/Method.hs
+++ b/src/Network/KRPC/Method.hs
@@ -35,6 +35,7 @@ import Data.String
 import Data.Typeable
 import Network.DatagramServer.Mainline
 import Network.DatagramServer.Types
+import Network.DHT.Types
 
 
 -- | Method datatype used to describe method name, parameters and
@@ -105,3 +106,5 @@ class ( Typeable req, Typeable resp)
   validateExchange :: dht req -> dht resp -> Bool
   validateExchange _ _ = True
 
+  makeQueryExtra :: DHTData dht ip -> Proxy req -> Proxy resp -> IO (QueryExtra dht)
+  makeResponseExtra :: DHTData dht ip -> req -> Proxy resp -> IO (ResponseExtra dht)