New improved iterative search algorithm.

1 files changed, 36 insertions, 0 deletions

diff --git a/src/Network/BitTorrent/DHT/Query.hs b/src/Network/BitTorrent/DHT/Query.hs
index 8215c95d..39ef9604 100644
--- a/src/Network/BitTorrent/DHT/Query.hs
+++ b/src/Network/BitTorrent/DHT/Query.hs
@@ -13,6 +13,7 @@
 {-# LANGUAGE FlexibleContexts    #-}
 {-# LANGUAGE ScopedTypeVariables #-}
 {-# LANGUAGE TemplateHaskell     #-}
+{-# LANGUAGE TupleSections       #-}
 module Network.BitTorrent.DHT.Query
        ( -- * Handler
          -- | To bind specific set of handlers you need to pass
@@ -40,6 +41,9 @@ module Network.BitTorrent.DHT.Query
        , Search
        , search
        , publish
+       , ioFindNode
+       , ioGetPeers
+       , isearch
 
        -- ** Routing table
        , insertNode
@@ -67,6 +71,8 @@ import Data.Either
 import Data.List as L
 import Data.Monoid
 import Data.Text as T
+import qualified Data.Set as Set
+         ;import Data.Set (Set)
 import Network
 import Text.PrettyPrint as PP hiding ((<>), ($$))
 import Text.PrettyPrint.HughesPJClass hiding ((<>),($$))
@@ -82,6 +88,7 @@ import Network.BitTorrent.DHT.Message
 import Network.BitTorrent.DHT.Routing as R
 import Network.BitTorrent.DHT.Session
 import Control.Concurrent.STM
+import qualified Network.BitTorrent.DHT.Search as Search
 
 {-----------------------------------------------------------------------
 --  Handlers
@@ -182,6 +189,35 @@ announceQ ih p NodeInfo {..} = do
 --  Iterative queries
 -----------------------------------------------------------------------}
 
+
+ioGetPeers :: Address ip => InfoHash -> DHT ip (NodeInfo ip -> IO ([NodeInfo ip], [PeerAddr ip]))
+ioGetPeers ih = do
+    session <- ask
+    return $ \ni -> runDHT session $ do
+        r <- try $ getPeersQ ih ni
+        case r of
+            Right e -> return $ either (,[]) ([],) e
+            Left e -> let _ = e :: QueryFailure in return ([],[])
+
+ioFindNode :: Address ip => InfoHash -> DHT ip (NodeInfo ip -> IO ([NodeInfo ip], [NodeInfo ip]))
+ioFindNode ih = do
+    session <- ask
+    return $ \ni -> runDHT session $ do
+        NodeFound ns <- FindNode (toNodeId ih) <@> nodeAddr ni
+        return $ L.partition (\n -> nodeId n /= toNodeId ih) ns
+
+isearch :: (Ord r, Ord ip) =>
+            (InfoHash -> DHT ip (NodeInfo ip -> IO ([NodeInfo ip], [r])))
+            -> InfoHash
+            -> DHT ip (Set r)
+isearch f ih = do
+    qry <- f ih
+    ns <- kclosest 8 ih <$> getTable
+    liftIO $ do s <- Search.newSearch qry (toNodeId ih) ns
+                Search.search s
+                atomically $ readTVar (Search.searchResults s)
+
+
 type Search    ip o = Conduit [NodeInfo ip] (DHT ip) [o ip]
 
 -- TODO: use reorder and filter (Traversal option) leftovers