Used partitionTransform to simplify the onion client.

1 files changed, 29 insertions, 24 deletions

diff --git a/dht/src/Network/Tox/Onion/Transport.hs b/dht/src/Network/Tox/Onion/Transport.hs
index 407cd387..913d339d 100644
--- a/dht/src/Network/Tox/Onion/Transport.hs
+++ b/dht/src/Network/Tox/Onion/Transport.hs
@@ -24,6 +24,7 @@ module Network.Tox.Onion.Transport
     , decrypt
     , peelSymmetric
     , OnionRoute(..)
+    , dummyRoute
     , N0
     , N1
     , N2
@@ -39,9 +40,11 @@ module Network.Tox.Onion.Transport
     , wrapOnion
     , wrapOnionPure
     , unwrapAnnounceResponse
+    , wrapIndirectHops
     ) where
 
-import Data.ByteString (ByteString)
+import qualified Data.ByteString as B
+         ;import Data.ByteString (ByteString)
 import Data.Maybe
 import Data.Serialize
 import Network.Socket
@@ -50,8 +53,10 @@ import Crypto.Tox hiding (encrypt,decrypt)
 import Network.Tox.TCP.NodeId (udpNodeInfo)
 import qualified Data.Tox.Relay as TCP
 import Data.Tox.Onion
+import Network.Address (nullAddress4)
 import Network.Tox.DHT.Transport (SendNodes(..))
 import Network.Tox.NodeId
+import Network.QueryResponse
 
 {-
 encodeOnionAddr :: TransportCrypto
@@ -88,28 +93,28 @@ wrapForRoute :: TransportCrypto
                       -> NodeInfo
                       -> OnionRoute
                       -> IO (Either TCP.RelayPacket (OnionRequest N0))
-wrapForRoute crypto msg ni r@OnionRoute{routeRelayPort=Nothing} = do
-    -- We needn't use the same nonce value here, but I think it is safe to do so.
-    let nonce = msgNonce msg
-    fwd <-  wrapOnion crypto (routeAliasA r)
-                             nonce
-                             (id2key . nodeId $ routeNodeA r)
-                             (nodeAddr $ routeNodeB r)
-        =<< wrapOnion crypto (routeAliasB r)
-                             nonce
-                             (id2key . nodeId $ routeNodeB r)
-                             (nodeAddr $ routeNodeC r)
-        =<< wrapOnion crypto (routeAliasC r)
-                             nonce
-                             (id2key . nodeId $ routeNodeC r)
-                             (nodeAddr ni)
-                             (NotForwarded msg)
-    return $ Right OnionRequest
-        { onionNonce    = nonce
-        , onionForward  = fwd
-        , pathFromOwner = NoReturnPath
-        }
-wrapForRoute crypto msg ni r@OnionRoute{routeRelayPort = Just tcpport} = do
+wrapForRoute crypto msg ni r@OnionRoute{routeRelayPort=Nothing} =
+    wrapIndirectHops crypto msg ni r $ \nonce saddr msg' -> do
+        fwd <- wrapOnion crypto (routeAliasA r)
+                                nonce
+                                (id2key . nodeId $ routeNodeA r)
+                                saddr
+                                msg'
+        return $ Right OnionRequest { onionNonce    = nonce
+                                    , onionForward  = fwd
+                                    , pathFromOwner = NoReturnPath
+                                    }
+wrapForRoute crypto msg ni r@OnionRoute { routeRelayPort = Just _ } =
+    wrapIndirectHops crypto msg ni r $ \nonce saddr fwd ->
+        return $ Left $ TCP.OnionPacket nonce $ Addressed (nodeAddr $ routeNodeB r) fwd
+
+wrapIndirectHops :: TransportCrypto
+                      -> OnionMessage Encrypted
+                      -> NodeInfo
+                      -> OnionRoute
+                      -> (Nonce24 -> SockAddr -> Forwarding N2 (OnionMessage Encrypted) -> IO a)
+                      -> IO a
+wrapIndirectHops crypto msg ni r fin = do
     let nonce = msgNonce msg
     fwd <- wrapOnion crypto (routeAliasB r)
                              nonce
@@ -120,7 +125,7 @@ wrapForRoute crypto msg ni r@OnionRoute{routeRelayPort = Just tcpport} = do
                              (id2key . nodeId $ routeNodeC r)
                              (nodeAddr ni)
                              (NotForwarded msg)
-    return $ Left $ TCP.OnionPacket nonce $ Addressed (nodeAddr $ routeNodeB r) fwd
+    fin nonce (nodeAddr $ routeNodeB r) fwd
 
 unwrapAnnounceResponse :: Maybe NodeId -> NodeInfo -> AnnounceResponse -> ([NodeInfo], [Rendezvous], Maybe Nonce32)
 unwrapAnnounceResponse alias ni (AnnounceResponse is_stored (SendNodes ns0)) | let ns = map udpNodeInfo ns0