2 files changed, 22 insertions, 5 deletions

diff --git a/OnionRouter.hs b/OnionRouter.hs
index 20279c5d..428b6d8e 100644
--- a/OnionRouter.hs
+++ b/OnionRouter.hs
@@ -210,6 +210,11 @@ randomIvalInteger (l,h) rng
                         (x,g') = next g -- next :: RandomGen g => g -> (Int, g)
                         v' = (v * b + (fromIntegral x - fromIntegral genlo))
 
+-- Repeatedly attempt to select 3 nodes as a secure onion route letting 1 second
+-- elapse between retries.
+--
+-- Only the DRG random seed is updated.  Hopefully another thread will change the
+-- trampolineNodes set so that selection can succeed.
 selectTrampolines :: OnionRouter -> IO [NodeInfo]
 selectTrampolines or = do
     myThreadId >>= flip labelThread ("OnionRouter.selectTrampolines")
@@ -226,6 +231,12 @@ selectTrampolines or = do
             myThreadId >>= flip labelThread ("OnionRouter")
             return ns
 
+-- Select 3 indices into the trampolineNodes set and returns the associated
+-- nodes provided they are suitable for use in an onion route.  Otherwise, it
+-- returns Left with the nodes that were selected.
+--
+-- The only write this function does to STM state is that the onionDRG random
+-- seed will be updated.
 selectTrampolines' :: OnionRouter -> STM (Either [NodeInfo] [NodeInfo])
 selectTrampolines' or = do
     cnt <- readTVar (trampolineCount or)
@@ -425,10 +436,14 @@ hookBucketList _ _ or (RoutingTransition ni Stranger) = do
         Just n -> do writeTVar (trampolineIds or) (HashMap.delete (nodeId ni) ns)
                      cnt <- pred <$> readTVar (trampolineCount or)
                      writeTVar (trampolineCount or) cnt
-                     if n == cnt
-                        then modifyTVar' (trampolineNodes or) (IntMap.delete n)
-                        else do lastnode <- (IntMap.! cnt) <$> readTVar (trampolineNodes or)
-                                modifyTVar' (trampolineNodes or) (IntMap.insert n lastnode . IntMap.delete cnt)
+                     case compare n cnt of
+                        EQ -> modifyTVar' (trampolineNodes or) (IntMap.delete n)
+                        LT -> do lastnode <- (IntMap.! cnt) <$> readTVar (trampolineNodes or)
+                                 modifyTVar' (trampolineNodes or)
+                                             (IntMap.insert n lastnode . IntMap.delete cnt)
+                                 modifyTVar' (trampolineIds or)
+                                             (HashMap.delete (nodeId ni) . HashMap.insert (nodeId lastnode) n)
+                        GT -> writeTChan (routeLog or) $ "BUG!! Trampoline maps are out of sync."
                      writeTChan (routeLog or) $ "ONION trampoline Stranger " ++ unwords [show n,show ni]
         Nothing -> return ()
 hookBucketList _ _ _ _ = return () -- ignore Applicant event.
diff --git a/examples/dhtd.hs b/examples/dhtd.hs
index 369650f9..fce976db 100644
--- a/examples/dhtd.hs
+++ b/examples/dhtd.hs
@@ -808,6 +808,8 @@ clientSession s@Session{..} sock cnum h = do
         ("onion", s) -> cmd0 $ join $ atomically $ do
             rm <- readTVar $ routeMap onionRouter
             ts <- readTVar $ trampolineNodes onionRouter
+            tcnt <- readTVar $ trampolineCount onionRouter
+            icnt <- HashMap.size <$> readTVar (trampolineIds onionRouter)
             rs <- mapM readTVar (pendingRoutes onionRouter)
             let showRecord :: Int -> Bool -> [String]
                 showRecord n True = [show n, "pending", ""]
@@ -817,7 +819,7 @@ clientSession s@Session{..} sock cnum h = do
                     | otherwise = [show n, "error!",""]
                 r = map (uncurry showRecord) $ IntMap.toAscList rs
             return $ do
-                hPutClientChunk h $ "trampolines: " ++ show (IntMap.size ts) ++ "\n"
+                hPutClientChunk h $ "trampolines: " ++ show (IntMap.size ts,tcnt) ++ "\n"
                 hPutClient h $ showColumns $ ["","responses","timeouts"]:r
 
         -- necrypto <FRIEND-TOXID>