    -- | Simple functions to run UDP clients and servers.
    module Network.Run.UDP (
        runUDPClient,
        runUDPServer,
        runUDPServerFork,
    ) where

    import Control.Concurrent (forkFinally, forkIO)
    import qualified Control.Exception as E
    import Control.Monad (forever, void)
    import Data.ByteString (ByteString)
    import qualified Data.List.NonEmpty as NE
    import Network.Socket
    import Network.Socket.ByteString

    import Network.Run.Core

    -- | Running a UDP client with a socket.
    --   The client action takes a socket and
    --   server's socket address.
    --   They should be used with 'sendTo'.
    runUDPClient :: HostName -> ServiceName -> (Socket -> SockAddr -> IO a) -> IO a
    runUDPClient :: forall a.
    HostName -> HostName -> (Socket -> SockAddr -> IO a) -> IO a
    runUDPClient HostName
    host HostName
    port Socket -> SockAddr -> IO a
    client = do
        addr <- SocketType
    -> Maybe HostName
    -> HostName
    -> [AddrInfoFlag]
    -> (NonEmpty AddrInfo -> AddrInfo)
    -> IO AddrInfo
    resolve SocketType
    Datagram (HostName -> Maybe HostName
    forall a. a -> Maybe a
    Just HostName
    host) HostName
    port [AddrInfoFlag
    AI_ADDRCONFIG] NonEmpty AddrInfo -> AddrInfo
    forall a. NonEmpty a -> a
    NE.head
        let sockAddr = AddrInfo -> SockAddr
    addrAddress AddrInfo
    addr
        E.bracket (openSocket addr) close $ \Socket
    sock -> Socket -> SockAddr -> IO a
    client Socket
    sock SockAddr
    sockAddr

    -- | Running a UDP server with an open socket in a single Haskell thread.
    runUDPServer :: Maybe HostName -> ServiceName -> (Socket -> IO a) -> IO a
    runUDPServer :: forall a. Maybe HostName -> HostName -> (Socket -> IO a) -> IO a
    runUDPServer Maybe HostName
    mhost HostName
    port Socket -> IO a
    server = do
        addr <- SocketType
    -> Maybe HostName
    -> HostName
    -> [AddrInfoFlag]
    -> (NonEmpty AddrInfo -> AddrInfo)
    -> IO AddrInfo
    resolve SocketType
    Datagram Maybe HostName
    mhost HostName
    port [AddrInfoFlag
    AI_PASSIVE] NonEmpty AddrInfo -> AddrInfo
    forall a. NonEmpty a -> a
    NE.head
        E.bracket (openServerSocket addr) close server

    -- | Running a UDP server with a connected socket in each Haskell thread.
    --   The first request is given to the server.
    --   Suppose that the server is serving on __addrS:portS__ and
    --   a client connects to the service from __addrC:portC__.
    --   A connected socket is created by binding to __*:portS__ and
    --   connecting to __addrC:portC__,
    --   resulting in __(UDP,addrS:portS,addrC:portC)__ where
    --   __addrS__ is given magically.
    --   This approach is fragile due to NAT rebidings.
    runUDPServerFork
        :: [HostName] -> ServiceName -> (Socket -> ByteString -> IO ()) -> IO ()
    runUDPServerFork :: [HostName] -> HostName -> (Socket -> ByteString -> IO ()) -> IO ()
    runUDPServerFork [] HostName
    _ Socket -> ByteString -> IO ()
    _ = () -> IO ()
    forall a. a -> IO a
    forall (m :: * -> *) a. Monad m => a -> m a
    return ()
    runUDPServerFork (HostName
    h : [HostName]
    hs) HostName
    port Socket -> ByteString -> IO ()
    server = do
        (HostName -> IO ThreadId) -> [HostName] -> IO ()
    forall (t :: * -> *) (m :: * -> *) a b.
    (Foldable t, Monad m) =>
    (a -> m b) -> t a -> m ()
    mapM_ (IO () -> IO ThreadId
    forkIO (IO () -> IO ThreadId)
    -> (HostName -> IO ()) -> HostName -> IO ThreadId
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . HostName -> IO ()
    forall {b}. HostName -> IO b
    run) [HostName]
    hs
        HostName -> IO ()
    forall {b}. HostName -> IO b
    run HostName
    h
      where
        run :: HostName -> IO b
    run HostName
    host = do
            HostName -> IO ()
    labelMe (HostName -> IO ()) -> HostName -> IO ()
    forall a b. (a -> b) -> a -> b
    $ HostName
    "UDP server for " HostName -> HostName -> HostName
    forall a. [a] -> [a] -> [a]
    ++ HostName
    h
            Maybe HostName -> HostName -> (Socket -> IO b) -> IO b
    forall a. Maybe HostName -> HostName -> (Socket -> IO a) -> IO a
    runUDPServer (HostName -> Maybe HostName
    forall a. a -> Maybe a
    Just HostName
    host) HostName
    port ((Socket -> IO b) -> IO b) -> (Socket -> IO b) -> IO b
    forall a b. (a -> b) -> a -> b
    $ \Socket
    lsock -> IO () -> IO b
    forall (f :: * -> *) a b. Applicative f => f a -> f b
    forever (IO () -> IO b) -> IO () -> IO b
    forall a b. (a -> b) -> a -> b
    $ do
                (bs0, peeraddr) <- Socket -> Int -> IO (ByteString, SockAddr)
    recvFrom Socket
    lsock Int
    2048
                let family = case SockAddr
    peeraddr of
                        SockAddrInet{} -> Family
    AF_INET
                        SockAddrInet6{} -> Family
    AF_INET6
                        SockAddr
    _ -> HostName -> Family
    forall a. HasCallStack => HostName -> a
    error HostName
    "family"
                    hints =
                        AddrInfo
    defaultHints
                            { addrSocketType = Datagram
                            , addrFamily = family
                            , addrFlags = [AI_PASSIVE]
                            }
                addr <- NE.head <$> getAddrInfo (Just hints) Nothing (Just port)
                s <- openServerSocket addr
                connect s peeraddr
                void $ forkFinally (labelMe "UDP server" >> server s bs0) (\Either SomeException ()
    _ -> Socket -> IO ()
    close Socket
    s)