{-# LANGUAGE CPP #-}
module Network.RPC.Curryer.StreamlyAdditions where
import Control.Monad.IO.Class
import Network.Socket (Socket, PortNumber, SocketOption, SockAddr(..), maxListenQueue, Family(..), SocketType(..), defaultProtocol, tupleToHostAddress, withSocketsDo, socket, setSocketOption, bind, getSocketName)
import qualified Network.Socket as Net
import Control.Exception (onException)
import Control.Concurrent.MVar
import Data.Word
import qualified Streamly.Internal.Data.Unfold as UF
import Streamly.Network.Socket hiding (acceptor)
#if MIN_VERSION_streamly(0,9,0)
import qualified Streamly.Internal.Data.Stream.StreamD as D
import Streamly.Internal.Data.Unfold.Type (Unfold(..))
#else
import qualified Streamly.Internal.Data.Stream as D
import Streamly.Internal.Data.Unfold (Unfold(..))
#endif

acceptorOnAddr
    :: MonadIO m
    => [(SocketOption, Int)]
    -> Maybe (MVar SockAddr)
    -> Unfold m ((Word8, Word8, Word8, Word8), PortNumber) Socket
acceptorOnAddr :: forall (m :: * -> *).
MonadIO m =>
[(SocketOption, Int)]
-> Maybe (MVar SockAddr)
-> Unfold m ((Word8, Word8, Word8, Word8), PortNumber) Socket
acceptorOnAddr [(SocketOption, Int)]
opts Maybe (MVar SockAddr)
mLock = forall a c (m :: * -> *) b.
(a -> c) -> Unfold m c b -> Unfold m a b
UF.lmap ((Word8, Word8, Word8, Word8), PortNumber)
-> (Int, SockSpec, SockAddr)
f (forall (m :: * -> *).
MonadIO m =>
Maybe (MVar SockAddr) -> Unfold m (Int, SockSpec, SockAddr) Socket
acceptor Maybe (MVar SockAddr)
mLock)
    where
    f :: ((Word8, Word8, Word8, Word8), PortNumber)
-> (Int, SockSpec, SockAddr)
f ((Word8, Word8, Word8, Word8)
addr, PortNumber
port) =
        (Int
maxListenQueue
        , SockSpec
            { sockFamily :: Family
sockFamily = Family
AF_INET
            , sockType :: SocketType
sockType = SocketType
Stream
            , sockProto :: ProtocolNumber
sockProto = ProtocolNumber
defaultProtocol -- TCP
            , sockOpts :: [(SocketOption, Int)]
sockOpts = [(SocketOption, Int)]
opts
            }
        , PortNumber -> HostAddress -> SockAddr
SockAddrInet PortNumber
port ((Word8, Word8, Word8, Word8) -> HostAddress
tupleToHostAddress (Word8, Word8, Word8, Word8)
addr)
        )

acceptor :: MonadIO m => Maybe (MVar SockAddr) -> Unfold m (Int, SockSpec, SockAddr) Socket
acceptor :: forall (m :: * -> *).
MonadIO m =>
Maybe (MVar SockAddr) -> Unfold m (Int, SockSpec, SockAddr) Socket
acceptor Maybe (MVar SockAddr)
mLock = forall (m :: * -> *) b c a.
Functor m =>
(b -> c) -> Unfold m a b -> Unfold m a c
UF.map forall a b. (a, b) -> a
fst (forall (m :: * -> *).
MonadIO m =>
Maybe (MVar SockAddr)
-> Unfold m (Int, SockSpec, SockAddr) (Socket, SockAddr)
listenTuples Maybe (MVar SockAddr)
mLock)

listenTuples :: MonadIO m
    => Maybe (MVar SockAddr)
    -> Unfold m (Int, SockSpec, SockAddr) (Socket, SockAddr)
listenTuples :: forall (m :: * -> *).
MonadIO m =>
Maybe (MVar SockAddr)
-> Unfold m (Int, SockSpec, SockAddr) (Socket, SockAddr)
listenTuples Maybe (MVar SockAddr)
mSockLock = forall (m :: * -> *) a b s.
(s -> m (Step s b)) -> (a -> m s) -> Unfold m a b
Unfold forall {m :: * -> *}.
MonadIO m =>
Socket -> m (Step Socket (Socket, SockAddr))
step forall {m :: * -> *}.
MonadIO m =>
(Int, SockSpec, SockAddr) -> m Socket
inject
 where
    inject :: (Int, SockSpec, SockAddr) -> m Socket
inject (Int
listenQLen, SockSpec
spec, SockAddr
addr) =
      forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ do
        Socket
sock <- Int -> SockSpec -> SockAddr -> IO Socket
initListener Int
listenQLen SockSpec
spec SockAddr
addr
        SockAddr
sockAddr <- Socket -> IO SockAddr
getSocketName Socket
sock
        case Maybe (MVar SockAddr)
mSockLock of
          Just MVar SockAddr
mvar -> forall a. MVar a -> a -> IO ()
putMVar MVar SockAddr
mvar SockAddr
sockAddr
          Maybe (MVar SockAddr)
Nothing -> forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
        forall (f :: * -> *) a. Applicative f => a -> f a
pure Socket
sock

    step :: Socket -> m (Step Socket (Socket, SockAddr))
step Socket
listener = do
        (Socket, SockAddr)
r <- forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (Socket -> IO (Socket, SockAddr)
Net.accept Socket
listener forall a b. IO a -> IO b -> IO a
`onException` Socket -> IO ()
Net.close Socket
listener)
        forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. a -> s -> Step s a
D.Yield (Socket, SockAddr)
r Socket
listener

initListener :: Int -> SockSpec -> SockAddr -> IO Socket
initListener :: Int -> SockSpec -> SockAddr -> IO Socket
initListener Int
listenQLen SockSpec
sockSpec SockAddr
addr =
  forall a. IO a -> IO a
withSocketsDo forall a b. (a -> b) -> a -> b
$ do
    Socket
sock <- Family -> SocketType -> ProtocolNumber -> IO Socket
socket (SockSpec -> Family
sockFamily SockSpec
sockSpec) (SockSpec -> SocketType
sockType SockSpec
sockSpec) (SockSpec -> ProtocolNumber
sockProto SockSpec
sockSpec)
    Socket -> IO ()
use Socket
sock forall a b. IO a -> IO b -> IO a
`onException` Socket -> IO ()
Net.close Socket
sock
    forall (m :: * -> *) a. Monad m => a -> m a
return Socket
sock

    where

    use :: Socket -> IO ()
use Socket
sock = do
        forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ (forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry (Socket -> SocketOption -> Int -> IO ()
setSocketOption Socket
sock)) (SockSpec -> [(SocketOption, Int)]
sockOpts SockSpec
sockSpec)
        Socket -> SockAddr -> IO ()
bind Socket
sock SockAddr
addr
        Socket -> Int -> IO ()
Net.listen Socket
sock Int
listenQLen