1 files changed, 24 insertions, 6 deletions

diff --git a/dht/OnionRouter.hs b/dht/OnionRouter.hs
index bdaf04b2..e6f647b5 100644
--- a/dht/OnionRouter.hs
+++ b/dht/OnionRouter.hs
@@ -70,7 +70,7 @@ data OnionRouter = OnionRouter
       -- | 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
+    , tcpMode :: TVar (Maybe Bool) -- Nothing: tcp disabled, False: use trampolinesUDP, True: use trampolinesTCP
       -- | The pseudo-random generator used to select onion routes.
     , onionDRG        :: TVar ChaChaDRG
       -- | Building onion routes happens in a dedicated thread. See 'forkRouteBuilder'.
@@ -154,11 +154,12 @@ newtype RouteEvent = BuildRoute RouteId
 
 newOnionRouter :: TransportCrypto
                     -> (String -> IO ())
+                    -> Bool -- is tcp enabled?
                     -> IO ( OnionRouter
                           , TVar ( ChaChaDRG
                                  , Word64Map (Either (Maybe (Bool,TCP.RelayPacket) -> IO ())
                                                      (Maybe (OnionMessage Identity) -> IO ()))))
-newOnionRouter crypto perror = do
+newOnionRouter crypto perror tcp_enabled = do
     drg0 <- drgNew
     (rlog,pq,rm) <- atomically $ do
         rlog <- newTChan
@@ -222,7 +223,7 @@ newOnionRouter crypto perror = do
                             tbl
                             (TCP.nodeSearch prober tcp)
                             (fmap (maybe False $ const True) . TCP.tcpPing (TCP.tcpClient tcp))
-        tcpmode <- newTVar True
+        tcpmode <- newTVar $ if tcp_enabled then Just True else Nothing
         let o = OnionRouter
                 { pendingRoutes  = pr
                 , onionDRG       = drg
@@ -369,9 +370,12 @@ selectTrampolines or = do
     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)
+        tset f = do
+            mm <- readTVar (tcpMode or)
+            -- TODO: better logic for deciding to use TCP or UDP trampolines.
+            if fromMaybe False mm
+                then left Left  . right Left  <$> f (trampolinesTCP or)
+                else left Right . right Right <$> f (trampolinesUDP or)
     atomically (tset $ internalSelectTrampolines (onionDRG or)) >>= \case
         Left ns -> do
             -- atomically $ writeTChan (routeLog or)
@@ -710,3 +714,17 @@ ipClass' :: SockAddr -> IPClass
 ipClass' (SockAddrInet _ addr)                = IPClass (fromBE32 addr .&. 0xFF000000)
 ipClass' (SockAddrInet6 _ _ (hiword,_,_,_) _) = IPClass hiword
 ipClass' _                                    = IPClass 0 -- unreachable.
+
+requestTCPMode :: OnionRouter -> Maybe Bool -> IO Bool
+requestTCPMode or wanted_mode = atomically $ requestTCPModeSTM or wanted_mode
+
+requestTCPModeSTM :: OnionRouter -> Maybe Bool -> STM Bool
+requestTCPModeSTM or wanted_mode = do
+    m <- readTVar (tcpMode or)
+    case m of
+        Nothing -> return False
+        Just oldmode -> case wanted_mode of
+            Just newmode -> do
+                writeTVar (tcpMode or) (Just newmode)
+                return newmode
+            Nothing -> return oldmode