{-|
Module      : Z.IO.Network.UDP
Description : UDP 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 UDP sender and receiver.

* Socket FD is created lazily if no local address is provided, that means various functions
  that need FD will throw bad FD exception if you 'initUDP' with no local address e.g. 'setTTL'.

* If you want to create a socket FD but don't care about which port or interface you're using,
  use 'SocketAddrIPv4' 'portAny' 'ipv4Any' when 'initUDP'.

* Prefer 'recvUDPLoop' because it can reuse receive buffer.

-}

module Z.IO.Network.UDP (
  -- * TCP Client
    UDP
  , initUDP
  , UDPConfig(..)
  , defaultUDPConfig
  , sendUDP
  , UDPRecvConfig(..)
  , defaultUDPRecvConfig
  , recvUDPLoop
  , recvUDP
  , getSockName
  -- * Connected UDP Client
  , ConnectedUDP
  , connectUDP
  , disconnectUDP
  , getPeerName
  , sendConnectedUDP
  -- * multicast and broadcast
  , setMembership
  , setSourceMembership
  , setMulticastLoop
  , setMulticastTTL
  , setMulticastInterface
  , setBroadcast
  , setTTL
  -- * Constants
  -- ** UDPFlag
  , UDPFlag
  , pattern UDP_DEFAULT
  , pattern UDP_IPV6ONLY
  , pattern UDP_REUSEADDR
  -- ** Membership
  , Membership
  , pattern JOIN_GROUP
  , pattern LEAVE_GROUP
  ) where

import Control.Concurrent
import Control.Monad
import qualified Data.Primitive.PrimArray       as A
import Data.IORef
import Data.Word
import Data.Int
import Data.Bits ((.&.))
import Foreign.Storable (peek, poke)
import Foreign.Ptr (plusPtr)
import Foreign.C
import GHC.Generics
import Z.Data.CBytes
import qualified Z.Data.Vector.Base             as V
import qualified Z.Data.Vector.Extra            as V
import qualified Z.Data.Text.Print              as T
import Z.Data.JSON                              (JSON)
import Z.IO.Network.SocketAddr
import Z.Foreign
import Z.IO.UV.FFI
import Z.IO.UV.Manager
import Z.IO.Exception
import Z.IO.Resource

-- | UDP socket.
--
-- UDP is not a sequential protocol, thus not an instance of 'Input\/Output'.
-- Message are received or sent individually, UDP socket is NOT thread safe!
-- Use 'MVar' 'UDP' in multiple threads.
--
data UDP = UDP
    { UDP -> Ptr UVHandle
_udpHandle  :: {-# UNPACK #-} !(Ptr UVHandle)
    , UDP -> Int
_udpSlot    :: {-# UNPACK #-} !UVSlot
    , UDP -> UVManager
_udpManager :: UVManager
    , UDP -> MutablePrimArray RealWorld Word8
_udpSendBuffer ::  {-# UNPACK #-} !(A.MutablePrimArray RealWorld Word8)
    , UDP -> IORef Bool
udpClosed  :: {-# UNPACK #-} !(IORef Bool)
    }

instance Show UDP where show :: UDP -> String
show = forall a. Print a => a -> String
T.toString
instance T.Print UDP where
    {-# INLINABLE toUTF8BuilderP #-}
    toUTF8BuilderP :: Int -> UDP -> Builder ()
toUTF8BuilderP Int
_ (UDP Ptr UVHandle
hdl Int
slot UVManager
uvm MutablePrimArray RealWorld Word8
_ IORef Bool
_) = do
        Builder ()
"UDP{udpHandle="    forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> forall a. Print a => a -> Builder ()
T.toUTF8Builder Ptr UVHandle
hdl
        Builder ()
",udpSlot="         forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> forall a. Print a => a -> Builder ()
T.toUTF8Builder Int
slot
        Builder ()
",udpManager="      forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> forall a. Print a => a -> Builder ()
T.toUTF8Builder UVManager
uvm
        Char -> Builder ()
T.char7 Char
'}'

-- | UDP options.
--
-- Though technically message length field in the UDP header is a max of 65535, but large packets
-- could be more likely dropped by routers,
-- usually a packet(IPV4) with a payload <= 508 bytes is considered safe.
data UDPConfig = UDPConfig
    { UDPConfig -> Int
udpSendMsgSize :: {-# UNPACK #-} !Int         -- ^ maximum size of sending buffer
    , UDPConfig -> Maybe (SocketAddr, UDPFlag)
udpLocalAddr :: Maybe (SocketAddr, UDPFlag)   -- ^ do we want bind a local address before receiving & sending?
                                                    --   set to Nothing to let OS pick a random one.
    } deriving (UDPConfig -> UDPConfig -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: UDPConfig -> UDPConfig -> Bool
$c/= :: UDPConfig -> UDPConfig -> Bool
== :: UDPConfig -> UDPConfig -> Bool
$c== :: UDPConfig -> UDPConfig -> Bool
Eq, Eq UDPConfig
UDPConfig -> UDPConfig -> Bool
UDPConfig -> UDPConfig -> Ordering
UDPConfig -> UDPConfig -> UDPConfig
forall a.
Eq a
-> (a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
min :: UDPConfig -> UDPConfig -> UDPConfig
$cmin :: UDPConfig -> UDPConfig -> UDPConfig
max :: UDPConfig -> UDPConfig -> UDPConfig
$cmax :: UDPConfig -> UDPConfig -> UDPConfig
>= :: UDPConfig -> UDPConfig -> Bool
$c>= :: UDPConfig -> UDPConfig -> Bool
> :: UDPConfig -> UDPConfig -> Bool
$c> :: UDPConfig -> UDPConfig -> Bool
<= :: UDPConfig -> UDPConfig -> Bool
$c<= :: UDPConfig -> UDPConfig -> Bool
< :: UDPConfig -> UDPConfig -> Bool
$c< :: UDPConfig -> UDPConfig -> Bool
compare :: UDPConfig -> UDPConfig -> Ordering
$ccompare :: UDPConfig -> UDPConfig -> Ordering
Ord, Int -> UDPConfig -> ShowS
[UDPConfig] -> ShowS
UDPConfig -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [UDPConfig] -> ShowS
$cshowList :: [UDPConfig] -> ShowS
show :: UDPConfig -> String
$cshow :: UDPConfig -> String
showsPrec :: Int -> UDPConfig -> ShowS
$cshowsPrec :: Int -> UDPConfig -> ShowS
Show, forall x. Rep UDPConfig x -> UDPConfig
forall x. UDPConfig -> Rep UDPConfig x
forall a.
(forall x. a -> Rep a x) -> (forall x. Rep a x -> a) -> Generic a
$cto :: forall x. Rep UDPConfig x -> UDPConfig
$cfrom :: forall x. UDPConfig -> Rep UDPConfig x
Generic)
      deriving anyclass (Int -> UDPConfig -> Builder ()
forall a. (Int -> a -> Builder ()) -> Print a
toUTF8BuilderP :: Int -> UDPConfig -> Builder ()
$ctoUTF8BuilderP :: Int -> UDPConfig -> Builder ()
T.Print, Value -> Converter UDPConfig
UDPConfig -> Value
UDPConfig -> Builder ()
forall a.
(Value -> Converter a)
-> (a -> Value) -> (a -> Builder ()) -> JSON a
encodeJSON :: UDPConfig -> Builder ()
$cencodeJSON :: UDPConfig -> Builder ()
toValue :: UDPConfig -> Value
$ctoValue :: UDPConfig -> Value
fromValue :: Value -> Converter UDPConfig
$cfromValue :: Value -> Converter UDPConfig
JSON)

-- | @UDPConfig 512 Nothing@
defaultUDPConfig :: UDPConfig
{-# INLINABLE defaultUDPConfig #-}
defaultUDPConfig :: UDPConfig
defaultUDPConfig = Int -> Maybe (SocketAddr, UDPFlag) -> UDPConfig
UDPConfig Int
512 forall a. Maybe a
Nothing

-- | Initialize a UDP socket.
--
initUDP :: UDPConfig -> Resource UDP
{-# INLINABLE initUDP #-}
initUDP :: UDPConfig -> Resource UDP
initUDP (UDPConfig Int
sbsiz Maybe (SocketAddr, UDPFlag)
maddr) = forall a. IO a -> (a -> IO ()) -> Resource a
initResource
    (do UVManager
uvm <- IO UVManager
getUVManager
        (Ptr UVHandle
hdl, Int
slot) <- forall a. HasCallStack => UVManager -> (Ptr UVLoop -> IO a) -> IO a
withUVManager UVManager
uvm forall a b. (a -> b) -> a -> b
$ \ Ptr UVLoop
loop -> do
            Ptr UVHandle
hdl <- Ptr UVLoop -> IO (Ptr UVHandle)
hs_uv_handle_alloc Ptr UVLoop
loop
            Int
slot <- HasCallStack => UVManager -> IO UVSlotUnsafe -> IO Int
getUVSlot UVManager
uvm (Ptr UVHandle -> IO UVSlotUnsafe
peekUVHandleData Ptr UVHandle
hdl)
            -- init uv struct
            forall a. (HasCallStack, Integral a) => IO a -> IO ()
throwUVIfMinus_ (Ptr UVLoop -> Ptr UVHandle -> IO UDPFlag
uv_udp_init Ptr UVLoop
loop Ptr UVHandle
hdl)
            forall (m :: * -> *) a. Monad m => a -> m a
return (Ptr UVHandle
hdl, Int
slot)

        -- bind the socket if address is available
        -- This is safe without lock UV manager
        forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ Maybe (SocketAddr, UDPFlag)
maddr forall a b. (a -> b) -> a -> b
$ \ (SocketAddr
addr, UDPFlag
flag) ->
            forall a. SocketAddr -> (MBA# SocketAddr -> IO a) -> IO a
withSocketAddrUnsafe SocketAddr
addr forall a b. (a -> b) -> a -> b
$ \ MBA# SocketAddr
p ->
                forall a. (HasCallStack, Integral a) => IO a -> IO ()
throwUVIfMinus_ (Ptr UVHandle -> MBA# SocketAddr -> UDPFlag -> IO UDPFlag
uv_udp_bind Ptr UVHandle
hdl MBA# SocketAddr
p UDPFlag
flag)

        MutablePrimArray RealWorld Word8
sbuf <- forall (m :: * -> *) a.
(PrimMonad m, Prim a) =>
Int -> m (MutablePrimArray (PrimState m) a)
A.newPinnedPrimArray (forall a. Ord a => a -> a -> a
max Int
0 Int
sbsiz)
        IORef Bool
closed <- forall a. a -> IO (IORef a)
newIORef Bool
False
        forall (m :: * -> *) a. Monad m => a -> m a
return (Ptr UVHandle
-> Int
-> UVManager
-> MutablePrimArray RealWorld Word8
-> IORef Bool
-> UDP
UDP Ptr UVHandle
hdl Int
slot UVManager
uvm MutablePrimArray RealWorld Word8
sbuf IORef Bool
closed))
    (\ (UDP Ptr UVHandle
hdl Int
_ UVManager
uvm MutablePrimArray RealWorld Word8
_  IORef Bool
closed) -> forall a. HasCallStack => UVManager -> IO a -> IO a
withUVManager' UVManager
uvm forall a b. (a -> b) -> a -> b
$ do
        Bool
c <- forall a. IORef a -> IO a
readIORef IORef Bool
closed
        -- hs_uv_handle_close won't return error
        forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless Bool
c forall a b. (a -> b) -> a -> b
$ forall a. IORef a -> a -> IO ()
writeIORef IORef Bool
closed Bool
True forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Ptr UVHandle -> IO ()
hs_uv_handle_close Ptr UVHandle
hdl)

checkUDPClosed :: HasCallStack => UDP -> IO ()
{-# INLINABLE checkUDPClosed #-}
checkUDPClosed :: HasCallStack => UDP -> IO ()
checkUDPClosed UDP
udp = do
    Bool
c <- forall a. IORef a -> IO a
readIORef (UDP -> IORef Bool
udpClosed UDP
udp)
    forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when Bool
c forall a. HasCallStack => IO a
throwECLOSED

-- | Get the local IP and port of the 'UDP'.
getSockName :: HasCallStack => UDP -> IO SocketAddr
{-# INLINABLE getSockName #-}
getSockName :: HasCallStack => UDP -> IO SocketAddr
getSockName udp :: UDP
udp@(UDP Ptr UVHandle
hdl Int
_ UVManager
_ MutablePrimArray RealWorld Word8
_ IORef Bool
_) = do
    HasCallStack => UDP -> IO ()
checkUDPClosed UDP
udp
    (MBA# SocketAddr -> IO ()) -> IO SocketAddr
withSocketAddrStorageUnsafe forall a b. (a -> b) -> a -> b
$ \ MBA# SocketAddr
paddr ->
        forall (f :: * -> *) a. Functor f => f a -> f ()
void forall a b. (a -> b) -> a -> b
$ forall a b. Prim a => a -> (MBA# SocketAddr -> IO b) -> IO (a, b)
withPrimUnsafe (forall a b. (Integral a, Num b) => a -> b
fromIntegral CSize
sizeOfSocketAddrStorage :: CInt) forall a b. (a -> b) -> a -> b
$ \ MBA# SocketAddr
plen ->
            forall a. (HasCallStack, Integral a) => IO a -> IO ()
throwUVIfMinus_ (Ptr UVHandle -> MBA# SocketAddr -> MBA# SocketAddr -> IO UDPFlag
uv_udp_getsockname Ptr UVHandle
hdl MBA# SocketAddr
paddr MBA# SocketAddr
plen)

-- | Wrapper for a connected 'UDP'.
newtype ConnectedUDP = ConnectedUDP UDP

instance Show ConnectedUDP where show :: ConnectedUDP -> String
show = forall a. Print a => a -> String
T.toString
instance T.Print ConnectedUDP where
    toUTF8BuilderP :: Int -> ConnectedUDP -> Builder ()
toUTF8BuilderP Int
_ (ConnectedUDP (UDP Ptr UVHandle
hdl Int
slot UVManager
uvm MutablePrimArray RealWorld Word8
_ IORef Bool
_)) = do
        Builder ()
"ConnectedUDP{udpHandle="    forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> forall a. Print a => a -> Builder ()
T.toUTF8Builder Ptr UVHandle
hdl
        Builder ()
",udpSlot="         forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> forall a. Print a => a -> Builder ()
T.toUTF8Builder Int
slot
        Builder ()
",udpManager="      forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> forall a. Print a => a -> Builder ()
T.toUTF8Builder UVManager
uvm
        Char -> Builder ()
T.char7 Char
'}'

-- | Associate the UDP handle to a remote address and port,
-- so every message sent by this handle is automatically sent to that destination
connectUDP :: HasCallStack => UDP -> SocketAddr -> IO ConnectedUDP
{-# INLINABLE connectUDP #-}
connectUDP :: HasCallStack => UDP -> SocketAddr -> IO ConnectedUDP
connectUDP udp :: UDP
udp@(UDP Ptr UVHandle
hdl Int
_ UVManager
_ MutablePrimArray RealWorld Word8
_ IORef Bool
_) SocketAddr
addr = do
    HasCallStack => UDP -> IO ()
checkUDPClosed UDP
udp
    forall a. SocketAddr -> (MBA# SocketAddr -> IO a) -> IO a
withSocketAddrUnsafe SocketAddr
addr forall a b. (a -> b) -> a -> b
$ \ MBA# SocketAddr
paddr ->
        forall a. (HasCallStack, Integral a) => IO a -> IO ()
throwUVIfMinus_ (Ptr UVHandle -> MBA# SocketAddr -> IO UDPFlag
uv_udp_connect Ptr UVHandle
hdl MBA# SocketAddr
paddr)
    forall (m :: * -> *) a. Monad m => a -> m a
return (UDP -> ConnectedUDP
ConnectedUDP UDP
udp)

-- | Disconnect the UDP handle from a remote address and port.
disconnectUDP :: HasCallStack => ConnectedUDP -> IO UDP
{-# INLINABLE disconnectUDP #-}
disconnectUDP :: HasCallStack => ConnectedUDP -> IO UDP
disconnectUDP (ConnectedUDP udp :: UDP
udp@(UDP Ptr UVHandle
hdl Int
_ UVManager
_ MutablePrimArray RealWorld Word8
_ IORef Bool
_)) = do
    HasCallStack => UDP -> IO ()
checkUDPClosed UDP
udp
    forall a. (HasCallStack, Integral a) => IO a -> IO ()
throwUVIfMinus_ (Ptr UVHandle -> Ptr SocketAddr -> IO UDPFlag
uv_udp_disconnect Ptr UVHandle
hdl forall a. Ptr a
nullPtr)
    forall (m :: * -> *) a. Monad m => a -> m a
return UDP
udp

-- | Get the remote IP and port on 'ConnectedUDP'.
getPeerName :: HasCallStack => ConnectedUDP -> IO SocketAddr
{-# INLINABLE getPeerName #-}
getPeerName :: HasCallStack => ConnectedUDP -> IO SocketAddr
getPeerName (ConnectedUDP udp :: UDP
udp@(UDP Ptr UVHandle
hdl Int
_ UVManager
_ MutablePrimArray RealWorld Word8
_ IORef Bool
_)) = do
    HasCallStack => UDP -> IO ()
checkUDPClosed UDP
udp
    (MBA# SocketAddr -> IO ()) -> IO SocketAddr
withSocketAddrStorageUnsafe forall a b. (a -> b) -> a -> b
$ \ MBA# SocketAddr
paddr ->
        forall (f :: * -> *) a. Functor f => f a -> f ()
void forall a b. (a -> b) -> a -> b
$ forall a b. Prim a => a -> (MBA# SocketAddr -> IO b) -> IO (a, b)
withPrimUnsafe (forall a b. (Integral a, Num b) => a -> b
fromIntegral CSize
sizeOfSocketAddrStorage :: CInt) forall a b. (a -> b) -> a -> b
$ \ MBA# SocketAddr
plen ->
            forall a. (HasCallStack, Integral a) => IO a -> IO ()
throwUVIfMinus_ (Ptr UVHandle -> MBA# SocketAddr -> MBA# SocketAddr -> IO UDPFlag
uv_udp_getpeername Ptr UVHandle
hdl MBA# SocketAddr
paddr MBA# SocketAddr
plen)

-- | Send a UDP message with a connected UDP.
--
-- WARNING: A 'InvalidArgument' with errno 'UV_EMSGSIZE' will be thrown
-- if message is larger than 'sendMsgSize'.
sendConnectedUDP :: HasCallStack => ConnectedUDP -> V.Bytes -> IO ()
{-# INLINABLE sendConnectedUDP #-}
sendConnectedUDP :: HasCallStack => ConnectedUDP -> Bytes -> IO ()
sendConnectedUDP (ConnectedUDP udp :: UDP
udp@(UDP Ptr UVHandle
hdl Int
_ UVManager
uvm MutablePrimArray RealWorld Word8
sbuf IORef Bool
_)) (V.PrimVector PrimArray Word8
ba Int
s Int
la) = forall a. IO a -> IO a
mask_ forall a b. (a -> b) -> a -> b
$ do
    HasCallStack => UDP -> IO ()
checkUDPClosed UDP
udp
    -- copy message to pinned buffer
    Int
lb <- forall (m :: * -> *) a.
(PrimMonad m, Prim a) =>
MutablePrimArray (PrimState m) a -> m Int
getSizeofMutablePrimArray MutablePrimArray RealWorld Word8
sbuf
    forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
la forall a. Ord a => a -> a -> Bool
> Int
lb) (forall a. (HasCallStack, Integral a) => IO a -> IO ()
throwUVIfMinus_ (forall (m :: * -> *) a. Monad m => a -> m a
return UDPFlag
UV_EMSGSIZE))
    forall (m :: * -> *) a.
(PrimMonad m, Prim a) =>
MutablePrimArray (PrimState m) a
-> Int -> PrimArray a -> Int -> Int -> m ()
copyPrimArray MutablePrimArray RealWorld Word8
sbuf Int
0 PrimArray Word8
ba Int
s Int
la
    forall a b. MutablePrimArray RealWorld a -> (Ptr a -> IO b) -> IO b
withMutablePrimArrayContents MutablePrimArray RealWorld Word8
sbuf forall a b. (a -> b) -> a -> b
$ \ Ptr Word8
pbuf -> do
        MVar Int
m <- forall a. HasCallStack => UVManager -> IO a -> IO a
withUVManager' UVManager
uvm forall a b. (a -> b) -> a -> b
$ do
            Int
reqSlot <- HasCallStack => UVManager -> IO UVSlotUnsafe -> IO Int
getUVSlot UVManager
uvm (Ptr UVHandle -> Ptr Word8 -> Int -> IO UVSlotUnsafe
hs_uv_udp_send_connected Ptr UVHandle
hdl Ptr Word8
pbuf Int
la)
            MVar Int
reqMVar <- UVManager -> Int -> IO (MVar Int)
getBlockMVar UVManager
uvm Int
reqSlot
            -- since we locked uv manager here, it won't affect next event
            Maybe Int
_ <- forall a. MVar a -> IO (Maybe a)
tryTakeMVar MVar Int
reqMVar
            forall (m :: * -> *) a. Monad m => a -> m a
return MVar Int
reqMVar
        -- we can't cancel uv_udp_send_t in current libuv
        -- and disaster will happen if buffer got collected.
        -- so we have to turn to uninterruptibleMask_'s help.
        -- i.e. sendUDP is an uninterruptible operation.
        -- OS will guarantee writing a socket will not
        -- hang forever anyway.
        forall a. (HasCallStack, Integral a) => IO a -> IO ()
throwUVIfMinus_  (forall a. IO a -> IO a
uninterruptibleMask_ forall a b. (a -> b) -> a -> b
$ forall a. MVar a -> IO a
takeMVar MVar Int
m)

-- | Send a UDP message to target address.
--
-- WARNING: A 'InvalidArgument' with errno 'UV_EMSGSIZE' will be thrown
-- if message is larger than 'sendMsgSize'.
sendUDP :: HasCallStack => UDP -> SocketAddr -> V.Bytes -> IO ()
{-# INLINABLE sendUDP #-}
sendUDP :: HasCallStack => UDP -> SocketAddr -> Bytes -> IO ()
sendUDP udp :: UDP
udp@(UDP Ptr UVHandle
hdl Int
_ UVManager
uvm MutablePrimArray RealWorld Word8
sbuf IORef Bool
_) SocketAddr
addr (V.PrimVector PrimArray Word8
ba Int
s Int
la) = forall a. IO a -> IO a
mask_ forall a b. (a -> b) -> a -> b
$ do
    HasCallStack => UDP -> IO ()
checkUDPClosed UDP
udp
    -- copy message to pinned buffer
    Int
lb <- forall (m :: * -> *) a.
(PrimMonad m, Prim a) =>
MutablePrimArray (PrimState m) a -> m Int
getSizeofMutablePrimArray MutablePrimArray RealWorld Word8
sbuf
    forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
la forall a. Ord a => a -> a -> Bool
> Int
lb) (forall a. (HasCallStack, Integral a) => IO a -> IO ()
throwUVIfMinus_ (forall (m :: * -> *) a. Monad m => a -> m a
return UDPFlag
UV_EMSGSIZE))
    forall (m :: * -> *) a.
(PrimMonad m, Prim a) =>
MutablePrimArray (PrimState m) a
-> Int -> PrimArray a -> Int -> Int -> m ()
copyPrimArray MutablePrimArray RealWorld Word8
sbuf Int
0 PrimArray Word8
ba Int
s Int
la
    forall a b. MutablePrimArray RealWorld a -> (Ptr a -> IO b) -> IO b
withMutablePrimArrayContents MutablePrimArray RealWorld Word8
sbuf forall a b. (a -> b) -> a -> b
$ \ Ptr Word8
pbuf -> do
        MVar Int
m <- forall a. HasCallStack => UVManager -> IO a -> IO a
withUVManager' UVManager
uvm forall a b. (a -> b) -> a -> b
$ do
            Int
reqSlot <- forall a. SocketAddr -> (MBA# SocketAddr -> IO a) -> IO a
withSocketAddrUnsafe SocketAddr
addr forall a b. (a -> b) -> a -> b
$ \ MBA# SocketAddr
paddr ->
                HasCallStack => UVManager -> IO UVSlotUnsafe -> IO Int
getUVSlot UVManager
uvm (Ptr UVHandle
-> MBA# SocketAddr -> Ptr Word8 -> Int -> IO UVSlotUnsafe
hs_uv_udp_send Ptr UVHandle
hdl MBA# SocketAddr
paddr Ptr Word8
pbuf Int
la)
            MVar Int
reqMVar <- UVManager -> Int -> IO (MVar Int)
getBlockMVar UVManager
uvm Int
reqSlot
            -- since we locked uv manager here, it won't affect next event
            Maybe Int
_ <- forall a. MVar a -> IO (Maybe a)
tryTakeMVar MVar Int
reqMVar
            forall (m :: * -> *) a. Monad m => a -> m a
return MVar Int
reqMVar
        -- we can't cancel uv_udp_send_t in current libuv
        -- and disaster will happen if buffer got collected.
        -- so we have to turn to uninterruptibleMask_'s help.
        -- i.e. sendUDP is an uninterruptible operation.
        -- OS will guarantee writing a socket will not
        -- hang forever anyway.
        forall a. (HasCallStack, Integral a) => IO a -> IO ()
throwUVIfMinus_  (forall a. IO a -> IO a
uninterruptibleMask_ forall a b. (a -> b) -> a -> b
$ forall a. MVar a -> IO a
takeMVar MVar Int
m)

-- | Set IP multicast loop flag. Makes multicast packets loop back to local sockets.
setMulticastLoop :: HasCallStack => UDP -> Bool -> IO ()
{-# INLINABLE setMulticastLoop #-}
setMulticastLoop :: HasCallStack => UDP -> Bool -> IO ()
setMulticastLoop udp :: UDP
udp@(UDP Ptr UVHandle
hdl Int
_ UVManager
_ MutablePrimArray RealWorld Word8
_ IORef Bool
_) Bool
loop = do
    HasCallStack => UDP -> IO ()
checkUDPClosed UDP
udp
    forall a. (HasCallStack, Integral a) => IO a -> IO ()
throwUVIfMinus_ (Ptr UVHandle -> UDPFlag -> IO UDPFlag
uv_udp_set_multicast_loop Ptr UVHandle
hdl (if Bool
loop then UDPFlag
1 else UDPFlag
0))

-- | Set the multicast ttl.
setMulticastTTL :: HasCallStack => UDP -> Int -> IO ()
{-# INLINABLE setMulticastTTL #-}
setMulticastTTL :: HasCallStack => UDP -> Int -> IO ()
setMulticastTTL udp :: UDP
udp@(UDP Ptr UVHandle
hdl Int
_ UVManager
_ MutablePrimArray RealWorld Word8
_ IORef Bool
_) Int
ttl = do
    HasCallStack => UDP -> IO ()
checkUDPClosed UDP
udp
    forall a. (HasCallStack, Integral a) => IO a -> IO ()
throwUVIfMinus_ (Ptr UVHandle -> UDPFlag -> IO UDPFlag
uv_udp_set_multicast_ttl Ptr UVHandle
hdl (forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
ttl'))
  where ttl' :: Int
ttl' = Int -> Int -> Int -> Int
V.rangeCut Int
ttl Int
1 Int
255

-- | Set the multicast interface to send or receive data on.
setMulticastInterface :: HasCallStack => UDP -> CBytes ->IO ()
{-# INLINABLE setMulticastInterface #-}
setMulticastInterface :: HasCallStack => UDP -> CBytes -> IO ()
setMulticastInterface udp :: UDP
udp@(UDP Ptr UVHandle
hdl Int
_ UVManager
_ MutablePrimArray RealWorld Word8
_ IORef Bool
_) CBytes
iaddr = do
    HasCallStack => UDP -> IO ()
checkUDPClosed UDP
udp
    forall a. CBytes -> (BA# Word8 -> IO a) -> IO a
withCBytesUnsafe CBytes
iaddr forall a b. (a -> b) -> a -> b
$ \ BA# Word8
iaddrp ->
        forall a. (HasCallStack, Integral a) => IO a -> IO ()
throwUVIfMinus_ (Ptr UVHandle -> BA# Word8 -> IO UDPFlag
uv_udp_set_multicast_interface Ptr UVHandle
hdl BA# Word8
iaddrp)

-- | Set broadcast on or off.
setBroadcast :: HasCallStack => UDP -> Bool -> IO ()
{-# INLINABLE setBroadcast #-}
setBroadcast :: HasCallStack => UDP -> Bool -> IO ()
setBroadcast udp :: UDP
udp@(UDP Ptr UVHandle
hdl Int
_ UVManager
_ MutablePrimArray RealWorld Word8
_ IORef Bool
_) Bool
b = do
    HasCallStack => UDP -> IO ()
checkUDPClosed UDP
udp
    forall a. (HasCallStack, Integral a) => IO a -> IO ()
throwUVIfMinus_ (Ptr UVHandle -> UDPFlag -> IO UDPFlag
uv_udp_set_broadcast Ptr UVHandle
hdl (if Bool
b then UDPFlag
1 else UDPFlag
0))

-- | Set the time to live.
setTTL :: HasCallStack
       => UDP
       -> Int       -- ^ 1 ~ 255
       -> IO ()
{-# INLINABLE setTTL #-}
setTTL :: HasCallStack => UDP -> Int -> IO ()
setTTL udp :: UDP
udp@(UDP Ptr UVHandle
hdl Int
_ UVManager
_ MutablePrimArray RealWorld Word8
_ IORef Bool
_) Int
ttl = do
    HasCallStack => UDP -> IO ()
checkUDPClosed UDP
udp
    forall a. (HasCallStack, Integral a) => IO a -> IO ()
throwUVIfMinus_ (Ptr UVHandle -> UDPFlag -> IO UDPFlag
uv_udp_set_ttl Ptr UVHandle
hdl (forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
ttl))

-- | Set membership for a multicast group.
setMembership :: HasCallStack
              => UDP
              -> CBytes             -- ^ Multicast address to set membership for.
              -> CBytes             -- ^ Interface address.
              -> Membership       -- ^ UV_JOIN_GROUP | UV_LEAVE_GROUP
              -> IO ()
{-# INLINABLE setMembership #-}
setMembership :: HasCallStack => UDP -> CBytes -> CBytes -> UDPFlag -> IO ()
setMembership udp :: UDP
udp@(UDP Ptr UVHandle
hdl Int
_ UVManager
_ MutablePrimArray RealWorld Word8
_ IORef Bool
_) CBytes
gaddr CBytes
iaddr UDPFlag
member = do
    HasCallStack => UDP -> IO ()
checkUDPClosed UDP
udp
    forall a. CBytes -> (BA# Word8 -> IO a) -> IO a
withCBytesUnsafe CBytes
gaddr forall a b. (a -> b) -> a -> b
$ \ BA# Word8
gaddrp ->
        forall a. CBytes -> (BA# Word8 -> IO a) -> IO a
withCBytesUnsafe CBytes
iaddr forall a b. (a -> b) -> a -> b
$ \ BA# Word8
iaddrp ->
            forall a. (HasCallStack, Integral a) => IO a -> IO ()
throwUVIfMinus_ (Ptr UVHandle -> BA# Word8 -> BA# Word8 -> UDPFlag -> IO UDPFlag
uv_udp_set_membership Ptr UVHandle
hdl BA# Word8
gaddrp BA# Word8
iaddrp UDPFlag
member)

-- | Set membership for a source-specific multicast group.
setSourceMembership :: HasCallStack
                    => UDP
                    -> CBytes           -- ^ Multicast address to set membership for.
                    -> CBytes           -- ^ Interface address.
                    -> CBytes           -- ^ Source address.
                    -> Membership     -- ^ UV_JOIN_GROUP | UV_LEAVE_GROUP
                    -> IO ()
{-# INLINABLE setSourceMembership #-}
setSourceMembership :: HasCallStack =>
UDP -> CBytes -> CBytes -> CBytes -> UDPFlag -> IO ()
setSourceMembership udp :: UDP
udp@(UDP Ptr UVHandle
hdl Int
_ UVManager
_ MutablePrimArray RealWorld Word8
_ IORef Bool
_) CBytes
gaddr CBytes
iaddr CBytes
source UDPFlag
member = do
    HasCallStack => UDP -> IO ()
checkUDPClosed UDP
udp
    forall a. CBytes -> (BA# Word8 -> IO a) -> IO a
withCBytesUnsafe CBytes
gaddr forall a b. (a -> b) -> a -> b
$ \ BA# Word8
gaddrp ->
        forall a. CBytes -> (BA# Word8 -> IO a) -> IO a
withCBytesUnsafe CBytes
iaddr forall a b. (a -> b) -> a -> b
$ \ BA# Word8
iaddrp ->
            forall a. CBytes -> (BA# Word8 -> IO a) -> IO a
withCBytesUnsafe CBytes
source forall a b. (a -> b) -> a -> b
$ \ BA# Word8
sourcep ->
                forall a. (HasCallStack, Integral a) => IO a -> IO ()
throwUVIfMinus_ (Ptr UVHandle
-> BA# Word8 -> BA# Word8 -> BA# Word8 -> UDPFlag -> IO UDPFlag
uv_udp_set_source_membership Ptr UVHandle
hdl BA# Word8
gaddrp BA# Word8
iaddrp BA# Word8
sourcep UDPFlag
member)

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

-- | Receiving buffering config.
--
data UDPRecvConfig = UDPRecvConfig
    { UDPRecvConfig -> Int32
recvMsgSize :: {-# UNPACK #-} !Int32      -- ^ maximum size of a received message
    , UDPRecvConfig -> Int
recvBatchSize :: {-# UNPACK #-} !Int      -- ^ how many messages we want to receive per uv loop,
                                                --   inside each uv_run, we do batch receiving,
                                                --   increase this number can improve receiving performance,
                                                --   at the cost of memory and potential GHC thread starving.
    } deriving (UDPRecvConfig -> UDPRecvConfig -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: UDPRecvConfig -> UDPRecvConfig -> Bool
$c/= :: UDPRecvConfig -> UDPRecvConfig -> Bool
== :: UDPRecvConfig -> UDPRecvConfig -> Bool
$c== :: UDPRecvConfig -> UDPRecvConfig -> Bool
Eq, Eq UDPRecvConfig
UDPRecvConfig -> UDPRecvConfig -> Bool
UDPRecvConfig -> UDPRecvConfig -> Ordering
UDPRecvConfig -> UDPRecvConfig -> UDPRecvConfig
forall a.
Eq a
-> (a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
min :: UDPRecvConfig -> UDPRecvConfig -> UDPRecvConfig
$cmin :: UDPRecvConfig -> UDPRecvConfig -> UDPRecvConfig
max :: UDPRecvConfig -> UDPRecvConfig -> UDPRecvConfig
$cmax :: UDPRecvConfig -> UDPRecvConfig -> UDPRecvConfig
>= :: UDPRecvConfig -> UDPRecvConfig -> Bool
$c>= :: UDPRecvConfig -> UDPRecvConfig -> Bool
> :: UDPRecvConfig -> UDPRecvConfig -> Bool
$c> :: UDPRecvConfig -> UDPRecvConfig -> Bool
<= :: UDPRecvConfig -> UDPRecvConfig -> Bool
$c<= :: UDPRecvConfig -> UDPRecvConfig -> Bool
< :: UDPRecvConfig -> UDPRecvConfig -> Bool
$c< :: UDPRecvConfig -> UDPRecvConfig -> Bool
compare :: UDPRecvConfig -> UDPRecvConfig -> Ordering
$ccompare :: UDPRecvConfig -> UDPRecvConfig -> Ordering
Ord, Int -> UDPRecvConfig -> ShowS
[UDPRecvConfig] -> ShowS
UDPRecvConfig -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [UDPRecvConfig] -> ShowS
$cshowList :: [UDPRecvConfig] -> ShowS
show :: UDPRecvConfig -> String
$cshow :: UDPRecvConfig -> String
showsPrec :: Int -> UDPRecvConfig -> ShowS
$cshowsPrec :: Int -> UDPRecvConfig -> ShowS
Show, ReadPrec [UDPRecvConfig]
ReadPrec UDPRecvConfig
Int -> ReadS UDPRecvConfig
ReadS [UDPRecvConfig]
forall a.
(Int -> ReadS a)
-> ReadS [a] -> ReadPrec a -> ReadPrec [a] -> Read a
readListPrec :: ReadPrec [UDPRecvConfig]
$creadListPrec :: ReadPrec [UDPRecvConfig]
readPrec :: ReadPrec UDPRecvConfig
$creadPrec :: ReadPrec UDPRecvConfig
readList :: ReadS [UDPRecvConfig]
$creadList :: ReadS [UDPRecvConfig]
readsPrec :: Int -> ReadS UDPRecvConfig
$creadsPrec :: Int -> ReadS UDPRecvConfig
Read, forall x. Rep UDPRecvConfig x -> UDPRecvConfig
forall x. UDPRecvConfig -> Rep UDPRecvConfig x
forall a.
(forall x. a -> Rep a x) -> (forall x. Rep a x -> a) -> Generic a
$cto :: forall x. Rep UDPRecvConfig x -> UDPRecvConfig
$cfrom :: forall x. UDPRecvConfig -> Rep UDPRecvConfig x
Generic)
      deriving anyclass (Int -> UDPRecvConfig -> Builder ()
forall a. (Int -> a -> Builder ()) -> Print a
toUTF8BuilderP :: Int -> UDPRecvConfig -> Builder ()
$ctoUTF8BuilderP :: Int -> UDPRecvConfig -> Builder ()
T.Print, Value -> Converter UDPRecvConfig
UDPRecvConfig -> Value
UDPRecvConfig -> Builder ()
forall a.
(Value -> Converter a)
-> (a -> Value) -> (a -> Builder ()) -> JSON a
encodeJSON :: UDPRecvConfig -> Builder ()
$cencodeJSON :: UDPRecvConfig -> Builder ()
toValue :: UDPRecvConfig -> Value
$ctoValue :: UDPRecvConfig -> Value
fromValue :: Value -> Converter UDPRecvConfig
$cfromValue :: Value -> Converter UDPRecvConfig
JSON)

-- | @UDPRecvConfig 512 6@
defaultUDPRecvConfig :: UDPRecvConfig
{-# INLINABLE defaultUDPRecvConfig #-}
defaultUDPRecvConfig :: UDPRecvConfig
defaultUDPRecvConfig = Int32 -> Int -> UDPRecvConfig
UDPRecvConfig Int32
512 Int
6


-- The buffer passing of UDP is a litte complicated here, to get maximum performance,
-- we do batch receiving. i.e. recv multiple messages inside libuv's event loop:
--
--   udpRecvLargeBuffer:
--
--   +---------+--------------+-----------+----------+--------+---------+------------
--   | buf siz | partial flag | addr flag |   addr   | buffer | buf siz | partial ...
--   +--4bytes-+----4bytes----+--4bytes---+-128bytes-+-bufsiz-+---------+------------
--   ^                                                        ^
--   |                                                        |
--   +---------------------+       +--------------------------+
--                         |       |
--                      +--+---+---+--+----
--   udpRecvBufferArray | buf0 | buf1 | ...
--                      +------+------+----
--
-- We allocate a large buffer (buffer_size * buffer_number),
-- each time we poke the udpRecvBufferArray and its last index (size - 1) to uv manager's buffer table.
--
-- On libuv side each alloc callback picks the last pointer from udpRecvBufferArray, decrease last index by 1
-- the read result is write into the `buf siz` cell, then followed with partial flag, if addr is not NULL
-- then addr flag is 1 (otherwise 0), following addr if not NULL, the buffer is already written when
-- recv callback is called.
--
-- On haskell side, we read buffer table's size, which is decreased by n(which is the times callback are called).
-- Then we poke those cells out.
--
newRecvBuf :: Int32 -> Int -> IO (A.MutablePrimArray RealWorld Word8, A.MutablePrimArray RealWorld (Ptr Word8))
{-# INLINABLE newRecvBuf #-}
newRecvBuf :: Int32
-> Int
-> IO
     (MutablePrimArray RealWorld Word8,
      MutablePrimArray RealWorld (Ptr Word8))
newRecvBuf Int32
bufSiz Int
bufArrSiz = do
    MutablePrimArray RealWorld Word8
rbuf <- forall (m :: * -> *) a.
(PrimMonad m, Prim a) =>
Int -> m (MutablePrimArray (PrimState m) a)
A.newPinnedPrimArray (forall a b. (Integral a, Num b) => a -> b
fromIntegral Int32
bufsiz' forall a. Num a => a -> a -> a
* Int
bufArrSiz')
    MutablePrimArray RealWorld (Ptr Word8)
rbufArr <- forall (m :: * -> *) a.
(PrimMonad m, Prim a) =>
Int -> m (MutablePrimArray (PrimState m) a)
A.newPinnedPrimArray Int
bufArrSiz'

    -- initialize buffer array with right index
    forall a b. MutablePrimArray RealWorld a -> (Ptr a -> IO b) -> IO b
withMutablePrimArrayContents MutablePrimArray RealWorld Word8
rbuf forall a b. (a -> b) -> a -> b
$ \ Ptr Word8
p ->
        forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ [Int
0..Int
bufArrSiz'forall a. Num a => a -> a -> a
-Int
1] forall a b. (a -> b) -> a -> b
$ \ Int
i -> do
            let bufNPtr :: Ptr Word8
bufNPtr = Ptr Word8
p forall a b. Ptr a -> Int -> Ptr b
`plusPtr` (Int
i forall a. Num a => a -> a -> a
* forall a b. (Integral a, Num b) => a -> b
fromIntegral Int32
bufsiz')
            forall a (m :: * -> *).
(Prim a, PrimMonad m) =>
MutablePrimArray (PrimState m) a -> Int -> a -> m ()
writePrimArray MutablePrimArray RealWorld (Ptr Word8)
rbufArr Int
i Ptr Word8
bufNPtr
    forall (m :: * -> *) a. Monad m => a -> m a
return (MutablePrimArray RealWorld Word8
rbuf, MutablePrimArray RealWorld (Ptr Word8)
rbufArr)
  where
    -- (message size + sockaddr flag + + flag size) + sockaddr_in size + buffer
    -- see diagram above
    bufsiz' :: Int32
bufsiz' = Int32
140 forall a. Num a => a -> a -> a
+ (forall a. Ord a => a -> a -> a
max Int32
0 Int32
bufSiz)
    bufArrSiz' :: Int
bufArrSiz' = forall a. Ord a => a -> a -> a
max Int
1 Int
bufArrSiz

-- | Recv UDP message within a loop
--
-- Loop receiving can be faster since it can reuse receiving buffer. Unlike TCP server
-- from "Z.IO.Network.TCP"server, UDP worker function is called on current haskell thread
-- instead of a forked one, if you have heavy computations to do within the worker function,
-- consider using 'forkBa', or a producer-consumer architecture
--
recvUDPLoop :: HasCallStack
            => UDPRecvConfig
            -> UDP
            -> ((Maybe SocketAddr, Bool, V.Bytes) -> IO a)
            -> IO ()
{-# INLINABLE recvUDPLoop #-}
recvUDPLoop :: forall a.
HasCallStack =>
UDPRecvConfig
-> UDP -> ((Maybe SocketAddr, Bool, Bytes) -> IO a) -> IO ()
recvUDPLoop (UDPRecvConfig Int32
bufSiz Int
bufArrSiz) udp :: UDP
udp@(UDP Ptr UVHandle
hdl Int
slot UVManager
uvm MutablePrimArray RealWorld Word8
_ IORef Bool
_) (Maybe SocketAddr, Bool, Bytes) -> IO a
worker = do
    forall a b c. IO a -> (a -> IO b) -> (a -> IO c) -> IO c
bracket
        (do Ptr UVHandle
check <- forall a. HasCallStack => IO (Ptr a) -> IO (Ptr a)
throwOOMIfNull forall a b. (a -> b) -> a -> b
$ IO (Ptr UVHandle)
hs_uv_check_alloc
            forall a. (HasCallStack, Integral a) => IO a -> IO ()
throwUVIfMinus_ (Ptr UVHandle -> Ptr UVHandle -> IO UDPFlag
hs_uv_check_init Ptr UVHandle
check Ptr UVHandle
hdl)
            forall (m :: * -> *) a. Monad m => a -> m a
return Ptr UVHandle
check)
        Ptr UVHandle -> IO ()
hs_uv_check_close forall a b. (a -> b) -> a -> b
$
        \ Ptr UVHandle
check -> do
            buf :: (MutablePrimArray RealWorld Word8,
 MutablePrimArray RealWorld (Ptr Word8))
buf@(MutablePrimArray RealWorld Word8
_, MutablePrimArray RealWorld (Ptr Word8)
rbufArr) <- Int32
-> Int
-> IO
     (MutablePrimArray RealWorld Word8,
      MutablePrimArray RealWorld (Ptr Word8))
newRecvBuf Int32
bufSiz Int
bufArrSiz
            forall a b. MutablePrimArray RealWorld a -> (Ptr a -> IO b) -> IO b
withMutablePrimArrayContents MutablePrimArray RealWorld (Ptr Word8)
rbufArr forall a b. (a -> b) -> a -> b
$ \ Ptr (Ptr Word8)
p -> do
                UVManager -> Int -> Ptr Word8 -> Int -> IO ()
pokeBufferTable UVManager
uvm Int
slot (forall a b. Ptr a -> Ptr b
castPtr Ptr (Ptr Word8)
p) (Int
bufArrSizforall a. Num a => a -> a -> a
-Int
1)
                -- init uv_check_t must come after poking buffer
                forall a. (HasCallStack, Integral a) => IO a -> IO ()
throwUVIfMinus_ forall a b. (a -> b) -> a -> b
$ Ptr UVHandle -> IO UDPFlag
hs_uv_udp_check_start Ptr UVHandle
check
            forall (f :: * -> *) a b. Applicative f => f a -> f b
forever forall a b. (a -> b) -> a -> b
$ do
                [(Maybe SocketAddr, Bool, Bytes)]
msgs <- HasCallStack =>
UDP
-> (MutablePrimArray RealWorld Word8,
    MutablePrimArray RealWorld (Ptr Word8))
-> Int32
-> IO [(Maybe SocketAddr, Bool, Bytes)]
recvUDPWith UDP
udp (MutablePrimArray RealWorld Word8,
 MutablePrimArray RealWorld (Ptr Word8))
buf Int32
bufSiz
                UVManager -> Int -> Int -> IO ()
pokeBufferSizeTable UVManager
uvm Int
slot (Int
bufArrSizforall a. Num a => a -> a -> a
-Int
1)
                forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ [(Maybe SocketAddr, Bool, Bytes)]
msgs (Maybe SocketAddr, Bool, Bytes) -> IO a
worker

-- | Recv messages from UDP socket, return source address if available, and a `Bool`
-- to indicate if the message is partial (larger than receive buffer size).
--
recvUDP :: HasCallStack => UDPRecvConfig -> UDP -> IO [(Maybe SocketAddr, Bool, V.Bytes)]
{-# INLINABLE recvUDP #-}
recvUDP :: HasCallStack =>
UDPRecvConfig -> UDP -> IO [(Maybe SocketAddr, Bool, Bytes)]
recvUDP (UDPRecvConfig Int32
bufSiz Int
bufArrSiz) udp :: UDP
udp@(UDP Ptr UVHandle
hdl Int
slot UVManager
uvm MutablePrimArray RealWorld Word8
_ IORef Bool
_)  = do
    forall a b c. IO a -> (a -> IO b) -> (a -> IO c) -> IO c
bracket
        (do Ptr UVHandle
check <- forall a. HasCallStack => IO (Ptr a) -> IO (Ptr a)
throwOOMIfNull forall a b. (a -> b) -> a -> b
$ IO (Ptr UVHandle)
hs_uv_check_alloc
            forall a. (HasCallStack, Integral a) => IO a -> IO ()
throwUVIfMinus_ (Ptr UVHandle -> Ptr UVHandle -> IO UDPFlag
hs_uv_check_init Ptr UVHandle
check Ptr UVHandle
hdl)
            forall (m :: * -> *) a. Monad m => a -> m a
return Ptr UVHandle
check)
        Ptr UVHandle -> IO ()
hs_uv_check_close forall a b. (a -> b) -> a -> b
$
        \ Ptr UVHandle
check -> do
            buf :: (MutablePrimArray RealWorld Word8,
 MutablePrimArray RealWorld (Ptr Word8))
buf@(MutablePrimArray RealWorld Word8
_, MutablePrimArray RealWorld (Ptr Word8)
rbufArr) <- Int32
-> Int
-> IO
     (MutablePrimArray RealWorld Word8,
      MutablePrimArray RealWorld (Ptr Word8))
newRecvBuf Int32
bufSiz Int
bufArrSiz
            forall a b. MutablePrimArray RealWorld a -> (Ptr a -> IO b) -> IO b
withMutablePrimArrayContents MutablePrimArray RealWorld (Ptr Word8)
rbufArr forall a b. (a -> b) -> a -> b
$ \ Ptr (Ptr Word8)
p -> do
                UVManager -> Int -> Ptr Word8 -> Int -> IO ()
pokeBufferTable UVManager
uvm Int
slot (forall a b. Ptr a -> Ptr b
castPtr Ptr (Ptr Word8)
p) (Int
bufArrSizforall a. Num a => a -> a -> a
-Int
1)
                -- init uv_check_t must come after poking buffer
                forall a. (HasCallStack, Integral a) => IO a -> IO ()
throwUVIfMinus_ forall a b. (a -> b) -> a -> b
$ Ptr UVHandle -> IO UDPFlag
hs_uv_udp_check_start Ptr UVHandle
check
            HasCallStack =>
UDP
-> (MutablePrimArray RealWorld Word8,
    MutablePrimArray RealWorld (Ptr Word8))
-> Int32
-> IO [(Maybe SocketAddr, Bool, Bytes)]
recvUDPWith UDP
udp (MutablePrimArray RealWorld Word8,
 MutablePrimArray RealWorld (Ptr Word8))
buf Int32
bufSiz

recvUDPWith :: HasCallStack
            => UDP
            -> (A.MutablePrimArray RealWorld Word8, A.MutablePrimArray RealWorld (Ptr Word8))
            -> Int32
            -> IO [(Maybe SocketAddr, Bool, V.Bytes)]
{-# INLINABLE recvUDPWith #-}
recvUDPWith :: HasCallStack =>
UDP
-> (MutablePrimArray RealWorld Word8,
    MutablePrimArray RealWorld (Ptr Word8))
-> Int32
-> IO [(Maybe SocketAddr, Bool, Bytes)]
recvUDPWith udp :: UDP
udp@(UDP Ptr UVHandle
hdl Int
slot UVManager
uvm MutablePrimArray RealWorld Word8
_ IORef Bool
_) (MutablePrimArray RealWorld Word8
rubf, MutablePrimArray RealWorld (Ptr Word8)
rbufArr) Int32
bufSiz =
    -- It's important to keep recv buffer alive, even if we don't directly use it
    forall a. IO a -> IO a
mask_ forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. MutablePrimArray RealWorld a -> (Ptr a -> IO b) -> IO b
withMutablePrimArrayContents MutablePrimArray RealWorld Word8
rubf forall a b. (a -> b) -> a -> b
$ \ Ptr Word8
_ -> do
        HasCallStack => UDP -> IO ()
checkUDPClosed UDP
udp

        Int
bufArrSiz <- forall (m :: * -> *) a.
(PrimMonad m, Prim a) =>
MutablePrimArray (PrimState m) a -> m Int
getSizeofMutablePrimArray MutablePrimArray RealWorld (Ptr Word8)
rbufArr
        -- we have to reset the buffer size, during receiving it'll be overwritten
        forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ [Int
0..Int
bufArrSizforall a. Num a => a -> a -> a
-Int
1] forall a b. (a -> b) -> a -> b
$ \ Int
i -> do
            Ptr Word8
p <- forall a (m :: * -> *).
(Prim a, PrimMonad m) =>
MutablePrimArray (PrimState m) a -> Int -> m a
readPrimArray MutablePrimArray RealWorld (Ptr Word8)
rbufArr Int
i
            forall a. Storable a => Ptr a -> a -> IO ()
poke (forall a b. Ptr a -> Ptr b
castPtr Ptr Word8
p :: Ptr Int32) Int32
bufSiz

        MVar Int
m <- UVManager -> Int -> IO (MVar Int)
getBlockMVar UVManager
uvm Int
slot

        forall a. (HasCallStack, Integral a) => IO a -> IO ()
throwUVIfMinus_ forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. HasCallStack => UVManager -> IO a -> IO a
withUVManager' UVManager
uvm forall a b. (a -> b) -> a -> b
$ do
            -- clean up
            Maybe Int
_ <- forall a. MVar a -> IO (Maybe a)
tryTakeMVar MVar Int
m
            Ptr UVHandle -> IO UDPFlag
hs_uv_udp_recv_start Ptr UVHandle
hdl

        Int
r <- forall a. MVar a -> IO a
takeMVar MVar Int
m forall a b. IO a -> IO b -> IO a
`onException` (do
                -- normally we call 'uv_udp_recv_stop' in C read callback
                -- but when exception raise, here's the place to stop
                -- stop a handle twice will be a libuv error, so we don't check result
                UDPFlag
_ <- forall a. HasCallStack => UVManager -> IO a -> IO a
withUVManager' UVManager
uvm (Ptr UVHandle -> IO UDPFlag
uv_udp_recv_stop Ptr UVHandle
hdl)
                forall (f :: * -> *) a. Functor f => f a -> f ()
void (forall a. MVar a -> IO (Maybe a)
tryTakeMVar MVar Int
m))

        forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM [Int
rforall a. Num a => a -> a -> a
+Int
1..Int
bufArrSizforall a. Num a => a -> a -> a
-Int
1] forall a b. (a -> b) -> a -> b
$ \ Int
i -> do
            Ptr Word8
p        <- forall a (m :: * -> *).
(Prim a, PrimMonad m) =>
MutablePrimArray (PrimState m) a -> Int -> m a
readPrimArray MutablePrimArray RealWorld (Ptr Word8)
rbufArr Int
i
            -- see the buffer struct diagram above
            Int
result   <- forall a. (HasCallStack, Integral a) => IO a -> IO a
throwUVIfMinus (forall a b. (Integral a, Num b) => a -> b
fromIntegral forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall a. Storable a => Ptr a -> IO a
peek @Int32 (forall a b. Ptr a -> Ptr b
castPtr Ptr Word8
p))
            Int32
flag     <- forall a. Storable a => Ptr a -> IO a
peek @Int32 (forall a b. Ptr a -> Ptr b
castPtr (Ptr Word8
p forall a b. Ptr a -> Int -> Ptr b
`plusPtr` Int
4))
            Int32
addrFlag <- forall a. Storable a => Ptr a -> IO a
peek @Int32 (forall a b. Ptr a -> Ptr b
castPtr (Ptr Word8
p forall a b. Ptr a -> Int -> Ptr b
`plusPtr` Int
8))
            !Maybe SocketAddr
addr <- if Int32
addrFlag forall a. Eq a => a -> a -> Bool
== Int32
1
                then forall a. a -> Maybe a
Just forall (m :: * -> *) a b. Monad m => (a -> b) -> m a -> m b
<$!> HasCallStack => Ptr SocketAddr -> IO SocketAddr
peekSocketAddr (forall a b. Ptr a -> Ptr b
castPtr (Ptr Word8
p forall a b. Ptr a -> Int -> Ptr b
`plusPtr` Int
12))
                else forall (m :: * -> *) a. Monad m => a -> m a
return forall a. Maybe a
Nothing
            let !partial :: Bool
partial = Int32
flag forall a. Bits a => a -> a -> a
.&. Int32
UV_UDP_PARTIAL forall a. Eq a => a -> a -> Bool
/= Int32
0
            MutablePrimArray RealWorld Word8
mba <- forall (m :: * -> *) a.
(PrimMonad m, Prim a) =>
Int -> m (MutablePrimArray (PrimState m) a)
A.newPrimArray Int
result
            forall (m :: * -> *) a.
(PrimMonad m, Prim a) =>
MutablePrimArray (PrimState m) a -> Int -> Ptr a -> Int -> m ()
copyPtrToMutablePrimArray MutablePrimArray RealWorld Word8
mba Int
0 (Ptr Word8
p forall a b. Ptr a -> Int -> Ptr b
`plusPtr` Int
140) Int
result
            PrimArray Word8
ba <- forall (m :: * -> *) a.
PrimMonad m =>
MutablePrimArray (PrimState m) a -> m (PrimArray a)
A.unsafeFreezePrimArray MutablePrimArray RealWorld Word8
mba
            forall (m :: * -> *) a. Monad m => a -> m a
return (Maybe SocketAddr
addr, Bool
partial, forall a. PrimArray a -> Int -> Int -> PrimVector a
V.PrimVector PrimArray Word8
ba Int
0 Int
result)