1 files changed, 43 insertions, 14 deletions

diff --git a/Tox.hs b/Tox.hs
index 34e5d6f3..8a4fccb5 100644
--- a/Tox.hs
+++ b/Tox.hs
@@ -11,6 +11,7 @@
 module Tox where
 
 import Control.Arrow
+import Control.Concurrent (MVar)
 import Control.Concurrent.STM
 import qualified Crypto.Cipher.Salsa    as Salsa
 import qualified Crypto.Cipher.XSalsa   as XSalsa
@@ -368,8 +369,9 @@ data Routing = Routing
     , committee6  :: TriadCommittee NodeId SockAddr
     }
 
+type ToxClient = Client String Method TransactionId NodeInfo (Message ByteString)
 
-newClient :: SockAddr -> IO (Client String Method TransactionId NodeInfo (Message ByteString))
+newClient :: SockAddr -> IO (ToxClient, Routing)
 newClient addr = do
     udp <- udpTransport addr
     secret <- generateSecretKey
@@ -402,9 +404,31 @@ newClient addr = do
         sched4 <- newTVar Int.empty
         sched6 <- newTVar Int.empty
         return $ Routing tentative_info sched4 tbl4 committee4 sched6 tbl6 committee6
-    let net = layerTransport (parsePacket secret cache)
-                             (encodePacket secret cache)
-                             udp
+
+    -- If we have 8-byte keys for IntMap, then use it for transaction lookups.
+    -- Otherwise, use ordinary Map.  The details of which will be hidden by an
+    -- existential closure (see mkclient below).
+    tblvar <-
+      if fitsInInt (Proxy :: Proxy Word64)
+        then do
+            let intmapT = transactionMethods (contramapT intKey intMapMethods) gen
+            intmap_var <- atomically $ newTVar (drg, mempty)
+            return $ Right (intmapT,intmap_var)
+         else do
+            let mapT = transactionMethods (contramapT nonceKey mapMethods) gen
+            map_var <- atomically $ newTVar (drg, mempty)
+            return $ Left (mapT,map_var)
+    let net = onInbound (updateRouting outgoingClient routing)
+              $ layerTransport (parsePacket secret cache)
+                               (encodePacket secret cache)
+              $ udp
+
+        -- Paranoid: It's safe to define /net/ and /client/ to be mutually
+        -- recursive since 'updateRouting' does not invoke 'awaitMessage' which
+        -- which was modified by 'onInbound'.  However, I'm going to avoid the
+        -- mutual reference just to be safe.
+        outgoingClient = client { clientNet = net { awaitMessage = return Nothing } }
+
         dispatch tbl = DispatchMethods
             { classifyInbound = classify
             , lookupHandler = handlers
@@ -421,7 +445,16 @@ newClient addr = do
             let (bs, g') = randomBytesGenerate 24 g
             writeTVar var (g',pending)
             return $ TransactionId nonce8 (Nonce24 bs)
-        client tbl var = Client
+
+        client = either mkclient mkclient tblvar
+
+        mkclient :: DRG g =>
+                    ( TransactionMethods (g,t (MVar (Message ByteString)))
+                                         TransactionId
+                                         (Message ByteString)
+                    , TVar (g, t (MVar (Message ByteString)))
+                    ) -> ToxClient
+        mkclient (tbl,var) = Client
             { clientNet           = net
             , clientDispatcher    = dispatch tbl
             , clientErrorReporter = printErrors stderr
@@ -433,15 +466,8 @@ newClient addr = do
                 R.thisNode <$> readTVar var
             , clientResponseId    = genNonce24 var
             }
-    if fitsInInt (Proxy :: Proxy Word64)
-      then do
-        let intmapT = transactionMethods (contramapT intKey intMapMethods) gen
-        intmap_var <- atomically $ newTVar (drg, mempty)
-        return (client intmapT intmap_var)
-      else do
-        let mapT = transactionMethods (contramapT nonceKey mapMethods) gen
-        map_var <- atomically $ newTVar (drg, mempty)
-        return (client mapT map_var)
+
+    return (client, routing)
 
 last8 :: ByteString -> Nonce8
 last8 bs
@@ -482,6 +508,9 @@ type Handler = MethodHandler String TransactionId NodeInfo (Message ByteString)
 
 handler typ f = Just $ MethodHandler decodePayload (encodePayload typ) f
 
+updateRouting :: ToxClient -> Routing -> NodeInfo -> Message ByteString -> IO ()
+updateRouting = error "todo updateRouting"
+
 data Ping = Ping deriving Show
 data Pong = Pong deriving Show