1 files changed, 17 insertions, 15 deletions

diff --git a/dht/src/Network/Tox.hs b/dht/src/Network/Tox.hs
index 97b97bad..69c56e24 100644
--- a/dht/src/Network/Tox.hs
+++ b/dht/src/Network/Tox.hs
@@ -278,10 +278,10 @@ newOnionClient crypto net r toks keydb orouter map_var store load = c
 newTox :: TVar Onion.AnnouncedKeys   -- ^ Store of announced keys we are a rendezvous for.
        -> [String]                   -- ^ Bind-address to listen on. Must provide at least one.
        -> ( ContactInfo extra -> SockAddr -> Session -> IO () )
-       -> Maybe SecretKey            -- ^ Optional DHT secret key to use.
-        -> ( Int -> Onion.OnionMessage Encrypted -> IO () ) -- ^ TCP-bound onion responses.
+       -> (TransportCrypto, ContactInfo extra)
+       -> ( Int -> Onion.OnionMessage Encrypted -> IO () ) -- ^ TCP-bound onion responses.
        -> IO (Tox extra)
-newTox keydb bindspecs onsess suppliedDHTKey tcp = do
+newTox keydb bindspecs onsess crypto tcp = do
     addrs <- mapM (`getBindAddress` True) bindspecs
     let tryBind addr next _ = udpTransport' addr `catchIOError` (next . Just)
         failedBind mbe = do
@@ -291,21 +291,14 @@ newTox keydb bindspecs onsess suppliedDHTKey tcp = do
             throwIO $ userError "Tox UDP listen port?"
     (udp,sock) <- foldr tryBind failedBind addrs Nothing
     addr <- getSocketName sock
-    (relay,sendTCP) <- tcpRelay addr (\a x -> sendMessage udp a $ S.runPut $ Onion.putRequest x)
-    tox <- newToxOverTransport keydb addr onsess suppliedDHTKey udp sendTCP
+    (relay,sendTCP) <- tcpRelay (fst crypto) addr (\a x -> sendMessage udp a $ S.runPut $ Onion.putRequest x)
+    tox <- newToxOverTransport keydb addr onsess crypto udp sendTCP
     return tox { toxAnnounceToLan = announceToLan sock (key2id $ transportPublic $ toxCryptoKeys tox)
                , toxRelayServer   = Just relay
                }
 
--- | This version of 'newTox' is useful for automated tests using 'testPairTransport'.
-newToxOverTransport :: TVar Onion.AnnouncedKeys
-        -> SockAddr
-        -> ( ContactInfo extra -> SockAddr -> Session -> IO () )
-        -> Maybe SecretKey
-        -> Onion.UDPTransport
-        -> ( Int -> Onion.OnionMessage Encrypted -> IO () ) -- ^ TCP-bound onion responses.
-        -> IO (Tox extra)
-newToxOverTransport keydb addr onNewSession suppliedDHTKey udp tcp = do
+newToxCrypto :: Maybe SecretKey -> IO (TransportCrypto, ContactInfo extra)
+newToxCrypto suppliedDHTKey = do
     roster <- newContactInfo
     crypto0 <- newCrypto
     let -- patch in supplied DHT key
@@ -316,12 +309,21 @@ newToxOverTransport keydb addr onNewSession suppliedDHTKey udp tcp = do
                 , transportPublic = toPublic k
                 }
         -- patch in newly allocated roster state.
-        crypto = crypto1 { userKeys = myKeyPairs roster }
     forM_ suppliedDHTKey $ \k -> do
         maybe (dput XMisc "failed to encode suppliedDHTKey")
               (dputB XMisc . C8.append "Using suppliedDHTKey: ")
               $ encodeSecret k
+    return (crypto1 { userKeys = myKeyPairs roster }, roster )
 
+-- | This version of 'newTox' is useful for automated tests using 'testPairTransport'.
+newToxOverTransport :: TVar Onion.AnnouncedKeys
+        -> SockAddr
+        -> ( ContactInfo extra -> SockAddr -> Session -> IO () )
+        -> (TransportCrypto, ContactInfo extra)
+        -> Onion.UDPTransport
+        -> ( Int -> Onion.OnionMessage Encrypted -> IO () ) -- ^ TCP-bound onion responses.
+        -> IO (Tox extra)
+newToxOverTransport keydb addr onNewSession (crypto,roster) udp tcp = do
     drg <- drgNew
     let lookupClose _ = return Nothing