1 files changed, 44 insertions, 39 deletions

diff --git a/dht/src/Network/Tox/Onion/Routes.hs b/dht/src/Network/Tox/Onion/Routes.hs
index 9ce4e316..2f13a513 100644
--- a/dht/src/Network/Tox/Onion/Routes.hs
+++ b/dht/src/Network/Tox/Onion/Routes.hs
@@ -88,7 +88,7 @@ data OnionRouter = OnionRouter
     , tcpProber :: TCP.TCPProber
     , tcpProberThread :: ThreadId
       -- | Kademlia table of TCP relays.
-    , tcpBucketRefresher :: BucketRefresher NodeId TCP.NodeInfo
+    , tcpBucketRefresher :: BucketRefresher NodeId TCP.NodeInfo Nonce8
     , tcpRelayPinger :: RelayPinger
       -- | Debug prints are written to this channel which is then flushed to
       -- 'routeLogger'.
@@ -601,44 +601,49 @@ hookQueries or t8 tmethods = TransactionMethods
         modifyTVar' (pendingQueries or) (W64.insert w8 pq)
         writeTChan (routeLog or) $ "ONION query add " ++ unwords [ show (Just $ pendingVersion pq,w8), ":=", show ni ]
         return (tid,d')
-    , dispatchResponse = \tid x d -> {-# SCC "hookQ.dispatchResponse" #-} do -- :: tid -> x -> d -> STM (d, IO ())
-        let Nonce8 w8 = t8 tid
-        mb <- W64.lookup w8 <$> readTVar (pendingQueries or)
-        modifyTVar' (pendingQueries or) (W64.delete w8)
-        forM_ mb $ \pq -> do
-            let od = pendingDestination pq
-                RouteId rid = fromMaybe (routeId (nodeId (onionNodeInfo od)))
-                                $ onionRouteSpec od
-            modifyArray (routeMap or) (fmap gotResponse) rid
-        writeTChan (routeLog or) $ "ONION query del " ++ show (fmap pendingVersion mb, w8)
-        dispatchResponse tmethods tid x d
-    , dispatchCancel = \tid d -> {-# SCC "hookQ.dispatchCancel" #-} do -- :: tid -> d -> STM d
-        let Nonce8 w8 = t8 tid
-        mb <- W64.lookup w8 <$> readTVar (pendingQueries or)
-        modifyTVar' (pendingQueries or) (W64.delete w8)
-        forM_ mb $ \pq -> do
-            let od = pendingDestination pq
-                RouteId rid = fromMaybe (routeId (nodeId (onionNodeInfo od)))
-                                $ onionRouteSpec od
-            mrr <- readArray (routeMap or) rid
-            forM_ mrr $ \rr -> do
-            when (routeVersion rr == pendingVersion pq) $ do
-            let expireRoute = modifyArray (pendingRoutes or) expire rid
-                expire ver | ver <= succ (pendingVersion pq) = succ (pendingVersion pq)
-                           | otherwise                       = ver
-            modifyArray (routeMap or) (fmap gotTimeout) rid
-            case rr of
-                RouteRecord{ responseCount = 0
-                           , timeoutCount  = c
-                           , routeVersion = v } | c >= 5 -> expireRoute
-                RouteRecord{ responseCount = 1
-                           , timeoutCount = c
-                           , routeVersion = v } | c >= 10 -> expireRoute
-                RouteRecord{ timeoutCount = c
-                           , routeVersion = v } | c >= 20 -> expireRoute
-                _ -> return ()
-        writeTChan (routeLog or) $ "ONION query can " ++ show (fmap pendingVersion mb, w8)
-        dispatchCancel tmethods tid d
+    , dispatchResponse = \tid rx d -> {-# SCC "hookQ.dispatchResponse" #-} do -- :: tid -> x -> d -> STM (d, IO ())
+        case rx of
+            Success x -> do
+                let Nonce8 w8 = t8 tid
+                mb <- W64.lookup w8 <$> readTVar (pendingQueries or)
+                modifyTVar' (pendingQueries or) (W64.delete w8)
+                forM_ mb $ \pq -> do
+                    let od = pendingDestination pq
+                        RouteId rid = fromMaybe (routeId (nodeId (onionNodeInfo od)))
+                                        $ onionRouteSpec od
+                    modifyArray (routeMap or) (fmap gotResponse) rid
+                writeTChan (routeLog or) $ "ONION query del " ++ show (fmap pendingVersion mb, w8)
+                dispatchResponse tmethods tid rx d
+            _ -> do -- Timed out or canceled...
+                let Nonce8 w8 = t8 tid
+                mb <- W64.lookup w8 <$> readTVar (pendingQueries or)
+                modifyTVar' (pendingQueries or) (W64.delete w8)
+                forM_ mb $ \pq -> do
+                    let od = pendingDestination pq
+                        RouteId rid = fromMaybe (routeId (nodeId (onionNodeInfo od)))
+                                        $ onionRouteSpec od
+                    mrr <- readArray (routeMap or) rid
+                    forM_ mrr $ \rr -> do
+                    when (routeVersion rr == pendingVersion pq) $ do
+                    let expireRoute = modifyArray (pendingRoutes or) expire rid
+                        expire ver | ver <= succ (pendingVersion pq) = succ (pendingVersion pq)
+                                   | otherwise                       = ver
+                    case rx of
+                        TimedOut -> do
+                            modifyArray (routeMap or) (fmap gotTimeout) rid
+                            case rr of
+                                RouteRecord{ responseCount = 0
+                                           , timeoutCount  = c
+                                           , routeVersion = v } | c >= 5 -> expireRoute
+                                RouteRecord{ responseCount = 1
+                                           , timeoutCount = c
+                                           , routeVersion = v } | c >= 10 -> expireRoute
+                                RouteRecord{ timeoutCount = c
+                                           , routeVersion = v } | c >= 20 -> expireRoute
+                                _ -> return ()
+                        _ -> return () -- Don't penalize route for canceled queries.
+                writeTChan (routeLog or) $ "ONION query can " ++ show (fmap pendingVersion mb, w8)
+                dispatchResponse tmethods tid rx d
     }