2 files changed, 247 insertions, 0 deletions

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