{-# LANGUAGE OverloadedStrings #-}

-- | Simple functions to run TCP clients and servers.
module Network.Run.TCP.Timeout (
    runTCPServer,
    TimeoutServer,

    -- * Generalized API
    runTCPServerWithSocket,
    openServerSocket,
    openServerSocketWithOptions,
    openServerSocketWithOpts,
) where

import Control.Concurrent (forkFinally)
import qualified Control.Exception as E
import Control.Monad (forever, void)
import Network.Socket
import qualified System.TimeManager as T

import Network.Run.Core

-- | A server type
type TimeoutServer a =
    T.Manager
    -- ^ A global timeout manager
    -> T.Handle
    -- ^ A thread-local timeout handler
    -> Socket
    -- ^ A connected socket
    -> IO a

-- | Running a TCP server with a connected socket.
runTCPServer
    :: Int
    -- ^ Timeout in second.
    -> Maybe HostName
    -> ServiceName
    -> TimeoutServer a
    -> IO a
runTCPServer :: forall a.
Int -> Maybe HostName -> HostName -> TimeoutServer a -> IO a
runTCPServer Int
tm Maybe HostName
mhost HostName
port TimeoutServer a
server = IO a -> IO a
forall a. IO a -> IO a
withSocketsDo (IO a -> IO a) -> IO a -> IO a
forall a b. (a -> b) -> a -> b
$ do
    AddrInfo
addr <- SocketType
-> Maybe HostName -> HostName -> [AddrInfoFlag] -> IO AddrInfo
resolve SocketType
Stream Maybe HostName
mhost HostName
port [AddrInfoFlag
AI_PASSIVE]
    IO Socket -> (Socket -> IO ()) -> (Socket -> IO a) -> IO a
forall a b c. IO a -> (a -> IO b) -> (a -> IO c) -> IO c
E.bracket (AddrInfo -> IO Socket
openTCPServerSocket AddrInfo
addr) Socket -> IO ()
close ((Socket -> IO a) -> IO a) -> (Socket -> IO a) -> IO a
forall a b. (a -> b) -> a -> b
$ \Socket
sock ->
        Int -> Socket -> TimeoutServer a -> IO a
forall a. Int -> Socket -> TimeoutServer a -> IO a
runTCPServerWithSocket Int
tm Socket
sock TimeoutServer a
server

-- | Running a TCP client with a connected socket for a given listen
-- socket.
runTCPServerWithSocket
    :: Int
    -- ^ Timeout in second.
    -> Socket
    -> TimeoutServer a
    -> IO a
runTCPServerWithSocket :: forall a. Int -> Socket -> TimeoutServer a -> IO a
runTCPServerWithSocket Int
tm Socket
sock TimeoutServer a
server = IO a -> IO a
forall a. IO a -> IO a
withSocketsDo (IO a -> IO a) -> IO a -> IO a
forall a b. (a -> b) -> a -> b
$ do
    Int -> (Manager -> IO a) -> IO a
forall a. Int -> (Manager -> IO a) -> IO a
T.withManager (Int
tm Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
1000000) ((Manager -> IO a) -> IO a) -> (Manager -> IO a) -> IO a
forall a b. (a -> b) -> a -> b
$ \Manager
mgr -> IO () -> IO a
forall (f :: * -> *) a b. Applicative f => f a -> f b
forever (IO () -> IO a) -> IO () -> IO a
forall a b. (a -> b) -> a -> b
$
        IO (Socket, SockAddr)
-> ((Socket, SockAddr) -> IO ())
-> ((Socket, SockAddr) -> IO ())
-> IO ()
forall a b c. IO a -> (a -> IO b) -> (a -> IO c) -> IO c
E.bracketOnError (Socket -> IO (Socket, SockAddr)
accept Socket
sock) (Socket -> IO ()
close (Socket -> IO ())
-> ((Socket, SockAddr) -> Socket) -> (Socket, SockAddr) -> IO ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Socket, SockAddr) -> Socket
forall a b. (a, b) -> a
fst) (((Socket, SockAddr) -> IO ()) -> IO ())
-> ((Socket, SockAddr) -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \(Socket
conn, SockAddr
_peer) ->
            IO ThreadId -> IO ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (IO ThreadId -> IO ()) -> IO ThreadId -> IO ()
forall a b. (a -> b) -> a -> b
$ IO (Maybe a)
-> (Either SomeException (Maybe a) -> IO ()) -> IO ThreadId
forall a. IO a -> (Either SomeException a -> IO ()) -> IO ThreadId
forkFinally (Manager -> Socket -> IO (Maybe a)
server' Manager
mgr Socket
conn) (IO () -> Either SomeException (Maybe a) -> IO ()
forall a b. a -> b -> a
const (IO () -> Either SomeException (Maybe a) -> IO ())
-> IO () -> Either SomeException (Maybe a) -> IO ()
forall a b. (a -> b) -> a -> b
$ Socket -> IO ()
gclose Socket
conn)
  where
    server' :: Manager -> Socket -> IO (Maybe a)
server' Manager
mgr Socket
conn = do
        HostName -> IO ()
labelMe HostName
"TCP timeout server"
        Manager -> IO () -> (Handle -> IO a) -> IO (Maybe a)
forall a. Manager -> IO () -> (Handle -> IO a) -> IO (Maybe a)
T.withHandle Manager
mgr (() -> IO ()
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return ()) ((Handle -> IO a) -> IO (Maybe a))
-> (Handle -> IO a) -> IO (Maybe a)
forall a b. (a -> b) -> a -> b
$ \Handle
th -> TimeoutServer a
server Manager
mgr Handle
th Socket
conn