1 files changed, 13 insertions, 9 deletions

diff --git a/src/Network/Tox/Transport.hs b/src/Network/Tox/Transport.hs
index e79e4d8b..217d5b1d 100644
--- a/src/Network/Tox/Transport.hs
+++ b/src/Network/Tox/Transport.hs
@@ -10,7 +10,8 @@ module Network.Tox.Transport (toxTransport, RouteId) where
 
 import Network.QueryResponse
 import Crypto.Tox
-import Network.Tox.DHT.Transport
+import Data.Tox.Relay as TCP
+import Network.Tox.DHT.Transport as UDP
 import Network.Tox.Onion.Transport
 import Network.Tox.Crypto.Transport
 import OnionRouter
@@ -20,20 +21,23 @@ import Network.Socket
 toxTransport ::
       TransportCrypto
           -> OnionRouter
-          -> (PublicKey -> IO (Maybe NodeInfo))
+          -> (PublicKey -> IO (Maybe UDP.NodeInfo))
           -> UDPTransport
-          -> (Int -> OnionMessage Encrypted -> IO ()) -- ^ TCP-bound callback.
+          -> (TCP.NodeInfo -> RelayPacket -> IO ())   -- ^ TCP server-bound callback.
+          -> (Int -> OnionMessage Encrypted -> IO ()) -- ^ TCP client-bound callback.
           -> IO ( Transport String SockAddr (CryptoPacket Encrypted)
-                , Transport String NodeInfo (DHTMessage Encrypted8)
+                , Transport String UDP.NodeInfo (DHTMessage Encrypted8)
                 , Transport String (OnionDestination RouteId) (OnionMessage Encrypted)
                 , Transport String AnnouncedRendezvous (PublicKey,OnionData)
                 , Transport String SockAddr (Handshake Encrypted))
-toxTransport crypto orouter closeLookup udp tcp = do
+toxTransport crypto orouter closeLookup udp tcp2server tcp2client = do
     (netcrypto, udp0) <- partitionTransport parseCrypto encodeCrypto udp
-    (dht,udp1) <- partitionTransportM (parseDHTAddr crypto) (fmap Just . encodeDHTAddr) $ forwardOnions crypto udp0 tcp
-    (onion1,udp2) <- partitionTransportM (parseOnionAddr $ lookupSender orouter)
-                                        (encodeOnionAddr crypto $ lookupRoute orouter)
-                                        udp1
+    (dht,udp1) <- partitionTransportM (parseDHTAddr crypto) (fmap Just . encodeDHTAddr)
+                                    $ forwardOnions crypto udp0 tcp2client
+    (onion1,udp2) <- partitionAndForkTransport tcp2server
+                             (parseOnionAddr $ lookupSender orouter)
+                             (encodeOnionAddr crypto $ lookupRoute orouter)
+                             udp1
     (dta,onion) <- partitionTransportM (parseDataToRoute crypto) (encodeDataToRoute crypto) onion1
     let handshakes = layerTransport parseHandshakes encodeHandshakes udp2
     return ( netcrypto