1 files changed, 19 insertions, 4 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.