{-# LANGUAGE CPP #-} {-# LANGUAGE ExistentialQuantification #-} {-# LANGUAGE FlexibleContexts #-} {-# LANGUAGE FlexibleInstances #-} {-# LANGUAGE LambdaCase #-} {-# LANGUAGE MultiParamTypeClasses #-} {-# LANGUAGE NamedFieldPuns #-} {-# LANGUAGE NondecreasingIndentation #-} {-# LANGUAGE OverloadedStrings #-} {-# LANGUAGE PartialTypeSignatures #-} {-# LANGUAGE RankNTypes #-} {-# LANGUAGE RecordWildCards #-} {-# LANGUAGE ScopedTypeVariables #-} {-# LANGUAGE TupleSections #-} {-# LANGUAGE TypeFamilies #-} import Control.Arrow import Control.Concurrent.STM import Control.DeepSeq import Control.Exception import Control.Monad import Data.Char import Data.Hashable import Data.List import qualified Data.Map as Map import Data.Maybe import qualified Data.Set as Set import Data.Time.Clock import GHC.Conc (threadStatus,ThreadStatus(..)) import GHC.Stats import Network.Socket 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 import Control.Concurrent.Lifted.Instrument #else import Control.Concurrent.Lifted import GHC.Conc (labelThread) #endif import Network.Address hiding (NodeId, NodeInfo(..)) import Network.BitTorrent.DHT.Search import Network.QueryResponse import Network.StreamServer import Kademlia import qualified Mainline import Network.DHT.Routing as R import Data.Aeson as J (ToJSON, FromJSON) 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 import qualified Data.MinMaxPSQ as MM showReport :: [(String,String)] -> String showReport kvs = do let colwidth = maximum $ map (length . fst) kvs (k,v) <- kvs concat [ printf " %-*s" (colwidth+1) k, v, "\n" ] marshalForClient :: String -> String marshalForClient s = show (length s) ++ ":" ++ s -- | Writes a message and signals ready for next command. hPutClient :: Handle -> String -> IO () hPutClient h s = hPutStr h ('.' : marshalForClient s) -- | Writes message, but signals there is more to come. hPutClientChunk :: Handle -> String -> IO () hPutClientChunk h s = hPutStr h (' ' : marshalForClient s) data DHTQuery nid ni = forall addr r tok. Ord addr => DHTQuery { qsearch :: Search nid addr tok ni r , qhandler :: ni -> nid -> IO ([ni], [r], tok) , qshowR :: r -> String , qshowTok :: tok -> Maybe String } data DHTSearch nid ni = forall addr tok r. DHTSearch { searchThread :: ThreadId , searchState :: SearchState nid addr tok ni r , searchShowTok :: tok -> Maybe String , searchResults :: TVar (Set.Set String) } data DHT = forall nid ni. ( Show ni , Read ni , ToJSON ni , FromJSON ni , Ord ni , Hashable ni , Show nid , Ord nid , Hashable nid ) => DHT { dhtBuckets :: TVar (BucketList ni) , dhtPing :: ni -> IO Bool , dhtQuery :: Map.Map String (DHTQuery nid ni) , dhtParseId :: String -> Either String nid , dhtSearches :: TVar (Map.Map (String,nid) (DHTSearch nid ni)) } nodesFileName :: String -> String nodesFileName netname = netname ++ "-nodes.json" saveNodes :: String -> DHT -> IO () saveNodes netname DHT{dhtBuckets} = do bkts <- atomically $ readTVar dhtBuckets let ns = map fst $ concat $ R.toList bkts bs = J.encode ns fname = nodesFileName netname L.writeFile fname bs loadNodes :: FromJSON ni => String -> IO [ni] loadNodes netname = do let fname = nodesFileName netname attempt <- tryIOError $ do J.decode <$> L.readFile fname >>= maybe (ioError $ userError "Nothing") return either (const $ return []) return attempt pingNodes :: String -> DHT -> IO Bool pingNodes netname DHT{dhtPing} = do let fname = nodesFileName netname attempt <- tryIOError $ do J.decode <$> L.readFile fname >>= maybe (ioError $ userError "Nothing") return case attempt of Left _ -> return False Right ns -> do fork $ do myThreadId >>= flip labelThread ("pinging."++fname) putStrLn $ "Forked "++show fname withTaskGroup 10 $ \g -> do mapTasks_ g (map dhtPing ns) putStrLn $ "Load finished "++show fname return True reportTable :: Show ni => BucketList ni -> [(String,String)] reportTable bkts = map (show *** show . fst) $ concat $ zipWith map (map (,) [0::Int ..]) $ R.toList $ bkts reportResult :: String -> (r -> String) -> (tok -> Maybe String) -> (ni -> String) -> Handle -> Either String ([ni],[r],tok) -> IO () reportResult meth showR showTok showN h (Left e) = hPutClient h e reportResult meth showR showTok showN h (Right (ns,rs,tok)) = do hPutClient h $ showReport report where report = intercalate [("","")] [ tok_r , node_r , result_r ] tok_r = maybe [] (pure . ("token:",)) $ showTok tok node_r = map ( ("n",) . showN ) ns result_r | (meth=="node") = [] | otherwise = map ( (take 1 meth,) . showR ) rs -- example: -- * 10 peer 141d6c6ee2810f46d28bbe8373d4f454a4122535 -- - 1 peer 141d6c6ee2810f46d28bbe8373d4f454a4122535 -- 22 node 141d6c6ee2810f46d28bbe8373d4f454a4122535 -- -- key: '*' in progress -- '-' stopped -- ' ' finished showSearches :: ( Show nid , Ord nid , Hashable nid , Ord ni , Hashable ni ) => Map.Map (String,nid) (DHTSearch nid ni) -> IO String showSearches searches = do tups <- forM (Map.toList searches) $ \((meth,nid),DHTSearch{..}) -> do (is'fin, cnt) <- atomically $ (,) <$> searchIsFinished searchState <*> (Set.size <$> readTVar searchResults) tstat <- threadStatus searchThread let stat = case tstat of _ | is'fin -> ' ' ThreadFinished -> '-' ThreadDied -> '-' _ -> '*' return (stat,show cnt,meth,show nid) let cnt'width = maximum $ map (\(_,cnt,_,_)->length cnt) tups mth'width = maximum $ map (\(_,_,mth,_)->length mth) tups return $ do -- List monad. (stat,cnt,meth,nid) <- tups printf " %c %-*s %-*s %s\n" stat cnt'width cnt mth'width meth nid forkSearch :: ( Ord nid , Hashable nid , Ord ni , Hashable ni , Show nid ) => String -> nid -> DHTQuery nid ni -> TVar (Map.Map (String,nid) (DHTSearch nid ni)) -> TVar (BucketList ni) -> ThreadId -> TVar (Maybe (IO ())) -> STM () forkSearch method nid DHTQuery{qsearch,qshowTok,qshowR} dhtSearches dhtBuckets tid kvar = do ns <- R.kclosest (searchSpace qsearch) searchK nid <$> readTVar dhtBuckets st <- newSearch qsearch nid ns results <- newTVar Set.empty let storeResult r = modifyTVar' results (Set.insert (qshowR r)) >> return True new = DHTSearch { searchThread = tid , searchState = st , searchShowTok = qshowTok , searchResults = results } modifyTVar' dhtSearches $ Map.insert (method,nid) new writeTVar kvar $ Just $ searchLoop qsearch nid storeResult st reportSearchResults meth h DHTSearch{searchShowTok,searchState,searchResults} = do (ns,rs) <- atomically $ do mm <- readTVar $ searchInformant searchState rset <- readTVar searchResults let ns = map (\(MM.Binding ni tok _) -> (ni,tok)) $ MM.toList mm rs = Set.toList rset return (ns,rs) let n'width = succ $ maximum $ map (length . show . fst) ns showN (n,tok) = take n'width (show n ++ repeat ' ') ++ (fromMaybe "" $ searchShowTok tok) ns' = map showN ns reportResult meth id (const Nothing) id h (Right (ns',rs,())) data Session = Session { netname :: String , dhts :: Map.Map String DHT , externalAddresses :: IO [SockAddr] , swarms :: Mainline.SwarmsDatabase , signalQuit :: MVar () } clientSession s@Session{..} sock cnum h = do line <- map toLower . dropWhile isSpace <$> hGetLine h let (c,args) = second (dropWhile isSpace) $ break isSpace line cmd0 :: IO () -> IO () cmd0 action = action >> clientSession s sock cnum h switchNetwork dest = do hPutClient h ("Network: "++dest) clientSession s{netname=dest} sock cnum h case (c,args) of ("stop", _) -> do hPutClient h "Terminating DHT Daemon." hClose h putMVar signalQuit () ("quit", _) -> hPutClient h "" >> hClose h ("pid", _) -> cmd0 $ do pid <- getProcessID hPutClient h (show pid) ("external-ip", _) -> cmd0 $ do unlines . map (either show show . Mainline.either4or6) <$> externalAddresses >>= hPutClient h #ifdef THREAD_DEBUG ("threads", _) -> cmd0 $ do ts <- threadsInformation tm <- getCurrentTime let r = map (\PerThread{..} -> (show lbl,show (diffUTCTime tm startTime))) ts hPutClient h $ showReport r #endif ("mem", s) -> cmd0 $ do case s of "gc" -> do hPutClient h "Performing garbage collection..." performMajorGC "" -> do is_enabled <- getGCStatsEnabled if is_enabled then do GCStats{..} <- getGCStats let r = [ ("bytesAllocated", show bytesAllocated) , ("numGcs", show numGcs) , ("maxBytesUsed", show maxBytesUsed) , ("numByteUsageSamples", show numByteUsageSamples) , ("cumulativeBytesUsed", show cumulativeBytesUsed) , ("bytesCopied", show bytesCopied) , ("currentBytesUsed", show currentBytesUsed) , ("currentBytesSlop", show currentBytesSlop) , ("maxBytesSlop", show maxBytesSlop) , ("peakMegabytesAllocated", show peakMegabytesAllocated) , ("mutatorCpuSeconds", show mutatorCpuSeconds) , ("mutatorWallSeconds", show mutatorWallSeconds) , ("gcCpuSeconds", show gcCpuSeconds) , ("gcWallSeconds", show gcWallSeconds) , ("cpuSeconds", show cpuSeconds) , ("wallSeconds", show wallSeconds) , ("parTotBytesCopied", show parTotBytesCopied) , ("parMaxBytesCopied", show parMaxBytesCopied) ] hPutClient h $ showReport r else hPutClient h "Run with +RTS -T to obtain live memory-usage information." _ -> hPutClient h "error." ("ls", _) | Just DHT{dhtBuckets} <- Map.lookup netname dhts -> cmd0 $ do bkts <- atomically $ readTVar dhtBuckets let r = reportTable bkts hPutClient h $ showReport $ r ++ [ ("buckets", show $ R.shape bkts) , ("node-id", show $ thisNode bkts) , ("network", netname) ] ("ping", s) | Just DHT{dhtPing} <- Map.lookup netname dhts -> cmd0 $ do case readEither s of Right addr -> do result <- dhtPing addr let rs = [" ", show result] hPutClient h $ unlines rs Left er -> hPutClient h er ("g", s) | Just DHT{..} <- Map.lookup netname dhts -> cmd0 $ do -- arguments: method -- nid -- (optional dest-ni) self <- atomically $ thisNode <$> readTVar dhtBuckets let (method,xs) = break isSpace $ dropWhile isSpace s (nidstr,ys) = break isSpace $ dropWhile isSpace xs destination = dropWhile isSpace ys goQuery qry = either (hPutClient h . ("Bad search target: "++)) (goTarget qry) $ dhtParseId nidstr goTarget DHTQuery{..} nid = go nid >>= reportResult method qshowR qshowTok show h where go | null destination = fmap Right . qhandler self | otherwise = case readEither destination of Right ni -> fmap (maybe (Left "Timeout.") Right) . flip (searchQuery qsearch) ni Left e -> const $ return $ Left ("Bad destination: "++e) maybe (hPutClient h ("Unsupported method: "++method)) goQuery $ Map.lookup method dhtQuery ("s", s) | Just dht@DHT{..} <- Map.lookup netname dhts -> cmd0 $ do let (method,xs) = break isSpace s (nidstr,ys) = break isSpace $ dropWhile isSpace xs presentSearches = hPutClient h =<< showSearches =<< atomically (readTVar dhtSearches) goTarget qry nid = do kvar <- atomically $ newTVar Nothing -- Forking a thread, but it may ubruptly quit if the following -- STM action decides not to add a new search. This is so that -- I can store the ThreadId into new DHTSearch structure. tid <- fork $ join $ atomically (readTVar kvar >>= maybe retry return) join $ atomically $ do schs <- readTVar dhtSearches case Map.lookup (method,nid) schs of Nothing -> do forkSearch method nid qry dhtSearches dhtBuckets tid kvar return $ presentSearches Just sch -> do writeTVar kvar (Just $ return ()) return $ reportSearchResults method h sch goQuery qry = either (hPutClient h . ("Bad search target: "++)) (goTarget qry) $ dhtParseId nidstr if null method then presentSearches else maybe (hPutClient h ("Unsupported method: "++method)) goQuery $ Map.lookup method dhtQuery ("x", s) | Just DHT{..} <- Map.lookup netname dhts -> cmd0 $ do let (method,xs) = break isSpace s (nidstr,ys) = break isSpace $ dropWhile isSpace xs go nid = join $ atomically $ do schs <- readTVar dhtSearches case Map.lookup (method,nid) schs of Nothing -> return $ hPutClient h "No match." Just DHTSearch{searchThread} -> do modifyTVar' dhtSearches (Map.delete (method,nid)) return $ do killThread searchThread hPutClient h "Removed search." either (hPutClient h . ("Bad search target: "++)) go $ dhtParseId nidstr ("save", _) | Just dht <- Map.lookup netname dhts -> cmd0 $ do saveNodes netname dht hPutClient h $ "Saved " ++ nodesFileName netname ++ "." ("load", _) | Just dht <- Map.lookup netname dhts -> cmd0 $ 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." readExternals :: [TVar (BucketList Mainline.NodeInfo)] -> IO [SockAddr] readExternals vars = do as <- atomically $ mapM (fmap (Mainline.nodeAddr . selfNode) . readTVar) vars let unspecified (SockAddrInet _ 0) = True unspecified (SockAddrInet6 _ _ (0,0,0,0) _) = True unspecified _ = False -- TODO: Filter to only global addresses? return $ filter (not . unspecified) as defaultPort = "6881" main = do args <- getArgs p <- case take 2 (dropWhile (/="-p") args) of ["-p",port] | not ("-" `isPrefixOf` port) -> return port ("-p":_) -> error "Port not specified! (-p PORT)" _ -> return defaultPort 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" quitTox <- return $ return () -- TODO: forkListener tox mainlineSearches <- atomically $ newTVar Map.empty 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 , dhtSearches = mainlineSearches } dhts = Map.fromList [ ("bt4", mainlineDHT Mainline.routing4) , ("bt6", mainlineDHT Mainline.routing6) ] waitForSignal <- do signalQuit <- newEmptyMVar let session = clientSession $ Session { netname = "bt4" -- initial default DHT , dhts = dhts -- all DHTs , signalQuit = signalQuit , swarms = swarms , externalAddresses = readExternals [ Mainline.routing4 btR , Mainline.routing6 btR ] } srv <- streamServer (withSession session) (SockAddrUnix "dht.sock") return $ do () <- takeMVar signalQuit quitListening srv let bkts4 = Mainline.routing4 btR btSaved4 <- loadNodes "bt4" :: IO [Mainline.NodeInfo] putStrLn $ "Loaded "++show (length btSaved4)++" nodes for bt4." fallbackNodes4 <- Mainline.bootstrapNodes Mainline.Want_IP4 fork $ do myThreadId >>= flip labelThread "bootstrap.Mainline4" bootstrap (Mainline.nodeSearch bt) bkts4 (Mainline.ping bt) btSaved4 fallbackNodes4 saveNodes "bt4" (dhts Map.! "bt4") btSaved6 <- loadNodes "bt6" putStrLn $ "Loaded "++show (length btSaved6)++" nodes for bt6." let bkts6 = Mainline.routing6 btR fallbackNodes6 <- Mainline.bootstrapNodes Mainline.Want_IP6 fork $ do myThreadId >>= flip labelThread "bootstrap.Mainline6" bootstrap (Mainline.nodeSearch bt) bkts6 (Mainline.ping bt) btSaved6 fallbackNodes6 saveNodes "bt6" (dhts Map.! "bt6") hPutStr stderr $ showReport $ map (("bootstrap (IPv4)",) . show) fallbackNodes4 ++ map (("bootstrap (IPv6)",) . show) fallbackNodes6 waitForSignal quitBt quitTox swarmsdb <- atomically $ readTVar (Mainline.contactInfo swarms) L.writeFile "bt-peers.dat" $ S.encodeLazy swarmsdb