1 files changed, 19 insertions, 28 deletions

diff --git a/src/Network/QueryResponse.hs b/src/Network/QueryResponse.hs
index 0fa1a05a..70d981e2 100644
--- a/src/Network/QueryResponse.hs
+++ b/src/Network/QueryResponse.hs
@@ -108,18 +108,18 @@ partitionTransportM parse encodex tr = do
                 }
     return (xtr, ytr)
 
-addHandler :: (addr -> x -> IO (Maybe (x -> x))) -> Transport err addr x -> Transport err addr x
-addHandler f tr = tr
+addHandler :: ErrorReporter addr x meth tid err -> (addr -> x -> IO (Maybe (x -> x))) -> Transport err addr x -> Transport err addr x
+addHandler err f tr = tr
     { awaitMessage = \kont -> fix $ \eat -> awaitMessage tr $ \m -> do
         case m of
             Just (Right (x, addr)) -> f addr x >>= maybe eat (kont . Just . Right . (, addr) . ($ x))
-            Just (Left  e        ) -> kont $ Just (Left e)
+            Just (Left  e        ) -> reportParseError err e >> kont (Just $ Left e)
             Nothing                -> kont $ Nothing
     }
 
 -- | Modify a 'Transport' to invoke an action upon every received packet.
 onInbound :: (addr -> x -> IO ()) -> Transport err addr x -> Transport err addr x
-onInbound f tr = addHandler (\addr x -> f addr x >> return (Just id)) tr
+onInbound f tr = addHandler ignoreErrors (\addr x -> f addr x >> return (Just id)) tr
 
 -- * Using a query\/response client.
 
@@ -153,16 +153,17 @@ sendQuery ::
         -> a                                    -- ^ The outbound query.
         -> addr                                 -- ^ Destination address of query.
         -> IO (Maybe b)                         -- ^ The response, or 'Nothing' if it timed out.
-sendQuery (Client net d err pending whoami _) meth q addr = do
+sendQuery (Client net d err pending whoami _) meth q addr0 = do
     mvar <- newEmptyMVar
-    tid <- atomically $ do
+    (tid,addr,expiry) <- atomically $ do
         tbl <- readTVar pending
-        (tid, tbl') <- dispatchRegister (tableMethods d) mvar tbl
+        (tid, tbl') <- dispatchRegister (tableMethods d) mvar addr0 tbl
+        (addr,expiry) <- methodTimeout meth tid addr0
         writeTVar pending tbl'
-        return tid
+        return (tid,addr,expiry)
     self <- whoami (Just addr)
     sendMessage net addr (wrapQuery meth tid self addr q)
-    mres <- timeout (1000000 * methodTimeout meth) $ takeMVar mvar
+    mres <- timeout expiry $ takeMVar mvar
     case mres of
         Just x -> return $ Just $ unwrapResponse meth x
         Nothing -> do
@@ -248,8 +249,10 @@ dispatchQuery (NoReply unwrapQ f) tid self x addr =
 -- peer-to-peer algorithm will define a 'MethodSerializer' for every 'MethodHandler' that
 -- might be returned by 'lookupHandler'.
 data MethodSerializer tid addr x meth a b = MethodSerializer
-    { -- | Seconds to wait for a response.
-      methodTimeout :: Int
+    { -- | Returns the microseconds to wait for a response to this query being
+      -- sent to the given address.  The /addr/ may also be modified to add
+      -- routing information.
+      methodTimeout :: tid -> addr -> STM (addr,Int)
       -- | A method identifier used for error reporting.  This needn't be the
       -- same as the /meth/ argument to 'MethodHandler', but it is suggested.
     , method :: meth
@@ -269,13 +272,13 @@ data MethodSerializer tid addr x meth a b = MethodSerializer
 --
 -- The type variable /d/ is used to represent the current state of the
 -- transaction generator and the table of pending transactions.
-data TransactionMethods d tid x = TransactionMethods
+data TransactionMethods d tid addr x = TransactionMethods
     {
       -- | Before a query is sent, this function stores an 'MVar' to which the
       -- response will be written too.  The returned /tid/ is a transaction id
       -- that can be used to forget the 'MVar' if the remote peer is not
       -- responding.
-      dispatchRegister :: MVar x -> d -> STM (tid, d)
+      dispatchRegister :: MVar x -> addr -> d -> STM (tid, d)
       -- | This method is invoked when an incoming packet /x/ indicates it is
       -- a response to the transaction with id /tid/.  The returned IO action
       -- is will write the packet to the correct 'MVar' thus completing the
@@ -318,27 +321,15 @@ instance Contravariant (TableMethods t) where
                      (\k t   -> del (f k) t)
                      (\k t   -> lookup (f k) t)
 
--- | Since 'Int' may be 32 or 64 bits, this function is provided as a
--- convenience to test if an integral type, such as 'Data.Word.Word64', can be
--- safely transformed into an 'Int' for use with 'IntMap'.
---
--- Returns 'True' if the proxied type can be losslessly converted to 'Int' using
--- 'fromIntegral'.
-fitsInInt :: forall word. (Bounded word, Integral word) => Proxy word -> Bool
-fitsInInt Proxy = (original == casted)
- where
-    original = div maxBound 2 :: word
-    casted   = fromIntegral (fromIntegral original :: Int) :: word
-
 -- | Construct 'TransactionMethods' methods out of 3 lookup table primitives and a
 -- function for generating unique transaction ids.
 transactionMethods ::
     TableMethods t tid  -- ^ Table methods to lookup values by /tid/.
     -> (g -> (tid,g))   -- ^ Generate a new unique /tid/ value and update the generator state /g/.
-    -> TransactionMethods (g,t (MVar x)) tid x
+    -> TransactionMethods (g,t (MVar x)) tid addr x
 transactionMethods (TableMethods insert delete lookup) generate = TransactionMethods
     { dispatchCancel = \tid (g,t) -> return (g, delete tid t)
-    , dispatchRegister = \v (g,t) ->
+    , dispatchRegister = \v _ (g,t) ->
         let (tid,g') = generate g
             t' = insert tid v t
         in return ( tid, (g',t') )
@@ -356,7 +347,7 @@ data DispatchMethods tbl err meth tid addr x = DispatchMethods
       -- | Lookup the handler for a inbound query.
     , lookupHandler :: meth -> Maybe (MethodHandler err tid addr x)
       -- | Methods for handling incoming responses.
-    , tableMethods :: TransactionMethods tbl tid x
+    , tableMethods :: TransactionMethods tbl tid addr x
     }
 
 -- | These methods indicate what should be done upon various conditions.  Write