1 files changed, 23 insertions, 8 deletions

diff --git a/OnionRouter.hs b/OnionRouter.hs
index 0b0b7900..bbc9ad8f 100644
--- a/OnionRouter.hs
+++ b/OnionRouter.hs
@@ -1,5 +1,6 @@
+{-# LANGUAGE LambdaCase               #-}
 {-# LANGUAGE NondecreasingIndentation #-}
-{-# LANGUAGE LambdaCase #-}
+{-# LANGUAGE RankNTypes               #-}
 module OnionRouter where
 
 import Control.Concurrent.Lifted.Instrument
@@ -15,6 +16,7 @@ import Network.Tox.Onion.Transport
 import qualified Network.Tox.TCP as TCP
 import qualified TCPProber as TCP
 
+import Control.Arrow
 import Control.Concurrent.STM
 import Control.Concurrent.STM.TArray
 import Control.Monad
@@ -64,6 +66,8 @@ data OnionRouter = OnionRouter
     , trampolinesUDP :: TrampolineSet NodeInfo
       -- | A set for TCP relays to use as trampolines when UDP is not available.
     , trampolinesTCP :: TrampolineSet TCP.NodeInfo
+      -- | True when we need to rely on TCP relays because UDP is apparently unavailable.
+    , tcpMode :: TVar Bool
       -- | The pseudo-random generator used to select onion routes.
     , onionDRG        :: TVar ChaChaDRG
       -- | Building onion routes happens in a dedicated thread. See 'forkRouteBuilder'.
@@ -184,6 +188,7 @@ newOnionRouter crypto perror = do
                             tbl
                             (TCP.nodeSearch prober tcp)
                             (fmap (maybe False $ const True) . TCP.tcpPing (TCP.tcpClient tcp))
+        tcpmode <- newTVar False
         let o = OnionRouter
                 { pendingRoutes  = pr
                 , onionDRG       = drg
@@ -201,6 +206,7 @@ newOnionRouter crypto perror = do
                     , setNodeClass = nodeClass . TCP.udpNodeInfo
                     , setIDs       = tti
                     }
+                , tcpMode = tcpmode
                 , tcpKademliaClient = tcp
                     { TCP.tcpClient =
                         let c = TCP.tcpClient tcp
@@ -322,20 +328,27 @@ 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 :: Show ni => OnionRouter -> TrampolineSet ni -> IO [ni]
-selectTrampolines or tset = do
+selectTrampolines :: OnionRouter -> IO (Either [TCP.NodeInfo] [NodeInfo])
+selectTrampolines or = do
     myThreadId >>= flip labelThread ("OnionRouter.selectTrampolines")
-    atomically (selectTrampolines' (onionDRG or) tset) >>= \case
+    let tset :: (forall x. TrampolineSet x -> STM (Either [x] [x]))
+                -> STM (Either (Either [TCP.NodeInfo] [NodeInfo])
+                               (Either [TCP.NodeInfo] [NodeInfo]))
+        tset f = bool (left Right . right Right <$> f (trampolinesUDP or))
+                      (left Left  . right Left  <$> f (trampolinesTCP or))
+                      =<< readTVar (tcpMode or)
+    atomically (tset $ selectTrampolines' (onionDRG or)) >>= \case
         Left ns -> do
             -- atomically $ writeTChan (routeLog or)
             routeLogger or $ unwords
-                ( "ONION Discarding insecure trampolines:" : (map show ns) )
+                ( "ONION Discarding insecure trampolines:" : (either (map show) (map show) ns))
             myThreadId >>= flip labelThread ("OnionRouter.selectTrampolines.sleep")
             case ns of
-                [_,_,_] -> threadDelay 1000000 -- wait 1 second if we failed the distinct3by predicate.
+                Left  [_,_,_] -> threadDelay 1000000 -- wait 1 second if we failed the distinct3by predicate.
+                Right [_,_,_] -> 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 tset
+            selectTrampolines or
         Right ns -> do
             myThreadId >>= flip labelThread ("OnionRouter")
             return ns
@@ -397,7 +410,7 @@ handleEvent :: (NodeId -> NodeInfo -> IO (Maybe [NodeInfo])) -> OnionRouter -> R
 handleEvent getnodes or e@(BuildRoute (RouteId rid)) = do
     routeLogger or $ "ONION Rebuilding RouteId " ++ show rid
     mb <- do
-        ts <- selectTrampolines or (trampolinesUDP or)
+        mts <- selectTrampolines or
         join . atomically $ do
         drg <- readTVar (onionDRG or)
         av <- newTVar Nothing
@@ -416,6 +429,7 @@ handleEvent getnodes or e@(BuildRoute (RouteId rid)) = do
                     csel = shiftR sel 4 .&. 0x3
                     sendq s q ni = fmap (listToMaybe . drop (fromIntegral s)) <$> getnodes q ni
                     sendqs = do
+                        let Right ts = mts
                         forkIO $ sendq asel aq (ts !! 0) >>= atomically . writeTVar av . Just
                         forkIO $ sendq bsel bq (ts !! 1) >>= atomically . writeTVar bv . Just
                         forkIO $ sendq csel cq (ts !! 2) >>= atomically . writeTVar cv . Just
@@ -433,6 +447,7 @@ handleEvent getnodes or e@(BuildRoute (RouteId rid)) = do
                               tm
                 return $ do
                     myThreadId >>= flip labelThread ("OnionRouter.sendqs")
+                    let Right ts = mts
                     nodes <- case ts of
                         [_,_,_] -> sendqs
                         _       -> return []