Multi Transports: TCP for DHT/Cookies/Handshakes.

1 files changed, 66 insertions, 14 deletions

diff --git a/dht/src/Network/Tox/TCP.hs b/dht/src/Network/Tox/TCP.hs
index 4b3a4594..dc4c9967 100644
--- a/dht/src/Network/Tox/TCP.hs
+++ b/dht/src/Network/Tox/TCP.hs
@@ -21,6 +21,7 @@ import Data.Functor.Contravariant
 import Data.Functor.Identity
 import Data.Hashable
 import qualified Data.HashMap.Strict as HashMap
+import qualified Data.IntMap.Strict as IntMap
 import Data.IP
 import Data.Maybe
 import Data.Monoid
@@ -48,9 +49,9 @@ import Network.Kademlia.Search hiding (sendQuery)
 import Network.QueryResponse
 import Network.QueryResponse.TCP
 import Network.Tox.TCP.NodeId ()
-import Network.Tox.DHT.Handlers (toxSpace)
+import Network.Tox.DHT.Transport (toxSpace)
 import Network.Tox.Onion.Transport hiding (encrypt,decrypt)
-import Network.Tox.Onion.Handlers (unwrapAnnounceResponse)
+import Network.Tox.Onion.Transport (unwrapAnnounceResponse)
 import qualified Network.Tox.NodeId as UDP
 import Text.XXD
 import Data.Proxy
@@ -72,8 +73,8 @@ nodeIP :: NodeInfo -> IP
 nodeIP ni = UDP.nodeIP $ udpNodeInfo ni
 
 tcpStream :: (Show y, Show x, Serialize y, Sized y, Serialize x, Sized x) =>
-                   TransportCrypto -> StreamHandshake NodeInfo x y
-tcpStream crypto = StreamHandshake
+                   TransportCrypto -> (NodeInfo -> IO st) -> StreamHandshake NodeInfo (st,x) y
+tcpStream crypto mkst = StreamHandshake
     { streamHello = \addr h -> do
         (skey, hello) <- atomically $ do
             n24 <- transportNewNonce crypto
@@ -113,6 +114,7 @@ tcpStream crypto = StreamHandshake
         nread <- newMVar (sessionBaseNonce $ runIdentity $ welcomeData welcome)
         let them = sessionPublicKey $ runIdentity $ welcomeData welcome
         hvar <- newMVar h
+        st <- mkst addr
         return SessionProtocol
             { streamGoodbye = do
                 dput XTCP $ "Closing " ++ show addr
@@ -138,7 +140,7 @@ tcpStream crypto = StreamHandshake
                                                     dput XTCP $ "TCP:" ++ show addr ++ " --> " ++ show x'
                                                     return ())
                                                r
-                                        return $ either (const Nothing) Just r
+                                        return $ either (const Nothing) (Just . (,) st) r
                  in bracket (takeMVar hvar) (putMVar hvar)
                      $ \h -> go h `catchIOError` \e -> do
                                         dput XTCP $ "TCP exception: " ++ show e
@@ -158,9 +160,26 @@ tcpStream crypto = StreamHandshake
     , streamAddr = nodeAddr
     }
 
-toxTCP :: TransportCrypto -> IO ( TCPCache (SessionProtocol RelayPacket RelayPacket)
-                                , TransportA err NodeInfo RelayPacket (Bool,RelayPacket) )
-toxTCP crypto = tcpTransport 30 (tcpStream crypto)
+newtype SessionData = SessionData (MVar (IntMap.IntMap NodeId))
+
+newSessionData :: NodeInfo -> IO SessionData
+newSessionData _ = SessionData <$> newMVar IntMap.empty
+
+getRelayedRemote :: SessionData -> ConId -> IO NodeId
+getRelayedRemote (SessionData keymapVar) (ConId i) = do
+    keymap <- takeMVar keymapVar
+    let k = fromMaybe UDP.zeroID $ IntMap.lookup (fromIntegral i) keymap
+    putMVar keymapVar keymap
+    return k
+
+setRelayedRemote :: SessionData -> ConId -> NodeId -> IO ()
+setRelayedRemote (SessionData keymapVar) (ConId conid) nid = do
+    keymap <- takeMVar keymapVar
+    putMVar keymapVar $ IntMap.insert (fromIntegral conid) nid keymap
+
+toxTCP :: TransportCrypto -> IO ( TCPCache (SessionProtocol (SessionData,RelayPacket) RelayPacket)
+                                , TransportA err NodeInfo (SessionData,RelayPacket) (Bool,RelayPacket) )
+toxTCP crypto = tcpTransport 30 (tcpStream crypto newSessionData)
 
 tcpSpace :: KademliaSpace NodeId NodeInfo
 tcpSpace = contramap udpNodeInfo toxSpace
@@ -292,6 +311,8 @@ keyToNonce k = unsafeDupablePerformIO $ withByteArray k $ \ptr -> do
     w8 <- peek ptr
     return $ Nonce8 w8
 
+type RelayCache = TCPCache (SessionProtocol (SessionData,RelayPacket) RelayPacket)
+
 -- | Create a new TCP relay client.  Because polymorphic existential record
 -- updates are currently hard with GHC, this function accepts parameters for
 -- generalizing the table-entry type for pending transactions.  Safe trivial
@@ -301,14 +322,18 @@ newClient :: TransportCrypto
                -> ((Maybe (Bool,RelayPacket) -> IO ()) -> a) -- ^ store mvar for query
                -> (a -> RelayPacket -> IO void)  -- ^ load mvar for query
                -> IO ( ( TVar (ChaChaDRG, Data.Word64Map.Word64Map a)
-                       , TCPCache (SessionProtocol RelayPacket RelayPacket) )
+                       , RelayCache
+                       , Transport String ViaRelay ByteString )
                      , Client String PacketNumber Nonce8 NodeInfo (Bool,RelayPacket))
 newClient crypto store load = do
-    (tcpcache,net) <- toxTCP crypto
+    (tcpcache,net0) <- toxTCP crypto
+    (relaynet,net1) <- partitionRelay net0
+    let net2 = {- XXX: Client type forces this pointless layering. -}
+               layerTransport ((Right .) . (,) . (,) False . snd) (,) net1
     drg <- drgNew
     map_var <- atomically $ newTVar (drg, Data.Word64Map.empty)
-    return $ (,) (map_var,tcpcache) Client
-        { clientNet = {- XXX: Client type forces this pointless layering. -} layerTransport ((Right .) . (,) . (,) False) (,) net
+    return $ (,) (map_var,tcpcache,relaynet) Client
+        { clientNet = net2
         , clientDispatcher = DispatchMethods
             { classifyInbound = (. snd) $ \case
                 RelayPing n             -> IsQuery PingPacket n
@@ -318,7 +343,7 @@ newClient crypto store load = do
                 OnionPacketResponse (OnionAnnounceResponse n8 n24 ciphered) -> IsResponse n8
                 OnionPacketResponse o@(OnionToRouteResponse _) -> IsUnsolicited $ handle2route o
                 OOBRecv k bs            -> IsUnsolicited $ handleOOB k bs
-                wut                   -> IsUnknown (show wut)
+                wut                     -> IsUnknown (show wut)
             , lookupHandler   = \case
                 PingPacket -> trace ("tcp-received-ping") $ Just MethodHandler
                     { methodParse     = \case (_,RelayPing n8) -> Right ()
@@ -330,7 +355,10 @@ newClient crypto store load = do
                             { methodParse = \x -> Left "tcp-lookuphandler?" -- :: x -> Either err a
                             , noreplyAction = \addr a -> dput XTCP $ "tcp-lookupHandler: "++show w
                             }
-            , tableMethods = transactionMethods' store (\x -> mapM_ (load x . snd)) (contramap (\(Nonce8 w64) -> w64) w64MapMethods)
+            , tableMethods = transactionMethods'
+                                store
+                                (\x -> mapM_ (load x . snd))
+                                (contramap (\(Nonce8 w64) -> w64) w64MapMethods)
                                 $ first (either error Nonce8 . decode) . randomBytesGenerate 8
             }
         , clientErrorReporter = logErrors
@@ -341,3 +369,27 @@ newClient crypto store load = do
             }
         , clientResponseId = return
         }
+
+data ViaRelay = ViaRelay (Maybe ConId) UDP.NodeId NodeInfo
+ deriving (Eq,Ord,Show)
+
+partitionRelay :: TransportA err NodeInfo (SessionData,RelayPacket) (Bool,RelayPacket)
+                      -> IO ( Transport err ViaRelay ByteString
+                            , TransportA err NodeInfo (SessionData,RelayPacket) (Bool,RelayPacket))
+partitionRelay tr = partitionTransportM parse encode tr
+  where
+    parse :: ((SessionData,RelayPacket), NodeInfo) -> IO (Either (ByteString, ViaRelay) ((SessionData,RelayPacket),NodeInfo))
+    parse ((st,RelayData bs conid), ni) = do
+        nid <- getRelayedRemote st conid
+        return $ Left (bs, ViaRelay (Just conid) nid ni)
+    parse ((_,OOBRecv key bs), ni) =
+        return $ Left (bs, ViaRelay Nothing (UDP.key2id key) ni)
+    parse passthrough@((st,RoutingResponse conid k),ni) = do
+        setRelayedRemote st conid (UDP.key2id k)
+        return $ Right passthrough
+    parse passthrough =
+        return $ Right passthrough
+
+    encode :: (ByteString, ViaRelay) -> IO (Maybe ((Bool,RelayPacket), NodeInfo))
+    encode (bs, ViaRelay (Just conid) _ ni) = return $ Just ((False,RelayData bs conid), ni)
+    encode (bs, ViaRelay Nothing    nid ni) = return $ Just ((False,OOBSend (UDP.id2key nid) bs), ni)