1 files changed, 10 insertions, 6 deletions

diff --git a/src/Network/Tox/Transport.hs b/src/Network/Tox/Transport.hs
index 30df93c8..57f07aad 100644
--- a/src/Network/Tox/Transport.hs
+++ b/src/Network/Tox/Transport.hs
@@ -8,6 +8,7 @@
 {-# LANGUAGE TypeOperators              #-}
 module Network.Tox.Transport (toxTransport, RouteId) where
 
+import Data.ByteString (ByteString)
 import Network.QueryResponse
 import Crypto.Tox
 import Network.Tox.DHT.Transport
@@ -22,21 +23,24 @@ toxTransport ::
           -> OnionRouter
           -> (PublicKey -> IO (Maybe NodeInfo))
           -> UDPTransport
-          -> IO ( Transport String NodeInfo (DHTMessage Encrypted8)
+          -> IO ( Transport String SockAddr (CryptoPacket Encrypted)
+                , Transport String NodeInfo (DHTMessage Encrypted8)
                 , Transport String (OnionDestination RouteId) (OnionMessage Encrypted)
                 , Transport String AnnouncedRendezvous (PublicKey,OnionData)
-                , Transport String SockAddr NetCrypto )
+                , Transport String SockAddr (Handshake Encrypted))
 toxTransport crypto orouter closeLookup udp = do
-    (dht,udp1) <- partitionTransportM (parseDHTAddr crypto) (fmap Just . encodeDHTAddr) $ forwardOnions crypto udp
+    (netcrypto, udp0) <- partitionTransport parseCrypto encodeCrypto udp
+    (dht,udp1) <- partitionTransportM (parseDHTAddr crypto) (fmap Just . encodeDHTAddr) $ forwardOnions crypto udp0
     (onion1,udp2) <- partitionTransportM (parseOnionAddr $ lookupSender orouter)
                                         (encodeOnionAddr crypto $ lookupRoute orouter)
                                         udp1
     (dta,onion) <- partitionTransportM (parseDataToRoute crypto) (encodeDataToRoute crypto) onion1
-    let netcrypto = layerTransport parseNetCrypto encodeNetCrypto udp2
-    return ( forwardDHTRequests crypto closeLookup dht
+    let handshakes = layerTransport parseHandshakes encodeHandshakes udp2
+    return ( netcrypto
+           , forwardDHTRequests crypto closeLookup dht
            , onion
            , dta
-           , netcrypto
+           , handshakes
            )