2 files changed, 31 insertions, 10 deletions

diff --git a/examples/dht.hs b/examples/dht.hs
index feeee9ff..16b12678 100644
--- a/examples/dht.hs
+++ b/examples/dht.hs
@@ -1,6 +1,6 @@
 {-# LANGUAGE NondecreasingIndentation #-}
 import Control.Monad
-import Control.Monad.Fix
+import Data.Function
 import Control.Monad.IO.Class
 import Data.Char
 import Network.Socket as Socket
@@ -34,16 +34,21 @@ hReadInt h = do
 -- and /.../ is the sequence of characters
 --
 -- Note: The first byte after the count is ignored and discarded.
-readResponse :: Handle -> IO String
+readResponse :: Handle -> IO (Char, String)
 readResponse h = do
+    c <- hGetChar h
     n <- hReadInt h
-    sequence $ replicate n (hGetChar h)
+    s <- sequence $ replicate n (hGetChar h)
+    return (c,s)
 
 -- | Send a command to the dhtd daemon and then print the response.
 sendCommand :: Handle -> String -> InputT IO ()
-sendCommand h cmd = do resp <- liftIO $ do hPutStrLn h cmd
-                                           readResponse h
-                       outputStrLn resp
+sendCommand h cmd = do liftIO $ hPutStrLn h cmd
+                       fix $ \again -> do
+                           (c, resp) <- liftIO $ readResponse h
+                           if c /= '.'
+                            then outputStr resp >> again
+                            else outputStrLn resp
 
 -- | Get one line of input and send it to the daemon, then run the
 -- passed continuation if it wasn't "quit".
diff --git a/examples/dhtd.hs b/examples/dhtd.hs
index 1b02a37a..68074e2b 100644
--- a/examples/dhtd.hs
+++ b/examples/dhtd.hs
@@ -33,6 +33,7 @@ import System.Mem
 import Data.Torrent (InfoHash)
 import Network.BitTorrent.Address
 import Network.BitTorrent.DHT
+import Network.BitTorrent.DHT.Search
 import Network.BitTorrent.DHT.Query
 import Network.BitTorrent.DHT.Message (FindNode(..),NodeFound(..),GetPeers(..),GotPeers(..))
 import Network.KRPC.Manager (QueryFailure(..))
@@ -49,6 +50,7 @@ import Data.Time.Clock
 #else
 import Control.Concurrent
 #endif
+import Control.Concurrent.STM
 import System.Environment
 
 mkNodeAddr :: SockAddr -> NodeAddr IPv4
@@ -130,7 +132,10 @@ marshalForClient :: String -> String
 marshalForClient s = show (length s) ++ ":" ++ s
 
 hPutClient :: Handle -> String -> IO ()
-hPutClient h s = hPutStr h (marshalForClient s)
+hPutClient h s = hPutStr h ('.' : marshalForClient s)
+
+hPutClientChunk :: Handle -> String -> IO ()
+hPutClientChunk h s = hPutStr h (' ' : marshalForClient s)
 
 clientSession :: Node IPv4 -> MVar () -> RestrictedSocket -> Int -> Handle -> IO ()
 clientSession st signalQuit sock n h = do
@@ -280,9 +285,20 @@ clientSession st signalQuit sock n h = do
         ("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
+                            (tid, s) <- isearch ioGetPeers ih
+                            flip fix Set.empty $ \again shown -> do
+                                (chk,fin) <- liftIO . atomically $ do
+                                    r <- (Set.\\ shown) <$> readTVar (searchResults s)
+                                    if not $ Set.null r
+                                        then (,) r <$> searchIsFinished s
+                                        else searchIsFinished s >>= check >> return (Set.empty,True)
+                                let ps = case Set.toList chk of
+                                            [] -> ""
+                                            _  -> unlines $ map (show . pPrint) $ Set.toList chk
+                                if fin then return $ hPutClient h ps
+                                       else do
+                                            liftIO $ hPutClientChunk h ps
+                                            again (shown `Set.union` chk)
                         Left er -> return $ hPutClient h er
 
         _    -> cmd0 $ hPutClient h "error."