1 files changed, 47 insertions, 18 deletions

diff --git a/examples/dhtd.hs b/examples/dhtd.hs
index f98b05bd..67888eb4 100644
--- a/examples/dhtd.hs
+++ b/examples/dhtd.hs
@@ -39,6 +39,7 @@ import GHC.Conc                  (labelThread)
 #endif
 
 import Network.Address hiding (NodeId, NodeInfo(..))
+import Network.BitTorrent.DHT.Search
 import Network.QueryResponse
 import Network.StreamServer
 import Kademlia
@@ -68,23 +69,31 @@ hPutClient h s = hPutStr h ('.' : marshalForClient s)
 hPutClientChunk :: Handle -> String -> IO ()
 hPutClientChunk h s = hPutStr h (' ' : marshalForClient s)
 
-data DHT = forall ni. ( Show     ni
-                      , Read     ni
-                      , ToJSON   ni
-                      , FromJSON ni
-                      ) =>
+data DHTQuery nid ni = forall addr r tok. DHTQuery
+    { qsearch  :: Search nid addr tok ni r
+    , qhandler :: ni -> nid -> IO ([ni], [r], tok)
+    , qshowR   :: r -> String
+    , qshowTok :: tok -> Maybe String
+    }
+
+data DHT = forall nid ni. ( Show     ni
+                          , Read     ni
+                          , ToJSON   ni
+                          , FromJSON ni
+                          ) =>
     DHT
     { dhtBuckets :: TVar (BucketList ni)
     , dhtPing    :: ni -> IO Bool
-    -- dhtClient :: Client err meth tid addr x
+    , dhtQuery   :: Map.Map String (DHTQuery nid ni)
+    , dhtParseId :: String -> Either String nid
     }
 
 nodesFileName :: String -> String
 nodesFileName netname = netname ++ "-nodes.json"
 
 saveNodes :: String -> DHT -> IO ()
-saveNodes netname (DHT var _) = do
-    bkts <- atomically $ readTVar var
+saveNodes netname DHT{dhtBuckets} = do
+    bkts <- atomically $ readTVar dhtBuckets
     let ns = map fst $ concat $ R.toList bkts
         bs = J.encode ns
         fname = nodesFileName netname
@@ -99,7 +108,7 @@ loadNodes netname = do
     either (const $ return []) return attempt
 
 pingNodes :: String -> DHT -> IO Bool
-pingNodes netname (DHT _ ping) = do
+pingNodes netname DHT{dhtPing} = do
     let fname = nodesFileName netname
     attempt <- tryIOError $ do
         J.decode <$> L.readFile fname
@@ -110,7 +119,7 @@ pingNodes netname (DHT _ ping) = do
                          myThreadId >>= flip labelThread ("pinging."++fname)
                          putStrLn $ "Forked "++show fname
                          withTaskGroup 10 $ \g -> do
-                           mapTasks_ g (map ping ns)
+                           mapTasks_ g (map dhtPing ns)
                          putStrLn $ "Load finished "++show fname
                        return True
 
@@ -189,9 +198,9 @@ clientSession s@Session{..} sock cnum h = do
                         _ -> hPutClient h "error."
 
 #endif
-        ("ls", _) | Just (DHT var _) <- Map.lookup netname dhts
+        ("ls", _) | Just DHT{dhtBuckets} <- Map.lookup netname dhts
                     -> cmd0 $ do
-                        bkts <- atomically $ readTVar var
+                        bkts <- atomically $ readTVar dhtBuckets
                         let r = reportTable bkts
                         hPutClient h $
                             showReport $
@@ -199,10 +208,10 @@ clientSession s@Session{..} sock cnum h = do
                                          , ("node-id", show $ thisNode bkts)
                                          , ("network", netname) ]
 
-        ("ping", s) | Just (DHT _ ping) <- Map.lookup netname dhts
+        ("ping", s) | Just DHT{dhtPing} <- Map.lookup netname dhts
                     -> cmd0 $ do
                         case readEither s of
-                          Right addr -> do result <- ping addr
+                          Right addr -> do result <- dhtPing addr
                                            let rs = [" ", show result]
                                            hPutClient h $ unlines rs
                           Left er -> hPutClient h er
@@ -238,15 +247,35 @@ main = do
                 _             -> return defaultPort
     addr <- getBindAddress p True{- ipv6 -}
 
-    (bt,btR) <- Mainline.newClient addr
+    (bt,btR,swarms) <- Mainline.newClient addr
     quitBt <- forkListener bt
 
     tox <- return $ error "TODO: Tox.newClient"
     quitTox <- return $ return () -- TODO: forkListener tox
 
-    let dhts = Map.fromList
-                [ ("bt4", DHT (Mainline.routing4 btR) (Mainline.ping bt))
-                , ("bt6", DHT (Mainline.routing6 btR) (Mainline.ping bt))
+    let mainlineDHT bkts = DHT
+            { dhtBuckets = bkts btR
+            , dhtPing    = Mainline.ping bt
+            , dhtQuery   = Map.fromList
+                            [ ("node", DHTQuery (Mainline.nodeSearch bt)
+                                                (\ni -> fmap Mainline.unwrapNodes
+                                                        . Mainline.findNodeH btR ni
+                                                        . flip Mainline.FindNode (Just Mainline.Want_Both))
+                                                show
+                                                (const Nothing))
+                            , ("peer", DHTQuery (Mainline.peerSearch bt)
+                                                (\ni -> fmap Mainline.unwrapPeers
+                                                        . Mainline.getPeersH btR swarms ni
+                                                        . flip Mainline.GetPeers (Just Mainline.Want_Both)
+                                                        . (read . show)) -- TODO: InfoHash -> NodeId
+                                                show
+                                                (Just . show))
+                            ]
+            , dhtParseId = readEither :: String -> Either String Mainline.NodeId
+            }
+        dhts = Map.fromList
+                [ ("bt4", mainlineDHT Mainline.routing4)
+                , ("bt6", mainlineDHT Mainline.routing6)
                 ]
 
     waitForSignal <- do