1 files changed, 42 insertions, 10 deletions

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