1 files changed, 18 insertions, 12 deletions

diff --git a/src/Network/Tox.hs b/src/Network/Tox.hs
index 35eaebb5..15b00780 100644
--- a/src/Network/Tox.hs
+++ b/src/Network/Tox.hs
@@ -167,11 +167,15 @@ w64Key (DHT.TransactionId (Nonce8 w) _) = w
 nonceKey :: DHT.TransactionId -> Nonce8
 nonceKey (DHT.TransactionId n _) = n
 
-myAddr :: DHT.Routing -> Maybe NodeInfo -> IO NodeInfo
-myAddr routing maddr = atomically $ do
+-- | Return my own address.
+myAddr :: TVar (R.BucketList NodeInfo) -- ^ IPv4 buckets
+       -> TVar (R.BucketList NodeInfo) -- ^ IPv6 buckets
+       -> Maybe NodeInfo             -- ^ Interested remote address
+       -> IO NodeInfo
+myAddr routing4 routing6 maddr = atomically $ do
     let var = case flip DHT.prefer4or6 Nothing <$> maddr of
-                Just Want_IP6 -> DHT.routing6 routing
-                _             -> DHT.routing4 routing
+                Just Want_IP6 -> routing4
+                _             -> routing6
     a <- readTVar var
     return $ R.thisNode a
 
@@ -179,7 +183,7 @@ newClient :: (DRG g, Show addr, Show meth) =>
               g -> Transport String addr x
                 -> (Client String meth DHT.TransactionId addr x -> x -> MessageClass String meth DHT.TransactionId addr x)
                 -> (Maybe addr -> IO addr)
-                -> (meth -> Maybe (MethodHandler String DHT.TransactionId addr x))
+                -> (Client String meth DHT.TransactionId addr x -> meth -> Maybe (MethodHandler String DHT.TransactionId addr x))
                 -> (forall d. TransactionMethods d DHT.TransactionId addr x -> TransactionMethods d DHT.TransactionId addr x)
                 -> (Client String meth DHT.TransactionId addr x -> Transport String addr x -> Transport String addr x)
                 -> IO (Client String meth DHT.TransactionId addr x)
@@ -207,7 +211,7 @@ newClient drg net classify selfAddr handlers modifytbl modifynet = do
         mkclient (tbl,var) handlers =
             let client = Client
                     { clientNet           = addHandler (reportParseError eprinter) (handleMessage client) $ modifynet client net
-                    , clientDispatcher    = dispatch tbl var handlers client
+                    , clientDispatcher    = dispatch tbl var (handlers client) client
                     , clientErrorReporter = eprinter { reportTimeout = reportTimeout ignoreErrors }
                     , clientPending       = var
                     , clientAddress       = selfAddr
@@ -308,13 +312,15 @@ newTox keydb addr mbSessionsState suppliedDHTKey = do
     drg <- drgNew
     let lookupClose _ = return Nothing
 
-    routing <- DHT.newRouting addr crypto updateIP updateIP
+    mkrouting <- DHT.newRouting addr crypto updateIP updateIP
     let ignoreErrors _ = return () -- Set this to (hPutStrLn stderr) to debug onion route building.
     orouter <- newOnionRouter ignoreErrors
     (dhtcrypt,onioncrypt,dtacrypt,cryptonet) <- toxTransport crypto orouter lookupClose udp
     let dhtnet0 = layerTransportM (DHT.decrypt crypto) (DHT.encrypt crypto) dhtcrypt
-    dhtclient <- newClient drg dhtnet0 DHT.classify (myAddr routing) (DHT.handlers crypto routing) id
-                    $ \client net -> onInbound (DHT.updateRouting client routing orouter) net
+        tbl4 = DHT.routing4 $ mkrouting (error "missing client")
+        tbl6 = DHT.routing6 $ mkrouting (error "missing client")
+    dhtclient <- newClient drg dhtnet0 DHT.classify (myAddr tbl4 tbl6) (DHT.handlers crypto . mkrouting) id
+                    $ \client net -> onInbound (DHT.updateRouting client (mkrouting client) orouter) net
 
     orouter <- forkRouteBuilder orouter $ \nid ni -> fmap (\(_,ns,_)->ns) <$> DHT.getNodes dhtclient nid ni
 
@@ -324,8 +330,8 @@ newTox keydb addr mbSessionsState suppliedDHTKey = do
     oniondrg <- drgNew
     let onionnet = layerTransportM (Onion.decrypt crypto) (Onion.encrypt crypto) onioncrypt
     onionclient <- newClient oniondrg onionnet (const Onion.classify)
-                    (getOnionAlias crypto $ R.thisNode <$> readTVar (DHT.routing4 routing))
-                    (Onion.handlers onionnet routing toks keydb)
+                    (getOnionAlias crypto $ R.thisNode <$> readTVar (DHT.routing4 $ mkrouting dhtclient))
+                    (const $ Onion.handlers onionnet (mkrouting dhtclient) toks keydb)
                     (hookQueries orouter DHT.transactionKey)
                     (const id)
 
@@ -337,7 +343,7 @@ newTox keydb addr mbSessionsState suppliedDHTKey = do
         , toxCrypto = addHandler (hPutStrLn stderr) (cryptoNetHandler sessionsState) cryptonet
         , toxCryptoSessions = sessionsState
         , toxCryptoKeys = crypto
-        , toxRouting = routing
+        , toxRouting = mkrouting dhtclient
         , toxTokens = toks
         , toxAnnouncedKeys = keydb
         , toxOnionRoutes = orouter