5 files changed, 539 insertions, 0 deletions

diff --git a/bittorrent.cabal b/bittorrent.cabal
index 0fddf5d0..fbd17397 100644
--- a/bittorrent.cabal
+++ b/bittorrent.cabal
@@ -83,6 +83,8 @@ library
                        Network.BitTorrent.DHT.Routing
                        Network.BitTorrent.DHT.Session
                        Network.BitTorrent.DHT.Token
+                       Network.StreamServer
+                       Network.SocketLike
   other-modules:       Paths_bittorrent
   if !flag(dht-only)
      exposed-modules:  Network.BitTorrent
@@ -336,6 +338,24 @@ benchmark bench
                      , criterion
   ghc-options:         -O2 -fforce-recomp
 
+executable dht
+  hs-source-dirs:      examples
+  main-is:             dht.hs
+  default-language:    Haskell2010
+  build-depends:       base, haskeline, network, bytestring
+
+executable dhtd
+  hs-source-dirs:      examples
+  main-is:             dhtd.hs
+  default-language:    Haskell2010
+  build-depends:       base, network, bytestring
+                     , mtl
+                     , lifted-base
+                     , pretty
+                     , data-default
+                     , monad-logger
+                     , bittorrent
+
 -- Utility to work with torrent files.
 executable mktorrent
   if !flag(examples)
diff --git a/examples/dht.hs b/examples/dht.hs
new file mode 100644
index 00000000..feeee9ff
--- /dev/null
+++ b/examples/dht.hs
@@ -0,0 +1,80 @@
+{-# LANGUAGE NondecreasingIndentation #-}
+import Control.Monad
+import Control.Monad.Fix
+import Control.Monad.IO.Class
+import Data.Char
+import Network.Socket as Socket
+import System.Console.Haskeline
+import System.Environment
+import System.Exit
+import System.IO
+import System.IO.Unsafe
+import qualified Data.ByteString as B
+
+-- | Reads one character.  If it is not a digit,
+-- then it is discarded and 'Nothing' is returned.
+hReadDigit :: Handle -> IO (Maybe Char)
+hReadDigit h = do c <- hGetChar h
+                  return $ guard (isDigit c) >> pure c
+
+-- | Expected input: "nnn:..."
+-- Here we read the digit sequence "nnn" and drop the colon
+-- as it is the first non-digit.
+hReadInt :: Handle -> IO Int
+hReadInt h = do
+        nstr <- fix $ \readDigits ->
+                        maybe (return []) -- dropped non-digit character
+                              (($ unsafeInterleaveIO readDigits) . fmap . (:))
+                             =<< hReadDigit h
+        readIO nstr :: IO Int
+
+
+-- | Read a length prefixed string from a handle.
+-- The format is "nnn:..." where /nnn/ is an ascii-encoded character count
+-- and /.../ is the sequence of characters
+--
+-- Note: The first byte after the count is ignored and discarded.
+readResponse :: Handle -> IO String
+readResponse h = do
+    n <- hReadInt h
+    sequence $ replicate n (hGetChar h)
+
+-- | Send a command to the dhtd daemon and then print the response.
+sendCommand :: Handle -> String -> InputT IO ()
+sendCommand h cmd = do resp <- liftIO $ do hPutStrLn h cmd
+                                           readResponse h
+                       outputStrLn resp
+
+-- | Get one line of input and send it to the daemon, then run the
+-- passed continuation if it wasn't "quit".
+interactiveMode :: Handle -> InputT IO () -> InputT IO ()
+interactiveMode h repl = do
+    minput <- getInputLine "dht> "
+    case minput of
+       Nothing     -> return ()
+       Just "quit" -> sendCommand h "quit" >> return ()
+       Just cmd    -> sendCommand h cmd >> repl
+
+main :: IO ()
+main = do
+    -- Open the control socket to the daemon.
+    h <- liftIO $ handle (\e -> do hPutStrLn stderr (show (e ::IOError))
+                                   exitFailure)
+                  $ do sock <- socket AF_UNIX Stream defaultProtocol
+                       connect sock (SockAddrUnix "dht.sock")
+                       socketToHandle sock ReadWriteMode
+
+    -- Haskeline's default looks only at our stdin and not our stdout.
+    -- That's a bad idea because we can take input from the command line.
+    behavior <- do
+        useTerminal <- and <$> mapM hIsTerminalDevice [stdin,stdout]
+        return $ if useTerminal then preferTerm else useFileHandle stdin
+
+    runInputTBehaviorWithPrefs behavior defaultPrefs defaultSettings $ do
+
+    -- A command may be specified on the command line
+    -- or else we enter an interactive shell.
+    args <- dropWhile isSpace . unwords <$> liftIO getArgs
+    case args of
+        (_:_) -> sendCommand h args >> sendCommand h "quit"
+        _     -> fix $ interactiveMode h
diff --git a/examples/dhtd.hs b/examples/dhtd.hs
new file mode 100644
index 00000000..6bf48595
--- /dev/null
+++ b/examples/dhtd.hs
@@ -0,0 +1,192 @@
+{-# LANGUAGE NondecreasingIndentation #-}
+{-# LANGUAGE FlexibleContexts      #-}
+{-# LANGUAGE LambdaCase            #-}
+{-# LANGUAGE OverloadedStrings     #-}
+{-# LANGUAGE ScopedTypeVariables   #-}
+{-# LANGUAGE TupleSections         #-}
+
+import Control.Arrow;
+import Control.Concurrent
+import Control.Exception.Lifted as Lifted
+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 qualified Data.ByteString as B (ByteString,writeFile,readFile)
+        ; import Data.ByteString (ByteString)
+import System.IO
+import System.IO.Error
+import Text.PrettyPrint.HughesPJClass
+import Text.Printf
+import Control.Monad.Reader.Class
+
+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
+
+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
+
+
+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 :: forall b. (NodeAddr IPv4 -> NodeId -> DHT IPv4 b) -> IO b
+godht f = do
+    a <- btBindAddr "8008" False
+    dht def { optTimeout = 5 } a (const $ const True) $ 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 cmd action = action >> clientSession st signalQuit sock n h
+    case line of
+
+        "quit" -> hPutClient h "goodbye." >> hClose h
+
+        "stop" -> do hPutClient h "Terminating DHT Daemon."
+                     hClose h
+                     putMVar signalQuit ()
+
+        "ls" -> cmd $ join $ runDHT st $ 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)]
+                        ]
+
+        _    -> cmd $ hPutClient h "error."
+
+main :: IO ()
+main = do
+    godht $ \a me0 -> do
+    printTable
+    bs <- liftIO bootstrapNodes
+            `onException`
+                 (Lifted.ioError $ userError "unable to resolve bootstrap nodes")
+    saved_nodes <- resume
+
+    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)
+               ]
+    snapshot >>= liftIO . B.writeFile "dht-nodes.dat"
+
+    waitForSignal
+
diff --git a/src/Network/SocketLike.hs b/src/Network/SocketLike.hs
new file mode 100644
index 00000000..2aa78e3e
--- /dev/null
+++ b/src/Network/SocketLike.hs
@@ -0,0 +1,104 @@
+{-# LANGUAGE GeneralizedNewtypeDeriving #-}
+-- |
+--
+-- A socket could be used indirectly via a 'System.IO.Handle' or a conduit from
+-- Michael Snoyman's conduit package.  But doing so presents an encapsulation
+-- problem.  Do we allow access to the underlying socket and trust that it wont
+-- be used in an unsafe way?  Or do we protect it at the higher level and deny
+-- access to various state information?
+--
+-- The 'SocketLike' class enables the approach that provides a safe wrapper to
+-- the underlying socket and gives access to various state information without
+-- enabling direct reads or writes.
+module Network.SocketLike
+    ( SocketLike(..)
+    , RestrictedSocket
+    , restrictSocket
+    , restrictHandleSocket
+    -- * Re-exports
+    --
+    -- | To make the 'SocketLike' methods less awkward to use, the types
+    -- 'CUInt', 'SockAddr', and 'PortNumber' are re-exported.
+    , CUInt
+    , PortNumber
+    , SockAddr(..)
+    ) where
+
+import Network.Socket
+    ( PortNumber
+    , SockAddr
+    )
+import Foreign.C.Types ( CUInt )
+
+import qualified Network.Socket as NS
+import System.IO (Handle,hClose,hIsOpen)
+
+-- | A safe (mostly read-only) interface to a 'NS.Socket'.  Note that despite
+-- how this class is named, it provides no access to typical 'NS.Socket' uses
+-- like sending or receiving network packets.
+class SocketLike sock where
+    -- | See 'NS.getSocketName'
+    getSocketName :: sock -> IO SockAddr
+    -- | See 'NS.getPeerName'
+    getPeerName   :: sock -> IO SockAddr
+    -- | See 'NS.getPeerCred'
+    getPeerCred   :: sock -> IO (CUInt, CUInt, CUInt)
+    -- | See 'NS.socketPort'
+    socketPort    :: sock -> IO PortNumber
+    -- | See 'NS.sIsConnected'
+    --
+    -- __Warning__: Don't rely on this method if it's possible the socket was
+    -- converted into a 'Handle'.
+    sIsConnected  :: sock -> IO Bool
+    -- | See 'NS.sIsBound'
+    sIsBound      :: sock -> IO Bool
+    -- | See 'NS.sIsListening'
+    sIsListening  :: sock -> IO Bool
+    -- | See 'NS.sIsReadable'
+    sIsReadable   :: sock -> IO Bool
+    -- | See 'NS.sIsWritable'
+    sIsWritable   :: sock -> IO Bool
+
+    -- | This is the only exposed write-access method to the
+    -- underlying state.  Usually implemented by 'NS.close'
+    sClose        :: sock -> IO ()
+
+instance SocketLike NS.Socket where
+    getSocketName = NS.getSocketName
+    getPeerName   = NS.getPeerName
+    getPeerCred   = NS.getPeerCred
+    socketPort    = NS.socketPort
+    sIsConnected  = NS.sIsConnected -- warning: this is always False if the socket
+                                    --          was converted to a Handle
+    sIsBound      = NS.sIsBound
+    sIsListening  = NS.sIsListening
+    sIsReadable   = NS.sIsReadable
+    sIsWritable   = NS.sIsWritable
+
+    sClose        = NS.sClose
+
+-- | An encapsulated socket.  Data reads and writes are not possible.
+data RestrictedSocket = Restricted (Maybe Handle) NS.Socket deriving Show
+
+instance SocketLike RestrictedSocket where
+    getSocketName (Restricted mb sock) = NS.getSocketName sock
+    getPeerName   (Restricted mb sock) = NS.getPeerName sock
+    getPeerCred   (Restricted mb sock) = NS.getPeerCred sock
+    socketPort    (Restricted mb sock) = NS.socketPort sock
+    sIsConnected  (Restricted mb sock) = maybe (NS.sIsConnected sock) (hIsOpen) mb
+    sIsBound      (Restricted mb sock) = NS.sIsBound sock
+    sIsListening  (Restricted mb sock) = NS.sIsListening sock
+    sIsReadable   (Restricted mb sock) = NS.sIsReadable sock
+    sIsWritable   (Restricted mb sock) = NS.sIsWritable sock
+    sClose        (Restricted mb sock) = maybe (NS.sClose sock) (\h -> hClose h >> NS.sClose sock) mb
+
+-- | Create a 'RestrictedSocket' that explicitly disallows sending or
+-- receiving data.
+restrictSocket :: NS.Socket -> RestrictedSocket
+restrictSocket socket = Restricted Nothing socket
+
+-- | Build a 'RestrictedSocket' for which 'sClose' will close the given
+-- 'Handle'.  It is intended that this 'Handle' was obtained via
+-- 'NS.socketToHandle'.
+restrictHandleSocket :: Handle ->  NS.Socket -> RestrictedSocket
+restrictHandleSocket h socket = Restricted (Just h) socket
diff --git a/src/Network/StreamServer.hs b/src/Network/StreamServer.hs
new file mode 100644
index 00000000..a6cead0e
--- /dev/null
+++ b/src/Network/StreamServer.hs
@@ -0,0 +1,143 @@
+-- | This module implements a bare-bones TCP or Unix socket server.
+{-# LANGUAGE TypeFamilies #-}
+{-# LANGUAGE TypeOperators #-}
+{-# LANGUAGE OverloadedStrings #-}
+{-# LANGUAGE RankNTypes #-}
+module Network.StreamServer
+    ( streamServer
+    , ServerHandle
+    , ServerConfig(..)
+    , withSession
+    , quitListening
+    , dummyServerHandle
+    ) where
+
+import Data.Monoid
+import Network.Socket as Socket
+import Data.ByteString.Char8
+    ( hGetNonBlocking
+    )
+import qualified Data.ByteString.Char8 as S
+    ( hPutStrLn
+    )
+import System.Directory (removeFile)
+import System.IO
+    ( IOMode(..)
+    , hSetBuffering
+    , BufferMode(..)
+    , hWaitForInput
+    , hClose
+    , hIsEOF
+    , hPutStrLn
+    , stderr
+    , hFlush
+    )
+import Control.Monad
+import Control.Monad.Fix (fix)
+import Control.Concurrent (forkIO, threadDelay, ThreadId, mkWeakThreadId)
+import Control.Exception (catch,handle,try,finally)
+import System.IO.Error (tryIOError)
+import System.Mem.Weak
+import System.IO.Error
+
+-- import Data.Conduit
+import Control.Monad.IO.Class (MonadIO (liftIO))
+import qualified Data.ByteString as S (ByteString)
+import System.IO (Handle)
+import Control.Concurrent.MVar (newMVar)
+
+import Network.SocketLike
+
+data ServerHandle = ServerHandle Socket (Weak ThreadId)
+
+
+-- | Create a useless do-nothing 'ServerHandle'.
+dummyServerHandle :: IO ServerHandle
+dummyServerHandle = do
+    mvar <- newMVar Closed
+    let sock = MkSocket 0 AF_UNSPEC NoSocketType 0 mvar
+    thread <- mkWeakThreadId <=< forkIO $ return ()
+    return (ServerHandle sock thread)
+
+removeSocketFile :: SockAddr -> IO ()
+removeSocketFile (SockAddrUnix fname) = removeFile fname
+removeSocketFile _                    = return ()
+
+-- | Terminate the server accept-loop.  Call this to shut down the server.
+quitListening :: ServerHandle -> IO ()
+quitListening (ServerHandle socket _) =
+    finally (Socket.getSocketName socket >>= removeSocketFile)
+            (Socket.close socket)
+
+
+-- | It's 'bshow' instead of 'show' to enable swapping in a 'ByteString'
+-- variation.  (This is not exported.)
+bshow :: Show a => a -> String
+bshow e = show e
+
+-- | Send a string to stderr.  Not exported.  Default 'serverWarn' when
+-- 'withSession' is used to configure the server.
+warnStderr :: String -> IO ()
+warnStderr str = hPutStrLn stderr str >> hFlush stderr
+
+data ServerConfig = ServerConfig
+    { serverWarn    :: String -> IO ()
+    -- ^ Action to report warnings and errors.
+    , serverSession :: RestrictedSocket -> Int -> Handle -> IO ()
+    -- ^ Action to handle interaction with a client
+    }
+
+-- | Initialize a 'ServerConfig' using the provided session handler.
+withSession :: (RestrictedSocket -> Int -> Handle -> IO ()) -> ServerConfig
+withSession session = ServerConfig warnStderr session
+
+-- | Launch a thread to listen at the given bind address and dispatch
+-- to session handler threads on every incomming connection. Supports
+-- IPv4 and IPv6, TCP and unix sockets.
+--
+-- The returned handle can be used with 'quitListening' to terminate the
+-- thread and prevent any new sessions from starting.  Currently active
+-- session threads will not be terminated or signaled in any way.
+streamServer :: ServerConfig -> SockAddr -> IO ServerHandle
+streamServer cfg addr = do
+    let warn   = serverWarn cfg
+        family = case addr of
+                    SockAddrInet  {} -> AF_INET
+                    SockAddrInet6 {} -> AF_INET6
+                    SockAddrUnix  {} -> AF_UNIX
+    sock <- socket family Stream 0
+    setSocketOption sock ReuseAddr 1
+    fix $ \loop ->
+        tryIOError (removeSocketFile addr) >> bind sock addr
+        `catchIOError` \e -> do warn $ "bind-error: " <> bshow addr <> " " <> bshow e
+                                threadDelay 5000000
+                                loop
+    listen sock maxListenQueue
+    thread <- mkWeakThreadId <=< forkIO $ acceptLoop cfg sock 0
+    return (ServerHandle sock thread)
+
+-- | Not exported.  This, combined with 'acceptException' form a mutually recursive
+-- loop that handles incomming connections.  To quit the loop, the socket must be
+-- closed by 'quitListening'.
+acceptLoop :: ServerConfig -> Socket -> Int -> IO ()
+acceptLoop cfg sock n = handle (acceptException cfg n sock) $ do
+    con <- accept sock
+    let conkey = n + 1
+    h <- socketToHandle (fst con) ReadWriteMode
+    forkIO $ serverSession cfg (restrictHandleSocket h (fst con)) conkey h
+    acceptLoop cfg sock (n + 1)
+
+acceptException :: ServerConfig -> Int -> Socket -> IOError -> IO ()
+acceptException cfg n sock ioerror = do
+    Socket.close sock
+    case show (ioeGetErrorType ioerror) of
+      "resource exhausted" -> do -- try again
+                                 serverWarn cfg $ ("acceptLoop: resource exhasted")
+                                 threadDelay 500000
+                                 acceptLoop cfg sock (n + 1)
+      "invalid argument"   -> do -- quit on closed socket
+                                 return ()
+      message              -> do -- unexpected exception
+                                 serverWarn cfg $ ("acceptLoop: "<>bshow message)
+                                 return ()
+