1 files changed, 27 insertions, 9 deletions

diff --git a/examples/dhtd.hs b/examples/dhtd.hs
index 4d9072b5..36b8e294 100644
--- a/examples/dhtd.hs
+++ b/examples/dhtd.hs
@@ -21,6 +21,7 @@ import Control.Concurrent.STM
 import Control.DeepSeq
 import Control.Exception
 import Control.Monad
+import Data.Bool
 import Data.Char
 import Data.Hashable
 import Data.List
@@ -125,6 +126,11 @@ data DHTSearch nid ni = forall addr tok r. DHTSearch
     , searchResults    :: TVar (Set.Set String)
     }
 
+data DHTPing ni = forall r. DHTPing
+    { pingQuery :: ni -> IO (Maybe r)
+    , pingShowResult :: r -> String
+    }
+
 data DHT = forall nid ni. ( Show     ni
                           , Read     ni
                           , ToJSON   ni
@@ -139,7 +145,7 @@ data DHT = forall nid ni. ( Show     ni
                           ) =>
     DHT
     { dhtBuckets       :: TVar (BucketList ni)
-    , dhtPing          :: ni -> IO Bool
+    , dhtPing          :: Map.Map String (DHTPing ni)
     , dhtQuery         :: Map.Map String (DHTQuery nid ni)
     , dhtAnnouncables  :: Map.Map String DHTAnnouncable
     , dhtParseId       :: String -> Either String nid
@@ -185,7 +191,7 @@ fallbackLoad fname = do
 
 
 pingNodes :: String -> DHT -> IO Bool
-pingNodes netname DHT{dhtPing} = do
+pingNodes netname DHT{dhtPing} | Just DHTPing{pingQuery=ping} <- Map.lookup "ping" dhtPing = do
     let fname = nodesFileName netname
     attempt <- tryIOError $ do
         J.decode <$> L.readFile fname
@@ -196,10 +202,11 @@ pingNodes netname DHT{dhtPing} = do
                          myThreadId >>= flip labelThread ("pinging."++fname)
                          putStrLn $ "Forked "++show fname
                          withTaskGroup ("withTaskGroup."++fname) 10 $ \g -> do
-                           mapM_ (\n -> forkTask g (show n) $ void $ dhtPing n)
+                           mapM_ (\n -> forkTask g (show n) $ void $ ping n)
                                  (ns `asTypeOf` [])
                          putStrLn $ "Load finished "++show fname
                        return True
+pingNodes _ _ = return False
 
 
 
@@ -402,11 +409,13 @@ clientSession s@Session{..} sock cnum h = do
                                          , ("node-id", show $ thisNode bkts)
                                          , ("network", netname) ]
 
-        ("ping", s) | Just DHT{dhtPing} <- Map.lookup netname dhts
+        (pinglike, s) | Just DHT{dhtPing} <- Map.lookup netname dhts
+                      , Just DHTPing{ pingQuery=ping
+                                    , pingShowResult=showr } <- Map.lookup pinglike dhtPing
                     -> cmd0 $ do
                         case readEither s of
-                          Right addr -> do result <- dhtPing addr
-                                           let rs = [" ", show result]
+                          Right addr -> do result <- ping addr
+                                           let rs = [" ", maybe "Timeout." showr result]
                                            hPutClient h $ unlines rs
                           Left er -> hPutClient h er
         ("k", s) | ""           <- strp s -> cmd0 $ do
@@ -658,7 +667,10 @@ main = do
             peerPort <- atomically $ newTVar 6881 -- BitTorrent client TCP port.
             let mainlineDHT bkts wantip = DHT
                     { dhtBuckets = bkts btR
-                    , dhtPing    = Mainline.ping bt
+                    , dhtPing    = Map.singleton "ping" $ DHTPing
+                                        { pingQuery = fmap (bool Nothing (Just ())) . Mainline.ping bt
+                                        , pingShowResult = show
+                                        }
                     , dhtQuery   = Map.fromList
                                     [ ("node", DHTQuery (Mainline.nodeSearch bt)
                                                         (\ni -> fmap Mainline.unwrapNodes
@@ -722,7 +734,10 @@ main = do
 
             let toxDHT bkts = DHT
                     { dhtBuckets = bkts (Tox.toxRouting tox)
-                    , dhtPing    = Tox.ping (Tox.toxDHT tox)
+                    , dhtPing    = Map.singleton "ping" $ DHTPing
+                                    { pingQuery = fmap (bool Nothing (Just ())) . Tox.ping (Tox.toxDHT tox)
+                                    , pingShowResult = show
+                                    }
                     , dhtQuery   = Map.fromList
                                     [ ("node", DHTQuery (Tox.nodeSearch $ Tox.toxDHT tox)
                                                         (\ni -> fmap Tox.unwrapNodes
@@ -844,13 +859,16 @@ main = do
     forM_ (Map.toList dhts)
           $ \(netname, dht@DHT { dhtBuckets = bkts
                              , dhtQuery = qrys
-                             , dhtPing = ping
+                             , dhtPing = pings
                              , dhtFallbackNodes = getBootstrapNodes }) -> do
         btSaved <- loadNodes netname --  :: IO [Mainline.NodeInfo]
         putStrLn $ "Loaded "++show (length btSaved)++" nodes for "++netname++"."
         fallbackNodes <- getBootstrapNodes
         let isNodesSearch :: ni :~: r -> Search nid addr tok ni r -> Search nid addr tok ni ni
             isNodesSearch Refl sch = sch
+            ping = maybe (const $ return False)
+                         (\DHTPing{pingQuery} -> fmap (maybe False (const True)) . pingQuery)
+                         $ Map.lookup "ping" pings
         fork $ do
             myThreadId >>= flip labelThread ("bootstrap."++netname)
             case Map.lookup "node" qrys of