{-# LANGUAGE MagicHash #-}
{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE ExistentialQuantification #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}

{-|
Module      : Z.IO.Network.TCP
Description : TCP servers and clients
Copyright   : (c) Dong Han, 2018
License     : BSD
Maintainer  : winterland1989@gmail.com
Stability   : experimental
Portability : non-portable

This module provides an API for creating TCP servers and clients.

-}

module Z.IO.Network.TCP (
  -- * TCP Client
    TCPClientConfig(..)
  , defaultTCPClientConfig
  , initTCPClient
  -- * TCP Server
  , TCPServerConfig(..)
  , defaultTCPServerConfig
  , startTCPServer
  ) where

import           Control.Concurrent.MVar
import           Control.Monad
import           Control.Monad.IO.Class
import           Data.Primitive.PrimArray
import           Foreign.Ptr
import           GHC.Ptr
import           Z.IO.Buffered
import           Z.IO.Exception
import           Z.IO.Network.SocketAddr
import           Z.IO.Resource
import           Z.IO.UV.FFI
import           Z.IO.UV.Manager
import           Data.Coerce

--------------------------------------------------------------------------------

-- | A TCP client configuration
--
data TCPClientConfig = TCPClientConfig
    { TCPClientConfig -> Maybe SocketAddr
tcpClientAddr :: Maybe SocketAddr -- ^ assign a local address, or let OS pick one
    , TCPClientConfig -> SocketAddr
tcpRemoteAddr :: SocketAddr      -- ^ remote target address
    , TCPClientConfig -> Bool
tcpNoDelay :: Bool             -- ^ if we want to use @TCP_NODELAY@
    }

-- | Default config, connect to @localhost:8888@.
defaultTCPClientConfig :: TCPClientConfig
defaultTCPClientConfig :: TCPClientConfig
defaultTCPClientConfig = Maybe SocketAddr -> SocketAddr -> Bool -> TCPClientConfig
TCPClientConfig Maybe SocketAddr
forall a. Maybe a
Nothing (PortNumber -> InetAddr -> SocketAddr
SocketAddrInet PortNumber
8888 InetAddr
inetLoopback) Bool
True

-- | init a TCP client 'Resource', which open a new connect when used.
--
initTCPClient :: HasCallStack => TCPClientConfig -> Resource UVStream
initTCPClient :: TCPClientConfig -> Resource UVStream
initTCPClient TCPClientConfig{Bool
Maybe SocketAddr
SocketAddr
tcpNoDelay :: Bool
tcpRemoteAddr :: SocketAddr
tcpClientAddr :: Maybe SocketAddr
tcpNoDelay :: TCPClientConfig -> Bool
tcpRemoteAddr :: TCPClientConfig -> SocketAddr
tcpClientAddr :: TCPClientConfig -> Maybe SocketAddr
..} = do
    UVManager
uvm <- IO UVManager -> Resource UVManager
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO IO UVManager
getUVManager
    UVStream
client <- HasCallStack => UVManager -> Resource UVStream
UVManager -> Resource UVStream
initTCPStream UVManager
uvm
    let hdl :: Ptr UVHandle
hdl = UVStream -> Ptr UVHandle
uvsHandle UVStream
client
    IO () -> Resource ()
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> Resource ())
-> ((Ptr SocketAddr -> IO ()) -> IO ())
-> (Ptr SocketAddr -> IO ())
-> Resource ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. SocketAddr -> (Ptr SocketAddr -> IO ()) -> IO ()
forall a. SocketAddr -> (Ptr SocketAddr -> IO a) -> IO a
withSocketAddr SocketAddr
tcpRemoteAddr ((Ptr SocketAddr -> IO ()) -> Resource ())
-> (Ptr SocketAddr -> IO ()) -> Resource ()
forall a b. (a -> b) -> a -> b
$ \ Ptr SocketAddr
targetPtr -> do
        Maybe SocketAddr -> (SocketAddr -> IO ()) -> IO ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ Maybe SocketAddr
tcpClientAddr ((SocketAddr -> IO ()) -> IO ()) -> (SocketAddr -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \ SocketAddr
tcpClientAddr' ->
            SocketAddr -> (Ptr SocketAddr -> IO ()) -> IO ()
forall a. SocketAddr -> (Ptr SocketAddr -> IO a) -> IO a
withSocketAddr SocketAddr
tcpClientAddr' ((Ptr SocketAddr -> IO ()) -> IO ())
-> (Ptr SocketAddr -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \ Ptr SocketAddr
localPtr ->
                -- bind is safe without withUVManager
                IO CInt -> IO ()
forall a. (HasCallStack, Integral a) => IO a -> IO ()
throwUVIfMinus_ (Ptr UVHandle -> Ptr SocketAddr -> CUInt -> IO CInt
uv_tcp_bind Ptr UVHandle
hdl Ptr SocketAddr
localPtr CUInt
0)
        -- nodelay is safe without withUVManager
        Bool -> IO () -> IO ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when Bool
tcpNoDelay (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ IO CInt -> IO ()
forall a. (HasCallStack, Integral a) => IO a -> IO ()
throwUVIfMinus_ (Ptr UVHandle -> CInt -> IO CInt
uv_tcp_nodelay Ptr UVHandle
hdl CInt
1)
        IO Int -> IO ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (IO Int -> IO ())
-> ((Ptr UVLoop -> IO UVSlotUnSafe) -> IO Int)
-> (Ptr UVLoop -> IO UVSlotUnSafe)
-> IO ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. HasCallStack =>
UVManager -> (Ptr UVLoop -> IO UVSlotUnSafe) -> IO Int
UVManager -> (Ptr UVLoop -> IO UVSlotUnSafe) -> IO Int
withUVRequest UVManager
uvm ((Ptr UVLoop -> IO UVSlotUnSafe) -> IO ())
-> (Ptr UVLoop -> IO UVSlotUnSafe) -> IO ()
forall a b. (a -> b) -> a -> b
$ \ Ptr UVLoop
_ -> Ptr UVHandle -> Ptr SocketAddr -> IO UVSlotUnSafe
hs_uv_tcp_connect Ptr UVHandle
hdl Ptr SocketAddr
targetPtr
    UVStream -> Resource UVStream
forall (m :: * -> *) a. Monad m => a -> m a
return UVStream
client

--------------------------------------------------------------------------------

-- | A TCP server configuration
--
data TCPServerConfig = TCPServerConfig
    { TCPServerConfig -> SocketAddr
tcpListenAddr       :: SocketAddr      -- ^ listening address
    , TCPServerConfig -> Int
tcpListenBacklog    :: Int           -- ^ listening socket's backlog size, should be large enough(>128)
    , TCPServerConfig -> Bool
tcpServerWorkerNoDelay :: Bool       -- ^ if we want to use @TCP_NODELAY@
    , TCPServerConfig -> UVStream -> IO ()
tcpServerWorker     :: UVStream -> IO ()  -- ^ worker which get an accepted TCP stream,
                                            -- the socket will be closed upon exception or worker finishes.
    }

-- | A default hello world server on localhost:8888
--
-- Test it with @main = startTCPServer defaultTCPServerConfig@, now try @nc -v 127.0.0.1 8888@
--
defaultTCPServerConfig :: TCPServerConfig
defaultTCPServerConfig :: TCPServerConfig
defaultTCPServerConfig = SocketAddr -> Int -> Bool -> (UVStream -> IO ()) -> TCPServerConfig
TCPServerConfig
    (PortNumber -> InetAddr -> SocketAddr
SocketAddrInet PortNumber
8888 InetAddr
inetAny)
    Int
256
    Bool
True
    (\ UVStream
uvs -> UVStream -> Ptr Word8 -> Int -> IO ()
forall o.
(Output o, HasCallStack) =>
o -> Ptr Word8 -> Int -> IO ()
writeOutput UVStream
uvs (Addr# -> Ptr Word8
forall a. Addr# -> Ptr a
Ptr Addr#
"hello world"#) Int
11)

-- | Start a server
--
-- Fork new worker thread upon a new connection.
--
startTCPServer :: HasCallStack => TCPServerConfig -> IO ()
startTCPServer :: TCPServerConfig -> IO ()
startTCPServer TCPServerConfig{Bool
Int
SocketAddr
UVStream -> IO ()
tcpServerWorker :: UVStream -> IO ()
tcpServerWorkerNoDelay :: Bool
tcpListenBacklog :: Int
tcpListenAddr :: SocketAddr
tcpServerWorker :: TCPServerConfig -> UVStream -> IO ()
tcpServerWorkerNoDelay :: TCPServerConfig -> Bool
tcpListenBacklog :: TCPServerConfig -> Int
tcpListenAddr :: TCPServerConfig -> SocketAddr
..} = do
    let backLog :: Int
backLog = Int -> Int -> Int
forall a. Ord a => a -> a -> a
max Int
tcpListenBacklog Int
128
    UVManager
serverUVManager <- IO UVManager
getUVManager
    Resource UVStream -> (UVStream -> IO ()) -> IO ()
forall (m :: * -> *) a b.
(MonadMask m, MonadIO m, HasCallStack) =>
Resource a -> (a -> m b) -> m b
withResource (HasCallStack => UVManager -> Resource UVStream
UVManager -> Resource UVStream
initTCPStream UVManager
serverUVManager) ((UVStream -> IO ()) -> IO ()) -> (UVStream -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \ (UVStream Ptr UVHandle
serverHandle Int
serverSlot UVManager
_ IORef Bool
_) ->
        IO (Ptr UVHandle)
-> (Ptr UVHandle -> IO ()) -> (Ptr UVHandle -> IO ()) -> IO ()
forall a b c. IO a -> (a -> IO b) -> (a -> IO c) -> IO c
bracket
            (IO (Ptr UVHandle) -> IO (Ptr UVHandle)
forall a. HasCallStack => IO (Ptr a) -> IO (Ptr a)
throwOOMIfNull (IO (Ptr UVHandle) -> IO (Ptr UVHandle))
-> IO (Ptr UVHandle) -> IO (Ptr UVHandle)
forall a b. (a -> b) -> a -> b
$ Ptr UVHandle -> IO (Ptr UVHandle)
hs_uv_accept_check_alloc Ptr UVHandle
serverHandle)
            Ptr UVHandle -> IO ()
hs_uv_accept_check_close ((Ptr UVHandle -> IO ()) -> IO ())
-> (Ptr UVHandle -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$
            \ Ptr UVHandle
check -> do
-- The buffer passing of accept is a litte complicated here, to get maximum performance,
-- we do batch accepting. i.e. recv multiple client inside libuv's event loop:
--
-- we poke uvmanager's buffer table as a Ptr Word8, with byte size (backLog*sizeof(UVFD))
-- inside libuv event loop, we cast the buffer back to int32_t* pointer.
-- each accept callback push a new socket fd to the buffer, and increase a counter(buffer_size_table).
-- backLog should be large enough(>128), so under windows we can't possibly filled it up within one
-- uv_run, under unix we hacked uv internal to provide a stop and resume function, when backLog is
-- reached, we will stop receiving.
--
-- once back to haskell side, we read all accepted sockets and fork worker threads.
-- if backLog is reached, we resume receiving from haskell side.
--
-- Step 1.
-- we allocate a new uv_check_t for given uv_stream_t, with predefined checking callback
-- see hs_accept_check_cb in hs_uv_stream.c
                IO CInt -> IO ()
forall a. (HasCallStack, Integral a) => IO a -> IO ()
throwUVIfMinus_ (IO CInt -> IO ()) -> IO CInt -> IO ()
forall a b. (a -> b) -> a -> b
$ Ptr UVHandle -> IO CInt
hs_uv_accept_check_init Ptr UVHandle
check
                SocketAddr -> (Ptr SocketAddr -> IO ()) -> IO ()
forall a. SocketAddr -> (Ptr SocketAddr -> IO a) -> IO a
withSocketAddr SocketAddr
tcpListenAddr ((Ptr SocketAddr -> IO ()) -> IO ())
-> (Ptr SocketAddr -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \ Ptr SocketAddr
addrPtr -> do
                    MVar Int
m <- UVManager -> Int -> IO (MVar Int)
getBlockMVar UVManager
serverUVManager Int
serverSlot
                    MutablePrimArray RealWorld Int32
acceptBuf <- Int -> IO (MutablePrimArray (PrimState IO) Int32)
forall (m :: * -> *) a.
(PrimMonad m, Prim a) =>
Int -> m (MutablePrimArray (PrimState m) a)
newPinnedPrimArray Int
backLog
                    let acceptBufPtr :: Ptr Word8
acceptBufPtr = Ptr Int32 -> Ptr Word8
coerce (MutablePrimArray RealWorld Int32 -> Ptr Int32
forall s a. MutablePrimArray s a -> Ptr a
mutablePrimArrayContents MutablePrimArray RealWorld Int32
acceptBuf :: Ptr UVFD)
-- Step 2.
-- we allocate a buffer to hold accepted FDs, pass it just like a normal reading buffer.
                    UVManager -> IO () -> IO ()
forall a. HasCallStack => UVManager -> IO a -> IO a
withUVManager' UVManager
serverUVManager (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ do
                        -- We use buffersize as accepted fd count(count backwards)
                        UVManager -> Int -> Ptr Word8 -> Int -> IO ()
pokeBufferTable UVManager
serverUVManager Int
serverSlot Ptr Word8
acceptBufPtr (Int
backLogInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1)
                        IO CInt -> IO ()
forall a. (HasCallStack, Integral a) => IO a -> IO ()
throwUVIfMinus_ (Ptr UVHandle -> Ptr SocketAddr -> CUInt -> IO CInt
uv_tcp_bind Ptr UVHandle
serverHandle Ptr SocketAddr
addrPtr CUInt
0)
                        IO CInt -> IO ()
forall a. (HasCallStack, Integral a) => IO a -> IO ()
throwUVIfMinus_ (Ptr UVHandle -> CInt -> IO CInt
hs_uv_listen Ptr UVHandle
serverHandle (CInt -> CInt -> CInt
forall a. Ord a => a -> a -> a
max CInt
4 (Int -> CInt
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
backLog)))

                    IO () -> IO ()
forall (f :: * -> *) a b. Applicative f => f a -> f b
forever (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ do
                        -- wait until accept some FDs
                        !Int
acceptCountDown <- MVar Int -> IO Int
forall a. MVar a -> IO a
takeMVar MVar Int
m
-- Step 3.
-- Copy buffer, fetch accepted FDs and fork worker threads.

                        -- we lock uv manager here in case of next uv_run overwrite current accept buffer
                        PrimArray Int32
acceptBufCopy <- UVManager -> IO (PrimArray Int32) -> IO (PrimArray Int32)
forall a. HasCallStack => UVManager -> IO a -> IO a
withUVManager' UVManager
serverUVManager (IO (PrimArray Int32) -> IO (PrimArray Int32))
-> IO (PrimArray Int32) -> IO (PrimArray Int32)
forall a b. (a -> b) -> a -> b
$ do
                            Maybe Int
_ <- MVar Int -> IO (Maybe Int)
forall a. MVar a -> IO (Maybe a)
tryTakeMVar MVar Int
m
                            UVManager -> Int -> Ptr Word8 -> Int -> IO ()
pokeBufferTable UVManager
serverUVManager Int
serverSlot Ptr Word8
acceptBufPtr (Int
backLogInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1)

                            -- if acceptCountDown count to -1, we should resume on haskell side
                            Bool -> IO () -> IO ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
acceptCountDown Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== -Int
1) (Ptr UVHandle -> IO ()
hs_uv_listen_resume Ptr UVHandle
serverHandle)

                            -- copy accepted FDs
                            let acceptCount :: Int
acceptCount = Int
backLog Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1 Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
acceptCountDown
                            MutablePrimArray RealWorld Int32
acceptBuf' <- Int -> IO (MutablePrimArray (PrimState IO) Int32)
forall (m :: * -> *) a.
(PrimMonad m, Prim a) =>
Int -> m (MutablePrimArray (PrimState m) a)
newPrimArray Int
acceptCount
                            MutablePrimArray (PrimState IO) Int32
-> Int
-> MutablePrimArray (PrimState IO) Int32
-> Int
-> Int
-> IO ()
forall (m :: * -> *) a.
(PrimMonad m, Prim a) =>
MutablePrimArray (PrimState m) a
-> Int -> MutablePrimArray (PrimState m) a -> Int -> Int -> m ()
copyMutablePrimArray MutablePrimArray RealWorld Int32
MutablePrimArray (PrimState IO) Int32
acceptBuf' Int
0 MutablePrimArray RealWorld Int32
MutablePrimArray (PrimState IO) Int32
acceptBuf (Int
acceptCountDownInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1) Int
acceptCount
                            MutablePrimArray (PrimState IO) Int32 -> IO (PrimArray Int32)
forall (m :: * -> *) a.
PrimMonad m =>
MutablePrimArray (PrimState m) a -> m (PrimArray a)
unsafeFreezePrimArray MutablePrimArray RealWorld Int32
MutablePrimArray (PrimState IO) Int32
acceptBuf'

                        -- fork worker thread
                        [Int] -> (Int -> IO ()) -> IO ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ [Int
0..PrimArray Int32 -> Int
forall a. Prim a => PrimArray a -> Int
sizeofPrimArray PrimArray Int32
acceptBufCopyInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1] ((Int -> IO ()) -> IO ()) -> (Int -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \ Int
i -> do
                            let fd :: Int32
fd = PrimArray Int32 -> Int -> Int32
forall a. Prim a => PrimArray a -> Int -> a
indexPrimArray PrimArray Int32
acceptBufCopy Int
i
                            if Int32
fd Int32 -> Int32 -> Bool
forall a. Ord a => a -> a -> Bool
< Int32
0
                            -- minus fd indicate a server error and we should close server
                            then IO Int32 -> IO ()
forall a. (HasCallStack, Integral a) => IO a -> IO ()
throwUVIfMinus_ (Int32 -> IO Int32
forall (m :: * -> *) a. Monad m => a -> m a
return Int32
fd)
                            -- It's important to use the worker thread's mananger instead of server's one!
                            else IO ThreadId -> IO ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (IO ThreadId -> IO ()) -> (IO () -> IO ThreadId) -> IO () -> IO ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. IO () -> IO ThreadId
forkBa (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ do
                                UVManager
uvm <- IO UVManager
getUVManager
                                Resource UVStream -> (UVStream -> IO ()) -> IO ()
forall (m :: * -> *) a b.
(MonadMask m, MonadIO m, HasCallStack) =>
Resource a -> (a -> m b) -> m b
withResource (HasCallStack =>
(Ptr UVLoop -> Ptr UVHandle -> IO ())
-> UVManager -> Resource UVStream
(Ptr UVLoop -> Ptr UVHandle -> IO ())
-> UVManager -> Resource UVStream
initUVStream (\ Ptr UVLoop
loop Ptr UVHandle
hdl -> do
                                    IO CInt -> IO ()
forall a. (HasCallStack, Integral a) => IO a -> IO ()
throwUVIfMinus_ (Ptr UVLoop -> Ptr UVHandle -> IO CInt
uv_tcp_init Ptr UVLoop
loop Ptr UVHandle
hdl)
                                    IO CInt -> IO ()
forall a. (HasCallStack, Integral a) => IO a -> IO ()
throwUVIfMinus_ (Ptr UVHandle -> Int32 -> IO CInt
hs_uv_tcp_open Ptr UVHandle
hdl Int32
fd)) UVManager
uvm) ((UVStream -> IO ()) -> IO ()) -> (UVStream -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \ UVStream
uvs -> do
                                    Bool -> IO () -> IO ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when Bool
tcpServerWorkerNoDelay (IO () -> IO ()) -> (IO CInt -> IO ()) -> IO CInt -> IO ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. IO CInt -> IO ()
forall a. (HasCallStack, Integral a) => IO a -> IO ()
throwUVIfMinus_ (IO CInt -> IO ()) -> IO CInt -> IO ()
forall a b. (a -> b) -> a -> b
$
                                        -- safe without withUVManager
                                        Ptr UVHandle -> CInt -> IO CInt
uv_tcp_nodelay (UVStream -> Ptr UVHandle
uvsHandle UVStream
uvs) CInt
1
                                    UVStream -> IO ()
tcpServerWorker UVStream
uvs

--------------------------------------------------------------------------------

initTCPStream :: HasCallStack => UVManager -> Resource UVStream
initTCPStream :: UVManager -> Resource UVStream
initTCPStream = HasCallStack =>
(Ptr UVLoop -> Ptr UVHandle -> IO ())
-> UVManager -> Resource UVStream
(Ptr UVLoop -> Ptr UVHandle -> IO ())
-> UVManager -> Resource UVStream
initUVStream (\ Ptr UVLoop
loop Ptr UVHandle
hdl ->
    IO CInt -> IO ()
forall a. (HasCallStack, Integral a) => IO a -> IO ()
throwUVIfMinus_ (Ptr UVLoop -> Ptr UVHandle -> IO CInt
uv_tcp_init Ptr UVLoop
loop Ptr UVHandle
hdl))