1 files changed, 27 insertions, 5 deletions

diff --git a/dht/src/Network/Tox.hs b/dht/src/Network/Tox.hs
index f136ab96..23dbfe27 100644
--- a/dht/src/Network/Tox.hs
+++ b/dht/src/Network/Tox.hs
@@ -1,12 +1,13 @@
 {-# LANGUAGE CPP                        #-}
-{-# LANGUAGE FlexibleContexts #-}
 {-# LANGUAGE DeriveDataTypeable         #-}
 {-# LANGUAGE DeriveFoldable             #-}
 {-# LANGUAGE DeriveFunctor              #-}
 {-# LANGUAGE DeriveGeneric              #-}
 {-# LANGUAGE DeriveTraversable          #-}
 {-# LANGUAGE ExistentialQuantification  #-}
+{-# LANGUAGE FlexibleContexts           #-}
 {-# LANGUAGE FlexibleInstances          #-}
+{-# LANGUAGE GADTs                      #-}
 {-# LANGUAGE GeneralizedNewtypeDeriving #-}
 {-# LANGUAGE LambdaCase                 #-}
 {-# LANGUAGE NamedFieldPuns             #-}
@@ -72,7 +73,7 @@ import qualified Network.Tox.Onion.Transport as Onion
 import Network.Tox.RelayPinger
 import System.Global6
 import Network.Tox.Transport
-import Network.Tox.TCP (tcpClient)
+import Network.Tox.TCP (tcpClient, ViaRelay(..))
 import Network.Tox.Onion.Routes
 import Network.Tox.ContactInfo
 import Text.XXD
@@ -131,7 +132,7 @@ myAddr routing4 routing6 maddr = atomically $ do
     a <- readTVar var
     return $ Multi.UDP ==> R.thisNode a
 
-newClient :: (DRG g, Show addr, Show meth) =>
+newClient :: forall g addr meth x. (DRG g, Show addr, Show meth) =>
               g -> Transport String addr x
                 -> (Client String meth DHT.TransactionId addr x
                     -> x
@@ -160,12 +161,25 @@ newClient drg net classify selfAddr handlers modifytbl modifynet = do
             let word64mapT = transactionMethods (contramap w64Key w64MapMethods) gen
             map_var <- atomically $ newTVar (drg, Data.Word64Map.empty)
             return $ Left (word64mapT,map_var)
-    let dispatch tbl var handlers client = DispatchMethods
+    let dispatch :: TransactionMethods tbl DHT.TransactionId addr x
+                      -> p
+                      -> (meth
+                          -> Maybe (MethodHandlerA String DHT.TransactionId addr x y))
+                      -> Client String meth DHT.TransactionId addr x
+                      -> DispatchMethodsA tbl String meth DHT.TransactionId addr x y
+        dispatch tbl var handlers client = DispatchMethods
                 { classifyInbound = classify client
                 , lookupHandler   = handlers -- var
                 , tableMethods    = modifytbl tbl
                 }
         eprinter = logErrors -- printErrors stderr
+        mkclient :: (TransactionMethods
+                         (g, pending) DHT.TransactionId addr x,
+                       TVar (g, pending))
+                      -> (ClientA String meth DHT.TransactionId addr x x
+                          -> meth
+                          -> Maybe (MethodHandlerA String DHT.TransactionId addr x x))
+                      -> ClientA String meth DHT.TransactionId addr x x
         mkclient (tbl,var) handlers =
             let client = Client
                     { clientNet           = addHandler (reportParseError eprinter) (handleMessage client) $ modifynet client net
@@ -337,7 +351,15 @@ newToxOverTransport keydb addr onNewSession (crypto,roster) udp tcp = do
     dhtclient <- newClient drg dhtnet0 DHT.classify (myAddr tbl4 tbl6) (DHT.handlers crypto . mkrouting) id
                     (\client net -> onInbound (updateOnUDP client) net)
 
-    hscache <- newHandshakeCache crypto (sendMessage handshakes)
+    hscache <- newHandshakeCache crypto $ \saddr hs -> do
+        saddr' <- case saddr of
+            Multi.SessionTCP :=> Identity (ViaRelay Nothing nid relay)
+              -> do let relayclient = relayClient $ tcpRelayPinger orouter
+                    msaddr <- Multi.tcpConnectionRequest relayclient (id2key nid) relay
+                    return $ maybe saddr Multi.sessionAddr msaddr
+            _ -> return saddr
+        sendMessage handshakes saddr' hs
+        return saddr'
     let sparams = SessionParams
             { spCrypto           = crypto
             , spSessions         = sessions