2 files changed, 62 insertions, 26 deletions

diff --git a/Mainline.hs b/Mainline.hs
index 7c54a096..921bea0f 100644
--- a/Mainline.hs
+++ b/Mainline.hs
@@ -92,6 +92,13 @@ instance S.Serialize NodeId where
 instance FiniteBits NodeId where
     finiteBitSize _ = 160
 
+instance Read NodeId where
+    readsPrec _ str
+        | (bs, xs) <- Base16.decode $ Char8.pack str
+        , B.length bs == 20
+                    = [ (NodeId bs, drop 40 str) ]
+        | otherwise = []
+
 zeroID :: NodeId
 zeroID = NodeId $ B.replicate 20 0
 
@@ -465,7 +472,7 @@ traced (TableMethods ins del lkup)
 
 type MainlineClient = Client String Method TransactionId NodeInfo (Message BValue)
 
-newClient :: SockAddr -> IO (MainlineClient, Routing)
+newClient :: SockAddr -> IO (MainlineClient, Routing, SwarmsDatabase)
 newClient addr = do
     udp <- udpTransport addr
     nid <- NodeId <$> getRandomBytes 20
@@ -564,7 +571,7 @@ newClient addr = do
                             (sched6 routing)
                             (refreshBucket (nodeSearch client) (routing6 routing) (nodeId tenative_info))
 
-    return (client, routing)
+    return (client, routing, swarms)
 
 -- | Modifies a purely random 'NodeId' to one that is related to a given
 -- routable address in accordance with BEP 42.
@@ -969,14 +976,14 @@ ping client addr =
 
 -- searchQuery        :: ni -> IO (Maybe [ni], [r], tok))
 getNodes :: MainlineClient -> NodeId -> NodeInfo -> IO (Maybe ([NodeInfo],[NodeInfo],()))
-getNodes = mainlineSend (Method "find_node") unwrap $ flip FindNode (Just Want_Both)
- where
-    unwrap (NodeFound ns4 ns6) = (ns4++ns6, ns4++ns6,())
+getNodes = mainlineSend (Method "find_node") unwrapNodes $ flip FindNode (Just Want_Both)
+
+unwrapNodes (NodeFound ns4 ns6) = (ns4++ns6, ns4++ns6,())
 
 getPeers :: MainlineClient -> NodeId -> NodeInfo -> IO (Maybe ([NodeInfo],[PeerAddr],Token))
-getPeers = mainlineSend (Method "get_peers") unwrap $ flip GetPeers (Just Want_Both) . coerce
- where
-    unwrap (GotPeers ps (NodeFound ns4 ns6) tok) = (ns4++ns6, ps, tok)
+getPeers = mainlineSend (Method "get_peers") unwrapPeers $ flip GetPeers (Just Want_Both) . coerce
+
+unwrapPeers (GotPeers ps (NodeFound ns4 ns6) tok) = (ns4++ns6, ps, tok)
 
 data TriadSlot = SlotA | SlotB | SlotC
  deriving (Eq,Ord,Enum,Show,Read)
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