To encode a query, we need the transaction id.

1 files changed, 13 insertions, 11 deletions

diff --git a/src/Network/QueryResponse.hs b/src/Network/QueryResponse.hs
index 58db3a71..b6fea4a9 100644
--- a/src/Network/QueryResponse.hs
+++ b/src/Network/QueryResponse.hs
@@ -97,17 +97,18 @@ data Client err tbl meth tid addr x ctx = Client
 
 -- | An incomming message can be classified into three cases.
 data MessageClass err meth tid
-    = IsQuery meth   -- ^ An unsolicited query is handled based on it's /meth/ value.
-    | IsResponse tid -- ^ A response to a outgoing query we associated with a /tid/ value.
-    | IsUnknown err  -- ^ None of the above.
+    = IsQuery meth tid -- ^ An unsolicited query is handled based on it's /meth/ value.  Any response
+                       -- should include the provided /tid/ value.
+    | IsResponse tid   -- ^ A response to a outgoing query we associated with a /tid/ value.
+    | IsUnknown err    -- ^ None of the above.
 
 -- | Handler for an inbound query of type _x_ from an address of type _addr_.
-data MethodHandler err addr x ctx = forall a b. MethodHandler
+data MethodHandler err tid addr x ctx = forall a b. MethodHandler
     { -- | Parse the query into a more specific type for this method.
       methodParse :: x -> Either err a
       -- | Serialize the response for transmission, given a context /ctx/ and the origin
       -- and destination addresses.
-    , methodSerialize :: ctx -> addr -> addr -> b -> x
+    , methodSerialize :: ctx -> tid -> addr -> addr -> b -> x
       -- | Fully typed action to perform upon the query.  The remote origin
       -- address of the query is provided to the handler.
     , methodAction :: addr -> a -> IO b
@@ -116,14 +117,15 @@ data MethodHandler err addr x ctx = forall a b. MethodHandler
 -- | Attempt to invoke a 'MethodHandler' upon a given inbound query.  If the
 -- parse is successful, the returned IO action will construct our reply.
 -- Otherwise, a parse err is returned.
-dispatchQuery :: MethodHandler err addr x ctx -- ^ Handler to invoke.
+dispatchQuery :: MethodHandler err tid addr x ctx -- ^ Handler to invoke.
                  -> ctx                       -- ^ Arbitrary context used during serialization.
+                 -> tid                       -- ^ The transaction id for this query\/response session.
                  -> addr                      -- ^ Our own address, to which the query was sent.
                  -> x                         -- ^ The query packet.
                  -> addr                      -- ^ The origin address of the query.
                  -> Either err (IO x)
-dispatchQuery (MethodHandler unwrapQ wrapR f) ctx self x addr =
-    fmap (\a -> wrapR ctx self addr <$> f addr a) $ unwrapQ x
+dispatchQuery (MethodHandler unwrapQ wrapR f) ctx tid self x addr =
+    fmap (\a -> wrapR ctx tid self addr <$> f addr a) $ unwrapQ x
 
 -- | These four parameters are required to implement an ougoing query.  A
 -- peer-to-peer algorithm will define a 'MethodSerializer' for every 'MethodHandler' that
@@ -234,7 +236,7 @@ data DispatchMethods tbl err meth tid addr x ctx = DispatchMethods
     { -- | Clasify an inbound packet as a query or response.
       classifyInbound :: x -> MessageClass err meth tid
       -- | Lookup the handler for a inbound query.
-    , lookupHandler :: meth -> Maybe (MethodHandler err addr x ctx)
+    , lookupHandler :: meth -> Maybe (MethodHandler err tid addr x ctx)
       -- | Methods for handling incomming responses.
     , tableMethods :: TableMethods tbl tid x
     }
@@ -276,13 +278,13 @@ handleMessage (Client net d err pending whoami) again = do
                                      again
     Just (Right (plain, addr)) -> do
         case classifyInbound d plain of
-            IsQuery meth -> case lookupHandler d meth of
+            IsQuery meth tid -> case lookupHandler d meth of
                                 Nothing -> reportMissingHandler err meth addr plain
                                 Just m  -> do
                                   (self,ctx) <- whoami
                                   either (reportParseError err)
                                          (>>= sendMessage net addr)
-                                         (dispatchQuery m ctx self plain addr)
+                                         (dispatchQuery m ctx tid self plain addr)
             IsResponse tid -> do
                 action <- atomically $ do
                     ts0 <- readTVar pending