4 files changed, 38 insertions, 31 deletions

diff --git a/dht/src/Data/Tox/Onion.hs b/dht/src/Data/Tox/Onion.hs
index e0d7c744..e19f71b6 100644
--- a/dht/src/Data/Tox/Onion.hs
+++ b/dht/src/Data/Tox/Onion.hs
@@ -198,7 +198,7 @@ parseOnionAddr lookupSender (msg,saddr)
                       (return . Left . \od -> (msg,od))
                       maddr
             Just (Right msg@(OnionToRouteResponse asym)) -> do
-                let ni = asymNodeInfo saddr asym
+                let ni = asymNodeInfo nodeInfo saddr asym
                 return $ Left (msg, OnionDestination SearchingAlias ni Nothing)
             _ -> return right
 
diff --git a/dht/src/Network/Tox.hs b/dht/src/Network/Tox.hs
index 4b66cfc8..5d27f34f 100644
--- a/dht/src/Network/Tox.hs
+++ b/dht/src/Network/Tox.hs
@@ -351,7 +351,7 @@ newToxOverTransport keydb addr onNewSession (crypto,roster) udp tcp = do
                 (fromMaybe (\_ _ -> return ()) tcp)
     sessions <- initSessions (sendMessage cryptonet)
 
-    let dhtnet0 = layerTransportM (DHT.decrypt crypto) (DHT.encrypt crypto) dhtcrypt
+    let dhtnet0 = layerTransportM (DHT.decrypt crypto nodeId) (DHT.encrypt crypto nodeId) dhtcrypt
         tbl4 = DHT.routing4 $ mkrouting (error "missing client")
         tbl6 = DHT.routing6 $ mkrouting (error "missing client")
         updateOnion bkts tr = hookBucketList DHT.toxSpace bkts orouter (trampolinesUDP orouter) tr
diff --git a/dht/src/Network/Tox/DHT/Transport.hs b/dht/src/Network/Tox/DHT/Transport.hs
index 7414343d..0583c9a3 100644
--- a/dht/src/Network/Tox/DHT/Transport.hs
+++ b/dht/src/Network/Tox/DHT/Transport.hs
@@ -58,8 +58,8 @@ import Data.Word
 import GHC.Generics
 import Network.Socket
 
-type DHTTransport = Transport String NodeInfo (DHTMessage Encrypted8)
-type HandleHi a = Arrival String NodeInfo (DHTMessage Encrypted8) -> IO a
+type DHTTransport ni = Transport String ni (DHTMessage Encrypted8)
+type HandleHi ni a = Arrival String ni (DHTMessage Encrypted8) -> IO a
 
 
 data DHTMessage (f :: * -> *)
@@ -95,32 +95,36 @@ mapMessage f (DHTLanDiscovery nid)     = Nothing
 instance Sized Ping where size = ConstSize 1
 instance Sized Pong where size = ConstSize 1
 
-parseDHTAddr :: TransportCrypto -> (ByteString, SockAddr) -> IO (Either (DHTMessage Encrypted8,NodeInfo) (ByteString,SockAddr))
-parseDHTAddr crypto (msg,saddr)
+parseDHTAddr :: Eq saddr =>
+                    STM [(saddr, (Int, PublicKey))]
+                    -> (NodeId -> saddr -> Either String ni)
+                    -> (ByteString, saddr)
+                    -> IO (Either (DHTMessage Encrypted8,ni) (ByteString,saddr))
+parseDHTAddr pendingCookies nodeInfo (msg,saddr)
     | Just (typ,bs) <- B.uncons msg
     , let right = return $ Right (msg,saddr)
           left  = either (const right) (return . Left)
     = case typ of
-    0x00 -> left $ direct bs saddr DHTPing
-    0x01 -> left $ direct bs saddr DHTPong
-    0x02 -> left $ direct bs saddr DHTGetNodes
-    0x04 -> left $ direct bs saddr DHTSendNodes
-    0x18 -> left $ direct bs saddr DHTCookieRequest
+    0x00 -> left $ direct nodeInfo bs saddr DHTPing
+    0x01 -> left $ direct nodeInfo bs saddr DHTPong
+    0x02 -> left $ direct nodeInfo bs saddr DHTGetNodes
+    0x04 -> left $ direct nodeInfo bs saddr DHTSendNodes
+    0x18 -> left $ direct nodeInfo bs saddr DHTCookieRequest
     0x19 -> do
-        cs <- atomically $ readTVar (pendingCookies crypto)
-        let ni = fromMaybe (noReplyAddr saddr) $ do
+        cs <- atomically pendingCookies
+        let ni = fromMaybe (noReplyAddr nodeInfo saddr) $ do
                     (cnt,key) <- lookup saddr cs <|> listToMaybe (map snd cs)
                     either (const Nothing) Just $ nodeInfo (key2id key) saddr
         left $ fanGet bs getCookie (uncurry DHTCookie) (const $ ni)
-    0x20 -> left $ fanGet bs getDHTReqest (uncurry DHTDHTRequest) (asymNodeInfo saddr . snd)
+    0x20 -> left $ fanGet bs getDHTReqest (uncurry DHTDHTRequest) (asymNodeInfo nodeInfo saddr . snd)
     0x21 -> left $ do
                nid <- runGet get bs
                ni <- nodeInfo nid saddr
                return (DHTLanDiscovery nid, ni)
     _ -> right
 
-encodeDHTAddr :: (DHTMessage Encrypted8,NodeInfo) -> IO (ByteString, SockAddr)
-encodeDHTAddr (msg,ni) = return (runPut $ putMessage msg, nodeAddr ni)
+encodeDHTAddr :: (ni -> saddr) -> (DHTMessage Encrypted8,ni) -> IO (ByteString, saddr)
+encodeDHTAddr nodeAddr (msg,ni) = return (runPut $ putMessage msg, nodeAddr ni)
 
 dhtMessageType :: ( Serialize (f DHTRequest)
                   , Serialize (f (Cookie Encrypted)), Serialize (f CookieRequest)
@@ -159,23 +163,24 @@ getDHT = getAsymm
 
 
 -- Throws an error if called with a non-internet socket.
-direct :: Sized a => ByteString
-                  -> SockAddr
+direct :: Sized a => (NodeId -> saddr -> Either String ni)
+                  -> ByteString
+                  -> saddr
                   -> (Asymm (Encrypted8 a) -> DHTMessage Encrypted8)
-                  -> Either String (DHTMessage Encrypted8, NodeInfo)
-direct bs saddr f = fanGet bs getDHT f (asymNodeInfo saddr)
+                  -> Either String (DHTMessage Encrypted8, ni)
+direct nodeInfo bs saddr f = fanGet bs getDHT f (asymNodeInfo nodeInfo saddr)
 
 -- Throws an error if called with a non-internet socket.
-asymNodeInfo :: SockAddr -> Asymm a -> NodeInfo
-asymNodeInfo saddr asym = either (error . mappend "asymNodeInfo: ") id $ nodeInfo (key2id $ senderKey asym) saddr
+asymNodeInfo :: (NodeId -> saddr -> Either String ni) -> saddr -> Asymm a -> ni
+asymNodeInfo nodeInfo saddr asym = either (error . mappend "asymNodeInfo: ") id $ nodeInfo (key2id $ senderKey asym) saddr
 
 
 fanGet :: ByteString -> Get x -> (x -> a) -> (x -> b) -> Either String (a,b)
 fanGet bs getIt f nid = fmap (f &&& nid) $ runGet getIt bs
 
 -- Throws an error if called with a non-internet socket.
-noReplyAddr :: SockAddr -> NodeInfo
-noReplyAddr saddr = either (error . mappend "noReplyAddr: ") id $ nodeInfo zeroID saddr
+noReplyAddr :: (NodeId -> saddr -> Either String ni) -> saddr -> ni
+noReplyAddr nodeInfo saddr = either (error . mappend "noReplyAddr: ") id $ nodeInfo zeroID saddr
 
 
 data DHTRequest
@@ -396,10 +401,10 @@ instance Serialize CookieRequest where
     get = CookieRequest <$> getPublicKey <* {- padding -} getPublicKey
     put (CookieRequest k) = putPublicKey k >> {- padding -} putPublicKey k
 
-forwardDHTRequests :: TransportCrypto -> (PublicKey -> IO (Maybe NodeInfo)) -> DHTTransport -> DHTTransport
+forwardDHTRequests :: TransportCrypto -> (PublicKey -> IO (Maybe ni)) -> DHTTransport ni -> DHTTransport ni
 forwardDHTRequests crypto closeLookup dht = dht { awaitMessage = await' }
  where
-    await' :: HandleHi a -> STM (IO a)
+    -- await' :: HandleHi ni a -> STM (IO a)
     await' pass = awaitMessage dht $ \case
         Arrival src m@(DHTDHTRequest target payload) | target /= transportPublic crypto
             -> do mni <- closeLookup target
@@ -408,8 +413,8 @@ forwardDHTRequests crypto closeLookup dht = dht { awaitMessage = await' }
                   join $ atomically (await' pass)
         m -> pass m
 
-encrypt :: TransportCrypto -> DHTMessage ((,) Nonce8) -> NodeInfo -> IO (DHTMessage Encrypted8, NodeInfo)
-encrypt crypto msg ni = do
+encrypt :: TransportCrypto -> (ni -> NodeId) -> DHTMessage ((,) Nonce8) -> ni -> IO (DHTMessage Encrypted8, ni)
+encrypt crypto nodeId msg ni = do
     let cipher n plain = Composed $ encryptMessage crypto (id2key $ nodeId ni) n plain
     m <- sequenceMessage $ transcode cipher msg
     return (m, ni)
@@ -423,8 +428,8 @@ encryptMessage crypto destKey n arg = do
     secret <- lookupSharedSecret crypto (transportSecret crypto) destKey n
     return $ E8 $ ToxCrypto.encrypt secret plain
 
-decrypt :: TransportCrypto -> DHTMessage Encrypted8 -> NodeInfo -> IO (Either String (DHTMessage ((,) Nonce8), NodeInfo))
-decrypt crypto msg ni = do
+decrypt :: TransportCrypto -> (ni -> NodeId) -> DHTMessage Encrypted8 -> ni -> IO (Either String (DHTMessage ((,) Nonce8), ni))
+decrypt crypto nodeId msg ni = do
     let decipher n c = Composed $ decryptMessage crypto n . left ((,) $ id2key $ nodeId ni) $ c
     msg' <- sequenceMessage $ transcode decipher msg
     return $ fmap (, ni) $ sequenceMessage msg'
diff --git a/dht/src/Network/Tox/Transport.hs b/dht/src/Network/Tox/Transport.hs
index fef8e8bb..7728ba7a 100644
--- a/dht/src/Network/Tox/Transport.hs
+++ b/dht/src/Network/Tox/Transport.hs
@@ -16,6 +16,7 @@ import Network.Tox.Onion.Transport
 import Network.Tox.Crypto.Transport
 import Network.Tox.Onion.Routes
 
+import Control.Concurrent.STM
 import Network.Socket
 
 toxTransport ::
@@ -33,7 +34,8 @@ toxTransport ::
                 , Transport String SockAddr (Handshake Encrypted))
 toxTransport crypto orouter closeLookup addr udp tcp2server tcp2client = do
     (netcrypto, udp0) <- partitionTransport parseCrypto encodeCrypto udp
-    (dht,udp1) <- partitionTransportM (parseDHTAddr crypto) (fmap Just . encodeDHTAddr)
+    (dht,udp1) <- partitionTransportM (parseDHTAddr (readTVar $ pendingCookies crypto) nodeInfo)
+                                      (fmap Just . encodeDHTAddr nodeAddr)
                                     $ forwardOnions crypto addr udp0 tcp2client
     (onion1,udp2) <- partitionAndForkTransport tcp2server
                              (parseOnionAddr $ lookupSender orouter)