1 files changed, 33 insertions, 1 deletions

diff --git a/examples/dhtd.hs b/examples/dhtd.hs
index 2abaecdd..6477fac4 100644
--- a/examples/dhtd.hs
+++ b/examples/dhtd.hs
@@ -29,6 +29,7 @@ import System.Environment
 import System.IO
 import System.Mem
 import System.Posix.Process
+import Text.PrettyPrint.HughesPJClass
 import Text.Printf
 import Text.Read
 #ifdef THREAD_DEBUG
@@ -50,6 +51,8 @@ import qualified Data.Aeson as J
 import qualified Data.ByteString.Lazy as L
 import Control.Concurrent.Async.Pool
 import System.IO.Error
+import qualified Data.Serialize as S
+import Network.BitTorrent.DHT.ContactInfo as Peers
 
 showReport :: [(String,String)] -> String
 showReport kvs = do
@@ -156,6 +159,7 @@ data Session = Session
     { netname           :: String
     , dhts              :: Map.Map String DHT
     , externalAddresses :: IO [SockAddr]
+    , swarms            :: Mainline.SwarmsDatabase
     , signalQuit        :: MVar ()
     }
 
@@ -177,7 +181,7 @@ clientSession s@Session{..} sock cnum h = do
                     pid <- getProcessID
                     hPutClient h (show pid)
         ("external-ip", _) -> cmd0 $ do
-                    unlines . map show <$> externalAddresses
+                    unlines . map (either show show . Mainline.either4or6) <$> externalAddresses
                         >>= hPutClient h
 #ifdef THREAD_DEBUG
         ("threads", _) -> cmd0 $ do
@@ -269,6 +273,23 @@ clientSession s@Session{..} sock cnum h = do
                         b <- pingNodes netname dht
                         if b then hPutClient h $ "Pinging " ++ nodesFileName netname ++ "."
                              else hPutClient h $ "Failed: " ++ nodesFileName netname ++ "."
+
+        ("swarms", s) -> cmd0 $ do
+                    let fltr = case s of
+                                    ('-':'v':cs) | all isSpace (take 1 cs)
+                                      -> const True
+                                    _ -> (\(h,c,n) -> c/=0 )
+                    ss <- atomically $ Peers.knownSwarms <$> readTVar (Mainline.contactInfo swarms)
+                    let r = map (\(h,c,n) -> (unwords [show h,show c], maybe "" show n))
+                            $ filter fltr ss
+                    hPutClient h $ showReport r
+
+        ("peers", s) -> cmd0 $ case readEither s of
+                        Right ih -> do
+                            ps <- atomically $ Peers.lookup ih <$> readTVar (Mainline.contactInfo swarms)
+                            hPutClient h $ showReport $ map (((,) "") . show . pPrint) ps
+                        Left er -> hPutClient h er
+
         (n, _) | n `elem` Map.keys dhts -> switchNetwork n
 
         _    -> cmd0 $ hPutClient h "error."
@@ -291,6 +312,13 @@ main = do
     addr <- getBindAddress p True{- ipv6 -}
 
     (bt,btR,swarms) <- Mainline.newClient addr
+
+    -- Restore peer database before forking the listener thread.
+    peerdb <- left show <$> tryIOError (L.readFile "bt-peers.dat")
+    either (hPutStrLn stderr . ("bt-peers.dat: "++))
+           (atomically . writeTVar (Mainline.contactInfo swarms))
+           (peerdb >>= S.decodeLazy)
+
     quitBt <- forkListener bt
 
     tox <- return $ error "TODO: Tox.newClient"
@@ -327,6 +355,7 @@ main = do
                 { netname           = "bt4" -- initial default DHT
                 , dhts              = dhts  -- all DHTs
                 , signalQuit        = signalQuit
+                , swarms            = swarms
                 , externalAddresses = readExternals
                                         [ Mainline.routing4 btR
                                         , Mainline.routing6 btR
@@ -362,3 +391,6 @@ main = do
 
     quitBt
     quitTox
+
+    swarmsdb <- atomically $ readTVar (Mainline.contactInfo swarms)
+    L.writeFile "bt-peers.dat" $ S.encodeLazy swarmsdb