Bug fix: distinct index selection for trampolines.

1 files changed, 17 insertions, 9 deletions

diff --git a/OnionRouter.hs b/OnionRouter.hs
index f1b0b41c..ad6323fe 100644
--- a/OnionRouter.hs
+++ b/OnionRouter.hs
@@ -17,6 +17,8 @@ import Control.Monad
 import Crypto.Random
 import Data.Array.MArray
 import Data.Bits
+import Data.Bool
+import Data.List
 import qualified Data.ByteString     as B
 import qualified Data.HashMap.Strict as HashMap
          ;import Data.HashMap.Strict (HashMap)
@@ -251,6 +253,20 @@ selectTrampolines or = do
             myThreadId >>= flip labelThread ("OnionRouter")
             return ns
 
+choose3 :: (Integral a, DRG drg) => drg -> a -> ([a], drg)
+choose3 drg0 cnt = ([a,b,c], drg)
+ where
+        (a, drg1) = randomR (0,cnt - 1) drg0
+        (b0, drg2) = randomR (0,cnt - 2) drg1
+        (c0, drg ) = randomR (0,cnt - 3) drg2
+        b | b0 < a    = b0
+          | otherwise = b0 + 1
+        [ac,bc] = sort [a,b]
+        c1 | c0 < ac    = c0
+           | otherwise  = c0 + 1
+        c | c1 < bc   = c1
+          | otherwise = c1 + 1
+
 -- 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.
@@ -262,15 +278,7 @@ selectTrampolines' or = do
     cnt <- readTVar (trampolineCount or)
     ts <- readTVar (trampolineNodes or)
     drg0 <- readTVar (onionDRG or)
-    let (a, drg1) = randomR (0,cnt - 1) drg0
-        (b0, drg2) = randomR (0,cnt - 2) drg1
-        (c0, drg ) = randomR (0,cnt - 3) drg2
-        b | b0 < a    = b0
-          | otherwise = b0 + 1
-        c1 | c0 < a     = c0
-           | otherwise  = c0 + 1
-        c | c1 < b    = c1
-          | otherwise = c1 + 1
+    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