2 files changed, 30 insertions, 11 deletions

diff --git a/src/Network/DatagramServer.hs b/src/Network/DatagramServer.hs
index d478d667..e004eef3 100644
--- a/src/Network/DatagramServer.hs
+++ b/src/Network/DatagramServer.hs
@@ -127,6 +127,7 @@ import Network.Socket hiding (listen)
 import Network.Socket.ByteString as BS
 import System.IO.Error
 import System.Timeout
+import Network.KRPC.Method
 
 
 {-----------------------------------------------------------------------
@@ -320,7 +321,15 @@ sendQuery sock addr q = handle sockError $ sendMessage sock addr q
 --  This function should throw 'QueryFailure' exception if quered node
 --  respond with @error@ message or the query timeout expires.
 --
-query :: forall h a b raw msg. (SerializableTo raw b, Show (QueryMethod msg), Ord (TransactionID msg), Serialize (TransactionID msg), SerializableTo raw a, WireFormat raw msg) => Manager raw msg -> QueryMethod msg -> SockAddr -> a -> IO b
+query :: forall h a b raw msg.
+    ( SerializableTo raw b
+    , Show (QueryMethod msg)
+    , Ord (TransactionID msg)
+    , Serialize (TransactionID 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)
 
 -- | Like 'query' but possibly returns your externally routable IP address.
@@ -330,6 +339,7 @@ query' :: forall h a b raw msg.
     , Ord (TransactionID 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 (,))
 
@@ -343,6 +353,7 @@ queryRaw :: 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 , raw)
 queryRaw mgr meth addr params = queryK mgr meth addr params (\raw x _ -> (x,raw))
 
@@ -353,6 +364,7 @@ queryK :: forall h a b x raw msg.
     , Show (QueryMethod msg)
     , Ord (TransactionID 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
@@ -373,16 +385,18 @@ queryK mgr@Manager{..} meth addr params kont = do
       `onException` unregisterQuery (tid, addr) pendingCalls
 
     timeout (optQueryTimeout options * 10 ^ (6 :: Int)) $ do
-      -- TODO: Loop with polymorphic sanity-check (to handle Tox's nonces).
-      (raw,res) <- readMVar ares -- MVar (KQueryArgs, KResult)
-      case res of
-        Left  (KError   c m _) -> throwIO $ QueryFailed c (T.decodeUtf8 m)
-        Right m -> case decodePayload m of
-                    Right r -> case envelopeClass (r :: msg b) of
-                                Response reflectedAddr -> pure $ kont raw (envelopePayload r) reflectedAddr
-                                Error (KError   c m _) -> throwIO $ QueryFailed c (T.decodeUtf8 m) -- XXX neccessary?
-                                Query _                -> throwIO $ QueryFailed ProtocolError "BUG!! UNREACHABLE"
-                    Left e  -> throwIO $ QueryFailed ProtocolError (T.pack e)
+      fix $ \loop -> do
+          (raw,res) <- readMVar ares -- MVar (KQueryArgs, KResult)
+          case res of
+            Left  (KError   c m _) -> throwIO $ QueryFailed c (T.decodeUtf8 m)
+            Right m -> case decodePayload m of
+                        Right r -> case envelopeClass (r :: msg b) of
+                                    Response reflectedAddr
+                                        | validateExchange q r -> return $ kont raw (envelopePayload r) reflectedAddr
+                                        | otherwise            -> loop
+                                    Error (KError   c m _) -> throwIO $ QueryFailed c (T.decodeUtf8 m) -- XXX neccessary?
+                                    Query _                -> throwIO $ QueryFailed ProtocolError "BUG!! UNREACHABLE"
+                        Left e  -> throwIO $ QueryFailed ProtocolError (T.pack e)
 
   case mres of
     Just res -> do
diff --git a/src/Network/KRPC/Method.hs b/src/Network/KRPC/Method.hs
index 84c7fe4c..2033f808 100644
--- a/src/Network/KRPC/Method.hs
+++ b/src/Network/KRPC/Method.hs
@@ -100,3 +100,8 @@ class ( Typeable req, Typeable resp)
   method = Method $ fromString $ L.map toLower $ show $ typeOf hole
     where
       hole = error "krpc.method: impossible" :: req
+
+
+  validateExchange :: dht req -> dht resp -> Bool
+  validateExchange _ _ = True
+