diff options
Diffstat (limited to 'examples/dhtd.hs')
-rw-r--r-- | examples/dhtd.hs | 24 |
1 files changed, 20 insertions, 4 deletions
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 | |||
33 | import Data.Torrent (InfoHash) | 33 | import Data.Torrent (InfoHash) |
34 | import Network.BitTorrent.Address | 34 | import Network.BitTorrent.Address |
35 | import Network.BitTorrent.DHT | 35 | import Network.BitTorrent.DHT |
36 | import Network.BitTorrent.DHT.Search | ||
36 | import Network.BitTorrent.DHT.Query | 37 | import Network.BitTorrent.DHT.Query |
37 | import Network.BitTorrent.DHT.Message (FindNode(..),NodeFound(..),GetPeers(..),GotPeers(..)) | 38 | import Network.BitTorrent.DHT.Message (FindNode(..),NodeFound(..),GetPeers(..),GotPeers(..)) |
38 | import Network.KRPC.Manager (QueryFailure(..)) | 39 | import Network.KRPC.Manager (QueryFailure(..)) |
@@ -49,6 +50,7 @@ import Data.Time.Clock | |||
49 | #else | 50 | #else |
50 | import Control.Concurrent | 51 | import Control.Concurrent |
51 | #endif | 52 | #endif |
53 | import Control.Concurrent.STM | ||
52 | import System.Environment | 54 | import System.Environment |
53 | 55 | ||
54 | mkNodeAddr :: SockAddr -> NodeAddr IPv4 | 56 | mkNodeAddr :: SockAddr -> NodeAddr IPv4 |
@@ -130,7 +132,10 @@ marshalForClient :: String -> String | |||
130 | marshalForClient s = show (length s) ++ ":" ++ s | 132 | marshalForClient s = show (length s) ++ ":" ++ s |
131 | 133 | ||
132 | hPutClient :: Handle -> String -> IO () | 134 | hPutClient :: Handle -> String -> IO () |
133 | hPutClient h s = hPutStr h (marshalForClient s) | 135 | hPutClient h s = hPutStr h ('.' : marshalForClient s) |
136 | |||
137 | hPutClientChunk :: Handle -> String -> IO () | ||
138 | hPutClientChunk h s = hPutStr h (' ' : marshalForClient s) | ||
134 | 139 | ||
135 | clientSession :: Node IPv4 -> MVar () -> RestrictedSocket -> Int -> Handle -> IO () | 140 | clientSession :: Node IPv4 -> MVar () -> RestrictedSocket -> Int -> Handle -> IO () |
136 | clientSession st signalQuit sock n h = do | 141 | clientSession st signalQuit sock n h = do |
@@ -280,9 +285,20 @@ clientSession st signalQuit sock n h = do | |||
280 | ("search-peers", s) -> cmd $ do | 285 | ("search-peers", s) -> cmd $ do |
281 | case readEither s of | 286 | case readEither s of |
282 | Right ih -> do | 287 | Right ih -> do |
283 | ps <- Set.toList <$> isearch ioGetPeers ih | 288 | (tid, s) <- isearch ioGetPeers ih |
284 | return $ do | 289 | flip fix Set.empty $ \again shown -> do |
285 | hPutClient h $ unlines $ map (show . pPrint) ps | 290 | (chk,fin) <- liftIO . atomically $ do |
291 | r <- (Set.\\ shown) <$> readTVar (searchResults s) | ||
292 | if not $ Set.null r | ||
293 | then (,) r <$> searchIsFinished s | ||
294 | else searchIsFinished s >>= check >> return (Set.empty,True) | ||
295 | let ps = case Set.toList chk of | ||
296 | [] -> "" | ||
297 | _ -> unlines $ map (show . pPrint) $ Set.toList chk | ||
298 | if fin then return $ hPutClient h ps | ||
299 | else do | ||
300 | liftIO $ hPutClientChunk h ps | ||
301 | again (shown `Set.union` chk) | ||
286 | Left er -> return $ hPutClient h er | 302 | Left er -> return $ hPutClient h er |
287 | 303 | ||
288 | _ -> cmd0 $ hPutClient h "error." | 304 | _ -> cmd0 $ hPutClient h "error." |