path: root/examples/dhtd.hs

{-# LANGUAGE NondecreasingIndentation #-}
{-# LANGUAGE FlexibleContexts      #-}
{-# LANGUAGE LambdaCase            #-}
{-# LANGUAGE OverloadedStrings     #-}
{-# LANGUAGE ScopedTypeVariables   #-}
{-# LANGUAGE TupleSections         #-}
{-# LANGUAGE RecordWildCards       #-}
{-# LANGUAGE CPP #-}

import Control.Arrow
import Control.Monad
import Control.Monad.Logger
import Control.Monad.Reader
import Data.Char
import Data.Default
import Data.List as L
import Data.Maybe
import Data.String
import qualified Data.ByteString as B (ByteString,writeFile,readFile)
        ; import Data.ByteString (ByteString)
import qualified Data.ByteString.Char8 as B8
import System.IO
import System.IO.Error
import Text.PrettyPrint.HughesPJClass
import Text.Printf
import Text.Read
import Control.Monad.Reader.Class
import System.Posix.Process (getProcessID)

import Data.Torrent (InfoHash)
import Network.BitTorrent.Address
import Network.BitTorrent.DHT
import qualified Network.BitTorrent.DHT.Routing as R
import Network.BitTorrent.DHT.Session
import Network.SocketLike
import Network.StreamServer
import Control.Exception.Lifted as Lifted
#ifdef THREAD_DEBUG
import Control.Concurrent.Lifted.Instrument
import Data.Time ()
import Data.Time.Clock
#else
import Control.Concurrent
#endif
import System.Environment

mkNodeAddr :: SockAddr -> NodeAddr IPv4
mkNodeAddr addr = NodeAddr (fromJust $ fromSockAddr addr)
                           (fromMaybe 0 $ sockAddrPort addr) -- FIXME

btBindAddr :: String -> Bool -> IO (NodeAddr IPv4)
btBindAddr s b = mkNodeAddr <$> getBindAddress s b

printReport :: MonadIO m => [(String,String)] -> m ()
printReport kvs = liftIO $ do
    putStrLn (showReport kvs)
    hFlush stdout

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" ]

showEnry :: Show a => (NodeInfo a, t) -> [Char]
showEnry (n,_) = intercalate " "
    [ show $ pPrint (nodeId n)
    , show $ nodeAddr n
    ]

printTable :: DHT IPv4 ()
printTable = do
    t <- showTable
    liftIO $ do
        putStrLn t
        hFlush stdout

showTable :: DHT IPv4 String
showTable = do
    nodes <- R.toList <$> getTable
    return $ showReport
           $ map (show *** showEnry)
           $ concat $ zipWith map (map (,) [0::Int ..]) nodes

bootstrapNodes :: IO [NodeAddr IPv4]
bootstrapNodes = mapMaybe fromAddr
                <$> mapM resolveHostName defaultBootstrapNodes

-- ExtendedCaps (Map.singleton

noDebugPrints :: LogSource -> LogLevel -> Bool
noDebugPrints _ = \case LevelDebug -> False
                        _          -> True

noLogging :: LogSource -> LogLevel -> Bool
noLogging _ _ = False

allNoise :: LogSource -> LogLevel ->  Bool
allNoise _ _ = True

resume :: DHT IPv4 (Maybe B.ByteString)
resume = do
    restore_attempt <- liftIO $ tryIOError $ B.readFile "dht-nodes.dat"
    saved_nodes <-
        either (const $ do liftIO $ putStrLn "Error reading dht-nodes.dat"
                           return Nothing)
               (return . Just)
            restore_attempt
    return saved_nodes

godht :: String -> (NodeAddr IPv4 -> NodeId -> DHT IPv4 b) -> IO b
godht p f = do
    a <- btBindAddr p False
    dht def { optTimeout = 5 } a noDebugPrints $ do
        me0 <- asks tentativeNodeId
        printReport [("tentative node-id",show $ pPrint me0)
                   ,("listen-address", show a)
                   ]
        f a me0

marshalForClient :: String -> String
marshalForClient s = show (length s) ++ ":" ++ s

hPutClient :: Handle -> String -> IO ()
hPutClient h s = hPutStr h (marshalForClient s)

clientSession :: Node IPv4 -> MVar () -> RestrictedSocket -> Int -> Handle -> IO ()
clientSession st signalQuit sock n h = do
    line <- map toLower . dropWhile isSpace <$> hGetLine h
    let cmd0 action = action >> clientSession st signalQuit sock n h
        cmd  action = cmd0 $ join $ runDHT st action
        (c,args) = second (dropWhile isSpace) $ break isSpace line
    case (c,args) of

        ("quit", _) -> hPutClient h "" >> hClose h

        ("stop", _) -> do hPutClient h "Terminating DHT Daemon."
                          hClose h
                          putMVar signalQuit ()

        ("ls", _) -> cmd $ do
                    tbl <- getTable
                    t <- showTable
                    me <- myNodeIdAccordingTo (read "8.8.8.8:6881")
                    ip <- routableAddress
                    return $ do
                    hPutClient h $ unlines
                        [ t
                        , showReport
                            [ ("node-id", show $ pPrint me)
                            , ("internet address", show ip)
                            , ("buckets", show $ R.shape tbl)]
                        ]
        ("external-ip", _) -> cmd $ do
                    ip <- routableAddress
                    return $ do
                    hPutClient h $ maybe "" (takeWhile (/=':') . show) ip

        ("swarms", s) -> cmd $ do
                    let fltr = case s of
                                    ('-':'v':cs) | all isSpace (take 1 cs)
                                      -> const True
                                    _ -> (\(h,c,n) -> c/=0 )
                    ss <- getSwarms
                    let r = map (\(h,c,n) -> (unwords [show h,show c], maybe "" show n))
                            $ filter fltr ss
                    return $ do
                    hPutClient h $ showReport r

        ("peers", s) -> cmd $ case readEither s of
                        Right ih -> do
                            ps <- allPeers ih
                            seq ih $ return $ do
                            hPutClient h $ showReport $ map (((,) "") . show . pPrint) ps
                        Left er -> return $ hPutClient h er

        ("pid", _) -> cmd $ return $ do
                    pid <- getProcessID
                    hPutClient h (show pid)
#ifdef THREAD_DEBUG
        ("threads", _) -> cmd $ return $ do
                    ts <- threadsInformation
                    tm <- getCurrentTime
                    let r = map (\PerThread{..} -> (show lbl,show (diffUTCTime tm startTime))) ts
                    hPutClient h $ showReport r
#endif
        ("closest", s) -> cmd $ do
                        let (ns,hs) = second (dropWhile isSpace) $ break isSpace s
                            parse | null hs = do
                                       ih <- readEither ns
                                       return (8 :: Int, ih :: InfoHash)
                                  | otherwise = do
                                       n <- readEither ns
                                       ih <- readEither hs
                                       return (n :: Int, ih :: InfoHash)
                        case parse of
                            Right (n,ih) -> do
                                tbl <- getTable
                                let nodes = R.kclosest n ih tbl
                                return $ do
                                hPutClient h $ unlines $ map (showEnry . (flip (,) 0)) nodes
                            Left er -> return $ hPutClient h er

        _    -> cmd0 $ hPutClient h "error."

defaultPort = error "TODO defaultPort"

main :: IO ()
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)"
                _             -> defaultPort
    godht p $ \a me0 -> do
    printTable
    bs <- liftIO bootstrapNodes
            `onException`
                 (Lifted.ioError $ userError "unable to resolve bootstrap nodes")
    saved_nodes <- resume

    peers'trial <- liftIO $ tryIOError $ B.readFile "bt-peers.dat"
    saved_peers <-
        either (const $ do liftIO $ putStrLn "Error reading bt-peers.dat"
                           return Nothing)
               (return . Just)
            peers'trial

    maybe (return ()) mergeSavedPeers saved_peers

    when (isJust saved_nodes) $ do
        b <- isBootstrapped
        tbl <- getTable
        bc <- optBucketCount <$> asks options
        printTable
        me <- case concat $ R.toList tbl of
                (n,_):_ -> myNodeIdAccordingTo (nodeAddr n)
                _       -> return me0
        printReport [("node-id",show $ pPrint me)
                   ,("listen-address", show a)
                   ,("bootstrapped", show b)
                   ,("buckets", show $ R.shape tbl)
                   ,("optBucketCount", show bc)
                   ,("dht-nodes.dat", "Running bootstrap...")
                   ]

    st <- ask
    waitForSignal <- liftIO $ do
        signalQuit <- newEmptyMVar
        srv <- streamServer (withSession $ clientSession st signalQuit) (SockAddrUnix "dht.sock")
        return $ liftIO $ do
            () <- takeMVar signalQuit
            quitListening srv

    bootstrap saved_nodes bs

    b <- isBootstrapped
    tbl <- getTable
    bc <- optBucketCount <$> asks options
    printTable
    ip <- routableAddress
    me <- case concat $ R.toList tbl of
            (n,_):_ -> myNodeIdAccordingTo (nodeAddr n)
            _       -> return me0
    printReport [("node-id",show $ pPrint me)
               ,("internet address", show ip)
               ,("listen-address", show a)
               ,("bootstrapped", show b)
               ,("buckets", show $ R.shape tbl)
               ,("optBucketCount", show bc)
               ]

    waitForSignal -- Await unix socket to signal termination.

    snapshot >>= liftIO . B.writeFile "dht-nodes.dat"
    savePeerStore >>= liftIO . B.writeFile "bt-peers.dat"