1 files changed, 51 insertions, 9 deletions

diff --git a/dht/src/Network/Tox/TCP.hs b/dht/src/Network/Tox/TCP.hs
index c4727a20..04119164 100644
--- a/dht/src/Network/Tox/TCP.hs
+++ b/dht/src/Network/Tox/TCP.hs
@@ -43,7 +43,7 @@ import Data.Tox.Relay
 import qualified Data.Word64Map
 import DebugTag
 import DPut
-import Network.Address (setPort,PortNumber,localhost4,fromSockAddr)
+import Network.Address (setPort,PortNumber,localhost4,fromSockAddr,nullAddress4)
 import Network.Kademlia.Routing
 import Network.Kademlia.Search hiding (sendQuery)
 import Network.QueryResponse
@@ -319,21 +319,25 @@ type RelayCache = TCPCache (SessionProtocol (SessionData,RelayPacket) RelayPacke
 -- defaults are 'id' and 'tryPutMVar'.  The resulting customized table state
 -- will be returned to the caller along with the new client.
 newClient :: TransportCrypto
-               -> ((Maybe (Bool,RelayPacket) -> IO ()) -> a) -- ^ store mvar for query
-               -> (a -> RelayPacket -> IO void)  -- ^ load mvar for query
+               -> ((Maybe (Bool,RelayPacket) -> IO ()) -> a)                    -- ^ store mvar for relay query
+               -> (a -> RelayPacket -> IO void)                                 -- ^ load mvar for relay query
+               -> (SockAddr -> Nonce8 -> IO (Maybe (OnionDestination RouteId))) -- ^ lookup sender of onion query
+               -> (UDP.NodeInfo -> RouteId -> IO (Maybe OnionRoute))            -- ^ lookup OnionRoute by id
                -> IO ( ( TVar (ChaChaDRG, Data.Word64Map.Word64Map a)
                        , RelayCache
-                       , Transport String ViaRelay ByteString )
+                       , Transport String ViaRelay ByteString
+                       , Transport String (OnionDestination RouteId) (OnionMessage Encrypted) )
                      , Client String PacketNumber Nonce8 NodeInfo (Bool,RelayPacket))
-newClient crypto store load = do
+newClient crypto store load lookupSender getRoute = do
     (tcpcache,net0) <- toxTCP crypto
     (relaynet,net1) <- partitionRelay net0
-    let net2 = {- XXX: Client type forces this pointless layering. -}
-               layerTransport ((Right .) . (,) . (,) False . snd) (,) net1
+    (onionnet,net2) <- partitionOnion crypto lookupSender getRoute net1
+    let net3 = {- XXX: Client type forces this pointless layering. -}
+               layerTransport ((Right .) . (,) . (,) False . snd) (,) net2
     drg <- drgNew
     map_var <- atomically $ newTVar (drg, Data.Word64Map.empty)
-    return $ (,) (map_var,tcpcache,relaynet) Client
-        { clientNet = net2
+    return $ (,) (map_var,tcpcache,relaynet,onionnet) Client
+        { clientNet = net3
         , clientDispatcher = DispatchMethods
             { classifyInbound = (. snd) $ \case
                 RelayPing n             -> IsQuery PingPacket n
@@ -393,3 +397,41 @@ partitionRelay tr = partitionTransportM parse encode tr
     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)
+
+
+partitionOnion :: TransportCrypto
+                  -> (SockAddr -> Nonce8 -> IO (Maybe (OnionDestination RouteId)))
+                  -> (UDP.NodeInfo -> RouteId -> IO (Maybe OnionRoute))
+                  -> TransportA err NodeInfo (SessionData,RelayPacket) (Bool,RelayPacket)
+                  -> IO ( Transport err (OnionDestination RouteId) (OnionMessage Encrypted)
+                        , TransportA err NodeInfo (SessionData,RelayPacket) (Bool,RelayPacket))
+partitionOnion crypto lookupSender getRoute tr = partitionTransportM parse encode tr
+ where
+    parse :: ((SessionData,RelayPacket), NodeInfo)
+                -> IO (Either (OnionMessage Encrypted   , OnionDestination RouteId)
+                              ((SessionData,RelayPacket), NodeInfo))
+    parse pass@((_,OnionPacketResponse msg@(OnionAnnounceResponse n8 _ _)), nodeA) = do
+        m <- lookupSender (nodeAddr nodeA) n8
+        case m of
+            Nothing -> return $ Right pass
+            Just od -> return $ Left (msg, od)
+    parse ((_,OnionPacketResponse msg@(OnionToRouteResponse asym)), nodeA) =
+        return $
+            let Right ni = UDP.nodeInfo (UDP.key2id $ senderKey asym) nullAddress4
+                -- -- We have this information, but currently, we're discarding it...
+                -- r = dummyRoute { routeNodeA     = udpNodeInfo nodeA
+                --                , routeRelayPort = Just $ tcpPort nodeA }
+                tryAllKeys = SearchingAlias -- We unfortunately don't know what toxid was used to encrypt this.
+                                            -- Toxcore only supports a single toxid per DHT node and in that case,
+                                            -- it is unambiguous.
+            in Left (msg, OnionDestination tryAllKeys ni Nothing)
+    parse pass = return $ Right pass
+
+    encode :: (OnionMessage Encrypted,OnionDestination RouteId) -> IO (Maybe ((Bool,RelayPacket),NodeInfo))
+    encode (msg,OnionDestination _ ni (Just rid)) = do
+        moroute <- getRoute ni rid
+        forM (moroute >>= \r -> (,) r <$> routeRelayPort r) $ \(oroute,tcpport) ->
+            wrapIndirectHops crypto msg ni oroute $ \nonce saddr fwd ->
+                return ( (True,OnionPacket nonce $ Addressed saddr fwd)
+                       , NodeInfo (routeNodeA oroute) tcpport )
+    encode _ = return Nothing