3 files changed, 145 insertions, 14 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
diff --git a/src/Network/BitTorrent/DHT/Routing.hs b/src/Network/BitTorrent/DHT/Routing.hs
index 5c9788dc..cf4a4de3 100644
--- a/src/Network/BitTorrent/DHT/Routing.hs
+++ b/src/Network/BitTorrent/DHT/Routing.hs
@@ -442,16 +442,16 @@ size = L.sum . shape
 depth :: Table ip -> BucketCount
 depth = L.length . shape
 
-lookupBucket :: NodeId -> Table ip -> Maybe (Bucket ip)
-lookupBucket nid = go 0
+lookupBucket :: NodeId -> Table ip -> [Bucket ip]
+lookupBucket nid = go 0 []
   where
-    go i (Zero table bucket)
-      |  testIdBit nid i  = pure bucket
-      |     otherwise     = go (succ i) table
-    go i (One  bucket table)
-      |  testIdBit nid i  = go (succ i) table
-      |     otherwise     = pure bucket
-    go _ (Tip _ _ bucket) = pure bucket
+    go i bs (Zero table bucket)
+      |  testIdBit nid i     = bucket : toBucketList table ++ bs
+      |     otherwise        = go (succ i) (bucket:bs) table
+    go i bs (One  bucket table)
+      |  testIdBit nid i     = go (succ i) (bucket:bs) table
+      |     otherwise        = bucket : toBucketList table ++ bs
+    go _ bs (Tip _ _ bucket) = bucket : bs
 
 compatibleNodeId :: Table ip -> IO NodeId
 compatibleNodeId tbl = genBucketSample prefix br
@@ -504,11 +504,14 @@ instance TableKey InfoHash where
 -- | Get a list of /K/ closest nodes using XOR metric. Used in
 -- 'find_node' and 'get_peers' queries.
 kclosest :: Eq ip => TableKey a => K -> a -> Table ip -> [NodeInfo ip]
-kclosest k (toNodeId -> nid)
-  = L.take k . rank nodeId nid
-  . L.map PSQ.key . PSQ.toList . fromMaybe PSQ.empty
-  . fmap bktNodes
-  . lookupBucket nid
+kclosest k (toNodeId -> nid) tbl = take k $ rank nodeId nid (L.concat bucket)
+                                          ++ rank nodeId nid (L.concat everyone)
+ where
+    (bucket,everyone) =
+        L.splitAt 1
+        . L.map (L.map PSQ.key . PSQ.toList . bktNodes)
+        . lookupBucket nid
+        $ tbl
 
 {-----------------------------------------------------------------------
 --  Routing
diff --git a/src/Network/BitTorrent/DHT/Search.hs b/src/Network/BitTorrent/DHT/Search.hs
new file mode 100644
index 00000000..1fe73c30
--- /dev/null
+++ b/src/Network/BitTorrent/DHT/Search.hs
@@ -0,0 +1,92 @@
+{-# LANGUAGE PatternSynonyms #-}
+{-# LANGUAGE RecordWildCards #-}
+module Network.BitTorrent.DHT.Search where
+
+import Control.Concurrent
+import Control.Concurrent.Async.Pool
+import Control.Concurrent.STM
+import Control.Exception
+import Control.Monad
+import Data.Bool
+import Data.Function
+import Data.List
+import qualified Data.Map.Strict     as Map
+         ;import Data.Map.Strict     (Map)
+import Data.Maybe
+import qualified Data.Set            as Set
+         ;import Data.Set            (Set)
+import System.IO
+
+import qualified Data.MinMaxPSQ   as MM
+         ;import Data.MinMaxPSQ   (MinMaxPSQ)
+import qualified Data.Wrapper.PSQ as PSQ
+         ;import Data.Wrapper.PSQ (pattern (:->), Binding, PSQ)
+import Network.BitTorrent.Address
+
+data IterativeSearch ip r = IterativeSearch
+    { searchTarget       :: NodeId
+    , searchQuery        :: NodeInfo ip -> IO ([NodeInfo ip], [r])
+    , searchPendingCount :: TVar Int
+    , searchQueued       :: TVar (MinMaxPSQ (NodeInfo ip) NodeDistance)
+    , searchInformant    :: TVar (MinMaxPSQ (NodeInfo ip) NodeDistance)
+    , searchVisited      :: TVar (Set (NodeAddr ip))
+    , searchResults      :: TVar (Set r)
+    }
+
+newSearch :: Eq ip => (NodeInfo ip -> IO ([NodeInfo ip], [r]))
+                      -> NodeId -> [NodeInfo ip] -> IO (IterativeSearch ip r)
+newSearch qry target ns = atomically $ do
+    c <- newTVar 0
+    q <- newTVar $ MM.fromList $ map (\n -> n :-> distance target (nodeId n)) ns
+    i <- newTVar MM.empty
+    v <- newTVar Set.empty
+    r <- newTVar Set.empty
+    return $ IterativeSearch target qry c q i v r
+
+searchAlpha :: Int
+searchAlpha = 3
+
+searchK :: Int
+searchK = 8
+
+sendQuery :: (Ord a, Ord t) =>
+            IterativeSearch t a
+            -> Binding (NodeInfo t) NodeDistance
+            -> IO ()
+sendQuery IterativeSearch{..} (ni :-> d) = do
+    (ns,rs) <- handle (\(SomeException e) -> return ([],[]))
+                      (searchQuery ni)
+    atomically $ do
+        modifyTVar searchPendingCount pred
+        vs <- readTVar searchVisited
+        -- We only queue a node if it is not yet visited
+        let insertFoundNode n q
+             | nodeAddr n `Set.member` vs = q
+             | otherwise                  = MM.insertTake searchK n (distance searchTarget $ nodeId n) q
+        modifyTVar searchQueued    $ \q -> foldr insertFoundNode q ns
+        modifyTVar searchInformant $ MM.insertTake searchK ni d
+        modifyTVar searchResults   $ \s -> foldr Set.insert s rs
+
+search ::
+    (Ord r, Ord ip) =>
+    IterativeSearch ip r -> IO ()
+search s@IterativeSearch{..} = withTaskGroup searchAlpha $ \g -> do
+    fix $ \again -> do
+        join $ atomically $ do
+            found <- MM.minView <$> readTVar searchQueued
+            cnt <- readTVar $ searchPendingCount
+            case found of
+                Nothing            -> retry
+                Just (ni :-> d, q) -> do
+                    informants <- readTVar searchInformant
+                    if MM.size informants < searchK
+                       && (cnt > 0 || not (MM.null q))
+                       || PSQ.prio (fromJust $ MM.findMax informants) > d
+                    then do
+                        writeTVar searchQueued q
+                        modifyTVar searchVisited $ Set.insert (nodeAddr ni)
+                        modifyTVar searchPendingCount succ
+                        return $ withAsync g (sendQuery s (ni :-> d)) (const again)
+                    else do
+                        check (cnt == 0)
+                        return $ return ()