1 files changed, 143 insertions, 0 deletions

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 ()
+