6 files changed, 222 insertions, 14 deletions

diff --git a/OnionRouter.hs b/OnionRouter.hs
index 55ba9c28..39b3a872 100644
--- a/OnionRouter.hs
+++ b/OnionRouter.hs
@@ -12,6 +12,7 @@ import Network.QueryResponse
 import Network.Tox.NodeId
 import Network.Tox.Onion.Transport
 import qualified Network.Tox.TCP as TCP
+import qualified TCPProber as TCP
 
 import Control.Concurrent.STM
 import Control.Concurrent.STM.TArray
@@ -83,6 +84,8 @@ data OnionRouter = OnionRouter
     , tcpKademliaClient :: TCP.TCPClient String () Nonce8
       -- | This thread maintains the TCP relay table.
     , tcpKademliaThread :: ThreadId
+    , tcpProber :: TCP.TCPProber
+    , tcpProberThread :: ThreadId
       -- | Kademlia table of TCP relays.
     , tcpBucketRefresher :: BucketRefresher NodeId TCP.NodeInfo
       -- | Debug prints are written to this channel which is then flushed to
@@ -181,7 +184,8 @@ newOnionRouter crypto perror = do
         tc <- newTVar 0
         pr <- newArray (0,11) 0
         rlog <- newTChan
-        refresher <- newBucketRefresher tbl (TCP.nodeSearch tcp) (fmap (maybe False $ const True) . TCP.tcpPing (TCP.tcpClient tcp))
+        prober <- TCP.newProber
+        refresher <- newBucketRefresher tbl (TCP.nodeSearch prober tcp) (fmap (maybe False $ const True) . TCP.tcpPing (TCP.tcpClient tcp))
         return OnionRouter
             { pendingRoutes = pr
             , onionDRG = drg
@@ -202,6 +206,8 @@ newOnionRouter crypto perror = do
             , routeLog = rlog
             , routeThread = error "forkRouteBuilder not invoked (missing onion route builder thread)."
             , tcpKademliaThread = error "forkRouteBuilder not invoked (missing TCP bucket maintenance thread)."
+            , tcpProber = prober
+            , tcpProberThread = error "forkRouteBuilder not invoked (missing TCP probe thread)."
             , routeLogger = perror
             }
     return or
diff --git a/TCPProber.hs b/TCPProber.hs
new file mode 100644
index 00000000..8d468e53
--- /dev/null
+++ b/TCPProber.hs
@@ -0,0 +1,162 @@
+{-# LANGUAGE LambdaCase #-}
+module TCPProber where
+
+import Control.Concurrent
+import GHC.Conc
+
+import Control.Arrow
+import Control.Concurrent.STM
+import Control.Monad
+import Data.Function
+import Data.IP
+import Data.Maybe
+import Data.Time.Clock.POSIX
+import Network.Socket
+import System.Timeout
+
+import Crypto.Tox
+import Data.Wrapper.PSQ as PSQ
+import Network.Kademlia.Search
+import Network.Tox.NodeId
+import qualified Network.Tox.TCP as TCP
+
+resolvePort :: TCP.RelayClient -> NodeInfo -> IO (Maybe PortNumber)
+resolvePort tcp ni = do
+    got <- newTVarIO Nothing
+    cnt <- newTVarIO 0
+    let n port = TCP.NodeInfo ni port
+        forkPort port = do
+            atomically $ modifyTVar' cnt succ
+            t <- forkIO $ do
+                m <- TCP.tcpPing tcp $ n port
+                atomically $ do
+                    mg <- readTVar got
+                    when (isNothing mg)
+                        $ forM_ m $ \_ -> writeTVar got $ Just port
+                    modifyTVar' cnt pred
+            labelThread t $ "probe." ++ show port ++ "." ++ show (nodeId ni)
+            return t
+        readResult = atomically $ do
+            m <- readTVar got
+            case m of
+                Just v  -> return $ Just v
+                Nothing -> readTVar cnt >>= check . (== 0) >> return Nothing
+    t443 <- forkPort 443
+    t80 <- forkPort 80
+    p <- timeout 1000000 readResult >>= \case
+            Just (Just p) -> do
+                killThread t443
+                killThread t80
+                return $ Just p
+            _ -> do
+                let uport = nodePort ni
+                tudp <- forM (guard $ uport `notElem` [443,80,3389,33445])
+                            $ \() -> forkPort uport
+                t3389 <- forkPort 3389
+                t33445 <- forkPort 33445
+                p <- readResult
+                mapM_ killThread [t443,t80,t3389,t33445]
+                mapM_ killThread (tudp :: Maybe ThreadId)
+                return p
+    return p
+
+data TCPProber = TCPProber
+    { probeQueue :: TVar (PSQ' NodeId POSIXTime SockAddr{-UDP-})
+    , probeSpill :: TVar (PSQ' NodeId POSIXTime SockAddr{-UDP-})
+    , probeSpillCount :: TVar Int -- Data.HashPSQ has O(n) size, so we keep a count.
+    , probeCache :: TVar (PSQ' NodeId POSIXTime (SockAddr{-UDP-},PortNumber{-TCP-}))
+    , probeCacheCount :: TVar Int
+    }
+
+newProber :: STM TCPProber
+newProber = do
+    q        <- newTVar PSQ.empty
+    spill    <- newTVar PSQ.empty
+    spillcnt <- newTVar 0
+    cache    <- newTVar PSQ.empty
+    cachecnt <- newTVar 0
+    return TCPProber
+        { probeQueue      = q
+        , probeSpill      = spill
+        , probeSpillCount = spillcnt
+        , probeCache      = cache
+        , probeCacheCount = cachecnt
+        }
+
+
+enqueueProbe :: TCPProber -> NodeInfo -> IO ()
+enqueueProbe prober ni = do
+    tm <- getPOSIXTime
+    atomically $ do
+        spill <- readTVar (probeSpill prober)
+        cache <- readTVar (probeCache prober)
+        q     <- readTVar (probeQueue prober)
+        let bump var x = modifyTVar' var $ insert' (nodeId ni) x tm
+        case PSQ.lookup (nodeId ni) cache of
+            Just (tm, x)                       -> bump (probeCache prober) x
+            Nothing | member (nodeId ni) spill -> bump (probeSpill prober) (nodeAddr ni)
+                    | member (nodeId ni) q     -> return ()
+                    | otherwise                -> bump (probeQueue prober) (nodeAddr ni)
+
+maxSpill :: Int
+maxSpill = 100
+
+maxCache :: Int
+maxCache = 50
+
+runProbeQueue :: TCPProber -> TCP.RelayClient -> Int -> IO ()
+runProbeQueue prober client maxjobs = do
+    jcnt <- newTVarIO 0
+    fix $ \loop -> do
+        (tm, mni) <- atomically $ do
+            j <- readTVar jcnt
+            check (j < maxjobs)
+            q <- readTVar $ probeQueue prober
+            case minView q of
+                Nothing                       -> retry
+                Just (Binding nid saddr tm,q') -> do
+                    writeTVar (probeQueue prober) q'
+                    return (tm, nodeInfo nid saddr)
+        forM_ mni $ \ni -> do
+            atomically $ modifyTVar' jcnt succ
+            t <- forkIO $ do
+                m <- resolvePort client ni
+                atomically $ case m of
+                    Nothing -> do
+                        pcnt <- readTVar (probeSpillCount prober)
+                        modifyTVar' (probeSpill prober) $ insert' (nodeId ni) (nodeAddr ni) tm
+                        if (pcnt == maxSpill)
+                            then modifyTVar' (probeSpill prober) deleteMin
+                            else modifyTVar' (probeSpillCount prober) succ
+                    Just p  -> do
+                        ccnt <- readTVar (probeCacheCount prober)
+                        modifyTVar' (probeCache prober) $ insert' (nodeId ni) (nodeAddr ni,p) tm
+                        if (ccnt == maxCache)
+                            then modifyTVar' (probeCache prober) deleteMin
+                            else modifyTVar' (probeCacheCount prober) succ
+                atomically $ modifyTVar' jcnt pred
+            labelThread t ("probe."++show ni)
+        loop
+
+getNodes :: TCPProber -> TCP.TCPClient err () Nonce8 -> NodeId -> TCP.NodeInfo -> IO (Maybe ([TCP.NodeInfo],[TCP.NodeInfo],Maybe ()))
+getNodes prober tcp seeking dst = do
+    r <- TCP.getUDPNodes' tcp seeking (TCP.udpNodeInfo dst)
+    let tcps (ns,_,mb) = (ns',ns',mb)
+         where ns' = do
+                    n <- ns
+                    [ TCP.NodeInfo n 0 ]
+    fmap join $ forM r $ \(ns,gw) -> do
+        let ts = tcps ns
+        if TCP.nodeId gw == TCP.nodeId dst
+            then return $ Just ts
+            else do
+                enqueueProbe prober (TCP.udpNodeInfo dst)
+                return $ Just ts
+        return $ Just ts
+
+nodeSearch :: TCPProber -> TCP.TCPClient err () Nonce8 -> Search NodeId (IP, PortNumber) () TCP.NodeInfo TCP.NodeInfo
+nodeSearch prober tcp = Search
+    { searchSpace       = TCP.tcpSpace
+    , searchNodeAddress = TCP.nodeIP &&& TCP.tcpPort
+    , searchQuery       = getNodes prober tcp
+    }
diff --git a/examples/dhtd.hs b/examples/dhtd.hs
index 959383dc..1223edc8 100644
--- a/examples/dhtd.hs
+++ b/examples/dhtd.hs
@@ -97,6 +97,7 @@ import qualified Network.Tox.Onion.Transport as Tox
 import qualified Network.Tox.Onion.Handlers as Tox
 import qualified Network.Tox.Crypto.Transport as Tox
 import qualified Network.Tox.TCP as TCP
+import qualified TCPProber as TCP
 import Data.Typeable
 import Network.Tox.ContactInfo as Tox
 import OnionRouter
@@ -1490,6 +1491,7 @@ initTox opts ssvar keysdb mbxmpp invc = case porttox opts of
                                         Want_IP4 -> toxStrap4
                                         Want_IP6 -> toxStrap6
                     }
+                tcpprober = tcpProber $ Tox.toxOnionRoutes tox
                 tcpclient = tcpKademliaClient $ Tox.toxOnionRoutes tox
                 tcpRefresher = tcpBucketRefresher $ Tox.toxOnionRoutes tox
                 tcpDHT = DHT
@@ -1500,9 +1502,9 @@ initTox opts ssvar keysdb mbxmpp invc = case porttox opts of
                                     , pingShowResult = show
                                     }
                     , dhtQuery = Map.singleton "node" DHTQuery
-                                    { qsearch = TCP.nodeSearch tcpclient
+                                    { qsearch = TCP.nodeSearch tcpprober tcpclient
                                     , qhandler = \ni nid -> do
-                                        ns <- R.kclosest (searchSpace $ TCP.nodeSearch tcpclient) searchK nid
+                                        ns <- R.kclosest (searchSpace $ TCP.nodeSearch tcpprober tcpclient) searchK nid
                                                 <$> atomically (readTVar $ refreshBuckets tcpRefresher)
                                         return (ns,ns,Just ())
                                     , qshowR = show -- TCP.NodeInfo
diff --git a/g b/g
index b5cdc09b..4ebe82ed 100755
--- a/g
+++ b/g
@@ -13,7 +13,7 @@ done
 
 warn="-freverse-errors -fwarn-unused-imports -Wmissing-signatures -fdefer-typed-holes"
 exts="-XOverloadedStrings -XRecordWildCards"
-defs="-DBENCODE_AESON -DTHREAD_DEBUG"
+defs="-DBENCODE_AESON -DTHREAD_DEBUG -UVERSION_lens"
 hidden="crypto-random crypto-api crypto-numbers cryptohash prettyclass"
 [ "$rootname" == "stack" ] && hidden=""
 hide=""
diff --git a/src/Network/Tox.hs b/src/Network/Tox.hs
index ea9bbe56..d42f447d 100644
--- a/src/Network/Tox.hs
+++ b/src/Network/Tox.hs
@@ -65,6 +65,7 @@ import Network.Tox.ContactInfo
 import Text.XXD
 import DPut
 import DebugTag
+import TCPProber
 import Network.Tox.Avahi
 import Network.Tox.Session
 import Network.Tox.Relay
@@ -375,6 +376,10 @@ forkTox tox with_avahi = do
     quitNC      <- forkListener "toxCrypto" (toxCrypto tox)
     quitTCP     <- forkListener "relay-client" (clientNet $ tcpClient $ tcpKademliaClient $ toxOnionRoutes tox)
     tcpKademlia <- forkPollForRefresh (tcpBucketRefresher $ toxOnionRoutes tox)
+    tcpprobe <- forkIO $ runProbeQueue (tcpProber $ toxOnionRoutes tox)
+                                       (tcpClient $ tcpKademliaClient $ toxOnionRoutes tox)
+                                       12
+    labelThread tcpprobe "tcp-probe"
     quitAvahi <- if with_avahi then do
                     forkPollForRefresh (DHT.refresher4 $ toxRouting tox)
                     forkPollForRefresh (DHT.refresher6 $ toxRouting tox)
@@ -392,6 +397,7 @@ forkTox tox with_avahi = do
                 quitOnion
                 quitTCP
                 killThread tcpKademlia
+                killThread tcpprobe
                 quitToRoute
                 quitHs
            , bootstrap (DHT.refresher4 $ toxRouting tox)
diff --git a/src/Network/Tox/TCP.hs b/src/Network/Tox/TCP.hs
index 48059108..e3f5012b 100644
--- a/src/Network/Tox/TCP.hs
+++ b/src/Network/Tox/TCP.hs
@@ -16,6 +16,7 @@ import Data.Functor.Identity
 import Data.Hashable
 import qualified Data.HashMap.Strict as HashMap
 import Data.IP
+import Data.Maybe
 import Data.Monoid
 import Data.Serialize
 import Data.Word
@@ -24,13 +25,14 @@ import Network.Socket (SockAddr(..))
 import qualified Text.ParserCombinators.ReadP as RP
 import System.IO.Error
 
+import ControlMaybe
 import Crypto.Tox
 import Data.ByteString (hPut,hGet,ByteString,length)
 import Data.Tox.Relay
 import qualified Data.Word64Map
 import DebugTag
 import DPut
-import Network.Address (setPort,PortNumber)
+import Network.Address (setPort,PortNumber,localhost4,fromSockAddr)
 import Network.Kademlia.Routing
 import Network.Kademlia.Search hiding (sendQuery)
 import Network.QueryResponse
@@ -116,7 +118,13 @@ tcpStream crypto = StreamHandshake
         dput XTCP $ "TCP:" ++ show addr ++ " <-- " ++ show hello
         hPut h $ encode $ encryptPayload (noncef $ helloNonce hello) hello
         welcomeE <- withSize $ fmap decode . hGet h . constSize
-        let Right welcome = welcomeE >>= \w -> decryptPayload (noncef $ welcomeNonce w) w
+        let mwelcome = welcomeE >>= \w -> decryptPayload (noncef $ welcomeNonce w) w
+            nil = SessionProtocol
+                { streamGoodbye = return ()
+                , streamDecode = return Nothing
+                , streamEncode = \y -> return ()
+                }
+        either (\_ -> return nil) id $ mwelcome <&> \welcome -> do
         dput XTCP $ "TCP:" ++ show addr ++ " --> " ++ show welcome
         noncef' <- lookupNonceFunction crypto skey (sessionPublicKey $ runIdentity $ welcomeData welcome)
         nsend <- newMVar (sessionBaseNonce $ runIdentity $ helloData hello)
@@ -166,12 +174,14 @@ toxTCP crypto = tcpTransport 30 (tcpStream crypto)
 tcpSpace :: KademliaSpace NodeId NodeInfo
 tcpSpace = contramap udpNodeInfo toxSpace
 
+{-
 nodeSearch :: TCPClient err () Nonce8 -> Search NodeId (IP, PortNumber) () NodeInfo NodeInfo
 nodeSearch tcp = Search
     { searchSpace       = tcpSpace
     , searchNodeAddress = nodeIP &&& tcpPort
-    , searchQuery       = getTCPNodes tcp
+    , searchQuery       = getNodes tcp
     }
+-}
 
 data TCPClient err meth tid = TCPClient
     { tcpCrypto     :: TransportCrypto
@@ -179,24 +189,44 @@ data TCPClient err meth tid = TCPClient
     , tcpGetGateway :: UDP.NodeInfo -> STM (Maybe NodeInfo)
     }
 
+{-
 getTCPNodes :: TCPClient err () Nonce8 -> NodeId -> NodeInfo -> IO (Maybe ([NodeInfo],[NodeInfo],Maybe ()))
 getTCPNodes tcp seeking dst = do
-    r <- getUDPNodes tcp seeking (udpNodeInfo dst)
+    r <- getUDPNodes' tcp seeking (udpNodeInfo dst)
     let tcps (ns,_,mb) = (ns',ns',mb)
          where ns' = do
                     n <- ns
                     [ NodeInfo n (fromIntegral 443) , NodeInfo n (fromIntegral 80) , NodeInfo n (UDP.nodePort n) ]
-    return $ tcps <$> r
+    fmap join $ forM r $ \(ns,gw) -> do
+        let ts = tcps ns
+        {-
+        if nodeId gw == nodeId dst
+            then return $ Just ts
+            else do
+                forkIO $ void $ tcpPing (tcpClient tcp) dst
+                return $ Just ts
+        -}
+        forM_ ((\(xs,_,_) -> xs) ts) (forkIO . void . tcpPing (tcpClient tcp))
+        return $ Just ts
+-}
 
 getUDPNodes :: TCPClient err () Nonce8 -> NodeId -> UDP.NodeInfo -> IO (Maybe ([UDP.NodeInfo], [UDP.NodeInfo], Maybe ()))
-getUDPNodes tcp seeking dst = do
-    mgateway <- atomically $ tcpGetGateway tcp dst
+getUDPNodes tcp seeking dst = fmap fst <$> getUDPNodes' tcp seeking dst
+
+getUDPNodes' :: TCPClient err () Nonce8 -> NodeId -> UDP.NodeInfo -> IO (Maybe (([UDP.NodeInfo], [UDP.NodeInfo], Maybe ()), NodeInfo))
+getUDPNodes' tcp seeking dst0 = do
+    mgateway <- atomically $ tcpGetGateway tcp dst0
     fmap join $ forM mgateway $ \gateway -> do
         (b,c,n24) <- atomically $ do
             b <- transportNewKey (tcpCrypto tcp)
             c <- transportNewKey (tcpCrypto tcp)
             n24 <- transportNewNonce (tcpCrypto tcp)
             return (b,c,n24)
+        let (dst,gateway') = if UDP.nodeId dst0 == nodeId gateway
+                                then ( dst0 { UDP.nodeIP = fromJust $ fromSockAddr localhost4 }
+                                     , gateway { udpNodeInfo = (udpNodeInfo gateway)
+                                                    { UDP.nodeIP = fromJust $ fromSockAddr localhost4 }})
+                                else (dst0,gateway)
         wrap2 <- lookupNonceFunction (tcpCrypto tcp) b (UDP.id2key $ UDP.nodeId dst)
         wrap1 <- lookupNonceFunction (tcpCrypto tcp) c (UDP.id2key $ nodeId gateway)
         wrap0 <- lookupNonceFunction (tcpCrypto tcp) (transportSecret $ tcpCrypto tcp) (UDP.id2key $ UDP.nodeId dst)
@@ -205,7 +235,7 @@ getUDPNodes tcp seeking dst = do
                     , method = () -- meth
                     , wrapQuery = \n8 src gateway x ->
                         OnionPacket n24 $ Addressed (UDP.nodeAddr dst)
-                            $ wrapOnionPure b (wrap2 n24) (nodeAddr gateway)
+                            $ wrapOnionPure b (wrap2 n24) (nodeAddr gateway')
                                 $ wrapOnionPure c (wrap1 n24) (UDP.nodeAddr dst)
                                     $ NotForwarded $ encryptPayload (wrap0 n24)
                                        $ OnionAnnounce Asymm
@@ -221,7 +251,7 @@ getUDPNodes tcp seeking dst = do
         r <- sendQuery (tcpClient tcp) meth (AnnounceRequest zeros32 seeking UDP.zeroID) gateway
         forM r $ \response -> do
             let (ns,_,mb) = either (const ([],[],Nothing)) (unwrapAnnounceResponse Nothing dst) $ response
-            return (ns,ns, const () <$> mb)
+            return ( (ns,ns, const () <$> mb), gateway )
 
 
 handleOOB :: PublicKey -> ByteString -> NodeInfo -> NodeInfo -> IO (Maybe (RelayPacket -> RelayPacket))
@@ -243,7 +273,9 @@ tcpPing client dst = sendQuery client meth () dst
                 , method         = ()
                 }
 
-newClient :: TransportCrypto -> IO (Client String () Nonce8 NodeInfo RelayPacket)
+type RelayClient = Client String () Nonce8 NodeInfo RelayPacket
+
+newClient :: TransportCrypto -> IO RelayClient
 newClient crypto = do
     net <- toxTCP crypto
     drg <- drgNew