8 files changed, 236 insertions, 23 deletions

diff --git a/bittorrent.cabal b/bittorrent.cabal
index 2d76e606..f6559752 100644
--- a/bittorrent.cabal
+++ b/bittorrent.cabal
@@ -91,6 +91,8 @@ library
                        Network.SocketLike
   other-modules:       Paths_bittorrent
                        Data.Wrapper.PSQ
+                       Data.MinMaxPSQ
+                       Network.BitTorrent.DHT.Search
   if !flag(dht-only)
      exposed-modules:  Network.BitTorrent
                        Network.BitTorrent.Client
@@ -143,6 +145,7 @@ library
                        -- Concurrency
                      , SafeSemaphore
                      , lifted-async
+                     , async-pool
 --                     , BoundedChan >= 1.0.1.0
                      , split-channel >= 0.2
                      , stm           >= 2.4
@@ -373,6 +376,7 @@ executable dhtd
                      , monad-logger
                      , bittorrent
                      , unix
+                     , containers
   if flag(thread-debug)
     build-depends: time
     cpp-options:   -DTHREAD_DEBUG
diff --git a/examples/dhtd.hs b/examples/dhtd.hs
index 05a3de26..1b02a37a 100644
--- a/examples/dhtd.hs
+++ b/examples/dhtd.hs
@@ -16,6 +16,7 @@ import Data.Default
 import Data.List as L
 import Data.Maybe
 import Data.String
+import qualified Data.Set as Set
 import qualified Data.ByteString as B (ByteString,writeFile,readFile)
         ; import Data.ByteString (ByteString)
 import qualified Data.ByteString.Char8 as B8
@@ -95,8 +96,9 @@ bootstrapNodes = mapMaybe fromAddr
 -- ExtendedCaps (Map.singleton
 
 noDebugPrints :: LogSource -> LogLevel -> Bool
-noDebugPrints _ = \case LevelDebug -> False
-                        _          -> True
+noDebugPrints _ = \case LevelDebug   -> False
+                        LevelOther _ -> False
+                        _            -> True
 
 noLogging :: LogSource -> LogLevel -> Bool
 noLogging _ _ = False
@@ -275,6 +277,14 @@ clientSession st signalQuit sock n h = do
                             hPutClient h $ showReport rs
                         Left er -> return $ hPutClient h er
 
+        ("search-peers", s) -> cmd $ do
+                    case readEither s of
+                        Right ih -> do
+                            ps <- Set.toList <$> isearch ioGetPeers ih
+                            return $ do
+                            hPutClient h $ unlines $ map (show . pPrint) ps
+                        Left er -> return $ hPutClient h er
+
         _    -> cmd0 $ hPutClient h "error."
 
 defaultPort = error "TODO defaultPort"
diff --git a/src/Data/MinMaxPSQ.hs b/src/Data/MinMaxPSQ.hs
new file mode 100644
index 00000000..96937604
--- /dev/null
+++ b/src/Data/MinMaxPSQ.hs
@@ -0,0 +1,64 @@
+{-# LANGUAGE BangPatterns #-}
+module Data.MinMaxPSQ where
+
+import Data.Ord
+import qualified Data.Wrapper.PSQ as PSQ
+         ;import Data.Wrapper.PSQ as PSQ hiding (insert, null, size)
+import Prelude hiding (null, take)
+
+data MinMaxPSQ k p = MinMaxPSQ !(PSQ k p) !(PSQ k (Down p))
+
+empty :: MinMaxPSQ k p
+empty = MinMaxPSQ PSQ.empty PSQ.empty
+
+null :: MinMaxPSQ k p -> Bool
+null (MinMaxPSQ nq xq) = PSQ.null nq
+
+size :: MinMaxPSQ k p -> Int
+size (MinMaxPSQ nq xq) = PSQ.size nq
+
+toList :: (Ord k, Ord p) => MinMaxPSQ k p -> [Binding k p]
+toList (MinMaxPSQ nq xq) = PSQ.toList nq
+
+fromList :: (Ord k, Ord p) => [Binding k p] -> MinMaxPSQ k p
+fromList kps = MinMaxPSQ (PSQ.fromList kps)
+                         (PSQ.fromList $ map (\(k :-> p) -> (k :-> Down p)) kps)
+
+findMin :: (Ord k, Ord p) => MinMaxPSQ k p -> Maybe (Binding k p)
+findMin (MinMaxPSQ nq xq) = PSQ.findMin nq
+
+findMax :: (Ord k, Ord p) => MinMaxPSQ k p -> Maybe (Binding k p)
+findMax (MinMaxPSQ nq xq) = fmap (\(k :-> Down p) -> k :-> p) $ PSQ.findMin xq
+
+insert :: (Ord k, Ord p) => k -> p -> MinMaxPSQ k p -> MinMaxPSQ k p
+insert k p (MinMaxPSQ nq xq) = MinMaxPSQ (PSQ.insert k p        nq)
+                                         (PSQ.insert k (Down p) xq)
+
+delete :: (Ord k, Ord p) => k -> MinMaxPSQ k p -> MinMaxPSQ k p
+delete k (MinMaxPSQ nq xq) = MinMaxPSQ (PSQ.delete k nq) (PSQ.delete k xq)
+
+deleteMin :: (Ord k, Ord p) => MinMaxPSQ k p -> MinMaxPSQ k p
+deleteMin (MinMaxPSQ nq xq) = case PSQ.minView nq of
+    Just (k :-> _, nq') -> MinMaxPSQ nq' (PSQ.delete k xq)
+    Nothing             -> MinMaxPSQ nq xq
+
+deleteMax :: (Ord k, Ord p) => MinMaxPSQ k p -> MinMaxPSQ k p
+deleteMax (MinMaxPSQ nq xq) = case PSQ.minView xq of
+    Just (k :-> _, xq') -> MinMaxPSQ (PSQ.delete k nq) xq'
+    Nothing             -> MinMaxPSQ nq xq
+
+minView :: (Ord k, Ord p) => MinMaxPSQ k p -> Maybe (Binding k p, MinMaxPSQ k p)
+minView (MinMaxPSQ nq xq) = fmap (\(k :-> p, nq') -> (k :-> p, MinMaxPSQ nq' (PSQ.delete k xq)))
+                            $ PSQ.minView nq
+
+maxView :: (Ord k, Ord p) => MinMaxPSQ k p -> Maybe (Binding k p, MinMaxPSQ k p)
+maxView (MinMaxPSQ nq xq) = fmap (\(k :-> Down p, xq') -> (k :-> p, MinMaxPSQ (PSQ.delete k nq) xq'))
+                            $ PSQ.minView xq
+
+insertTake :: (Ord k, Ord p) => Int -> k -> p -> MinMaxPSQ k p -> MinMaxPSQ k p
+insertTake n k p q = take n $ insert k p q
+
+take :: (Ord k, Ord p) => Int -> MinMaxPSQ k p -> MinMaxPSQ k p
+take !n !q | (size q <= n) = q
+           | null q        = q
+           | otherwise     = take n $ deleteMax q
diff --git a/src/Data/Wrapper/PSQ.hs b/src/Data/Wrapper/PSQ.hs
index e8fa2d98..2c08011b 100644
--- a/src/Data/Wrapper/PSQ.hs
+++ b/src/Data/Wrapper/PSQ.hs
@@ -15,10 +15,14 @@ type Binding k p = (k,p,())
 pattern (:->) :: k -> p -> Binding k p
 pattern k :-> p <- (k,p,()) where k :-> p = (k,p,())
 
-key :: Binding k v -> k
+key :: Binding k p -> k
 key (k,p,v) = k
 {-# INLINE key #-}
 
+prio :: Binding k p -> p
+prio (k,p,v) = p
+{-# INLINE prio #-}
+
 insert :: (Ord k, Ord p) => k -> p -> PSQ k p -> PSQ k p
 insert k p q = OrdPSQ.insert k p () q
 {-# INLINE insert #-}
diff --git a/src/Network/BitTorrent/Address.hs b/src/Network/BitTorrent/Address.hs
index 381ff50b..a8e12b35 100644
--- a/src/Network/BitTorrent/Address.hs
+++ b/src/Network/BitTorrent/Address.hs
@@ -696,10 +696,10 @@ newtype NodeDistance = NodeDistance BS.ByteString
   deriving (Eq, Ord)
 
 instance Pretty NodeDistance where
-  pPrint (NodeDistance bs) = foldMap bitseq $ BS.unpack bs
-    where
-      listBits w = L.map (testBit w) (L.reverse [0..finiteBitSize w - 1])
-      bitseq = foldMap (int . fromEnum) . listBits
+  pPrint (NodeDistance bs) = text $ BC.unpack (Base16.encode bs)
+
+instance Show NodeDistance where
+  show (NodeDistance bs) = BC.unpack (Base16.encode bs)
 
 -- | distance(A,B) = |A xor B| Smaller values are closer.
 distance :: NodeId -> NodeId -> NodeDistance
@@ -751,7 +751,7 @@ genBucketSample' gen (NodeId self) (q,m,b)
 data NodeAddr a = NodeAddr
   { nodeHost ::                !a
   , nodePort :: {-# UNPACK #-} !PortNumber
-  } deriving (Eq, Typeable, Functor, Foldable, Traversable)
+  } deriving (Eq, Ord, Typeable, Functor, Foldable, Traversable)
 
 instance Show a => Show (NodeAddr a) where
   showsPrec i NodeAddr {..}
@@ -1285,7 +1285,7 @@ getBindAddress listenPortString enabled6 = do
     let (x6s,x4s) = partition (\s -> addrFamily s == AF_INET6) xs
         listenAddr =
           case if enabled6 then x6s++x4s else x4s of
-            AddrInfo { addrAddress = listenAddr } : _ -> listenAddr
+            AddrInfo { addrAddress = addr } : _ -> addr
             _ -> if enabled6
                   then SockAddrInet6 (parsePort listenPortString) 0 iN6ADDR_ANY 0
                   else SockAddrInet (parsePort listenPortString) iNADDR_ANY
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 ()