1 files changed, 19 insertions, 9 deletions

diff --git a/dht/src/Network/Tox/Onion/Routes.hs b/dht/src/Network/Tox/Onion/Routes.hs
index c054b99e..d61c721e 100644
--- a/dht/src/Network/Tox/Onion/Routes.hs
+++ b/dht/src/Network/Tox/Onion/Routes.hs
@@ -14,6 +14,7 @@ import Network.QueryResponse
 import Network.QueryResponse.TCP
 import Network.Tox.NodeId
 import Network.Tox.Onion.Transport as Onion
+import Network.Tox.RelayPinger
 import qualified Data.Tox.Relay as TCP
 import qualified Network.Tox.TCP as TCP
 import qualified TCPProber as TCP
@@ -88,6 +89,7 @@ data OnionRouter = OnionRouter
     , tcpProberThread :: ThreadId
       -- | Kademlia table of TCP relays.
     , tcpBucketRefresher :: BucketRefresher NodeId TCP.NodeInfo
+    , tcpRelayPinger :: RelayPinger
       -- | Debug prints are written to this channel which is then flushed to
       -- 'routeLogger'.
     , routeLog        :: TChan String
@@ -207,7 +209,8 @@ newOnionRouter crypto perror tcp_enabled = do
                 writeTChan rlog $ unwords ["Selected TCP Gateway:",show ni,"via",show gw]
                 return gw
             }
-    or <- atomically $ do
+
+    or0 <- atomically $ do
         -- chan <- newTChan
         drg <- newTVar drg0
         -- forM_ [0..11] $ \n -> writeTChan chan $ BuildRoute (RouteId n)
@@ -243,12 +246,7 @@ newOnionRouter crypto perror tcp_enabled = do
                     }
                 , tcpMode = tcpmode
                 , tcpKademliaClient = tcp
-                    { TCP.tcpClient =
-                        let c = TCP.tcpClient tcp
-                        in c { clientNet = addHandler perror (handleMessage c)
-                                            $ onInbound (updateTCP o)
-                                            $ clientNet c }
-                    }
+                , tcpRelayPinger     = error "forkRouteBuilder: no RelayPinger specified"
                 , tcpBucketRefresher = refresher
                 , routeLog           = rlog
                 , routeThread        = error "forkRouteBuilder not invoked (missing onion route builder thread)."
@@ -259,10 +257,21 @@ newOnionRouter crypto perror tcp_enabled = do
                 , routeLogger        = perror
                 }
         return o
+    pinger <- forkRelayPinger (kademSpace $ refreshKademlia $ tcpBucketRefresher or0)
+                              (TCP.tcpClient tcp)
+    let or = or0 { tcpRelayPinger    = pinger
+                 , tcpKademliaClient = tcp
+                    { TCP.tcpClient =
+                        let c = TCP.tcpClient tcp
+                        in c { clientNet = addHandler perror (handleMessage c)
+                                            $ onInbound (updateTCP or)
+                                            $ clientNet c }
+                    }
+                 }
     return (or,tcptbl)
 
 updateTCP :: OnionRouter -> TCP.NodeInfo -> p -> IO ()
-updateTCP or addr x =
+updateTCP or addr x = do
     let refresher = tcpBucketRefresher or
         kademlia0 = refreshKademlia refresher
         kademlia = kademlia0 { kademIO = (kademIO kademlia0)
@@ -275,7 +284,8 @@ updateTCP or addr x =
                                     tblTransition (kademIO kademlia0) tr
                                 }
                              }
-    in insertNode kademlia addr
+    atomically $ bumpRelay (tcpRelayPinger or) addr
+    insertNode kademlia addr
 
 selectGateway :: TVar (R.BucketList TCP.NodeInfo) -> NodeInfo -> STM (Maybe TCP.NodeInfo)
 selectGateway tbl ni = do