1 files changed, 17 insertions, 6 deletions

diff --git a/dht/ToxManager.hs b/dht/ToxManager.hs
index 403167c2..0363f7f0 100644
--- a/dht/ToxManager.hs
+++ b/dht/ToxManager.hs
@@ -35,6 +35,7 @@ import DebugTag
 import Foreign.Storable
 import HandshakeCache
 import Network.Address
+import Network.Kademlia.Bootstrap
 import qualified Network.Kademlia.Routing    as R
          ;import Network.Kademlia.Routing    as R
 import Network.Kademlia.Search
@@ -51,7 +52,7 @@ import Network.Tox.NodeId
 import qualified Network.Tox.Onion.Handlers  as Tox
 import qualified Network.Tox.Onion.Transport as Tox
          ;import Network.Tox.Onion.Transport (OnionData (..))
-import Network.Tox.Onion.Routes (tcpKademliaClient)
+import Network.Tox.Onion.Routes (tcpKademliaClient,tcpBucketRefresher)
 import qualified Network.Tox.TCP as TCP
 import Presence
 import Text.Read
@@ -310,6 +311,11 @@ connectViaRelay tx theirKey theirDhtKey ann tkey now = do
 
 gotDhtPubkey :: Tox.DHTPublicKey -> ToxToXMPP -> PublicKey -> IO ()
 gotDhtPubkey theirDhtKey tx theirKey = do
+  dput XNetCrypto $ unlines $
+    [ "Recieved DHTKey from " ++ show (Tox.key2id theirKey)
+    , "  DHT: " ++ show target
+    ] ++ let Tox.SendNodes ns = Tox.dhtpkNodes theirDhtKey
+         in [ "  relay " ++ show n | n <- ns ]
   contact <- atomically $ getContact theirKey (txAccount tx) >>= mapM (readTVar . contactLastSeenAddr)
   forM_ contact $ \lastSeen -> do
     case lastSeen of
@@ -591,13 +597,18 @@ checkSoliciting presence me them contact = do
 
 nearNodes :: Tox extra -> NodeId -> STM [NodeInfo]
 nearNodes tox nid = do
+  let refresher = tcpBucketRefresher $ toxOnionRoutes tox
   bkts4 <- readTVar $ routing4 $ toxRouting tox
   bkts6 <- readTVar $ routing6 $ toxRouting tox
-  let nss =
-        map
-          (R.kclosest (searchSpace $ toxQSearch tox) (searchK $ toxQSearch tox) nid)
-          [bkts4, bkts6]
-  return $ foldr interweave [] nss
+  bktsTCP <- readTVar $ refreshBuckets refresher
+  let udpSpace = searchSpace $ toxQSearch tox
+      k = searchK $ toxQSearch tox
+      nss = map (R.kclosest udpSpace k nid) [bkts4, bkts6]
+      ns_UDP = foldr interweave [] nss
+      ns_TCP = case refresher of
+                BucketRefresher{refreshSearch=sch}
+                    -> R.kclosest (searchSpace sch) k nid bktsTCP
+  return $ ns_UDP ++ map TCP.udpNodeInfo ns_TCP
 
 startConnecting0 :: ToxToXMPP -> PublicKey -> Contact -> String -> IO ()
 startConnecting0 tx them contact reason = do