1 files changed, 38 insertions, 35 deletions

diff --git a/src/Network/QueryResponse.hs b/src/Network/QueryResponse.hs
index 5b8bcda4..58db3a71 100644
--- a/src/Network/QueryResponse.hs
+++ b/src/Network/QueryResponse.hs
@@ -39,9 +39,9 @@ import Data.Maybe
 --  > r <- sendQuery client method q
 --  > -- Quit client.
 --  > quitServer
-forkListener :: Client err tbl meth tid addr x -> IO (IO ())
+forkListener :: Client err tbl meth tid addr x ctx -> IO (IO ())
 forkListener client = do
-    thread_id <- fork $ do
+    thread_id <- forkIO $ do
         myThreadId >>= flip labelThread "listener"
         fix $ handleMessage client
     return $ do
@@ -52,12 +52,12 @@ forkListener client = do
 -- out if 'forkListener' was never invoked to spawn a thread receive and
 -- dispatch the response.
 sendQuery ::
-    forall err a b tbl x meth tid addr.
-        Client err tbl meth tid addr x  -- ^ A query/response implementation.
-        -> Method addr x meth a b       -- ^ Information for marshalling the query.
-        -> a                            -- ^ The outbound query.
-        -> addr                         -- ^ Destination address of query.
-        -> IO (Maybe b)                 -- ^ The response, or 'Nothing' if it timed out.
+    forall err a b tbl x ctx meth tid addr.
+        Client err tbl meth tid addr x ctx -- ^ A query/response implementation.
+        -> MethodSerializer addr x ctx meth a b      -- ^ Information for marshalling the query.
+        -> 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
     mvar <- newEmptyMVar
     tid <- atomically $ do
@@ -65,8 +65,8 @@ sendQuery (Client net d err pending whoami) meth q addr = do
         let (tid, tbl') = dispatchRegister (tableMethods d) mvar tbl
         writeTVar pending tbl'
         return tid
-    self <- whoami
-    sendMessage net addr (wrapQuery meth self addr q)
+    (self,ctx) <- whoami
+    sendMessage net addr (wrapQuery meth ctx self addr q)
     mres <- timeout (methodTimeout meth) $ takeMVar mvar
     case mres of
         Just x -> return $ Just $ unwrapResponse meth x
@@ -78,20 +78,21 @@ sendQuery (Client net d err pending whoami) meth q addr = do
 -- * Implementing a query\/response 'Client'.
 
 -- | All inputs required to implement a query\/response client.
-data Client err tbl meth tid addr x = Client
+data Client err tbl meth tid addr x ctx = Client
     { -- | The 'Transport' used to dispatch and receive packets.
       clientNet :: Transport err addr x
       -- | Methods for handling inbound packets.
-    , clientDispatcher :: DispatchMethods tbl err meth tid addr x
+    , clientDispatcher :: DispatchMethods tbl err meth tid addr x ctx
       -- | Methods for reporting various conditions.
     , clientErrorReporter :: ErrorReporter addr x meth tid err
       -- | State necessary for routing inbound responses and assigning unique
       -- /tid/ values for outgoing queries.
     , clientPending :: TVar tbl
-      -- | An action yielding this client\'s own address.  It is invoked once on
-      -- each outbound and inbound packet.  It is valid for this to always
-      -- return the same value.
-    , clientMyAddress :: IO addr
+      -- | An action yielding this client\'s own address along with some
+      -- context neccessary for serializing outgoing packets.  It is invoked
+      -- once on each outbound and inbound packet.  It is valid for this to
+      -- always return the same value.
+    , clientContext :: IO (addr,ctx)
     }
 
 -- | An incomming message can be classified into three cases.
@@ -101,12 +102,12 @@ data MessageClass err meth tid
     | 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 = forall a b. MethodHandler
+data MethodHandler err 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 type for transmission.  Origin and destination
-      -- addresses for the packet are supplied in case they are required.
-    , methodSerialize :: addr -> addr -> b -> x
+      -- | Serialize the response for transmission, given a context /ctx/ and the origin
+      -- and destination addresses.
+    , methodSerialize :: ctx -> 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
@@ -115,18 +116,19 @@ data MethodHandler err addr x = 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 -- ^ Handler to invoke.
-                 -> addr                  -- ^ Our own address, to which the query was sent.
-                 -> x                     -- ^ The query packet.
-                 -> addr                  -- ^ The origin address of the query.
+dispatchQuery :: MethodHandler err addr x ctx -- ^ Handler to invoke.
+                 -> ctx                       -- ^ Arbitrary context used during serialization.
+                 -> 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) self x addr =
-    fmap (\a -> wrapR self addr <$> f addr a) $ unwrapQ x
+dispatchQuery (MethodHandler unwrapQ wrapR f) ctx self x addr =
+    fmap (\a -> wrapR ctx 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 'Method' for every 'MethodHandler' that
+-- peer-to-peer algorithm will define a 'MethodSerializer' for every 'MethodHandler' that
 -- might be returned by 'lookupHandler'.
-data Method addr x meth a b = Method
+data MethodSerializer addr x ctx meth a b = MethodSerializer
     { -- | Seconds to wait for a response.
       methodTimeout :: Int
       -- | A method identifier used for error reporting.  This needn't be the
@@ -134,8 +136,9 @@ data Method addr x meth a b = Method
     , method :: meth
       -- | Serialize the outgoing query /a/ into a transmitable packet /x/.
       -- The /addr/ arguments are, respectively, our own origin address and the
-      -- destination of the request.
-    , wrapQuery :: addr -> addr -> a -> x
+      -- destination of the request.  The /ctx/ argument is useful for attaching
+      -- auxillary notations on all outgoing packets.
+    , wrapQuery :: ctx -> addr -> addr -> a -> x
       -- | Parse an inbound packet /x/ into a response /b/ for this query.
     , unwrapResponse :: x -> b
     }
@@ -227,11 +230,11 @@ transactionTableMethods insert delete lookup generate = TableMethods
     }
 
 -- | A set of methods neccessary for dispatching incomming packets.
-data DispatchMethods tbl err meth tid addr x = DispatchMethods
+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)
+    , lookupHandler :: meth -> Maybe (MethodHandler err addr x ctx)
       -- | Methods for handling incomming responses.
     , tableMethods :: TableMethods tbl tid x
     }
@@ -264,7 +267,7 @@ data ErrorReporter addr x meth tid err = ErrorReporter
 -- 'fix' in a forked thread that loops until 'awaitMessage' returns 'Nothing'
 -- or throws an exception.
 handleMessage ::
-    Client err tbl meth tid addr x
+    Client err tbl meth tid addr x ctx
     -> IO ()
     -> IO ()
 handleMessage (Client net d err pending whoami) again = do
@@ -276,10 +279,10 @@ handleMessage (Client net d err pending whoami) again = do
             IsQuery meth -> case lookupHandler d meth of
                                 Nothing -> reportMissingHandler err meth addr plain
                                 Just m  -> do
-                                  self <- whoami
+                                  (self,ctx) <- whoami
                                   either (reportParseError err)
                                          (>>= sendMessage net addr)
-                                         (dispatchQuery m self plain addr)
+                                         (dispatchQuery m ctx self plain addr)
             IsResponse tid -> do
                 action <- atomically $ do
                     ts0 <- readTVar pending