2 files changed, 12 insertions, 11 deletions

diff --git a/dht/src/Network/Tox/Crypto/Transport.hs b/dht/src/Network/Tox/Crypto/Transport.hs
index a18b550d..127b2dac 100644
--- a/dht/src/Network/Tox/Crypto/Transport.hs
+++ b/dht/src/Network/Tox/Crypto/Transport.hs
@@ -91,21 +91,21 @@ import GHC.TypeNats
 showCryptoMsg :: Word32 -> CryptoMessage -> [Char]
 showCryptoMsg _  msg = show msg
 
-parseCrypto :: (ByteString, SockAddr) -> Either (CryptoPacket Encrypted, SockAddr) (ByteString, SockAddr)
+parseCrypto :: (ByteString, addr ) -> Either (CryptoPacket Encrypted, addr) (ByteString, addr)
 parseCrypto (bbs,saddr) = case B.uncons bbs of
     Just (0x1b,bs) -> case runGet get bs of
         Right pkt -> Left (pkt, saddr) -- Successful parse, handle this packet.
         Left  _   -> Right (bs,saddr)  -- Failed parse, strip first byte and pass it on.
     _             -> Right (bbs,saddr) -- Type-code mismatch, pass it on.
 
-encodeCrypto :: (CryptoPacket Encrypted, SockAddr) -> Maybe (ByteString, SockAddr)
+encodeCrypto :: (CryptoPacket Encrypted, addr) -> Maybe (ByteString, addr)
 encodeCrypto (x,saddr) = Just (B.cons 0x1b (runPut $ put x),saddr)
 
-parseHandshakes :: ByteString -> SockAddr -> Either String (Handshake Encrypted, SockAddr)
+parseHandshakes :: ByteString -> addr -> Either String (Handshake Encrypted, addr)
 parseHandshakes (B.uncons -> Just (0x1a,pkt)) saddr = left ("parseHandshakes: "++) $ (,saddr) <$> runGet get pkt
 parseHandshakes bs                            _     = Left $ "parseHandshakes_: " ++ show (B.unpack $ B.take 1 bs)
 
-encodeHandshakes :: Handshake Encrypted -> SockAddr -> (ByteString, SockAddr)
+encodeHandshakes :: Handshake Encrypted -> addr -> (ByteString, addr)
 encodeHandshakes x saddr = (B.cons 0x1a (runPut $ put x),saddr)
 
 {-
diff --git a/dht/src/Network/Tox/Transport.hs b/dht/src/Network/Tox/Transport.hs
index 0ca9b758..12886245 100644
--- a/dht/src/Network/Tox/Transport.hs
+++ b/dht/src/Network/Tox/Transport.hs
@@ -65,14 +65,15 @@ toxTransport crypto orouter closeLookup addr udp relaynet tcp2server tcp2client
                              (encodeOnionAddr crypto $ lookupRoute orouter)
                              udp1
     (dta,onion) <- partitionTransportM (parseDataToRoute crypto) (encodeDataToRoute crypto) onion1
+    (netcryptoTCP, relaynet1) <- partitionTransport parseCrypto encodeCrypto relaynet0
+    multi_netcrypto <- mergeTransports $ DMap.fromList
+        [ Multi.SessionUDP :=> ByAddress netcrypto
+        , Multi.SessionTCP :=> ByAddress netcryptoTCP ]
     let handshakes = layerTransport parseHandshakes encodeHandshakes udp2
-        promoteUDP :: TransportA err SockAddr x y -> TransportA err Multi.SessionAddress x y
-        promoteUDP net = layerTransport (\msg saddr -> Right (msg,Multi.SessionUDP ==> saddr))
-                                        (\msg (Multi.SessionUDP :=> Identity saddr) -> (msg,saddr))
-                                        net
-        -- TODO: Enable sessions over TCP
-        multi_netcrypto = promoteUDP netcrypto
-        multi_handshakes = promoteUDP handshakes
+        handshakesTCP = layerTransport parseHandshakes encodeHandshakes relaynet1
+    multi_handshakes <- mergeTransports $ DMap.fromList
+        [ Multi.SessionUDP :=> ByAddress handshakes
+        , Multi.SessionTCP :=> ByAddress handshakesTCP ]
     return ( multi_netcrypto
            , forwardDHTRequests crypto (fmap (fmap (Multi.UDP ==>)) . closeLookup) dht
            , onion