1 files changed, 11 insertions, 5 deletions

diff --git a/dht/src/Network/Tox/Transport.hs b/dht/src/Network/Tox/Transport.hs
index b017f7be..e90917f6 100644
--- a/dht/src/Network/Tox/Transport.hs
+++ b/dht/src/Network/Tox/Transport.hs
@@ -18,18 +18,24 @@ import Network.Tox.Onion.Transport
 import Network.Tox.Crypto.Transport
 import Network.Tox.Onion.Routes
 
+import Control.Applicative
 import Control.Concurrent.STM
 import qualified Data.ByteString    as B
 import qualified Data.Dependent.Map as DMap
 import Data.Dependent.Sum
 import Data.Functor.Identity
+import Data.Maybe
 import Network.Socket
 
-pendingCookiesUDP :: TransportCrypto -> STM [(SockAddr, (Int, PublicKey))]
-pendingCookiesUDP crypto = readTVar $ pendingCookies crypto
+pendingCookiesUDP :: TransportCrypto -> SockAddr -> STM (Maybe UDP.NodeInfo)
+pendingCookiesUDP crypto saddr = do
+    cs <- readTVar $ pendingCookies crypto
+    return $ do
+        (_,key) <- lookup saddr cs <|> listToMaybe (map snd cs)
+        either (const Nothing) Just $ nodeInfo (key2id key) saddr
 
-pendingCookiesTCP :: TransportCrypto -> STM [(ViaRelay, (Int, PublicKey))]
-pendingCookiesTCP crypto = return [] -- TODO
+pendingCookiesTCP :: ViaRelay -> STM (Maybe ViaRelay)
+pendingCookiesTCP ni = return $ Just ni
 
 toxTransport ::
       TransportCrypto
@@ -52,7 +58,7 @@ toxTransport crypto orouter closeLookup addr udp relaynet _ tcp2client = do
     -- rlynet0 = layerTransportM (DHT.decrypt crypto Multi.relayNodeId) (DHT.encrypt crypto Multi.relayNodeId) relaynet
     (netcryptoTCP, relaynet0) <- partitionTransport parseCrypto encodeCrypto relaynet
     (dhtTCP,relaynet1) <- partitionTransportM
-                            (parseDHTAddr (pendingCookiesTCP crypto) (\nid viarelay -> Right viarelay))
+                            (parseDHTAddr pendingCookiesTCP (\nid viarelay -> Right viarelay))
                             (fmap Just . encodeDHTAddr id)
                             relaynet0
     let _ = dhtTCP :: Transport String ViaRelay (DHTMessage Encrypted8)