Generalized selectTrampolines

1 files changed, 25 insertions, 13 deletions

diff --git a/OnionRouter.hs b/OnionRouter.hs
index 39b3a872..bfd517fd 100644
--- a/OnionRouter.hs
+++ b/OnionRouter.hs
@@ -291,24 +291,31 @@ randomIvalInteger (l,h) rng
 --
 -- 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
+selectTrampolines :: Show ni => OnionRouter -> TrampolineSet ni -> IO [ni]
+selectTrampolines or tset = do
     myThreadId >>= flip labelThread ("OnionRouter.selectTrampolines")
-    atomically (selectTrampolines' or) >>= \case
+    atomically (selectTrampolines' tset) >>= \case
         Left ns -> do
             -- atomically $ writeTChan (routeLog or)
             routeLogger or $ unwords
-                ( "ONION Discarding insecure trampolines:" : (map (show . nodeAddr) ns) )
+                ( "ONION Discarding insecure trampolines:" : (map show ns) )
             myThreadId >>= flip labelThread ("OnionRouter.selectTrampolines.sleep")
             case ns of
                 [_,_,_] -> threadDelay 1000000 -- wait 1 second if we failed the distinct3by predicate.
                 _       -> threadDelay 5000000 -- wait 5 seconds if insufficient nodes.
             myThreadId >>= flip labelThread ("OnionRouter.selectTrampolines")
-            selectTrampolines or
+            selectTrampolines or tset
         Right ns -> do
             myThreadId >>= flip labelThread ("OnionRouter")
             return ns
 
+data TrampolineSet ni = TrampolineSet
+    { setNodes     :: TVar (IntMap ni)
+    , setCount     :: TVar Int
+    , setDRG       :: TVar ChaChaDRG
+    , setNodeClass :: ni -> IPClass
+    }
+
 choose3 :: (Integral a, DRG drg) => drg -> a -> ([a], drg)
 choose3 drg0 cnt = ([a,b,c], drg)
  where
@@ -329,25 +336,30 @@ choose3 drg0 cnt = ([a,b,c], drg)
 --
 -- 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)
-    ts <- readTVar (trampolineNodes or)
-    drg0 <- readTVar (onionDRG or)
+selectTrampolines' :: TrampolineSet ni -> STM (Either [ni] [ni])
+selectTrampolines' TrampolineSet{..} = do
+    cnt <- readTVar setCount
+    ts <- readTVar setNodes
+    drg0 <- readTVar setDRG
     let ([a,b,c],drg) = choose3 drg0 cnt
         ns = mapMaybe (\n -> IntMap.lookup n ts) [a,b,c]
     ns' <- case ns of
-            [an,bn,cn] | distinct3by nodeClass an bn cn
+            [an,bn,cn] | distinct3by setNodeClass an bn cn
                        -> return $ Right ns
             _          -> return $ Left ns
-    writeTVar (onionDRG or) drg
+    writeTVar setDRG drg
     return ns'
 
 handleEvent :: (NodeId -> NodeInfo -> IO (Maybe [NodeInfo])) -> OnionRouter -> RouteEvent -> IO ()
 handleEvent getnodes or e@(BuildRoute (RouteId rid)) = do
     routeLogger or $ "ONION Rebuilding RouteId " ++ show rid
     mb <- do
-        ts <- selectTrampolines or
+        ts <- selectTrampolines or TrampolineSet
+                    { setNodes     = trampolineNodes or
+                    , setCount     = trampolineCount or
+                    , setDRG       = onionDRG or
+                    , setNodeClass = nodeClass
+                    }
         join . atomically $ do
         drg <- readTVar (onionDRG or)
         av <- newTVar Nothing