{-# LANGUAGE CPP #-}

module Network.Run.Core (
    resolve,
    openSocket,
    openClientSocket,
    openClientSocketWithOptions,
    openClientSocketWithOpts,
    openServerSocket,
    openServerSocketWithOptions,
    openServerSocketWithOpts,
    openTCPServerSocket,
    openTCPServerSocketWithOptions,
    openTCPServerSocketWithOpts,
    gclose,
    labelMe,
) where

import qualified Data.List.NonEmpty as NE
import Control.Arrow
import Control.Concurrent
import qualified Control.Exception as E
import Control.Monad (when)
import Foreign (Storable)
import GHC.Conc.Sync
import Network.Socket

resolve
    :: SocketType
    -> Maybe HostName
    -> ServiceName
    -> [AddrInfoFlag]
    -> IO AddrInfo
resolve :: SocketType
-> Maybe HostName -> HostName -> [AddrInfoFlag] -> IO AddrInfo
resolve SocketType
socketType Maybe HostName
mhost HostName
port [AddrInfoFlag]
flags =
    NonEmpty AddrInfo -> AddrInfo
forall a. NonEmpty a -> a
NE.head (NonEmpty AddrInfo -> AddrInfo)
-> IO (NonEmpty AddrInfo) -> IO AddrInfo
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Maybe AddrInfo
-> Maybe HostName -> Maybe HostName -> IO (NonEmpty AddrInfo)
forall (t :: * -> *).
GetAddrInfo t =>
Maybe AddrInfo
-> Maybe HostName -> Maybe HostName -> IO (t AddrInfo)
getAddrInfo (AddrInfo -> Maybe AddrInfo
forall a. a -> Maybe a
Just AddrInfo
hints) Maybe HostName
mhost (HostName -> Maybe HostName
forall a. a -> Maybe a
Just HostName
port)
  where
    hints :: AddrInfo
hints =
        AddrInfo
defaultHints
            { addrSocketType = socketType
            , addrFlags = flags
            }

#if !MIN_VERSION_network(3,1,2)
openSocket :: AddrInfo -> IO Socket
openSocket addr = socket (addrFamily addr) (addrSocketType addr) (addrProtocol addr)
#endif

-- | This is the same as
--
-- @
-- 'openClientSocketWithOptions' []
-- @
openClientSocket :: AddrInfo -> IO Socket
openClientSocket :: AddrInfo -> IO Socket
openClientSocket = [(SocketOption, Int)] -> AddrInfo -> IO Socket
openClientSocketWithOptions []

-- | Open a client socket with the given options
--
-- The options are set before 'connect'. This is equivalent to
--
-- @
-- 'openClientSocketWithOpts' . 'map' ('second' 'SockOptValue')
-- @
openClientSocketWithOptions :: [(SocketOption, Int)] -> AddrInfo -> IO Socket
openClientSocketWithOptions :: [(SocketOption, Int)] -> AddrInfo -> IO Socket
openClientSocketWithOptions = [(SocketOption, SockOptValue)] -> AddrInfo -> IO Socket
openClientSocketWithOpts ([(SocketOption, SockOptValue)] -> AddrInfo -> IO Socket)
-> ([(SocketOption, Int)] -> [(SocketOption, SockOptValue)])
-> [(SocketOption, Int)]
-> AddrInfo
-> IO Socket
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((SocketOption, Int) -> (SocketOption, SockOptValue))
-> [(SocketOption, Int)] -> [(SocketOption, SockOptValue)]
forall a b. (a -> b) -> [a] -> [b]
map ((Int -> SockOptValue)
-> (SocketOption, Int) -> (SocketOption, SockOptValue)
forall b c d. (b -> c) -> (d, b) -> (d, c)
forall (a :: * -> * -> *) b c d.
Arrow a =>
a b c -> a (d, b) (d, c)
second Int -> SockOptValue
forall a. Storable a => a -> SockOptValue
SockOptValue)

-- | Open a client socket with the given options
--
-- This must be used rather than 'openClientSocketWithOptions' for options such
-- as 'Network.Socket.Linger' which require a composite value
-- ('Network.Socket.StructLinger').
--
-- The options are set before 'connect'.
openClientSocketWithOpts :: [(SocketOption, SockOptValue)] -> AddrInfo -> IO Socket
openClientSocketWithOpts :: [(SocketOption, SockOptValue)] -> AddrInfo -> IO Socket
openClientSocketWithOpts [(SocketOption, SockOptValue)]
opts AddrInfo
addr = IO Socket
-> (Socket -> IO ()) -> (Socket -> IO Socket) -> IO Socket
forall a b c. IO a -> (a -> IO b) -> (a -> IO c) -> IO c
E.bracketOnError (AddrInfo -> IO Socket
openSocket AddrInfo
addr) Socket -> IO ()
close ((Socket -> IO Socket) -> IO Socket)
-> (Socket -> IO Socket) -> IO Socket
forall a b. (a -> b) -> a -> b
$ \Socket
sock -> do
    ((SocketOption, SockOptValue) -> IO ())
-> [(SocketOption, SockOptValue)] -> IO ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ ((SocketOption -> SockOptValue -> IO ())
-> (SocketOption, SockOptValue) -> IO ()
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry ((SocketOption -> SockOptValue -> IO ())
 -> (SocketOption, SockOptValue) -> IO ())
-> (SocketOption -> SockOptValue -> IO ())
-> (SocketOption, SockOptValue)
-> IO ()
forall a b. (a -> b) -> a -> b
$ Socket -> SocketOption -> SockOptValue -> IO ()
setSockOptValue Socket
sock) [(SocketOption, SockOptValue)]
opts
    Socket -> SockAddr -> IO ()
connect Socket
sock (SockAddr -> IO ()) -> SockAddr -> IO ()
forall a b. (a -> b) -> a -> b
$ AddrInfo -> SockAddr
addrAddress AddrInfo
addr
    Socket -> IO Socket
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return Socket
sock

-- | Open socket for server use
--
-- This is the same as:
--
-- @
-- 'openServerSocketWithOptions' []
-- @
openServerSocket :: AddrInfo -> IO Socket
openServerSocket :: AddrInfo -> IO Socket
openServerSocket = [(SocketOption, Int)] -> AddrInfo -> IO Socket
openServerSocketWithOptions []

-- | Open socket for server use, and set the provided options before binding.
--
-- This is equivalent to
--
-- @
-- 'openServerSocketWithOpts' . 'map' ('second' 'SockOptValue')
-- @
openServerSocketWithOptions :: [(SocketOption, Int)] -> AddrInfo -> IO Socket
openServerSocketWithOptions :: [(SocketOption, Int)] -> AddrInfo -> IO Socket
openServerSocketWithOptions = [(SocketOption, SockOptValue)] -> AddrInfo -> IO Socket
openServerSocketWithOpts ([(SocketOption, SockOptValue)] -> AddrInfo -> IO Socket)
-> ([(SocketOption, Int)] -> [(SocketOption, SockOptValue)])
-> [(SocketOption, Int)]
-> AddrInfo
-> IO Socket
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((SocketOption, Int) -> (SocketOption, SockOptValue))
-> [(SocketOption, Int)] -> [(SocketOption, SockOptValue)]
forall a b. (a -> b) -> [a] -> [b]
map ((Int -> SockOptValue)
-> (SocketOption, Int) -> (SocketOption, SockOptValue)
forall b c d. (b -> c) -> (d, b) -> (d, c)
forall (a :: * -> * -> *) b c d.
Arrow a =>
a b c -> a (d, b) (d, c)
second Int -> SockOptValue
forall a. Storable a => a -> SockOptValue
SockOptValue)

-- | Open socket for server use, and set the provided options before binding.
--
-- In addition to the given options, the socket is configured to
--
-- * allow reuse of local addresses (SO_REUSEADDR)
-- * automatically be closed during a successful @execve@ (FD_CLOEXEC)
-- * bind to the address specified
openServerSocketWithOpts :: [(SocketOption, SockOptValue)] -> AddrInfo -> IO Socket
openServerSocketWithOpts :: [(SocketOption, SockOptValue)] -> AddrInfo -> IO Socket
openServerSocketWithOpts [(SocketOption, SockOptValue)]
opts AddrInfo
addr = IO Socket
-> (Socket -> IO ()) -> (Socket -> IO Socket) -> IO Socket
forall a b c. IO a -> (a -> IO b) -> (a -> IO c) -> IO c
E.bracketOnError (AddrInfo -> IO Socket
openSocket AddrInfo
addr) Socket -> IO ()
close ((Socket -> IO Socket) -> IO Socket)
-> (Socket -> IO Socket) -> IO Socket
forall a b. (a -> b) -> a -> b
$ \Socket
sock -> do
    Socket -> SocketOption -> Int -> IO ()
setSocketOption Socket
sock SocketOption
ReuseAddr Int
1
#if !defined(openbsd_HOST_OS)
    Bool -> IO () -> IO ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (AddrInfo -> Family
addrFamily AddrInfo
addr Family -> Family -> Bool
forall a. Eq a => a -> a -> Bool
== Family
AF_INET6) (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ Socket -> SocketOption -> Int -> IO ()
setSocketOption Socket
sock SocketOption
IPv6Only Int
1
#endif
    ((SocketOption, SockOptValue) -> IO ())
-> [(SocketOption, SockOptValue)] -> IO ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ ((SocketOption -> SockOptValue -> IO ())
-> (SocketOption, SockOptValue) -> IO ()
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry ((SocketOption -> SockOptValue -> IO ())
 -> (SocketOption, SockOptValue) -> IO ())
-> (SocketOption -> SockOptValue -> IO ())
-> (SocketOption, SockOptValue)
-> IO ()
forall a b. (a -> b) -> a -> b
$ Socket -> SocketOption -> SockOptValue -> IO ()
setSockOptValue Socket
sock) [(SocketOption, SockOptValue)]
opts
    Socket -> (CInt -> IO ()) -> IO ()
forall r. Socket -> (CInt -> IO r) -> IO r
withFdSocket Socket
sock CInt -> IO ()
setCloseOnExecIfNeeded
    Socket -> SockAddr -> IO ()
bind Socket
sock (SockAddr -> IO ()) -> SockAddr -> IO ()
forall a b. (a -> b) -> a -> b
$ AddrInfo -> SockAddr
addrAddress AddrInfo
addr
    Socket -> IO Socket
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return Socket
sock

-- | Open TCP socket for server use
--
-- This is the same as:
--
-- @
-- 'openTCPServerSocketWithOptions' []
-- @
openTCPServerSocket :: AddrInfo -> IO Socket
openTCPServerSocket :: AddrInfo -> IO Socket
openTCPServerSocket = [(SocketOption, Int)] -> AddrInfo -> IO Socket
openTCPServerSocketWithOptions []

-- | Open socket for server use, and set the provided options before binding.
--
-- This is equivalent to
--
-- @
-- 'openTCPServerSocketWithOpts' . 'map' ('second' 'SockOptValue')
-- @
openTCPServerSocketWithOptions :: [(SocketOption, Int)] -> AddrInfo -> IO Socket
openTCPServerSocketWithOptions :: [(SocketOption, Int)] -> AddrInfo -> IO Socket
openTCPServerSocketWithOptions = [(SocketOption, SockOptValue)] -> AddrInfo -> IO Socket
openTCPServerSocketWithOpts ([(SocketOption, SockOptValue)] -> AddrInfo -> IO Socket)
-> ([(SocketOption, Int)] -> [(SocketOption, SockOptValue)])
-> [(SocketOption, Int)]
-> AddrInfo
-> IO Socket
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((SocketOption, Int) -> (SocketOption, SockOptValue))
-> [(SocketOption, Int)] -> [(SocketOption, SockOptValue)]
forall a b. (a -> b) -> [a] -> [b]
map ((Int -> SockOptValue)
-> (SocketOption, Int) -> (SocketOption, SockOptValue)
forall b c d. (b -> c) -> (d, b) -> (d, c)
forall (a :: * -> * -> *) b c d.
Arrow a =>
a b c -> a (d, b) (d, c)
second Int -> SockOptValue
forall a. Storable a => a -> SockOptValue
SockOptValue)

-- | Open socket for server use, and set the provided options before binding.
--
-- In addition to the given options, the socket is configured to
--
-- * allow reuse of local addresses (SO_REUSEADDR)
-- * automatically be closed during a successful @execve@ (FD_CLOEXEC)
-- * bind to the address specified
-- * listen with queue length with 1024
openTCPServerSocketWithOpts :: [(SocketOption, SockOptValue)] -> AddrInfo -> IO Socket
openTCPServerSocketWithOpts :: [(SocketOption, SockOptValue)] -> AddrInfo -> IO Socket
openTCPServerSocketWithOpts [(SocketOption, SockOptValue)]
opts AddrInfo
addr = do
    Socket
sock <- [(SocketOption, SockOptValue)] -> AddrInfo -> IO Socket
openServerSocketWithOpts [(SocketOption, SockOptValue)]
opts AddrInfo
addr
    Socket -> Int -> IO ()
listen Socket
sock Int
1024
    Socket -> IO Socket
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return Socket
sock

gclose :: Socket -> IO ()
#if MIN_VERSION_network(3,1,1)
gclose :: Socket -> IO ()
gclose Socket
sock = Socket -> Int -> IO ()
gracefulClose Socket
sock Int
5000
#else
gclose = close
#endif

labelMe :: String -> IO ()
labelMe :: HostName -> IO ()
labelMe HostName
name = do
    ThreadId
tid <- IO ThreadId
myThreadId
    ThreadId -> HostName -> IO ()
labelThread ThreadId
tid HostName
name